AddOnlyThreadSafeSeparateChainingDictionary in C#

AddOnlyThreadSafeSeparateChainingDictionary.cs

/*
This is free and unencumbered software released into the public domain.

Anyone is free to copy, modify, publish, use, compile, sell, or
distribute this software, either in source code form or as a compiled
binary, for any purpose, commercial or non-commercial, and by any
means.

In jurisdictions that recognize copyright laws, the author or authors
of this software dedicate any and all copyright interest in the
software to the public domain. We make this dedication for the benefit
of the public at large and to the detriment of our heirs and
successors. We intend this dedication to be an overt act of
relinquishment in perpetuity of all present and future rights to this
software under copyright law.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

For more information, please refer to <http://unlicense.org/>
*/

#nullable enable
using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.CompilerServices;
using System.Threading;

// ReSharper disable LocalVariableHidesMember


public sealed class AddOnlyThreadSafeSeparateChainingDictionary<TKey, TValue>
    where TKey : notnull
{
    sealed class Table
    {
        public readonly int[] Buckets; // -1 init
        public readonly TKey?[] Keys;
        public readonly TValue?[] Values;
        public readonly int[] Nexts; // -1 init
        public readonly int[] Hashes; //-1 init

        public readonly int Mask; // capacity-1 (power of 2)

        public readonly int ResizeStartCount;
        public int Count;
        public int ActualCount => Math.Min(Count, Mask + 1);

        public Table(int capacityPow2)
        {
            int cap = RoundUpPower2(Math.Max(2, capacityPow2));
            Mask = cap - 1;
            ResizeStartCount = (int)(Mask * 0.75);
            Buckets = new int[cap];
            Array.Fill(Buckets, -1);

            Keys = new TKey?[cap];
            Values = new TValue?[cap];
            Hashes = new int[cap];
            Nexts = new int[cap];
            Buckets.CopyTo(Hashes);
            Buckets.CopyTo(Nexts);
            Count = 0;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public int BucketIndex(int hashCode) => hashCode & Mask;
    }

    Table table;
    int activeWriters; // writer critical section count
    readonly object activeWritersLock = new object();
    bool waitingActiveWriters;
    int resizingFlag; // 0/1
    readonly object resizingLock = new object();
    int resizingFullLockFlag; // 0/1
    readonly object resizingFullLock = new object();
    object[] locks;
    readonly IEqualityComparer<TKey> comparer;

    public AddOnlyThreadSafeSeparateChainingDictionary(int capacity = 16, IEqualityComparer<TKey>? comparer = null)
    {
        var concurrencyLevel = Environment.ProcessorCount;
        var locks = new object[concurrencyLevel];
        locks[0] = locks;
        for (int i = 1; i < locks.Length; i++)
        {
            locks[i] = new object();
        }

        this.locks = locks;
        table = new Table(capacity);
        this.comparer = comparer ?? EqualityComparer<TKey>.Default;
    }

    public int Count => table.ActualCount;


    public bool TryGetValue(TKey key, [NotNullWhen(true)] out TValue? value)
    {
        int hash = comparer.GetHashCode(key) & int.MaxValue;
        value = default!;
        return TryGetValue(Volatile.Read(ref table), key, hash, out _, ref value);
    }

    bool TryGetValue(Table t, TKey key, int hash, out int firstEntry, ref TValue value)
    {
        int b = t.BucketIndex(hash);
        var nexts = t.Nexts;
        var endE = -1;
        firstEntry = Volatile.Read(ref t.Buckets[b]);
        for (int e = firstEntry; e != endE; e = nexts[e])
        {
            if (hash == t.Hashes[e] && comparer.Equals(t.Keys[e]!, key))
            {
                value = t.Values[e]!;
                return true;
            }
        }

        return false;
    }

    public TValue GetOrAdd(TKey key, Func<TKey, TValue> valueFactory)
    {
        if (valueFactory == null) throw new ArgumentNullException(nameof(valueFactory));

        int hash = comparer.GetHashCode(key) & int.MaxValue;
        var value = default(TValue)!;
        var t = Volatile.Read(ref table);
        if (TryGetValue(t, key, hash, out int entry, ref value))
        {
            return value;
        }

        TryAddInternal(t, key, hash, entry, ref value, true,
            static (key, valueFactory) => Unsafe.As<Func<TKey, TValue>>(valueFactory)!(key)!,
            valueFactory);
        return value;
    }

    public TValue GetOrAdd(TKey key, Func<TKey, object?, TValue> valueFactory, object arg)
    {
        int hash = comparer.GetHashCode(key) & int.MaxValue;
        var value = default(TValue)!;
        var t = Volatile.Read(ref table);
        if (TryGetValue(t, key, hash, out int entry, ref value))
        {
            return value;
        }

        TryAddInternal(t, key, hash, entry, ref value, true, valueFactory, arg);
        return value;
    }

    public TValue GetOrAdd(TKey key, TValue value)
    {
        int hash = comparer.GetHashCode(key) & int.MaxValue;
        var t = Volatile.Read(ref table);
        if (TryGetValue(t, key, hash, out int entry, ref value))
        {
            return value;
        }

        TryAddInternal(t, key, hash, entry, ref value, true);
        return value;
    }

    public bool TryAdd(TKey key, TValue value)
    {
        int hash = comparer.GetHashCode(key) & int.MaxValue;
        var t = Volatile.Read(ref table);

        if (TryGetValue(t, key, hash, out int entry, ref value))
        {
            return false;
        }

        return TryAddInternal(t, key, hash, entry, ref value, true);
    }

    bool TryAddInternal(Table t, TKey key, int hash, int endEntry, ref TValue value, bool skip,
        Func<TKey, object?, TValue>? valueFactory = null, object? state = null)
    {
        for (;;)
        {
            if (!skip)
            {
                t = Volatile.Read(ref table);
                if (TryGetValue(t, key, hash, out endEntry, ref value))
                {
                    return false;
                }
            }
            else
            {
                skip = false;
            }

            int b = t.BucketIndex(hash);


            lock (locks[b % locks.Length])
            {
                var nexts = t.Nexts;
                // retry search in lock
                for (int e = Volatile.Read(ref t.Buckets[b]); e != endEntry; e = nexts[e])
                {
                    if (hash == t.Hashes[e] && comparer.Equals(t.Keys[e]!, key))
                    {
                        value = t.Values[e]!;
                        return false;
                    }
                }

                Interlocked.Increment(ref activeWriters);

                if (Volatile.Read(ref resizingFullLockFlag) != 0)
                {
                    if (Interlocked.Decrement(ref activeWriters) == 1 && Volatile.Read(ref resizingFlag) == 1)
                    {
                        PulseExitWrite();
                    }

                    WaitResizeFullLock();

                    continue;
                }

                try
                {
                    if (t != Volatile.Read(ref table))
                    {
                        continue;
                    }

                    int idx = Interlocked.Increment(ref t.Count) - 1;
                    if (idx < t.Keys.Length)
                    {
                        int oldHead = Volatile.Read(ref t.Buckets[b]);
                        if (valueFactory != null)
                        {
                            value = valueFactory(key, state);
                        }

                        t.Keys[idx] = key;
                        t.Values[idx] = value;
                        t.Nexts[idx] = oldHead;
                        t.Hashes[idx] = hash;
                        Volatile.Write(ref t.Buckets[b], idx);

                        if (idx == t.ResizeStartCount)
                        {
                            StartResize();
                        }

                        return true;
                    }
                }
                finally
                {
                    if (Interlocked.Decrement(ref activeWriters) == 1 && Volatile.Read(ref resizingFlag) == 1)
                    {
                        PulseExitWrite();
                    }
                }
            }

            if (t != Volatile.Read(ref table)) continue;

            WaitResize();
        }
    }

    void PulseExitWrite()
    {
        lock (activeWritersLock)
        {
            if (waitingActiveWriters)
            {
                Monitor.Pulse(activeWritersLock);
            }
        }
    }

    void WaitResizeFullLock()
    {
        lock (resizingFullLock)
        {
            while (Volatile.Read(ref resizingFullLockFlag) != 0)
            {
                Monitor.Wait(resizingFullLock);
            }
        }
    }

    void WaitResize()
    {
        lock (resizingLock)
        {
            while (Volatile.Read(ref resizingFlag) == 0 &&
                   Volatile.Read(ref table).Count >= Volatile.Read(ref table).Keys.Length)
            {
                Monitor.Wait(resizingLock);
            }

            while (Volatile.Read(ref resizingFlag) == 1)
            {
                Monitor.Wait(resizingLock);
            }
        }
    }

    // ----------------------------
    // Resize: quiescent resize + table swap
    // ----------------------------
    void StartResize()
    {
        lock (resizingLock)
        {
            while (Volatile.Read(ref resizingFlag) == 1)
            {
                Monitor.Wait(resizingLock);
            }

            Volatile.Write(ref resizingFlag, 1);
        }

        try
        {
            var oldT = Volatile.Read(ref table);

            var newT = new Table(oldT.Keys.Length << 1);

            int limit = oldT.ActualCount;
            for (int i = 0; i < limit; i++)
            {
                int hash = oldT.Hashes[i];
                if (hash < 0)
                {
                    limit = i;
                    break;
                }

                int b = newT.BucketIndex(hash);
                int oldHead = newT.Buckets[b];
                var k = oldT.Keys[i]!;
                var v = oldT.Values[i];
                int idx = newT.Count++;
                newT.Keys[idx] = k;
                newT.Values[idx] = v;
                newT.Nexts[idx] = oldHead;
                newT.Hashes[idx] = hash;
                newT.Buckets[b] = idx;
            }

            lock (resizingFullLock)
            {
                resizingFullLockFlag = 1;
            }

            if (Volatile.Read(ref activeWriters) > 1)
            {
                lock (activeWritersLock)
                {
                    waitingActiveWriters = true;
                    while (Volatile.Read(ref activeWriters) > 1)
                    {
                        Monitor.Wait(activeWritersLock);
                    }
                }
            }

            var newLimit = oldT.ActualCount;
            for (int i = limit; i < newLimit; i++)
            {
                ref int hash = ref oldT.Hashes[i];
                var k = oldT.Keys[i]!;
                var v = oldT.Values[i];
                int b = newT.BucketIndex(hash);
                int idx = newT.Count++;
                int oldHead = newT.Buckets[b];
                newT.Keys[idx] = k;
                newT.Values[idx] = v;
                newT.Hashes[idx] = hash;
                newT.Nexts[idx] = oldHead;
                newT.Buckets[b] = idx;
            }

            Volatile.Write(ref table, newT);
            lock (resizingFullLock)
            {
                Volatile.Write(ref resizingFullLockFlag, 0);
                Monitor.PulseAll(resizingFullLock);
            }
        }
        finally
        {
            lock (resizingLock)
            {
                Volatile.Write(ref resizingFlag, 0);
                Monitor.PulseAll(resizingLock);
            }
        }
    }

    static int RoundUpPower2(int x)
    {
        if (x <= 1) return 1;
        x--;
        x |= x >> 1;
        x |= x >> 2;
        x |= x >> 4;
        x |= x >> 8;
        x |= x >> 16;
        x++;
        return x;
    }
}