// SPDX-License-Identifier: GPL-2.0-or-later
//
// NVMe: PCIe SSDs, written from the NVM Express Base Specification 1.4.
//
// The internal disk of nearly every laptop made since about 2016, and of most
// desktops: an M.2 stick that is not SATA at all, but a PCI device with its
// own command set. There is no ATA here and no SCSI.
//
// The model is two rings in memory per queue: a submission queue the driver
// writes 64-byte commands into, and a completion queue the controller writes
// 16-byte results into. A doorbell register tells the controller how far the
// submission queue has been filled; another tells it how far the completion
// queue has been read. Queue 0 is the admin queue, which creates the others.
//
// Like the AHCI driver beside it: one I/O queue, one command at a time, and
// **completion by polling**. A completion queue is created with interrupts
// disabled, so nothing here depends on interrupt routing -- which on real
// hardware without an AML interpreter is the part that cannot be trusted.
// A flash read completes in tens of microseconds, so the spin is short, and
// it is yielded once a wait goes long.

#include "types.h"
#include "defs.h"
#include "param.h"
#include "x86.h"
#include "spinlock.h"
#include "sleeplock.h"
#include "fs.h"
#include "file.h"
#include "proc.h"
#include "pci.h"
#include "dma.h"
#include "blkdev.h"

// --- controller registers (section 3.1) --------------------------------------

#define NVME_CAP    0x00               // 64 bits
#define   CAP_MQES(c)   ((uint32)((c) & 0xFFFF))         // max entries, 0-based
#define   CAP_TO(c)     ((uint32)(((c) >> 24) & 0xFF))   // ready timeout, 500 ms
#define   CAP_DSTRD(c)  ((uint32)(((c) >> 32) & 0xF))    // doorbell stride
#define   CAP_MPSMIN(c) ((uint32)(((c) >> 48) & 0xF))    // min page, 4K << n
#define NVME_VS     0x08
#define NVME_INTMS  0x0C
#define NVME_CC     0x14
#define   CC_EN         (1u << 0)
#define   CC_IOSQES     (6u << 16)      // 64-byte submission entries
#define   CC_IOCQES     (4u << 20)      // 16-byte completion entries
#define NVME_CSTS   0x1C
#define   CSTS_RDY      (1u << 0)
#define   CSTS_CFS      (1u << 1)       // controller fatal status
#define NVME_AQA    0x24
#define NVME_ASQ    0x28
#define NVME_ACQ    0x30
#define NVME_DB     0x1000

// --- commands ----------------------------------------------------------------

#define ADM_CREATE_SQ   0x01
#define ADM_CREATE_CQ   0x05
#define ADM_IDENTIFY    0x06
#define NVM_FLUSH       0x00
#define NVM_WRITE       0x01
#define NVM_READ        0x02

#define CNS_NAMESPACE   0
#define CNS_CONTROLLER  1

struct sqe {
  uint32 cdw0;                    // opcode 7:0, command ID 31:16
  uint32 nsid;
  uint64 rsvd;
  uint64 mptr;
  uint64 prp1, prp2;
  uint32 cdw10, cdw11, cdw12, cdw13, cdw14, cdw15;
};

struct cqe {
  uint32 result;
  uint32 rsvd;
  uint16 sqhd, sqid;
  uint16 cid;
  uint16 status;                  // bit 0 is the phase tag
};

#define QDEPTH 32                 // entries per queue, admin and I/O alike

struct queue {
  struct sqe *sq;
  volatile struct cqe *cq;
  int    qid;
  int    tail, head;
  int    phase;                   // what a new completion's phase bit will be
};

#define NVME_MAX_CTRL  4
#define NVME_MAX_NS    8

struct nvme {
  volatile uint8 *r;
  uint32 dstrd;
  uint16 cid;
  struct queue admin, io;
  uint8 *xfer;                    // bounce buffer, BSIZE, one page
  struct sleeplock lock;          // one command at a time
  int    index;
};

struct nvme_ns {
  struct nvme *c;
  uint32 nsid;
  uint64 nsectors;
  uint32 secsize;
  char   name[12];
};

static struct nvme ctrls[NVME_MAX_CTRL];
static struct nvme_ns nss[NVME_MAX_NS];
static int nctrl, nns;

static inline uint32 rd(volatile uint8 *b, int o)        { return *(volatile uint32 *)(b + o); }
static inline void   wr(volatile uint8 *b, int o, uint32 v) { *(volatile uint32 *)(b + o) = v; }

