// SPDX-License-Identifier: Apache-2.0
//
// vgpud -- the virtio-gpu driver, outside the kernel.
//
// docs/plan/hybrid.md's first eviction. The same device the kernel drives in
// kernel/drivers/gpu/virtio_gpu/, driven instead by a process that asks the
// kernel for four things and nothing else:
//
//   PCIIOC_CLAIM   take the device (the kernel driver is asked to let go)
//   PCIIOC_CFG     read config space, to find the vendor capabilities
//   mmap(BAR)      reach the registers
//   PCIIOC_DMA     memory the device can be pointed at
//
// The point is not that this is better than the in-kernel driver -- it is the
// same code in a worse language for the job. The point is that it is a real
// driver rather than a test program, so the interface has something to be
// judged by, and what it asks for is what a foreign driver would ask for too.
//
// Note what is *not* here: no syscall knows this is a GPU. The kernel handed
// over a PCI device and this decided it was a GPU, which is the property that
// makes the interface general rather than a hole shaped like one thing.

#include "types.h"
#include "user.h"
#include "fcntl.h"
#include "posix.h"
#include "pciuser.h"
#include "virtio.h"

// --- small output, since lib/ has no printf ---------------------------------

static void
puts(const char *s)
{
    int n = 0;
    while (s[n]) n++;
    write(1, (void *)s, n);
}

static void
putn(long long v)
{
    char b[24];
    int i = 0;
    if (v < 0) { write(1, "-", 1); v = -v; }
    if (v == 0) { write(1, "0", 1); return; }
    while (v > 0) { b[i++] = '0' + (int)(v % 10); v /= 10; }
    while (i-- > 0) write(1, &b[i], 1);
}

static void
puthex(unsigned long long v, int digits)
{
    static const char h[] = "0123456789abcdef";
    char b[17];
    for (int i = digits - 1; i >= 0; i--) { b[i] = h[v & 0xF]; v >>= 4; }
    write(1, b, digits);
}

static void
die(const char *why)
{
    puts("vgpud: "); puts(why); puts("\n");
    exit();
}

// --- register access --------------------------------------------------------

static inline unsigned char  r8 (volatile unsigned char *p, int o) { return p[o]; }
static inline void           w8 (volatile unsigned char *p, int o, unsigned char v) { p[o] = v; }
static inline unsigned short r16(volatile unsigned char *p, int o) { return *(volatile unsigned short *)(p + o); }
static inline void           w16(volatile unsigned char *p, int o, unsigned short v) { *(volatile unsigned short *)(p + o) = v; }
static inline unsigned int   r32(volatile unsigned char *p, int o) { return *(volatile unsigned int *)(p + o); }
static inline void           w32(volatile unsigned char *p, int o, unsigned int v) { *(volatile unsigned int *)(p + o) = v; }
static inline void           w64(volatile unsigned char *p, int o, unsigned long long v)
{
    // Two 32-bit writes, low half first. The virtio spec requires it and the
    // device latches on the high write; a single 64-bit store is a different
    // bus transaction and some implementations simply drop it.
    w32(p, o, (unsigned int)v);
    w32(p, o + 4, (unsigned int)(v >> 32));
}

// --- the device -------------------------------------------------------------

static int pcifd;
static int index_;

static volatile unsigned char *common;
static volatile unsigned char *notify;
static volatile unsigned char *isr;
static unsigned int notify_mult;

static struct pciuser_info info;

static unsigned int
cfg32(unsigned int off)
{
    struct pciuser_cfg c;
    c.index = (unsigned int)index_;
    c.offset = off & ~3u;
    c.value = 0;
    if (ioctl(pcifd, PCIIOC_CFG, &c) < 0)
        die("config read failed");
    return c.value;
}

static unsigned char
cfg8(unsigned int off)
{
    return (unsigned char)(cfg32(off) >> ((off & 3) * 8));
}

