OffsetAllocator.ts

// TODO: remove all new AllocationNode() and statically allocate them in module scope

const NUM_TOP_BINS = 32;
const BINS_PER_LEAF = 8;
const TOP_BINS_INDEX_SHIFT = 3;
const LEAF_BINS_INDEX_MASK = 0x7;
const NUM_LEAF_BINS = NUM_TOP_BINS * BINS_PER_LEAF;

export const NO_SPACE = 0xffffffff;
const UNUSED_NODE_INDEX = 0xffffffff;

export type NodeIndex = number;

export interface Allocation {
    offset: number;
    nodeId: NodeIndex;
}

interface StorageReport {
    totalFreeSpace: number;
    largestFreeRegion: number;
}

interface StorageReportFull {
    freeRegions: Array<{ size: number; count: number }>;
}

function lzcnt(v: number): number {
    return Math.clz32(v);
}

function tzcnt(v: number): number {
    v >>>= 0;
    if (v === 0) return 32;
    v &= -v;
    return 31 - Math.clz32(v);
}

function findLowestSetBitAfter(bitMask: number, startBitIndex: number): number {
    const maskBeforeStartIndex = (1 << startBitIndex) - 1;
    const maskAfterStartIndex = ~maskBeforeStartIndex;
    const bitsAfter = bitMask & maskAfterStartIndex;
    return bitsAfter === 0 ? NO_SPACE : tzcnt(bitsAfter);
}

const MANTISSA_BITS = 3;
const MANTISSA_VALUE = 1 << MANTISSA_BITS;
const MANTISSA_MASK = MANTISSA_VALUE - 1;

export function uintToFloatRoundUp(size: number): number {
    let exp = 0;
    let mantissa = 0;

    if (size < MANTISSA_VALUE) {
        mantissa = size;
    } else {
        const leadingZeros = lzcnt(size);
        const highestSetBit = 31 - leadingZeros;
        const mantissaStartBit = highestSetBit - MANTISSA_BITS;
        exp = mantissaStartBit + 1;
        mantissa = (size >> mantissaStartBit) & MANTISSA_MASK;
        const lowBitsMask = (1 << mantissaStartBit) - 1;
        if ((size & lowBitsMask) !== 0) mantissa++;
    }

    return (exp << MANTISSA_BITS) + mantissa;
}

export function uintToFloatRoundDown(size: number): number {
    let exp = 0;
    let mantissa = 0;

    if (size < MANTISSA_VALUE) {
        mantissa = size;
    } else {
        const leadingZeros = lzcnt(size);
        const highestSetBit = 31 - leadingZeros;
        const mantissaStartBit = highestSetBit - MANTISSA_BITS;
        exp = mantissaStartBit + 1;
        mantissa = (size >> mantissaStartBit) & MANTISSA_MASK;
    }

    return (exp << MANTISSA_BITS) | mantissa;
}

export function floatToUint(floatValue: number): number {
    const exponent = floatValue >>> MANTISSA_BITS;
    const mantissa = floatValue & MANTISSA_MASK;
    return (exponent === 0 ? mantissa : (mantissa | MANTISSA_VALUE) << (exponent - 1)) >>> 0;
}

class AllocationNode {
    static size = 7 * 4;

    dataOffset: number;
    dataSize: number;
    binListPrev: NodeIndex;
    binListNext: NodeIndex;
    neighborPrev: NodeIndex;
    neighborNext: NodeIndex;
    used: boolean;

    constructor() {
        this.dataOffset = 0;
        this.dataSize = 0;
        this.binListPrev = UNUSED_NODE_INDEX;
        this.binListNext = UNUSED_NODE_INDEX;
        this.neighborPrev = UNUSED_NODE_INDEX;
        this.neighborNext = UNUSED_NODE_INDEX;
        this.used = false;
    }

    fromData(view: DataView, offset: number = 0): AllocationNode {
        this.dataOffset = view.getUint32(offset, true);
        this.dataSize = view.getUint32(offset + 4, true);
        this.binListPrev = view.getUint32(offset + 8, true);
        this.binListNext = view.getUint32(offset + 12, true);
        this.neighborPrev = view.getUint32(offset + 16, true);
        this.neighborNext = view.getUint32(offset + 20, true);
        this.used = Boolean(view.getUint32(offset + 24, true));
        return this;
    }

