メモリ連続の辞書構造を作ってみる

概要

C# には標準のキーと値のコレクション System.Collections.Generic.Dictionary<TKey,TValue> があります。

var dic = new Dictionary<int, string>();
Dictionary<int, string>.KeyCollection keys = dic.Keys;
Dictionary<int, string>.ValueCollection values = dic.Values;

ここで keys と values を Span<T> で取りたい、というのが今回の話です。

ちなみに

キーまたはインデックスからアクセスできるキーと値のペアのコレクションが欲しい場合は .net9 preview6 で System.Collections.Generic.OrderedDictionary<TKey, TValue> が追加されたので、そちらを利用できます。

OrderedDictionary<string, int> d = new()
{
    ["a"] = 1,
    ["b"] = 2,
    ["c"] = 3,
};

d.Add("d", 4);
d.RemoveAt(0);
d.RemoveAt(2);

SpanDictionary

名前は適当です。Remove() で要素を削除しても、要素の並びが歯抜けにならず連続になります。
都合上 Remove() は O(n) で２桁くらい遅いので、基本的にはおまけです。

// 使い方
var dic = new SpanDictionary<int, string>();
dic.Add(1, "one");
dic.Add(5, "five");
dic.Add(3, "three");

ReadOnlySpan<int> keys = dic.Keys;
Assert.Equal([1, 5, 3], keys);

Span<string> values = dic.Values;
Assert.Equal(["one", "five", "three"], values);

dic.Clear();

ref var refOne = ref dic[1];
refOne = "ONE";
Assert.Equal("ONE", dic[1]);

ref var refTwo = ref dic.AddOrRef(2, out var exist);
if (!exist)
    refTwo = "TWO";
Assert.Equal("TWO", dic[2]);

Assert.False(dic.TryGetValue(3, out var three));
Assert.Equal(default, three);

dic.Clear();
dic.Add(1, "one");
dic.Add(2, "two");
dic.Add(3, "three");
dic.Remove(2);

Assert.Equal([1, 3], dic.Keys);
Assert.Equal(["one", "three"], dic.Values);

コード

SpanDictionary

using System.Diagnostics;

