ln.bson/ln.bson.storage/index/IndexTree.cs

using System;
using System.Collections.Generic;
using System.Threading;

namespace ln.bson.storage.index
{
    public class IndexTree
    {
        private IndexLeaf _root = new IndexLeaf();

        public void Query(string path, QueryOperator queryOperator, byte[] value, ISet<long> resultSet) => Query(path.Split('.'), queryOperator, value, resultSet);
        public void Query(Span<string> path, QueryOperator queryOperator, byte[] value, ISet<long> resultSet)
        {
            if (!_root.TryGetLeaf(path, out IndexLeaf indexLeaf))
                throw new KeyNotFoundException();
            
            indexLeaf.Index.Query(queryOperator, value, resultSet);
        }

        public void AddDocument(long key, BsonDocument bsonDocument) => _root.AddValue(key, bsonDocument);
        public void RemoveDocument(long key) => _root.RemoveValue(key);
        public void ReplaceDocument(long oldKey, long newKey, BsonDocument newDocument) =>
            _root.ReplaceValue(oldKey, newKey, newDocument);

        public bool EnsureIndex(string path, Comparison<byte[]> comparison) =>
            EnsureIndex(path, comparison, out IBsonIndex bsonIndex);
        public bool EnsureIndex(string path, Comparison<byte[]> comparison, out IBsonIndex bsonIndex)
        {
            _root.EnsureLeaf(path.Split('.',StringSplitOptions.None), out IndexLeaf indexLeaf);
            return indexLeaf.EnsureIndex(comparison, out bsonIndex);
        }

        public void Clear() => _root.Clear();

        class IndexLeaf
        {
            private Dictionary<string, IndexLeaf> _children = new Dictionary<string, IndexLeaf>();
            private IBsonIndex _index;

            public IBsonIndex Index => _index;

            public void EnsureLeaf(Span<string> path, out IndexLeaf indexLeaf)
            {
                if (path.IsEmpty)
                    indexLeaf = this;
                else
                {
                    if (!_children.TryGetValue(path[0], out IndexLeaf childLeaf))
                    {
                        childLeaf = new IndexLeaf();
                        _children.Add(path[0], childLeaf);
                    }
                    childLeaf.EnsureLeaf(path.Slice(1), out indexLeaf);
                }
            }
            public bool TryGetLeaf(Span<string> path, out IndexLeaf indexLeaf)
            {
                if (path.IsEmpty)
                {
                    indexLeaf = this;
                    return true;
                }
                else
                {
                    if (!_children.TryGetValue(path[0], out IndexLeaf childLeaf))
                    {
                        indexLeaf = null;
                        return false;
                    }

                    return childLeaf.TryGetLeaf(path.Slice(1), out indexLeaf);
                }
            }

            public bool EnsureIndex() => EnsureIndex(BsonIndex.CompareBytesLittleEndian, out IBsonIndex bsonIndex);
            public bool EnsureIndex(Comparison<byte[]> comparison, out IBsonIndex bsonIndex)
            {
                if (_index is null)
                {
                    bsonIndex = _index = new BsonIndex(comparison);
                    return true;
                }

                bsonIndex = _index;
                return false;
            }
            
            public void AddValue(long key, BsonValue bsonValue)
            {
                if (_index is not null)
                    _index.AddValue(key, bsonValue);

                if (bsonValue is BsonDocument bsonDocument)
                {
                    foreach (KeyValuePair<string, IndexLeaf> child in _children)
                    {
                        if (bsonDocument.Contains(child.Key))
                            child.Value.AddValue(key, bsonDocument[child.Key]);
                    }
                }
            }

            public void RemoveValue(long key)
            {
                if (_index is not null)
                    _index.RemoveValue(key);
                
                foreach (KeyValuePair<string, IndexLeaf> child in _children)
                    child.Value.RemoveValue(key);
            }

            public void ReplaceValue(long oldKey, long newKey, BsonValue newValue)
            {
                if (_index is not null)
                    _index.ReplaceValue(oldKey, newKey, newValue);

                if (newValue is BsonDocument newDocument)
                {
                    foreach (KeyValuePair<string, IndexLeaf> child in _children)
                    {
                        if (newDocument.Contains(child.Key))
                            child.Value.ReplaceValue(oldKey, newKey, newDocument[child.Key]);
                        else
                            child.Value.RemoveValue(oldKey);
                    }
                }
                else
                {
                    foreach (KeyValuePair<string, IndexLeaf> child in _children)
                        child.Value.RemoveValue(oldKey);                    
                }
            }

            public void Clear()
            {
                if (_index is not null)
                    _index.Clear();
                
                foreach (KeyValuePair<string, IndexLeaf> child in _children)
                    child.Value.Clear();
            }
            
        }

    }
}