Simplify block allocation logic

Recyclable.Collections/RecyclableDictionary.cs

@@ -0,0 +1,346 @@
+using System.Collections;
+using System.Collections.Generic;
+using System.Numerics;
+using System.Runtime.CompilerServices;
+
+namespace Recyclable.Collections
+{
+    internal sealed class RecyclableDictionary<TKey, TValue> : IEnumerable<KeyValuePair<TKey, TValue>>, IDisposable
+        where TKey : notnull
+    {
+        private static readonly bool _needsClearing = !typeof(TKey).IsValueType || !typeof(TValue).IsValueType;
+
+        private readonly int _blockSize;
+        private readonly int _blockSizeMinus1;
+        private readonly byte _blockShift;
+
+        private Entry[][] _entries;
+        private int[] _buckets;
+        private int _count;
+        private bool _disposed;
+
+        private struct Entry
+        {
+            public int HashCode;
+            public int Next;
+            public TKey Key;
+            public TValue Value;
+        }
+
+        public RecyclableDictionary(int blockSize = RecyclableDefaults.BlockSize)
+        {
+            if (!BitOperations.IsPow2((uint)blockSize))
+            {
+                blockSize = (int)BitOperations.RoundUpToPowerOf2((uint)blockSize);
+            }
+
+            _blockSize = blockSize;
+            _blockSizeMinus1 = blockSize - 1;
+            _blockShift = (byte)(31 - BitOperations.LeadingZeroCount((uint)blockSize));
+
+            _entries = new Entry[4][];
+            _entries[0] = new Entry[blockSize];
+            _buckets = new int[4];
+            Array.Fill(_buckets, -1);
+        }
+
+        public int Count => _count;
+
+        public TValue this[TKey key]
+        {
+            get
+            {
+                var index = FindIndex(key, out _);
+                if (index >= 0)
+                {
+                    return GetEntry(index).Value;
+                }
+
+                ThrowHelper.ThrowArgumentOutOfRangeException(nameof(key), "Key not found");
+                return default!;
+            }
+            set
+            {
+                var index = FindIndex(key, out _);
+                if (index >= 0)
+                {
+                    ref var entry = ref GetEntry(index);
+                    entry.Value = value;
+                    return;
+                }
+
+                Add(key, value);
+            }
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public void Add(TKey key, TValue value)
+        {
+            var hash = key.GetHashCode() & int.MaxValue;
+            if (_count >= (_buckets.Length * 3) / 4)
+            {
+                ResizeBuckets(_buckets.Length * 2);
+            }
+
+            EnsureCapacity(_count + 1);
+
+            int bucket = hash & (_buckets.Length - 1);
+            int index = _buckets[bucket];
+            while (index >= 0)
+            {
+                ref var check = ref GetEntry(index);
+                if (check.HashCode == hash && EqualityComparer<TKey>.Default.Equals(check.Key, key))
+                {
+                    throw new ArgumentException("An element with the same key already exists.", nameof(key));
+                }
+
+                index = check.Next;
+            }
+
+            int newIndex = _count++;
+            ref var entry = ref GetEntry(newIndex);
+            entry.HashCode = hash;
+            entry.Key = key;
+            entry.Value = value;
+            entry.Next = _buckets[bucket];
+            _buckets[bucket] = newIndex;
+
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool ContainsKey(TKey key) => FindIndex(key, out _) >= 0;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool TryGetValue(TKey key, out TValue value)
+        {
+            var index = FindIndex(key, out _);
+            if (index >= 0)
+            {
+                value = GetEntry(index).Value;
+                return true;
+            }
+
+            value = default!;
+            return false;
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public bool Remove(TKey key)
+        {
+            if (_count == 0)
+            {
+                return false;
+            }
+
+            var hash = key.GetHashCode() & int.MaxValue;
+            int bucket = hash & (_buckets.Length - 1);
+            int prev = -1;
+            int index = _buckets[bucket];
+            while (index >= 0)
+            {
+                ref var entry = ref GetEntry(index);
+                if (entry.HashCode == hash && EqualityComparer<TKey>.Default.Equals(entry.Key, key))
+                {
+                    if (prev < 0)
+                    {
+                        _buckets[bucket] = entry.Next;
+                    }
+                    else
+                    {
+                        GetEntry(prev).Next = entry.Next;
+                    }
+
+                    int lastIndex = _count - 1;
+                    if (index != lastIndex)
+                    {
+                        MoveEntry(lastIndex, index);
+                    }
+
+                    ClearEntry(lastIndex);
+                    _count--;
+
+                    return true;
+                }
+
+                prev = index;
+                index = entry.Next;
+            }
+
+            return false;
+        }
+
+        public void Clear()
+        {
+            if (_needsClearing)
+            {
+                for (var i = 0; i < _count; i++)
+                {
+                    ClearEntry(i);
+                }
+            }
+
+            _count = 0;
+            Array.Fill(_buckets, -1);
+        }
+
+        public TKey GetKey(int index) => GetEntry(index).Key;
+
+        public TValue GetValue(int index) => GetEntry(index).Value;
+
+        public Enumerator GetEnumerator() => new(this);
+
+        IEnumerator<KeyValuePair<TKey, TValue>> IEnumerable<KeyValuePair<TKey, TValue>>.GetEnumerator() => new Enumerator(this);
+
+        IEnumerator IEnumerable.GetEnumerator() => new Enumerator(this);
+
+        public void Dispose()
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            Clear();
+            _entries = Array.Empty<Entry[]>();
+            _buckets = Array.Empty<int>();
+            _disposed = true;
+            GC.SuppressFinalize(this);
+        }
+
+        private ref Entry GetEntry(int index)
+        {
+            return ref _entries[index >> _blockShift][index & _blockSizeMinus1];
+        }
+
+        private void ClearEntry(int index)
+        {
+            ref var entry = ref GetEntry(index);
+            entry.HashCode = 0;
+            entry.Next = -1;
+            if (_needsClearing)
+            {
+                entry.Key = default!;
+                entry.Value = default!;
+            }
+        }
+
+        private void MoveEntry(int sourceIndex, int targetIndex)
+        {
+            ref var source = ref GetEntry(sourceIndex);
+            ref var target = ref GetEntry(targetIndex);
+            target = source;
+
+            int bucket = source.HashCode & (_buckets.Length - 1);
+            int cur = _buckets[bucket];
+            int prev = -1;
+            while (cur != sourceIndex)
+            {
+                prev = cur;
+                cur = GetEntry(cur).Next;
+            }
+
+            if (prev < 0)
+            {
+                _buckets[bucket] = targetIndex;
+            }
+            else
+            {
+                GetEntry(prev).Next = targetIndex;
+            }
+        }
+
+        private void EnsureCapacity(int min)
+        {
+            int requiredBlocks = (min + _blockSizeMinus1) >> _blockShift;
+            if (_entries.Length < requiredBlocks)
+            {
+                Array.Resize(ref _entries, requiredBlocks);
+            }
+
+            for (int i = 0; i < requiredBlocks; i++)
+            {
+                _entries[i] ??= new Entry[_blockSize];
+            }
+        }
+
+        private void ResizeBuckets(int newSize)
+        {
+            var newBuckets = new int[newSize];
+            Array.Fill(newBuckets, -1);
+            for (int i = 0; i < _count; i++)
+            {
+                ref var entry = ref GetEntry(i);
+                int bucket = entry.HashCode & (newSize - 1);
+                entry.Next = newBuckets[bucket];
+                newBuckets[bucket] = i;
+            }
+
+            _buckets = newBuckets;
+        }
+
+        private int FindIndex(TKey key, out int hash)
+        {
+            if (_count == 0)
+            {
+                hash = 0;
+                return -1;
+            }
+
+            hash = key.GetHashCode() & int.MaxValue;
+            int bucket = hash & (_buckets.Length - 1);
+            int index = _buckets[bucket];
+            while (index >= 0)
+            {
+                ref var entry = ref GetEntry(index);
+                if (entry.HashCode == hash && EqualityComparer<TKey>.Default.Equals(entry.Key, key))
+                {
+                    return index;
+                }
+
+                index = entry.Next;
+            }
+
+            return -1;
+        }
+
+        public struct Enumerator : IEnumerator<KeyValuePair<TKey, TValue>>
+        {
+            private readonly RecyclableDictionary<TKey, TValue> _dictionary;
+            private int _index;
+            private KeyValuePair<TKey, TValue> _current;
+
+            internal Enumerator(RecyclableDictionary<TKey, TValue> dictionary)
+            {
+                _dictionary = dictionary;
+                _index = 0;
+                _current = default;
+            }
+
+            public KeyValuePair<TKey, TValue> Current => _current;
+
+            object IEnumerator.Current => _current;
+
+            public bool MoveNext()
+            {
+                if (_index >= _dictionary._count)
+                {
+                    return false;
+                }
+
+                var entry = _dictionary.GetEntry(_index++);
+                _current = new KeyValuePair<TKey, TValue>(entry.Key, entry.Value);
+                return true;
+            }
+
+            public void Reset()
+            {
+                _index = 0;
+                _current = default;
+            }
+
+            public void Dispose()
+            {
+            }
+        }
+    }
+}

Recyclable.CollectionsTests/RecyclableDictionaryTests.cs

@@ -0,0 +1,85 @@
+using FluentAssertions;
+using Recyclable.Collections;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Recyclable.CollectionsTests
+{
+    public class RecyclableDictionaryTests
+    {
+        private static readonly (int, string)[] _testData = new[] { (1, "a"), (2, "b"), (3, "c"), (4, "d"), (5, "e") };
+
+        [Fact]
+        public void AddShouldStoreItems()
+        {
+            using var dict = new RecyclableDictionary<int, string>();
+            foreach (var (key, value) in _testData)
+            {
+                dict.Add(key, value);
+            }
+
+            _ = dict.Should().HaveCount(_testData.Length);
+            for (var i = 0; i < _testData.Length; i++)
+            {
+                dict.GetKey(i).Should().Be(_testData[i].Item1);
+                dict.GetValue(i).Should().Be(_testData[i].Item2);
+            }
+        }
+
+        [Fact]
+        public void IndexerShouldUpdateValue()
+        {
+            using var dict = new RecyclableDictionary<int, string>();
+            dict.Add(1, "a");
+
+            dict[1] = "b";
+
+            _ = dict[1].Should().Be("b");
+        }
+
+        [Fact]
+        public void RemoveShouldSwapWithLast()
+        {
+            using var dict = new RecyclableDictionary<int, string>();
+            foreach (var (key, value) in _testData)
+            {
+                dict.Add(key, value);
+            }
+
+            dict.Remove(3).Should().BeTrue();
+            _ = dict.ContainsKey(3).Should().BeFalse();
+            _ = dict.Should().HaveCount(_testData.Length - 1);
+        }
+
+        [Fact]
+        public void TryGetValueShouldReturnValue()
+        {
+            using var dict = new RecyclableDictionary<int, string>();
+            dict.Add(1, "a");
+
+            var found = dict.TryGetValue(1, out var value);
+
+            _ = found.Should().BeTrue();
+            _ = value.Should().Be("a");
+        }
+
+        [Fact]
+        public void EnumeratorShouldYieldItems()
+        {
+            using var dict = new RecyclableDictionary<int, string>();
+            foreach (var (key, value) in _testData)
+            {
+                dict.Add(key, value);
+            }
+
+            var actual = new List<KeyValuePair<int, string>>();
+            var enumerator = dict.GetEnumerator();
+            while (enumerator.MoveNext())
+            {
+                actual.Add(enumerator.Current);
+            }
+
+            _ = actual.Should().ContainInConsecutiveOrder(_testData.Select(x => new KeyValuePair<int, string>(x.Item1, x.Item2)));
+        }
+    }
+}