/// <summary>
/// メモリ連続の辞書
/// </summary>
/// <typeparam name="K"></typeparam>
/// <typeparam name="V"></typeparam>
[DebuggerDisplay("Count = {Count}")]
[DebuggerTypeProxy(typeof(DebugView<,>))]
file class SpanDictionary<K, V>
{
    private record struct Entry
    {
        internal int Bucket;
        internal int HashCode;
        internal int Next;
    }

    private static int GetHashCodeKey(in K key) => key!.GetHashCode() & int.MaxValue;

    private Entry[] _entries = Array.Empty<Entry>();
    private K[] _keys = Array.Empty<K>();
    private V[] _values = Array.Empty<V>();

    /// <summary>
    /// キーのスパン
    /// </summary>
    public ReadOnlySpan<K> Keys => this._keys.AsSpan(0, this._count);

    /// <summary>
    /// 値のスパン
    /// </summary>
    public Span<V> Values => this._values.AsSpan(0, this._count);

    private int _count;
    /// <summary>
    /// 要素数
    /// </summary>
    public int Count => this._count;

    /// <summary>
    /// 容量
    /// </summary>
    public int Capacity => this._entries.Length;

    /// <summary>
    /// コンストラクタ
    /// </summary>
    /// <param name="capacity"></param>
    /// <exception cref="System.ArgumentOutOfRangeException"></exception>
    public SpanDictionary(int capacity)
    {
        ArgumentOutOfRangeException.ThrowIfNegative(capacity);
        this.Initialize(capacity);
    }

    /// <summary>
    /// コンストラクタ
    /// </summary>
    public SpanDictionary() : this(0) { }

    /// <summary>
    /// 取得 キーが存在しない場合 default 値を追加 O(1)
    /// </summary>
    /// <param name="key"></param>
    /// <returns></returns>
    public ref V this[in K key] => ref this.Insert(key, false);

    private void Initialize(int capacity)
    {
        int prime = HashUtl.GetPrime(capacity);
        this._entries = new Entry[prime];
        this._keys = new K[prime];
        this._values = new V[prime];
    }

    /// <summary>
    /// 追加 O(1) 既にキーが存在する場合例外
    /// </summary>
    /// <param name="key"></param>
    /// <param name="value"></param>
    /// <exception cref="System.InvalidOperationException"></exception>
    public void Add(in K key, in V value) => this.Insert(key, true) = value;

    /// <summary>
    /// 追加 O(1) 既にキーが存在する場合例外
    /// </summary>
    /// <param name="key"></param>
    /// <returns></returns>
    /// <exception cref="System.InvalidOperationException"></exception>
    public ref V Add(in K key) => ref this.Insert(key, true);

    private void Resize() => this.Resize(HashUtl.ExpandPrime(this._count), false);

    private void Resize(int newSize, bool forceNewHashCodes)
    {
        var entries2 = new Entry[newSize];
        this._entries.AsSpan().CopyTo(entries2);
        var span = entries2.AsSpan(0, this._entries.Length);
        for (int n = 0; n < span.Length; ++n)
            span[n].Bucket = default;

        if (forceNewHashCodes)
        {
            for (int n = 0; n < this._count; n++)
            {
                ref var entry = ref entries2[n];
                if (entry.Next > 0)
                {
                    entry.HashCode = GetHashCodeKey(this._keys[n]);
                }
            }
        }
        for (int n = 0; n < this._count; n++)
        {
            ref var entry = ref entries2[n];
            int num = entry.HashCode % newSize;
            ref var tmp = ref entries2[num];
            entry.Next = tmp.Bucket;
            tmp.Bucket = n + 1;
        }
        this._entries = entries2;

        var keys2 = new K[newSize];
        var keySpan = this._keys.AsSpan(0, this._count);
        keySpan.CopyTo(keys2);
        keySpan.Clear();
        this._keys = keys2;

        var values2 = new V[newSize];
        var valueSpan = this._values.AsSpan(0, this._count);
        valueSpan.CopyTo(values2);
        valueSpan.Clear();
        this._values = values2;
    }

    private unsafe ref V Insert(in K key, bool add, bool* exist)
    {
        int num = GetHashCodeKey(key);
        int num2 = num % this._entries.Length;
        ref var entry2 = ref this._entries[num2];
        var n = entry2.Bucket;
        while (n > 0)
        {
            ref var tmp = ref this._entries[n - 1];
            if (tmp.HashCode == num && EqualityComparer<K>.Default.Equals(this._keys[n - 1], key))
            {
                if (add)
                    throw new InvalidOperationException($"Exist key {key}");
                *exist = true;
                return ref this._values[n - 1];
            }
            n = tmp.Next;
        }
        if (this._count == this._entries.Length)
        {
            this.Resize();
            num2 = num % this._entries.Length;
        }
        ref var entry22 = ref this._entries[num2];
        int num4 = this._count;
        this._count++;
        ref var entry4 = ref this._entries[num4];
        this._keys[num4] = key;
        entry4.HashCode = num;
        entry4.Next = entry22.Bucket;

        entry22.Bucket = num4 + 1;
        *exist = false;
        return ref this._values[num4];
    }

    private unsafe ref V Insert(in K key, bool add)
    {
        bool _ = false;
        return ref this.Insert(key, add, &_);
    }

    /// <summary>
    /// 値を取得 O(1)
    /// キーは新たに作成されない、作成する場合 [] を使用
    /// </summary>
    /// <param name="key"></param>
    /// <param name="value"></param>
    /// <returns></returns>
    public bool TryGetValue(in K key, out V value)
    {
        int num = GetHashCodeKey(key);
        int num2 = num % this._entries.Length;
        ref var entry2 = ref this._entries[num2];
        var n = entry2.Bucket;
        while (n > 0)
        {
            ref var tmp = ref this._entries[n - 1];
            if (tmp.HashCode == num && EqualityComparer<K>.Default.Equals(this._keys[n - 1], key))
            {
                value = this._values[n - 1];
                return true;
            }
            n = tmp.Next;
        }
        value = default!;
        return false;
    }

    /// <summary>
    /// キーが存在する場合 true O(1)
    /// </summary>
    /// <param name="key"></param>
    /// <returns></returns>
    public bool ContainsKey(in K key) => this.TryGetValue(key, out _);

    /// <summary>
    /// 値を取得 O(1)
    /// キーを新たに作成、既に存在している場合 exist は true
    /// </summary>
    /// <param name="key"></param>
    /// <param name="exist"></param>
    /// <returns></returns>
    public unsafe ref V AddOrRef(in K key, out bool exist)
    {
        bool tmp = false;
        ref var result = ref this.Insert(key, false, &tmp);
        exist = tmp;
        return ref result!;
    }

    /// <summary>
    /// クリア
    /// </summary>
    public void Clear()
    {
        if (this._count > 0)
        {
            this._entries.AsSpan().Clear();
            this._keys.AsSpan(0, this._count).Clear();
            this._values.AsSpan(0, this._count).Clear();
            this._count = 0;
        }
    }

    /// <summary>
    /// 要素を削除 O(n)
    /// </summary>
    /// <param name="key"></param>
    /// <returns></returns>
    public bool Remove(in K key)
    {
        int num = GetHashCodeKey(key);
        int num2 = num % this._entries.Length;
        int num3 = 0;
        int n = this._entries[num2].Bucket;
        while (n > 0)
        {
            var n2 = n - 1;
            ref var entry = ref this._entries[n2];
            if (entry.HashCode == num && EqualityComparer<K>.Default.Equals(this._keys[n2], key))
            {
                if (num3 is 0)
                    this._entries[num2].Bucket = entry.Next;
                else
                    this._entries[num3].Next = entry.Next;

                var rightIndex = this._count - 1;
                // コンパクション
                for (int m = n2; m < rightIndex; ++m)
                {
                    ref var left = ref this._entries[m];
                    ref var right = ref this._entries[m + 1];
                    var bucket = left.Bucket;
                    left = right;
                    left.Bucket = bucket;
                    this._keys[m] = this._keys[m + 1];
                    this._values[m] = this._values[m + 1];
                }
                for (int m = 0; m < this._entries.Length; ++m)
                {
                    ref var left = ref this._entries[m];
                    if (left.Bucket > n2)
                        --left.Bucket;
                    if (left.Next > n2)
                        --left.Next;
                }
                ref var last = ref this._entries[rightIndex];
                var bucket2 = last.Bucket;
                last = default;
                last.Bucket = bucket2;
                this._keys[rightIndex] = default!;
                this._values[rightIndex] = default!;

                --this._count;
                return true;
            }
            num3 = n2;
            n = entry.Next;
        }
        return false;
    }
}

