ImmutableKeyedCollection.cs

using System;
using System.Collections;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Collections.ObjectModel;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using System.Threading;

/// <summary>
/// Creation methods for <see cref="ImmutableKeyedCollection{TKey, TItem}"/>, an immutable counterpart of <see
/// cref="KeyedCollection{TKey,TItem}"/>.
/// </summary>
public static class ImmutableKeyedCollection
{
    public static ImmutableKeyedCollection<TKey, TItem> CreateRange<TKey, TItem>(Func<TItem, TKey> keySelector, params IEnumerable<TItem> items)
        where TKey : notnull
    {
        return CreateRange(keySelector, null, items);
    }

    public static ImmutableKeyedCollection<TKey, TItem> CreateRange<TKey, TItem>(Func<TItem, TKey> keySelector, params ImmutableArray<TItem> items)
        where TKey : notnull
    {
        return CreateRange(keySelector, null, items);
    }

    public static ImmutableKeyedCollection<TKey, TItem> CreateRange<TKey, TItem>(Func<TItem, TKey> keySelector, IEqualityComparer<TKey>? keyComparer, params IEnumerable<TItem> items)
        where TKey : notnull
    {
        return CreateRange(keySelector, keyComparer, items.ToImmutableArray());
    }

    public static ImmutableKeyedCollection<TKey, TItem> CreateRange<TKey, TItem>(Func<TItem, TKey> keySelector, IEqualityComparer<TKey>? keyComparer, params ImmutableArray<TItem> items)
        where TKey : notnull
    {
        var keys = ImmutableArray.CreateRange(items, keySelector);
        return new(items, keys, keyComparer);
    }
}

/// <summary>
/// Preserves item order while also allowing fast key lookups. Immutable counterpart of <see
/// cref="KeyedCollection{TKey,TItem}"/>.
/// </summary>
public sealed class ImmutableKeyedCollection<TKey, TItem> : IReadOnlyList<TItem>, IReadOnlyDictionary<TKey, TItem>
    where TKey : notnull
{
    private const int MinHashTableSize = 15;

    private readonly ImmutableArray<TItem> items;
    private readonly ImmutableArray<TKey> keys;
    private readonly IEqualityComparer<TKey> keyComparer;


    /// <summary>
    /// Use <see cref="GetCachedIndexesByKey"/> rather than referencing this field directly.
    /// </summary>
    private Dictionary<TKey, int>? cachedIndexesByKey;

    public ImmutableKeyedCollection(ImmutableArray<TItem> items, ImmutableArray<TKey> keys, IEqualityComparer<TKey>? keyComparer)
    {
        if (items.IsDefault) items = [];
        if (keys.IsDefault) keys = [];
        if (items.Length != keys.Length) throw new ArgumentException("There must be one key for each item, matched by index.", nameof(keys));

        this.items = items;
        this.keys = keys;
        this.keyComparer = keyComparer ?? EqualityComparer<TKey>.Default;
    }

    private Dictionary<TKey, int> GetCachedIndexesByKey()
    {
        Debug.Assert(keys.Length >= MinHashTableSize);

        var cachedIndexesByKey = Volatile.Read(ref this.cachedIndexesByKey);
        if (cachedIndexesByKey is null)
        {
            var newCache = new Dictionary<TKey, int>(keyComparer);

            for (var i = 0; i < keys.Length; i++)
            {
                // If the keys aren't unique, behave like we would when we're under the minimum hash table size (first
                // index wins).
                _ = newCache.TryAdd(keys[i], i);
            }

            var otherCacheWhichWonTheRace = Interlocked.CompareExchange(ref this.cachedIndexesByKey, newCache, null);

            // If another thread was racing to create this dictionary, the dictionaries are interchangeable. It would be
            // better for all threads to agree on one dictionary so that the extra dictionaries are garbage-collected as
            // soon as possible. So, use the dictionary which beat the one we were building, unless there was none, in
            // which case we are the thread which won.
            cachedIndexesByKey = otherCacheWhichWonTheRace ?? newCache;
        }

        return cachedIndexesByKey;
    }

    public int IndexOf(TKey key)
    {
        if (keys.Length < MinHashTableSize)
            return keys.IndexOf(key, keyComparer);

        return GetCachedIndexesByKey().GetValueOrDefault(key, -1);
    }

    public bool TryGetValue(TKey key, [MaybeNullWhen(false)] out TItem item)
    {
        var index = IndexOf(key);
        if (index != -1)
        {
            item = items[index];
            return true;
        }
        item = default;
        return false;
    }

    public TItem this[TKey key] => IndexOf(key) switch
    {
        -1 => throw new KeyNotFoundException(),
        var index => items[index],
    };

    public bool ContainsKey(TKey key) => IndexOf(key) != -1;

    public IEnumerable<TKey> Keys => keys;

    IEnumerable<TItem> IReadOnlyDictionary<TKey, TItem>.Values => items;

    public int Count => items.Length;

    public TItem this[int index] => items[index];

    IEnumerator<KeyValuePair<TKey, TItem>> IEnumerable<KeyValuePair<TKey, TItem>>.GetEnumerator()
    {
        return keys.Zip(items, KeyValuePair.Create).GetEnumerator();
    }

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

    public IEnumerator<TItem> GetEnumerator() => ((IEnumerable<TItem>)items).GetEnumerator();
}