PersistentMap/PersistentMap/Nodes.cs

using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;

namespace PersistentMap;

[Flags]
public enum NodeFlags : byte
{
    None = 0,
    IsLeaf = 1 << 0,
    IsRoot = 1 << 1,
    HasPrefixes = 1 << 2
}

[StructLayout(LayoutKind.Sequential, Pack = 1)]
public struct NodeHeader
{
    // 6 Bytes: OwnerId for Copy-on-Write (CoW)
    public OwnerId Owner;

    // 1 Byte: Number of items currently used
    public byte Count;

    // 1 Byte: Type flags (Leaf, Root, etc.)
    public NodeFlags Flags;

    public NodeHeader(OwnerId owner, byte count, NodeFlags flags)
    {
        Owner = owner;
        Count = count;
        Flags = flags;
    }
}

// Constraint: Internal Nodes fixed at 32 children.
// This removes the need for a separate array allocation for children references.
[InlineArray(32)]
public struct NodeBuffer<V>
{
    private Node<V>? _element0;
}

public abstract class Node<K>
{
    public NodeHeader Header;
        
    // FIX: Change to 'internal' so BTreeFunctions can shift the array directly.
    internal long[]? _prefixes;

    protected Node(OwnerId owner, NodeFlags flags)
    {
        Header = new NodeHeader(owner, 0, flags);
    }

    public abstract Span<K> GetKeys(); 

    // Keep this for Search (Read-Only): it limits the view to valid items only.
    public Span<long> Prefixes => _prefixes.AsSpan(0, Header.Count);
        
    public bool IsLeaf => (Header.Flags & NodeFlags.IsLeaf) != 0;

    public abstract Node<K> EnsureEditable(OwnerId transactionId);
        
    public void SetCount(int newCount) => Header.Count = (byte)newCount;
    
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public LeafNode<K, V> AsLeaf<V>()
    {
        // Zero-overhead cast. Assumes you checked IsLeaf or know logic flow.
        return Unsafe.As<LeafNode<K, V>>(this);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public InternalNode<K> AsInternal()
    {
        // Zero-overhead cast. Assumes you checked !IsLeaf or know logic flow.
        return Unsafe.As<InternalNode<K>>(this);
    }
}

public sealed class LeafNode<K, V> : Node<K>
{
    public const int Capacity = 64;
    public const int MergeThreshold = 8;

    // Leaf stores Keys and Values
    public K[] Keys;
    public LeafNode<K, V>? Next; // For range scans
    public V[] Values;

    public LeafNode(OwnerId owner) : base(owner, NodeFlags.IsLeaf | NodeFlags.HasPrefixes)
    {
        Keys = new K[Capacity];
        Values = new V[Capacity];
        _prefixes = new long[Capacity];
    }

    // Copy Constructor for CoW
    private LeafNode(LeafNode<K, V> original, OwnerId newOwner)
        : base(newOwner, original.Header.Flags)
    {
        Header.Count = original.Header.Count;
        Next = original.Next;

        // Allocate new arrays
        Keys = new K[Capacity];
        Values = new V[Capacity];
        _prefixes = new long[Capacity];

        // Copy data
        Array.Copy(original.Keys, Keys, original.Header.Count);
        Array.Copy(original.Values, Values, original.Header.Count);
        if (original._prefixes != null)
            Array.Copy(original._prefixes, _prefixes, original.Header.Count);
    }

    public override Node<K> EnsureEditable(OwnerId transactionId)
    {
        // CASE 1: Persistent Mode (transactionId is None).
        // We MUST create a copy, because we cannot distinguish "Shared Immutable Node (0)"
        // from "New Mutable Node (0)" based on ID alone.
        // However, since BTreeFunctions only calls this once before descending, 
        // we won't copy the same fresh node twice.
        if (transactionId == OwnerId.None)
        {
            return new LeafNode<K, V>(this, OwnerId.None);
        }

        // CASE 2: Transient Mode.
        // If we own the node, return it.
        if (Header.Owner == transactionId)
        {
            return this;
        }

        // CASE 3: CoW needed (Ownership mismatch).
        return new LeafNode<K, V>(this, transactionId);
    }

    public override Span<K> GetKeys()
    {
        return Keys.AsSpan(0, Header.Count);
    }

    public Span<V> GetValues()
    {
        return Values.AsSpan(0, Header.Count);
    }
}

public sealed class InternalNode<K> : Node<K>
{
    public const int Capacity = 32;

    // Inline buffer for children (no array object overhead)
    public NodeBuffer<K> Children;

