Add count

This commit is contained in:
Linus Björnstam 2026-02-11 12:56:48 +01:00
parent 79b5ab98aa
commit 280177a9cb
4 changed files with 151 additions and 59 deletions

View file

@ -37,55 +37,139 @@ namespace PersistentMap
}
}
public static Node<K> Set<K, V, TStrategy>(Node<K> root, K key, V value, TStrategy strategy, OwnerId owner)
where TStrategy : IKeyStrategy<K>
// Public API
public static Node<K> Set<K, V>(Node<K> root, K key, V value, IKeyStrategy<K> strategy, OwnerId owner, out bool countChanged)
{
root = root.EnsureEditable(owner);
var splitResult = InsertRecursive(root, key, value, strategy, owner, out countChanged);
if (splitResult != null)
{
var newRoot = new InternalNode<K>(owner);
newRoot.Children[0] = root;
newRoot.Keys[0] = splitResult.Separator;
newRoot.Children[1] = splitResult.NewNode;
newRoot.SetCount(1);
newRoot._prefixes![0] = strategy.GetPrefix(splitResult.Separator);
return newRoot;
}
return root;
}
// Recursive Helper
private static SplitResult<K>? InsertRecursive<K, V>(Node<K> node, K key, V value, IKeyStrategy<K> strategy, OwnerId owner, out bool added)
{
if (node.IsLeaf)
{
var leaf = node.AsLeaf<V>();
int index = FindIndex(leaf, key, strategy);
if (index < leaf.Header.Count && strategy.Compare(leaf.Keys[index], key) == 0)
{
// Root CoW check
root = root.EnsureEditable(owner);
var splitResult = InsertRecursive(root, key, value, strategy, owner);
if (splitResult != null)
{
// The root split. Create a new root.
var newRoot = new InternalNode<K>(owner);
newRoot.Children[0] = root;
newRoot.Keys[0] = splitResult.Separator;
newRoot.Children[1] = splitResult.NewNode;
newRoot.SetCount(1);
// Prefixes for internal nodes are derived from the separator keys
newRoot._prefixes![0] = strategy.GetPrefix(splitResult.Separator);
return newRoot;
}
return root;
leaf.Values[index] = value;
added = false; // Key existed, value updated. Count does not change.
return null;
}
public static Node<K> Remove<K, V, TStrategy>(Node<K> root, K key, TStrategy strategy, OwnerId owner)
where TStrategy : IKeyStrategy<K>
added = true; // New key. Count +1.
if (leaf.Header.Count < LeafNode<K, V>.Capacity)
{
root = root.EnsureEditable(owner);
bool rebalanceNeeded = RemoveRecursive<K,V,TStrategy>(root, key, strategy, owner);
// If root is internal and became empty (count 0), replace with its only child
if (rebalanceNeeded)
{
if (!root.IsLeaf)
{
// CHANGE: Use AsInternal()
var internalRoot = root.AsInternal();
if (internalRoot.Header.Count == 0)
{
return internalRoot.Children[0]!;
}
}
}
return root;
InsertIntoLeaf(leaf, index, key, value, strategy);
return null;
}
else
{
return SplitLeaf(leaf, index, key, value, strategy, owner);
}
}
else
{
var internalNode = node.AsInternal();
int index = FindRoutingIndex(internalNode, key, strategy);
var child = internalNode.Children[index]!.EnsureEditable(owner);
internalNode.Children[index] = child;
var split = InsertRecursive(child, key, value, strategy, owner, out added);
if (split != null)
{
if (internalNode.Header.Count < InternalNode<K>.Capacity - 1)
{
InsertIntoInternal(internalNode, index, split.Separator, split.NewNode, strategy);
return null;
}
else
{
return SplitInternal(internalNode, index, split.Separator, split.NewNode, strategy, owner);
}
}
return null;
}
}
// Public API
public static Node<K> Remove<K, V, TStrategy>(Node<K> root, K key, TStrategy strategy, OwnerId owner, out bool countChanged)
where TStrategy : IKeyStrategy<K>
{
root = root.EnsureEditable(owner);
bool rebalanceNeeded = RemoveRecursive<K, V, TStrategy>(root, key, strategy, owner, out countChanged);
if (rebalanceNeeded)
{
if (!root.IsLeaf)
{
var internalRoot = root.AsInternal();
if (internalRoot.Header.Count == 0)
{
return internalRoot.Children[0]!;
}
}
}
return root;
}
// Recursive Helper
private static bool RemoveRecursive<K, V, TStrategy>(Node<K> node, K key, TStrategy strategy, OwnerId owner, out bool removed)
where TStrategy : IKeyStrategy<K>
{
if (node.IsLeaf)
{
var leaf = node.AsLeaf<V>();
int index = FindIndex(leaf, key, strategy);
if (index < leaf.Header.Count && strategy.Compare(leaf.Keys[index], key) == 0)
{
RemoveFromLeaf(leaf, index);
removed = true; // Item removed. Count -1.
return leaf.Header.Count < LeafNode<K, V>.MergeThreshold;
}
removed = false; // Item not found.
return false;
}
else
{
var internalNode = node.AsInternal();
int index = FindRoutingIndex(internalNode, key, strategy);
var child = internalNode.Children[index]!.EnsureEditable(owner);
internalNode.Children[index] = child;
bool childUnderflow = RemoveRecursive<K, V, TStrategy>(child, key, strategy, owner, out removed);
if (removed && childUnderflow)
{
return HandleUnderflow<K, V, TStrategy>(internalNode, index, strategy, owner);
}
return false;
}
}
// ---------------------------------------------------------
// Internal Helpers: Search
@ -158,7 +242,7 @@ namespace PersistentMap
public SplitResult(Node<K> newNode, K separator) { NewNode = newNode; Separator = separator; }
}
private static SplitResult<K>? InsertRecursive<K, V, TStrategy>(Node<K> node, K key, V value, TStrategy strategy, OwnerId owner)
private static SplitResult<K>? InsertRecur2sive<K, V, TStrategy>(Node<K> node, K key, V value, TStrategy strategy, OwnerId owner)
where TStrategy : IKeyStrategy<K>
{
// DELETE the single FindIndex call at the top.
@ -197,7 +281,7 @@ namespace PersistentMap
var child = internalNode.Children[childIndex]!.EnsureEditable(owner);
internalNode.Children[childIndex] = child;
var split = InsertRecursive(child, key, value, strategy, owner);
var split = InsertRecursive(child, key, value, strategy, owner, out _);
if (split != null)
{

View file

@ -15,8 +15,8 @@ public sealed class PersistentMap<K, V, TStrategy> : BaseOrderedMap<K, V, TStrat
// OPTIMIZATION: Use OwnerId.None (0).
// This signals EnsureEditable to always copy the root path,
// producing a new tree of nodes that also have OwnerId.None.
var newRoot = BTreeFunctions.Set(_root, key, value, _strategy, OwnerId.None);
return new PersistentMap<K, V, TStrategy>(newRoot, _strategy, Count + 1);
var newRoot = BTreeFunctions.Set(_root, key, value, _strategy, OwnerId.None, out bool countChanged);
return new PersistentMap<K, V, TStrategy>(newRoot, _strategy, countChanged ? Count + 1 : Count);
}
public static PersistentMap<K, V, TStrategy> Empty(TStrategy strategy)
@ -32,7 +32,8 @@ public sealed class PersistentMap<K, V, TStrategy> : BaseOrderedMap<K, V, TStrat
public PersistentMap<K, V, TStrategy> Remove(K key)
{
var newRoot = BTreeFunctions.Remove<K,V, TStrategy>(_root, key, _strategy, OwnerId.None);
var newRoot = BTreeFunctions.Remove<K,V, TStrategy>(_root, key, _strategy, OwnerId.None, out bool removed);
if (!removed) return this;
return new PersistentMap<K, V, TStrategy>(newRoot, _strategy, Count - 1);
}

View file

@ -15,12 +15,14 @@ public sealed class TransientMap<K, V, TStrategy> : BaseOrderedMap<K, V, TStrate
public void Set(K key, V value)
{
_root = BTreeFunctions.Set(_root, key, value, _strategy, _transactionId);
_root = BTreeFunctions.Set(_root, key, value, _strategy, _transactionId, out bool countChanged);
if (countChanged) Count++;
}
public void Remove(K key)
{
_root = BTreeFunctions.Remove<K,V, TStrategy>(_root, key, _strategy, _transactionId);
_root = BTreeFunctions.Remove<K,V, TStrategy>(_root, key, _strategy, _transactionId, out bool removed);
if (removed) Count--;
}
public PersistentMap<K, V, TStrategy> ToPersistent()

View file

@ -21,7 +21,8 @@ public class BTreeFuzzTests
// CONFIGURATION
const int Iterations = 100_000; // High enough to trigger all splits/merges
const int KeyRange = 5000; // Small enough to cause frequent collisions
int Seed = Environment.TickCount;
const bool showOps = true;
int Seed = 2135974; // Environment.TickCount;
// ORACLES
var reference = new SortedDictionary<int, int>();
@ -42,6 +43,7 @@ public class BTreeFuzzTests
if (isInsert)
{
// ACTION: INSERT
if (showOps)Console.WriteLine("insert");
reference[key] = val;
subject.Set(key, val);
}
@ -50,6 +52,7 @@ public class BTreeFuzzTests
// ACTION: REMOVE
if (reference.ContainsKey(key))
{
if (showOps)Console.WriteLine("remove");
reference.Remove(key);
subject.Remove(key);
}
@ -63,10 +66,10 @@ public class BTreeFuzzTests
// 2. VERIFY CONSISTENCY (Expensive but necessary)
// We check consistency every 1000 ops or if the tree is small,
// to keep the test fast enough.
if (i % 1000 == 0 || reference.Count < 100)
{
//if (i % 1000 == 0 || reference.Count < 100)
//{
AssertConsistency(reference, subject);
}
//}
}
// Final check
@ -132,10 +135,11 @@ public class BTreeFuzzTests
private void AssertConsistency(SortedDictionary<int, int> expected, TransientMap<int, int, IntStrategy> actual)
{
// 1. Count
// if (expected.Count != actual.Count)
// {
// throw new Exception($"Count Mismatch! Expected {expected.Count}, Got {actual.Count}");
// }
if (expected.Count != actual.Count)
{
Console.WriteLine("BP");
throw new Exception($"Count Mismatch! Expected {expected.Count}, Got {actual.Count}");
}
// 2. Full Scan Verification
using var enumerator = actual.GetEnumerator();
@ -148,6 +152,7 @@ public class BTreeFuzzTests
if (enumerator.Current.Key != kvp.Key || enumerator.Current.Value != kvp.Value)
{
Console.WriteLine("BP");
throw new Exception($"Content Mismatch! Expected [{kvp.Key}:{kvp.Value}], Got [{enumerator.Current.Key}:{enumerator.Current.Value}]");
}
}