Files
natsdotnet/src/NATS.Server/Internal/SubjectTree/Nodes.cs

847 lines
23 KiB
C#

// Go reference: server/stree/node.go, leaf.go, node4.go, node10.go, node16.go, node48.go, node256.go
namespace NATS.Server.Internal.SubjectTree;
/// <summary>
/// Internal node interface for the Adaptive Radix Tree.
/// </summary>
internal interface INode
{
/// <summary>
/// Gets whether this node is a terminal subject node that directly stores a subscription value.
/// </summary>
bool IsLeaf { get; }
/// <summary>
/// Gets structural metadata for branch nodes, including compressed path prefix and child count.
/// </summary>
NodeMeta? Base { get; }
/// <summary>
/// Sets the compressed path fragment represented by this node.
/// </summary>
/// <param name="pre">Subject bytes shared by all descendants below this node.</param>
void SetPrefix(ReadOnlySpan<byte> pre);
/// <summary>
/// Adds a child edge for the next subject byte in the adaptive radix tree.
/// </summary>
/// <param name="c">Subject byte used to route lookups to the child node.</param>
/// <param name="n">Child node that owns the remaining subject suffix for this edge.</param>
void AddChild(byte c, INode n);
/// <summary>
/// Returns the child node for the given key byte, or null if not found.
/// The returned wrapper allows in-place replacement of the child reference.
/// </summary>
/// <param name="c">Subject byte to look up in the node's child index.</param>
ChildRef? FindChild(byte c);
/// <summary>
/// Removes the child edge for the provided subject byte.
/// </summary>
/// <param name="c">Subject byte whose child mapping should be removed.</param>
void DeleteChild(byte c);
/// <summary>
/// Gets whether this node has reached its capacity and must grow to the next node shape.
/// </summary>
bool IsFull { get; }
/// <summary>
/// Expands this node to a larger branching factor to accept more distinct subject bytes.
/// </summary>
INode Grow();
/// <summary>
/// Attempts to shrink this node to a smaller branching representation when sparse.
/// </summary>
INode? Shrink();
/// <summary>
/// Matches a subject split into tokens against this node's compressed path fragment.
/// </summary>
/// <param name="parts">Remaining subject tokens to match from this node downward.</param>
/// <returns>The remaining tokens after consuming this node, and whether the fragment matched.</returns>
(ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts);
/// <summary>
/// Gets a short node kind name used by diagnostics and debugging tools.
/// </summary>
string Kind { get; }
/// <summary>
/// Iterates child nodes until the callback returns <see langword="false" />.
/// </summary>
/// <param name="f">Callback invoked for each child node in this branch.</param>
void Iter(Func<INode, bool> f);
/// <summary>
/// Returns the current child nodes for traversal or inspection.
/// </summary>
INode?[] Children();
/// <summary>
/// Gets the number of active child edges in this node.
/// </summary>
ushort NumChildren { get; }
/// <summary>
/// Gets the compressed path bytes represented by this node.
/// </summary>
byte[] Path();
}
/// <summary>
/// Wrapper that allows in-place replacement of a child reference in a node.
/// This is analogous to Go's *node pointer.
/// </summary>
internal sealed class ChildRef(Func<INode?> getter, Action<INode?> setter)
{
/// <summary>
/// Gets or replaces the child node reference stored at a specific branch slot.
/// </summary>
public INode? Node
{
get => getter();
set => setter(value);
}
}
/// <summary>
/// Base metadata for internal (non-leaf) nodes.
/// </summary>
internal sealed class NodeMeta
{
/// <summary>
/// Gets or sets the compressed subject prefix shared by descendants of this branch node.
/// </summary>
public byte[] Prefix { get; set; } = [];
/// <summary>
/// Gets or sets the number of child edges currently populated for this branch node.
/// </summary>
public ushort Size { get; set; }
}
#region Leaf Node
/// <summary>
/// Leaf node holding a value and suffix.
/// Go reference: server/stree/leaf.go
/// </summary>
internal sealed class Leaf<T> : INode
{
public T Value;
public byte[] Suffix;
/// <summary>
/// Initializes a terminal subject-tree node that stores a value for an exact suffix match.
/// </summary>
/// <param name="suffix">Remaining subject bytes that must match to resolve this leaf.</param>
/// <param name="value">Subscription payload or state associated with the matched subject.</param>
public Leaf(ReadOnlySpan<byte> suffix, T value)
{
Value = value;
Suffix = Parts.CopyBytes(suffix);
}
/// <inheritdoc />
public bool IsLeaf => true;
/// <inheritdoc />
public NodeMeta? Base => null;
/// <inheritdoc />
public bool IsFull => true;
/// <inheritdoc />
public ushort NumChildren => 0;
/// <inheritdoc />
public string Kind => "LEAF";
/// <summary>
/// Checks whether the provided subject bytes exactly match this leaf suffix.
/// </summary>
/// <param name="subject">Subject bytes remaining after traversing parent branch prefixes.</param>
/// <returns><see langword="true" /> when the subject resolves to this exact leaf.</returns>
public bool Match(ReadOnlySpan<byte> subject) => subject.SequenceEqual(Suffix);
/// <summary>
/// Replaces the stored suffix when leaf content is split or merged during tree updates.
/// </summary>
/// <param name="suffix">New exact-match suffix bytes for this leaf.</param>
public void SetSuffix(ReadOnlySpan<byte> suffix) => Suffix = Parts.CopyBytes(suffix);
/// <inheritdoc />
public byte[] Path() => Suffix;
/// <inheritdoc />
public INode?[] Children() => [];
/// <inheritdoc />
public void Iter(Func<INode, bool> f) { }
/// <inheritdoc />
public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
=> Parts.MatchPartsAgainstFragment(parts, Suffix);
// These should not be called on a leaf.
/// <inheritdoc />
public void SetPrefix(ReadOnlySpan<byte> pre) => throw new InvalidOperationException("setPrefix called on leaf");
/// <inheritdoc />
public void AddChild(byte c, INode n) => throw new InvalidOperationException("addChild called on leaf");
/// <inheritdoc />
public ChildRef? FindChild(byte c) => throw new InvalidOperationException("findChild called on leaf");
/// <inheritdoc />
public INode Grow() => throw new InvalidOperationException("grow called on leaf");
/// <inheritdoc />
public void DeleteChild(byte c) => throw new InvalidOperationException("deleteChild called on leaf");
/// <inheritdoc />
public INode? Shrink() => throw new InvalidOperationException("shrink called on leaf");
}
#endregion
#region Node4
/// <summary>
/// Node with up to 4 children.
/// Go reference: server/stree/node4.go
/// </summary>
internal sealed class Node4 : INode
{
private readonly INode?[] _child = new INode?[4];
private readonly byte[] _key = new byte[4];
internal readonly NodeMeta Meta = new();
/// <summary>
/// Initializes a small branch node for up to four subject-byte fan-out edges.
/// </summary>
/// <param name="prefix">Compressed subject prefix represented by this branch.</param>
public Node4(ReadOnlySpan<byte> prefix)
{
SetPrefix(prefix);
}
/// <inheritdoc />
public bool IsLeaf => false;
/// <inheritdoc />
public NodeMeta? Base => Meta;
/// <inheritdoc />
public ushort NumChildren => Meta.Size;
/// <inheritdoc />
public bool IsFull => Meta.Size >= 4;
/// <inheritdoc />
public string Kind => "NODE4";
/// <inheritdoc />
public byte[] Path() => Meta.Prefix;
/// <inheritdoc />
public void SetPrefix(ReadOnlySpan<byte> pre)
{
Meta.Prefix = pre.ToArray();
}
/// <inheritdoc />
public void AddChild(byte c, INode n)
{
if (Meta.Size >= 4) throw new InvalidOperationException("node4 full!");
_key[Meta.Size] = c;
_child[Meta.Size] = n;
Meta.Size++;
}
/// <inheritdoc />
public ChildRef? FindChild(byte c)
{
for (int i = 0; i < Meta.Size; i++)
{
if (_key[i] == c)
{
var idx = i;
return new ChildRef(() => _child[idx], v => _child[idx] = v);
}
}
return null;
}
/// <inheritdoc />
public void DeleteChild(byte c)
{
for (int i = 0; i < Meta.Size; i++)
{
if (_key[i] == c)
{
var last = Meta.Size - 1;
if (i < last)
{
_key[i] = _key[last];
_child[i] = _child[last];
_key[last] = 0;
_child[last] = null;
}
else
{
_key[i] = 0;
_child[i] = null;
}
Meta.Size--;
return;
}
}
}
/// <inheritdoc />
public INode Grow()
{
var nn = new Node10(Meta.Prefix);
for (int i = 0; i < 4; i++)
{
nn.AddChild(_key[i], _child[i]!);
}
return nn;
}
/// <inheritdoc />
public INode? Shrink()
{
if (Meta.Size == 1) return _child[0];
return null;
}
/// <inheritdoc />
public void Iter(Func<INode, bool> f)
{
for (int i = 0; i < Meta.Size; i++)
{
if (!f(_child[i]!)) return;
}
}
/// <inheritdoc />
public INode?[] Children()
{
var result = new INode?[Meta.Size];
Array.Copy(_child, result, Meta.Size);
return result;
}
/// <inheritdoc />
public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
=> Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}
#endregion
#region Node10
/// <summary>
/// Node with up to 10 children. Optimized for numeric subject tokens (0-9).
/// Go reference: server/stree/node10.go
/// </summary>
internal sealed class Node10 : INode
{
private readonly INode?[] _child = new INode?[10];
private readonly byte[] _key = new byte[10];
internal readonly NodeMeta Meta = new();
/// <summary>
/// Initializes a branch node tuned for numeric token fan-out, common in ordered stream subjects.
/// </summary>
/// <param name="prefix">Compressed subject prefix represented by this branch.</param>
public Node10(ReadOnlySpan<byte> prefix)
{
SetPrefix(prefix);
}
/// <inheritdoc />
public bool IsLeaf => false;
/// <inheritdoc />
public NodeMeta? Base => Meta;
/// <inheritdoc />
public ushort NumChildren => Meta.Size;
/// <inheritdoc />
public bool IsFull => Meta.Size >= 10;
/// <inheritdoc />
public string Kind => "NODE10";
/// <inheritdoc />
public byte[] Path() => Meta.Prefix;
/// <inheritdoc />
public void SetPrefix(ReadOnlySpan<byte> pre)
{
Meta.Prefix = pre.ToArray();
}
/// <inheritdoc />
public void AddChild(byte c, INode n)
{
if (Meta.Size >= 10) throw new InvalidOperationException("node10 full!");
_key[Meta.Size] = c;
_child[Meta.Size] = n;
Meta.Size++;
}
/// <inheritdoc />
public ChildRef? FindChild(byte c)
{
for (int i = 0; i < Meta.Size; i++)
{
if (_key[i] == c)
{
var idx = i;
return new ChildRef(() => _child[idx], v => _child[idx] = v);
}
}
return null;
}
/// <inheritdoc />
public void DeleteChild(byte c)
{
for (int i = 0; i < Meta.Size; i++)
{
if (_key[i] == c)
{
var last = Meta.Size - 1;
if (i < last)
{
_key[i] = _key[last];
_child[i] = _child[last];
_key[last] = 0;
_child[last] = null;
}
else
{
_key[i] = 0;
_child[i] = null;
}
Meta.Size--;
return;
}
}
}
/// <inheritdoc />
public INode Grow()
{
var nn = new Node16(Meta.Prefix);
for (int i = 0; i < 10; i++)
{
nn.AddChild(_key[i], _child[i]!);
}
return nn;
}
/// <inheritdoc />
public INode? Shrink()
{
if (Meta.Size > 4) return null;
var nn = new Node4([]);
for (int i = 0; i < Meta.Size; i++)
{
nn.AddChild(_key[i], _child[i]!);
}
return nn;
}
/// <inheritdoc />
public void Iter(Func<INode, bool> f)
{
for (int i = 0; i < Meta.Size; i++)
{
if (!f(_child[i]!)) return;
}
}
/// <inheritdoc />
public INode?[] Children()
{
var result = new INode?[Meta.Size];
Array.Copy(_child, result, Meta.Size);
return result;
}
/// <inheritdoc />
public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
=> Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}
#endregion
#region Node16
/// <summary>
/// Node with up to 16 children.
/// Go reference: server/stree/node16.go
/// </summary>
internal sealed class Node16 : INode
{
private readonly INode?[] _child = new INode?[16];
private readonly byte[] _key = new byte[16];
internal readonly NodeMeta Meta = new();
/// <summary>
/// Initializes a medium branch node for moderate subject fan-out without index indirection.
/// </summary>
/// <param name="prefix">Compressed subject prefix represented by this branch.</param>
public Node16(ReadOnlySpan<byte> prefix)
{
SetPrefix(prefix);
}
/// <inheritdoc />
public bool IsLeaf => false;
/// <inheritdoc />
public NodeMeta? Base => Meta;
/// <inheritdoc />
public ushort NumChildren => Meta.Size;
/// <inheritdoc />
public bool IsFull => Meta.Size >= 16;
/// <inheritdoc />
public string Kind => "NODE16";
/// <inheritdoc />
public byte[] Path() => Meta.Prefix;
/// <inheritdoc />
public void SetPrefix(ReadOnlySpan<byte> pre)
{
Meta.Prefix = pre.ToArray();
}
/// <inheritdoc />
public void AddChild(byte c, INode n)
{
if (Meta.Size >= 16) throw new InvalidOperationException("node16 full!");
_key[Meta.Size] = c;
_child[Meta.Size] = n;
Meta.Size++;
}
/// <inheritdoc />
public ChildRef? FindChild(byte c)
{
for (int i = 0; i < Meta.Size; i++)
{
if (_key[i] == c)
{
var idx = i;
return new ChildRef(() => _child[idx], v => _child[idx] = v);
}
}
return null;
}
/// <inheritdoc />
public void DeleteChild(byte c)
{
for (int i = 0; i < Meta.Size; i++)
{
if (_key[i] == c)
{
var last = Meta.Size - 1;
if (i < last)
{
_key[i] = _key[last];
_child[i] = _child[last];
_key[last] = 0;
_child[last] = null;
}
else
{
_key[i] = 0;
_child[i] = null;
}
Meta.Size--;
return;
}
}
}
/// <inheritdoc />
public INode Grow()
{
var nn = new Node48(Meta.Prefix);
for (int i = 0; i < 16; i++)
{
nn.AddChild(_key[i], _child[i]!);
}
return nn;
}
/// <inheritdoc />
public INode? Shrink()
{
if (Meta.Size > 10) return null;
var nn = new Node10([]);
for (int i = 0; i < Meta.Size; i++)
{
nn.AddChild(_key[i], _child[i]!);
}
return nn;
}
/// <inheritdoc />
public void Iter(Func<INode, bool> f)
{
for (int i = 0; i < Meta.Size; i++)
{
if (!f(_child[i]!)) return;
}
}
/// <inheritdoc />
public INode?[] Children()
{
var result = new INode?[Meta.Size];
Array.Copy(_child, result, Meta.Size);
return result;
}
/// <inheritdoc />
public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
=> Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}
#endregion
#region Node48
/// <summary>
/// Node with up to 48 children. Uses a 256-byte index array (1-indexed) to map keys to child slots.
/// Go reference: server/stree/node48.go
/// </summary>
internal sealed class Node48 : INode
{
internal readonly INode?[] Child = new INode?[48];
internal readonly byte[] Key = new byte[256]; // 1-indexed: 0 means no entry
internal readonly NodeMeta Meta = new();
/// <summary>
/// Initializes a high fan-out branch node that trades memory for faster byte-key lookups.
/// </summary>
/// <param name="prefix">Compressed subject prefix represented by this branch.</param>
public Node48(ReadOnlySpan<byte> prefix)
{
SetPrefix(prefix);
}
/// <inheritdoc />
public bool IsLeaf => false;
/// <inheritdoc />
public NodeMeta? Base => Meta;
/// <inheritdoc />
public ushort NumChildren => Meta.Size;
/// <inheritdoc />
public bool IsFull => Meta.Size >= 48;
/// <inheritdoc />
public string Kind => "NODE48";
/// <inheritdoc />
public byte[] Path() => Meta.Prefix;
/// <inheritdoc />
public void SetPrefix(ReadOnlySpan<byte> pre)
{
Meta.Prefix = pre.ToArray();
}
/// <inheritdoc />
public void AddChild(byte c, INode n)
{
if (Meta.Size >= 48) throw new InvalidOperationException("node48 full!");
Child[Meta.Size] = n;
Key[c] = (byte)(Meta.Size + 1); // 1-indexed
Meta.Size++;
}
/// <inheritdoc />
public ChildRef? FindChild(byte c)
{
var i = Key[c];
if (i == 0) return null;
var idx = i - 1;
return new ChildRef(() => Child[idx], v => Child[idx] = v);
}
/// <inheritdoc />
public void DeleteChild(byte c)
{
var i = Key[c];
if (i == 0) return;
i--; // Adjust for 1-indexing
var last = (byte)(Meta.Size - 1);
if (i < last)
{
Child[i] = Child[last];
for (int ic = 0; ic < 256; ic++)
{
if (Key[ic] == last + 1)
{
Key[ic] = (byte)(i + 1);
break;
}
}
}
Child[last] = null;
Key[c] = 0;
Meta.Size--;
}
/// <inheritdoc />
public INode Grow()
{
var nn = new Node256(Meta.Prefix);
for (int c = 0; c < 256; c++)
{
var i = Key[c];
if (i > 0)
{
nn.AddChild((byte)c, Child[i - 1]!);
}
}
return nn;
}
/// <inheritdoc />
public INode? Shrink()
{
if (Meta.Size > 16) return null;
var nn = new Node16([]);
for (int c = 0; c < 256; c++)
{
var i = Key[c];
if (i > 0)
{
nn.AddChild((byte)c, Child[i - 1]!);
}
}
return nn;
}
/// <inheritdoc />
public void Iter(Func<INode, bool> f)
{
foreach (var c in Child)
{
if (c != null && !f(c)) return;
}
}
/// <inheritdoc />
public INode?[] Children()
{
var result = new INode?[Meta.Size];
Array.Copy(Child, result, Meta.Size);
return result;
}
/// <inheritdoc />
public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
=> Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}
#endregion
#region Node256
/// <summary>
/// Node with up to 256 children. Direct array indexed by byte value.
/// Go reference: server/stree/node256.go
/// </summary>
internal sealed class Node256 : INode
{
internal readonly INode?[] Child = new INode?[256];
internal readonly NodeMeta Meta = new();
/// <summary>
/// Initializes the maximum fan-out branch node with direct byte-to-child indexing.
/// </summary>
/// <param name="prefix">Compressed subject prefix represented by this branch.</param>
public Node256(ReadOnlySpan<byte> prefix)
{
SetPrefix(prefix);
}
/// <inheritdoc />
public bool IsLeaf => false;
/// <inheritdoc />
public NodeMeta? Base => Meta;
/// <inheritdoc />
public ushort NumChildren => Meta.Size;
/// <inheritdoc />
public bool IsFull => false; // node256 is never full
/// <inheritdoc />
public string Kind => "NODE256";
/// <inheritdoc />
public byte[] Path() => Meta.Prefix;
/// <inheritdoc />
public void SetPrefix(ReadOnlySpan<byte> pre)
{
Meta.Prefix = pre.ToArray();
}
/// <inheritdoc />
public void AddChild(byte c, INode n)
{
Child[c] = n;
Meta.Size++;
}
/// <inheritdoc />
public ChildRef? FindChild(byte c)
{
if (Child[c] == null) return null;
return new ChildRef(() => Child[c], v => Child[c] = v);
}
/// <inheritdoc />
public void DeleteChild(byte c)
{
if (Child[c] != null)
{
Child[c] = null;
Meta.Size--;
}
}
/// <inheritdoc />
public INode Grow() => throw new InvalidOperationException("grow can not be called on node256");
/// <inheritdoc />
public INode? Shrink()
{
if (Meta.Size > 48) return null;
var nn = new Node48([]);
for (int c = 0; c < 256; c++)
{
if (Child[c] != null)
{
nn.AddChild((byte)c, Child[c]!);
}
}
return nn;
}
/// <inheritdoc />
public void Iter(Func<INode, bool> f)
{
for (int i = 0; i < 256; i++)
{
if (Child[i] != null)
{
if (!f(Child[i]!)) return;
}
}
}
/// <inheritdoc />
public INode?[] Children()
{
// Return the full 256 array, same as Go
return (INode?[])Child.Clone();
}
/// <inheritdoc />
public (ReadOnlyMemory<byte>[] RemainingParts, bool Matched) MatchParts(ReadOnlyMemory<byte>[] parts)
=> Parts.MatchPartsAgainstFragment(parts, Meta.Prefix);
}
#endregion