Complicated Hello World

using System.Collections;

namespace HelloWorld;

public class StringTrieMap<T> : TrieMap<string, char, StringKeyHandler, T> { }

public class StringKeyHandler : ITrieKeyHandler<string, char>
{
    public static IEnumerable<char> Decompose(string key) => key;

    public static string CreateKeyObject(IEnumerable<char> elements) =>
        string.Concat(elements);
}

public class TrieMap<TKey, TKeyElement, TKeyHandler, TValue>
    : ICollection<KeyValuePair<TKey, TValue>>,
        IDictionary<TKey, TValue>,
        IEnumerable<KeyValuePair<TKey, TValue>>,
        IReadOnlyCollection<KeyValuePair<TKey, TValue>>,
        IReadOnlyDictionary<TKey, TValue>
    where TKeyElement : notnull
    where TKeyHandler : ITrieKeyHandler<TKey, TKeyElement>
{
    public int Count => _root.Count;
    public bool IsReadOnly => false;

    public TValue this[TKey key]
    {
        get
        {
            ArgumentNullException.ThrowIfNull(key, nameof(key));

            return _root.Get(TKeyHandler.Decompose(key));
        }
        set
        {
            ArgumentNullException.ThrowIfNull(key, nameof(key));

            _root.Replace(
                TKeyHandler.Decompose(key).GetEnumerator(),
                value,
                _comparer
            );
        }
    }

    public ICollection<TKey> Keys
    {
        get
        {
            var keys = new List<TKey>();
            var keyStack = new Stack<TKeyElement>();

            foreach (var _ in _root.EnumerateKeys(keyStack))
            {
                keys.Add(TKeyHandler.CreateKeyObject(keyStack.Reverse()));
            }

            return keys;
        }
    }

    IEnumerable<TKey> IReadOnlyDictionary<TKey, TValue>.Keys
    {
        get
        {
            var keyStack = new Stack<TKeyElement>();

            foreach (var _ in WrapEnumerable(_root.EnumerateKeys(keyStack)))
            {
                yield return TKeyHandler.CreateKeyObject(keyStack.Reverse());
            }
        }
    }

    public ICollection<TValue> Values =>
        _root.EnumerateKeyValuePairs(new()).ToList();

    IEnumerable<TValue> IReadOnlyDictionary<TKey, TValue>.Values =>
        WrapEnumerable(_root.EnumerateKeyValuePairs(new()));

    private TrieNode<
        TKeyElement,
        TValue,
        TValue,
        TrieMapStorageAccessor<TValue>
    > _root;
    private IComparer<TKeyElement> _comparer;
    private int _version = 0;

    public TrieMap(IComparer<TKeyElement>? comparer = null)
    {
        _comparer = comparer ?? Comparer<TKeyElement>.Default;
        _root = new(_comparer);
    }

    public TrieMap(
        IEnumerable<KeyValuePair<TKey, TValue>> collection,
        IComparer<TKeyElement>? comparer = null
    )
        : this(comparer)
    {
        foreach (var item in collection)
        {
            Add(item.Key, item.Value);
        }
    }

    public void CopyTo(KeyValuePair<TKey, TValue>[] array, int arrayIndex)
    {
        ArgumentNullException.ThrowIfNull(array, nameof(array));
        ArgumentOutOfRangeException.ThrowIfNegative(
            arrayIndex,
            nameof(arrayIndex)
        );

        if (array.Length - arrayIndex < Count)
        {
            throw new ArgumentException("Not enough room to copy to.");
        }

        var i = arrayIndex;

        foreach (var item in this)
        {
            array[i++] = item;
        }
    }

    public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
    {
        var keyStack = new Stack<TKeyElement>();

        foreach (
            var value in WrapEnumerable(_root.EnumerateKeyValuePairs(keyStack))
        )
        {
            yield return new KeyValuePair<TKey, TValue>(
                TKeyHandler.CreateKeyObject(keyStack.Reverse()),
                value
            );
        }
    }

    IEnumerator IEnumerable.GetEnumerator() =>
        ((IEnumerable<KeyValuePair<TKey, TValue>>)this).GetEnumerator();

    public bool TryGetValue(TKey key, out TValue value)
    {
        ArgumentNullException.ThrowIfNull(key, nameof(key));

        return _root.TryGetValue(TKeyHandler.Decompose(key), out value);
    }

    public bool Contains(KeyValuePair<TKey, TValue> item)
    {
        ArgumentNullException.ThrowIfNull(item, nameof(item));
        ArgumentNullException.ThrowIfNull(item.Key, nameof(item.Key));

        return _root.Contains(TKeyHandler.Decompose(item.Key), item.Value);
    }

    public bool ContainsKey(TKey key)
    {
        ArgumentNullException.ThrowIfNull(key, nameof(key));

        return _root.ContainsKey(TKeyHandler.Decompose(key));
    }

    public void Add(KeyValuePair<TKey, TValue> item)
    {
        ArgumentNullException.ThrowIfNull(item, nameof(item));

        Add(item.Key, item.Value);
    }

    public void Add(TKey key, TValue value)
    {
        ArgumentNullException.ThrowIfNull(key, nameof(key));

        var oldCount = Count;
        _root.Add(
            TKeyHandler.Decompose(key).GetEnumerator(),
            value,
            _comparer
        );
        UpdateVersion(oldCount);
    }

    public bool Remove(KeyValuePair<TKey, TValue> item)
    {
        ArgumentNullException.ThrowIfNull(item, nameof(item));
        ArgumentNullException.ThrowIfNull(item.Key, nameof(item.Key));

        var oldCount = Count;
        _root.Remove(
            TKeyHandler.Decompose(item.Key).GetEnumerator(),
            item.Value
        );
        return UpdateVersion(oldCount);
    }

    public bool Remove(TKey key)
    {
        ArgumentNullException.ThrowIfNull(key, nameof(key));

        var oldCount = Count;
        _root.RemoveKey(TKeyHandler.Decompose(key).GetEnumerator());
        return UpdateVersion(oldCount);
    }

    public void Clear()
    {
        var oldCount = Count;
        _root.Clear();
        UpdateVersion(oldCount);
    }

    private IEnumerable<T> WrapEnumerable<T>(IEnumerable<T> enumerable)
    {
        var initialVersion = _version;

        foreach (var value in enumerable)
        {
            yield return _version == initialVersion
                ? value
                : throw new InvalidOperationException(
                    "Collection was modified after the enumerator was instantiated."
                );
        }
    }

    private bool UpdateVersion(int oldCount)
    {
        if (_root.Count != oldCount)
        {
            ++_version;
            return true;
        }
        else
        {
            return false;
        }
    }
}

