SimplifyLoad/SimplifyLoadBorrow: replace address casts of head objects

Replaces address casts of heap objects
```
  %1 = unchecked_addr_cast %0 : $*SomeClass to $*OtherClass
  %2 = load [copy] %1
```
with ref-casts of the loaded value
```
  %1 = load [copy] %0
  %2 = unchecked_ref_cast %1 : $SomeClass to $OtherClass
```

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyLoad.swift

@@ -26,6 +26,9 @@ extension LoadInst : OnoneSimplifiable, SILCombineSimplifiable {
     if replaceLoadOfGlobalLet(context) {
       return
     }
+    if tryRemoveAddressCast(context) {
+      return
+    }
     removeIfDead(context)
   }
 
@@ -116,6 +119,27 @@ extension LoadInst : OnoneSimplifiable, SILCombineSimplifiable {
     return true
   }
 
+  /// Replaces address casts of heap objects
+  /// ```
+  ///   %1 = unchecked_addr_cast %0 : $*SomeClass to $*OtherClass
+  ///   %2 = load [copy] %1
+  /// ```
+  /// with ref-casts of the loaded value
+  /// ```
+  ///   %1 = load [copy] %0
+  ///   %2 = unchecked_ref_cast %1 : $SomeClass to $OtherClass
+  /// ```
+  private func tryRemoveAddressCast(_ context: SimplifyContext) -> Bool {
+    guard let addrCast = address.isAddressCastOfHeapObjects else {
+      return false
+    }
+    let builder = Builder(before: self, context)
+    let newLoad = builder.createLoad(fromAddress: addrCast.fromAddress, ownership: loadOwnership)
+    let cast = builder.createUncheckedRefCast(from: newLoad, to: addrCast.type.objectType)
+    replace(with: cast, context)
+    return true
+  }
+
   private func isZeroLoadFromEmptyCollection() -> Bool {
     if !type.isBuiltinInteger {
       return false
@@ -364,6 +388,18 @@ private extension Value {
   }
 }
 
+extension Value {
+  var isAddressCastOfHeapObjects: UncheckedAddrCastInst? {
+    if let addrCast = self as? UncheckedAddrCastInst,
+       addrCast.fromAddress.type.isHeapObjectReferenceType,
+       addrCast.type.isHeapObjectReferenceType
+    {
+      return addrCast
+    }
+    return nil
+  }
+}
+
 private extension Instruction {
   var isShiftRightByAtLeastOne: Bool {
     guard let bi = self as? BuiltinInst,

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyLoadBorrow.swift

@@ -32,20 +32,64 @@ extension LoadBorrowInst : Simplifiable, SILCombineSimplifiable {
     //   %3 = some_forwarding_instruction %1 // zero or more forwarding instructions
     // ```
     //
-    tryCombineWithCopy(context)
+    if tryCombineWithCopy(context) {
+      return
+    }
+
+    if tryRemoveAddrCast(context) {
+      return
+    }
+
+    tryForwardStoreBorrow(context)
   }
 
-  private func tryCombineWithCopy(_ context: SimplifyContext) {
+  private func tryCombineWithCopy(_ context: SimplifyContext) -> Bool {
     let forwardedValue = lookThroughOwnedConvertibaleForwardingChain()
     guard let singleUser = forwardedValue.uses.ignore(usersOfType: EndBorrowInst.self).singleUse?.instruction,
           let copy = singleUser as? CopyValueInst,
           copy.parentBlock == self.parentBlock else {
-      return
+      return false
     }
     let builder = Builder(before: self, context)
     let loadCopy = builder.createLoad(fromAddress: address, ownership: .copy)
     let forwardedOwnedValue = replaceGuaranteed(value: self, withOwnedValue: loadCopy, context)
     copy.replace(with: forwardedOwnedValue, context)
     context.erase(instructionIncludingAllUsers: self)
+    return true
+  }
+
+  /// Replaces address casts of heap objects
+  /// ```
+  ///   %1 = unchecked_addr_cast %0 : $*SomeClass to $*OtherClass
+  ///   %2 = load_borrow %1
+  ///   // ... uses of %2
+  ///   end_borrow %2
+  /// ```
+  /// with ref-casts of the loaded value
+  /// ```
+  ///   %1 = load_borrow %0
+  ///   %2 = unchecked_ref_cast %1 : $SomeClass to $OtherClass
+  ///   // ... uses of %2
+  ///   end_borrow %2
+  /// ```
+  private func tryRemoveAddrCast(_ context: SimplifyContext) -> Bool {
+    guard let addrCast = address.isAddressCastOfHeapObjects else {
+      return false
+    }
+    let builder = Builder(before: self, context)
+    let newLoad = builder.createLoadBorrow(fromAddress: addrCast.fromAddress)
+    let cast = builder.createUncheckedRefCast(from: newLoad, to: addrCast.type.objectType)
+    replace(with: newLoad, context)
+    newLoad.uses.filter{ !$0.endsLifetime }.ignore(user: cast).replaceAll(with: cast, context)
+    return true
+  }
+
+  private func tryForwardStoreBorrow(_ context: SimplifyContext) {
+    guard let storeBorrow = address as? StoreBorrowInst else {
+      return
+    }
+    let builder = Builder(before: self, context)
+    let beginBorrow = builder.createBeginBorrow(of: storeBorrow.source)
+    replace(with: beginBorrow, context)
   }
 }

test/SILOptimizer/pre_specialize_layouts.swift

@@ -158,7 +158,8 @@ internal func testEmitIntoClient<T>(t: T) {
 // OPT:   [[R8:%.*]] = unchecked_addr_cast [[R7]] : $*AnyObject to $*SomeClass
 // OPT:   [[F6:%.*]] = function_ref @$s30pre_specialized_module_layouts14InternalThing2V9computedZxvryXl_Ts5 : $@yield_once @convention(method) (@guaranteed InternalThing2<AnyObject>) -> @yields @in_guaranteed AnyObject
 // OPT:   ([[R9:%.*]], {{%.*}}) = begin_apply [[F6]]({{.*}}) : $@yield_once @convention(method) (@guaranteed InternalThing2<AnyObject>) -> @yields @in_guaranteed AnyObject
-// OPT:   [[R10:%.*]] = unchecked_addr_cast [[R9]] : $*AnyObject to $*SomeClass
+// OPT:   [[LD:%.*]] = load [[R9]]
+// OPT:   [[R10:%.*]] = unchecked_ref_cast [[LD]] : $AnyObject to $SomeClass
 // OPT: } // end sil function '$s30pre_specialized_module_layouts16useInternalThingyyxlFAA9SomeClassC_Tg5'
 
 public func usePrespecializedEntryPoints(wrapperStruct: ReferenceWrapperStruct, overaligned: OveralignedReferenceWrapperStruct, array: [Int], stride96: Stride96) {

test/SILOptimizer/simplify_load.sil

@@ -31,6 +31,10 @@ class D : C {
   @_hasStorage var d: Builtin.Int64
 }
 
+struct AO {
+  let a: AnyObject
+}
+
 sil_global hidden [let] @gb : $B
 sil_global hidden [let] @gb2 : $(B, Int64)
 
@@ -543,3 +547,45 @@ bb0:
   return %4
 }
 
+// CHECK-LABEL: sil [ossa] @load_of_addr_cast :
+// CHECK:         %1 = load [copy] %0
+// CHECK:         %2 = unchecked_ref_cast %1 : $B to $AnyObject
+// CHECK:         fix_lifetime %2
+// CHECK:         destroy_value %2
+// CHECK:       } // end sil function 'load_of_addr_cast'
+sil [ossa] @load_of_addr_cast : $@convention(thin) (@inout B) -> () {
+bb0(%0 : $*B):
+  %1 = unchecked_addr_cast %0 to $*AnyObject
+  %2 = load [copy] %1
+  fix_lifetime %2
+  destroy_value %2
+  %6 = tuple ()
+  return %6
+}
+
+// CHECK-LABEL: sil [ossa] @dont_remove_non_class_addr_cast1 :
+// CHECK:         %1 = unchecked_addr_cast %0
+// CHECK:       } // end sil function 'dont_remove_non_class_addr_cast1'
+sil [ossa] @dont_remove_non_class_addr_cast1 : $@convention(thin) (@inout AO) -> () {
+bb0(%0 : $*AO):
+  %1 = unchecked_addr_cast %0 to $*AnyObject
+  %2 = load [copy] %1
+  fix_lifetime %2
+  destroy_value %2
+  %6 = tuple ()
+  return %6
+}
+
+// CHECK-LABEL: sil [ossa] @dont_remove_non_class_addr_cast2 :
+// CHECK:         %1 = unchecked_addr_cast %0
+// CHECK:       } // end sil function 'dont_remove_non_class_addr_cast2'
+sil [ossa] @dont_remove_non_class_addr_cast2 : $@convention(thin) (@inout AnyObject) -> () {
+bb0(%0 : $*AnyObject):
+  %1 = unchecked_addr_cast %0 to $*AO
+  %2 = load [copy] %1
+  fix_lifetime %2
+  destroy_value %2
+  %6 = tuple ()
+  return %6
+}
+

test/SILOptimizer/simplify_load_borrow.sil

@@ -93,3 +93,20 @@ bb0(%0 : $*B):
   %4 = tuple ()
   return %4
 }
+
+// CHECK-LABEL: sil [ossa] @load_of_addr_cast :
+// CHECK:         %1 = load_borrow %0
+// CHECK:         %2 = unchecked_ref_cast %1 to $AnyObject
+// CHECK:         fix_lifetime %2
+// CHECK:         end_borrow %1
+// CHECK:       } // end sil function 'load_of_addr_cast'
+sil [ossa] @load_of_addr_cast : $@convention(thin) (@inout B) -> () {
+bb0(%0 : $*B):
+  %1 = unchecked_addr_cast %0 to $*AnyObject
+  %2 = load_borrow %1
+  fix_lifetime %2
+  end_borrow %2
+  %6 = tuple ()
+  return %6
+}
+