Merge pull request #21239 from atrick/5.0-fix-devirt-enum

[5.0] Fix SILCombine to devirtualize existentials wrapped in enum.

Author: eeckstein

include/swift/SILOptimizer/Utils/Existential.h

@@ -35,7 +35,7 @@ SILValue findInitExistentialFromGlobalAddrAndApply(GlobalAddrInst *GAI,
 /// alloc_stack user \p ASIUser.
 /// If the value is copied from another stack location, \p isCopied is set to
 /// true.
-SILValue getAddressOfStackInit(AllocStackInst *ASI, SILInstruction *ASIUser,
+SILValue getAddressOfStackInit(SILValue allocStackAddr, SILInstruction *ASIUser,
                                bool &isCopied);
 
 /// Find the init_existential, which could be used to determine a concrete

include/swift/SILOptimizer/Utils/Local.h

@@ -127,6 +127,13 @@ SILValue castValueToABICompatibleType(SILBuilder *B, SILLocation Loc,
                                       SILValue Value,
                                       SILType SrcTy,
                                       SILType DestTy);
+/// Peek through trivial Enum initialization, typically for pointless
+/// Optionals.
+///
+/// The returned InitEnumDataAddr dominates the given
+/// UncheckedTakeEnumDataAddrInst.
+InitEnumDataAddrInst *
+findInitAddressForTrivialEnum(UncheckedTakeEnumDataAddrInst *UTEDAI);
 
 /// Returns a project_box if it is the next instruction after \p ABI and
 /// and has \p ABI as operand. Otherwise it creates a new project_box right

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -729,6 +729,13 @@ static bool canReplaceCopiedArg(FullApplySite Apply,
   SILValue existentialAddr =
       cast<InitExistentialAddrInst>(InitExistential)->getOperand();
 
+  // If we peeked through an InitEnumDataAddr or some such, then don't assume we
+  // can reuse the copied value. It's likely destroyed by
+  // UncheckedTakeEnumDataInst before the copy.
+  auto *ASI = dyn_cast<AllocStackInst>(existentialAddr);
+  if (!ASI)
+    return false;
+
   // Return true only if the given value is guaranteed to be initialized across
   // the given call site.
   //

lib/SILOptimizer/Utils/Existential.cpp

@@ -10,12 +10,13 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "swift/SILOptimizer/Utils/Existential.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ProtocolConformance.h"
-#include "swift/SILOptimizer/Utils/Existential.h"
-#include "swift/SILOptimizer/Utils/CFG.h"
-#include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/BasicBlockUtils.h"
+#include "swift/SIL/InstructionUtils.h"
+#include "swift/SILOptimizer/Utils/CFG.h"
+#include "swift/SILOptimizer/Utils/Local.h"
 #include "llvm/ADT/SmallPtrSet.h"
 
 using namespace swift;
@@ -122,11 +123,14 @@ SILValue swift::findInitExistentialFromGlobalAddrAndApply(GlobalAddrInst *GAI,
 /// alloc_stack user \p ASIUser.
 /// If the value is copied from another stack location, \p isCopied is set to
 /// true.
-SILValue swift::getAddressOfStackInit(AllocStackInst *ASI,
+///
+/// allocStackAddr may either itself be an AllocStackInst or an
+/// InitEnumDataAddrInst that projects the value of an AllocStackInst.
+SILValue swift::getAddressOfStackInit(SILValue allocStackAddr,
                                       SILInstruction *ASIUser, bool &isCopied) {
   SILInstruction *SingleWrite = nullptr;
   // Check that this alloc_stack is initialized only once.
-  for (auto Use : ASI->getUses()) {
+  for (auto Use : allocStackAddr->getUses()) {
     auto *User = Use->getUser();
 
     // Ignore instructions which don't write to the stack location.
@@ -138,7 +142,7 @@ SILValue swift::getAddressOfStackInit(AllocStackInst *ASI,
       continue;
     }
     if (auto *CAI = dyn_cast<CopyAddrInst>(User)) {
-      if (CAI->getDest() == ASI) {
+      if (CAI->getDest() == allocStackAddr) {
         if (SingleWrite)
           return SILValue();
         SingleWrite = CAI;
@@ -168,8 +172,13 @@ SILValue swift::getAddressOfStackInit(AllocStackInst *ASI,
 
   // A very simple dominance check. As ASI is an operand of ASIUser,
   // SingleWrite dominates ASIUser if it is in the same block as ASI or ASIUser.
+  //
+  // If allocStack holds an Optional, then ASI is an InitEnumDataAddrInst
+  // projection and not strictly an operand of ASIUser. We rely on the guarantee
+  // that this InitEnumDataAddrInst must occur before the InjectEnumAddrInst
+  // that was the source of the existential address.
   SILBasicBlock *BB = SingleWrite->getParent();
-  if (BB != ASI->getParent() && BB != ASIUser->getParent())
+  if (BB != allocStackAddr->getParentBlock() && BB != ASIUser->getParent())
     return SILValue();
 
   if (auto *CAI = dyn_cast<CopyAddrInst>(SingleWrite)) {
@@ -179,6 +188,21 @@ SILValue swift::getAddressOfStackInit(AllocStackInst *ASI,
     SILValue CAISrc = CAI->getSrc();
     if (auto *ASI = dyn_cast<AllocStackInst>(CAISrc))
       return getAddressOfStackInit(ASI, CAI, isCopied);
+
+    // Recognize a stack location holding an Optional.
+    //   %stack_adr = alloc_stack
+    //   %data_adr  = init_enum_data_addr %stk_adr
+    //   %enum_adr  = inject_enum_addr %stack_adr
+    //   %copy_src  = unchecked_take_enum_data_addr %enum_adr
+    // Replace %copy_src with %data_adr and recurse.
+    //
+    // TODO: a general Optional elimination sil-combine could
+    // supersede this check.
+    if (auto *UTEDAI = dyn_cast<UncheckedTakeEnumDataAddrInst>(CAISrc)) {
+      if (InitEnumDataAddrInst *IEDAI = findInitAddressForTrivialEnum(UTEDAI))
+        return getAddressOfStackInit(IEDAI, CAI, isCopied);
+    }
+
     // Check if the CAISrc is a global_addr.
     if (auto *GAI = dyn_cast<GlobalAddrInst>(CAISrc)) {
       return findInitExistentialFromGlobalAddrAndCopyAddr(GAI, CAI);

lib/SILOptimizer/Utils/Local.cpp

@@ -653,6 +653,51 @@ ProjectBoxInst *swift::getOrCreateProjectBox(AllocBoxInst *ABI, unsigned Index){
   return B.createProjectBox(ABI->getLoc(), ABI, Index);
 }
 
+// Peek through trivial Enum initialization, typically for pointless
+// Optionals.
+//
+// Given an UncheckedTakeEnumDataAddrInst, check that there are no
+// other uses of the Enum value and return the address used to initialized the
+// enum's payload:
+//
+//   %stack_adr = alloc_stack
+//   %data_adr  = init_enum_data_addr %stk_adr
+//   %enum_adr  = inject_enum_addr %stack_adr
+//   %copy_src  = unchecked_take_enum_data_addr %enum_adr
+//   dealloc_stack %stack_adr
+//   (No other uses of %stack_adr.)
+InitEnumDataAddrInst *
+swift::findInitAddressForTrivialEnum(UncheckedTakeEnumDataAddrInst *UTEDAI) {
+  auto *ASI = dyn_cast<AllocStackInst>(UTEDAI->getOperand());
+  if (!ASI)
+    return nullptr;
+
+  SILInstruction *singleUser = nullptr;
+  for (auto use : ASI->getUses()) {
+    auto *user = use->getUser();
+    if (user == UTEDAI)
+      continue;
+
+    // As long as there's only one UncheckedTakeEnumDataAddrInst and one
+    // InitEnumDataAddrInst, we don't care how many InjectEnumAddr and
+    // DeallocStack users there are.
+    if (isa<InjectEnumAddrInst>(user) || isa<DeallocStackInst>(user))
+      continue;
+
+    if (singleUser)
+      return nullptr;
+
+    singleUser = user;
+  }
+  if (!singleUser)
+    return nullptr;
+
+  // Assume, without checking, that the returned InitEnumDataAddr dominates the
+  // given UncheckedTakeEnumDataAddrInst, because that's how SIL is defined. I
+  // don't know where this is actually verified.
+  return dyn_cast<InitEnumDataAddrInst>(singleUser);
+}
+
 //===----------------------------------------------------------------------===//
 //                       String Concatenation Optimizer
 //===----------------------------------------------------------------------===//

test/SILOptimizer/devirt_specialized_conformance.swift

@@ -42,3 +42,35 @@ func driver() {
 }
 
 driver()
+
+// <rdar://problem/46322928> Failure to devirtualize a protocol method
+// applied to an opened existential blocks implemention of
+// DataProtocol.
+public protocol ContiguousBytes {
+    func withUnsafeBytes<R>(_ body: (UnsafeRawBufferPointer) throws -> R) rethrows -> R
+}
+
+extension Array : ContiguousBytes {}
+extension ContiguousArray : ContiguousBytes {}
+
+@inline(never)
+func takesPointer(_ p: UnsafeRawBufferPointer) {}
+
+// In specialized testWithUnsafeBytes<A>(_:), the conditional case and call to withUnsafeBytes must be eliminated.
+// Normally, we expect Array.withUnsafeBytes to be inlined so we would see:
+//   [[TAKES_PTR:%.*]] = function_ref @$s30devirt_specialized_conformance12takesPointeryySWF : $@convention(thin) (UnsafeRawBufferPointer) -> ()
+//   apply [[TAKES_PTR]](%{{.*}}) : $@convention(thin) (UnsafeRawBufferPointer) -> ()
+// But the inlining isn't consistent across builds with and without debug info.
+//
+// CHECK-LABEL: sil shared [noinline] @$s30devirt_specialized_conformance19testWithUnsafeBytesyyxlFSayypG_Tg5 : $@convention(thin) (@guaranteed Array<Any>) -> () {
+// CHECK: bb0
+// CHECK-NOT: witness_method
+// CHECK-LABEL: } // end sil function '$s30devirt_specialized_conformance19testWithUnsafeBytesyyxlFSayypG_Tg5'
+@inline(never)
+func testWithUnsafeBytes<T>(_ t: T) {
+  if let cb = t as? ContiguousBytes {
+    cb.withUnsafeBytes { takesPointer($0) }
+  }
+}
+
+testWithUnsafeBytes([])

test/SILOptimizer/sil_combine_concrete_existential.sil

@@ -437,3 +437,65 @@ bb0:
   %v = tuple ()
   return %v : $()
 }
+
+// <rdar://problem/46322928> Failure to devirtualize a protocol method
+// applied to an opened existential blocks implemention of
+// DataProtocol.
+public protocol ContiguousBytes {
+  func withUnsafeBytes<R>(_ body: (UnsafeRawBufferPointer) throws -> R) rethrows -> R
+}
+
+extension Array : ContiguousBytes {
+}
+
+// specialized thunk for @callee_guaranteed (@unowned UnsafeRawBufferPointer) -> (@error @owned Error)
+sil [transparent] [reabstraction_thunk] @testDevirtExistentialEnumThunk : $@convention(thin) (UnsafeRawBufferPointer) -> (@out (), @error Error)
+
+// Test that the witness method's opened archetype is replaced by a
+// concrete conformance.
+//
+// Note that the enum is destroyed by unchecked_take_enum_data_addr
+// before the apply (and the extracted data is destroy by the copy
+// itself), so SILCombine cannot replace the call's self
+// argument. Instead, we expect a later devirtualizer pass to just
+// insert a cast of the opened existential.
+//
+// CHECK-LABEL: sil shared [noinline] @testDevirtExistentialEnum : $@convention(thin) (@guaranteed Array<Int>) -> () {
+// CHECK: [[ALLOC_EXISTENTIAL:%.*]] = alloc_stack $ContiguousBytes
+// CHECK: [[ALLOC_ENUM:%.*]] = alloc_stack $Optional<ContiguousBytes>
+// CHECK: [[ENUM:%.*]] = init_enum_data_addr [[ALLOC_ENUM]] : $*Optional<ContiguousBytes>, #Optional.some!enumelt.1
+// CHECK: [[INIT_DATA:%.*]] = init_existential_addr [[ENUM]] : $*ContiguousBytes, $Array<Int>
+// CHECK: store %0 to [[INIT_DATA]] : $*Array<Int>
+// CHECK: inject_enum_addr [[ALLOC_ENUM]] : $*Optional<ContiguousBytes>, #Optional.some!enumelt.1
+// CHECK: [[TAKE_DATA:%.*]] = unchecked_take_enum_data_addr [[ALLOC_ENUM]] : $*Optional<ContiguousBytes>, #Optional.some!enumelt.1
+// CHECK: copy_addr [take] [[TAKE_DATA]] to [initialization] [[ALLOC_EXISTENTIAL]] : $*ContiguousBytes
+// CHECK: [[OPENED:%.*]] = open_existential_addr immutable_access [[ALLOC_EXISTENTIAL]] : $*ContiguousBytes to $*@opened("8402EC1A-F35D-11E8-950A-D0817AD9F6DD") ContiguousBytes
+// CHECK: [[THUNK:%.*]] = function_ref @testDevirtExistentialEnumThunk : $@convention(thin) (UnsafeRawBufferPointer) -> (@out (), @error Error)
+// CHECK: [[TTF:%.*]] = thin_to_thick_function [[THUNK:%.*]] : $@convention(thin) (UnsafeRawBufferPointer) -> (@out (), @error Error) to $@noescape @callee_guaranteed (UnsafeRawBufferPointer) -> (@out (), @error Error)
+// CHECK: [[WM:%.*]] = witness_method $Array<Int>, #ContiguousBytes.withUnsafeBytes!1 : <Self where Self : ContiguousBytes><R> (Self) -> ((UnsafeRawBufferPointer) throws -> R) throws -> R : $@convention(witness_method: ContiguousBytes) <τ_0_0 where τ_0_0 : ContiguousBytes><τ_1_0> (@noescape @callee_guaranteed (UnsafeRawBufferPointer) -> (@out τ_1_0, @error Error), @in_guaranteed τ_0_0) -> (@out τ_1_0, @error Error)
+// CHECK: apply [nothrow] [[WM]]<@opened("8402EC1A-F35D-11E8-950A-D0817AD9F6DD") ContiguousBytes, ()>(%{{.*}}, [[TTF]], [[OPENED]]) : $@convention(witness_method: ContiguousBytes) <τ_0_0 where τ_0_0 : ContiguousBytes><τ_1_0> (@noescape @callee_guaranteed (UnsafeRawBufferPointer) -> (@out τ_1_0, @error Error), @in_guaranteed τ_0_0) -> (@out τ_1_0, @error Error)
+// CHECK: destroy_addr [[ALLOC_EXISTENTIAL]] : $*ContiguousBytes
+// CHECK-LABEL: } // end sil function 'testDevirtExistentialEnum'
+sil shared [noinline] @testDevirtExistentialEnum : $@convention(thin) (@guaranteed Array<Int>) -> () {
+bb0(%0 : $Array<Int>):
+  %3 = alloc_stack $ContiguousBytes
+  %4 = alloc_stack $Optional<ContiguousBytes>
+  %5 = init_enum_data_addr %4 : $*Optional<ContiguousBytes>, #Optional.some!enumelt.1
+  %6 = init_existential_addr %5 : $*ContiguousBytes, $Array<Int>
+  store %0 to %6 : $*Array<Int>
+  inject_enum_addr %4 : $*Optional<ContiguousBytes>, #Optional.some!enumelt.1
+  %9 = unchecked_take_enum_data_addr %4 : $*Optional<ContiguousBytes>, #Optional.some!enumelt.1
+  copy_addr [take] %9 to [initialization] %3 : $*ContiguousBytes
+  dealloc_stack %4 : $*Optional<ContiguousBytes>
+  %12 = open_existential_addr immutable_access %3 : $*ContiguousBytes to $*@opened("8402EC1A-F35D-11E8-950A-D0817AD9F6DD") ContiguousBytes
+  %13 = alloc_stack $()
+  %14 = function_ref @testDevirtExistentialEnumThunk : $@convention(thin) (UnsafeRawBufferPointer) -> (@out (), @error Error)
+  %15 = thin_to_thick_function %14 : $@convention(thin) (UnsafeRawBufferPointer) -> (@out (), @error Error) to $@noescape @callee_guaranteed (UnsafeRawBufferPointer) -> (@out (), @error Error)
+  %16 = witness_method $@opened("8402EC1A-F35D-11E8-950A-D0817AD9F6DD") ContiguousBytes, #ContiguousBytes.withUnsafeBytes!1 : <Self where Self : ContiguousBytes><R> (Self) -> ((UnsafeRawBufferPointer) throws -> R) throws -> R, %12 : $*@opened("8402EC1A-F35D-11E8-950A-D0817AD9F6DD") ContiguousBytes : $@convention(witness_method: ContiguousBytes) <τ_0_0 where τ_0_0 : ContiguousBytes><τ_1_0> (@noescape @callee_guaranteed (UnsafeRawBufferPointer) -> (@out τ_1_0, @error Error), @in_guaranteed τ_0_0) -> (@out τ_1_0, @error Error)
+  %21 = apply [nothrow] %16<@opened("8402EC1A-F35D-11E8-950A-D0817AD9F6DD") ContiguousBytes, ()>(%13, %15, %12) : $@convention(witness_method: ContiguousBytes) <τ_0_0 where τ_0_0 : ContiguousBytes><τ_1_0> (@noescape @callee_guaranteed (UnsafeRawBufferPointer) -> (@out τ_1_0, @error Error), @in_guaranteed τ_0_0) -> (@out τ_1_0, @error Error)
+  dealloc_stack %13 : $*()
+  destroy_addr %3 : $*ContiguousBytes
+  dealloc_stack %3 : $*ContiguousBytes
+  %25 = tuple ()
+  return %25 : $()
+}