[pred-memopts] Rather than asserting on recursive initialization, just return false and bail.

Until the beginning of the ownership transition, DI and predictable mem opts
used the same memory use collector. I split them partially since I need to turn
on ownership for predictable mem opts at one time, but also b/c there was a huge
amount of special code that would only trigger if it was used by DI or used by
predictable mem opts. After I did the copy some of the asserts that were needed
for DI remained in the predictable mem opts code. When pred-memopts was only run
in the mandatory pipeline keeping these assertions were ok, but pred-memopts was
recently added to the perf pipeline meaning that it may see code that breaks
these DI invariants (and thus hit this assertion).

We should remove this limitation on predictable-memopts but that would require
some scheduled time to read the code (more than I have to fix this bug = p). So
instead I changed the code to just bail in these cases.

rdar://40032102
(cherry picked from commit 269f8e8)

Author: gottesmm

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.cpp

@@ -247,36 +247,46 @@ namespace {
 
     /// This is the main entry point for the use walker.  It collects uses from
     /// the address and the refcount result of the allocation.
-    void collectFrom() {
-      if (auto *ABI = TheMemory.getContainer())
-        collectContainerUses(ABI);
-      else
-        collectUses(TheMemory.getAddress(), 0);
-
-      // Collect information about the retain count result as well.
-      for (auto UI : TheMemory.MemoryInst->getUses()) {
-        auto *User = UI->getUser();
-
-        // If this is a release or dealloc_stack, then remember it as such.
-        if (isa<StrongReleaseInst>(User) || isa<DeallocStackInst>(User) ||
-            isa<DeallocBoxInst>(User)) {
-          Releases.push_back(User);
-        }
-      }
-    }
+    LLVM_NODISCARD bool collectFrom();
 
   private:
-    void collectUses(SILValue Pointer, unsigned BaseEltNo);
-    void collectContainerUses(AllocBoxInst *ABI);
+    LLVM_NODISCARD bool collectUses(SILValue Pointer, unsigned BaseEltNo);
+    LLVM_NODISCARD bool collectContainerUses(AllocBoxInst *ABI);
     void addElementUses(unsigned BaseEltNo, SILType UseTy,
                         SILInstruction *User, DIUseKind Kind);
-    void collectTupleElementUses(TupleElementAddrInst *TEAI,
-                                 unsigned BaseEltNo);
-    void collectStructElementUses(StructElementAddrInst *SEAI,
-                                  unsigned BaseEltNo);
+    LLVM_NODISCARD bool collectTupleElementUses(TupleElementAddrInst *TEAI,
+                                                unsigned BaseEltNo);
+    LLVM_NODISCARD bool collectStructElementUses(StructElementAddrInst *SEAI,
+                                                 unsigned BaseEltNo);
   };
 } // end anonymous namespace
 
+bool ElementUseCollector::collectFrom() {
+  bool shouldOptimize = false;
+
+  if (auto *ABI = TheMemory.getContainer()) {
+    shouldOptimize = collectContainerUses(ABI);
+  } else {
+    shouldOptimize = collectUses(TheMemory.getAddress(), 0);
+  }
+
+  if (!shouldOptimize)
+    return false;
+
+  // Collect information about the retain count result as well.
+  for (auto UI : TheMemory.MemoryInst->getUses()) {
+    auto *User = UI->getUser();
+
+    // If this is a release or dealloc_stack, then remember it as such.
+    if (isa<StrongReleaseInst>(User) || isa<DeallocStackInst>(User) ||
+        isa<DeallocBoxInst>(User)) {
+      Releases.push_back(User);
+    }
+  }
+
+  return true;
+}
+
 /// addElementUses - An operation (e.g. load, store, inout use, etc) on a value
 /// acts on all of the aggregate elements in that value.  For example, a load
 /// of $*(Int,Int) is a use of both Int elements of the tuple.  This is a helper
