Merge pull request #24925 from aschwaighofer/dynamically_replaceable_closure_fixes_5.1

[5.1] Fix closurescope analysis and capture promotion

Author: aschwaighofer

include/swift/SIL/SILInstruction.h

@@ -1806,6 +1806,8 @@ class ApplyInstBase<Impl, Base, false> : public Base {
   /// issues. It is at the bottom of the file.
   SILFunction *getCalleeFunction() const;
 
+  bool isCalleeDynamicallyReplaceable() const;
+
   /// Gets the referenced function if the callee is a function_ref instruction.
   SILFunction *getReferencedFunction() const {
     if (auto *FRI = dyn_cast<FunctionRefBaseInst>(getCallee()))
@@ -7683,6 +7685,27 @@ SILValue ApplyInstBase<Impl, Base, false>::getCalleeOrigin() const {
   }
 }
 
+template <class Impl, class Base>
+bool ApplyInstBase<Impl, Base, false>::isCalleeDynamicallyReplaceable() const {
+  SILValue Callee = getCalleeOrigin();
+
+  while (true) {
+    if (auto *FRI = dyn_cast<FunctionRefInst>(Callee))
+      return false;
+
+    if (auto *FRI = dyn_cast<DynamicFunctionRefInst>(Callee))
+      return true;
+    if (auto *FRI = dyn_cast<PreviousDynamicFunctionRefInst>(Callee))
+      return true;
+
+    if (auto *PAI = dyn_cast<PartialApplyInst>(Callee)) {
+      Callee = PAI->getCalleeOrigin();
+      continue;
+    }
+    return false;
+  }
+}
+
 template <class Impl, class Base>
 SILFunction *ApplyInstBase<Impl, Base, false>::getCalleeFunction() const {
   SILValue Callee = getCalleeOrigin();

lib/SILOptimizer/Analysis/ClosureScope.cpp

@@ -82,7 +82,8 @@ class ClosureScopeData {
   void recordScope(PartialApplyInst *PAI) {
     // Only track scopes of non-escaping closures.
     auto closureTy = PAI->getCallee()->getType().castTo<SILFunctionType>();
-    if (!isNonEscapingClosure(closureTy))
+    if (!isNonEscapingClosure(closureTy) ||
+        PAI->isCalleeDynamicallyReplaceable())
       return;
 
     auto closureFunc = PAI->getCalleeFunction();

lib/SILOptimizer/IPO/CapturePromotion.cpp

@@ -976,7 +976,7 @@ bool isPartialApplyNonEscapingUser(Operand *CurrentOp, PartialApplyInst *PAI,
   unsigned Index = OpNo - 1 + closureConv.getNumSILArguments();
 
   auto *Fn = PAI->getReferencedFunction();
-  if (!Fn || !Fn->isDefinition()) {
+  if (!Fn || !Fn->isDefinition() || Fn->isDynamicallyReplaceable()) {
     LLVM_DEBUG(llvm::dbgs() << "        FAIL! Not a direct function definition "
                           "reference.\n");
     return false;

test/SILOptimizer/access_enforcement_selection.sil

@@ -99,7 +99,24 @@ bb0:
   return %9 : $()
 }
 
+sil [dynamically_replacable] [ossa] @closureCapturingByStorageAddress2 : $@convention(thin) (@inout_aliasable Builtin.Int64) -> ()
 
+// Make sure that we handle dynamic_function_ref.
+sil hidden [ossa] @escapeAsArgumentToPartialApplyDynamic : $@convention(thin) () -> () {
+bb0:
+  %2 = alloc_box ${ var Builtin.Int64 }, var, name "x"
+  %3 = project_box %2 : ${ var Builtin.Int64 }, 0
+  %4 = function_ref @takesInoutAndClosure : $@convention(thin) (@inout Builtin.Int64, @guaranteed @callee_guaranteed () -> ()) -> ()
+  %5 = dynamic_function_ref @closureCapturingByStorageAddress2 : $@convention(thin) (@inout_aliasable Builtin.Int64) -> ()
+  %6 = partial_apply [callee_guaranteed] %5(%3) : $@convention(thin) (@inout_aliasable Builtin.Int64) -> ()
+  %7 = begin_access [modify] [unknown] %3 : $*Builtin.Int64
+  %8 = apply %4(%7, %6) : $@convention(thin) (@inout Builtin.Int64, @guaranteed @callee_guaranteed () -> ()) -> ()
+  end_access %7 : $*Builtin.Int64
+  destroy_value %6 : $@callee_guaranteed () -> ()
+  destroy_value %2 : ${ var Builtin.Int64 }
+  %9 = tuple ()
+  return %9 : $()
+}
 // Test static enforcement of copied boxes.
 // FIXME: Oops... We make this dynamic.
 //

test/SILOptimizer/capture_promotion.sil

@@ -474,3 +474,59 @@ bb0:
   %tuple = tuple()
   return %tuple : $()
 }
+
+sil private [dynamically_replacable] @closure0_guaranteed2 : $@convention(thin) (@guaranteed <τ_0_0> { var τ_0_0 } <Foo>, @guaranteed <τ_0_0> { var τ_0_0 } <Baz>, @guaranteed <τ_0_0> { var τ_0_0 } <Int>) -> Int {
+bb0(%0 : $<τ_0_0> { var τ_0_0 } <Foo>, %2 : $<τ_0_0> { var τ_0_0 } <Baz>, %4 : $<τ_0_0> { var τ_0_0 } <Int>):
+  %1 = project_box %0 : $<τ_0_0> { var τ_0_0 } <Foo>, 0
+  %3 = project_box %2 : $<τ_0_0> { var τ_0_0 } <Baz>, 0
+  %5 = project_box %4 : $<τ_0_0> { var τ_0_0 } <Int>, 0
+  %6 = tuple ()
+  %7 = function_ref @dummy_func : $@convention(thin) (Int, Int, Int) -> Int
+  %8 = load %1 : $*Foo
+  strong_retain %8 : $Foo
+  %10 = class_method %8 : $Foo, #Foo.foo!1 : (Foo) -> () -> Int, $@convention(method) (@guaranteed Foo) -> Int
+  %11 = apply %10(%8) : $@convention(method) (@guaranteed Foo) -> Int
+  %12 = struct_element_addr %3 : $*Baz, #Baz.x
+  %13 = load %12 : $*Int
+  %14 = load %5 : $*Int
+  %15 = apply %7(%11, %13, %14) : $@convention(thin) (Int, Int, Int) -> Int
+  return %15 : $Int
+}
+
+// This should not crash.
+sil @test_capture_promotion_dynamic_function_ref : $@convention(thin) () -> () {
+bb0:
+  %0 = tuple ()
+  %1 = alloc_box $<τ_0_0> { var τ_0_0 } <Foo>
+  %1a = project_box %1 : $<τ_0_0> { var τ_0_0 } <Foo>, 0
+  %2 = function_ref @foo_allocating_init : $@convention(thin) (@thick Foo.Type) -> @owned Foo
+  %3 = metatype $@thick Foo.Type
+  %4 = apply %2(%3) : $@convention(thin) (@thick Foo.Type) -> @owned Foo
+  store %4 to %1a : $*Foo
+  %6 = alloc_box $<τ_0_0> { var τ_0_0 } <Baz>
+  %6a = project_box %6 : $<τ_0_0> { var τ_0_0 } <Baz>, 0
+  %7 = function_ref @baz_init : $@convention(thin) (@thin Baz.Type) -> @owned Baz
+  %8 = metatype $@thin Baz.Type
+  %9 = apply %7(%8) : $@convention(thin) (@thin Baz.Type) -> @owned Baz
+  store %9 to %6a : $*Baz
+  %11 = alloc_box $<τ_0_0> { var τ_0_0 } <Int>
+  %11a = project_box %11 : $<τ_0_0> { var τ_0_0 } <Int>, 0
+  %12 = function_ref @convert_from_integer_literal : $@convention(thin) (Builtin.Word, @thin Int.Type) -> Int
+  %13 = metatype $@thin Int.Type
+  %14 = integer_literal $Builtin.Word, 3
+  %15 = apply %12(%14, %13) : $@convention(thin) (Builtin.Word, @thin Int.Type) -> Int
+  store %15 to %11a : $*Int
+
+  %17 = dynamic_function_ref @closure0_guaranteed2 : $@convention(thin) (@guaranteed <τ_0_0> { var τ_0_0 } <Foo>, @guaranteed <τ_0_0> { var τ_0_0 } <Baz>, @guaranteed <τ_0_0> { var τ_0_0 } <Int>) -> Int
+  %21 = partial_apply [callee_guaranteed] [on_stack] %17(%1, %6, %11) : $@convention(thin) (@guaranteed <τ_0_0> { var τ_0_0 } <Foo>, @guaranteed <τ_0_0> { var τ_0_0 } <Baz>, @guaranteed <τ_0_0> { var τ_0_0 } <Int>) -> Int
+  %md = mark_dependence %21 :$@noescape @callee_guaranteed () -> Int on %1 : $<τ_0_0> { var τ_0_0 } <Foo>
+  %md2 = mark_dependence %md :$@noescape @callee_guaranteed () -> Int on %6 : $<τ_0_0> { var τ_0_0 } <Baz>
+  %md3 = mark_dependence %md2 :$@noescape @callee_guaranteed () -> Int on %11 : $<τ_0_0> { var τ_0_0 } <Int>
+  apply %md3() : $@noescape @callee_guaranteed () -> Int
+  dealloc_stack %21 : $@noescape @callee_guaranteed () -> Int
+  strong_release %11 : $<τ_0_0> { var τ_0_0 } <Int>
+  strong_release %6 : $<τ_0_0> { var τ_0_0 } <Baz>
+  strong_release %1 : $<τ_0_0> { var τ_0_0 } <Foo>
+  %tuple = tuple()
+  return %tuple : $()
+}