file class DebugView<K, V>
{
    private readonly (K, V)[] _array;

    public DebugView(SpanDictionary<K, V> value)
    {
        this._array = new (K, V)[value.Keys.Length];

        for (int n = 0; n < this._array.Length; ++n)
            this._array[n] = (value.Keys[n], value.Values[n]);
    }

    [DebuggerBrowsable(DebuggerBrowsableState.RootHidden)]
    public (K, V)[] Items => this._array;
}

file static class HashUtl
{
    private static bool IsPrime(int candidate)
    {
        if ((candidate & 1) != 0)
        {
            int num = (int)MathF.Sqrt(candidate);
            for (int i = 3; i <= num; i += 2)
                if (candidate % i == 0)
                    return false;
            return true;
        }
        return candidate == 2;
    }

    internal static int GetPrime(int min)
    {
        for (int j = min | 1; j < int.MaxValue; j += 2)
            if (IsPrime(j) && (j - 1) % 101 != 0)
                return j;
        return min;
    }

    internal static int ExpandPrime(int oldSize)
    {
        int num = 2 * oldSize;
        if (num > 2146435069 && 2146435069 > oldSize)
            return 2146435069;
        return GetPrime(num);
    }
}

テストコード

using Xunit;

public class _SpanDictionaryTest
{
    void 概要()
    {
        var dic = new Dictionary<int, string>();
        Dictionary<int, string>.KeyCollection keys = dic.Keys;
        Dictionary<int, string>.ValueCollection values = dic.Values;

        OrderedDictionary<string, int> d = new()
        {
            ["a"] = 1,
            ["b"] = 2,
            ["c"] = 3,
        };

        d.Add("d", 4);
        d.RemoveAt(0);
        d.RemoveAt(2);
    }

    [Fact]
    void 使い方()
    {
        // 使い方
        var dic = new SpanDictionary<int, string>();
        dic.Add(1, "one");
        dic.Add(5, "five");
        dic.Add(3, "three");

        ReadOnlySpan<int> keys = dic.Keys;
        Assert.Equal([1, 5, 3], keys);

        Span<string> values = dic.Values;
        Assert.Equal(["one", "five", "three"], values);

        dic.Clear();

        ref var refOne = ref dic[1];
        refOne = "ONE";
        Assert.Equal("ONE", dic[1]);

        ref var refTwo = ref dic.AddOrRef(2, out var exist);
        if (!exist)
            refTwo = "TWO";
        Assert.Equal("TWO", dic[2]);

        Assert.False(dic.TryGetValue(3, out var three));
        Assert.Equal(default, three);

        dic.Remove(1);

        Assert.Equal([2], dic.Keys);
        Assert.Equal(["TWO"], dic.Values);
    }

