[SILOptimizer] fix KeyPathProjector memory management

I was inconsistently providing initialized or uninitialized memory
to the callback when projecting a settable address, depending on
component type. We should always provide an uninitialized address.

Author: NobodyNada

include/swift/SILOptimizer/Utils/KeyPathProjector.h

@@ -58,6 +58,9 @@ class KeyPathProjector {
   /// \param accessType The access type of the projected address.
   /// \param callback A callback to invoke with the projected adddress.
   ///     The projected address is only valid from within \p callback.
+  ///     If accessType is Get or Modify, the projected addres is an
+  ///     initialized address type. If accessType is set, the projected
+  ///     address points to uninitialized memory.
   virtual void project(AccessType accessType,
                std::function<void(SILValue addr)> callback) = 0;
 

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -219,13 +219,13 @@ bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
     return false;
 
   SILValue keyPath, rootAddr, valueAddr;
-  bool isModify = false;
+  bool isSet = false;
   if (callee->getName() == "swift_setAtWritableKeyPath" ||
       callee->getName() == "swift_setAtReferenceWritableKeyPath") {
     keyPath = AI->getArgument(1);
     rootAddr = AI->getArgument(0);
     valueAddr = AI->getArgument(2);
-    isModify = true;
+    isSet = true;
   } else if (callee->getName() == "swift_getAtKeyPath") {
     keyPath = AI->getArgument(2);
     rootAddr = AI->getArgument(1);
@@ -240,13 +240,13 @@ bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
     return false;
 
   KeyPathProjector::AccessType accessType;
-  if (isModify) accessType = KeyPathProjector::AccessType::Set;
+  if (isSet) accessType = KeyPathProjector::AccessType::Set;
   else accessType = KeyPathProjector::AccessType::Get;
 
   projector->project(accessType, [&](SILValue projectedAddr) {
-    if (isModify) {
+    if (isSet) {
       Builder.createCopyAddr(AI->getLoc(), valueAddr, projectedAddr,
-                             IsTake, IsNotInitialization);
+                             IsTake, IsInitialization);
     } else {
       Builder.createCopyAddr(AI->getLoc(), projectedAddr, valueAddr,
                              IsNotTake, IsInitialization);

lib/SILOptimizer/Utils/KeyPathProjector.cpp

@@ -34,6 +34,11 @@ class RootProjector : public KeyPathProjector {
 
     void project(AccessType accessType,
                  std::function<void (SILValue)> callback) override {
+        if (accessType == AccessType::Set) {
+          // We're setting the identity key path (\.self). The callback
+          // expects an uninitialized address, so destroy the old value.
+          builder.emitDestroyAddr(loc, root);
+        }
         callback(root);
     }
 
@@ -105,7 +110,11 @@ class StoredPropertyProjector : public ComponentProjector {
         parentAccessType = AccessType::Modify;
 
       parent->project(parentAccessType, [&](SILValue parentValue) {
-        callback(builder.createStructElementAddr(loc, parentValue, storedProperty));
+        auto addr = builder.createStructElementAddr(loc, parentValue, storedProperty);
+        // If we're setting, destroy the old value (the callback expects uninitialized memory)
+        if (accessType == AccessType::Set)
+          builder.createDestroyAddr(loc, addr);
+        callback(addr);
       });
     } else {
       // Accessing a class member -> reading the class
@@ -143,6 +152,9 @@ class StoredPropertyProjector : public ComponentProjector {
           addr = beginAccess;
         }
 
+        // If we're setting, destroy the old value (the callback expects uninitialized memory)
+        if (accessType == AccessType::Set)
+          builder.createDestroyAddr(loc, addr);
         callback(addr);
 
         // if a child hasn't started a new access (i.e. beginAccess is unchanged),
@@ -178,8 +190,11 @@ class TupleElementProjector : public ComponentProjector {
       parentAccessType = AccessType::Modify;
 
     parent->project(parentAccessType, [&](SILValue parentValue) {
-      callback(builder.createTupleElementAddr(loc, parentValue,
-                                              component.getTupleIndex()));
+      auto addr = builder.createTupleElementAddr(loc, parentValue, component.getTupleIndex());
+      // If we're setting, destroy the old value (the callback expects uninitialized memory)
+      if (accessType == AccessType::Set)
+        builder.createDestroyAddr(loc, addr);
+      callback(addr);
     });
   }
 };

test/SILOptimizer/optimize_keypath.swift

@@ -130,7 +130,8 @@ func testGenStructRead<T>(_ s: GenStruct<T>) -> T {
 
 // CHECK-LABEL: sil {{.*}}testGenStructWrite
 // CHECK: [[A:%[0-9]+]] = struct_element_addr %0
-// CHECK: copy_addr %1 to [[A]]
+// CHECK: destroy_addr [[A]]
+// CHECK: copy_addr {{.*}} to [initialization] [[A]]
 // CHECK: return
 @inline(never)
 @_semantics("optimize.sil.specialize.generic.never")
@@ -183,7 +184,8 @@ func testDerivedClass2Read(_ c: DerivedClass2) -> Int {
 // CHECK: [[S:%[0-9]+]] = alloc_stack $T
 // CHECK: [[E:%[0-9]+]] = ref_element_addr %0
 // CHECK: [[A:%[0-9]+]] = begin_access [modify] [dynamic] [[E]]
-// CHECK: copy_addr [take] [[S]] to [[A]]
+// CHECK: destroy_addr [[A]]
+// CHECK: copy_addr [take] [[S]] to [initialization] [[A]]
 // CHECK: end_access [[A]]
 // CHECK: return
 @inline(never)