// Memory-backed block device.
//
// Serves an entire SCHFS image out of RAM, so the system can boot and run a
// shell with the physical disk untouched. That is what makes destructive work
// on the real disk survivable: reformatting, running mkfs, or bringing up a new
// on-disk format cannot take away the root you are running from. It is also the
// first piece of the installer -- an install environment is exactly "a root in
// RAM plus a disk you are allowed to destroy".
//
// The image is linked into the kernel as a binary blob (see the ramdisk_obj
// rule in the Makefile) and is served in place: the blob lives in writable
// kernel memory, so writes to the ramdisk simply modify it. They are lost on
// reboot, which is the expected behaviour for this kind of root.
//
// Enabled by building with RAMROOT=1. Without it this file still compiles, the
// device reports zero size, and the kernel boots from the IDE disk as before.
//
// The device is larger than the image. The image is only what the root holds,
// with a little slack; past its end the device goes on to half of the memory
// free at mount time, and those blocks are pages taken from kalloc the first
// time each is written -- a block nobody has written reads as zeros and costs
// nothing. SchFS grows the volume to the device's size when it mounts (see
// grow_to in kernel/fs/schfs/vol.c), so `apk add` on a live system has as
// much room as the machine has memory to give, rather than whatever empty
// space was built into the ISO. The ISO used to carry 60 MB of zeros for it.
//
// Pages are not given back when SchFS frees a block: there is no discard.
// A freed block is reused by the next write, so the cost is the high-water
// mark, not a leak.
//
// **A compressed image.** The blob can also be SCFF (include/fs/scff.h): the
// image cut into chunks, each compressed with zstd. Then no block is served in
// place. The page directory covers the image too, and a block with no page
// has not been unpacked yet: reading it decompresses its whole chunk, and
// every block of that chunk nobody has written gets a page of its own. The
// kernel and the ISO are the compressed size -- the 360 MB root with Xorg in
// it is 104 MB -- and the root costs memory only for what has been read or
// written. A block of zeros points at one shared zero page and costs nothing.

#include "core/types.h"
#include "core/defs.h"
#include "core/param.h"
#include "arch/memlayout.h"
#include "arch/mmu.h"
#include "core/spinlock.h"
#include "core/sleeplock.h"
#include "fs/fs.h"
#include "fs/buf.h"
#include "fs/scff.h"
#include "../scff/scff_zstd.h"

#ifdef RAMDISK_ROOT
// Supplied by `ld -r -b binary` over build/fs-ram.img, renamed by objcopy.
extern uchar _binary_ramdisk_start[];
extern uchar _binary_ramdisk_size[];
#endif

static uchar *rd_base;
static uint64 rd_size;     // bytes
static uint   rd_blocks;   // blocks in the image
static uint   rd_first;    // the first block with a page: rd_blocks, or 0 for SCFF

// SCFF only. rd_zlock is held across a chunk's decompression, and makes
// rd_chunk -- the one chunk-sized buffer -- its holder's.
static struct scff_header *rd_scff;
static uint64 *rd_index;
static struct sleeplock rd_zlock;
static uchar rd_chunk[SCFF_CHUNK_MAX] __attribute__((aligned(PGSIZE)));
static uchar rd_zero[BSIZE] __attribute__((aligned(PGSIZE)));
static uint64 rd_unpacked;  // chunks decompressed

// rd_dir[i] is a page of pointers, entry j the page holding block
// rd_first + i*RD_PTRS + j, or 0 for one never written. 32768 pages of
// 512 pointers reach 64 GB, which is as far as half of any machine this runs
// on is likely to go, and costs 256 KB of pointers here.
#define RD_PTRS   (PGSIZE / sizeof(uchar *))
#define RD_DIRS   32768
static uchar **rd_dir[RD_DIRS];
static uint64 rd_total;     // blocks, image and the rest; 0 until sized
static uint64 rd_pages;     // pages taken for blocks past the image
static struct spinlock rd_lock;

