DI: New way of modeling factory initializers

Previously, we were only able to detect factory initializers
dispatched through class_method. This didn't work for
factory initializers defined in protocol extensions.

The end result would be that we would strong_release an
uninitialized class instance, which could cause crashes.

Fix DI to correctly release the old instance using
dealloc_partial_ref instead.

Fixes <rdar://problem/27713221>.

Author: slavapestov

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -1107,15 +1107,30 @@ static bool isSelfInitUse(SILInstruction *I) {
  }
 
   // This is a self.init call if structured like this:
+  //
   // (call_expr type='SomeClass'
   //   (dot_syntax_call_expr type='() -> SomeClass' self
   //     (other_constructor_ref_expr implicit decl=SomeClass.init)
   //     (decl_ref_expr type='SomeClass', "self"))
   //   (...some argument...)
-  if (auto AE = dyn_cast<ApplyExpr>(LocExpr)) {
-    if ((AE = dyn_cast<ApplyExpr>(AE->getFn())) &&
-        isa<OtherConstructorDeclRefExpr>(AE->getFn()))
-      return true;
+  //
+  // Or like this:
+  //
+  // (call_expr type='SomeClass'
+  //   (dot_syntax_call_expr type='() -> SomeClass' self
+  //     (decr_ref_expr implicit decl=SomeClass.init)
+  //     (decl_ref_expr type='SomeClass', "self"))
+  //   (...some argument...)
+  //
+  if (auto *AE = dyn_cast<ApplyExpr>(LocExpr)) {
+    if ((AE = dyn_cast<ApplyExpr>(AE->getFn()))) {
+      if (isa<OtherConstructorDeclRefExpr>(AE->getFn()))
+        return true;
+      if (auto *DRE = dyn_cast<DeclRefExpr>(AE->getFn()))
+        if (auto *CD = dyn_cast<ConstructorDecl>(DRE->getDecl()))
+          if (CD->isFactoryInit())
+            return true;
+    }
   }
   return false;
 }
@@ -1128,59 +1143,19 @@ static bool isSelfInitUse(SILArgument *Arg) {
   // of a throwing delegated init.
   auto *BB = Arg->getParent();
   auto *Pred = BB->getSinglePredecessor();
-  if (!Pred || !isa<TryApplyInst>(Pred->getTerminator()))
-    return false;
-  return isSelfInitUse(Pred->getTerminator());
-}
-
-
-/// Determine if this value_metatype instruction is part of a call to
-/// self.init when delegating to a factory initializer.
-///
-/// FIXME: This is only necessary due to our broken model for factory
-/// initializers.
-static bool isSelfInitUse(ValueMetatypeInst *Inst) {
-  // "Inst" is a ValueMetatype instruction.  Check to see if it is
-  // used by an apply that came from a call to self.init.
-  for (auto UI : Inst->getUses()) {
-    auto *User = UI->getUser();
-
-    // Check whether we're looking up a factory initializer with
-    // class_method.
-    if (auto *CMI = dyn_cast<ClassMethodInst>(User)) {
-      // Only works for allocating initializers...
-      auto Member = CMI->getMember();
-      if (Member.kind != SILDeclRef::Kind::Allocator)
-        return false;
-
-      // ... of factory initializers.
-      auto ctor = dyn_cast_or_null<ConstructorDecl>(Member.getDecl());
-      return ctor && ctor->isFactoryInit();
-    }
-
-    if (auto apply = ApplySite::isa(User)) {
-      auto *LocExpr = apply.getLoc().getAsASTNode<ApplyExpr>();
-      if (!LocExpr)
-        return false;
-
-      LocExpr = dyn_cast<ApplyExpr>(LocExpr->getFn());
-      if (!LocExpr || !isa<OtherConstructorDeclRefExpr>(LocExpr->getFn()))
-        return false;
-
-      return true;
-    }
-
-    // Ignore the thick_to_objc_metatype instruction.
-    if (isa<ThickToObjCMetatypeInst>(User)) {
-      continue;
-    }
 
+  // The two interesting cases are where self.init throws, in which case
+  // the argument came from a try_apply, or if self.init is failable,
+  // in which case we have a switch_enum.
+  if (!Pred ||
+      (!isa<TryApplyInst>(Pred->getTerminator()) &&
+       !isa<SwitchEnumInst>(Pred->getTerminator())))
     return false;
-  }
 
-  return false;
+  return isSelfInitUse(Pred->getTerminator());
 }
 
