Merge pull request swiftlang#33923 from hborla/assign-by-wrapper-lowering

[Property Wrappers] Refine assign_by_wrapper lowering

Author: hborla

include/swift/SIL/SILInstruction.h

@@ -3950,6 +3950,16 @@ class AssignByWrapperInst
     : public AssignInstBase<SILInstructionKind::AssignByWrapperInst, 4> {
   friend SILBuilder;
 
+public:
+  /// The assignment destination for the property wrapper
+  enum class Destination {
+    BackingWrapper,
+    WrappedValue,
+  };
+
+private:
+  Destination AssignDest = Destination::WrappedValue;
+
   AssignByWrapperInst(SILDebugLocation DebugLoc, SILValue Src, SILValue Dest,
                        SILValue Initializer, SILValue Setter,
                        AssignOwnershipQualifier Qualifier =
@@ -3964,8 +3974,17 @@ class AssignByWrapperInst
     return AssignOwnershipQualifier(
       SILInstruction::Bits.AssignByWrapperInst.OwnershipQualifier);
   }
-  void setOwnershipQualifier(AssignOwnershipQualifier qualifier) {
+
+  Destination getAssignDestination() const { return AssignDest; }
+
+  void setAssignInfo(AssignOwnershipQualifier qualifier, Destination dest) {
+    using Qualifier = AssignOwnershipQualifier;
+    assert(qualifier == Qualifier::Init && dest == Destination::BackingWrapper ||
+           qualifier == Qualifier::Reassign && dest == Destination::BackingWrapper ||
+           qualifier == Qualifier::Reassign && dest == Destination::WrappedValue);
+
     SILInstruction::Bits.AssignByWrapperInst.OwnershipQualifier = unsigned(qualifier);
+    AssignDest = dest;
   }
 };
 

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h

@@ -254,6 +254,10 @@ enum DIUseKind {
   /// value.
   Assign,
 
+  /// The instruction is an assignment of a wrapped value with an already initialized
+  /// backing property wrapper.
+  AssignWrappedValue,
+
   /// The instruction is a store to a member of a larger struct value.
   PartialStore,
 

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -533,6 +533,7 @@ LifetimeChecker::LifetimeChecker(const DIMemoryObjectInfo &TheMemory,
     case DIUseKind::Escape:
       continue;
     case DIUseKind::Assign:
+    case DIUseKind::AssignWrappedValue:
     case DIUseKind::IndirectIn:
     case DIUseKind::InitOrAssign:
     case DIUseKind::InOutArgument:
@@ -750,6 +751,7 @@ void LifetimeChecker::doIt() {
       continue;
 
     case DIUseKind::Assign:
+    case DIUseKind::AssignWrappedValue:
       // Instructions classified as assign are only generated when lowering
       // InitOrAssign instructions in regions known to be initialized.  Since
       // they are already known to be definitely init, don't reprocess them.
@@ -1048,16 +1050,15 @@ void LifetimeChecker::handleStoreUse(unsigned UseID) {
   // an initialization or assign in the uses list so that clients know about it.
   if (isFullyUninitialized) {
     Use.Kind = DIUseKind::Initialization;
+  } else if (isFullyInitialized && isa<AssignByWrapperInst>(Use.Inst)) {
+    // If some fields are uninitialized, re-write assign_by_wrapper to assignment
+    // of the backing wrapper. If all fields are initialized, assign to the wrapped
+    // value.
+    auto allFieldsInitialized =
+        getAnyUninitializedMemberAtInst(Use.Inst, 0, TheMemory.getNumElements()) == -1;
+    Use.Kind = allFieldsInitialized ? DIUseKind::AssignWrappedValue : DIUseKind::Assign;
   } else if (isFullyInitialized) {
-    // Only re-write assign_by_wrapper to assignment if all fields have been
-    // initialized.
-    if (isa<AssignByWrapperInst>(Use.Inst) &&
-        getAnyUninitializedMemberAtInst(Use.Inst, 0,
-                                        TheMemory.getNumElements()) != -1) {
-      Use.Kind = DIUseKind::Initialization;
-    } else {
-      Use.Kind = DIUseKind::Assign;
-    }
+    Use.Kind = DIUseKind::Assign;
   } else {
     // If it is initialized on some paths, but not others, then we have an
     // inconsistent initialization, which needs dynamic control logic in the
@@ -1918,7 +1919,7 @@ void LifetimeChecker::updateInstructionForInitState(DIMemoryUse &Use) {
       Use.Kind == DIUseKind::SelfInit)
     InitKind = IsInitialization;
   else {
-    assert(Use.Kind == DIUseKind::Assign);
+    assert(Use.Kind == DIUseKind::Assign || Use.Kind == DIUseKind::AssignWrappedValue);
     InitKind = IsNotInitialization;
   }
 
@@ -1967,14 +1968,21 @@ void LifetimeChecker::updateInstructionForInitState(DIMemoryUse &Use) {
     Use.Inst = nullptr;
     NonLoadUses.erase(Inst);
 
-    if (TheMemory.isClassInitSelf() &&
-        Use.Kind == DIUseKind::SelfInit) {
-      assert(InitKind == IsInitialization);
-      AI->setOwnershipQualifier(AssignOwnershipQualifier::Reinit);
-    } else {
-      AI->setOwnershipQualifier((InitKind == IsInitialization
-                                 ? AssignOwnershipQualifier::Init
-                                 : AssignOwnershipQualifier::Reassign));
+    switch (Use.Kind) {
+    case DIUseKind::Initialization:
+      AI->setAssignInfo(AssignOwnershipQualifier::Init,
+                        AssignByWrapperInst::Destination::BackingWrapper);
+      break;
+    case DIUseKind::Assign:
+      AI->setAssignInfo(AssignOwnershipQualifier::Reassign,
+                        AssignByWrapperInst::Destination::BackingWrapper);
+      break;
+    case DIUseKind::AssignWrappedValue:
+      AI->setAssignInfo(AssignOwnershipQualifier::Reassign,
+                        AssignByWrapperInst::Destination::WrappedValue);
+      break;
+    default:
+      llvm_unreachable("Wrong use kind for assign_by_wrapper");
     }
 
     return;

lib/SILOptimizer/Mandatory/RawSILInstLowering.cpp

@@ -179,8 +179,8 @@ static void lowerAssignByWrapperInstruction(SILBuilderWithScope &b,
   SILLocation loc = inst->getLoc();
   SILBuilderWithScope forCleanup(std::next(inst->getIterator()));
 
-  switch (inst->getOwnershipQualifier()) {
-    case AssignOwnershipQualifier::Init: {
+  switch (inst->getAssignDestination()) {
+    case AssignByWrapperInst::Destination::BackingWrapper: {
       SILValue initFn = inst->getInitializer();
       CanSILFunctionType fTy = initFn->getType().castTo<SILFunctionType>();
       SILFunctionConventions convention(fTy, inst->getModule());
@@ -193,15 +193,18 @@ static void lowerAssignByWrapperInstruction(SILBuilderWithScope &b,
         getAssignByWrapperArgs(args, src, convention, b, forCleanup);
         SILValue wrappedSrc = b.createApply(loc, initFn, SubstitutionMap(),
                                             args);
-        b.createTrivialStoreOr(loc, wrappedSrc, dest,
-                               StoreOwnershipQualifier::Init);
+        if (inst->getOwnershipQualifier() == AssignOwnershipQualifier::Init ||
+            inst->getDest()->getType().isTrivial(*inst->getFunction())) {
+          b.createTrivialStoreOr(loc, wrappedSrc, dest, StoreOwnershipQualifier::Init);
+        } else {
+          b.createStore(loc, wrappedSrc, dest, StoreOwnershipQualifier::Assign);
+        }
       }
       // TODO: remove the unused setter function, which usually is a dead
       // partial_apply.
       break;
     }
-    case AssignOwnershipQualifier::Unknown:
-    case AssignOwnershipQualifier::Reassign: {
+    case AssignByWrapperInst::Destination::WrappedValue: {
       SILValue setterFn = inst->getSetter();
       CanSILFunctionType fTy = setterFn->getType().castTo<SILFunctionType>();
       SILFunctionConventions convention(fTy, inst->getModule());
@@ -220,8 +223,6 @@ static void lowerAssignByWrapperInstruction(SILBuilderWithScope &b,
       // partial_apply.
       break;
     }
-    case AssignOwnershipQualifier::Reinit:
-      llvm_unreachable("wrong qualifier for assign_by_wrapper");
   }
   inst->eraseFromParent();
 }

test/SILOptimizer/di_property_wrappers_leak.swift

@@ -0,0 +1,40 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %s -o %t/a.out
+// RUN: %target-codesign %t/a.out
+// RUN: %target-run %t/a.out
+
+// REQUIRES: executable_test
+
+import StdlibUnittest
+
+@propertyWrapper
+struct Wrapper<T> {
+  var wrappedValue: T
+}
+
+struct TestWrappedValueLeak {
+  @Wrapper var wrapped: LifetimeTracked = LifetimeTracked(0)
+  var str: String
+
+  init() {
+    wrapped = LifetimeTracked(42)
+    str = ""
+    wrapped = LifetimeTracked(27)
+  }
+
+  init(conditionalInit: Bool) {
+    if (conditionalInit) {
+      wrapped = LifetimeTracked(42)
+    }
+    str = ""
+    wrapped = LifetimeTracked(27)
+  }
+}
+
+TestSuite("Property Wrapper DI").test("test wrapped value leak") {
+  _ = TestWrappedValueLeak()
+  _ = TestWrappedValueLeak(conditionalInit: true)
+  _ = TestWrappedValueLeak(conditionalInit: false)
+}
+
+runAllTests()