void
ramdiskinit(void)
{
  initlock(&rd_lock, "ramdisk");
#ifdef RAMDISK_ROOT
  rd_base  = _binary_ramdisk_start;
  rd_size  = (uint64)_binary_ramdisk_size;
  rd_blocks = (uint)(rd_size / BSIZE);
  rd_first = rd_blocks;

  if(rd_size >= sizeof(struct scff_header) && memcmp(rd_base, SCFF_MAGIC, 4) == 0){
    struct scff_header *h = (struct scff_header *)rd_base;
    uint64 chunks = h->chunk_bytes ? (h->image_bytes + h->chunk_bytes - 1) / h->chunk_bytes : 0;
    if(h->version != SCFF_VERSION || h->codec != SCFF_CODEC_ZSTD ||
       h->chunk_bytes < BSIZE || h->chunk_bytes > SCFF_CHUNK_MAX ||
       (h->chunk_bytes & (h->chunk_bytes - 1)) || h->image_bytes % BSIZE ||
       chunks != h->nchunks || h->stored_bytes > rd_size ||
       h->index_off + (h->nchunks + 1) * sizeof(uint64) > rd_size)
      panic("ramdiskinit: the SCFF header does not describe this image");
    rd_index = (uint64 *)(rd_base + h->index_off);
    for(uint i = 0; i < h->nchunks; i++)
      if(rd_index[i] > rd_index[i + 1] || rd_index[i + 1] > h->stored_bytes ||
         rd_index[i + 1] - rd_index[i] > h->chunk_bytes)
        panic("ramdiskinit: the SCFF index is damaged");
    if(scff_zstd_init() < 0)
      panic("ramdiskinit: no room for the zstd decoder");
    initsleeplock(&rd_zlock, "scff");
    rd_scff = h;
    rd_blocks = (uint)(h->image_bytes / BSIZE);
    rd_first = 0;
    kinfo("ramdisk: SCFF, %u MB in %u MB, %u chunks of %u KB\n",
            (uint)(h->image_bytes >> 20), (uint)(rd_size >> 20),
            h->nchunks, h->chunk_bytes >> 10);
  }

  if(rd_blocks == 0)
    panic("ramdiskinit: embedded image is empty");

  // A truncated image is worth catching here rather than as a confusing
  // "invalid file system" panic later: the superblock lives in block 1.
  if(rd_blocks < 2)
    panic("ramdiskinit: embedded image smaller than a superblock");

  if(!rd_scff)
    kinfo("ramdisk: %u blocks (%u KB) at %p\n",
            rd_blocks, (uint)(rd_size / 1024), rd_base);
#else
  rd_base   = 0;
  rd_size   = 0;
  rd_blocks = 0;
#endif
}

int
ramdisk_present(void)
{
  return rd_blocks != 0;
}

// The device's size in blocks, fixed the first time anyone asks -- which is
// SchFS mounting the root, after every page has reached the allocator. Asked
// earlier, during ramdiskinit, the free count would be the first few
// megabytes kinit1 handed over and nothing more.
uint64
ramdisk_blocks(void)
{
  if (rd_blocks == 0)
    return 0;
  acquire(&rd_lock);
  if (rd_total == 0) {
    uint64 extra = kalloc_free_bytes() / 2 / BSIZE;
    uint64 room = (uint64)RD_DIRS * RD_PTRS - (rd_blocks - rd_first);
    if (extra > room)
      extra = room;
    rd_total = rd_blocks + extra;
    kinfo("ramdisk: %u MB image, %u MB of memory to grow into\n",
            (uint)(rd_size >> 20), (uint)(extra * BSIZE >> 20));
  }
  release(&rd_lock);
  return rd_total;
}

// The page behind a block, made when `make` says so. One that is only the
// shared zero page gets a page of its own before anything writes to it.
static uchar *
rd_page(uint64 blk, int make)
{
  uint64 i = (blk - rd_first) / RD_PTRS, j = (blk - rd_first) % RD_PTRS;
  uchar *pg = 0;

  acquire(&rd_lock);
  if (rd_dir[i] == 0 && make) {
    if ((rd_dir[i] = (uchar **)kalloc()) != 0)
      memset(rd_dir[i], 0, PGSIZE);
  }
  if (rd_dir[i]) {
    pg = rd_dir[i][j];
    if ((pg == 0 || pg == rd_zero) && make && (pg = (uchar *)kalloc()) != 0) {
      if (rd_dir[i][j] == rd_zero)
        memset(pg, 0, BSIZE);
      rd_dir[i][j] = pg;
      rd_pages++;
    }
  }
  release(&rd_lock);
  if (make && pg == 0)
    panic("ramdisk: out of memory for the live root");
  return pg;
}

