LocalCache.cs

public struct LocalCache<TKey, TValue> : IEnumerable<KeyValuePair<TKey, TValue>>, IDisposable
    where TKey : IEquatable<TKey>
{
    private const int RentArrayLength = 16;

    [ThreadStatic] private static Dictionary<TKey, TValue>? _cachedDictionary;

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static Pair[] RentArray() => ArrayPool<Pair>.Shared.Rent(RentArrayLength);

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static Dictionary<TKey, TValue> RentDictionary()
    {
        var result = _cachedDictionary ?? new Dictionary<TKey, TValue>(64);
        _cachedDictionary = null;
        return result;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static void ReturnArray(Pair[] array, int length)
    {
        array.AsSpan(0, length).Clear();
        ArrayPool<Pair>.Shared.Return(array);
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static void ReturnDictionary(Dictionary<TKey, TValue> dictionary)
    {
        dictionary.Clear();
        _cachedDictionary = dictionary;
    }

    public int Count
    {
        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        get => _count;
    }

    private Pair[] _array;
    private int _count;
    private Dictionary<TKey, TValue>? _dictionary;

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public LocalCache()
    {
        _array = RentArray();
        _dictionary = null;
        _count = 0;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public LocalCache(int capacity)
    {
        if (capacity < RentArrayLength)
        {
            _array = RentArray();
            _dictionary = null;
        }
        else
        {
            _array = [];
            _dictionary = RentDictionary();
        }

        _count = 0;
    }

    public void Add(TKey key, TValue value, bool replaceIfExists = false)
    {
        if (_dictionary != null)
        {
            if (replaceIfExists) _dictionary[key] = value;
            else _dictionary.Add(key, value);
        }
        else if (_count < _array.Length)
        {
            // проверка дубликата
            ref var exists = ref GetValueRef(key);
            if (Unsafe.IsNullRef(in exists))
            {
                _array[_count] = new Pair(key, value);
            }
            else
            {
                if (replaceIfExists) exists = value;
                else Throw("An item with the same key has already been added");
            }
        }
        else
        {
            _dictionary = RentDictionary();

            foreach (var pair in _array)
            {
                _dictionary.Add(pair.Key, pair.Value);
            }

            ReturnArray(_array, _count);
            _array = [];
            _dictionary.Add(key, value);
        }

        _count++;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly bool ContainsKey(TKey key)
    {
        ref readonly var value = ref GetValueRef(key);
        return !Unsafe.IsNullRef(in value);
    }

    public void Dispose()
    {
        if (_array is { Length: > 0 })
        {
            ReturnArray(_array, _count);
            _array = [];
        }

        if (_dictionary != null)
        {
            ReturnDictionary(_dictionary);
            _dictionary = null;
        }

        _count = 0;
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public readonly Enumerator GetEnumerator() => new(_array, _dictionary, _count);

    public readonly ref TValue GetValueRef(TKey key)
    {
        if (_count == 0) return ref Unsafe.NullRef<TValue>();
        if (_dictionary != null)
        {
            return ref CollectionsMarshal.GetValueRefOrNullRef(_dictionary, key);
        }

        foreach (ref var pair in _array.AsSpan(0, _count))
        {
            if (pair.Key.Equals(key)) return ref pair.Value;
        }

        return ref Unsafe.NullRef<TValue>();
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public bool TryGetValue(TKey key, [MaybeNullWhen(false)] out TValue value)
    {
        ref readonly var exists = ref GetValueRef(key);

        if (Unsafe.IsNullRef(in exists))
        {
            value = default!;
            return false;
        }

        value = exists;
        return true;
    }

    public readonly ref TValue this[TKey key]
    {
        get
        {
            ref var value = ref GetValueRef(key);
            if (Unsafe.IsNullRef(in value)) Throw("Key not found");
            return ref value;
        }
    }

    public struct Enumerator : IEnumerator<KeyValuePair<TKey, TValue>>
    {
        public readonly KeyValuePair<TKey, TValue> Current
        {
            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            get
            {
                if (!_isArray) return _dictionaryEnumerator.Current;

                var pair = _arrayEnumerator.Current; // извлекаем один раз
                return new KeyValuePair<TKey, TValue>(pair.Key, pair.Value);
            }
        }

        private readonly bool _isArray;
        private ArraySegment<Pair>.Enumerator _arrayEnumerator;
        private Dictionary<TKey, TValue>.Enumerator _dictionaryEnumerator;

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        internal Enumerator(Pair[] array, Dictionary<TKey, TValue>? dictionary, int count)
        {
            if (dictionary == null)
            {
                _isArray = true;
                _arrayEnumerator = new ArraySegment<Pair>(array, 0, count).GetEnumerator();
                _dictionaryEnumerator = default;
            }
            else
            {
                _isArray = false;
                _arrayEnumerator = default;
                _dictionaryEnumerator = dictionary.GetEnumerator();
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public bool MoveNext() => _isArray
            ? _arrayEnumerator.MoveNext()
            : _dictionaryEnumerator.MoveNext();

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public void Dispose()
        {
            if (_isArray) _arrayEnumerator.Dispose();
            else _dictionaryEnumerator.Dispose();
        }

        readonly void IEnumerator.Reset()
        {
        }

        readonly object IEnumerator.Current => Current;
    }

    [method: MethodImpl(MethodImplOptions.AggressiveInlining)]
    internal struct Pair(TKey key, TValue value)
    {
        public readonly TKey Key = key;
        public TValue Value = value;
    }

    [DoesNotReturn]
    private static void Throw(string message) => throw new Exception(message);

    readonly IEnumerator<KeyValuePair<TKey, TValue>> IEnumerable<KeyValuePair<TKey, TValue>>.GetEnumerator() =>
        GetEnumerator();

    readonly IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
}