TechbloxModdingAPI/TechbloxModdingAPI/Utility/WeakDictionary.cs

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

namespace TechbloxModdingAPI.Utility
{
    public class WeakDictionary<TKey, TValue> : IDictionary<TKey, TValue> where TValue : class
    {
        private Dictionary<TKey, WeakReference<TValue>> _dictionary = new Dictionary<TKey, WeakReference<TValue>>();

        public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
        {
            using var enumerator = _dictionary.GetEnumerator();
            while (enumerator.MoveNext())
            {
                if (enumerator.Current.Value.TryGetTarget(out var value))
                    yield return new KeyValuePair<TKey, TValue>(enumerator.Current.Key, value);
            }
        }

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

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

        public void Clear()
        {
            _dictionary.Clear();
        }

        public bool Contains(KeyValuePair<TKey, TValue> item)
        {
            return TryGetValue(item.Key, out var value) && item.Value == value;
        }

        public void CopyTo(KeyValuePair<TKey, TValue>[] array, int arrayIndex)
        {
            throw new System.NotImplementedException();
        }

        public bool Remove(KeyValuePair<TKey, TValue> item)
        {
            return Contains(item) && Remove(item.Key);
        }

        public int Count => _dictionary.Count;
        public bool IsReadOnly => false;

        public bool ContainsKey(TKey key)
        {
            return TryGetValue(key, out _);
        }

        public void Add(TKey key, TValue value)
        {
            _dictionary.Add(key, new WeakReference<TValue>(value));
        }

        public bool Remove(TKey key)
        {
            return _dictionary.Remove(key);
        }

        public bool TryGetValue(TKey key, out TValue value)
        {
            value = null;
            bool ret = _dictionary.TryGetValue(key, out var reference) && reference.TryGetTarget(out value);
            if (!ret) _dictionary.Remove(key);
            return ret;
        }

        public TValue this[TKey key]
        {
            get => TryGetValue(key, out var value)
                ? value
                : throw new KeyNotFoundException($"Key {key} not found in WeakDictionary.");
            set => _dictionary[key] = new WeakReference<TValue>(value);
        }

        public ICollection<TKey> Keys => _dictionary.Keys;
        public ICollection<TValue> Values => new ValueCollection(this);

        public class ValueCollection : ICollection<TValue>, IReadOnlyCollection<TValue>
        {
            private WeakDictionary<TKey, TValue> _dictionary;
            internal ValueCollection(WeakDictionary<TKey, TValue> dictionary)
            {
                _dictionary = dictionary;
            }

            public IEnumerator<TValue> GetEnumerator()
            {
                using var enumerator = _dictionary.GetEnumerator();
                while (enumerator.MoveNext())
                {
                    yield return enumerator.Current.Value;
                }
            }

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

            public void Add(TValue item)
            {
                throw new NotSupportedException("The value collection is read only.");
            }

            public void Clear()
            {
                throw new NotSupportedException("The value collection is read only.");
            }

            public bool Contains(TValue item)
            {
                return _dictionary.Any(kv => kv.Value == item);
            }

            public void CopyTo(TValue[] array, int arrayIndex)
            {
                throw new NotImplementedException();
            }

            public bool Remove(TValue item)
            {
                throw new NotSupportedException("The value collection is read only.");
            }

            public int Count => _dictionary.Count;
            public bool IsReadOnly => true;
        }
    }
}