C# Trie (improved)

using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;

namespace MyCollections
{
    public class Trie<TValue> : IDictionary<string, TValue>
    {
        private int _count;
        private readonly TrieNode _rootNode;

        public Trie()
        {
            _rootNode = new TrieNode('\0', null);
        }

        #region Core implementation

        private IEnumerable<KeyValuePair<string, TValue>> Enumerate(TrieNode root, string prefix)
        {
            string key;
            if (root == _rootNode)
                key = prefix + string.Empty;
            else
                key = prefix + root.PartialKey;

            if (root.HasValue)
                yield return new KeyValuePair<string, TValue>(key, root.Value);

            foreach (var kvp in root.Children)
            {
                foreach (var entry in Enumerate(kvp.Value, key))
                {
                    yield return entry;
                }
            }
        }

        private static TrieNode FindNode(TrieNode root, char[] key, int depth, bool create, out bool created)
        {
            created = false;

            if (depth == key.Length)
                return root;

            TrieNode node;
            if (root.Children.TryGetValue(key[depth], out node))
                return FindNode(node, key, depth + 1, create, out created);

            if (create)
            {
                var current = root;
                node = null;
                for (int i = depth; i < key.Length; i++)
                {
                    node = new TrieNode(key[i], current);
                    current.Children.Add(key[i], node);
                    current = node;
                }
                created = true;
                return node;
            }

            return null;
        }

        private TrieNode FindNode(string key, bool create, out bool created)
        {
            return FindNode(_rootNode, key.ToCharArray(), 0, create, out created);
        }

        private TrieNode FindNode(string key, bool create)
        {
            bool created;
            return FindNode(_rootNode, key.ToCharArray(), 0, create, out created);
        }

        private static void Prune(TrieNode removedNode)
        {
            var current = removedNode;
            while (current != null)
            {
                if (current.Children.Any() || current.HasValue)
                    break;
                if (current.Parent != null)
                {
                    current.Parent.Children.Remove(current.PartialKey);
                }
                current = current.Parent;
            }
        }

        class TrieNode
        {
            private readonly char _partialKey;
            private TValue _value;
            private bool _hasValue;
            private readonly TrieNode _parent;
            private readonly SortedDictionary<char, TrieNode> _children;

            public TrieNode(char partialKey, TrieNode parent)
            {
                _partialKey = partialKey;
                _parent = parent;
                _children = new SortedDictionary<char, TrieNode>();
            }

            public TrieNode Parent
            {
                get { return _parent; }
            }

            public char PartialKey
            {
                get { return _partialKey; }
            }

            public bool HasValue
            {
                get { return _hasValue; }
            }

            public TValue Value
            {
                get { return _value; }
                set
                {
                    _value = value;
                    _hasValue = true;
                }
            }

            public void ClearValue()
            {
                _hasValue = false;
            }

            public IDictionary<char, TrieNode> Children
            {
                get { return _children; }
            }
        }

        #endregion

        #region IDictionary<string, TValue> implementation

        public IEnumerator<KeyValuePair<string, TValue>> GetEnumerator()
        {
            return Enumerate(_rootNode, null).GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        void ICollection<KeyValuePair<string, TValue>>.Add(KeyValuePair<string, TValue> item)
        {
            Add(item.Key, item.Value);
        }

        public void Clear()
        {
            _rootNode.ClearValue();
            _rootNode.Children.Clear();
        }

        bool ICollection<KeyValuePair<string, TValue>>.Contains(KeyValuePair<string, TValue> item)
        {
            TValue value;
            if (TryGetValue(item.Key, out value))
                return Equals(value, item.Value);
            return false;
        }

        void ICollection<KeyValuePair<string, TValue>>.CopyTo(KeyValuePair<string, TValue>[] array, int arrayIndex)
        {
            if (_count + arrayIndex > array.Length)
                throw new ArgumentException("The destination array is not large enough");
            int index = arrayIndex;
            foreach (var item in Enumerate(_rootNode, null))
            {
                array[index] = item;
                index++;
            }
        }

        bool ICollection<KeyValuePair<string, TValue>>.Remove(KeyValuePair<string, TValue> item)
        {
            var node = FindNode(item.Key, false);
            if (node == null)
                return false;
            if (node.HasValue && Equals(node.Value, item.Value))
            {
                node.ClearValue();
                _count--;
                return true;
            }
            return false;
        }

        public int Count
        {
            get { return _count; }
        }

        public bool IsReadOnly
        {
            get { return false; }
        }

        public bool ContainsKey(string key)
        {
            var node = FindNode(key, false);
            if (node != null)
                return node.HasValue;
            return false;
        }

        public void Add(string key, TValue value)
        {
            var node = FindNode(key, true);
            if (node.HasValue)
                throw new ArgumentException("An element with the same key already exists in the trie.");
            node.Value = value;
            _count++;
        }

        public bool Remove(string key)
        {
            var node = FindNode(key, false);
            if (node == null)
                return false;

            if (node.HasValue)
            {
                node.ClearValue();
                Prune(node);
                _count--;
            }
            return false;
        }

        public bool TryGetValue(string key, out TValue value)
        {
            value = default(TValue);
            var node = FindNode(key, false);
            if (node == null)
                return false;
            if (node.HasValue)
            {
                value = node.Value;
                return true;
            }
            return false;
        }

        public TValue this[string key]
        {
            get
            {
                TValue value;
                if (TryGetValue(key, out value))
                    return value;
                throw new KeyNotFoundException();
            }
            set
            {
                bool created;
                var node = FindNode(key, true, out created);
                node.Value = value;
                if (created)
                    _count++;
            }
        }

        public ICollection<string> Keys
        {
            get { return Enumerate(_rootNode, null).Select(kvp => kvp.Key).ToList(); }
        }

        public ICollection<TValue> Values
        {
            get { return Enumerate(_rootNode, null).Select(kvp => kvp.Value).ToList(); }
        }

        #endregion

        #region Public trie-specific methods

        public bool ContainsPrefix(string prefix)
        {
            var node = FindNode(prefix, false);
            return (node.HasValue || node.Children.Any());
        }

        public bool RemovePrefix(string prefix)
        {
            var node = FindNode(prefix, false);
            if (node == null)
                return false;

            int count = Enumerate(node, null).Count();
            if (count == 0)
                return false;

            node.Children.Clear();
            node.ClearValue();
            Prune(node);
            _count -= count;

            return true;
        }

        public IEnumerable<KeyValuePair<string, TValue>> FindPrefix(string prefix)
        {
            var node = FindNode(prefix, false);
            if (node == null)
                yield break;

            string prefix2 = null;
            if (prefix.Length > 0)
                prefix2 = prefix.Substring(0, prefix.Length - 1);
            foreach (var kvp in Enumerate(node, prefix2))
            {
                yield return kvp;
            }
        }

        #endregion
    }
}