Merge pull request #61963 from nate-chandler/opaque-values/1/20221107

[AddressLowering] Handle try_apply with enum'd return.

Author: nate-chandler

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -3172,6 +3172,12 @@ class DefRewriter : SILInstructionVisitor<DefRewriter> {
 protected:
   // Set the storage address for an opaque block arg and mark it rewritten.
   void rewriteArg(SILPhiArgument *arg) {
+    if (auto *tai =
+            dyn_cast_or_null<TryApplyInst>(arg->getTerminatorForResult())) {
+      CallArgRewriter(tai, pass).rewriteArguments();
+      ApplyRewriter(tai, pass).convertApplyWithIndirectResults();
+      return;
+    }
     LLVM_DEBUG(llvm::dbgs() << "REWRITE ARG "; arg->dump());
     if (storage.storageAddress)
       LLVM_DEBUG(llvm::dbgs() << "  STORAGE "; storage.storageAddress->dump());
@@ -3371,6 +3377,9 @@ static void rewriteFunction(AddressLoweringState &pass) {
       DefRewriter::rewriteValue(valueDef, pass);
       valueAndStorage.storage.markRewritten();
     }
+    // The def of interest may have been changed by rewriteValue.  Get that
+    // redefinition back out of the ValueStoragePair.
+    valueDef = valueAndStorage.value;
     // Rewrite a use of any non-address value mapped to storage (does not
     // include the already rewritten uses of indirect arguments).
     if (valueDef->getType().isAddress())

test/SILOptimizer/address_lowering.sil

@@ -1247,6 +1247,47 @@ bb2(%9 : $Error):
   throw %9 : $Error
 }
 
+// CHECK-LABEL: sil [ossa] @f261_testTryApplyIntoEnum : $@convention(thin) <T> () -> (@out Optional<T>, @error any Error) {
+// CHECK:       {{bb[0-9]+}}([[OUT_ADDR:%[^,]+]] : $*Optional<T>):
+// CHECK:         [[SOME_ADDR:%[^,]+]] = init_enum_data_addr %0 : $*Optional<T>, #Optional.some!enumelt
+// CHECK:         try_apply undef<T>([[SOME_ADDR]]) {{.*}}, normal [[NORMAL:bb[0-9]+]], error {{bb[0-9]+}}
+// CHECK:       [[NORMAL]]
+// CHECK:         inject_enum_addr [[OUT_ADDR]] : $*Optional<T>, #Optional.some!enumelt
+// CHECK-LABEL: } // end sil function 'f261_testTryApplyIntoEnum'
+sil [ossa] @f261_testTryApplyIntoEnum : $@convention(thin) <T> () -> (@out Optional<T>, @error any Error) {
+entry:
+  try_apply undef<T>() : $@convention(thin) <T> () -> (@out T, @error any Error), normal okay, error nopers
+
+okay(%instance : @owned $T):
+  %some = enum $Optional<T>, #Optional.some!enumelt, %instance : $T
+  return %some : $Optional<T>
+
+nopers(%error : @owned $any Error):
+  throw %error : $any Error
+}
+
+// CHECK-LABEL: sil [ossa] @f262_testTryApplyIntoPhi : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[ADDR:%[^,]+]] :
+// CHECK:         try_apply undef<R>([[ADDR]]) {{.*}}, normal [[NORMAL:bb[0-9]+]]
+// CHECK:       [[NORMAL]](%2 : $()):                                    
+// CHECK:         br [[EXIT:bb[0-9]+]]                                          
+// CHECK:       [[EXIT]]:                                              
+// CHECK:         return {{%[^,]+}} : $()                                 
+// CHECK-LABEL: } // end sil function 'f262_testTryApplyIntoPhi'
+sil [ossa] @f262_testTryApplyIntoPhi : $@convention(thin) <R> () -> (@out R, @error any Error) {
+entry:
+  try_apply undef<R>() : $@convention(thin) <U> () -> (@out U, @error any Error), normal good, error bad
+
+good(%instance : @owned $R):
+  br exit(%instance : $R)
+
+exit(%instance_2 : @owned $R):
+  return %instance_2 : $R
+
+bad(%41 : @owned $any Error):
+  throw %41 : $any Error
+}
+
 // CHECK-LABEL: sil [ossa] @fixeeLifetime : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[ADDR:%[^,]+]] : $*T):
 // CHECK:         fix_lifetime [[ADDR]]