// SPDX-License-Identifier: Apache-2.0
//
// SchFS space manager: which blocks are in use.
//
// There is no free-space map on disk, and that is the interesting decision
// here. The committed tree already records exactly which blocks are live —
// every tree node is named by its parent's pointer, every data block by an
// extent record — so the map is *derived* at mount by walking the tree once,
// held in RAM, and never written back.
//
// The alternative, persisting it, buys a fast mount and costs the one problem
// copy-on-write is otherwise free of: the map has to be updated atomically with
// the tree it describes, and writing the map needs blocks, which changes the
// map. ZFS spends real complexity on that recursion and btrfs resolves it with
// a fixpoint loop. Deriving the map instead cannot disagree with the filesystem
// no matter what a crash interrupts, because it is computed from it — there is
// no second source of truth to get out of step.
//
// What it costs is a metadata scan at mount, proportional to the number of tree
// blocks rather than to the size of the volume. Nothing is written, so it is
// also safe on a read-only mount. If mount time ever becomes the thing that
// hurts, the map can be cached on disk later purely as a hint, validated
// against the uberblock's generation and discarded when it does not match — an
// addition, not a format change.
//
// ---------------------------------------------------------------------------
// Deferred frees
//
// A block released during the transaction group being built is still reachable
// from the last *committed* root, which is where a crash lands. Handing it
// straight back would let the filesystem overwrite the very copy it is supposed
// to fall back to. So frees go into a second bitmap and only rejoin the free
// map when the commit that supersedes them lands. One transaction of delay, no
// more: once group N is published, group N-1's tree is no longer anybody's
// fallback.

#include "schfs_int.h"

#ifdef SCHFS_HOSTED
#include <string.h>
#else
#include "types.h"
#include "defs.h"
#endif

/* ------------------------------------------------------------------ */
/* Bit vectors, paged                                                  */
/*                                                                     */
/* The bitmap has to be one flat array to be indexable, and the kernel */
/* hands out memory a page at a time with no promise that two pages    */
/* are adjacent. So it is a vector of page pointers instead: one page  */
/* covers SCHFS_SPACE_BITS_PER_PAGE blocks, and the page vector itself */
/* is a fixed-size array, which is what puts a hard ceiling on volume  */
/* size. See docs/architecture/schfs.md for the resulting limits.      */
/* ------------------------------------------------------------------ */

static int
bit_test(uint64_t *const *pages, uint64_t b)
{
    uint64_t pg  = b / SCHFS_SPACE_BITS_PER_PAGE;
    uint64_t off = b % SCHFS_SPACE_BITS_PER_PAGE;
    return (int)((pages[pg][off / 64] >> (off % 64)) & 1);
}

static void
bit_set(uint64_t *const *pages, uint64_t b)
{
    uint64_t pg  = b / SCHFS_SPACE_BITS_PER_PAGE;
    uint64_t off = b % SCHFS_SPACE_BITS_PER_PAGE;
    pages[pg][off / 64] |= 1ULL << (off % 64);
}

// The whole 64-block word containing b, so a scan can skip a fully used region
// sixty-four blocks at a time instead of one. On a nearly full volume that is
// the difference between an allocation costing microseconds and milliseconds.
static uint64_t
bit_word(uint64_t *const *pages, uint64_t b)
{
    uint64_t pg  = b / SCHFS_SPACE_BITS_PER_PAGE;
    uint64_t off = b % SCHFS_SPACE_BITS_PER_PAGE;
    return pages[pg][off / 64];
}

static int
popcount64(uint64_t v)
{
    int n = 0;
    while (v) { v &= v - 1; n++; }
    return n;
}

/* ------------------------------------------------------------------ */
/* Setup                                                               */
/* ------------------------------------------------------------------ */

int
schfs_space_init(struct schfs_space *s, uint64_t nblocks,
                 void *(*palloc)(void))
{
    memset(s, 0, sizeof(*s));

    if (nblocks <= SCHFS_FIRST_BLOCK)
        return -3;
    if (nblocks > SCHFS_MAX_BLOCKS)
        return -4;                   /* larger than this build can map */

    uint64_t npages = (nblocks + SCHFS_SPACE_BITS_PER_PAGE - 1)
                    / SCHFS_SPACE_BITS_PER_PAGE;

    for (uint64_t i = 0; i < npages; i++) {
        void *p = palloc();
        if (p == 0)
            return -5;               /* caller frees what was taken */
        memset(p, 0, SCHFS_BLKSIZE);
        s->inuse[i] = (uint64_t *)p;

        p = palloc();
        if (p == 0)
            return -5;
        memset(p, 0, SCHFS_BLKSIZE);
        s->pend[i] = (uint64_t *)p;
        s->npages = (int)(i + 1);
    }

    s->nblocks     = nblocks;
    s->free_blocks = nblocks - SCHFS_FIRST_BLOCK;
    s->hint        = SCHFS_FIRST_BLOCK;

    // The boot block and the uberblock ring are never allocatable, and neither
    // is anything past the end of the volume: the last bitmap page covers more
    // blocks than exist, and a stray allocation out there would be written to a
    // sector the device does not have.
    for (uint64_t b = 0; b < SCHFS_FIRST_BLOCK; b++)
        bit_set(s->inuse, b);
    for (uint64_t b = nblocks; b < (uint64_t)npages * SCHFS_SPACE_BITS_PER_PAGE; b++)
        bit_set(s->inuse, b);

    return 0;
}

