Fix Escape Analysis for OSSA's pattern of array.uninitialized (#39789)

Author: meg-gupta

include/swift/SILOptimizer/Analysis/EscapeAnalysis.h

@@ -1059,8 +1059,8 @@ class EscapeAnalysis : public BottomUpIPAnalysis {
   // returns an instantiated array struct and unsafe pointer to the elements.
   struct ArrayUninitCall {
     SILValue arrayStorageRef;
-    TupleExtractInst *arrayStruct = nullptr;
-    TupleExtractInst *arrayElementPtr = nullptr;
+    SILValue arrayStruct = nullptr;
+    SILValue arrayElementPtr = nullptr;
 
     bool isValid() const {
       return arrayStorageRef && arrayStruct && arrayElementPtr;
@@ -1072,9 +1072,9 @@ class EscapeAnalysis : public BottomUpIPAnalysis {
   ArrayUninitCall
   canOptimizeArrayUninitializedCall(ApplyInst *ai);
 
-  /// Return true of this tuple_extract is the result of an optimizable
+  /// Return true of this tuple_extract/destructure_tuple is the result of an optimizable
   /// @_semantics("array.uninitialized") call.
-  bool canOptimizeArrayUninitializedResult(TupleExtractInst *tei);
+  bool canOptimizeArrayUninitializedResult(SILInstruction *tei);
 
   /// Handle a call to "@_semantics(array.uninitialized") precisely by mapping
   /// each call result to a separate graph node and relating them to the

lib/SILOptimizer/Analysis/EscapeAnalysis.cpp

@@ -250,6 +250,14 @@ SILValue EscapeAnalysis::getPointerBase(SILValue value) {
     }
     return pointerOperand;
   }
+  case ValueKind::MultipleValueInstructionResult: {
+    if (auto *dt = dyn_cast<DestructureTupleInst>(value)) {
+      if (canOptimizeArrayUninitializedResult(dt))
+        return SILValue();
+      return dt->getOperand();
+    }
+    return SILValue();
+  }
   default:
     return SILValue();
   }
@@ -1939,6 +1947,11 @@ EscapeAnalysis::canOptimizeArrayUninitializedCall(ApplyInst *ai) {
         continue;
       }
     }
+    if (auto *dt = dyn_cast<DestructureTupleInst>(use->getUser())) {
+      call.arrayStruct = dt->getResult(0);
+      call.arrayElementPtr = dt->getResult(1);
+      continue;
+    }
     // If there are any other uses, such as a release_value, erase the previous
     // call info and bail out.
     call.arrayStruct = nullptr;
@@ -1957,8 +1970,9 @@ EscapeAnalysis::canOptimizeArrayUninitializedCall(ApplyInst *ai) {
 }
 
 bool EscapeAnalysis::canOptimizeArrayUninitializedResult(
-    TupleExtractInst *tei) {
-  ApplyInst *ai = dyn_cast<ApplyInst>(tei->getOperand());
+    SILInstruction *extract) {
+  assert(isa<TupleExtractInst>(extract) || isa<DestructureTupleInst>(extract));
+  ApplyInst *ai = dyn_cast<ApplyInst>(extract->getOperand(0));
   if (!ai)
     return false;
 
@@ -2362,11 +2376,17 @@ void EscapeAnalysis::analyzeInstruction(SILInstruction *I,
       // This is a tuple_extract which extracts the second result of an
       // array.uninitialized call (otherwise getPointerBase should have already
       // looked through it).
-      auto *TEI = cast<TupleExtractInst>(I);
-      assert(canOptimizeArrayUninitializedResult(TEI)
+      assert(canOptimizeArrayUninitializedResult(I)
              && "tuple_extract should be handled as projection");
       return;
     }
+    case SILInstructionKind::DestructureTupleInst: {
+      if (canOptimizeArrayUninitializedResult(I)) {
+        return;
+      }
+      setAllEscaping(I, ConGraph);
+      return;
+    }
     case SILInstructionKind::UncheckedRefCastAddrInst: {
       auto *URCAI = cast<UncheckedRefCastAddrInst>(I);
       CGNode *SrcNode = ConGraph->getNode(URCAI->getSrc());

test/SILOptimizer/escape_analysis.sil

@@ -1331,15 +1331,15 @@ bb(%0 : $*Array<X>, %1 : $@callee_owned (@inout X) -> (@out (), @error Error)):
   return %r : $()
 }
 
-// CHECK-LABEL: CG of arraysemantics_createUninitialized
+// CHECK-LABEL: CG of arraysemantics_createUninitialized1
 // CHECK-NEXT:    Arg [ref] %0 Esc: A, Succ: 
 // CHECK-NEXT:    Val [ref] %2 Esc: , Succ: (%6)
 // CHECK-NEXT:    Val [ref] %5 Esc: , Succ: %2
 // CHECK-NEXT:    Con [int] %6 Esc: R, Succ: (%6.1)
 // CHECK-NEXT:    Con %6.1 Esc: R, Succ: %0
 // CHECK-NEXT:    Ret [ref] return Esc: , Succ: %5
 // CHECK-NEXT:  End
-sil @arraysemantics_createUninitialized : $@convention(thin) (@owned X) -> @owned Array<X> {
+sil @arraysemantics_createUninitialized1 : $@convention(thin) (@owned X) -> @owned Array<X> {
 bb0(%0 : $X):
   %1 = function_ref @swift_bufferAllocate : $@convention(thin) () -> @owned AnyObject
   %2 = apply %1() : $@convention(thin) () -> @owned AnyObject
@@ -1353,6 +1353,27 @@ bb0(%0 : $X):
   return %5 : $Array<X>
 }
 
+// CHECK-LABEL: CG of arraysemantics_createUninitialized2
+// CHECK-NEXT:   Arg [ref] %0 Esc: A, Succ: 
+// CHECK-NEXT:   Val [ref] %2 Esc: , Succ: (%6)
+// CHECK-NEXT:   Val [ref] %5 Esc: , Succ: %2
+// CHECK-NEXT:   Con [int] %6 Esc: R, Succ: (%6.1)
+// CHECK-NEXT:   Con %6.1 Esc: R, Succ: %0
+// CHECK-NEXT:   Ret [ref] return Esc: , Succ: %5
+// CHECK-NEXT: End
+sil [ossa] @arraysemantics_createUninitialized2 : $@convention(thin) (@owned X) -> @owned Array<X> {
+bb0(%0 : @owned $X):
+  %1 = function_ref @swift_bufferAllocate : $@convention(thin) () -> @owned AnyObject
+  %2 = apply %1() : $@convention(thin) () -> @owned AnyObject
+  %3 = function_ref @createUninitialized : $@convention(method) (@owned AnyObject) -> (@owned Array<X>, UnsafeMutablePointer<X>)
+  %4 = apply %3(%2) : $@convention(method) (@owned AnyObject) -> (@owned Array<X>, UnsafeMutablePointer<X>)
+  (%5, %6) = destructure_tuple %4 : $(Array<X>, UnsafeMutablePointer<X>)
+  %7 = struct_extract %6 : $UnsafeMutablePointer<X>, #UnsafeMutablePointer._rawValue
+  %8 = pointer_to_address %7 : $Builtin.RawPointer to $*X
+  store %0 to [init] %8 : $*X
+  return %5 : $Array<X>
+}
+
 sil [_semantics "array.withUnsafeMutableBufferPointer"] @withUnsafeMutableBufferPointer : $@convention(method) (@owned @callee_owned (@inout X) -> (@out (), @error Error), @inout Array<X>) -> (@out (), @error Error)
 sil [_semantics "array.props.isNativeTypeChecked"] @is_native_type_checked : $@convention(method) (@guaranteed Array<X>) -> Bool
 sil [_semantics "array.check_subscript"] @check_subscript : $@convention(method) (Int32, Bool, @guaranteed Array<X>) -> ()