[5.0][stdlib] Support throwing yields in all _modify accessors (#21171)

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

This may also reduce ARC traffic by a little bit.

(cherry picked from commit d32fd28)

* [stdlib] _modify: Use defer to ensure invariants are restored when yield throws

_modify mustn’t leave Collection storage in an inconsistent state when the code to which we’re yielding happens to throw. In practice this means that any cleanup code must happen in defer blocks before the yield.

Fix Dictionary to do just that and add tests to cover this functionality (as well as some other aspects of _modify).

(cherry picked from commit ce96f1e)

# Conflicts:
#	validation-test/stdlib/Dictionary.swift

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -794,45 +794,8 @@ 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)
-        }
-      }
-      _fixLifetime(self)
+      defer { _fixLifetime(self) }
+      yield &_variant[key]
     }
   }
 }
@@ -935,8 +898,8 @@ extension Dictionary {
         native._insert(at: bucket, key: key, value: value)
       }
       let address = native._values + bucket.offset
+      defer { _fixLifetime(self) }
       yield &address.pointee
-      _fixLifetime(self)
     }
   }
 
@@ -1309,14 +1272,10 @@ extension Dictionary {
       return Values(_dictionary: self)
     }
     _modify {
-      var values: Values
-      do {
-        var temp = _Variant(dummy: ())
-        swap(&temp, &_variant)
-        values = Values(_variant: temp)
-      }
+      var values = Values(_variant: _Variant(dummy: ()))
+      swap(&values._variant, &_variant)
+      defer { self._variant = values._variant }
       yield &values
-      self._variant = values._variant
     }
   }
 
@@ -1499,8 +1458,8 @@ extension Dictionary {
         let native = _variant.ensureUniqueNative()
         let bucket = native.validatedBucket(for: position)
         let address = native._values + bucket.offset
+        defer { _fixLifetime(self) }
         yield &address.pointee
-        _fixLifetime(self)
       }
     }
 
@@ -1883,8 +1842,8 @@ extension Dictionary.Index {
       }
       let dummy = _HashTable.Index(bucket: _HashTable.Bucket(offset: 0), age: 0)
       _variant = .native(dummy)
+      defer { _variant = .cocoa(cocoa) }
       yield &cocoa
-      _variant = .cocoa(cocoa)
     }
   }
 #endif

stdlib/public/core/DictionaryVariant.swift

@@ -95,8 +95,8 @@ extension Dictionary._Variant {
     _modify {
       var native = _NativeDictionary<Key, Value>(object.unflaggedNativeInstance)
       self = .init(dummy: ())
+      defer { object = .init(native: native._storage) }
       yield &native
-      object = .init(native: native._storage)
     }
   }
 
@@ -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,52 @@ 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()
+        defer {
+          // This is in a defer block because yield might throw, and we need to
+          // preserve Dictionary's storage invariants when that happens.
+          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)
+          }
+        }
+        yield &value
+      } else {
+        var value: Value? = nil
+        defer {
+          // This is in a defer block because yield might throw, and we need to
+          // preserve Dictionary invariants when that happens.
+          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)
+          }
+        }
+        yield &value
+      }
+    }
+  }
+}
+
 // This function has a highly visible name to make it stand out in stack traces.
 @usableFromInline
 @inline(never)

stdlib/public/core/Set.swift

@@ -1367,8 +1367,8 @@ extension Set.Index {
       }
       let dummy = _HashTable.Index(bucket: _HashTable.Bucket(offset: 0), age: 0)
       _variant = .native(dummy)
+      defer { _variant = .cocoa(cocoa) }
       yield &cocoa
-      _variant = .cocoa(cocoa)
     }
   }
 #endif

stdlib/public/core/SetVariant.swift

@@ -92,8 +92,12 @@ extension Set._Variant {
     _modify {
       var native = _NativeSet<Element>(object.unflaggedNativeInstance)
       self = .init(dummy: ())
+      defer {
+        // This is in a defer block because yield might throw, and we need to
+        // preserve Set's storage invariants when that happens.
+        object = .init(native: native._storage)
+      }
       yield &native
-      object = .init(native: native._storage)
     }
   }