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 readonly Dictionary<TKey, WeakReference<TValue>> _dictionary = new();

        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() => 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) =>
            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) => Contains(item) && Remove(item.Key);
        public int Count => _dictionary.Count;
        public bool IsReadOnly => false;
        public bool ContainsKey(TKey key) => TryGetValue(key, out _);
        public void Add(TKey key, TValue value) => _dictionary.Add(key, new WeakReference<TValue>(value));
        public bool Remove(TKey key) => _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 => new KeyCollection(this);
        public ICollection<TValue> Values => new ValueCollection(this);

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

            public IEnumerator<TKey> GetEnumerator()
            {
                using var enumerator = _dictionary.GetEnumerator();
                while (enumerator.MoveNext())
                    yield return enumerator.Current.Key;
            }
            
            IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
            public void Add(TKey item) => throw new NotSupportedException("The key collection is read only.");
            public void Clear() => throw new NotSupportedException("The key collection is read only.");
            public bool Contains(TKey item) => _dictionary.ContainsKey(item);
            public void CopyTo(TKey[] array, int arrayIndex) => throw new NotImplementedException();
            public bool Remove(TKey item) => throw new NotSupportedException("The key collection is read only.");
            public int Count => _dictionary.Count;
            public bool IsReadOnly => true;
        }

        public class ValueCollection : ICollection<TValue>, IReadOnlyCollection<TValue>
        {
            private readonly 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() => 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) => _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;
        }
    }
}