// SPDX-License-Identifier: GPL-2.0-or-later
//
// 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 indexable, and the kernel hands out memory a   */
/* page at a time with no promise that two pages are adjacent. So it   */
/* is pages of bits found through pages of pointers -- see struct      */
/* schfs_bits -- and a missing page reads as zeros. See                */
/* docs/architecture/schfs.md for the resulting limits.                */
/* ------------------------------------------------------------------ */

static uint64_t *
bits_page(const struct schfs_bits *v, uint64_t pg)
{
    uint64_t **d = v->dir[pg / SCHFS_SPACE_PTRS_PER_PAGE];
    return d ? d[pg % SCHFS_SPACE_PTRS_PER_PAGE] : 0;
}

// The page for `pg`, taken from `palloc` (zeroed) if it is not there yet.
// 0 when there is no memory for it.
static uint64_t *
bits_page_get(struct schfs_bits *v, uint64_t pg, void *(*palloc)(void))
{
    uint64_t ***d = &v->dir[pg / SCHFS_SPACE_PTRS_PER_PAGE];
    if (!*d) {
        if (!palloc || !(*d = (uint64_t **)palloc()))
            return 0;
        memset(*d, 0, SCHFS_BLKSIZE);
    }
    uint64_t **slot = &(*d)[pg % SCHFS_SPACE_PTRS_PER_PAGE];
    if (!*slot) {
        if (!(*slot = (uint64_t *)palloc()))
            return 0;
        memset(*slot, 0, SCHFS_BLKSIZE);
    }
    return *slot;
}

// Give every page back, pointer pages included.
static void
bits_free(struct schfs_bits *v, void (*pfree)(void *))
{
    for (int i = 0; i < SCHFS_SPACE_DIRS; i++) {
        if (!v->dir[i])
            continue;
        if (pfree) {
            for (int j = 0; j < (int)SCHFS_SPACE_PTRS_PER_PAGE; j++)
                if (v->dir[i][j])
                    pfree(v->dir[i][j]);
            pfree(v->dir[i]);
        }
        v->dir[i] = 0;
    }
}

static int
bit_test(const struct schfs_bits *v, uint64_t b)
{
    uint64_t *pg = bits_page(v, b / SCHFS_SPACE_BITS_PER_PAGE);
    uint64_t off = b % SCHFS_SPACE_BITS_PER_PAGE;
    return pg ? (int)((pg[off / 64] >> (off % 64)) & 1) : 0;
}

