[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).

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