// Find a virtio configuration structure.
//
// A virtio 1.0 device does not put its registers at a fixed offset. It walks
// you there: a vendor capability in config space names which BAR, and how far
// into it, each structure lives. So a driver that cannot read config space
// cannot find its own registers -- which is why PCIIOC_CFG exists, and why the
// test program that came before this never needed it.
static volatile unsigned char *
find_cfg(int want, unsigned int *out_mult)
{
    unsigned char off = cfg8(0x34) & 0xFC;

    while (off) {
        unsigned char id   = cfg8(off);
        unsigned char next = cfg8(off + 1);

        if (id == PCI_CAP_ID_VNDR) {
            unsigned char type = cfg8(off + 3);
            unsigned char bar  = cfg8(off + 4);
            unsigned int  o    = cfg32(off + 8);
            unsigned int  len  = cfg32(off + 12);

            if (type == want && bar < PCIUSER_MAXBAR && info.bar[bar].base) {
                if (out_mult)
                    *out_mult = cfg32(off + 16);

                // Map from the start of the BAR, not from the capability's
                // offset: two capabilities usually live in one BAR, and
                // mapping each separately would hand out two mappings of
                // overlapping pages.
                unsigned long long want_len = (unsigned long long)o + len;
                if (want_len < 4096) want_len = 4096;
                if (want_len > info.bar[bar].size) want_len = info.bar[bar].size;

                volatile unsigned char *p =
                    (volatile unsigned char *)mmap(0, want_len,
                                                   PROT_READ | PROT_WRITE,
                                                   MAP_SHARED, pcifd,
                                                   info.bar[bar].base);
                if (p == 0 || (long)p == -1)
                    return 0;
                return p + o;
            }
        }
        off = next & 0xFC;
    }
    return 0;
}

// --- the queue --------------------------------------------------------------

#define CTRL_QUEUE 0

static struct vring_desc  *desc;
static struct vring_avail *avail;
static struct vring_used  *used;
static unsigned long long  ring_phys;
static unsigned short      qsize, notify_off, last_used, next_desc;

static void      *cmdbuf,  *respbuf;
static unsigned long long cmdphys, respphys;

static void
queue_setup(void)
{
    w16(common, VCOMMON_Q_SELECT, CTRL_QUEUE);
    qsize = r16(common, VCOMMON_Q_SIZE);
    if (qsize == 0)
        die("control queue has no slots");
    if (qsize > 64)
        qsize = 64;                 // more than one command in flight is not
                                    // something this driver has any use for
    w16(common, VCOMMON_Q_SIZE, qsize);
    notify_off = r16(common, VCOMMON_Q_NOTIFY_OFF);

    // One contiguous block for the three rings, laid out the way the spec
    // requires: descriptors, then available, then used, each aligned. The
    // device reads all three by physical address, so they have to come from
    // the kernel's DMA arena rather than from anything this process could
    // allocate for itself.
    unsigned int desc_sz  = 16 * qsize;
    unsigned int avail_sz = 6 + 2 * qsize;

    struct pciuser_dma d;
    d.size = 16384;
    if (ioctl(pcifd, PCIIOC_DMA, &d) < 0)
        die("no DMA memory for the queue");

    unsigned long long base  = d.uaddr;
    unsigned long long basep = d.paddr;

    unsigned int off_avail = (desc_sz + 1) & ~1u;
    unsigned int off_used  = (off_avail + avail_sz + 3) & ~3u;

    desc  = (struct vring_desc  *)(base);
    avail = (struct vring_avail *)(base + off_avail);
    used  = (struct vring_used  *)(base + off_used);
    ring_phys = basep;

    for (unsigned int i = 0; i < 16384; i++)
        ((volatile unsigned char *)base)[i] = 0;

    // Do not interrupt us. docmd() polls the used ring, so an interrupt is
    // not merely unwanted, it is dangerous: this process holds the device but
    // has no handler registered, and the line is shared with the e1000. A
    // level-triggered line nobody acknowledges re-fires for ever -- 77,000
    // times a second, which is what the first working run of this driver did
    // before this line existed. The storm detector in kernel/dev/irq.c named
    // the line and the neighbour it was drowning.
    avail->flags = VRING_AVAIL_F_NO_INTERRUPT;
    __sync_synchronize();

    w64(common, VCOMMON_Q_DESC,   basep);
    w64(common, VCOMMON_Q_DRIVER, basep + off_avail);
    w64(common, VCOMMON_Q_DEVICE, basep + off_used);
    w16(common, VCOMMON_Q_ENABLE, 1);

    // Command and response buffers, also from the arena for the same reason.
    struct pciuser_dma b;
    b.size = 8192;
    if (ioctl(pcifd, PCIIOC_DMA, &b) < 0)
        die("no DMA memory for command buffers");
    cmdbuf   = (void *)b.uaddr;
    cmdphys  = b.paddr;
    respbuf  = (void *)(b.uaddr + 4096);
    respphys = b.paddr + 4096;
}

