[AddressLowering] Rewrite checked_casts as seen.

Instead of waiting until after rewriting everything else, rewrite them
as the terminator results they produce are encountered.  This enables
forming projections in the correct locations.

Author: nate-chandler

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -419,11 +419,6 @@ struct AddressLoweringState {
   // parameters are rewritten.
   SmallBlotSetVector<FullApplySite, 16> indirectApplies;
 
-  // checked_cast_br instructions with loadable source type and opaque target
-  // type need to be rewritten in a post-pass, once all the uses of the opaque
-  // target value are rewritten to their address forms.
-  SmallVector<CheckedCastBranchInst *, 8> opaqueResultCCBs;
-
   // All function-exiting terminators (return or throw instructions).
   SmallVector<TermInst *, 8> exitingInsts;
 
@@ -611,15 +606,6 @@ void OpaqueValueVisitor::mapValueStorage() {
       if (auto apply = FullApplySite::isa(&inst))
         checkForIndirectApply(apply);
 
-      // Collect all checked_cast_br instructions that have a loadable source
-      // type and opaque target type
-      if (auto *ccb = dyn_cast<CheckedCastBranchInst>(&inst)) {
-        if (!ccb->getSourceLoweredType().isAddressOnly(*ccb->getFunction()) &&
-            ccb->getTargetLoweredType().isAddressOnly(*ccb->getFunction())) {
-          pass.opaqueResultCCBs.push_back(ccb);
-        }
-      }
-
       for (auto result : inst.getResults()) {
         if (isPseudoCallResult(result) || isPseudoReturnValue(result))
           continue;
@@ -3177,6 +3163,10 @@ class DefRewriter : SILInstructionVisitor<DefRewriter> {
       CallArgRewriter(tai, pass).rewriteArguments();
       ApplyRewriter(tai, pass).convertApplyWithIndirectResults();
       return;
+    } else if (auto *ccbi = dyn_cast_or_null<CheckedCastBranchInst>(
+                   arg->getTerminatorForResult())) {
+      CheckedCastBrRewriter(ccbi, pass).rewrite();
+      return;
     }
     LLVM_DEBUG(llvm::dbgs() << "REWRITE ARG "; arg->dump());
     if (storage.storageAddress)
@@ -3412,12 +3402,6 @@ static void rewriteFunction(AddressLoweringState &pass) {
     }
   }
 
-  // Rewrite all checked_cast_br instructions with loadable source type and
-  // opaque target type now
-  for (auto *ccb : pass.opaqueResultCCBs) {
-    CheckedCastBrRewriter(ccb, pass).rewrite();
-  }
-
   // Rewrite this function's return value now that all opaque values within the
   // function are rewritten. This still depends on a valid ValueStorage
   // projection operands.

test/SILOptimizer/address_lowering.sil

@@ -1613,6 +1613,36 @@ bb3:
   return %31 : $()
 }
 
+// Test the result being stored into an @out enum.
+// CHECK-LABEL: sil [ossa] @test_checked_cast_br4 : $@convention(method) <Element><T> (@owned TestGeneric<Element>) -> @out Optional<T> {
+// CHECK:       {{bb[0-9]+}}([[OUT_ADDR:%[^,]+]] : $*Optional<T>, [[INSTANCE:%[^,]+]] : @owned $TestGeneric<Element>):
+// CHECK:         [[TEMP:%[^,]+]] = alloc_stack $Optional<T>
+// CHECK:         [[CAST_DEST_ADDR:%[^,]+]] = init_enum_data_addr [[TEMP]] : $*Optional<T>, #Optional.some!enumelt
+// CHECK:         [[CAST_SOURCE_ADDR:%[^,]+]] = alloc_stack $TestGeneric<Element>
+// CHECK:         store [[INSTANCE]] to [init] [[CAST_SOURCE_ADDR]]
+// CHECK:         checked_cast_addr_br take_on_success TestGeneric<Element> in [[CAST_SOURCE_ADDR]] : $*TestGeneric<Element> to T in [[CAST_DEST_ADDR]] : $*T, [[SUCCESS:bb[0-9]+]]
+// CHECK:       [[SUCCESS]]:
+// CHECK:         dealloc_stack [[CAST_SOURCE_ADDR]]
+// CHECK:         inject_enum_addr [[TEMP]] : $*Optional<T>, #Optional.some!enumelt
+// CHECK:         copy_addr [take] [[TEMP]] to [init] [[OUT_ADDR]]
+// CHECK-LABEL: } // end sil function 'test_checked_cast_br4'
+sil [ossa] @test_checked_cast_br4 : $@convention(method) <Element><T> (@owned TestGeneric<Element>) -> @out Optional<T> {
+bb0(%3 : @owned $TestGeneric<Element>):
+  checked_cast_br %3 : $TestGeneric<Element> to T, bb1, bb2
+
+bb1(%5 : @owned $T):
+  %6 = enum $Optional<T>, #Optional.some!enumelt, %5 : $T
+  br bb3(%6 : $Optional<T>)
+
+bb2(%8 : @owned $TestGeneric<Element>):
+  destroy_value %8 : $TestGeneric<Element>
+  %10 = enum $Optional<T>, #Optional.none!enumelt
+  br bb3(%10 : $Optional<T>)
+
+bb3(%12 : @owned $Optional<T>):
+  return %12 : $Optional<T>
+}
+
 // CHECK-LABEL: sil hidden [ossa] @test_unchecked_bitwise_cast :
 // CHECK: bb0(%0 : $*U, %1 : $*T, %2 : $@thick U.Type):
 // CHECK:   [[STK:%.*]] = alloc_stack $T