Optimize dictionary access with CollectionsMarshal.GetValueRefOrAddDefault within Trie and Lookup

Author: akade

Akade.IndexedSet/DataStructures/Lookup.cs

@@ -1,4 +1,6 @@
-namespace Akade.IndexedSet.DataStructures;
+using System.Runtime.InteropServices;
+
+namespace Akade.IndexedSet.DataStructures;
 
 /// <summary>
 /// Modifiable lookup based on <see cref="Dictionary{TKey, TValue}"/> with <see cref="HashSet{T}"/> as value collection per key.
@@ -11,10 +13,8 @@ internal class Lookup<TKey, TValue>
 
     public bool Add(TKey key, TValue value)
     {
-        if (!_values.TryGetValue(key, out HashSet<TValue>? keySet))
-        {
-            keySet = _values[key] = [];
-        }
+        ref HashSet<TValue>? keySet = ref CollectionsMarshal.GetValueRefOrAddDefault(_values, key, out _);
+        keySet ??= [];
 
         if (keySet.Add(value))
         {

Akade.IndexedSet/DataStructures/Trie.cs

@@ -1,6 +1,7 @@
 using Akade.IndexedSet.Extensions;
 using Akade.IndexedSet.Utils;
 using System.Diagnostics.CodeAnalysis;
+using System.Runtime.InteropServices;
 
 namespace Akade.IndexedSet.DataStructures;
 
@@ -177,11 +178,10 @@ internal bool Add(ReadOnlySpan<char> key, TElement element)
             else
             {
                 _children ??= [];
-                if (!_children.TryGetValue(key[0], out TrieNode? trieNode))
-                {
-                    _children[key[0]] = trieNode = new();
-                }
 
+                ref TrieNode? trieNode = ref CollectionsMarshal.GetValueRefOrAddDefault(_children, key[0], out _);
+
+                trieNode ??= new();
                 return trieNode.Add(key[1..], element);
             }
         }