// Only on a page that exists: the inuse map has all of its pages from init,
// and the pend map takes one in schfs_space_free before it sets anything.
static void
bit_set(const struct schfs_bits *v, uint64_t b)
{
    uint64_t *pg = bits_page(v, b / SCHFS_SPACE_BITS_PER_PAGE);
    uint64_t off = b % SCHFS_SPACE_BITS_PER_PAGE;
    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(const struct schfs_bits *v, uint64_t b)
{
    uint64_t *pg = bits_page(v, b / SCHFS_SPACE_BITS_PER_PAGE);
    uint64_t off = b % SCHFS_SPACE_BITS_PER_PAGE;
    return pg ? pg[off / 64] : 0;
}

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

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

// Bring an already-sized map back to "empty volume": every block free except
// the ones that are never allocatable. Split out of schfs_space_init because
// mounting may have to try more than one uberblock, and each candidate root
// needs a scan starting from a clean map -- while the pages themselves were
// taken once and must not be taken again.
void
schfs_space_reset(struct schfs_space *s, uint64_t nblocks)
{
    for (uint64_t i = 0; i < s->npages; i++)
        memset(bits_page(&s->inuse, i), 0, SCHFS_BLKSIZE);
    bits_free(&s->pend, s->pfree);

    s->nblocks     = nblocks;
    s->free_blocks = nblocks - SCHFS_FIRST_BLOCK;
    s->pending     = 0;
    // 1% of the volume, at least 1 MB and at most 64 MB. A tree update touches
    // a path of at most SCHFS_MAX_LEVEL nodes per record, and a group rewrites
    // a node once however often it changes, so this is generous; btrfs keeps
    // between 16 MB and 512 MB for the same job.
    s->reserve = nblocks / 100;
    if (s->reserve < 256)
        s->reserve = 256;
    if (s->reserve > 16384)
        s->reserve = 16384;
    s->hint        = SCHFS_FIRST_BLOCK;
    // Cleared, not preserved: schfs_space_scan sets RETAIN_ALL when it finds a
    // snapshot, so carrying one candidate root's answer into the next root's
    // scan would pin blocks for a snapshot that tree may not even have.
    s->flags       = 0;
    s->snap_gen    = 0;

    // 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 < s->npages * SCHFS_SPACE_BITS_PER_PAGE; b++)
        bit_set(&s->inuse, b);
}

void
schfs_space_release(struct schfs_space *s)
{
    bits_free(&s->inuse, s->pfree);
    bits_free(&s->pend, s->pfree);
    // Zeroed rather than merely emptied, so calling this twice is harmless and
    // a released map cannot be mistaken for a sized one.
    memset(s, 0, sizeof(*s));
}

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

    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;

    // Every inuse page now; pend pages as frees need them. A failure part way
    // leaves what was taken reachable from the map, so release finds it.
    for (uint64_t i = 0; i < npages; i++) {
        if (!bits_page_get(&s->inuse, i, palloc)) {
            schfs_space_release(s);
            return -5;
        }
        s->npages = i + 1;
    }

    schfs_space_reset(s, nblocks);
    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, uint64_t gen)
{
    // A snapshot may still be holding this block, and if so it must stay.
    //
    // 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.
    //
    // ZFS's rule says which: a snapshot names the tree committed before the
    // group it was created in, so it can hold only blocks written before that
    // group. A block born in or after the newest snapshot's group is in none
    // of them and goes back to the free map as usual; only the older ones are
    // kept, until a mount derives the map from what the snapshots still reach.
    // This used to keep everything while any snapshot existed, so a volume
    // with one never reused a block it freed until it was remounted.
    if ((s->flags & SCHFS_SPACE_RETAIN_ALL) && gen < s->snap_gen)
        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 */
    // No page to note it in, and no memory for one: the block stays in use
    // until the next mount derives the map again. Lost space for a while,
    // never a block handed out twice.
    if (!bits_page_get(&s->pend, blk / SCHFS_SPACE_BITS_PER_PAGE, s->palloc))
        return;
    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 (uint64_t p = 0; p < s->npages; p++) {
        uint64_t *pend  = bits_page(&s->pend, p);
        uint64_t *inuse = bits_page(&s->inuse, p);
        if (!pend)
            continue;
        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;
        }
    }
    // The pages stay for the next group: a volume being written frees in the
    // same places again, and taking and dropping them every commit is waste.
    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);
    }

    // Shared runs, at the length they were *allocated*.
    //
    // Not redundant with the loop above, and the difference matters. A file
    // that shares a run may since have been truncated inside it, so its extent
    // record describes a shorter run than exists -- and the tail would be
    // derived as free while the reference record still holds it. The allocator
    // would then hand out a block that a later release frees a second time.
    for (int i = 0; i < n; i++) {
        if (it[i].key.type != SCHFS_KEY_EXTREF)
            continue;
        if (it[i].size != sizeof(struct schfs_extref))
            continue;
        if ((int)it[i].offset + (int)it[i].size > SCHFS_BLKSIZE)
            continue;
        const struct schfs_extref *x =
            (const struct schfs_extref *)((const char *)blk + it[i].offset);
        schfs_space_mark(s, it[i].key.offset, x->nblocks);
    }
}

// Inodes with no links left, for the mount to free. See schfs_vol.orphans.
static void
collect_orphans(struct schfs_vol *v, 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_INODE)
            continue;
        if (it[i].size != sizeof(struct schfs_inode))
            continue;
        if ((int)it[i].offset + (int)it[i].size > SCHFS_BLKSIZE)
            continue;
        const struct schfs_inode *in =
            (const struct schfs_inode *)((const char *)blk + it[i].offset);
        if (in->nlink != 0 || it[i].key.objectid == v->root_objectid)
            continue;
        if (v->norphans >= SCHFS_ORPHANS_MAX) { v->orphans_more = 1; continue; }
        v->orphans[v->norphans++] = it[i].key.objectid;
    }
}

// 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 newest;                 /* the largest `created` seen */
    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 (sn->created > sr->newest)
            sr->newest = sn->created;
        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);
                collect_orphans(v, 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;
    sr.newest = 0;
    v->norphans = 0;
    v->orphans_more = 0;

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

    if (sr.n > 0 || sr.overflow) {
        v->space.flags   |= SCHFS_SPACE_RETAIN_ALL;
        v->space.snap_gen = sr.newest;
    }

    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;
}