// Give every block of SCFF chunk c that has no page one: its contents, or the
// zero page. Blocks already there -- written since boot, or unpacked by
// whoever held the lock before us -- are left as they are.
static void
rd_unpack(uint c)
{
  uint per = rd_scff->chunk_bytes / BSIZE;
  uint64 first = (uint64)c * per;
  uint n = rd_blocks - first < per ? (uint)(rd_blocks - first) : per;
  uint64 len = rd_index[c + 1] - rd_index[c];
  uchar *src = rd_base + rd_index[c];
  uint64 want = (uint64)n * BSIZE;
  uchar *data = rd_chunk;

  acquiresleep(&rd_zlock);
  if (len == 0) {
    data = 0;                              // all zeros
  } else if (len == want) {
    data = src;                            // stored as it is
  } else {
    long got = scff_zstd_decompress(rd_chunk, sizeof(rd_chunk), src, len);
    if (got != (long)want)
      panic("ramdisk: an SCFF chunk does not decompress");
    rd_unpacked++;
  }

  for (uint k = 0; k < n; k++) {
    uint64 blk = first + k, i = blk / RD_PTRS, j = blk % RD_PTRS;
    uchar *b = data ? data + (uint64)k * BSIZE : 0;
    int zero = 1;
    if (b)
      for (uint w = 0; w < BSIZE / 8 && zero; w++)
        zero = ((uint64 *)b)[w] == 0;

    uchar *pg = zero ? rd_zero : (uchar *)kalloc();
    if (pg == 0)
      panic("ramdisk: out of memory for the live root");
    if (!zero)
      memmove(pg, b, BSIZE);

    acquire(&rd_lock);
    if (rd_dir[i] == 0 && (rd_dir[i] = (uchar **)kalloc()) != 0)
      memset(rd_dir[i], 0, PGSIZE);
    if (rd_dir[i] == 0)
      panic("ramdisk: out of memory for the live root");
    if (rd_dir[i][j] == 0) {
      rd_dir[i][j] = pg;
      if (!zero)
        rd_pages++;
      pg = 0;
    }
    release(&rd_lock);
    if (pg && pg != rd_zero)
      kfree((char *)pg);                   // someone was there first
  }
  releasesleep(&rd_zlock);
}

// Sync buf with the in-memory image. Mirrors iderw()'s contract: on a read,
// fill b->data and set B_VALID; on a write, copy out and clear B_DIRTY.
void
ramdiskrw(struct buf *b)
{
  if(!holdingsleep(&b->lock))
    panic("ramdiskrw: buf not locked");
  if((b->flags & (B_VALID|B_DIRTY)) == B_VALID)
    panic("ramdiskrw: nothing to do");
  if(rd_base == 0)
    panic("ramdiskrw: no ramdisk");
  if(b->blockno >= rd_first){
    if(b->blockno >= ramdisk_blocks())
      panic("ramdiskrw: block out of range");
    uchar *pg = rd_page(b->blockno, (b->flags & B_DIRTY) != 0);
    if(pg == 0 && rd_scff && b->blockno < rd_blocks && !(b->flags & B_DIRTY)){
      rd_unpack((uint)(b->blockno / (rd_scff->chunk_bytes / BSIZE)));
      pg = rd_page(b->blockno, 0);
    }
    if(b->flags & B_DIRTY){
      memmove(pg, b->data, BSIZE);
      b->flags &= ~B_DIRTY;
    } else if(pg){
      memmove(b->data, pg, BSIZE);
    } else {
      memset(b->data, 0, BSIZE);    // never written
    }
    b->flags |= B_VALID;
    return;
  }

  uchar *p = rd_base + (uint64)b->blockno * BSIZE;

  if(b->flags & B_DIRTY){
    memmove(p, b->data, BSIZE);
    b->flags &= ~B_DIRTY;
  } else {
    memmove(b->data, p, BSIZE);
  }
  b->flags |= B_VALID;
}