feat: port avl module - SequenceSet AVL tree (36 features, 17 tests)
This commit is contained in:
Joseph Doherty
@@ -0,0 +1,631 @@
|
||||
using System.Buffers.Binary;
|
||||
|
||||
namespace ZB.MOM.NatsNet.Server.Internal.DataStructures;
|
||||
|
||||
/// <summary>
|
||||
/// Memory and encoding-optimized set for storing unsigned 64-bit integers.
|
||||
/// Implemented as an AVL tree where each node holds bitmasks for set membership.
|
||||
/// Approximately 80-100x more memory-efficient than a <see cref="HashSet{T}"/>.
|
||||
/// Not thread-safe.
|
||||
/// </summary>
|
||||
public sealed class SequenceSet
|
||||
{
|
||||
private const int BitsPerBucket = 64;
|
||||
private const int NumBuckets = 32;
|
||||
internal const int NumEntries = NumBuckets * BitsPerBucket; // 2048
|
||||
|
||||
private const byte MagicByte = 22;
|
||||
private const byte CurrentVersion = 2;
|
||||
private const int HdrLen = 2;
|
||||
private const int MinLen = 2 + 8; // magic + version + num_nodes(4) + num_entries(4)
|
||||
|
||||
private Node? _root;
|
||||
private int _size;
|
||||
private int _nodes;
|
||||
private bool _changed;
|
||||
|
||||
// --- Errors ---
|
||||
|
||||
public static readonly Exception ErrBadEncoding = new InvalidDataException("ss: bad encoding");
|
||||
public static readonly Exception ErrBadVersion = new InvalidDataException("ss: bad version");
|
||||
public static readonly Exception ErrSetNotEmpty = new InvalidOperationException("ss: set not empty");
|
||||
|
||||
// --- Internal access for testing ---
|
||||
|
||||
internal Node? Root => _root;
|
||||
|
||||
// --- Public API ---
|
||||
|
||||
/// <summary>Inserts a sequence number into the set. Tree is balanced inline.</summary>
|
||||
public void Insert(ulong seq)
|
||||
{
|
||||
_root = Node.Insert(_root, seq, ref _changed, ref _nodes);
|
||||
if (_changed)
|
||||
{
|
||||
_changed = false;
|
||||
_size++;
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>Returns true if the sequence is a member of the set.</summary>
|
||||
public bool Exists(ulong seq)
|
||||
{
|
||||
for (var n = _root; n != null;)
|
||||
{
|
||||
if (seq < n.Base)
|
||||
{
|
||||
n = n.Left;
|
||||
}
|
||||
else if (seq >= n.Base + NumEntries)
|
||||
{
|
||||
n = n.Right;
|
||||
}
|
||||
else
|
||||
{
|
||||
return n.ExistsBit(seq);
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Sets the initial minimum sequence when known. More effectively utilizes space.
|
||||
/// The set must be empty.
|
||||
/// </summary>
|
||||
public void SetInitialMin(ulong min)
|
||||
{
|
||||
if (!IsEmpty)
|
||||
throw (InvalidOperationException)ErrSetNotEmpty;
|
||||
|
||||
_root = new Node(min);
|
||||
_nodes = 1;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Removes the sequence from the set. Returns true if the sequence was present.
|
||||
/// </summary>
|
||||
public bool Delete(ulong seq)
|
||||
{
|
||||
if (_root == null) return false;
|
||||
|
||||
_root = Node.Delete(_root, seq, ref _changed, ref _nodes);
|
||||
if (_changed)
|
||||
{
|
||||
_changed = false;
|
||||
_size--;
|
||||
if (_size == 0)
|
||||
Empty();
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/// <summary>Returns the number of items in the set.</summary>
|
||||
public int Size => _size;
|
||||
|
||||
/// <summary>Returns the number of nodes in the AVL tree.</summary>
|
||||
public int Nodes => _nodes;
|
||||
|
||||
/// <summary>Clears all items from the set.</summary>
|
||||
public void Empty()
|
||||
{
|
||||
_root = null;
|
||||
_size = 0;
|
||||
_nodes = 0;
|
||||
}
|
||||
|
||||
/// <summary>Returns true if the set contains no items.</summary>
|
||||
public bool IsEmpty => _root == null;
|
||||
|
||||
/// <summary>
|
||||
/// Invokes the callback for each item in ascending order.
|
||||
/// Stops early if the callback returns false.
|
||||
/// </summary>
|
||||
public void Range(Func<ulong, bool> f) => Node.Iter(_root, f);
|
||||
|
||||
/// <summary>Returns the heights of the left and right subtrees of the root.</summary>
|
||||
public (int Left, int Right) Heights()
|
||||
{
|
||||
if (_root == null) return (0, 0);
|
||||
return (_root.Left?.Height ?? 0, _root.Right?.Height ?? 0);
|
||||
}
|
||||
|
||||
/// <summary>Returns min, max, and count of set items.</summary>
|
||||
public (ulong Min, ulong Max, ulong Count) State()
|
||||
{
|
||||
if (_root == null) return (0, 0, 0);
|
||||
var (min, max) = MinMax();
|
||||
return (min, max, (ulong)_size);
|
||||
}
|
||||
|
||||
/// <summary>Returns the minimum and maximum values in the set.</summary>
|
||||
public (ulong Min, ulong Max) MinMax()
|
||||
{
|
||||
if (_root == null) return (0, 0);
|
||||
|
||||
ulong min = 0;
|
||||
for (var l = _root; l != null; l = l.Left)
|
||||
if (l.Left == null) min = l.Min();
|
||||
|
||||
ulong max = 0;
|
||||
for (var r = _root; r != null; r = r.Right)
|
||||
if (r.Right == null) max = r.Max();
|
||||
|
||||
return (min, max);
|
||||
}
|
||||
|
||||
/// <summary>Returns a deep clone of this set.</summary>
|
||||
public SequenceSet Clone()
|
||||
{
|
||||
var css = new SequenceSet { _nodes = _nodes, _size = _size };
|
||||
css._root = Node.Clone(_root);
|
||||
return css;
|
||||
}
|
||||
|
||||
/// <summary>Unions one or more sequence sets into this set.</summary>
|
||||
public void Union(params SequenceSet[] sets)
|
||||
{
|
||||
foreach (var sa in sets)
|
||||
{
|
||||
Node.NodeIter(sa._root, n =>
|
||||
{
|
||||
for (var nb = 0; nb < NumBuckets; nb++)
|
||||
{
|
||||
var b = n.Bits[nb];
|
||||
for (var pos = 0UL; b != 0; pos++)
|
||||
{
|
||||
if ((b & 1) == 1)
|
||||
{
|
||||
var seq = n.Base + ((ulong)nb * BitsPerBucket) + pos;
|
||||
Insert(seq);
|
||||
}
|
||||
b >>= 1;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// <summary>Returns the union of all given sets.</summary>
|
||||
public static SequenceSet? UnionSets(params SequenceSet[] sets)
|
||||
{
|
||||
if (sets.Length == 0) return null;
|
||||
|
||||
// Clone the largest set first for efficiency.
|
||||
Array.Sort(sets, (a, b) => b.Size.CompareTo(a.Size));
|
||||
var ss = sets[0].Clone();
|
||||
for (var i = 1; i < sets.Length; i++)
|
||||
{
|
||||
sets[i].Range(n =>
|
||||
{
|
||||
ss.Insert(n);
|
||||
return true;
|
||||
});
|
||||
}
|
||||
return ss;
|
||||
}
|
||||
|
||||
/// <summary>Returns the number of bytes needed to encode this set.</summary>
|
||||
public int EncodeLen() => MinLen + (Nodes * ((NumBuckets + 1) * 8 + 2));
|
||||
|
||||
/// <summary>
|
||||
/// Encodes this set into a compact binary representation.
|
||||
/// Reuses the provided buffer if it is large enough.
|
||||
/// </summary>
|
||||
public byte[] Encode(byte[]? buf)
|
||||
{
|
||||
var nn = Nodes;
|
||||
var encLen = EncodeLen();
|
||||
|
||||
if (buf == null || buf.Length < encLen)
|
||||
buf = new byte[encLen];
|
||||
|
||||
buf[0] = MagicByte;
|
||||
buf[1] = CurrentVersion;
|
||||
|
||||
var i = HdrLen;
|
||||
BinaryPrimitives.WriteUInt32LittleEndian(buf.AsSpan(i), (uint)nn);
|
||||
BinaryPrimitives.WriteUInt32LittleEndian(buf.AsSpan(i + 4), (uint)_size);
|
||||
i += 8;
|
||||
|
||||
Node.NodeIter(_root, n =>
|
||||
{
|
||||
BinaryPrimitives.WriteUInt64LittleEndian(buf.AsSpan(i), n.Base);
|
||||
i += 8;
|
||||
foreach (var b in n.Bits)
|
||||
{
|
||||
BinaryPrimitives.WriteUInt64LittleEndian(buf.AsSpan(i), b);
|
||||
i += 8;
|
||||
}
|
||||
BinaryPrimitives.WriteUInt16LittleEndian(buf.AsSpan(i), (ushort)n.Height);
|
||||
i += 2;
|
||||
});
|
||||
|
||||
return buf[..i];
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Decodes a sequence set from the binary representation.
|
||||
/// Returns the set and the number of bytes consumed.
|
||||
/// Throws <see cref="InvalidDataException"/> on malformed input.
|
||||
/// </summary>
|
||||
public static (SequenceSet Set, int BytesRead) Decode(ReadOnlySpan<byte> buf)
|
||||
{
|
||||
if (buf.Length < MinLen || buf[0] != MagicByte)
|
||||
throw (InvalidDataException)ErrBadEncoding;
|
||||
|
||||
return buf[1] switch
|
||||
{
|
||||
1 => Decodev1(buf),
|
||||
2 => Decodev2(buf),
|
||||
_ => throw (InvalidDataException)ErrBadVersion
|
||||
};
|
||||
}
|
||||
|
||||
// --- Internal tree helpers ---
|
||||
|
||||
/// <summary>Inserts a pre-built node directly into the tree (used during Decode).</summary>
|
||||
internal void InsertNode(Node n)
|
||||
{
|
||||
_nodes++;
|
||||
if (_root == null)
|
||||
{
|
||||
_root = n;
|
||||
return;
|
||||
}
|
||||
for (var p = _root; p != null;)
|
||||
{
|
||||
if (n.Base < p.Base)
|
||||
{
|
||||
if (p.Left == null) { p.Left = n; return; }
|
||||
p = p.Left;
|
||||
}
|
||||
else
|
||||
{
|
||||
if (p.Right == null) { p.Right = n; return; }
|
||||
p = p.Right;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static (SequenceSet Set, int BytesRead) Decodev2(ReadOnlySpan<byte> buf)
|
||||
{
|
||||
var index = 2;
|
||||
var nn = (int)BinaryPrimitives.ReadUInt32LittleEndian(buf[index..]);
|
||||
var sz = (int)BinaryPrimitives.ReadUInt32LittleEndian(buf[(index + 4)..]);
|
||||
index += 8;
|
||||
|
||||
var expectedLen = MinLen + (nn * ((NumBuckets + 1) * 8 + 2));
|
||||
if (buf.Length < expectedLen)
|
||||
throw (InvalidDataException)ErrBadEncoding;
|
||||
|
||||
var ss = new SequenceSet { _size = sz };
|
||||
var nodes = new Node[nn];
|
||||
|
||||
for (var i = 0; i < nn; i++)
|
||||
{
|
||||
var n = new Node(BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]));
|
||||
index += 8;
|
||||
for (var bi = 0; bi < NumBuckets; bi++)
|
||||
{
|
||||
n.Bits[bi] = BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]);
|
||||
index += 8;
|
||||
}
|
||||
n.Height = (int)BinaryPrimitives.ReadUInt16LittleEndian(buf[index..]);
|
||||
index += 2;
|
||||
nodes[i] = n;
|
||||
ss.InsertNode(n);
|
||||
}
|
||||
|
||||
return (ss, index);
|
||||
}
|
||||
|
||||
private static (SequenceSet Set, int BytesRead) Decodev1(ReadOnlySpan<byte> buf)
|
||||
{
|
||||
const int v1NumBuckets = 64;
|
||||
|
||||
var index = 2;
|
||||
var nn = (int)BinaryPrimitives.ReadUInt32LittleEndian(buf[index..]);
|
||||
var sz = (int)BinaryPrimitives.ReadUInt32LittleEndian(buf[(index + 4)..]);
|
||||
index += 8;
|
||||
|
||||
var expectedLen = MinLen + (nn * ((v1NumBuckets + 1) * 8 + 2));
|
||||
if (buf.Length < expectedLen)
|
||||
throw (InvalidDataException)ErrBadEncoding;
|
||||
|
||||
var ss = new SequenceSet();
|
||||
for (var i = 0; i < nn; i++)
|
||||
{
|
||||
var baseVal = BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]);
|
||||
index += 8;
|
||||
for (var nb = 0UL; nb < v1NumBuckets; nb++)
|
||||
{
|
||||
var n = BinaryPrimitives.ReadUInt64LittleEndian(buf[index..]);
|
||||
for (var pos = 0UL; n != 0; pos++)
|
||||
{
|
||||
if ((n & 1) == 1)
|
||||
{
|
||||
var seq = baseVal + (nb * BitsPerBucket) + pos;
|
||||
ss.Insert(seq);
|
||||
}
|
||||
n >>= 1;
|
||||
}
|
||||
index += 8;
|
||||
}
|
||||
// Skip encoded height.
|
||||
index += 2;
|
||||
}
|
||||
|
||||
if (ss.Size != sz)
|
||||
throw (InvalidDataException)ErrBadEncoding;
|
||||
|
||||
return (ss, index);
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// Internal Node class
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
internal sealed class Node
|
||||
{
|
||||
public ulong Base;
|
||||
public readonly ulong[] Bits = new ulong[NumBuckets];
|
||||
public Node? Left;
|
||||
public Node? Right;
|
||||
public int Height;
|
||||
|
||||
public Node(ulong baseVal)
|
||||
{
|
||||
Base = baseVal;
|
||||
Height = 1;
|
||||
}
|
||||
|
||||
// Sets the bit for seq. seq must be within [Base, Base+NumEntries).
|
||||
public void SetBit(ulong seq, ref bool inserted)
|
||||
{
|
||||
var offset = seq - Base;
|
||||
var i = (int)(offset / BitsPerBucket);
|
||||
var mask = 1UL << (int)(offset % BitsPerBucket);
|
||||
if ((Bits[i] & mask) == 0)
|
||||
{
|
||||
Bits[i] |= mask;
|
||||
inserted = true;
|
||||
}
|
||||
}
|
||||
|
||||
public bool ExistsBit(ulong seq)
|
||||
{
|
||||
var offset = seq - Base;
|
||||
var i = (int)(offset / BitsPerBucket);
|
||||
var mask = 1UL << (int)(offset % BitsPerBucket);
|
||||
return (Bits[i] & mask) != 0;
|
||||
}
|
||||
|
||||
// Clears the bit for seq. Returns true if the node is now empty.
|
||||
public bool ClearBit(ulong seq, ref bool deleted)
|
||||
{
|
||||
var offset = seq - Base;
|
||||
var i = (int)(offset / BitsPerBucket);
|
||||
var mask = 1UL << (int)(offset % BitsPerBucket);
|
||||
if ((Bits[i] & mask) != 0)
|
||||
{
|
||||
Bits[i] &= ~mask;
|
||||
deleted = true;
|
||||
}
|
||||
foreach (var b in Bits)
|
||||
if (b != 0) return false;
|
||||
return true;
|
||||
}
|
||||
|
||||
public ulong Min()
|
||||
{
|
||||
for (var i = 0; i < NumBuckets; i++)
|
||||
{
|
||||
if (Bits[i] != 0)
|
||||
return Base + (ulong)(i * BitsPerBucket) + (ulong)System.Numerics.BitOperations.TrailingZeroCount(Bits[i]);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
public ulong Max()
|
||||
{
|
||||
for (var i = NumBuckets - 1; i >= 0; i--)
|
||||
{
|
||||
if (Bits[i] != 0)
|
||||
return Base + (ulong)(i * BitsPerBucket) +
|
||||
(ulong)(BitsPerBucket - System.Numerics.BitOperations.LeadingZeroCount(Bits[i] >> 1));
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Static AVL helpers
|
||||
|
||||
public static int BalanceFactor(Node? n)
|
||||
{
|
||||
if (n == null) return 0;
|
||||
return (n.Left?.Height ?? 0) - (n.Right?.Height ?? 0);
|
||||
}
|
||||
|
||||
private static int MaxH(Node? n)
|
||||
{
|
||||
if (n == null) return 0;
|
||||
return Math.Max(n.Left?.Height ?? 0, n.Right?.Height ?? 0);
|
||||
}
|
||||
|
||||
public static Node Insert(Node? n, ulong seq, ref bool inserted, ref int nodes)
|
||||
{
|
||||
if (n == null)
|
||||
{
|
||||
var baseVal = (seq / NumEntries) * NumEntries;
|
||||
var newNode = new Node(baseVal);
|
||||
newNode.SetBit(seq, ref inserted);
|
||||
nodes++;
|
||||
return newNode;
|
||||
}
|
||||
|
||||
if (seq < n.Base)
|
||||
n.Left = Insert(n.Left, seq, ref inserted, ref nodes);
|
||||
else if (seq >= n.Base + NumEntries)
|
||||
n.Right = Insert(n.Right, seq, ref inserted, ref nodes);
|
||||
else
|
||||
n.SetBit(seq, ref inserted);
|
||||
|
||||
n.Height = MaxH(n) + 1;
|
||||
|
||||
var bf = BalanceFactor(n);
|
||||
if (bf > 1)
|
||||
{
|
||||
if (BalanceFactor(n.Left) < 0)
|
||||
n.Left = n.Left!.RotateLeft();
|
||||
return n.RotateRight();
|
||||
}
|
||||
if (bf < -1)
|
||||
{
|
||||
if (BalanceFactor(n.Right) > 0)
|
||||
n.Right = n.Right!.RotateRight();
|
||||
return n.RotateLeft();
|
||||
}
|
||||
return n;
|
||||
}
|
||||
|
||||
public static Node? Delete(Node? n, ulong seq, ref bool deleted, ref int nodes)
|
||||
{
|
||||
if (n == null) return null;
|
||||
|
||||
if (seq < n.Base)
|
||||
n.Left = Delete(n.Left, seq, ref deleted, ref nodes);
|
||||
else if (seq >= n.Base + NumEntries)
|
||||
n.Right = Delete(n.Right, seq, ref deleted, ref nodes);
|
||||
else if (n.ClearBit(seq, ref deleted))
|
||||
{
|
||||
nodes--;
|
||||
if (n.Left == null)
|
||||
n = n.Right;
|
||||
else if (n.Right == null)
|
||||
n = n.Left;
|
||||
else
|
||||
{
|
||||
n.Right = n.Right.InsertNodePrev(n.Left);
|
||||
n = n.Right;
|
||||
}
|
||||
}
|
||||
|
||||
if (n == null) return null;
|
||||
|
||||
n.Height = MaxH(n) + 1;
|
||||
|
||||
var bf = BalanceFactor(n);
|
||||
if (bf > 1)
|
||||
{
|
||||
if (BalanceFactor(n.Left) < 0)
|
||||
n.Left = n.Left!.RotateLeft();
|
||||
return n.RotateRight();
|
||||
}
|
||||
if (bf < -1)
|
||||
{
|
||||
if (BalanceFactor(n.Right) > 0)
|
||||
n.Right = n.Right!.RotateRight();
|
||||
return n.RotateLeft();
|
||||
}
|
||||
return n;
|
||||
}
|
||||
|
||||
private Node RotateLeft()
|
||||
{
|
||||
var r = Right;
|
||||
if (r != null)
|
||||
{
|
||||
Right = r.Left;
|
||||
r.Left = this;
|
||||
Height = MaxH(this) + 1;
|
||||
r.Height = MaxH(r) + 1;
|
||||
}
|
||||
else
|
||||
{
|
||||
Right = null;
|
||||
Height = MaxH(this) + 1;
|
||||
}
|
||||
return r!;
|
||||
}
|
||||
|
||||
private Node RotateRight()
|
||||
{
|
||||
var l = Left;
|
||||
if (l != null)
|
||||
{
|
||||
Left = l.Right;
|
||||
l.Right = this;
|
||||
Height = MaxH(this) + 1;
|
||||
l.Height = MaxH(l) + 1;
|
||||
}
|
||||
else
|
||||
{
|
||||
Left = null;
|
||||
Height = MaxH(this) + 1;
|
||||
}
|
||||
return l!;
|
||||
}
|
||||
|
||||
// Inserts nn into this subtree assuming nn.Base < all nodes in this subtree.
|
||||
public Node InsertNodePrev(Node nn)
|
||||
{
|
||||
if (Left == null)
|
||||
Left = nn;
|
||||
else
|
||||
Left = Left.InsertNodePrev(nn);
|
||||
|
||||
Height = MaxH(this) + 1;
|
||||
|
||||
var bf = BalanceFactor(this);
|
||||
if (bf > 1)
|
||||
{
|
||||
if (BalanceFactor(Left) < 0)
|
||||
Left = Left!.RotateLeft();
|
||||
return RotateRight();
|
||||
}
|
||||
if (bf < -1)
|
||||
{
|
||||
if (BalanceFactor(Right) > 0)
|
||||
Right = Right!.RotateRight();
|
||||
return RotateLeft();
|
||||
}
|
||||
return this;
|
||||
}
|
||||
|
||||
// Iterates nodes in tree order (pre-order: root → left → right).
|
||||
public static void NodeIter(Node? n, Action<Node> f)
|
||||
{
|
||||
if (n == null) return;
|
||||
f(n);
|
||||
NodeIter(n.Left, f);
|
||||
NodeIter(n.Right, f);
|
||||
}
|
||||
|
||||
// Iterates items in ascending order (in-order traversal).
|
||||
// Returns false if the callback returns false.
|
||||
public static bool Iter(Node? n, Func<ulong, bool> f)
|
||||
{
|
||||
if (n == null) return true;
|
||||
|
||||
if (!Iter(n.Left, f)) return false;
|
||||
|
||||
for (var num = n.Base; num < n.Base + NumEntries; num++)
|
||||
{
|
||||
if (n.ExistsBit(num))
|
||||
if (!f(num)) return false;
|
||||
}
|
||||
|
||||
return Iter(n.Right, f);
|
||||
}
|
||||
|
||||
public static Node? Clone(Node? src)
|
||||
{
|
||||
if (src == null) return null;
|
||||
var n = new Node(src.Base) { Height = src.Height };
|
||||
src.Bits.CopyTo(n.Bits, 0);
|
||||
n.Left = Clone(src.Left);
|
||||
n.Right = Clone(src.Right);
|
||||
return n;
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user