// Submit one command and wait for the answer.
//
// Polled, not interrupt-driven. This driver issues a handful of commands at
// startup and one flush per frame, and a poll on a used-ring index that the
// device updates by DMA is both simpler and lower-latency than waking a
// process. The interrupt path exists (PCIIOC_IRQ) and belongs in whatever
// wants to sleep between frames.
static int
docmd(unsigned int cmdlen, unsigned int resplen)
{
    unsigned short head = next_desc;

    desc[head].addr  = cmdphys;
    desc[head].len   = cmdlen;
    desc[head].flags = VRING_DESC_F_NEXT;
    desc[head].next  = (unsigned short)(head + 1);

    desc[head + 1].addr  = respphys;
    desc[head + 1].len   = resplen;
    desc[head + 1].flags = VRING_DESC_F_WRITE;
    desc[head + 1].next  = 0;

    avail->ring[avail->idx % qsize] = head;
    __sync_synchronize();
    avail->idx++;
    __sync_synchronize();

    *(volatile unsigned short *)(notify + notify_off * notify_mult) = CTRL_QUEUE;

    // Spin with a bound. A device that never answers is a bug somewhere, and
    // hanging for ever is a worse way to report it than saying so.
    for (int spin = 0; spin < 100000000; spin++) {
        __sync_synchronize();
        if (used->idx != last_used) {
            last_used = used->idx;
            next_desc = (unsigned short)((next_desc + 2) % qsize);
            return 0;
        }
    }
    return -1;
}

static void
hdr_init(void *p, unsigned int type)
{
    struct virtio_gpu_ctrl_hdr *h = (struct virtio_gpu_ctrl_hdr *)p;
    for (unsigned int i = 0; i < sizeof(*h); i++) ((unsigned char *)h)[i] = 0;
    h->type = type;
}

static int
resp_ok(void)
{
    struct virtio_gpu_ctrl_hdr *h = (struct virtio_gpu_ctrl_hdr *)respbuf;
    return h->type == VIRTIO_GPU_RESP_OK_NODATA ||
           h->type == VIRTIO_GPU_RESP_OK_DISPLAY_INFO;
}

// --- bringing the device up -------------------------------------------------

static int width, height;
static unsigned long long fb_uaddr, fb_paddr;
static unsigned int fb_bytes;

#define RESOURCE_ID 1

