[AddressLowering] Addr'd some cast_brs.

All checked casts are emitted as checked_cast_br
instructions.  More than just the instructions which produce or consume
an opaque value must be rewritten as checked_cast_addr_br
instructions.  In particular, all those instructions for which
canIRGenUseScalarCheckedCastInstructions returns false must be rewritten
as checked_cast_addr_br instructions.

Note the instructions to rewrite like that while visiting values and
then rewrite them near the end of rewriting.

Author: nate-chandler

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -139,6 +139,7 @@
 #include "swift/Basic/Range.h"
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/DebugUtils.h"
+#include "swift/SIL/DynamicCasts.h"
 #include "swift/SIL/OwnershipUtils.h"
 #include "swift/SIL/PrettyStackTrace.h"
 #include "swift/SIL/PrunedLiveness.h"
@@ -423,6 +424,9 @@ struct AddressLoweringState {
   // rewritten.
   SmallVector<UnconditionalCheckedCastInst *, 2> nonopaqueResultUCCs;
 
+  // checked_cast_br instructions to loadable type which need to be rewritten.
+  SmallVector<CheckedCastBranchInst *, 2> nonopaqueResultCCBs;
+
   // All function-exiting terminators (return or throw instructions).
   SmallVector<TermInst *, 8> exitingInsts;
 
@@ -654,6 +658,19 @@ void OpaqueValueVisitor::visitValue(SILValue value) {
               ucci->getTargetFormalType())) {
         pass.nonopaqueResultUCCs.push_back(ucci);
       }
+    } else if (auto *arg = dyn_cast<SILArgument>(value)) {
+      if (auto *ccbi = dyn_cast_or_null<CheckedCastBranchInst>(
+              arg->getTerminatorForResult())) {
+        if (ccbi->getSuccessBB() != arg->getParent())
+          return;
+        if (ccbi->getSourceLoweredType().isAddressOnly(*pass.function))
+          return;
+        if (!canIRGenUseScalarCheckedCastInstructions(
+                pass.function->getModule(), ccbi->getSourceFormalType(),
+                ccbi->getTargetFormalType())) {
+          pass.nonopaqueResultCCBs.push_back(ccbi);
+        }
+      }
     }
     return;
   }
@@ -3453,6 +3470,10 @@ static void rewriteFunction(AddressLoweringState &pass) {
     rewriteNonopaqueUnconditionalCheckedCast(ucci, pass);
   }
 
