PersistentMap/PersistentOrderedMap/Iterator.cs

using System.Collections;
using System.Runtime.CompilerServices;

namespace PersistentOrderedMap;



public struct BTreeEnumerable<K, V, TStrategy> : IEnumerable<KeyValuePair<K, V>>
where TStrategy : IKeyStrategy<K>
{
    private readonly Node<K> _root;
    private readonly TStrategy _strategy;
    private readonly K _min, _max;
    private readonly bool _hasMin, _hasMax;

    public BTreeEnumerable(Node<K> root, TStrategy strategy, bool hasMin, K min, bool hasMax, K max)
    {
        _root = root; _strategy = strategy;
        _hasMin = hasMin; _min = min;
        _hasMax = hasMax; _max = max;
    }

    public BTreeEnumerator<K, V, TStrategy> GetEnumerator()
    {
        return new BTreeEnumerator<K, V, TStrategy>(_root, _strategy, _hasMin, _min, _hasMax, _max);
    }

    IEnumerator<KeyValuePair<K, V>> IEnumerable<KeyValuePair<K, V>>.GetEnumerator() => GetEnumerator();
    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
}

// Fixed-size buffer for the path. 
// Depth 16 * 32 (branching factor) = Exabytes of capacity.
// The B tree is, theoretically, not limited by the size of an int, like
// all int-indexed data structures (or bit partitioned ones).
// This should be enough for anyone.
[InlineArray(16)]
internal struct IterNodeBuffer<K>
{
    private Node<K> _element0;
}

[InlineArray(16)]
internal struct IterIndexBuffer<K>
{
    private int _element0;
}

public struct BTreeEnumerator<K, V, TStrategy> : IEnumerator<KeyValuePair<K, V>>
        where TStrategy : IKeyStrategy<K>
    {
        private readonly TStrategy _strategy;
        private readonly Node<K> _root;
        
        // --- BOUNDS ---
        private readonly bool _hasMax;
        private readonly K _maxKey;
        private readonly bool _hasMin;
        private readonly K _minKey;

        // --- INLINE STACK ---
        private IterNodeBuffer<K> _nodeStack; 
        private IterIndexBuffer<K> _indexStack;
        private int _depth;

        // --- STATE ---
        private LeafNode<K, V>? _currentLeaf;
        private int _currentLeafIndex;
        private KeyValuePair<K, V> _current;

        // Unified Constructor
        // We use boolean flags because 'K' might be a struct where 'null' is impossible.
        public BTreeEnumerator(Node<K> root, TStrategy strategy, bool hasMin, K minKey, bool hasMax, K maxKey)
        {
            _root = root;
            _strategy = strategy;
            _hasMax = hasMax;
            _maxKey = maxKey;
            _hasMin = hasMin;
            _minKey = minKey;
            
            _nodeStack = new IterNodeBuffer<K>();
            _indexStack = new IterIndexBuffer<K>(); // Explicit struct init
            _depth = 0;
            _currentLeaf = null;
            _currentLeafIndex = -1;
            _current = default;

            if (root != null)
            {
                if (hasMin)
                {
                    Seek(minKey);
                }
                else
                {
                    DiveLeft();
                }
            }
        }

        // Logic 1: Unbounded Start (Go to very first item)
        private void DiveLeft()
        {
            Node<K> node = _root;
            _depth = 0;

            while (!node.IsLeaf)
            {
                var internalNode = node.AsInternal();
                _nodeStack[_depth] = internalNode;
                _indexStack[_depth] = 0; // Always take left-most child
                _depth++;
                node = internalNode.Children[0]!;
            }

            _currentLeaf = node.AsLeaf<V>();
            _currentLeafIndex = -1; // Position before the first element (0)
        }

        // Logic 2: Bounded Start (Go to specific key)
        private void Seek(K key)
        {
            Node<K> node = _root;
            _depth = 0;
            long keyPrefix = _strategy.UsesPrefixes ? _strategy.GetPrefix(key) : 0;
    

            // Dive using Routing
            while (!node.IsLeaf)
            {
                var internalNode = node.AsInternal();
                int idx = BTreeFunctions.FindRoutingIndex<K, TStrategy>(internalNode, key, keyPrefix,  _strategy);
                
                _nodeStack[_depth] = internalNode;
                _indexStack[_depth] = idx;
                _depth++;

                node = internalNode.Children[idx]!;
            }

            // Find index in Leaf
            _currentLeaf = node.AsLeaf<V>();
            int index = BTreeFunctions.FindIndex<K,V, TStrategy>(_currentLeaf, key, keyPrefix, _strategy);
            
            // Set position to (index - 1) so that the first MoveNext() lands on 'index'
            _currentLeafIndex = index - 1;
        }

        public bool MoveNext()
        {
            if (_currentLeaf == null) return false;

            // 1. Try to advance in current leaf
            if (++_currentLeafIndex < _currentLeaf.Header.Count)
            {
                // OPTIMIZATION: Check Max Bound (if active)
                if (_hasMax)
                {
                     // If Current Key > Max Key, we are done.
                     if (_strategy.Compare(_currentLeaf.Keys[_currentLeafIndex], _maxKey) > 0)
                     {
                         _currentLeaf = null; // Close iterator
                         return false;
                     }
                }

                _current = new KeyValuePair<K, V>(_currentLeaf.Keys[_currentLeafIndex], _currentLeaf.Values[_currentLeafIndex]);
                return true;
            }

            // 2. Leaf exhausted. Find next leaf.
            if (FindNextLeaf())
            {
                // Found new leaf, index reset to 0. 
                // Check Max Bound immediately for the first item
                if (_hasMax)
                {
                    if (_strategy.Compare(_currentLeaf!.Keys[0], _maxKey) > 0)
                    {
                        _currentLeaf = null;
                        return false;
                    }
                }

                _current = new KeyValuePair<K, V>(_currentLeaf.Keys[0], _currentLeaf.Values[0]);
                return true;
            }

            return false;
        }

        private bool FindNextLeaf()
        {
            while (_depth > 0)
            {
                _depth--;
                var internalNode = _nodeStack[_depth].AsInternal();
                int currentIndex = _indexStack[_depth];

                if (currentIndex < internalNode.Header.Count)
                {
                    int nextIndex = currentIndex + 1;
                    _indexStack[_depth] = nextIndex;
                    _depth++;

                    Node<K> node = internalNode.Children[nextIndex]!;
                    while (!node.IsLeaf)
                    {
                        _nodeStack[_depth] = node;
                        _indexStack[_depth] = 0;
                        _depth++;
                        node = node.AsInternal().Children[0]!;
                    }

                    _currentLeaf = node.AsLeaf<V>();
                    _currentLeafIndex = 0;
                    return true;
                }
            }
            return false;
        }

        public KeyValuePair<K, V> Current => _current;
        object IEnumerator.Current => _current;
        public void Reset()
        {
            // 1. Clear current state
            _depth = 0;
            _currentLeaf = null;
            _currentLeafIndex = -1;
            _current = default;

            // 2. Re-initialize based on how the iterator was created
            if (_root != null)
            {
                if (_hasMin)
                {
                    // If we had a start range, find it again
                    Seek(_minKey);
                }
                else
                {
                    // Otherwise, go back to the very first leaf
                    DiveLeft();
                }
            }
        }
        public void Dispose() { }
    }