internal class TrieMapStorageAccessor<TValue>
    : ITrieNodeStorageAccessor<TValue, TValue>
{
    public static void Add(
        TValue value,
        ref TValue? storage,
        ref bool hasStorage,
        ref int count
    )
    {
        if (hasStorage)
        {
            throw new ArgumentException(
                "An element with the given key already exists."
            );
        }

        storage = value;
        hasStorage = true;
        ++count;
    }

    public static void Remove(
        TValue value,
        ref TValue? storage,
        ref bool hasStorage,
        ref int count
    )
    {
        if (!hasStorage || StorageEquals(storage!, value))
        {
            return;
        }

        hasStorage = false;
        --count;
        storage = default;
    }

    public static bool Contains(TValue storage, TValue value) =>
        Equals(storage, value);

    public static IEnumerable<TValue> Enumerate(TValue storage)
    {
        yield return storage;
    }

    public static TValue Clone(TValue storage) => storage;

    public static bool StorageEquals(TValue a, TValue b) => Equals(a, b);

    public static int Count(TValue storage) => 1;
}

internal class TrieNode<TKeyElement, TValue, TStore, TStorageAccessor>
    where TKeyElement : notnull
    where TStorageAccessor : ITrieNodeStorageAccessor<TValue, TStore>
{
    public int Count => _count;

    private TStore? _storage;
    private bool _hasStorage;
    private int _count;
    private SortedDictionary<
        TKeyElement,
        TrieNode<TKeyElement, TValue, TStore, TStorageAccessor>
    > _children;

    public TrieNode(IComparer<TKeyElement> comparer)
    {
        _storage = default;
        _hasStorage = false;
        _count = 0;
        _children = new(comparer);
    }

    public TrieNode(
        TrieNode<TKeyElement, TValue, TStore, TStorageAccessor> other,
        IComparer<TKeyElement> comparer
    )
        : this(comparer)
    {
        if (other._hasStorage)
        {
            _storage = TStorageAccessor.Clone(other._storage!);
            _hasStorage = true;
        }

        foreach (var (keyElement, child) in other._children)
        {
            _children[keyElement] = new TrieNode<
                TKeyElement,
                TValue,
                TStore,
                TStorageAccessor
            >(child, comparer);
        }

        _count = other.Count;
    }

    public bool Contains(IEnumerable<TKeyElement> key, TValue value)
    {
        var node = this;

        foreach (var keyElement in key)
        {
            if (!node._children.TryGetValue(keyElement, out node))
            {
                return false;
            }
        }

        return node._hasStorage
            && TStorageAccessor.Contains(node._storage!, value);
    }

    public bool ContainsKey(IEnumerable<TKeyElement> key)
    {
        var node = this;

        foreach (var keyElement in key)
        {
            if (!node._children.TryGetValue(keyElement, out node))
            {
                return false;
            }
        }

        return node._hasStorage;
    }

    // This method is currently unused, but I might need it later
    public IEnumerable<TValue> GetAll(IEnumerable<TKeyElement> key)
    {
        var node = this;

        foreach (var keyElement in key)
        {
            if (!node._children.TryGetValue(keyElement, out node))
            {
                yield break;
            }
        }

        if (!node._hasStorage)
        {
            yield break;
        }

        foreach (var value in TStorageAccessor.Enumerate(node._storage!))
        {
            yield return value;
        }
    }

    public IEnumerable<Nothing> EnumerateKeys(Stack<TKeyElement> prefix)
    {
        if (_hasStorage)
        {
            yield return default;
        }

        foreach (var (keyElement, child) in _children)
        {
            prefix.Push(keyElement);

            foreach (var value in child.EnumerateKeys(prefix))
            {
                yield return value;
            }

            prefix.Pop();
        }
    }

    public IEnumerable<TValue> EnumerateKeyValuePairs(
        Stack<TKeyElement> prefix
    )
    {
        if (_hasStorage)
        {
            foreach (var value in TStorageAccessor.Enumerate(_storage!))
            {
                yield return value;
            }
        }

        foreach (var (keyElement, child) in _children)
        {
            prefix.Push(keyElement);

            foreach (var value in child.EnumerateKeyValuePairs(prefix))
            {
                yield return value;
            }

            prefix.Pop();
        }
    }

    public TStore Get(IEnumerable<TKeyElement> key)
    {
        var node = this;

        foreach (var keyElement in key)
        {
            if (!node._children.TryGetValue(keyElement, out node))
            {
                throw new KeyNotFoundException(
                    "The given key was not present."
                );
            }
        }

        if (!node._hasStorage)
        {
            throw new KeyNotFoundException("The given key was not present.");
        }

        return node._storage!;
    }

    public bool TryGetValue(IEnumerable<TKeyElement> key, out TStore value)
    {
        var node = this;

        foreach (var keyElement in key)
        {
            if (!node._children.TryGetValue(keyElement, out node))
            {
                value = default!;
                return false;
            }
        }

        if (!node._hasStorage)
        {
            throw new KeyNotFoundException("The given key was not present.");
        }

        value = node._storage!;
        return true;
    }

    public void Clear()
    {
        _hasStorage = false;
        _count = 0;
        _storage = default;
        _children.Clear();
    }

    public void Add(
        IEnumerator<TKeyElement> key,
        TValue value,
        IComparer<TKeyElement> comparer
    )
    {
        if (!key.MoveNext())
        {
            TStorageAccessor.Add(
                value,
                ref _storage,
                ref _hasStorage,
                ref _count
            );
            return;
        }

        var keyElement = key.Current;

        if (!_children.TryGetValue(keyElement, out var child))
        {
            child = new TrieNode<
                TKeyElement,
                TValue,
                TStore,
                TStorageAccessor
            >(comparer);
            _children[keyElement] = child;
        }

        var oldCount = child._count;
        child.Add(key, value, comparer);
        _count += child._count - oldCount;
    }

    public void Replace(
        IEnumerator<TKeyElement> key,
        TStore storage,
        IComparer<TKeyElement> comparer
    )
    {
        if (!key.MoveNext())
        {
            if (_hasStorage)
            {
                _count -= TStorageAccessor.Count(_storage!);
                _storage = storage;
            }
            else
            {
                _storage = storage;
                _hasStorage = true;
            }

            _count += TStorageAccessor.Count(_storage);

            return;
        }

        var keyElement = key.Current;

        if (!_children.TryGetValue(keyElement, out var child))
        {
            child = new TrieNode<
                TKeyElement,
                TValue,
                TStore,
                TStorageAccessor
            >(comparer);
            _children[keyElement] = child;
        }

        var oldCount = child._count;
        child.Replace(key, storage, comparer);
        _count += child._count - oldCount;
    }

    public void Remove(IEnumerator<TKeyElement> key, TValue value)
    {
        if (!key.MoveNext())
        {
            TStorageAccessor.Remove(
                value,
                ref _storage,
                ref _hasStorage,
                ref _count
            );

            return;
        }

        var keyElement = key.Current;

        if (!_children.TryGetValue(keyElement, out var child))
        {
            return;
        }

        var oldCount = child._count;
        child.Remove(key, value);
        _count += child._count - oldCount;

        if (child._count == 0)
        {
            _children.Remove(keyElement);
        }
    }

    public void RemoveKey(IEnumerator<TKeyElement> key)
    {
        if (!key.MoveNext())
        {
            if (_hasStorage)
            {
                _hasStorage = false;
                _count -= TStorageAccessor.Count(_storage!);
                _storage = default;
            }

            return;
        }

        var keyElement = key.Current;

        if (!_children.TryGetValue(keyElement, out var child))
        {
            return;
        }

        var oldCount = child._count;
        child.RemoveKey(key);
        _count += child._count - oldCount;

        if (child._count == 0)
        {
            _children.Remove(keyElement);
        }
    }
}