    [Fact]
    void Keys_Values_Count()
    {
        var dic = new SpanDictionary<int, string>();
        Assert.Equal(0, dic.Count);
        Assert.Equal([], dic.Keys);
        Assert.Equal([], dic.Values);

        dic.Add(1, "one");
        Assert.Equal(1, dic.Count);
        Assert.Equal([1], dic.Keys);
        Assert.Equal(["one"], dic.Values);

        dic.Add(2, "two");
        Assert.Equal(2, dic.Count);
        Assert.Equal([1, 2], dic.Keys);
        Assert.Equal(["one", "two"], dic.Values);
    }

    [Fact]
    void Constructor_Capacity()
    {
        Assert.Throws<ArgumentOutOfRangeException>(() => new SpanDictionary<int, string>(-1));

        {
            var dic = new SpanDictionary<int, string>(10);
            Assert.Equal([], dic.Keys);
            Assert.Equal([], dic.Values);
            Assert.True(dic.Capacity >= 10);
        }
        {
            var dic = new SpanDictionary<int, string>();
            Assert.Equal([], dic.Keys);
            Assert.Equal([], dic.Values);
        }
    }

    [Fact]
    void Indexer()
    {
        var dic = new SpanDictionary<int, string>();
        dic[1] = "one";
        Assert.Equal("one", dic[1]);

        dic[1] = "ONE";
        Assert.Equal("ONE", dic[1]);

        dic[2] = "two";
        Assert.Equal("two", dic[2]);

        dic[3] = "three";
        Assert.Equal("three", dic[3]);
    }

    [Fact]
    void Add_K_V()
    {
        Assert.Throws<InvalidOperationException>(() =>
        {
            var dic = new SpanDictionary<int, string>();
            dic.Add(1, "one");
            dic.Add(1, "one");
        });

        var dic = new SpanDictionary<int, string>();
        dic.Add(1, "one");
        Assert.Equal("one", dic[1]);

        dic.Add(2, "two");
        Assert.Equal("two", dic[2]);

        dic.Add(3, "three");
        Assert.Equal("three", dic[3]);
    }

    [Fact]
    void Add_K()
    {
        Assert.Throws<InvalidOperationException>(() =>
        {
            var dic = new SpanDictionary<int, string>();
            dic.Add(1) = "one";
            dic.Add(1) = "one";
        });

        var dic = new SpanDictionary<int, string>();
        dic.Add(1) = "one";
        Assert.Equal("one", dic[1]);

        dic.Add(2) = "two";
        Assert.Equal("two", dic[2]);

        dic.Add(3) = "three";
        Assert.Equal("three", dic[3]);
    }

    [Fact]
    void 容量が大きくなること()
    {
        var dic = new SpanDictionary<int, string>();
        var capacity = dic.Capacity;
        for (int n = 0; n < capacity + 1; ++n)
            dic.Add(n, n.ToString());

        Assert.True(dic.Capacity > capacity);
    }

    [Fact]
    void TryGetValue()
    {
        var dic = new SpanDictionary<int, string>();
        dic.Add(1, "one");

        Assert.True(dic.TryGetValue(1, out var one));
        Assert.Equal("one", one);

        Assert.False(dic.TryGetValue(2, out var two));
        Assert.Equal(default, two);
    }

    [Fact]
    void ContainsKey()
    {
        var dic = new SpanDictionary<int, string>();
        dic.Add(1, "one");

        Assert.True(dic.ContainsKey(1));
        Assert.False(dic.ContainsKey(2));
    }

    [Fact]
    void AddOrRef()
    {
        var dic = new SpanDictionary<int, string>();
        ref var one = ref dic.AddOrRef(1, out var exist);
        Assert.False(exist);
        one = "one";
        Assert.Equal("one", dic[1]);

        ref var one2 = ref dic.AddOrRef(1, out exist);
        Assert.True(exist);
        one2 = "ONE";
        Assert.Equal("ONE", dic[1]);
    }

    [Fact]
    void Clear()
    {
        var dic = new SpanDictionary<int, string>();
        dic.Add(1, "one");
        dic.Add(2, "two");
        dic.Add(3, "three");

        dic.Clear();
        Assert.Equal(0, dic.Count);
        Assert.Equal([], dic.Keys);
        Assert.Equal([], dic.Values);
    }

