Merge pull request #39004 from eeckstein/fix-closure-specializer

ClosureSpecializer: avoid an infinite optimization loop.

Author: eeckstein

lib/SILOptimizer/IPO/ClosureSpecializer.cpp

@@ -57,6 +57,7 @@
 
 #define DEBUG_TYPE "closure-specialization"
 #include "swift/Basic/Range.h"
+#include "swift/Demangling/Demangler.h"
 #include "swift/SIL/InstructionUtils.h"
 #include "swift/SIL/SILCloner.h"
 #include "swift/SIL/SILFunction.h"
@@ -162,6 +163,14 @@ class ClosureSpecCloner : public SILClonerWithScopes<ClosureSpecCloner> {
 namespace {
 struct ClosureInfo;
 
+static SILFunction *getClosureCallee(SILInstruction *inst) {
+  if (auto *PAI = dyn_cast<PartialApplyInst>(inst))
+    return cast<FunctionRefInst>(PAI->getCallee())->getReferencedFunction();
+
+  auto *TTTFI = cast<ThinToThickFunctionInst>(inst);
+  return cast<FunctionRefInst>(TTTFI->getCallee())->getReferencedFunction();
+}
+
 class CallSiteDescriptor {
   ClosureInfo *CInfo;
   FullApplySite AI;
@@ -188,11 +197,7 @@ class CallSiteDescriptor {
   }
 
   SILFunction *getClosureCallee() const {
-    if (auto *PAI = dyn_cast<PartialApplyInst>(getClosure()))
-      return cast<FunctionRefInst>(PAI->getCallee())->getReferencedFunction();
-
-    auto *TTTFI = cast<ThinToThickFunctionInst>(getClosure());
-    return cast<FunctionRefInst>(TTTFI->getCallee())->getReferencedFunction();
+    return ::getClosureCallee(getClosure());
   }
 
   bool closureHasRefSemanticContext() const {
@@ -1065,6 +1070,59 @@ static bool canSpecializeFullApplySite(FullApplySiteKind kind) {
   llvm_unreachable("covered switch");
 }
 
+static int getSpecializationLevelRecursive(StringRef funcName, Demangler &parent) {
+  using namespace Demangle;
+
+  Demangler demangler;
+  demangler.providePreallocatedMemory(parent);
+
+  // Check for this kind of node tree:
+  //
+  // kind=Global
+  //   kind=FunctionSignatureSpecialization
+  //     kind=SpecializationPassID, index=1
+  //     kind=FunctionSignatureSpecializationParam
+  //       kind=FunctionSignatureSpecializationParamKind, index=5
+  //       kind=FunctionSignatureSpecializationParamPayload, text="..."
+  //
+  Node *root = demangler.demangleSymbol(funcName);
+  if (!root)
+    return 0;
+  if (root->getKind() != Node::Kind::Global)
+    return 0;
+  Node *funcSpec = root->getFirstChild();
+  if (!funcSpec || funcSpec->getNumChildren() < 2)
+    return 0;
+  if (funcSpec->getKind() != Node::Kind::FunctionSignatureSpecialization)
+    return 0;
+  Node *param = funcSpec->getChild(1);
+  if (param->getKind() != Node::Kind::FunctionSignatureSpecializationParam)
+    return 0;
+  if (param->getNumChildren() < 2)
+    return 0;
+  Node *kindNd = param->getChild(0);
+  if (kindNd->getKind() != Node::Kind::FunctionSignatureSpecializationParamKind)
+    return 0;
+  auto kind = FunctionSigSpecializationParamKind(kindNd->getIndex());
+  if (kind != FunctionSigSpecializationParamKind::ConstantPropFunction)
+    return 0;
+    
+  Node *payload = param->getChild(1);
+  if (payload->getKind() != Node::Kind::FunctionSignatureSpecializationParamPayload)
+    return 1;
+  // Check if the specialized function is a specialization itself.
+  return 1 + getSpecializationLevelRecursive(payload->getText(), demangler);
+}
+
+/// If \p function is a function-signature specialization for a constant-
+/// propagated function argument, returns 1.
+/// If \p function is a specialization of such a specialization, returns 2.
+/// And so on.
+static int getSpecializationLevel(SILFunction *f) {
+  Demangle::StackAllocatedDemangler<1024> demangler;
+  return getSpecializationLevelRecursive(f->getName(), demangler);
+}
+
 bool SILClosureSpecializerTransform::gatherCallSites(
     SILFunction *Caller,
     llvm::SmallVectorImpl<std::unique_ptr<ClosureInfo>> &ClosureCandidates,
@@ -1252,6 +1310,24 @@ bool SILClosureSpecializerTransform::gatherCallSites(
           continue;
         }
 
+        // Avoid an infinite specialization loop caused by repeated runs of
+        // ClosureSpecializer and CapturePropagation.
+        // CapturePropagation propagates constant function-literals. Such
+        // function specializations can then be optimized again by the
+        // ClosureSpecializer and so on.
+        // This happens if a closure argument is called _and_ referenced in
+        // another closure, which is passed to a recursive call. E.g.
+        //
+        // func foo(_ c: @escaping () -> ()) {
+        //   c()
+        //   foo({ c() })
+        // }
+        //
+        // A limit of 2 is good enough and will not be exceed in "regular"
+        // optimization scenarios.
+        if (getSpecializationLevel(getClosureCallee(ClosureInst)) > 2)
+          continue;
+
         // Compute the final release points of the closure. We will insert
         // release of the captured arguments here.
         if (!CInfo)

test/SILOptimizer/closure_specialize_loop.swift

@@ -0,0 +1,16 @@
+// RUN: %{python} %S/../Inputs/timeout.py 10 %target-swift-frontend -O -parse-as-library %s -emit-sil | %FileCheck %s
+
+public func callit() {
+    testit { false }
+}
+
+// Check if the compiler terminates and does not full into an infinite optimization
+// loop between the ClosureSpecializer and CapturePropagation.
+
+// CHECK-LABEL: sil @$s23closure_specialize_loop6testit1cySbyc_tF
+public func testit(c: @escaping () -> Bool) {
+  if c() {
+    testit { !c() }
+  }
+}
+