Merge pull request #20918 from lorentey/dictionary-modify-layering

[stdlib] Refactor Dictionary.subscript._modify for better layering

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -794,44 +794,7 @@ extension Dictionary {
       }
     }
     _modify {
-      // FIXME: This code should be moved to _variant, with Dictionary.subscript
-      // yielding `&_variant[key]`.
-
-      let (bucket, found) = _variant.mutatingFind(key)
-
-      // FIXME: Mark this entry as being modified in hash table metadata
-      // so that lldb can recognize it's not valid.
-
-      // Move the old value (if any) out of storage, wrapping it into an
-      // optional before yielding it.
-      let native = _variant.asNative
-      if found {
-        var value: Value? = (native._values + bucket.offset).move()
-        yield &value
-        if let value = value {
-          // **Mutation**
-          //
-          // Initialize storage to new value.
-          (native._values + bucket.offset).initialize(to: value)
-        } else {
-          // **Removal**
-          //
-          // We've already deinitialized the value; deinitialize the key too and
-          // register the removal.
-          (native._keys + bucket.offset).deinitialize(count: 1)
-          native._delete(at: bucket)
-        }
-      } else {
-        var value: Value? = nil
-        yield &value
-        if let value = value {
-          // **Insertion**
-          //
-          // Insert the new entry at the correct place.  Note that
-          // `mutatingFind` already ensured that we have enough capacity.
-          native._insert(at: bucket, key: key, value: value)
-        }
-      }
+      yield &_variant[key]
       _fixLifetime(self)
     }
   }

stdlib/public/core/DictionaryVariant.swift

@@ -278,6 +278,31 @@ extension Dictionary._Variant: _DictionaryBuffer {
   }
 }
 
+extension Dictionary._Variant {
+  @inlinable
+  internal subscript(key: Key) -> Value? {
+    @inline(__always)
+    get {
+      return lookup(key)
+    }
+    @inline(__always)
+    _modify {
+#if _runtime(_ObjC)
+      guard isNative else {
+        let cocoa = asCocoa
+        var native = _NativeDictionary<Key, Value>(
+          cocoa, capacity: cocoa.count + 1)
+        self = .init(native: native)
+        yield &native[key, isUnique: true]
+        return
+      }
+#endif
+      let isUnique = isUniquelyReferenced()
+      yield &asNative[key, isUnique: isUnique]
+    }
+  }
+}
+
 extension Dictionary._Variant {
   /// Same as find(_:), except assume a corresponding key/value pair will be
   /// inserted if it doesn't already exist, and mutated if it does exist. When

stdlib/public/core/NativeDictionary.swift

@@ -387,6 +387,44 @@ extension _NativeDictionary: _DictionaryBuffer {
   }
 }
 
+extension _NativeDictionary {
+  @inlinable
+  subscript(key: Key, isUnique isUnique: Bool) -> Value? {
+    @inline(__always)
+    get {
+      // Dummy definition; don't use.
+      return lookup(key)
+    }
+    @inline(__always)
+    _modify {
+      let (bucket, found) = mutatingFind(key, isUnique: isUnique)
+      if found {
+        // Move the old value out of storage, wrapping it into an optional
+        // before yielding it.
+        var value: Value? = (_values + bucket.offset).move()
+        yield &value
+        if let value = value {
+          // **Mutation.** Initialize storage to new value.
+          (_values + bucket.offset).initialize(to: value)
+        } else {
+          // **Removal.** We've already deinitialized the value; deinitialize
+          // the key too and register the removal.
+          (_keys + bucket.offset).deinitialize(count: 1)
+          _delete(at: bucket)
+        }
+      } else {
+        var value: Value? = nil
+        yield &value
+        if let value = value {
+          // **Insertion.** Insert the new entry at the correct place.  Note
+          // that `mutatingFind` already ensured that we have enough capacity.
+          _insert(at: bucket, key: key, value: value)
+        }
+      }
+    }
+  }
+}
+
 // This function has a highly visible name to make it stand out in stack traces.
 @usableFromInline
 @inline(never)