+
 void ElementUseCollector::
 collectClassSelfUses(SILValue ClassPointer, SILType MemorySILType,
                      llvm::SmallDenseMap<VarDecl*, unsigned> &EltNumbering) {
@@ -1233,18 +1208,12 @@ collectClassSelfUses(SILValue ClassPointer, SILType MemorySILType,
       recordFailableInitCall(User);
     }
 
-    // If this is a ValueMetatypeInst, check to see if it's part of a
-    // self.init call to a factory initializer in a delegating
-    // initializer.
-    if (auto *VMI = dyn_cast<ValueMetatypeInst>(User)) {
-      if (isSelfInitUse(VMI))
-        Kind = DIUseKind::SelfInit;
-      else
-        // Otherwise, this is a simple reference to "type(of:)", which is
-        // always fine, even if self is uninitialized.
-        continue;
-    }
-    
+    // If this is a ValueMetatypeInst, this is a simple reference
+    // to "type(of:)", which is always fine, even if self is
+    // uninitialized.
+    if (isa<ValueMetatypeInst>(User))
+      continue;
+
     // If this is a partial application of self, then this is an escape point
     // for it.
     if (isa<PartialApplyInst>(User))
@@ -1289,27 +1258,26 @@ void ElementUseCollector::collectDelegatingClassInitSelfUses() {
       if (auto apply = dyn_cast<ApplyInst>(cast<AssignInst>(User)->getSrc())) {
         if (auto fn = apply->getCalleeFunction()) {
           if (fn->getRepresentation()
-                == SILFunctionTypeRepresentation::CFunctionPointer)
+                == SILFunctionTypeRepresentation::CFunctionPointer) {
             Uses.push_back(DIMemoryUse(User, DIUseKind::SelfInit, 0, 1));
+            continue;
+          }
         }
       }
 
-      continue;
     }
 
     // Stores *to* the allocation are writes.  If the value being stored is a
     // call to self.init()... then we have a self.init call.
     if (auto *AI = dyn_cast<AssignInst>(User)) {
       if (auto *AssignSource = dyn_cast<SILInstruction>(AI->getOperand(0)))
         if (isSelfInitUse(AssignSource)) {
-          assert(isa<ArchetypeType>(TheMemory.getType()));
           Uses.push_back(DIMemoryUse(User, DIUseKind::SelfInit, 0, 1));
           continue;
         }
       if (auto *AssignSource = dyn_cast<SILArgument>(AI->getOperand(0)))
         if (AssignSource->getParent() == AI->getParent())
           if (isSelfInitUse(AssignSource)) {
-            assert(isa<ArchetypeType>(TheMemory.getType()));
             Uses.push_back(DIMemoryUse(User, DIUseKind::SelfInit, 0, 1));
             continue;
           }
@@ -1379,16 +1347,10 @@ void ElementUseCollector::collectDelegatingClassInitSelfUses() {
           recordFailableInitCall(User);
         }
 
-        // If this is a ValueMetatypeInst, check to see if it's part of a
-        // self.init call to a factory initializer in a delegating
-        // initializer.
-        if (auto *VMI = dyn_cast<ValueMetatypeInst>(User)) {
-          if (isSelfInitUse(VMI))
-            Kind = DIUseKind::SelfInit;
-          else
-            // Otherwise, this is a simple reference to "type(of:)", which is
-            // always fine, even if self is uninitialized.
-            continue;
+        // A simple reference to "type(of:)" is always fine,
+        // even if self is uninitialized.
+        if (isa<ValueMetatypeInst>(User)) {
+          continue;
         }
 
         Uses.push_back(DIMemoryUse(User, Kind, 0, 1));

lib/SILOptimizer/Mandatory/DefiniteInitialization.cpp

@@ -35,39 +35,70 @@ static void diagnose(SILModule &M, SILLocation loc, ArgTypes... args) {
   M.getASTContext().Diags.diagnose(loc.getSourceLoc(), Diagnostic(args...));
 }
 
+enum class PartialInitializationKind {
+  /// The box contains a fully-initialized value.
+  IsNotInitialization,
+
+  /// The box contains a class instance that we own, but the instance has
+  /// not been initialized and should be freed with a special SIL
+  /// instruction made for this purpose.
+  IsReinitialization,
+
+  /// The box contains an undefined value that should be ignored.
+  IsInitialization,
+};
+
 /// Emit the sequence that an assign instruction lowers to once we know
 /// if it is an initialization or an assignment.  If it is an assignment,
 /// a live-in value can be provided to optimize out the reload.
 static void LowerAssignInstruction(SILBuilder &B, AssignInst *Inst,
-                                   IsInitialization_t isInitialization) {
-  DEBUG(llvm::dbgs() << "  *** Lowering [isInit=" << (bool)isInitialization
+                                   PartialInitializationKind isInitialization) {
+  DEBUG(llvm::dbgs() << "  *** Lowering [isInit=" << unsigned(isInitialization)
                      << "]: " << *Inst << "\n");
 
   ++NumAssignRewritten;
 
   SILValue Src = Inst->getSrc();
+  SILLocation Loc = Inst->getLoc();
 
-  // If this is an initialization, or the storage type is trivial, we
-  // can just replace the assignment with a store.
-
-  // Otherwise, if it has trivial type, we can always just replace the
-  // assignment with a store.  If it has non-trivial type and is an
-  // initialization, we can also replace it with a store.
-  if (isInitialization == IsInitialization ||
+  if (isInitialization == PartialInitializationKind::IsInitialization ||
       Inst->getDest()->getType().isTrivial(Inst->getModule())) {
-    B.createStore(Inst->getLoc(), Src, Inst->getDest());
+
+    // If this is an initialization, or the storage type is trivial, we
+    // can just replace the assignment with a store.
+    assert(isInitialization != PartialInitializationKind::IsReinitialization);
+    B.createStore(Loc, Src, Inst->getDest());
+  } else if (isInitialization == PartialInitializationKind::IsReinitialization) {
+
+    // We have a case where a convenience initializer on a class
+    // delegates to a factory initializer from a protocol extension.
+    // Factory initializers give us a whole new instance, so the existing
+    // instance, which has not been initialized and never will be, must be
+    // freed using dealloc_partial_ref.
+    SILValue Pointer = B.createLoad(Loc, Inst->getDest());
+    B.createStore(Loc, Src, Inst->getDest());
+
+    auto MetatypeTy = CanMetatypeType::get(
+        Inst->getDest()->getType().getSwiftRValueType(),
+        MetatypeRepresentation::Thick);
+    auto SILMetatypeTy = SILType::getPrimitiveObjectType(MetatypeTy);
+    SILValue Metatype = B.createValueMetatype(Loc, SILMetatypeTy, Pointer);
+
+    B.createDeallocPartialRef(Loc, Pointer, Metatype);
   } else {
+    assert(isInitialization == PartialInitializationKind::IsNotInitialization);
+
     // Otherwise, we need to replace the assignment with the full
-    // load/store/release dance.  Note that the new value is already
+    // load/store/release dance. Note that the new value is already
     // considered to be retained (by the semantics of the storage type),
     // and we're transferring that ownership count into the destination.
 
     // This is basically TypeLowering::emitStoreOfCopy, except that if we have
     // a known incoming value, we can avoid the load.
-    SILValue IncomingVal = B.createLoad(Inst->getLoc(), Inst->getDest());
+    SILValue IncomingVal = B.createLoad(Loc, Inst->getDest());
     B.createStore(Inst->getLoc(), Src, Inst->getDest());
 
-    B.emitReleaseValueOperation(Inst->getLoc(), IncomingVal);
+    B.emitReleaseValueOperation(Loc, IncomingVal);
   }
 
   Inst->eraseFromParent();
@@ -1632,12 +1663,24 @@ void LifetimeChecker::updateInstructionForInitState(DIMemoryUse &InstInfo) {
     InstInfo.Inst = nullptr;
     NonLoadUses.erase(Inst);
 
+    PartialInitializationKind PartialInitKind;
+
+    if (TheMemory.isClassInitSelf() &&
+        Kind == DIUseKind::SelfInit) {
+      assert(InitKind == IsInitialization);
+      PartialInitKind = PartialInitializationKind::IsReinitialization;
+    } else {
+      PartialInitKind = (InitKind == IsInitialization
+                         ? PartialInitializationKind::IsInitialization
+                         : PartialInitializationKind::IsNotInitialization);
+    }
+
     unsigned FirstElement = InstInfo.FirstElement;
     unsigned NumElements = InstInfo.NumElements;
 
     SmallVector<SILInstruction*, 4> InsertedInsts;
     SILBuilderWithScope B(Inst, &InsertedInsts);
-    LowerAssignInstruction(B, AI, InitKind);
+    LowerAssignInstruction(B, AI, PartialInitKind);
 
     // If lowering of the assign introduced any new loads or stores, keep track
     // of them.
@@ -1646,7 +1689,11 @@ void LifetimeChecker::updateInstructionForInitState(DIMemoryUse &InstInfo) {
         NonLoadUses[I] = Uses.size();
         Uses.push_back(DIMemoryUse(I, Kind, FirstElement, NumElements));
       } else if (isa<LoadInst>(I)) {
-        Uses.push_back(DIMemoryUse(I, Load, FirstElement, NumElements));
+        // If we have a re-initialization, the value must be a class,
+        // and the load is just there so we can free the uninitialized
+        // object husk; it's not an actual use of 'self'.
+        if (PartialInitKind != PartialInitializationKind::IsReinitialization)
+          Uses.push_back(DIMemoryUse(I, Load, FirstElement, NumElements));
       }
     }
     return;
@@ -2541,7 +2588,8 @@ static bool lowerRawSILOperations(SILFunction &Fn) {
       // Unprocessed assigns just lower into assignments, not initializations.
       if (auto *AI = dyn_cast<AssignInst>(Inst)) {
         SILBuilderWithScope B(AI);
-        LowerAssignInstruction(B, AI, IsNotInitialization);
+        LowerAssignInstruction(B, AI,
+            PartialInitializationKind::IsNotInitialization);
         // Assign lowering may split the block. If it did,
         // reset our iteration range to the block after the insertion.
         if (B.getInsertionBB() != &BB)