Immutable Trie Map in C#

TrieMap.cs

using System;
using System.Collections.Generic;

namespace mastermind.core.immutable
{
    public interface ITrieMap<K,V>
    {
        int GetHashCode(K key);
        V this[K key] { get; }
        IEnumerable<LeafNode<K, V>> All(); 
        ITrieMap<K, V> Add(K key, V value);
        ITrieMap<K, V> Remove(K key);
    }

    public interface ITrieNode<K, V>
    {
        bool TryGetValue(int hashCode, out V value);
        V GetValue(int hashCode);
        int GetValueHashCode(int hashCode);
        IEnumerable<LeafNode<K, V>> All();
        ITrieNode<K,V> GetNode(int hashCode);
        ITrieNode<K, V> Add(int hashCode, K key, V value, int width = 0, int depth = 0);
        ITrieNode<K, V> Remove(int hashCode);
    }

    public static class TrieIndexer
    {
        public static int GetIndex(int hash, int width, int depth)
        {
            if (width == 4)
                return (hash >> (depth * 2)) & 3;
            if (width == 8)
                return (hash >> (depth * 3)) & 7;
            if(width == 16)
                return (hash >> (depth * 4)) & 15;
            if(width == 32)
                return (hash >> (depth * 5)) & 31;
            throw new Exception("Width of "+width+" not supported in TrieMap");
        }
    }

    public class LeafNode<K, V> : ITrieNode<K, V>
    {
        private static readonly EqualityComparer<V> _valueEqualityComparer = EqualityComparer<V>.Default;

        public readonly int Hash;
        public readonly K Key;
        public readonly V Value;

        public LeafNode(int hash, K key, V value)
        {
            Hash = hash;
            Key = key;
            Value = value;
        }

        public bool TryGetValue(int hash, out V value)
        {
            bool success = Hash == hash;
            value = success ? Value : default(V);
            return success;
        }

        public V GetValue(int hash)
        {
            return Hash == hash
                ? Value
                : default (V);
        }

        public int GetValueHashCode(int hash)
        {
            return Hash == hash
                ? Value.GetHashCode()
                : 0;
        }

        public IEnumerable<LeafNode<K, V>> All()
        {
            yield return this;
        }

        public ITrieNode<K, V> GetNode(int hash)
        {
            return this;
        }

        public ITrieNode<K, V> Add(int hash, K key, V value, int width, int depth = 0)
        {
            if(hash == Hash)
                return _valueEqualityComparer.Equals(value, Value)
                    ? this
                    : new LeafNode<K, V>(hash, key, value);

            var targetDepth = depth;
            while (TrieIndexer.GetIndex(hash, width, targetDepth) == TrieIndexer.GetIndex(Hash, width, targetDepth))
                targetDepth += 1;

            var newBranch = new BranchNode<K, V>(width, targetDepth);
            newBranch.Values[TrieIndexer.GetIndex(Hash, width, targetDepth)] = this;
            newBranch.Values[TrieIndexer.GetIndex(hash, width, targetDepth)] = new LeafNode<K, V>(hash, key, value);

            for (int i = targetDepth - 1; i >= depth; i--)
            {
                var branch = new BranchNode<K, V>(width, i);
                branch.Values[TrieIndexer.GetIndex(Hash, width, i)] = newBranch;
                newBranch = branch;
            }
            return newBranch;
        }

        public ITrieNode<K, V> Remove(int hash)
        {
            if(hash == Hash)
                return null;
            return this;
        }
    }

    class BranchNode<K,V> : ITrieNode<K,V>
    {
        public readonly ITrieNode<K,V>[] Values;
        public readonly int Depth;

        public BranchNode(int width, int depth)
        {
            Values = new ITrieNode<K, V>[width];
            Depth = depth;
        }

        public IEnumerable<LeafNode<K, V>> All()
        {
            foreach (var value in Values)
            {
                if (value != null)
                {
                    foreach (var trieNode in value.All())
                    {
                        yield return trieNode;
                    }
                }
            }
        }

        public ITrieNode<K, V> GetNode(int hash)
        {
            int index = TrieIndexer.GetIndex(hash, Values.Length, Depth);
            var node = Values[index];
            return node != null
                ? node.GetNode(hash)
                : null;
        }

        public ITrieNode<K, V> Add(int hash, K key, V value, int width = 0, int depth = 0)
        {
            width = Values.Length;
            int index = TrieIndexer.GetIndex(hash, width, depth);
            var existingNode = Values[index];
            var newNode = existingNode != null
                ? existingNode.Add(hash, key, value, width, Depth + 1)
                : new LeafNode<K,V>(hash, key, value);

            if (newNode == existingNode)
                return this;
            
            var newBranch = new BranchNode<K, V>(width, Depth);
            Array.Copy(Values, newBranch.Values, width);
            newBranch.Values[index] = newNode;
            return newBranch;
        }

        public ITrieNode<K,V> Remove(int hash)
        {
            int width = Values.Length;
            int index = TrieIndexer.GetIndex(hash, width, Depth);
            var oldValue = Values[index];
            if (oldValue == null)
                return this;
            var newValue = Values[index].Remove(hash);
            if (newValue == oldValue)
                return this;
            var newBranch = new BranchNode<K, V>(width, Depth);
            Array.Copy(Values, newBranch.Values, width);
            newBranch.Values[index] = newValue;
            return newBranch;
        }

        public bool TryGetValue(int hash, out V value)
        {
            var node = GetNode(hash);
            if (node != null)
                return node.TryGetValue(hash, out value);
            value = default (V);
            return false;
        }

        public V GetValue(int hash)
        {
            var node = GetNode(hash);
            return (node != null)
                ? node.GetValue(hash)
                : default (V);
        }

        public int GetValueHashCode(int hash)
        {
            var node = GetNode(hash);
            return (node != null)
                ? node.GetValueHashCode(hash)
                : 0;
        }
    }

    public class TrieMap<K,V> : ITrieMap<K,V>
    {
        public static IEqualityComparer<K> EqualityComparer { get; set; }
        private readonly ITrieNode<K,V> _trieRoot;
		
        public static ITrieMap<K, V> Empty(int width, IEqualityComparer<K> equalityComparer)
        {
            return new TrieMap<K, V>(new BranchNode<K, V>(width, 0), equalityComparer);
        }

        private TrieMap(ITrieNode<K, V> trieRoot, IEqualityComparer<K> equalityComparer)
        {
            _trieRoot = trieRoot;
            EqualityComparer = equalityComparer;
        }

        public int GetHashCode(K key)
        {
            return _trieRoot.GetValueHashCode(EqualityComparer.GetHashCode(key));
        }

        public V this[K key]
        {
            get { return _trieRoot.GetValue(EqualityComparer.GetHashCode(key)); }
        }

        public IEnumerable<LeafNode<K, V>> All()
        {
            return _trieRoot.All();
        }

        public ITrieMap<K,V> Add(K key, V value)
        {
            int hashCode = EqualityComparer.GetHashCode(key);
            var newRoot = _trieRoot.Add(hashCode, key, value);
            return _trieRoot == newRoot
                ? this
                : new TrieMap<K, V>(newRoot, EqualityComparer);
        }

        public ITrieMap<K, V> Remove(K key)
        {
            var newRoot = _trieRoot.Remove(EqualityComparer.GetHashCode(key));
            return newRoot == _trieRoot
                ? this
                : new TrieMap<K, V>(newRoot, EqualityComparer);
        }
    }
}