    toData(view: DataView, offset: number = 0) {
        view.setUint32(offset, this.dataOffset, true);
        view.setUint32(offset + 4, this.dataSize, true);
        view.setUint32(offset + 8, this.binListPrev, true);
        view.setUint32(offset + 12, this.binListNext, true);
        view.setUint32(offset + 16, this.neighborPrev, true);
        view.setUint32(offset + 20, this.neighborNext, true);
        view.setUint32(offset + 24, Number(this.used), true);
    }
}

export class Allocator {
    private size: number;
    private maxAllocs: number;
    private freeStorage!: number;
    private usedBinsTop!: number;
    private usedBins!: Uint8Array;
    private binIndicesBuffer: ArrayBuffer;
    private binIndicesView: DataView;
    private freeNodesBuffer: ArrayBuffer;
    private freeNodesView: DataView;
    private nodesBuffer: ArrayBuffer;
    private nodesView: DataView;
    private freeOffset!: number;

    constructor(size: number, maxAllocs: number = 24 * 1024) {
        this.size = size;
        this.maxAllocs = maxAllocs;
        this.binIndicesBuffer = new ArrayBuffer(NUM_LEAF_BINS * 4);
        this.binIndicesView = new DataView(this.binIndicesBuffer);
        this.freeNodesBuffer = new ArrayBuffer(maxAllocs * 4);
        this.freeNodesView = new DataView(this.freeNodesBuffer);
        this.nodesBuffer = new ArrayBuffer(maxAllocs * AllocationNode.size);
        this.nodesView = new DataView(this.nodesBuffer);
        this.reset();
    }

    reset() {
        this.freeStorage = 0;
        this.usedBinsTop = 0;
        this.freeOffset = this.maxAllocs - 1;

        this.usedBins = new Uint8Array(NUM_TOP_BINS);

        for (let i = 0; i < NUM_LEAF_BINS; i++) {
            this.binIndicesView.setUint32(i * 4, UNUSED_NODE_INDEX, true);
        }

        for (let i = 0; i < this.maxAllocs; i++) {
            this.freeNodesView.setUint32(i * 4, this.maxAllocs - i - 1, true);
        }

        const node = new AllocationNode();
        for (let i = 0; i < this.maxAllocs; i++) {
            node.toData(this.nodesView, i * AllocationNode.size);
        }

        this.insertNodeIntoBin(this.size, 0);
    }

    allocate(size: number): Allocation {
        if (this.freeOffset === 0) {
            return { offset: NO_SPACE, nodeId: NO_SPACE };
        }

        const minBinIndex = uintToFloatRoundUp(size);
        const minTopBinIndex = minBinIndex >> TOP_BINS_INDEX_SHIFT;
        const minLeafBinIndex = minBinIndex & LEAF_BINS_INDEX_MASK;

        let topBinIndex = minTopBinIndex;
        let leafBinIndex = NO_SPACE;

        if (this.usedBinsTop & (1 << topBinIndex)) {
            leafBinIndex = findLowestSetBitAfter(this.usedBins[topBinIndex], minLeafBinIndex);
        }

        if (leafBinIndex === NO_SPACE) {
            topBinIndex = findLowestSetBitAfter(this.usedBinsTop, minTopBinIndex + 1);

            if (topBinIndex === NO_SPACE) {
                return { offset: NO_SPACE, nodeId: NO_SPACE };
            }

            leafBinIndex = tzcnt(this.usedBins[topBinIndex]);
        }

        const binIndex = (topBinIndex << TOP_BINS_INDEX_SHIFT) | leafBinIndex;

        const nodeIndex = this.binIndicesView.getUint32(binIndex * 4, true);
        const node = new AllocationNode().fromData(this.nodesView, nodeIndex * AllocationNode.size);
        const nodeTotalSize = node.dataSize;
        node.dataSize = size;
        node.used = true;
        this.binIndicesView.setUint32(binIndex * 4, node.binListNext, true);
        if (node.binListNext !== UNUSED_NODE_INDEX) {
            const nextNode = new AllocationNode().fromData(this.nodesView, node.binListNext * AllocationNode.size);
            nextNode.binListPrev = UNUSED_NODE_INDEX;
            nextNode.toData(this.nodesView, node.binListNext * AllocationNode.size);
        }
        this.freeStorage -= nodeTotalSize;

        if (this.binIndicesView.getUint32(binIndex * 4, true) === UNUSED_NODE_INDEX) {
            this.usedBins[topBinIndex] &= ~(1 << leafBinIndex);

            if (this.usedBins[topBinIndex] === 0) {
                this.usedBinsTop &= ~(1 << topBinIndex);
            }
        }

        const reminderSize = nodeTotalSize - size;
        if (reminderSize > 0) {
            const newNodeIndex = this.insertNodeIntoBin(reminderSize, node.dataOffset + size);

            if (node.neighborNext !== UNUSED_NODE_INDEX) {
                const nextNeighbor = new AllocationNode().fromData(this.nodesView, node.neighborNext * AllocationNode.size);
                nextNeighbor.neighborPrev = newNodeIndex;
                nextNeighbor.toData(this.nodesView, node.neighborNext * AllocationNode.size);
            }
            const newNode = new AllocationNode().fromData(this.nodesView, newNodeIndex * AllocationNode.size);
            newNode.neighborPrev = nodeIndex;
            newNode.neighborNext = node.neighborNext;
            newNode.toData(this.nodesView, newNodeIndex * AllocationNode.size);
            node.neighborNext = newNodeIndex;
        }

        node.toData(this.nodesView, nodeIndex * AllocationNode.size);

        return { offset: node.dataOffset, nodeId: nodeIndex };
    }