@@ -295,8 +305,8 @@ void ElementUseCollector::addElementUses(unsigned BaseEltNo, SILType UseTy,
 /// Given a tuple_element_addr or struct_element_addr, compute the new
 /// BaseEltNo implicit in the selected member, and recursively add uses of
 /// the instruction.
-void ElementUseCollector::
-collectTupleElementUses(TupleElementAddrInst *TEAI, unsigned BaseEltNo) {
+bool ElementUseCollector::collectTupleElementUses(TupleElementAddrInst *TEAI,
+                                                  unsigned BaseEltNo) {
 
   // If we're walking into a tuple within a struct or enum, don't adjust the
   // BaseElt.  The uses hanging off the tuple_element_addr are going to be
@@ -314,20 +324,20 @@ collectTupleElementUses(TupleElementAddrInst *TEAI, unsigned BaseEltNo) {
       BaseEltNo += getElementCountRec(Module, EltTy);
     }
   }
-  
-  collectUses(TEAI, BaseEltNo);
+
+  return collectUses(TEAI, BaseEltNo);
 }
 
-void ElementUseCollector::collectStructElementUses(StructElementAddrInst *SEAI,
+bool ElementUseCollector::collectStructElementUses(StructElementAddrInst *SEAI,
                                                    unsigned BaseEltNo) {
   // Generally, we set the "InStructSubElement" flag and recursively process
   // the uses so that we know that we're looking at something within the
   // current element.
   llvm::SaveAndRestore<bool> X(InStructSubElement, true);
-  collectUses(SEAI, BaseEltNo);
+  return collectUses(SEAI, BaseEltNo);
 }
 
-void ElementUseCollector::collectContainerUses(AllocBoxInst *ABI) {
+bool ElementUseCollector::collectContainerUses(AllocBoxInst *ABI) {
   for (Operand *UI : ABI->getUses()) {
     auto *User = UI->getUser();
 
@@ -340,16 +350,21 @@ void ElementUseCollector::collectContainerUses(AllocBoxInst *ABI) {
       continue;
 
     if (auto project = dyn_cast<ProjectBoxInst>(User)) {
-      collectUses(project, project->getFieldIndex());
+      if (!collectUses(project, project->getFieldIndex()))
+        return false;
       continue;
     }
 
     // Other uses of the container are considered escapes of the values.
-    for (unsigned field : indices(ABI->getBoxType()->getLayout()->getFields()))
+    for (unsigned field :
+         indices(ABI->getBoxType()->getLayout()->getFields())) {
       addElementUses(field,
                      ABI->getBoxType()->getFieldType(ABI->getModule(), field),
                      User, DIUseKind::Escape);
+    }
   }
+
+  return true;
 }
 
 // Returns true when the instruction represents added instrumentation for
@@ -367,7 +382,7 @@ static bool isSanitizerInstrumentation(SILInstruction *Instruction,
   return false;
 }
 
-void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
+bool ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
   assert(Pointer->getType().isAddress() &&
          "Walked through the pointer to the value?");
   SILType PointeeType = Pointer->getType().getObjectType();
@@ -383,19 +398,22 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
 
     // struct_element_addr P, #field indexes into the current element.
     if (auto *SEAI = dyn_cast<StructElementAddrInst>(User)) {
-      collectStructElementUses(SEAI, BaseEltNo);
+      if (!collectStructElementUses(SEAI, BaseEltNo))
+        return false;
       continue;
     }
 
     // Instructions that compute a subelement are handled by a helper.
     if (auto *TEAI = dyn_cast<TupleElementAddrInst>(User)) {
-      collectTupleElementUses(TEAI, BaseEltNo);
+      if (!collectTupleElementUses(TEAI, BaseEltNo))
+        return false;
       continue;
     }
 
     // Look through begin_access.
     if (auto I = dyn_cast<BeginAccessInst>(User)) {
-      collectUses(I, BaseEltNo);
+      if (!collectUses(I, BaseEltNo))
+        return false;
       continue;
     }
 
@@ -439,7 +457,7 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
       continue;
     }
 
-    if (auto *SWI = dyn_cast<StoreWeakInst>(User))
+    if (auto *SWI = dyn_cast<StoreWeakInst>(User)) {
       if (UI->getOperandNumber() == 1) {
         DIUseKind Kind;
         if (InStructSubElement)
@@ -451,8 +469,9 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
         Uses.push_back(DIMemoryUse(User, Kind, BaseEltNo, 1));
         continue;
       }
+    }
 
-    if (auto *SUI = dyn_cast<StoreUnownedInst>(User))
+    if (auto *SUI = dyn_cast<StoreUnownedInst>(User)) {
       if (UI->getOperandNumber() == 1) {
         DIUseKind Kind;
         if (InStructSubElement)
@@ -464,6 +483,7 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
         Uses.push_back(DIMemoryUse(User, Kind, BaseEltNo, 1));
         continue;
       }
+    }
 
     if (auto *CAI = dyn_cast<CopyAddrInst>(User)) {
       // If this is a copy of a tuple, we should scalarize it so that we don't
@@ -505,7 +525,13 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
       // If this is an out-parameter, it is like a store.
       unsigned NumIndirectResults = substConv.getNumIndirectSILResults();
       if (ArgumentNumber < NumIndirectResults) {
-        assert(!InStructSubElement && "We're initializing sub-members?");
+        // We do not support initializing sub members. This is an old
+        // restriction from when this code was used by Definite
+        // Initialization. With proper code review, we can remove this, but for
+        // now, lets be conservative.
+        if (InStructSubElement) {
+          return false;
+        }
         addElementUses(BaseEltNo, PointeeType, User,
                        DIUseKind::Initialization);
         continue;
@@ -550,29 +576,38 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
     // be explicitly initialized by a copy_addr or some other use of the
     // projected address).
     if (auto I = dyn_cast<InitEnumDataAddrInst>(User)) {
-      assert(!InStructSubElement &&
-             "init_enum_data_addr shouldn't apply to struct subelements");
+      // If we are in a struct already, bail. With proper analysis, we should be
+      // able to do this optimization.
+      if (InStructSubElement) {
+        return false;
+      }
+
       // Keep track of the fact that we're inside of an enum.  This informs our
       // recursion that tuple stores are not scalarized outside, and that stores
       // should not be treated as partial stores.
       llvm::SaveAndRestore<bool> X(InEnumSubElement, true);
-      collectUses(I, BaseEltNo);
+      if (!collectUses(I, BaseEltNo))
+        return false;
       continue;
     }
 
     // init_existential_addr is modeled as an initialization store.
     if (isa<InitExistentialAddrInst>(User)) {
-      assert(!InStructSubElement &&
-             "init_existential_addr should not apply to struct subelements");
+      // init_existential_addr should not apply to struct subelements.
+      if (InStructSubElement) {
+        return false;
+      }
       Uses.push_back(DIMemoryUse(User, DIUseKind::Initialization,
                                  BaseEltNo, 1));
       continue;
     }
 
     // inject_enum_addr is modeled as an initialization store.
     if (isa<InjectEnumAddrInst>(User)) {
-      assert(!InStructSubElement &&
-             "inject_enum_addr the subelement of a struct unless in a ctor");
+      // inject_enum_addr the subelement of a struct unless in a ctor.
+      if (InStructSubElement) {
+        return false;
+      }
       Uses.push_back(DIMemoryUse(User, DIUseKind::Initialization,
                                  BaseEltNo, 1));
       continue;
@@ -667,16 +702,22 @@ void ElementUseCollector::collectUses(SILValue Pointer, unsigned BaseEltNo) {
     // Now that we've scalarized some stuff, recurse down into the newly created
     // element address computations to recursively process it.  This can cause
     // further scalarization.
-    for (auto EltPtr : ElementAddrs)
-      collectTupleElementUses(cast<TupleElementAddrInst>(EltPtr), BaseEltNo);
+    if (llvm::any_of(ElementAddrs, [&](SILValue V) {
+          return !collectTupleElementUses(cast<TupleElementAddrInst>(V),
+                                          BaseEltNo);
+        })) {
+      return false;
+    }
   }
+
+  return true;
 }
 
 /// collectDIElementUsesFrom - Analyze all uses of the specified allocation
 /// instruction (alloc_box, alloc_stack or mark_uninitialized), classifying them
 /// and storing the information found into the Uses and Releases lists.
-void swift::collectDIElementUsesFrom(
+bool swift::collectDIElementUsesFrom(
     const DIMemoryObjectInfo &MemoryInfo, SmallVectorImpl<DIMemoryUse> &Uses,
     SmallVectorImpl<SILInstruction *> &Releases) {
-  ElementUseCollector(MemoryInfo, Uses, Releases).collectFrom();
+  return ElementUseCollector(MemoryInfo, Uses, Releases).collectFrom();
 }

lib/SILOptimizer/Mandatory/DIMemoryUseCollector.h

@@ -21,12 +21,14 @@
 #define SWIFT_SILOPTIMIZER_MANDATORY_DIMEMORYUSECOLLECTOR_H
 
 #include "swift/Basic/LLVM.h"
-#include "llvm/ADT/APInt.h"
 #include "swift/SIL/SILInstruction.h"
 #include "swift/SIL/SILType.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/Support/Compiler.h"
 
 namespace swift {
-  class SILBuilder;
+
+class SILBuilder;
 
 /// DIMemoryObjectInfo - This struct holds information about the memory object
 /// being analyzed that is required to correctly break it down into elements.
@@ -187,9 +189,10 @@ struct DIMemoryUse {
 /// collectDIElementUsesFrom - Analyze all uses of the specified allocation
 /// instruction (alloc_box, alloc_stack or mark_uninitialized), classifying them
 /// and storing the information found into the Uses and Releases lists.
-void collectDIElementUsesFrom(const DIMemoryObjectInfo &MemoryInfo,
-                              SmallVectorImpl<DIMemoryUse> &Uses,
-                              SmallVectorImpl<SILInstruction *> &Releases);
+LLVM_NODISCARD bool
+collectDIElementUsesFrom(const DIMemoryObjectInfo &MemoryInfo,
+                         SmallVectorImpl<DIMemoryUse> &Uses,
+                         SmallVectorImpl<SILInstruction *> &Releases);
 
 } // end namespace swift
 

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -1372,9 +1372,14 @@ static bool optimizeMemoryAllocations(SILFunction &Fn) {
       // Set up the datastructure used to collect the uses of the allocation.
       SmallVector<DIMemoryUse, 16> Uses;
       SmallVector<SILInstruction*, 4> Releases;
-      
-      // Walk the use list of the pointer, collecting them.
-      collectDIElementUsesFrom(MemInfo, Uses, Releases);
+
+      // Walk the use list of the pointer, collecting them. If we are not able
+      // to optimize, skip this value. *NOTE* We may still scalarize values
+      // inside the value.
+      if (!collectDIElementUsesFrom(MemInfo, Uses, Releases)) {
+        ++I;
+        continue;
+      }
 
       Changed |= AllocOptimize(Alloc, Uses, Releases).doIt();
 

test/SILOptimizer/predictable_memopt.sil

@@ -3,7 +3,6 @@
 import Builtin
 import Swift
 
-
 // CHECK-LABEL: sil @simple_reg_promotion
 // CHECK: bb0(%0 : $Int):
 // CHECK-NEXT: return %0 : $Int
@@ -838,3 +837,47 @@ bb2:
   destroy_addr %1 : $*Builtin.NativeObject
   unreachable
 }
+
+
+class K {
+  init()
+}
+
+sil @init_k : $@convention(thin) () -> @out K
+
+struct S {
+  var k: K
+}
+
+// CHECK-LABEL: sil @recursive_struct_destroy_with_apply : $@convention(thin) () -> S {
+// CHECK: alloc_stack
+// CHECK: } // end sil function 'recursive_struct_destroy_with_apply'
+sil @recursive_struct_destroy_with_apply : $@convention(thin) () -> S {
+bb0:
+  %0 = alloc_stack $S
+  %1 = struct_element_addr %0 : $*S, #S.k
+  %2 = function_ref @init_k : $@convention(thin) () -> @out K
+  %3 = apply %2(%1) : $@convention(thin) () -> @out K
+  %4 = load %0 : $*S
+  dealloc_stack %0 : $*S
+  return %4 : $S
+}
+
+struct SWithOpt {
+  var k: Optional<K>
+}
+
+// CHECK-LABEL: sil @recursive_struct_destroy_with_enum_init : $@convention(thin) (@owned K) -> @owned SWithOpt {
+// CHECK: alloc_stack
+// CHECK: } // end sil function 'recursive_struct_destroy_with_enum_init'
+sil @recursive_struct_destroy_with_enum_init : $@convention(thin) (@owned K) -> @owned SWithOpt {
+bb0(%arg : $K):
+  %0 = alloc_stack $SWithOpt
+  %1 = struct_element_addr %0 : $*SWithOpt, #SWithOpt.k
+  %2 = init_enum_data_addr %1 : $*Optional<K>, #Optional.some!enumelt.1
+  store %arg to %2 : $*K
+  inject_enum_addr %1 : $*Optional<K>, #Optional.some!enumelt.1
+  %4 = load %0 : $*SWithOpt
+  dealloc_stack %0 : $*SWithOpt
+  return %4 : $SWithOpt
+}