// 64-bit registers as two halves, low first: the specification allows a
// controller to support only 32-bit accesses.
static uint64
rd64(volatile uint8 *b, int o)
{
  return rd(b, o) | (uint64)rd(b, o + 4) << 32;
}

static void
wr64(volatile uint8 *b, int o, uint64 v)
{
  wr(b, o, (uint32)v);
  wr(b, o + 4, (uint32)(v >> 32));
}

static int
sq_doorbell(struct nvme *c, int qid)
{
  return NVME_DB + (2 * qid) * (4 << c->dstrd);
}

static int
cq_doorbell(struct nvme *c, int qid)
{
  return NVME_DB + (2 * qid + 1) * (4 << c->dstrd);
}

static int
queue_init(struct queue *q, int qid)
{
  q->sq = dma_alloc(QDEPTH * sizeof(struct sqe), 4096);
  q->cq = dma_alloc(QDEPTH * sizeof(struct cqe), 4096);
  if (!q->sq || !q->cq)
    return -1;
  memset(q->sq, 0, QDEPTH * sizeof(struct sqe));
  memset((void *)q->cq, 0, QDEPTH * sizeof(struct cqe));
  q->qid = qid;
  q->tail = q->head = 0;
  q->phase = 1;
  return 0;
}

// Submit one command and wait for its completion. Returns 0, or -1 with the
// status printed.
static int
submit(struct nvme *c, struct queue *q, struct sqe *cmd, uint32 *result)
{
  uint16 cid = ++c->cid;

  cmd->cdw0 = (cmd->cdw0 & 0xFFFF) | ((uint32)cid << 16);
  q->sq[q->tail] = *cmd;
  q->tail = (q->tail + 1) % QDEPTH;
  __sync_synchronize();
  wr(c->r, sq_doorbell(c, q->qid), q->tail);

  // Spin for the first few milliseconds, which is when nearly every command
  // finishes; after that give the processor away between looks. yield() only
  // from a process -- IDENTIFY and the root probe run before there is one.
  uint64 start = rdtsc(), limit = tsc_hz * 10;     // ten seconds
  volatile struct cqe *e = &q->cq[q->head];
  while ((e->status & 1) != q->phase) {
    uint64 now = rdtsc() - start;
    if (now > limit || (rd(c->r, NVME_CSTS) & CSTS_CFS)) {
      cprintf("nvme%d: command %x %s\n", c->index, cmd->cdw0 & 0xFF,
              now > limit ? "timed out" : "failed: controller fatal");
      return -1;
    }
    if (now > tsc_hz / 200 && myproc())
      yield();
    else
      asm volatile("pause");
  }
  __sync_synchronize();

  uint16 status = e->status >> 1;
  uint16 ecid = e->cid;
  if (result)
    *result = e->result;

  if (++q->head == QDEPTH) {
    q->head = 0;
    q->phase ^= 1;
  }
  wr(c->r, cq_doorbell(c, q->qid), q->head);

  if (ecid != cid) {
    cprintf("nvme%d: completion for command %d, expected %d\n", c->index,
            ecid, cid);
    return -1;
  }
  if (status) {
    // Status code type in 10:8, status code in 7:0.
    cprintf("nvme%d: command %x failed, type %d code %x\n", c->index,
            cmd->cdw0 & 0xFF, (status >> 8) & 7, status & 0xFF);
    return -1;
  }
  return 0;
}

static int
identify(struct nvme *c, int cns, uint32 nsid)
{
  struct sqe cmd;

  memset(&cmd, 0, sizeof(cmd));
  cmd.cdw0  = ADM_IDENTIFY;
  cmd.nsid  = nsid;
  cmd.prp1  = dma_phys(c->xfer);
  cmd.cdw10 = cns;
  return submit(c, &c->admin, &cmd, 0);
}

// --- block device ------------------------------------------------------------

