Cleanup custom dictionary implementations

Jint.Benchmark/DictionaryBenchmark.cs

@@ -31,7 +31,7 @@ public void HybridDictionary()
         var hybridDictionary = new HybridDictionary<object>();
         for (var i = 0; i < N; i++)
         {
-            hybridDictionary.Add(_keys[i], _keys);
+            hybridDictionary[_keys[i]] = _keys;
         }
 
         foreach (var key in _keys)
@@ -69,4 +69,4 @@ public void StringDictionarySlim()
             dictionary.ContainsKey(key);
         }
     }
-}
+}

Jint/Collections/DictionaryBase.cs

@@ -0,0 +1,37 @@
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+
+namespace Jint.Collections;
+
+internal abstract class DictionaryBase<TValue> : IEngineDictionary<Key, TValue>
+{
+    public ref TValue this[Key key]
+    {
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        get => ref GetValueRefOrAddDefault(key, out _);
+    }
+
+    public bool TryGetValue(Key key, [NotNullWhen(true)] out TValue? value)
+    {
+        value = default;
+        ref var temp = ref GetValueRefOrNullRef(key);
+        if (Unsafe.IsNullRef(ref temp))
+        {
+            return false;
+        }
+
+        value = temp!;
+        return true;
+    }
+
+    public bool ContainsKey(Key key)
+    {
+        ref var valueRefOrNullRef = ref GetValueRefOrNullRef(key);
+        return !Unsafe.IsNullRef(ref valueRefOrNullRef);
+    }
+
+    public abstract int Count { get; }
+
+    public abstract ref TValue GetValueRefOrNullRef(Key key);
+    public abstract ref TValue GetValueRefOrAddDefault(Key key, out bool exists);
+}

Jint/Collections/DictionarySlim.cs

@@ -18,7 +18,7 @@ namespace Jint.Collections;
 /// 3) It does not accept an equality comparer (assumes Object.GetHashCode() and Object.Equals() or overridden implementation are cheap and sufficient).
 /// </summary>
 [DebuggerDisplay("Count = {Count}")]
-internal class DictionarySlim<TKey, TValue> : IReadOnlyCollection<KeyValuePair<TKey, TValue>> where TKey : IEquatable<TKey>
+internal sealed class DictionarySlim<TKey, TValue> : IReadOnlyCollection<KeyValuePair<TKey, TValue>> where TKey : IEquatable<TKey>
 {
     // We want to initialize without allocating arrays. We also want to avoid null checks.
     // Array.Empty would give divide by zero in modulo operation. So we use static one element arrays.

Jint/Collections/HybridDictionary.cs

@@ -5,7 +5,7 @@
 
 namespace Jint.Collections;
 
-internal class HybridDictionary<TValue> : IEnumerable<KeyValuePair<Key, TValue>>
+internal sealed class HybridDictionary<TValue> : IEngineDictionary<Key, TValue>, IEnumerable<KeyValuePair<Key, TValue>>
 {
     private const int CutoverPoint = 9;
     private const int InitialDictionarySize = 13;
@@ -28,40 +28,34 @@ public HybridDictionary(int initialSize, bool checkExistingKeys)
         }
     }
 
-    protected HybridDictionary(StringDictionarySlim<TValue> dictionary)
+    public HybridDictionary(StringDictionarySlim<TValue> dictionary)
     {
         _checkExistingKeys = true;
         _dictionary = dictionary;
     }
 
-    public TValue this[Key key]
+    public ref TValue this[Key key]
     {
         get
-        {
-            TryGetValue(key, out var value);
-            return value;
-        }
-        set
         {
             if (_dictionary != null)
             {
-                _dictionary[key] = value;
+                return ref _dictionary[key];
             }
-            else if (_list != null)
+
+            if (_list != null)
             {
                 if (_list.Count >= CutoverPoint - 1)
                 {
-                    SwitchToDictionary(key, value, tryAdd: false);
+                    return ref SwitchToDictionary(key);
                 }
-                else
-                {
-                    _list[key] = value;
-                }
-            }
-            else
-            {
-                _list = new ListDictionary<TValue>(key, value, _checkExistingKeys);
+
+                return ref _list[key];
             }
+
+            var head = new ListDictionary<TValue>.DictionaryNode { Key = key, Value = default };
+            _list = new ListDictionary<TValue>(head, _checkExistingKeys);
+            return ref head.Value;
         }
     }
 
@@ -81,20 +75,48 @@ public bool TryGetValue(Key key, out TValue value)
         return false;
     }
 
