.NET-Community-Toolkit/CommunityToolkit.Mvvm/Messaging/Internals/Microsoft.Collections.Extensions/DictionarySlim{TKey,TValue}.cs

// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

// The DictionarySlim<TKey, TValue> type is originally from CoreFX labs, see
// the source repository on GitHub at https://github.com/dotnet/corefxlab.

using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Diagnostics.Contracts;
using System.Runtime.CompilerServices;

namespace Microsoft.Collections.Extensions;

/// <summary>
/// A lightweight Dictionary with three principal differences compared to <see cref="Dictionary{TKey, TValue}"/>
///
/// 1) It is possible to do "get or add" in a single lookup. For value types, this also saves a copy of the value.
/// 2) It assumes it is cheap to equate values.
/// 3) It assumes the keys implement <see cref="IEquatable{TKey}"/> and they are cheap and sufficient.
/// </summary>
/// <typeparam name="TKey">The type of keys in the dictionary.</typeparam>
/// <typeparam name="TValue">The type of values in the dictionary.</typeparam>
/// <remarks>
/// 1) This avoids having to do separate lookups (<see cref="Dictionary{TKey, TValue}.TryGetValue(TKey, out TValue)"/>
/// followed by <see cref="Dictionary{TKey, TValue}.Add(TKey, TValue)"/>.
/// There is not currently an API exposed to get a value by ref without adding if the key is not present.
/// 2) This means it can save space by not storing hash codes.
/// 3) This means it can avoid storing a comparer, and avoid the likely virtual call to a comparer.
/// </remarks>
[DebuggerDisplay("Count = {Count}")]
internal class DictionarySlim<TKey, TValue> : IDictionarySlim<TKey, TValue>
    where TKey : IEquatable<TKey>
    where TValue : class
{
    /// <summary>
    /// A reusable array of <see cref="Entry"/> items with a single value.
    /// This is used when a new <see cref="DictionarySlim{TKey,TValue}"/> instance is
    /// created, or when one is cleared. The first item being added to this collection
    /// will immediately cause the first resize (see <see cref="AddKey"/> for more info).
    /// </summary>
    private static readonly Entry[] InitialEntries = new Entry[1];

    /// <summary>
    /// The current number of items stored in the map.
    /// </summary>
    private int count;

    /// <summary>
    /// The 1-based index for the start of the free list within <see cref="entries"/>.
    /// </summary>
    private int freeList = -1;

    /// <summary>
    /// The array of 1-based indices for the <see cref="Entry"/> items stored in <see cref="entries"/>.
    /// </summary>
    private int[] buckets;

    /// <summary>
    /// The array of currently stored key-value pairs (ie. the lists for each hash group).
    /// </summary>
    private Entry[] entries;

    /// <summary>
    /// A type representing a map entry, ie. a node in a given list.
    /// </summary>
    private struct Entry
    {
        /// <summary>
        /// The key for the value in the current node.
        /// </summary>
        public TKey Key;

        /// <summary>
        /// The value in the current node, if present.
        /// </summary>
        public TValue? Value;

        /// <summary>
        /// The 0-based index for the next node in the current list.
        /// </summary>
        public int Next;
    }

    /// <summary>
    /// Initializes a new instance of the <see cref="DictionarySlim{TKey, TValue}"/> class.
    /// </summary>
    public DictionarySlim()
    {
        this.buckets = HashHelpers.SizeOneIntArray;
        this.entries = InitialEntries;
    }

    /// <inheritdoc/>
    public int Count => this.count;

    /// <inheritdoc/>
    public TValue this[TKey key]
    {
        get
        {
            Entry[] entries = this.entries;

            for (int i = this.buckets[key.GetHashCode() & (this.buckets.Length - 1)] - 1;
                 (uint)i < (uint)entries.Length;
                 i = entries[i].Next)
            {
                if (key.Equals(entries[i].Key))
                {
                    return entries[i].Value!;
                }
            }

            ThrowArgumentExceptionForKeyNotFound(key);

            return default!;
        }
    }

    /// <inheritdoc/>
    public void Clear()
    {
        this.count = 0;
        this.freeList = -1;
        this.buckets = HashHelpers.SizeOneIntArray;
        this.entries = InitialEntries;
    }

    /// <summary>
    /// Checks whether or not the dictionary contains a pair with a specified key.
    /// </summary>
    /// <param name="key">The key to look for.</param>
    /// <returns>Whether or not the key was present in the dictionary.</returns>
    public bool ContainsKey(TKey key)
    {
        Entry[] entries = this.entries;

        for (int i = this.buckets[key.GetHashCode() & (this.buckets.Length - 1)] - 1;
             (uint)i < (uint)entries.Length;
             i = entries[i].Next)
        {
            if (key.Equals(entries[i].Key))
            {
                return true;
            }
        }

        return false;
    }

    /// <summary>
    /// Gets the value if present for the specified key.
    /// </summary>
    /// <param name="key">The key to look for.</param>
    /// <param name="value">The value found, otherwise <see langword="default"/>.</param>
    /// <returns>Whether or not the key was present.</returns>
    public bool TryGetValue(TKey key, [NotNullWhen(true)] out TValue? value)
    {
        Entry[] entries = this.entries;

        for (int i = this.buckets[key.GetHashCode() & (this.buckets.Length - 1)] - 1;
             (uint)i < (uint)entries.Length;
             i = entries[i].Next)
        {
            if (key.Equals(entries[i].Key))
            {
                value = entries[i].Value!;

                return true;
            }
        }

        value = default!;

        return false;
    }

    /// <inheritdoc/>
    public bool TryRemove(TKey key)
    {
        Entry[] entries = this.entries;
        int bucketIndex = key.GetHashCode() & (this.buckets.Length - 1);
        int entryIndex = this.buckets[bucketIndex] - 1;
        int lastIndex = -1;

        while (entryIndex != -1)
        {
            Entry candidate = entries[entryIndex];

            if (candidate.Key.Equals(key))
            {
                if (lastIndex != -1)
                {
                    entries[lastIndex].Next = candidate.Next;
                }
                else
                {
                    this.buckets[bucketIndex] = candidate.Next + 1;
                }

                entries[entryIndex] = default;
                entries[entryIndex].Next = -3 - this.freeList;

                this.freeList = entryIndex;
                this.count--;

                return true;
            }

            lastIndex = entryIndex;
            entryIndex = candidate.Next;
        }

        return false;
    }

    /// <summary>
    /// Gets the value for the specified key, or, if the key is not present,
    /// adds an entry and returns the value by ref. This makes it possible to
    /// add or update a value in a single look up operation.
    /// </summary>
    /// <param name="key">Key to look for</param>
    /// <returns>Reference to the new or existing value</returns>
    public ref TValue? GetOrAddValueRef(TKey key)
    {
        Entry[] entries = this.entries;
        int bucketIndex = key.GetHashCode() & (this.buckets.Length - 1);

        for (int i = this.buckets[bucketIndex] - 1;
             (uint)i < (uint)entries.Length;
             i = entries[i].Next)
        {
            if (key.Equals(entries[i].Key))
            {
                return ref entries[i].Value;
            }
        }

        return ref AddKey(key, bucketIndex);
    }

    /// <summary>
    /// Creates a slot for a new value to add for a specified key.
    /// </summary>
    /// <param name="key">The key to use to add the new value.</param>
    /// <param name="bucketIndex">The target bucked index to use.</param>
    /// <returns>A reference to the slot for the new value to add.</returns>
    [MethodImpl(MethodImplOptions.NoInlining)]
    private ref TValue? AddKey(TKey key, int bucketIndex)
    {
        Entry[] entries = this.entries;
        int entryIndex;

        if (this.freeList != -1)
        {
            entryIndex = this.freeList;

            this.freeList = -3 - entries[this.freeList].Next;
        }
        else
        {
            if (this.count == entries.Length || entries.Length == 1)
            {
                entries = Resize();
                bucketIndex = key.GetHashCode() & (this.buckets.Length - 1);
            }

            entryIndex = this.count;
        }

        entries[entryIndex].Key = key;
        entries[entryIndex].Next = this.buckets[bucketIndex] - 1;

        this.buckets[bucketIndex] = entryIndex + 1;
        this.count++;

        return ref entries[entryIndex].Value;
    }

    /// <summary>
    /// Resizes the current dictionary to reduce the number of collisions
    /// </summary>
    [MethodImpl(MethodImplOptions.NoInlining)]
    private Entry[] Resize()
    {
        int count = this.count;
        int newSize = this.entries.Length * 2;

        if ((uint)newSize > int.MaxValue)
        {
            ThrowInvalidOperationExceptionForMaxCapacityExceeded();
        }

        DictionarySlim<TKey, TValue>.Entry[]? entries = new Entry[newSize];

        Array.Copy(this.entries, 0, entries, 0, count);

        int[]? newBuckets = new int[entries.Length];

        while (count-- > 0)
        {
            int bucketIndex = entries[count].Key.GetHashCode() & (newBuckets.Length - 1);

            entries[count].Next = newBuckets[bucketIndex] - 1;

            newBuckets[bucketIndex] = count + 1;
        }

        this.buckets = newBuckets;
        this.entries = entries;

        return entries;
    }

    /// <inheritdoc cref="IEnumerable{T}.GetEnumerator"/>
    [Pure]
    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    public Enumerator GetEnumerator() => new(this);

    /// <summary>
    /// Enumerator for <see cref="DictionarySlim{TKey,TValue}"/>.
    /// </summary>
    public ref struct Enumerator
    {
        private readonly Entry[] entries;
        private int index;
        private int count;

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        internal Enumerator(DictionarySlim<TKey, TValue> dictionary)
        {
            this.entries = dictionary.entries;
            this.index = 0;
            this.count = dictionary.count;
        }

        /// <inheritdoc cref="IEnumerator.MoveNext"/>
        public bool MoveNext()
        {
            if (this.count == 0)
            {
                return false;
            }

            this.count--;

            Entry[] entries = this.entries;

            while (entries[this.index].Next < -1)
            {
                this.index++;
            }

            // We need to preemptively increment the current index so that we still correctly keep track
            // of the current position in the dictionary even if the users doesn't access any of the
            // available properties in the enumerator. As this is a possibility, we can't rely on one of
            // them to increment the index before MoveNext is invoked again. We ditch the standard enumerator
            // API surface here to expose the Key/Value properties directly and minimize the memory copies.
            // For the same reason, we also removed the KeyValuePair<TKey, TValue> field here, and instead
            // rely on the properties lazily accessing the target instances directly from the current entry
            // pointed at by the index property (adjusted backwards to account for the increment here).
            this.index++;

            return true;
        }

        /// <summary>
        /// Gets the current key.
        /// </summary>
        public TKey Key
        {
            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            get => this.entries[this.index - 1].Key;
        }

        /// <summary>
        /// Gets the current value.
        /// </summary>
        public TValue Value
        {
            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            get => this.entries[this.index - 1].Value!;
        }
    }

    /// <summary>
    /// Throws an <see cref="ArgumentException"/> when trying to load an element with a missing key.
    /// </summary>
    private static void ThrowArgumentExceptionForKeyNotFound(TKey key)
    {
        throw new ArgumentException($"The target key {key} was not present in the dictionary");
    }

    /// <summary>
    /// Throws an <see cref="InvalidOperationException"/> when trying to resize over the maximum capacity.
    /// </summary>
    private static void ThrowInvalidOperationExceptionForMaxCapacityExceeded()
    {
        throw new InvalidOperationException("Max capacity exceeded");
    }
}