static int
nvme_rw(void *priv, uint64 lba, uint32 count, void *buf, int write)
{
  struct nvme_ns *ns = priv;
  struct nvme *c = ns->c;
  uint32 bytes = count * ns->secsize;
  struct sqe cmd;
  int rc;

  if (bytes > BSIZE || count == 0)
    return -1;

  // The bounce buffer is one page and page-aligned, so one PRP entry covers
  // it and PRP2 is never needed.
  memset(&cmd, 0, sizeof(cmd));
  cmd.cdw0  = write ? NVM_WRITE : NVM_READ;
  cmd.nsid  = ns->nsid;
  cmd.prp1  = dma_phys(c->xfer);
  cmd.cdw10 = (uint32)lba;
  cmd.cdw11 = (uint32)(lba >> 32);
  cmd.cdw12 = count - 1;              // a 0-based count

  acquiresleep(&c->lock);
  if (write)
    memmove(c->xfer, buf, bytes);
  rc = submit(c, &c->io, &cmd, 0);
  if (rc == 0 && !write)
    memmove(buf, c->xfer, bytes);
  releasesleep(&c->lock);
  return rc;
}

static int
nvme_flush(void *priv)
{
  struct nvme_ns *ns = priv;
  struct sqe cmd;

  memset(&cmd, 0, sizeof(cmd));
  cmd.cdw0 = NVM_FLUSH;
  cmd.nsid = ns->nsid;
  acquiresleep(&ns->c->lock);
  int rc = submit(ns->c, &ns->c->io, &cmd, 0);
  releasesleep(&ns->c->lock);
  return rc;
}

static const struct blkdev_ops nvme_ops = {
  .rw    = nvme_rw,
  .flush = nvme_flush,
};

// --- bring-up ----------------------------------------------------------------

static int
wait_ready(struct nvme *c, int want, uint32 to)
{
  // CAP.TO is the worst case the controller promises, in 500 ms units.
  for (uint32 i = 0; i < (to ? to : 1) * 500; i++) {
    uint32 csts = rd(c->r, NVME_CSTS);
    if (csts == 0xFFFFFFFF)
      return -1;                       // the device is not answering at all
    if (((csts & CSTS_RDY) != 0) == want)
      return 0;
    microdelay(1000);
  }
  return -1;
}

// An ASCII field from an identify page: fixed width, space padded.
static void
ident_string(const uint8 *p, int len, char *out, int cap)
{
  int n = 0;
  for (int i = 0; i < len && n < cap - 1; i++)
    out[n++] = p[i] >= 0x20 && p[i] < 0x7F ? p[i] : ' ';
  while (n > 0 && out[n - 1] == ' ')
    n--;
  out[n] = 0;
}

static void
add_namespace(struct nvme *c, uint32 nsid)
{
  if (nns >= NVME_MAX_NS)
    return;
  if (identify(c, CNS_NAMESPACE, nsid) < 0)
    return;

  uint8 *id = c->xfer;
  uint64 nsze = *(uint64 *)id;
  if (nsze == 0)
    return;                            // allocated but not active
  uint32 lbaf = *(uint32 *)(id + 128 + 4 * (id[26] & 0xF));
  uint32 lbads = (lbaf >> 16) & 0xFF;
  if (lbads < 9 || lbads > 12) {
    cprintf("nvme%d: namespace %d has %d-byte sectors; not used\n",
            c->index, nsid, 1 << lbads);
    return;
  }
  if (lbaf & 0xFFFF) {
    // Metadata interleaved with the data would land in our bounce buffer.
    cprintf("nvme%d: namespace %d carries metadata; not used\n",
            c->index, nsid);
    return;
  }

  struct nvme_ns *ns = &nss[nns++];
  ns->c = c;
  ns->nsid = nsid;
  ns->nsectors = nsze;
  ns->secsize = 1u << lbads;

  // nvme0n1: controller 0, namespace 1 -- the name Linux gives the same disk.
  char *p = ns->name;
  *p++ = 'n'; *p++ = 'v'; *p++ = 'm'; *p++ = 'e';
  *p++ = '0' + c->index;
  *p++ = 'n';
  if (nsid >= 10)
    *p++ = '0' + nsid / 10 % 10;
  *p++ = '0' + nsid % 10;
  *p = 0;

  cprintf("nvme: %s: %d MB, %d-byte sectors\n", ns->name,
          (int)((nsze << lbads) >> 20), ns->secsize);
  blkdev_register(ns->name, &nvme_ops, ns, ns->nsectors, ns->secsize);
}