static void
bring_up(void)
{
    // Reset, then the status handshake in the order the spec sets out. A
    // device that finds the sequence wrong fails the feature step rather than
    // telling you which line was at fault, so it is written out plainly.
    w8(common, VCOMMON_STATUS, 0);
    while (r8(common, VCOMMON_STATUS) != 0)
        ;
    w8(common, VCOMMON_STATUS, VSTATUS_ACK);
    w8(common, VCOMMON_STATUS, VSTATUS_ACK | VSTATUS_DRIVER);

    w32(common, VCOMMON_DFEAT_SEL, 1);
    unsigned int hi = r32(common, VCOMMON_DFEAT);
    if (!(hi & (1u << (VIRTIO_F_VERSION_1 - 32)))) {
        w8(common, VCOMMON_STATUS, VSTATUS_FAILED);
        die("device is not virtio 1.0");
    }

    w32(common, VCOMMON_GFEAT_SEL, 0);
    w32(common, VCOMMON_GFEAT, 0);
    w32(common, VCOMMON_GFEAT_SEL, 1);
    w32(common, VCOMMON_GFEAT, 1u << (VIRTIO_F_VERSION_1 - 32));

    w8(common, VCOMMON_STATUS, VSTATUS_ACK | VSTATUS_DRIVER | VSTATUS_FEATURES_OK);
    if (!(r8(common, VCOMMON_STATUS) & VSTATUS_FEATURES_OK)) {
        w8(common, VCOMMON_STATUS, VSTATUS_FAILED);
        die("device refused the feature set");
    }

    queue_setup();

    w8(common, VCOMMON_STATUS,
       VSTATUS_ACK | VSTATUS_DRIVER | VSTATUS_FEATURES_OK | VSTATUS_DRIVER_OK);

    // What size is the screen?
    hdr_init(cmdbuf, VIRTIO_GPU_CMD_GET_DISPLAY_INFO);
    if (docmd(sizeof(struct virtio_gpu_ctrl_hdr),
              sizeof(struct virtio_gpu_resp_display_info)) < 0)
        die("the device never answered GET_DISPLAY_INFO");
    if (!resp_ok())
        die("GET_DISPLAY_INFO refused");

    struct virtio_gpu_resp_display_info *di =
        (struct virtio_gpu_resp_display_info *)respbuf;
    width  = (int)di->pmodes[0].r.width;
    height = (int)di->pmodes[0].r.height;
    if (width <= 0 || height <= 0)
        die("the device reports no display");

    puts("vgpud: display is "); putn(width); puts("x"); putn(height); puts("\n");

    // A 2D resource the size of the screen.
    struct virtio_gpu_resource_create_2d *c =
        (struct virtio_gpu_resource_create_2d *)cmdbuf;
    hdr_init(c, VIRTIO_GPU_CMD_RESOURCE_CREATE_2D);
    c->resource_id = RESOURCE_ID;
    c->format      = VIRTIO_GPU_FORMAT_B8G8R8X8_UNORM;
    c->width       = (unsigned int)width;
    c->height      = (unsigned int)height;
    if (docmd(sizeof(*c), sizeof(struct virtio_gpu_ctrl_hdr)) < 0 || !resp_ok())
        die("RESOURCE_CREATE_2D refused");

    // Memory behind it. This is the framebuffer: the process writes here and
    // the device reads it, which is the entire reason PCIIOC_DMA reports a
    // physical address.
    fb_bytes = (unsigned int)(width * height * 4);

    struct pciuser_dma d;
    d.size = fb_bytes;
    if (ioctl(pcifd, PCIIOC_DMA, &d) < 0) {
        puts("vgpud: wanted "); putn(fb_bytes / 1024);
        puts(" KB of DMA memory for the framebuffer\n");
        die("not enough DMA memory -- see DMA_ARENA_BYTES");
    }
    fb_uaddr = d.uaddr;
    fb_paddr = d.paddr;

    struct virtio_gpu_resource_attach_backing *a =
        (struct virtio_gpu_resource_attach_backing *)cmdbuf;
    hdr_init(a, VIRTIO_GPU_CMD_RESOURCE_ATTACH_BACKING);
    a->resource_id = RESOURCE_ID;
    a->nr_entries  = 1;
    // The entries follow the struct rather than being in it: the count is
    // variable, so the command is a header with an array glued on the end.
    struct virtio_gpu_mem_entry *e = (struct virtio_gpu_mem_entry *)(a + 1);
    e->addr    = fb_paddr;
    e->length  = fb_bytes;
    e->padding = 0;
    if (docmd(sizeof(*a) + sizeof(*e), sizeof(struct virtio_gpu_ctrl_hdr)) < 0 ||
        !resp_ok())
        die("RESOURCE_ATTACH_BACKING refused");

    // Put it on the screen.
    struct virtio_gpu_set_scanout *so = (struct virtio_gpu_set_scanout *)cmdbuf;
    hdr_init(so, VIRTIO_GPU_CMD_SET_SCANOUT);
    so->r.x = so->r.y = 0;
    so->r.width  = (unsigned int)width;
    so->r.height = (unsigned int)height;
    so->scanout_id  = 0;
    so->resource_id = RESOURCE_ID;
    if (docmd(sizeof(*so), sizeof(struct virtio_gpu_ctrl_hdr)) < 0 || !resp_ok())
        die("SET_SCANOUT refused");

    puts("vgpud: scanout is ours, framebuffer at physical 0x");
    puthex(fb_paddr, 16); puts("\n");
}

// Push what has been drawn to the host. Two commands, and both are needed:
// the transfer copies from the backing into the device's copy of the resource,
// the flush tells it to put that on screen.
static void
present(void)
{
    struct virtio_gpu_transfer_to_host_2d *t =
        (struct virtio_gpu_transfer_to_host_2d *)cmdbuf;
    hdr_init(t, VIRTIO_GPU_CMD_TRANSFER_TO_HOST_2D);
    t->r.x = t->r.y = 0;
    t->r.width  = (unsigned int)width;
    t->r.height = (unsigned int)height;
    t->offset      = 0;
    t->resource_id = RESOURCE_ID;
    t->padding     = 0;
    if (docmd(sizeof(*t), sizeof(struct virtio_gpu_ctrl_hdr)) < 0)
        die("TRANSFER_TO_HOST_2D never answered");

    struct virtio_gpu_resource_flush *f =
        (struct virtio_gpu_resource_flush *)cmdbuf;
    hdr_init(f, VIRTIO_GPU_CMD_RESOURCE_FLUSH);
    f->r.x = f->r.y = 0;
    f->r.width  = (unsigned int)width;
    f->r.height = (unsigned int)height;
    f->resource_id = RESOURCE_ID;
    f->padding     = 0;
    if (docmd(sizeof(*f), sizeof(struct virtio_gpu_ctrl_hdr)) < 0)
        die("RESOURCE_FLUSH never answered");
}