    // Internal stores Keys (separators) and Children
    public K[] Keys;

    public InternalNode(OwnerId owner) : base(owner, NodeFlags.HasPrefixes)
    {
        Keys = new K[Capacity];
        _prefixes = new long[Capacity];
        // Children buffer is a struct, zero-initialized by default
    }

    // Copy Constructor for CoW
    private InternalNode(InternalNode<K> original, OwnerId newOwner)
        : base(newOwner, original.Header.Flags)
    {
        Header.Count = original.Header.Count;

        Keys = new K[Capacity];
        _prefixes = new long[Capacity];

        // Copy Keys and Prefixes
        Array.Copy(original.Keys, Keys, original.Header.Count);
        if (original._prefixes != null)
            Array.Copy(original._prefixes, _prefixes, original.Header.Count);

        // Copy Children (Manual loop required for InlineArray in generic context usually, 
        // but here we can iterate the span)
        var srcChildren = original.GetChildren();
        for (var i = 0; i < srcChildren.Length; i++) Children[i] = srcChildren[i];
    }

    public override Node<K> EnsureEditable(OwnerId transactionId)
    {
        if (transactionId == OwnerId.None)
        {
            return new InternalNode<K>(this, OwnerId.None);
        }

        if (Header.Owner == transactionId)
        {
            return this;
        }

        return new InternalNode<K>(this, transactionId);
    }
    public override Span<K> GetKeys()
    {
        return Keys.AsSpan(0, Header.Count);
    }

    // Exposes the InlineArray as a Span
    public Span<Node<K>?> GetChildren()
    {
        return MemoryMarshal.CreateSpan<Node<K>?>(ref Children[0]!, Header.Count + 1);
    }

    public void SetChild(int index, Node<K> node)
    {
        Children[index] = node;
    }
}

[StructLayout(LayoutKind.Auto, Pack = 1)]
public readonly struct OwnerId(uint id, ushort gen) : IEquatable<OwnerId>
{
    private const int BatchSize = 100;

    // The max of allocated IDs globally.
    // Starts at 0, so the first batch reserves IDs 1 to 100.
    private static long _globalHighWaterMark;


    // These fields are unique to each thread. They initialize to 0/default.
    // The current ID value this thread is handing out.
    [ThreadStatic] private static long _localCurrentId;

    // How many IDs are left in this thread's current batch.
    [ThreadStatic] private static int _localRemaining;

    // ---------------------------------------------------------
    // Instance Data (6 Bytes)
    // ---------------------------------------------------------
    private readonly uint Id = id; // 4 bytes
    private readonly ushort Gen = gen; // 2 bytes

    /// <summary>
    ///     Generates the next unique OwnerId.
    ///     mostly non-blocking (thread-local), hits Interlocked only once per 100 IDs.
    /// </summary>
    public static OwnerId Next()
    {
        // We have IDs remaining in our local batch.
        // This executes with zero locking overhead.
        if (_localRemaining > 0)
        {
            _localRemaining--;
            var val = ++_localCurrentId;
            return new OwnerId((uint)val, (ushort)(val >> 32));
        }

        // SLOW PATH: We ran out (or this is the thread's first call).
        return NextBatch();
    }

    private static OwnerId NextBatch()
    {
        // Atomically reserve a new block of IDs from the global counter.
        // Only one thread contends for this cache line at a time.
        var reservedEnd = Interlocked.Add(ref _globalHighWaterMark, BatchSize);

        // Calculate the start of our new range.
        var reservedStart = reservedEnd - BatchSize + 1;

        // Reset the local cache.
        // We set _localCurrentId to (start - 1) so that the first increment
        // inside the logic below lands exactly on 'reservedStart'.
        _localCurrentId = reservedStart - 1;
        _localRemaining = BatchSize;

        // Perform the generation logic (same as Fast Path)
        _localRemaining--;
        var val = ++_localCurrentId;
        return new OwnerId((uint)val, (ushort)(val >> 32));
    }

    public static readonly OwnerId None = new(0, 0);

    public bool IsNone => Id == 0 && Gen == 0;

    public bool Equals(OwnerId other)
    {
        return Id == other.Id && Gen == other.Gen;
    }

    public override bool Equals(object? obj)
    {
        return obj is OwnerId other && Equals(other);
    }

    public override int GetHashCode()
    {
        return HashCode.Combine(Id, Gen);
    }

    public static bool operator ==(OwnerId left, OwnerId right)
    {
        return left.Equals(right);
    }

    public static bool operator !=(OwnerId left, OwnerId right)
    {
        return !left.Equals(right);
    }
}