DI: self-consumed analysis rework, volume 2

lib/SILOptimizer/Mandatory/DIMemoryUseCollectorOwnership.cpp

@@ -426,52 +426,8 @@ bool DIMemoryUse::onlyTouchesTrivialElements(
 //                      DIElementUseInfo Implementation
 //===----------------------------------------------------------------------===//
 
-void DIElementUseInfo::trackFailableInitCall(
-    const DIMemoryObjectInfo &MemoryInfo, SILInstruction *I) {
-  // If we have a store to self inside the normal BB, we have a 'real'
-  // try_apply. Otherwise, this is a 'try? self.init()' or similar,
-  // and there is a store after.
-  if (auto *TAI = dyn_cast<TryApplyInst>(I)) {
-    trackFailureBlock(MemoryInfo, TAI, TAI->getNormalBB());
-    return;
-  }
-
-  if (auto *AI = dyn_cast<ApplyInst>(I)) {
-    // See if this is an optional initializer.
-    for (auto Op : AI->getUses()) {
-      SILInstruction *User = Op->getUser();
-
-      if (!isa<SelectEnumInst>(User) && !isa<SelectEnumAddrInst>(User))
-        continue;
-
-      auto value = cast<SingleValueInstruction>(User);
-
-      if (!value->hasOneUse())
-        continue;
-
-      User = value->use_begin()->getUser();
-      if (auto *CBI = dyn_cast<CondBranchInst>(User)) {
-        trackFailureBlock(MemoryInfo, CBI, CBI->getTrueBB());
-        return;
-      }
-    }
-  }
-}
-
-/// We have to detect if the self box contents were consumed. Do this by
-/// checking for a store into the self box in the success branch.  Once we rip
-/// this out of SILGen, DI will be able to figure this out in a more logical
-/// manner.
-void DIElementUseInfo::trackFailureBlock(const DIMemoryObjectInfo &TheMemory,
-                                         TermInst *TI, SILBasicBlock *BB) {
-  for (auto &II : *BB) {
-    if (auto *SI = dyn_cast<StoreInst>(&II)) {
-      if (SI->getDest() == TheMemory.MemoryInst) {
-        FailableInits.push_back(TI);
-        return;
-      }
-    }
-  }
+void DIElementUseInfo::trackStoreToSelf(SILInstruction *I) {
+  StoresToSelf.push_back(I);
 }
 
 //===----------------------------------------------------------------------===//
@@ -1103,10 +1059,14 @@ void ElementUseCollector::collectClassSelfUses() {
     return;
   }
 
+  // The number of stores of the initial 'self' argument into the self box
+  // that we saw.
+  unsigned StoresOfArgumentToSelf = 0;
+
   // Okay, given that we have a proper setup, we walk the use chains of the self
   // box to find any accesses to it. The possible uses are one of:
   //
-  //   1) The initialization store (TheStore).
+  //   1) The initialization store.
   //   2) Loads of the box, which have uses of self hanging off of them.
   //   3) An assign to the box, which happens at super.init.
   //   4) Potential escapes after super.init, if self is closed over.
@@ -1115,10 +1075,29 @@ void ElementUseCollector::collectClassSelfUses() {
   for (auto *Op : MUI->getUses()) {
     SILInstruction *User = Op->getUser();
 
-    // Stores to self are initializations store or the rebind of self as
-    // part of the super.init call.  Ignore both of these.
-    if (isa<StoreInst>(User) && Op->getOperandNumber() == 1)
-      continue;
+    // Stores to self.
+    if (auto *SI = dyn_cast<StoreInst>(User)) {
+      if (Op->getOperandNumber() == 1) {
+        // The initial store of 'self' into the box at the start of the
+        // function. Ignore it.
+        if (auto *Arg = dyn_cast<SILArgument>(SI->getSrc())) {
+          if (Arg->getParent() == MUI->getParent()) {
+            StoresOfArgumentToSelf++;
+            continue;
+          }
+        }
+
+        // A store of a load from the box is ignored.
+        // FIXME: SILGen should not emit these.
+        if (auto *LI = dyn_cast<LoadInst>(SI->getSrc()))
+          if (LI->getOperand() == MUI)
+            continue;
+
+        // Any other store needs to be recorded.
+        UseInfo.trackStoreToSelf(SI);
+        continue;
+      }
+    }
 
     // Ignore end_borrows. These can only come from us being the source of a
     // load_borrow.
@@ -1149,6 +1128,9 @@ void ElementUseCollector::collectClassSelfUses() {
     // and super.init must be called.
     trackUse(DIMemoryUse(User, DIUseKind::Load, 0, TheMemory.NumElements));
   }
+
+  assert(StoresOfArgumentToSelf == 1 &&
+         "The 'self' argument should have been stored into the box exactly once");
 }
 
 static bool isSuperInitUse(SILInstruction *User) {
@@ -1406,7 +1388,6 @@ void ElementUseCollector::collectClassSelfUses(
         (isSelfInitUse(User) || isSuperInitUse(User))) {
       if (isSelfOperand(Op, User)) {
         Kind = DIUseKind::SelfInit;
-        UseInfo.trackFailableInitCall(TheMemory, User);
       }
     }
 
@@ -1459,6 +1440,10 @@ void DelegatingInitElementUseCollector::collectClassInitSelfUses() {
   assert(TheMemory.NumElements == 1 && "delegating inits only have 1 bit");
   auto *MUI = cast<MarkUninitializedInst>(TheMemory.MemoryInst);
 
+  // The number of stores of the initial 'self' argument into the self box
+  // that we saw.
+  unsigned StoresOfArgumentToSelf = 0;
+
   // We walk the use chains of the self MUI to find any accesses to it.  The
   // possible uses are:
   //   1) The initialization store.
@@ -1475,14 +1460,33 @@ void DelegatingInitElementUseCollector::collectClassInitSelfUses() {
     if (isa<EndBorrowInst>(User))
       continue;
 
-    // Stores to self are initializations store or the rebind of self as
-    // part of the super.init call.  Ignore both of these.
-    if (isa<StoreInst>(User) && Op->getOperandNumber() == 1)
-      continue;
+    // Stores to self.
+    if (auto *SI = dyn_cast<StoreInst>(User)) {
+      if (Op->getOperandNumber() == 1) {
+        // The initial store of 'self' into the box at the start of the
+        // function. Ignore it.
+        if (auto *Arg = dyn_cast<SILArgument>(SI->getSrc())) {
+          if (Arg->getParent() == MUI->getParent()) {
+            StoresOfArgumentToSelf++;
+            continue;
+          }
+        }
+
+        // A store of a load from the box is ignored.
+        // FIXME: SILGen should not emit these.
+        if (auto *LI = dyn_cast<LoadInst>(SI->getSrc()))
+          if (LI->getOperand() == MUI)
+            continue;
+
+        // Any other store needs to be recorded.
+        UseInfo.trackStoreToSelf(SI);
+        continue;
+      }
+    }
 
     // For class initializers, the assign into the self box may be
     // captured as SelfInit or SuperInit elsewhere.
-    if (TheMemory.isClassInitSelf() && isa<AssignInst>(User) &&
+    if (isa<AssignInst>(User) &&
         Op->getOperandNumber() == 1) {
       // If the source of the assignment is an application of a C
       // function, there is no metatype argument, so treat the
@@ -1491,6 +1495,7 @@ void DelegatingInitElementUseCollector::collectClassInitSelfUses() {
         if (auto *Fn = AI->getCalleeFunction()) {
           if (Fn->getRepresentation() ==
               SILFunctionTypeRepresentation::CFunctionPointer) {
+            UseInfo.trackStoreToSelf(User);
             UseInfo.trackUse(DIMemoryUse(User, DIUseKind::SelfInit, 0, 1));
             continue;
           }
@@ -1503,13 +1508,15 @@ void DelegatingInitElementUseCollector::collectClassInitSelfUses() {
     if (auto *AI = dyn_cast<AssignInst>(User)) {
       if (auto *AssignSource = AI->getOperand(0)->getDefiningInstruction()) {
         if (isSelfInitUse(AssignSource) || isSuperInitUse(AssignSource)) {
+          UseInfo.trackStoreToSelf(User);
           UseInfo.trackUse(DIMemoryUse(User, DIUseKind::SelfInit, 0, 1));
           continue;
         }
       }
       if (auto *AssignSource = dyn_cast<SILArgument>(AI->getOperand(0))) {
         if (AssignSource->getParent() == AI->getParent() &&
             (isSelfInitUse(AssignSource) || isSuperInitUse(AssignSource))) {
+          UseInfo.trackStoreToSelf(User);
           UseInfo.trackUse(DIMemoryUse(User, DIUseKind::SelfInit, 0, 1));
           continue;
         }
@@ -1534,6 +1541,9 @@ void DelegatingInitElementUseCollector::collectClassInitSelfUses() {
     UseInfo.trackUse(DIMemoryUse(User, DIUseKind::Escape, 0, 1));
   }
 
+  assert(StoresOfArgumentToSelf == 1 &&
+         "The 'self' argument should have been stored into the box exactly once");
+
   // The MUI must be used on an alloc_box or alloc_stack instruction. If we have
   // an alloc_stack, there is nothing further to do.
   if (isa<AllocStackInst>(MUI->getOperand()))
@@ -1573,13 +1583,17 @@ void DelegatingInitElementUseCollector::collectValueTypeInitSelfUses(
     // 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 (AI->getDest() == I)
+      if (AI->getDest() == I) {
+        UseInfo.trackStoreToSelf(AI);
         Kind = DIUseKind::InitOrAssign;
+      }
     }
 
     if (auto *CAI = dyn_cast<CopyAddrInst>(User)) {
-      if (CAI->getDest() == I)
+      if (CAI->getDest() == I) {
+        UseInfo.trackStoreToSelf(CAI);
         Kind = DIUseKind::InitOrAssign;
+      }
     }
 
     // Look through begin_access
@@ -1665,7 +1679,6 @@ void DelegatingInitElementUseCollector::collectDelegatingClassInitSelfLoadUses(
         (isSelfInitUse(User) || isSuperInitUse(User))) {
       if (isSelfOperand(Op, User)) {
         Kind = DIUseKind::SelfInit;
-        UseInfo.trackFailableInitCall(TheMemory, User);
       }
     }
 
@@ -1701,7 +1714,7 @@ void swift::ownership::collectDIElementUsesFrom(
   // collector.
   if (MemoryInfo.isDelegatingInit()) {
     DelegatingInitElementUseCollector UseCollector(MemoryInfo, UseInfo);
-    if (MemoryInfo.getType()->hasReferenceSemantics()) {
+    if (MemoryInfo.isClassInitSelf()) {
       UseCollector.collectClassInitSelfUses();
     } else {
       UseCollector.collectValueTypeInitSelfUses();

lib/SILOptimizer/Mandatory/DIMemoryUseCollectorOwnership.h

@@ -150,6 +150,15 @@ class DIMemoryObjectInfo {
     return false;
   }
 
+  /// True if this memory object is the 'self' of a non-root class init method.
+  bool isNonRootClassSelf() const {
+    if (isClassInitSelf())
+      if (auto MUI = dyn_cast<MarkUninitializedInst>(MemoryInst))
+        return !MUI->isRootSelf();
+
+    return false;
+  }
+
   /// isDelegatingInit - True if this is a delegating initializer, one that
   /// calls 'self.init'.
   bool isDelegatingInit() const {
@@ -269,18 +278,13 @@ struct DIMemoryUse {
 struct DIElementUseInfo {
   SmallVector<DIMemoryUse, 16> Uses;
   SmallVector<SILInstruction *, 4> Releases;
-  TinyPtrVector<TermInst *> FailableInits;
+  TinyPtrVector<SILInstruction *> StoresToSelf;
 
   void trackUse(DIMemoryUse Use) { Uses.push_back(Use); }
 
   void trackDestroy(SILInstruction *Destroy) { Releases.push_back(Destroy); }
 
-  void trackFailableInitCall(const DIMemoryObjectInfo &TheMemory,
-                             SILInstruction *I);
-
-private:
-  void trackFailureBlock(const DIMemoryObjectInfo &TheMemory, TermInst *TI,
-                         SILBasicBlock *BB);
+  void trackStoreToSelf(SILInstruction *I);
 };
 
 /// collectDIElementUsesFrom - Analyze all uses of the specified allocation