public interface ITrieKeyHandler<TKey, TKeyElement>
{
    static abstract IEnumerable<TKeyElement> Decompose(TKey key);

    static abstract TKey CreateKeyObject(IEnumerable<TKeyElement> elements);
}

internal interface ITrieNodeStorageAccessor<TValue, TStore>
{
    static abstract void Add(
        TValue value,
        ref TStore? storage,
        ref bool hasStorage,
        ref int count
    );

    static abstract void Remove(
        TValue value,
        ref TStore? storage,
        ref bool hasStorage,
        ref int count
    );

    static abstract bool Contains(TStore storage, TValue value);

    static abstract IEnumerable<TValue> Enumerate(TStore storage);

    static abstract TStore Clone(TStore storage);

    static abstract bool StorageEquals(TStore a, TStore b);

    static abstract int Count(TStore storage);
}

internal struct Nothing { }

internal class Program
{
    private static void Main()
    {
        var wordsWithPunctuation =
            (IReadOnlyDictionary<string, char?>)
                new StringTrieMap<char?>()
                {
                    { "Hello", ',' },
                    { "World", null }
                };

        // IMPORTANT: Only works because hello world is in alphabetic order
        Console.WriteLine(
            string.Join(
                ' ',
                wordsWithPunctuation.Select(item =>
                    item.Value is null ? item.Key : $"{item.Key}{item.Value}"
                )
            )
        );
    }
}