-    public void SetOrUpdateValue<TState>(Key key, Func<TValue, TState, TValue> updater, TState state)
+    public ref TValue GetValueRefOrNullRef(Key key)
     {
         if (_dictionary != null)
         {
-            _dictionary.SetOrUpdateValue(key, updater, state);
+            return ref _dictionary.GetValueRefOrNullRef(key);
+        }
+
+        if (_list != null)
+        {
+            return ref _list.GetValueRefOrNullRef(key);
         }
-        else if (_list != null)
+
+        return ref Unsafe.NullRef<TValue>();
+    }
+
+    public ref TValue GetValueRefOrAddDefault(Key key, out bool exists)
+    {
+        if (_dictionary != null)
         {
-            _list.SetOrUpdateValue(key, updater, state);
+            return ref _dictionary.GetValueRefOrAddDefault(key, out exists);
         }
-        else
+
+        if (_list != null)
         {
-            _list = new ListDictionary<TValue>(key, updater(default, state), _checkExistingKeys);
+            return ref _list.GetValueRefOrAddDefault(key, out exists);
         }
+
+        var head = new ListDictionary<TValue>.DictionaryNode
+        {
+            Key = key,
+        };
+
+        _list = new ListDictionary<TValue>(head, _checkExistingKeys);
+        exists = false;
+        return ref head.Value;
+    }
+
+    [MethodImpl(MethodImplOptions.AggressiveInlining)]
+    public void SetOrUpdateValue<TState>(Key key, Func<TValue, TState, TValue> updater, TState state)
+    {
+        ref var currentValue = ref GetValueRefOrAddDefault(key, out _);
+        currentValue = updater(currentValue, state);
     }
 
     private bool SwitchToDictionary(Key key, TValue value, bool tryAdd)
@@ -115,11 +137,24 @@ private bool SwitchToDictionary(Key key, TValue value, bool tryAdd)
             dictionary[key] = value;
             result = true;
         }
+
         _dictionary = dictionary;
         _list = null;
         return result;
     }
 
+    private ref TValue SwitchToDictionary(Key key)
+    {
+        var dictionary = new StringDictionarySlim<TValue>(InitialDictionarySize);
+        foreach (var pair in _list)
+        {
+            dictionary[pair.Key] = pair.Value;
+        }
+        _dictionary = dictionary;
+        _list = null;
+        return ref dictionary[key];
+    }
+
     public int Count
     {
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
@@ -132,26 +167,25 @@ public bool TryAdd(Key key, TValue value)
         {
             return _dictionary.TryAdd(key, value);
         }
-        else
-        {
-            _list ??= new ListDictionary<TValue>(key, value, _checkExistingKeys);
 
-            if (_list.Count + 1 >= CutoverPoint)
-            {
-                return SwitchToDictionary(key, value, tryAdd: true);
-            }
-            else
-            {
-                return _list.Add(key, value, tryAdd: true);
-            }
+        _list ??= new ListDictionary<TValue>(key, value, _checkExistingKeys);
+
+        if (_list.Count + 1 >= CutoverPoint)
+        {
+            return SwitchToDictionary(key, value, tryAdd: true);
         }
+
+        return _list.Add(key, value, tryAdd: true);
     }
 
-    public void Add(Key key, TValue value)
+    /// <summary>
+    /// Adds a new item and expects key to not exist.
+    /// </summary>
+    public void AddDangerous(Key key, TValue value)
     {
         if (_dictionary != null)
         {
-            _dictionary.GetOrAddValueRef(key) = value;
+            _dictionary.AddDangerous(key, value);
         }
         else
         {
@@ -163,11 +197,11 @@ public void Add(Key key, TValue value)
             {
                 if (_list.Count + 1 >= CutoverPoint)
                 {
-                    SwitchToDictionary(key, value, tryAdd: false);
+                    SwitchToDictionary(key) = value;
                 }
                 else
                 {
-                    _list.Add(key, value);
+                    _list.AddDangerous(key, value);
                 }
             }
         }
@@ -181,17 +215,8 @@ public void Clear()
 
     public bool ContainsKey(Key key)
     {
-        if (_dictionary != null)
-        {
-            return _dictionary.ContainsKey(key);
-        }
-
-        if (_list != null)
-        {
-            return _list.ContainsKey(key);
-        }
-
-        return false;
+        ref var valueRefOrNullRef = ref GetValueRefOrNullRef(key);
+        return !Unsafe.IsNullRef(ref valueRefOrNullRef);
     }
 
     IEnumerator<KeyValuePair<Key, TValue>> IEnumerable<KeyValuePair<Key, TValue>>.GetEnumerator()
@@ -207,7 +232,6 @@ IEnumerator<KeyValuePair<Key, TValue>> IEnumerable<KeyValuePair<Key, TValue>>.Ge
         }
 
         return System.Linq.Enumerable.Empty<KeyValuePair<Key, TValue>>().GetEnumerator();
-
     }
 
     IEnumerator IEnumerable.GetEnumerator()

Jint/Collections/IEngineDictionary.cs

@@ -0,0 +1,19 @@
+using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
+
+// ReSharper disable once CheckNamespace
+namespace Jint;
+
+/// <summary>
+/// Contract for custom dictionaries that Jint uses.
+/// </summary>
+internal interface IEngineDictionary<in TKey, TValue>
+{
+    int Count { get; }
+
+    ref TValue this[TKey name] { get; }
+
+    public ref TValue GetValueRefOrNullRef(TKey key);
+
+    public ref TValue GetValueRefOrAddDefault(TKey key, out bool exists);
+}