    free(nodeIndex: NodeIndex) {
        if (nodeIndex === NO_SPACE) return;

        const node = new AllocationNode().fromData(this.nodesView, nodeIndex * AllocationNode.size);

        let offset = node.dataOffset;
        let size = node.dataSize;

        if (node.neighborPrev !== UNUSED_NODE_INDEX) {
            const prevNode = new AllocationNode().fromData(this.nodesView, node.neighborPrev * AllocationNode.size);
            if (!prevNode.used) {
                offset = prevNode.dataOffset;
                size += prevNode.dataSize;

                this.removeNodeFromBin(node.neighborPrev);

                node.neighborPrev = prevNode.neighborPrev;
            }
        }

        if (node.neighborNext !== UNUSED_NODE_INDEX) {
            const nextNode = new AllocationNode().fromData(this.nodesView, node.neighborNext * AllocationNode.size);
            if (!nextNode.used) {
                size += nextNode.dataSize;

                this.removeNodeFromBin(node.neighborNext);

                node.neighborNext = nextNode.neighborNext;
            }
        }

        const neighborNext = node.neighborNext;
        const neighborPrev = node.neighborPrev;

        this.freeNodesView.setUint32(++this.freeOffset * 4, nodeIndex, true);

        const combinedNodeIndex = this.insertNodeIntoBin(size, offset);

        if (neighborNext !== UNUSED_NODE_INDEX) {
            const combinedNode = new AllocationNode().fromData(this.nodesView, combinedNodeIndex * AllocationNode.size);
            combinedNode.neighborNext = neighborNext;
            combinedNode.toData(this.nodesView, combinedNodeIndex * AllocationNode.size);

            const nextNeighbor = new AllocationNode().fromData(this.nodesView, neighborNext * AllocationNode.size);
            nextNeighbor.neighborPrev = combinedNodeIndex;
            nextNeighbor.toData(this.nodesView, neighborNext * AllocationNode.size);
        }
        if (neighborPrev !== UNUSED_NODE_INDEX) {
            const combinedNode = new AllocationNode().fromData(this.nodesView, combinedNodeIndex * AllocationNode.size);
            combinedNode.neighborPrev = neighborPrev;
            combinedNode.toData(this.nodesView, combinedNodeIndex * AllocationNode.size);

            const prevNeighbor = new AllocationNode().fromData(this.nodesView, neighborPrev * AllocationNode.size);
            prevNeighbor.neighborNext = combinedNodeIndex;
            prevNeighbor.toData(this.nodesView, neighborPrev * AllocationNode.size);
        }
    }

    allocationSize(allocation: Allocation): number {
        return allocation.nodeId === NO_SPACE ? 0 : new AllocationNode().fromData(this.nodesView, allocation.nodeId * AllocationNode.size).dataSize;
    }

    storageReport(): StorageReport {
        let largestFreeRegion = 0;
        let freeStorage = 0;

        if (this.freeOffset > 0) {
            freeStorage = this.freeStorage;
            if (this.usedBinsTop) {
                const topBinIndex = 31 - lzcnt(this.usedBinsTop);
                const leafBinIndex = 31 - lzcnt(this.usedBins[topBinIndex]);
                largestFreeRegion = floatToUint((topBinIndex << TOP_BINS_INDEX_SHIFT) | leafBinIndex);
            }
        }

        return { totalFreeSpace: freeStorage, largestFreeRegion };
    }

