riversong-code-showcase/Source/Engine/Collections/MultiDictionary.cs

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

namespace DanieleMarotta.RiversongCodeShowcase
{
    public abstract class MultiDictionary<K, V, C> : IMultiDictionary<K, V, C> where C : ICollection<V>
    {
        private readonly Dictionary<K, C> _collections;

        protected MultiDictionary(IEqualityComparer<K> keyComparer = null)
        {
            _collections = keyComparer != null ? new Dictionary<K, C>(keyComparer) : new Dictionary<K, C>();
        }

        public C this[K key] => _collections[key];

        public Dictionary<K, C>.KeyCollection Keys => _collections.Keys;

        public int KeyCount => _collections.Count;

        public int ValueCount { get; private set; }

        protected abstract C CreateCollection();

        protected abstract bool AddToCollection(V value, C collection);

        public int Count(K key)
        {
            return _collections.TryGetValue(key, out var collection) ? collection.Count : 0;
        }

        public bool Add(K key, V value)
        {
            if (!_collections.TryGetValue(key, out var collection))
            {
                collection = CreateCollection();
                _collections.Add(key, collection);
            }

            if (AddToCollection(value, collection))
            {
                ValueCount++;
                return true;
            }

            return false;
        }

        public bool Remove(K key, V value)
        {
            if (_collections.TryGetValue(key, out var collection) && collection.Remove(value))
            {
                ValueCount--;
                if (collection.Count == 0) _collections.Remove(key);
                return true;
            }
            return false;
        }

        public bool Remove(V value)
        {
            var valueRemoved = false;
            var removeCollection = false;
            K collectionKeyToRemove = default;
            foreach (var kvp in _collections)
            {
                var collection = kvp.Value;
                if (collection.Remove(value))
                {
                    valueRemoved = true;
                    ValueCount--;
                    if (collection.Count == 0)
                    {
                        removeCollection = true;
                        collectionKeyToRemove = kvp.Key;
                    }
                    break;
                }
            }
            if (removeCollection) _collections.Remove(collectionKeyToRemove);
            return valueRemoved;
        }

        public void Clear(K key)
        {
            if (_collections.TryGetValue(key, out var collection))
            {
                ValueCount -= collection.Count;
                collection.Clear();
            }
        }

        public void Clear()
        {
            _collections.Clear();
            ValueCount = 0;
        }

        public bool ContainsKey(K key)
        {
            return _collections.ContainsKey(key);
        }

        public bool ContainsValue(K key, V value)
        {
            return _collections.TryGetValue(key, out var collection) && collection.Contains(value);
        }

        public bool ContainsValue(V value)
        {
            foreach (var collection in _collections.Values)
                if (collection.Contains(value))
                    return true;
            return false;
        }

        public bool TryGetValues(K key, out C values)
        {
            return _collections.TryGetValue(key, out values);
        }

        public IEnumerator<KeyValuePair<K, V>> GetEnumerator()
        {
            foreach (var kvp in _collections)
            {
                var key = kvp.Key;
                var collection = kvp.Value;
                foreach (var value in collection) yield return new KeyValuePair<K, V>(key, value);
            }
        }

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