SILOptimizer: fix a crash in ArrayElementValuePropagation.

This optimization is doing 2 things: replacing getElement calls and replacing append(contentOf:) calls.
Now if the argument to a append(contentOf:) is a previously replaced getElement call, we ended up in a use-after-free crash.

The main change here is to do the transformations immediately after gathering the data (and that means: separately) and not collecting all the data and do both transformation afterwards
The pass does not use any Analysis (where invalidation could be a problem). Also iterator invalidation is not a problem here.

SR-10003
rdar://problem/48445856

Author: eeckstein

lib/SILOptimizer/Transforms/ArrayElementValuePropagation.cpp

@@ -76,20 +76,6 @@ class ArrayAllocation {
 
 public:
 
-  /// Specifies the value with which a get-element call can be replaced.
-  struct GetElementReplacement {
-    ApplyInst *GetElementCall;
-    SILValue Replacement;
-  };
-
-  /// Specifies the set of elements with which an append-contentof call can be
-  /// replaced.
-  struct AppendContentOfReplacement {
-    ApplyInst *AppendContentOfCall;
-    llvm::SmallVector<SILValue, 4> ReplacementValues;
-    SILValue Array;
-  };
-
   ArrayAllocation() {}
 
   /// Analyzes an array allocation call.
@@ -99,14 +85,11 @@ class ArrayAllocation {
   /// or append(contentof) calls.
   bool analyze(ApplyInst *Alloc);
 
-  /// Gets the list of get_element calls which can be replaced.
-  void getGetElementReplacements(
-    llvm::SmallVectorImpl<GetElementReplacement> &Replacements);
+  /// Replace getElement calls with the actual values.
+  bool replaceGetElements();
 
-  /// Gets the list of append(contentof) calls which can be replaced by a
-  /// set of values.
-  void getAppendContentOfReplacements(
-    llvm::SmallVectorImpl<AppendContentOfReplacement> &Replacements);
+  /// Replace append(contentsOf:) with multiple append(element:)
+  bool replaceAppendContentOf();
 };
 
 /// Map the indices of array element initialization stores to their values.
@@ -222,6 +205,10 @@ bool ArrayAllocation::recursivelyCollectUses(ValueBase *Def) {
 }
 
 bool ArrayAllocation::analyze(ApplyInst *Alloc) {
+  GetElementCalls.clear();
+  AppendContentsOfCalls.clear();
+  ElementValueMap.clear();
+
   ArraySemanticsCall Uninitialized(Alloc, "array.uninitialized");
   if (!Uninitialized)
     return false;
@@ -245,8 +232,16 @@ bool ArrayAllocation::analyze(ApplyInst *Alloc) {
   return true;
 }
 
-void ArrayAllocation::getGetElementReplacements(
-                   llvm::SmallVectorImpl<GetElementReplacement> &Replacements) {
+/// Replace getElement calls with the actual values.
+///
+/// \code
+///    store %x to %element_address
+///    ...
+///    %e = apply %getElement(%array, %constant_index)
+/// \endcode
+/// The value %e is replaced with %x.
+bool ArrayAllocation::replaceGetElements() {
+  bool Changed = false;
   for (auto *GetElementCall : GetElementCalls) {
     ArraySemanticsCall GetElement(GetElementCall);
     assert(GetElement.getKind() == ArrayCallKind::kGetElement);
@@ -261,85 +256,107 @@ void ArrayAllocation::getGetElementReplacements(
     if (EltValueIt == ElementValueMap.end())
       continue;
 
-    Replacements.push_back({GetElementCall, EltValueIt->second});
+    Changed |= GetElement.replaceByValue(EltValueIt->second);
   }
+  return Changed;
 }
 
-void ArrayAllocation::getAppendContentOfReplacements(
-              llvm::SmallVectorImpl<AppendContentOfReplacement> &Replacements) {
+/// Replace append(contentsOf:) with multiple append(element:)
+///
+/// \code
+///    store %x to %source_array_element_address_0
+///    store %y to %source_array_element_address_1
+///    ...
+///    apply %append_contentsOf(%dest_array, %source_array)
+/// \endcode
+/// is replaced by
+/// \code
+///    store %x to %source_array_element_address_0
+///    store %y to %source_array_element_address_1
+///    ...
+///    apply %reserveCapacityForAppend(%dest_array, %number_of_values)
+///    apply %append_element(%dest_array, %x)
+///    apply %append_element(%dest_array, %y)
+///    ...
+/// \endcode
+/// The source_array and its initialization code can then be deleted (if not
+/// used otherwise).
+bool ArrayAllocation::replaceAppendContentOf() {
   if (AppendContentsOfCalls.empty())
-    return;
+    return false;
+  if (ElementValueMap.empty())
+    return false;
+
+  // Check if there is a store to each element.
   if (!replacementsAreValid())
-    return;
+    return false;
 
   llvm::SmallVector<SILValue, 4> ElementValueVector;
-  for (unsigned i = 0; i < ElementValueMap.size(); ++i)
-    ElementValueVector.push_back(ElementValueMap[i]);
+  for (unsigned i = 0; i < ElementValueMap.size(); ++i) {
+    SILValue V = ElementValueMap[i];
+    ElementValueVector.push_back(V);
+  }
 
-  for (auto *Call : AppendContentsOfCalls)
-    Replacements.push_back({Call, ElementValueVector, ArrayValue});
-}
+  SILFunction *Fn = AppendContentsOfCalls[0]->getFunction();
+  SILModule &M = Fn->getModule();
+  ASTContext &Ctx = M.getASTContext();
 
-// =============================================================================
-//                                 Driver
-// =============================================================================
+  LLVM_DEBUG(llvm::dbgs() << "Array append contentsOf calls replaced in "
+             << Fn->getName() << "\n");
 
-class ArrayElementPropagation : public SILFunctionTransform {
-public:
-  ArrayElementPropagation() {}
+  // Get the needed Array helper functions.
+  FuncDecl *AppendFnDecl = Ctx.getArrayAppendElementDecl();
+  if (!AppendFnDecl)
+    return false;
 
-  bool replaceAppendCalls(
-                  ArrayRef<ArrayAllocation::AppendContentOfReplacement> Repls) {
-    auto &Fn = *getFunction();
-    auto &M = Fn.getModule();
-    auto &Ctx = M.getASTContext();
+  FuncDecl *ReserveFnDecl = Ctx.getArrayReserveCapacityDecl();
+  if (!ReserveFnDecl)
+    return false;
 
-    if (Repls.empty())
-      return false;
+  auto Mangled = SILDeclRef(AppendFnDecl, SILDeclRef::Kind::Func).mangle();
+  SILFunction *AppendFn = M.findFunction(Mangled, SILLinkage::PublicExternal);
+  if (!AppendFn)
+    return false;
 
-    LLVM_DEBUG(llvm::dbgs() << "Array append contentsOf calls replaced in "
-                            << Fn.getName() << " (" << Repls.size() << ")\n");
+  Mangled = SILDeclRef(ReserveFnDecl, SILDeclRef::Kind::Func).mangle();
+  SILFunction *ReserveFn = M.findFunction(Mangled, SILLinkage::PublicExternal);
+  if (!ReserveFn)
+    return false;
 
-    FuncDecl *AppendFnDecl = Ctx.getArrayAppendElementDecl();
-    if (!AppendFnDecl)
-      return false;
+  bool Changed = false;
+  // Usually there is only a single append(contentsOf:) call. But there might
+  // be multiple - with the same source array to append.
+  for (ApplyInst *AppendContentOfCall : AppendContentsOfCalls) {
+    ArraySemanticsCall AppendContentsOf(AppendContentOfCall);
+    assert(AppendContentsOf && "Must be AppendContentsOf call");
 
-    FuncDecl *ReserveFnDecl = Ctx.getArrayReserveCapacityDecl();
-    if (!ReserveFnDecl)
-      return false;
+    NominalTypeDecl *AppendSelfArray = AppendContentsOf.getSelf()->getType().
+    getASTType()->getAnyNominal();
 
-    auto Mangled = SILDeclRef(AppendFnDecl, SILDeclRef::Kind::Func).mangle();
-    SILFunction *AppendFn = M.findFunction(Mangled, SILLinkage::PublicExternal);
-    if (!AppendFn)
-      return false;
-    
-    Mangled = SILDeclRef(ReserveFnDecl, SILDeclRef::Kind::Func).mangle();
-    SILFunction *ReserveFn = M.findFunction(Mangled, SILLinkage::PublicExternal);
-    if (!ReserveFn)
-      return false;
+    // In case if it's not an Array, but e.g. an ContiguousArray
+    if (AppendSelfArray != Ctx.getArrayDecl())
+      continue;
 
-    for (const ArrayAllocation::AppendContentOfReplacement &Repl : Repls) {
-      ArraySemanticsCall AppendContentsOf(Repl.AppendContentOfCall);
-      assert(AppendContentsOf && "Must be AppendContentsOf call");
+    SILType ArrayType = ArrayValue->getType();
+    auto *NTD = ArrayType.getASTType()->getAnyNominal();
+    SubstitutionMap ArraySubMap = ArrayType.getASTType()
+    ->getContextSubstitutionMap(M.getSwiftModule(), NTD);
 
-      NominalTypeDecl *AppendSelfArray = AppendContentsOf.getSelf()->getType().
-        getASTType()->getAnyNominal();
+    AppendContentsOf.replaceByAppendingValues(M, AppendFn, ReserveFn,
+                                              ElementValueVector,
+                                              ArraySubMap);
+    Changed = true;
+  }
+  return Changed;
+}
 
-      // In case if it's not an Array, but e.g. an ContiguousArray
-      if (AppendSelfArray != Ctx.getArrayDecl())
-        continue;
+// =============================================================================
+//                                 Driver
+// =============================================================================
 
-      SILType ArrayType = Repl.Array->getType();
-      auto *NTD = ArrayType.getASTType()->getAnyNominal();
-      SubstitutionMap ArraySubMap = ArrayType.getASTType()
-        ->getContextSubstitutionMap(M.getSwiftModule(), NTD);
-      
-      AppendContentsOf.replaceByAppendingValues(M, AppendFn, ReserveFn,
-                                                Repl.ReplacementValues,
-                                                ArraySubMap);
-    }
-    return true;
-  }
+class ArrayElementPropagation : public SILFunctionTransform {
+public:
+  ArrayElementPropagation() {}
 
   void run() override {
     auto &Fn = *getFunction();
@@ -348,39 +365,28 @@ class ArrayElementPropagation : public SILFunctionTransform {
     if (Fn.hasOwnership())
       return;
 
-    // Propagate the elements an of array value to its users.
-    llvm::SmallVector<ArrayAllocation::GetElementReplacement, 16>
-      GetElementReplacements;
-    llvm::SmallVector<ArrayAllocation::AppendContentOfReplacement, 4>
-      AppendContentsOfReplacements;
+    bool Changed = false;
 
     for (auto &BB :Fn) {
       for (auto &Inst : BB) {
         if (auto *Apply = dyn_cast<ApplyInst>(&Inst)) {
           ArrayAllocation ALit;
-          if (ALit.analyze(Apply)) {
-            ALit.getGetElementReplacements(GetElementReplacements);
-            ALit.getAppendContentOfReplacements(AppendContentsOfReplacements);
+          if (!ALit.analyze(Apply))
+            continue;
+
+          // First optimization: replace getElemente calls.
+          if (ALit.replaceGetElements()) {
+            Changed = true;
+            // Re-do the analysis if the SIL changed.
+            if (!ALit.analyze(Apply))
+              continue;
           }
+          // Second optimization: replace append(contentsOf:) calls.
+          Changed |= ALit.replaceAppendContentOf();
         }
       }
     }
 
-    LLVM_DEBUG(if (!GetElementReplacements.empty()) {
-      llvm::dbgs() << "Array elements replaced in " << Fn.getName() << " ("
-                   << GetElementReplacements.size() << ")\n";
-    });
-    
-    bool Changed = false;
-    
-    // Perform the actual replacement of the get_element call by its value.
-    for (ArrayAllocation::GetElementReplacement &Repl : GetElementReplacements) {
-      ArraySemanticsCall GetElement(Repl.GetElementCall);
-      Changed |= GetElement.replaceByValue(Repl.Replacement);
-    }
-
-    Changed |= replaceAppendCalls(AppendContentsOfReplacements);
-
     if (Changed) {
       PM->invalidateAnalysis(
           &Fn, SILAnalysis::InvalidationKind::CallsAndInstructions);

test/SILOptimizer/array_element_propagation_crash.swift

@@ -0,0 +1,20 @@
+// RUN: %empty-directory(%t) 
+// RUN: %target-build-swift -O %s -o %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: executable_test
+
+// Check that ArrayElementPropagation does not crash when compiling this function.
+@inline(never)
+func testit(_ arr: inout [Int]) {
+  let a = [28]
+  arr += [a[0]]
+}
+
+var a = [27]
+testit(&a)
+
+// As a bonus, also check if the code works as expected.
+// CHECK: [27, 28]
+print(a)
+