+  for (auto *ccbi : pass.nonopaqueResultCCBs) {
+    CheckedCastBrRewriter(ccbi, 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

@@ -1643,6 +1643,60 @@ bb3(%12 : @owned $Optional<T>):
   return %12 : $Optional<T>
 }
 
+// Verify that checked_cast_br neither from nor to an address-only value but
+// whose operands require rewriting to checked_cast_addr_br gets rewritten.
+// CHECK-LABEL: sil [ossa] @test_checked_cast_br5 : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[SRC:%[^,]+]] :
+// CHECK:         [[SRC_ADDR:%[^,]+]] = alloc_stack $AnyObject
+// CHECK:         store [[SRC]] to [init] [[SRC_ADDR]]
+// CHECK:         [[DEST_ADDR:%[^,]+]] = alloc_stack $@thick any AnyObject.Type
+// CHECK:         checked_cast_addr_br take_on_success AnyObject in [[SRC_ADDR]] {{.*}} to any AnyObject.Type in [[DEST_ADDR]] {{.*}}, [[SUCCESS:bb[0-9]+]], [[FAILURE:bb[0-9]+]]
+// CHECK:       [[FAILURE]]:
+// CHECK:         load [take] [[SRC_ADDR]] : $*AnyObject
+// CHECK:       [[SUCCESS]]:
+// CHECK:         load [trivial] [[DEST_ADDR]] : $*@thick any AnyObject.Type
+// CHECK-LABEL: } // end sil function 'test_checked_cast_br5'
+sil [ossa] @test_checked_cast_br5 : $@convention(thin) (@owned AnyObject) -> () {
+entry(%instance : @owned $AnyObject):
+  checked_cast_br %instance : $AnyObject to any AnyObject.Type, success, failure
+
+success(%metatype : $@thick any AnyObject.Type):
+  br exit
+
+failure(%original : @owned $AnyObject):
+  destroy_value %original : $AnyObject
+  br exit
+
+exit:
+  %retval = tuple ()
+  return %retval : $()
+}
+
+// Verify that checked_cast_br which both has an address-only value and also
+// produces a value whose type would anyway require rewriting but is not
+// address-only gets lowered properly.
+// CHECK-LABEL: sil [ossa] @test_checked_cast_br6 : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[SRC:%[^,]+]] :
+// CHECK:         [[DEST_ADDR:%[^,]+]] = alloc_stack $@thick any AnyObject.Type     
+// CHECK:         checked_cast_addr_br take_on_success T in [[SRC]] {{.*}} to any AnyObject.Type in [[DEST_ADDR]] {{.*}}, [[SUCCESS]], [[FAILURE]] 
+// CHECK:       [[FAILURE]]:                                              
+// CHECK:         destroy_addr [[SRC]] : $*T                           
+// CHECK:       [[SUCCESS]]:                                              
+// CHECK:         [[REGISTER_3:%[^,]+]] = load [trivial] [[DEST_ADDR]] : $*@thick any AnyObject.Type
+// CHECK-LABEL: } // end sil function 'test_checked_cast_br6'
+sil [ossa] @test_checked_cast_br6 : $@convention(thin) <T> (@in T) -> () {
+entry(%instance : @owned $T):
+  checked_cast_br %instance : $T to any AnyObject.Type, success, failure
+success(%metatype : $@thick any AnyObject.Type):
+  br exit
+failure(%original : @owned $T):
+  destroy_value %original : $T
+  br exit
+exit:
+  %retval = tuple ()
+  return %retval : $()
+}
+
 // CHECK-LABEL: sil hidden [ossa] @test_unchecked_bitwise_cast :
 // CHECK: bb0(%0 : $*U, %1 : $*T, %2 : $@thick U.Type):
 // CHECK:   [[STK:%.*]] = alloc_stack $T

test/SILOptimizer/opaque_values_Onone.swift

@@ -212,3 +212,41 @@ func castAnyObjectToMeta(_ ao: any AnyObject) -> AnyObject.Type {
 func castGenericToMeta<T>(_ t: T) -> AnyObject.Type {
   t as! AnyObject.Type
 }
+
+// CHECK-LABEL: sil hidden @maybeCastAnyObjectToMeta : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[SRC:%[^,]+]] :
+// CHECK:         [[SRC_ADDR:%[^,]+]] = alloc_stack $AnyObject
+// CHECK:         store [[SRC]] to [[SRC_ADDR]] : $*AnyObject
+// CHECK:         [[DEST_ADDR:%[^,]+]] = alloc_stack $@thick any AnyObject.Type
+// CHECK:         checked_cast_addr_br take_on_success AnyObject in [[SRC_ADDR]] {{.*}} to any AnyObject.Type in [[DEST_ADDR]] {{.*}}, [[SUCCESS:bb[0-9]+]], [[FAILURE:bb[0-9]+]]
+// CHECK:       [[SUCCESS]]:
+// CHECK:         load [[DEST_ADDR]] : $*@thick any AnyObject.Type
+// CHECK:       [[FAILURE]]:
+// CHECK:         load [[SRC_ADDR]] : $*AnyObject
+// CHECK-LABEL: } // end sil function 'maybeCastAnyObjectToMeta'
+@_silgen_name("maybeCastAnyObjectToMeta")
+func maybeCastAnyObjectToMeta(_ ao: any AnyObject) -> AnyObject.Type? {
+  if let m = ao as? AnyObject.Type {
+    return m
+  }
+  return nil
+}
+
+// CHECK-LABEL: sil hidden @maybeCastGenericToMeta : {{.*}} {
+// CHECK:       {{bb[0-9]+}}([[SRC:%[^,]+]] :
+// CHECK:         [[SRC_ADDR:%[^,]+]] = alloc_stack $T
+// CHECK:         copy_addr [[SRC]] to [init] [[SRC_ADDR]]
+// CHECK:         [[DEST_ADDR:%[^,]+]] = alloc_stack $@thick any AnyObject.Type
+// CHECK:         checked_cast_addr_br take_on_success T in [[SRC_ADDR]] {{.*}} to any AnyObject.Type in [[DEST_ADDR]] {{.*}}, [[SUCCESS:bb[0-9]+]], [[FAILURE:bb[0-9]+]]
+// CHECK:       [[SUCCESS]]:
+// CHECK:         load [[DEST_ADDR]]
+// CHECK:       [[FAILURE]]:
+// CHECK:         destroy_addr [[SRC_ADDR]]
+// CHECK-LABEL: } // end sil function 'maybeCastGenericToMeta'
+@_silgen_name("maybeCastGenericToMeta")
+func maybeCastGenericToMeta<T>(_ t: T) -> AnyObject.Type? {
+  if let m = t as? AnyObject.Type {
+    return m
+  }
+  return nil
+}