static int
nvme_probe(struct pci_dev *d)
{
  if (d->prog_if != 0x02)
    return -1;                         // not the NVM Express interface
  if (nctrl >= NVME_MAX_CTRL)
    return -1;
  if (!d->bar[0].base || d->bar[0].is_io) {
    cprintf("nvme: no memory BAR; not driving this controller\n");
    return -1;
  }

  struct nvme *c = &ctrls[nctrl];
  c->index = nctrl;

  pci_enable(d, PCI_CMD_MEMORY | PCI_CMD_BUS_MASTER);
  c->r = ioremap(d->bar[0].base, d->bar[0].size ? d->bar[0].size : 0x2000);
  if (!c->r)
    return -1;

  uint64 cap = rd64(c->r, NVME_CAP);
  uint32 vs = rd(c->r, NVME_VS);
  c->dstrd = CAP_DSTRD(cap);
  if (CAP_MPSMIN(cap) != 0) {
    cprintf("nvme: controller wants pages over 4 KB; not driving it\n");
    return -1;
  }
  if (CAP_MQES(cap) + 1 < QDEPTH) {
    cprintf("nvme: queues of %d entries are too small\n", CAP_MQES(cap) + 1);
    return -1;
  }

  // The firmware booted from this controller and left it running with its
  // own queues. Disabling it is a controller reset: everything the firmware
  // set up is forgotten, and ours replace it.
  wr(c->r, NVME_CC, rd(c->r, NVME_CC) & ~CC_EN);
  if (wait_ready(c, 0, CAP_TO(cap)) < 0) {
    cprintf("nvme: controller would not stop\n");
    return -1;
  }

  c->xfer = dma_alloc(BSIZE, 4096);
  if (!c->xfer || queue_init(&c->admin, 0) < 0 || queue_init(&c->io, 1) < 0)
    return -1;
  initsleeplock(&c->lock, "nvme");

  wr(c->r, NVME_AQA, ((QDEPTH - 1) << 16) | (QDEPTH - 1));
  wr64(c->r, NVME_ASQ, dma_phys(c->admin.sq));
  wr64(c->r, NVME_ACQ, dma_phys((void *)c->admin.cq));
  wr(c->r, NVME_INTMS, 0xFFFFFFFF);    // polled; see the top of the file
  wr(c->r, NVME_CC, CC_EN | CC_IOSQES | CC_IOCQES);
  if (wait_ready(c, 1, CAP_TO(cap)) < 0) {
    cprintf("nvme: controller would not start (csts %x)\n",
            rd(c->r, NVME_CSTS));
    return -1;
  }

  if (identify(c, CNS_CONTROLLER, 0) < 0)
    return -1;
  char model[41];
  ident_string(c->xfer + 24, 40, model, sizeof(model));
  uint32 nn = *(uint32 *)(c->xfer + 516);
  kinfo("nvme%d: %s, NVMe %d.%d, %d namespace%s\n", c->index, model,
        vs >> 16, (vs >> 8) & 0xFF, nn, nn == 1 ? "" : "s");

  // The I/O queue pair: completion queue first, since the submission queue
  // names it. Physically contiguous (PC = 1), interrupts off (IEN = 0).
  struct sqe cmd;
  memset(&cmd, 0, sizeof(cmd));
  cmd.cdw0  = ADM_CREATE_CQ;
  cmd.prp1  = dma_phys((void *)c->io.cq);
  cmd.cdw10 = ((QDEPTH - 1) << 16) | 1;
  cmd.cdw11 = 1;
  if (submit(c, &c->admin, &cmd, 0) < 0)
    return -1;

  memset(&cmd, 0, sizeof(cmd));
  cmd.cdw0  = ADM_CREATE_SQ;
  cmd.prp1  = dma_phys(c->io.sq);
  cmd.cdw10 = ((QDEPTH - 1) << 16) | 1;
  cmd.cdw11 = (1u << 16) | 1;          // completions to queue 1; contiguous
  if (submit(c, &c->admin, &cmd, 0) < 0)
    return -1;

  nctrl++;

  // Namespace IDs are 1-based and need not be dense; a consumer SSD has
  // exactly one. Looking at the first few is enough for anything that is
  // not an enterprise array.
  for (uint32 nsid = 1; nsid <= nn && nsid <= NVME_MAX_NS; nsid++)
    add_namespace(c, nsid);
  return 0;
}

static const struct pci_match nvme_match[] = {
  { PCI_ANY, PCI_ANY, 0x01, 0x08 },     // mass storage, non-volatile memory
  { 0, 0, 0, 0 },
};

const struct pci_driver nvme_driver = {
  .name    = "nvme",
  .devname = "nvme",
  .match   = nvme_match,
  .probe   = nvme_probe,
};