Fix infinite recursion in ClosureSpecialize

Fixes getSpecializationLevelRecursive to handle recursive manglings caused by interleaving CapturePropagation and ClosureSpecialize passes.

For some reason, only the first closure parameter was checked for recursion. We need to handle patterns like this:

  kind=FunctionSignatureSpecialization
    kind=SpecializationPassID, index=3
    kind=FunctionSignatureSpecializationParam
    kind=FunctionSignatureSpecializationParam
      kind=FunctionSignatureSpecializationParamKind, index=0
      kind=FunctionSignatureSpecializationParamPayload, text="$s4test10ExpressionO8contains5whereS3bXE_tFSbACXEfU_S2bXEfU_36$s4test12IndirectEnumVACycfcS2bXEfU_Tf3npf_n"

I fixed the logic so we now check for recursion on all closure parameters and bail out on unrecognized mangling formats.

For reference, see summary.sil in
Infinitely recursive closure specialization #61955
#61955

Fixes rdar://101589190 (Swift Compiler hangs when building this code for release)

Author: atrick

lib/SILOptimizer/IPO/ClosureSpecializer.cpp

@@ -1073,6 +1073,8 @@ static bool canSpecializeFullApplySite(FullApplySiteKind kind) {
   llvm_unreachable("covered switch");
 }
 
+const int SpecializationLevelLimit = 2;
+
 static int getSpecializationLevelRecursive(StringRef funcName, Demangler &parent) {
   using namespace Demangle;
 
@@ -1098,23 +1100,49 @@ static int getSpecializationLevelRecursive(StringRef funcName, Demangler &parent
     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);
+
+  // Handle up to six parameters, otherwise assume infinite recursion
+  if (funcSpec->getNumChildren() > 6)
+    return SpecializationLevelLimit + 1;
+
+  // Match any function specialization. We check for constant propagation at the
+  // parameter level.
+  Node *param = funcSpec->getChild(0);
+  if (param->getKind() != Node::Kind::SpecializationPassID)
+    return SpecializationLevelLimit + 1; // unrecognized format
+
+  unsigned maxParamLevel = 0;
+  for (unsigned paramIdx = 1; paramIdx < funcSpec->getNumChildren();
+       ++paramIdx) {
+    Node *param = funcSpec->getChild(paramIdx);
+    if (param->getKind() != Node::Kind::FunctionSignatureSpecializationParam)
+      return SpecializationLevelLimit + 1; // unrecognized format
+
+    // A parameter is recursive if it has a kind with index and type payload
+    if (param->getNumChildren() < 2)
+      continue;
+
+    Node *kindNd = param->getChild(0);
+    if (kindNd->getKind()
+        != Node::Kind::FunctionSignatureSpecializationParamKind) {
+      return SpecializationLevelLimit + 1; // unrecognized format
+    }
+    auto kind = FunctionSigSpecializationParamKind(kindNd->getIndex());
+    if (kind != FunctionSigSpecializationParamKind::ConstantPropFunction)
+      continue;
+    Node *payload = param->getChild(1);
+    if (payload->getKind()
+        != Node::Kind::FunctionSignatureSpecializationParamPayload) {
+      return SpecializationLevelLimit + 1; // unrecognized format
+    }
+    // Check if the specialized function is a specialization itself.
+    unsigned paramLevel =
+      1 + getSpecializationLevelRecursive(payload->getText(), demangler);
+    if (paramLevel > maxParamLevel)
+      maxParamLevel = paramLevel;
+  }
+  llvm::dbgs() << "maxParamLevel = " << maxParamLevel << "\n";
+  return maxParamLevel;
 }
 
 /// If \p function is a function-signature specialization for a constant-
@@ -1328,9 +1356,10 @@ bool SILClosureSpecializerTransform::gatherCallSites(
         //
         // A limit of 2 is good enough and will not be exceed in "regular"
         // optimization scenarios.
-        if (getSpecializationLevel(getClosureCallee(ClosureInst)) > 2)
+        if (getSpecializationLevel(getClosureCallee(ClosureInst))
+            > SpecializationLevelLimit) {
           continue;
-
+        }
         // Compute the final release points of the closure. We will insert
         // release of the captured arguments here.
         if (!CInfo)
@@ -1395,6 +1424,8 @@ bool SILClosureSpecializerTransform::specialize(SILFunction *Caller,
       if (!NewF) {
         NewF = ClosureSpecCloner::cloneFunction(FuncBuilder, CSDesc, NewFName);
         addFunctionToPassManagerWorklist(NewF, CSDesc.getApplyCallee());
+        LLVM_DEBUG(llvm::dbgs() << "\nThe rewritten callee is:\n";
+                   NewF->dump());
       }
 
       // Rewrite the call
@@ -1404,6 +1435,10 @@ bool SILClosureSpecializerTransform::specialize(SILFunction *Caller,
       Changed = true;
     }
   }
+  LLVM_DEBUG(if (Changed) {
+      llvm::dbgs() << "\nThe rewritten caller is:\n";
+      Caller->dump();
+    });
   return Changed;
 }