    [Fact]
    void Remove()
    {
        var dic = new SpanDictionary<int, string>();
        dic.Add(1, "one");
        dic.Add(2, "two");
        dic.Add(3, "three");

        Assert.True(dic.Remove(2));
        Assert.Equal([1, 3], dic.Keys);
        Assert.Equal(["one", "three"], dic.Values);

        Assert.False(dic.Remove(2));
        Assert.False(dic.Remove(4));
    }

    static (int, string)[] _testData;

    static _SpanDictionaryTest()
    {
        // 固定乱数
        var rand = new Random(0);
        _testData = new (int, string)[10000];
        foreach (ref var n in _testData.AsSpan())
        {
            var num = rand.Next();
            n = (num, num.ToString());
        }
    }

    static Action AddSpanDictionary()
    {
        return () =>
        {
            var dic = new SpanDictionary<int, string>();
            for (int n = 0; n < 10000; ++n)
                dic.Add(n, _testData[n].Item2);
        };
    }

    static Action AddDictionary()
    {
        return () =>
        {
            var dic = new Dictionary<int, string>();
            for (int n = 0; n < 10000; ++n)
                dic.Add(n, _testData[n].Item2);
        };
    }

    static Action AddRandomSpanDictionary()
    {
        return () =>
        {
            var dic = new SpanDictionary<int, string>();
            foreach (var (key, value) in _testData)
                dic.Add(key, value);
        };
    }

    static Action AddRandomDictionary()
    {
        return () =>
        {
            var dic = new Dictionary<int, string>();
            foreach (var (key, value) in _testData)
                dic.Add(key, value);
        };
    }

    static Action RemoveSpanDictionary()
    {
        return () =>
        {
            var dic = new SpanDictionary<int, string>();
            for (int n = 0; n < 10000; ++n)
                dic.Add(n, _testData[n].Item2);
            for (int n = 0; n < 10000; ++n)
                dic.Remove(n);
        };
    }

    static Action RemoveDictionary()
    {
        return () =>
        {
            var dic = new Dictionary<int, string>();
            for (int n = 0; n < 10000; ++n)
                dic.Add(n, _testData[n].Item2);
            for (int n = 0; n < 10000; ++n)
                dic.Remove(n);
        };
    }

    static Action RemoveRandomSpanDictionary()
    {
        return () =>
        {
            var dic = new SpanDictionary<int, string>();
            for (int n = 0; n < 10000; ++n)
                dic.Add(n, _testData[n].Item2);

            var rand = new Random(0);
            for (int n = 0; n < 10000; ++n)
                dic.Remove(rand.Next(0, dic.Count));
        };
    }

    static Action RemoveRandomDictionary()
    {
        return () =>
        {
            var dic = new Dictionary<int, string>();
            for (int n = 0; n < 10000; ++n)
                dic.Add(n, _testData[n].Item2);

            var rand = new Random(0);
            for (int n = 0; n < 10000; ++n)
                dic.Remove(rand.Next(0, dic.Count));
        };
    }

    static Action ForeachSpanDictionary()
    {
        var dic = new SpanDictionary<int, string>();
        for (int n = 0; n < 10000; ++n)
            dic.Add(n, _testData[n].Item2);
        return () =>
        {
            foreach (var n in dic.Values)
                n.GetHashCode();
        };
    }

    static Action ForeachDictionary()
    {
        var dic = new Dictionary<int, string>();
        for (int n = 0; n < 10000; ++n)
            dic.Add(n, _testData[n].Item2);
        return () =>
        {
            foreach (var n in dic.Values)
                n.GetHashCode();
        };
    }
}

パフォーマンス

Test	Score	%	CG0
AddSpanDictionary	196	100.0%	66
AddDictionary	378	192.9%	148
AddRandomSpanDictionary	148	100.0%	48
AddRandomDictionary	292	197.3%	161
RemoveSpanDictionary	1	100.0%	0
RemoveDictionary	401	40,100.0%	200
RemoveRandomSpanDictionary	1	100.0%	0
RemoveRandomDictionary	229	22,900.0%	131
ForeachSpanDictionary	1,528	100.0%	0
ForeachDictionary	1,295	84.8%	0

実行環境: Windows11 x64 .NET Runtime 9.0.0
Score は高いほどパフォーマンスがよいです。
GC0 はガベージコレクション回数を表します（少ないほどパフォーマンスがよい）。

• Remove() が非常に遅いので、使い物になりません。
• なんか全体的にパフォーマンス悪いねェ・・・