Merge pull request swiftlang#21045 from lorentey/modify-vs-throw

[stdlib] Dictionary: Support throwing yields in _modify accessors

Author: lorentey

stdlib/public/core/Dictionary.swift

@@ -794,8 +794,8 @@ extension Dictionary {
       }
     }
     _modify {
+      defer { _fixLifetime(self) }
       yield &_variant[key]
-      _fixLifetime(self)
     }
   }
 }
@@ -898,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)
     }
   }
 
@@ -1272,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
     }
   }
 
@@ -1462,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)
       }
     }
 
@@ -1846,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)
     }
   }
 

stdlib/public/core/NativeDictionary.swift

@@ -402,24 +402,32 @@ extension _NativeDictionary {
         // 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)
+        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
-        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)
+        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
       }
     }
   }

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)
     }
   }