void
schfs_space_mark(struct schfs_space *s, uint64_t blk, uint64_t n)
{
    for (uint64_t i = 0; i < n; i++) {
        uint64_t b = blk + i;
        if (b < SCHFS_FIRST_BLOCK || b >= s->nblocks)
            continue;
        if (bit_test(s->inuse, b))
            continue;                /* already counted; extents may be shared */
        bit_set(s->inuse, b);
        s->free_blocks--;
    }
}

/* ------------------------------------------------------------------ */
/* Allocation                                                          */
/* ------------------------------------------------------------------ */

uint64_t
schfs_space_alloc_run(struct schfs_space *s, uint64_t n)
{
    if (n == 0 || n > s->free_blocks)
        return 0;

    // Two sweeps: from the rotating hint to the end of the volume, then from
    // the start back to it. Starting every search at block 0 would make the
    // front of the volume the only place anything is ever written and turn
    // allocation into a linear scan of everything already used.
    for (int pass = 0; pass < 2; pass++) {
        uint64_t lo = pass == 0 ? s->hint : SCHFS_FIRST_BLOCK;
        uint64_t hi = pass == 0 ? s->nblocks : s->hint;
        uint64_t run = 0, first = 0;

        for (uint64_t b = lo; b < hi; ) {
            if (b % 64 == 0 && bit_word(s->inuse, b) == ~0ULL) {
                run = 0;
                b += 64;
                continue;
            }
            if (bit_test(s->inuse, b)) {
                run = 0;
                b++;
                continue;
            }
            if (run == 0)
                first = b;
            run++;
            b++;
            if (run == n) {
                for (uint64_t i = 0; i < n; i++)
                    bit_set(s->inuse, first + i);
                s->free_blocks -= n;
                s->hint = first + n < s->nblocks ? first + n : SCHFS_FIRST_BLOCK;
                return first;
            }
        }
    }
    return 0;
}

uint64_t
schfs_space_alloc(struct schfs_space *s)
{
    return schfs_space_alloc_run(s, 1);
}

void
schfs_space_free(struct schfs_space *s, uint64_t blk)
{
    // A snapshot exists, so this block may still be reachable from it.
    //
    // Copy-on-write frees the old block whenever one is rewritten, and that is
    // exactly the block a snapshot is holding. Freeing it would hand it out
    // again and overwrite the data the snapshot exists to keep -- silently,
    // with the damage only visible later when the snapshot is read.
    //
    // Knowing *which* freed blocks a snapshot still needs requires reference
    // counting per extent, which is a real design and not a flag. Until that
    // exists, the safe answer is to free nothing while a snapshot is present:
    // the free map is derived at mount anyway, so everything genuinely
    // unreachable is reclaimed on the next mount rather than lost. Conservative
    // in the only direction that is safe -- it can under-report free space, and
    // can never hand out a block that is still in use.
    if (s->flags & SCHFS_SPACE_RETAIN_ALL)
        return;

    if (blk < SCHFS_FIRST_BLOCK || blk >= s->nblocks)
        return;
    if (!bit_test(s->inuse, blk))
        return;                      /* never allocated, or freed twice */
    if (bit_test(s->pend, blk))
        return;                      /* already queued this group */
    bit_set(s->pend, blk);
    s->pending++;
}

void
schfs_space_commit(struct schfs_space *s)
{
    uint64_t words = SCHFS_SPACE_BITS_PER_PAGE / 64;
    for (int p = 0; p < s->npages; p++) {
        uint64_t *pend  = s->pend[p];
        uint64_t *inuse = s->inuse[p];
        for (uint64_t w = 0; w < words; w++) {
            if (pend[w] == 0)
                continue;
            s->free_blocks += popcount64(pend[w] & inuse[w]);
            inuse[w] &= ~pend[w];
            pend[w] = 0;
        }
    }
    s->pending = 0;
}

/* ------------------------------------------------------------------ */
/* Deriving the map from the committed tree                            */
/* ------------------------------------------------------------------ */

// Every extent record in a leaf names a run of data blocks. A range shared by
// several files (a reflink, a snapshot) appears more than once, which is why
// schfs_space_mark is idempotent rather than counting.
static void
mark_leaf_extents(struct schfs_space *s, const void *blk)
{
    const struct schfs_item *it = SCHFS_ITEMS(blk);
    int n = SCHFS_HDR(blk)->nritems;

    for (int i = 0; i < n; i++) {
        if (it[i].key.type != SCHFS_KEY_EXTENT)
            continue;
        if (it[i].size != sizeof(struct schfs_extent))
            continue;
        if ((int)it[i].offset + (int)it[i].size > SCHFS_BLKSIZE)
            continue;
        const struct schfs_extent *e =
            (const struct schfs_extent *)((const char *)blk + it[i].offset);
        schfs_space_mark(s, e->disk_block, e->nblocks);
    }
}