    storageReportFull(): StorageReportFull {
        const report: StorageReportFull = { freeRegions: [] };
        for (let i = 0; i < NUM_LEAF_BINS; i++) {
            let count = 0;
            let nodeIndex = this.binIndicesView.getUint32(i * 4, true);
            while (nodeIndex !== UNUSED_NODE_INDEX) {
                nodeIndex = new AllocationNode().fromData(this.nodesView, nodeIndex * AllocationNode.size).binListNext;
                count++;
            }
            report.freeRegions[i] = { size: floatToUint(i), count };
        }
        return report;
    }

    private insertNodeIntoBin(size: number, dataOffset: number): NodeIndex {
        // Round down to bin index to ensure that bin >= alloc
        const binIndex = uintToFloatRoundDown(size);

        const topBinIndex = binIndex >> TOP_BINS_INDEX_SHIFT;
        const leafBinIndex = binIndex & LEAF_BINS_INDEX_MASK;

        // Bin was empty before?
        if (this.binIndicesView.getUint32(binIndex * 4, true) === UNUSED_NODE_INDEX) {
            // Set bin mask bits
            this.usedBins[topBinIndex] |= 1 << leafBinIndex;
            this.usedBinsTop |= 1 << topBinIndex;
        }

        // Take a freelist node and insert on top of the bin linked list (next = old top)
        const topNodeIndex = this.binIndicesView.getUint32(binIndex * 4, true);
        const nodeIndex = this.freeNodesView.getUint32(this.freeOffset * 4, true);
        this.freeOffset--;

        // Create and initialize the new node
        const newNode = new AllocationNode();
        newNode.dataOffset = dataOffset;
        newNode.dataSize = size;
        newNode.binListNext = topNodeIndex;
        newNode.binListPrev = UNUSED_NODE_INDEX;
        newNode.neighborNext = UNUSED_NODE_INDEX;
        newNode.neighborPrev = UNUSED_NODE_INDEX;
        newNode.used = false;
        newNode.toData(this.nodesView, nodeIndex * AllocationNode.size);

        // Update the previous top node if it exists
        if (topNodeIndex !== UNUSED_NODE_INDEX) {
            const topNode = new AllocationNode().fromData(this.nodesView, topNodeIndex * AllocationNode.size);
            topNode.binListPrev = nodeIndex;
            topNode.toData(this.nodesView, topNodeIndex * AllocationNode.size);
        }

        // Update the bin index to point to the new top node
        this.binIndicesView.setUint32(binIndex * 4, nodeIndex, true);

        // Update free storage
        this.freeStorage += size;

        return nodeIndex;
    }

    private removeNodeFromBin(nodeIndex: NodeIndex) {
        const node = new AllocationNode().fromData(this.nodesView, nodeIndex * AllocationNode.size);

        if (node.binListPrev !== UNUSED_NODE_INDEX) {
            const prevNode = new AllocationNode().fromData(this.nodesView, node.binListPrev * AllocationNode.size);
            prevNode.binListNext = node.binListNext;
            prevNode.toData(this.nodesView, node.binListPrev * AllocationNode.size);
            if (node.binListNext !== UNUSED_NODE_INDEX) {
                const nextNode = new AllocationNode().fromData(this.nodesView, node.binListNext * AllocationNode.size);
                nextNode.binListPrev = node.binListPrev;
                nextNode.toData(this.nodesView, node.binListNext * AllocationNode.size);
            }
        } else {
            const binIndex = uintToFloatRoundDown(node.dataSize);
            const topBinIndex = binIndex >> TOP_BINS_INDEX_SHIFT;
            const leafBinIndex = binIndex & LEAF_BINS_INDEX_MASK;

            this.binIndicesView.setUint32(binIndex * 4, node.binListNext, true);
            if (node.binListNext !== UNUSED_NODE_INDEX) {
                const nextNode = new AllocationNode().fromData(this.nodesView, node.binListNext * AllocationNode.size);
                nextNode.binListPrev = UNUSED_NODE_INDEX;
                nextNode.toData(this.nodesView, node.binListNext * AllocationNode.size);
            }

            if (this.binIndicesView.getUint32(binIndex * 4, true) === UNUSED_NODE_INDEX) {
                this.usedBins[topBinIndex] &= ~(1 << leafBinIndex);

                if (this.usedBins[topBinIndex] === 0) {
                    this.usedBinsTop &= ~(1 << topBinIndex);
                }
            }
        }

        this.freeNodesView.setUint32(++this.freeOffset * 4, nodeIndex, true);
        this.freeStorage -= node.dataSize;
    }
}