// --- something to look at ---------------------------------------------------

// Not a compositor. A pattern that is obviously deliberate, so a screenshot
// answers "did a userspace process put this here" without ambiguity -- a black
// screen and a working screen look identical if all you draw is black.
static void
draw_test_pattern(void)
{
    volatile unsigned int *fb = (volatile unsigned int *)fb_uaddr;

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            unsigned int c;
            int band = (y * 8) / height;
            switch (band) {
            case 0: c = 0x0C0E13; break;                     /* SCH_BG */
            case 1: c = 0x7C5CFF; break;                     /* SCH_ACCENT */
            case 2: c = 0x4B1D9A; break;                     /* PSOGM purple */
            case 3: c = (unsigned int)((x * 255 / width) * 0x010101); break;
            case 4: c = (unsigned int)(x * 255 / width) << 16; break;
            case 5: c = (unsigned int)(x * 255 / width) << 8;  break;
            case 6: c = (unsigned int)(x * 255 / width);       break;
            default:
                // A checkerboard, so a stride bug is visible as a shear rather
                // than as a plausible-looking gradient.
                c = (((x / 16) + (y / 16)) & 1) ? 0xE6E8EC : 0x0C0E13;
                break;
            }
            fb[y * width + x] = c;
        }
    }
}

// --- main -------------------------------------------------------------------

int
main(int argc, char *argv[])
{
    pcifd = open("/dev/pci", O_RDWR);
    if (pcifd < 0)
        die("cannot open /dev/pci");

    int count = 0;
    if (ioctl(pcifd, PCIIOC_COUNT, &count) < 0)
        die("PCIIOC_COUNT failed");

    // Find it by what it is, not by where it is. A driver that takes a bus
    // address on the command line is one nobody can start automatically.
    index_ = -1;
    for (int i = 0; i < count; i++) {
        info.index = (unsigned int)i;
        if (ioctl(pcifd, PCIIOC_INFO, &info) < 0)
            continue;
        if (info.vendor == VIRTIO_VENDOR && info.device == VIRTIO_DEV_GPU) {
            index_ = i;
            break;
        }
    }
    if (index_ < 0)
        die("no virtio-gpu on this machine");

    puts("vgpud: virtio-gpu at "); puthex(info.bus, 2); puts(":");
    puthex(info.slot, 2); puts("."); puthex(info.func, 1);
    if (info.driver[0]) { puts(", currently held by "); puts(info.driver); }
    puts("\n");

    // The kernel driver is asked to let go here. It puts the VBE framebuffer
    // back on screen before returning, so there is never a moment with no
    // display -- see vgpu_release() and vga_fb_restore().
    if (ioctl(pcifd, PCIIOC_CLAIM, &index_) < 0)
        die("claim refused -- the kernel driver would not release it");
    puts("vgpud: claimed\n");

    // The info is re-read because claiming is what enables memory space and
    // bus mastering, and a BAR read before that can legitimately be zero.
    info.index = (unsigned int)index_;
    if (ioctl(pcifd, PCIIOC_INFO, &info) < 0)
        die("PCIIOC_INFO failed after claim");

    common = find_cfg(VIRTIO_PCI_CAP_COMMON_CFG, 0);
    notify = find_cfg(VIRTIO_PCI_CAP_NOTIFY_CFG, &notify_mult);
    isr    = find_cfg(VIRTIO_PCI_CAP_ISR_CFG, 0);
    if (!common || !notify)
        die("no modern virtio capabilities -- not a 1.0 device");

    bring_up();
    draw_test_pattern();
    present();

    puts("vgpud: drew a test pattern from userspace\n");

    if (argc > 1 && argv[1][0] == '-' && argv[1][1] == 'x') {
        // Exit without releasing, to show what a crashed driver leaves behind:
        // the kernel takes the device back and the screen is whatever the
        // device was last told to scan out.
        puts("vgpud: exiting without cleaning up (simulated crash)\n");
        exit();
    }

    // Nothing to serve yet -- the compositor is what will connect here, over
    // mechanism 1, which does not exist. Holding the device is the honest
    // thing to do until then: releasing it would put the console back and make
    // it look like this never worked.
    for (;;)
        sleep(100);
}