// Snapshot roots found while walking, so they can be walked in turn.
//
// Collected rather than descended into immediately: the walk keeps one block
// buffer and a path, and re-entering it from inside itself would need a second
// of each.
#define SNAP_ROOTS_MAX 64

struct snap_roots {
    int      n;
    int      overflow;
    uint64_t nr[SNAP_ROOTS_MAX];
    uint32_t crc[SNAP_ROOTS_MAX];
};

static void
collect_snap_roots(struct snap_roots *sr, const void *blk)
{
    const struct schfs_item *it = SCHFS_ITEMS(blk);
    int n = SCHFS_HDR(blk)->nritems;

    for (int i = 0; i < n; i++) {
        if (it[i].key.type != SCHFS_KEY_SNAP)
            continue;
        if (it[i].size != sizeof(struct schfs_snap))
            continue;
        if ((int)it[i].offset + (int)it[i].size > SCHFS_BLKSIZE)
            continue;
        const struct schfs_snap *sn =
            (const struct schfs_snap *)((const char *)blk + it[i].offset);
        if (sr->n >= SNAP_ROOTS_MAX) { sr->overflow = 1; continue; }
        sr->nr[sr->n]  = sn->root;
        sr->crc[sr->n] = sn->root_crc;
        sr->n++;
    }
}

// Walk the whole tree, marking every block it reaches.
//
// Iterative, with one block buffer and a path of block numbers, for the same
// reason everything else here is: a recursive walk keeps a block per level and
// the kernel stack is one block in total. Climbing re-reads the parent, so an
// internal node is read once per child — mount-time cost that the buffer cache
// largely absorbs, in exchange for bounded memory.
// Walk one root, marking everything it reaches. Snapshot records encountered
// on the way are collected into `sr` when it is non-null.
static int
scan_root(struct schfs_vol *v, uint64_t root, uint32_t root_crc,
          struct snap_roots *sr)
{
    struct schfs_space *s = &v->space;
    void *blk = v->ctx.a;

    struct {
        uint64_t nr;
        uint32_t crc;
        int      idx;                /* next child to descend into */
    } path[SCHFS_MAX_LEVEL + 1];

    int d = 0;
    path[0].nr  = root;
    path[0].crc = root_crc;
    path[0].idx = 0;
    schfs_space_mark(s, root, 1);

    for (;;) {
        if (schfs_node_read(&v->dev, path[d].nr, path[d].crc, blk) != 0)
            return -2;

        if (SCHFS_HDR(blk)->level == 0) {
            mark_leaf_extents(s, blk);
            if (sr)
                collect_snap_roots(sr, blk);
            d--;
        } else if (path[d].idx < SCHFS_HDR(blk)->nritems) {
            struct schfs_kptr *kp = &SCHFS_KPTRS(blk)[path[d].idx];
            path[d].idx++;
            if (d + 1 > SCHFS_MAX_LEVEL)
                return -2;
            d++;
            path[d].nr  = kp->blocknr;
            path[d].crc = kp->crc;
            path[d].idx = 0;
            schfs_space_mark(s, path[d].nr, 1);
            continue;
        } else {
            d--;                     /* every child visited */
        }

        if (d < 0)
            return 0;
    }
}

// Walk the live tree, then every snapshot's tree.
//
// The snapshots have to be walked or the derivation is actively dangerous: a
// block reachable only from a snapshot would look free, be handed out, and
// overwrite the very data the snapshot exists to keep. Nothing would complain
// at the time -- the corruption surfaces later, when the snapshot is read.
//
// The live tree is walked first because that is where the snapshot records
// live; a snapshot's own tree may contain older snapshot records, which are
// deliberately not followed. They name roots the live tree no longer keeps,
// and a snapshot that was deleted should not be resurrected by an older
// snapshot's memory of it.
int
schfs_space_scan(struct schfs_vol *v)
{
    struct snap_roots sr;
    sr.n = 0;
    sr.overflow = 0;

    int r = scan_root(v, v->tree.root, v->tree.root_crc, &sr);
    if (r != 0)
        return r;

    if (sr.n > 0)
        v->space.flags |= SCHFS_SPACE_RETAIN_ALL;

    for (int i = 0; i < sr.n; i++) {
        r = scan_root(v, sr.nr[i], sr.crc[i], 0);
        if (r != 0)
            return r;
    }

    // Refusing is the only safe answer. Carrying on would mean deriving a free
    // map that considers a live snapshot's blocks available.
    if (sr.overflow)
        return -2;

    return 0;
}