Mangling: add a closure specialization mangling for arguments which specialize for the same closure as a previous argument

For example:
```
  %1 = partial_apply %closure
  apply %f(%1, %1)    // first argument: `.closure(%1)`
                      // second argument: `.previousArgumentIndex(0)`
```

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -1427,7 +1427,8 @@ private struct PullbackClosureInfo {
 
   func specializedCalleeName(_ context: FunctionPassContext) -> String {
     let closureArgs = Array(self.closureArgDescriptors.map {
-      (argumentIndex: $0.closureArgIndex, argumentValue: $0.closure)
+      (argumentIndex: $0.closureArgIndex,
+       argumentValue: FunctionPassContext.ClosureArgumentMangling.closure($0.closure))
     })
     return context.mangle(withClosureArguments: closureArgs, from: pullbackFn)
   }

SwiftCompilerSources/Sources/Optimizer/PassManager/FunctionPassContext.swift

@@ -111,10 +111,34 @@ struct FunctionPassContext : MutatingContext {
     return String(taking: bridgedPassContext.mangleOutlinedVariable(function.bridged))
   }
 
-  func mangle(withClosureArguments closureArgs: [(argumentIndex: Int, argumentValue: Value)],
+  enum ClosureArgumentMangling {
+    case closure(SingleValueInstruction)
+
+    /// The argument specializes for the same closure as a previous argument, e.g.
+    /// ```
+    ///   %1 = partial_apply %closure
+    ///   apply %f(%1, %1)    // first argument: `.closure(%1)`
+    ///                       // second argument: `.previousArgumentIndex(0)`
+    case previousArgumentIndex(Int)
+  }
+
+  func mangle(withClosureArguments closureArgs: [(argumentIndex: Int, argumentValue: ClosureArgumentMangling)],
               from applySiteCallee: Function
   ) -> String {
-    closureArgs.withBridgedArrayRef{ bridgedClosureArgs in
+    let bridgedArgManglings = closureArgs.map {
+      switch $0.argumentValue {
+      case .closure(let closure):
+        return BridgedPassContext.ClosureArgMangling(argIdx: $0.argumentIndex,
+                                                     inst: Optional<Instruction>(closure).bridged,
+                                                     otherArgIdx: -1)
+      case .previousArgumentIndex(let idx):
+        return BridgedPassContext.ClosureArgMangling(argIdx: $0.argumentIndex,
+                                                     inst: OptionalBridgedInstruction(),
+                                                     otherArgIdx: idx)
+      }
+    }
+
+    return bridgedArgManglings.withBridgedArrayRef{ bridgedClosureArgs in
       String(taking: bridgedPassContext.mangleWithClosureArgs(bridgedClosureArgs, applySiteCallee.bridged))
     }
   }

docs/ABI/Mangling.rst

@@ -1362,6 +1362,7 @@ Some kinds need arguments, which precede ``Tf``.
   ARG-SPEC-KIND ::= 'n'                      // Unmodified argument
   ARG-SPEC-KIND ::= 'c'                      // Consumes n 'type' arguments which are closed over types in argument order
                                              // and one 'identifier' argument which is the closure symbol name
+  ARG-SPEC-KIND ::= 'C' NATURAL-ZERO         // the same closure as a previous argument <n>
   ARG-SPEC-KIND ::= 'p' CONST-PROP           // Constant propagated argument
   ARG-SPEC-KIND ::= 'e' 'D'? 'G'? 'X'?       // Generic argument, with optional dead, owned=>guaranteed or exploded-specifier
   ARG-SPEC-KIND ::= 'd' 'G'? 'X'?            // Dead argument, with optional owned=>guaranteed or exploded-specifier

include/swift/Demangling/Demangle.h

@@ -117,6 +117,7 @@ enum class FunctionSigSpecializationParamKind : unsigned {
   InOutToOut = 8,
   ConstantPropKeyPath = 9,
   ConstantPropStruct = 10,
+  ClosurePropPreviousArg = 11,
 
   // Option Set Flags use bits 6-31. This gives us 26 bits to use for option
   // flags.

include/swift/SILOptimizer/OptimizerBridging.h

@@ -187,9 +187,15 @@ struct BridgedPassContext {
   bool specializeAppliesInFunction(BridgedFunction function, bool isMandatory) const;
   BridgedOwnedString mangleOutlinedVariable(BridgedFunction function) const;
   BridgedOwnedString mangleAsyncRemoved(BridgedFunction function) const;
+
+  struct ClosureArgMangling {
+    SwiftInt argIdx;
+    OptionalBridgedInstruction inst;
+    SwiftInt otherArgIdx;
+  };
+
   BridgedOwnedString mangleWithDeadArgs(BridgedArrayRef bridgedDeadArgIndices, BridgedFunction function) const;
-  BridgedOwnedString mangleWithClosureArgs(BridgedArrayRef closureArgIndices,
-                                                               BridgedFunction applySiteCallee) const;
+  BridgedOwnedString mangleWithClosureArgs(BridgedArrayRef closureArgManglings, BridgedFunction applySiteCallee) const;
   BridgedOwnedString mangleWithConstCaptureArgs(BridgedArrayRef bridgedConstArgs,
                                                 BridgedFunction applySiteCallee) const;
   BridgedOwnedString mangleWithBoxToStackPromotedArgs(BridgedArrayRef bridgedPromotedArgIndices,

include/swift/SILOptimizer/Utils/SpecializationMangler.h

@@ -68,6 +68,7 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
     BoxToValue = 3,
     BoxToStack = 4,
     InOutToOut = 5,
+    ClosurePropPreviousArg = 6, // the same closure as a previous `ClosureProp` argument
 
     First_Option = 0,
     Last_Option = 31,
@@ -82,8 +83,20 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
     LastOptionSetEntry = 32768,
   };
 
-  using ArgInfo = std::pair<ArgumentModifierIntBase,
-                            NullablePtr<SILInstruction>>;
+  struct ArgInfo {
+    ArgumentModifier kind;
+    union {
+      SILInstruction *inst;
+      unsigned otherArgIdx; // only for `ClosurePropPreviousArg`
+    };
+
+    ArgInfo(ArgumentModifier kind) : kind(kind), inst(nullptr) {}
+
+    void append(ArgumentModifier modifier) {
+      kind = ArgumentModifier(ArgumentModifierIntBase(kind) | ArgumentModifierIntBase(modifier));
+    }
+  };
+
   // Information for each SIL argument in the original function before
   // specialization. This includes SIL indirect result argument required for
   // the original function type at the current stage of compilation.
@@ -99,9 +112,8 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
   void setArgumentConstantProp(unsigned OrigArgIdx, SILInstruction *constInst);
   void appendStringAsIdentifier(StringRef str);
 
-  void setArgumentClosureProp(unsigned OrigArgIdx, PartialApplyInst *PAI);
-  void setArgumentClosureProp(unsigned OrigArgIdx,
-                              ThinToThickFunctionInst *TTTFI);
+  void setArgumentClosureProp(unsigned OrigArgIdx, SILInstruction *closure);
+  void setArgumentClosurePropPreviousArg(unsigned OrigArgIdx, unsigned otherArgIdx);
   void setArgumentDead(unsigned OrigArgIdx);
   void setArgumentOwnedToGuaranteed(unsigned OrigArgIdx);
   void setArgumentGuaranteedToOwned(unsigned OrigArgIdx);
@@ -120,8 +132,7 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
 private:
   void mangleConstantProp(SILInstruction *constInst);
   void mangleClosureProp(SILInstruction *Inst);
-  void mangleArgument(ArgumentModifierIntBase ArgMod,
-                      NullablePtr<SILInstruction> Inst);
+  void mangleArgument(ArgInfo argInfo);
   void mangleReturnValue(ReturnValueModifierIntBase RetMod);
 };
 

lib/Demangling/Demangler.cpp

@@ -3430,6 +3430,18 @@ NodePointer Demangler::demangleFuncSpecParam(Node::Kind Kind) {
       return addChild(Param, createNode(
         Node::Kind::FunctionSignatureSpecializationParamKind,
         uint64_t(FunctionSigSpecializationParamKind::ClosureProp)));
+    case 'C': {
+      // Consumes an identifier and multiple type parameters.
+      // The parameters will be added later.
+      addChild(Param, createNode(
+        Node::Kind::FunctionSignatureSpecializationParamKind,
+        uint64_t(FunctionSigSpecializationParamKind::ClosurePropPreviousArg)));
+      int prevArgIdx = demangleNatural();
+      if (prevArgIdx < 0)
+        return nullptr;
+      return addChild(Param, createNode(
+         Node::Kind::FunctionSignatureSpecializationParamPayload, (Node::IndexType)prevArgIdx));
+    }
     case 'p': {
       for (;;) {
         switch (nextChar()) {

lib/Demangling/NodePrinter.cpp

@@ -1247,6 +1247,15 @@ void NodePrinter::printFunctionSigSpecializationParams(NodePointer Node,
       }
       Printer << "]";
       break;
+    case FunctionSigSpecializationParamKind::ClosurePropPreviousArg:
+      if (Idx + 2 > End)
+        return;
+      Printer << "[";
+      print(Node->getChild(Idx++), depth + 1);
+      Printer << " ";
+      print(Node->getChild(Idx++), depth + 1);
+      Printer << "]";
+      break;
     default:
       assert(
        ((V & unsigned(FunctionSigSpecializationParamKind::OwnedToGuaranteed)) ||
@@ -1889,11 +1898,15 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
   case Node::Kind::FunctionSignatureSpecializationParam:
     printer_unreachable("should be handled in printSpecializationPrefix");
   case Node::Kind::FunctionSignatureSpecializationParamPayload: {
-    std::string demangledName = demangleSymbolAsString(Node->getText());
-    if (demangledName.empty()) {
-      Printer << Node->getText();
-    } else {
-      Printer << demangledName;
+    if (Node->hasText()) {
+      std::string demangledName = demangleSymbolAsString(Node->getText());
+      if (demangledName.empty()) {
+        Printer << Node->getText();
+      } else {
+        Printer << demangledName;
+      }
+    } else if (Node->hasIndex()) {
+      Printer << Node->getIndex();
     }
     return nullptr;
   }
@@ -1971,6 +1984,9 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
     case FunctionSigSpecializationParamKind::ClosureProp:
       Printer << "Closure Propagated";
       return nullptr;
+    case FunctionSigSpecializationParamKind::ClosurePropPreviousArg:
+      Printer << "Same As Argument";
+      return nullptr;
     case FunctionSigSpecializationParamKind::ExistentialToGeneric:
     case FunctionSigSpecializationParamKind::Dead:
     case FunctionSigSpecializationParamKind::OwnedToGuaranteed:

lib/Demangling/Remangler.cpp

@@ -1602,6 +1602,9 @@ Remangler::mangleFunctionSignatureSpecializationParam(Node *node,
     case FunctionSigSpecializationParamKind::ClosureProp:
       Buffer << 'c';
       break;
+    case FunctionSigSpecializationParamKind::ClosurePropPreviousArg:
+      Buffer << 'C' << node->getChild(idx++)->getIndex();
+      break;
     case FunctionSigSpecializationParamKind::BoxToValue:
       Buffer << 'i';
       break;

lib/SILOptimizer/Utils/OptimizerBridging.cpp

@@ -274,32 +274,21 @@ BridgedOwnedString BridgedPassContext::mangleWithDeadArgs(BridgedArrayRef bridge
 }
 
 BridgedOwnedString BridgedPassContext::mangleWithClosureArgs(
-  BridgedArrayRef bridgedClosureArgs, BridgedFunction applySiteCallee
+  BridgedArrayRef closureArgManglings, BridgedFunction applySiteCallee
 ) const {
-
-  struct ClosureArgElement {
-    SwiftInt argIdx;
-    BridgeValueExistential argValue;
-  };
-
   auto pass = Demangle::SpecializationPass::ClosureSpecializer;
   auto serializedKind = applySiteCallee.getFunction()->getSerializedKind();
   Mangle::FunctionSignatureSpecializationMangler mangler(applySiteCallee.getFunction()->getASTContext(),
       pass, serializedKind, applySiteCallee.getFunction());
 
-  auto closureArgs = bridgedClosureArgs.unbridged<ClosureArgElement>();
-
-  for (ClosureArgElement argElmt : closureArgs) {
-    auto closureArg = argElmt.argValue.value.getSILValue();
-    auto closureArgIndex = argElmt.argIdx;
+  auto closureArgs = closureArgManglings.unbridged<ClosureArgMangling>();
 
-    if (auto *PAI = dyn_cast<PartialApplyInst>(closureArg)) {
-      mangler.setArgumentClosureProp(closureArgIndex,
-                                     const_cast<PartialApplyInst *>(PAI));
+  for (ClosureArgMangling argElmt : closureArgs) {
+    auto closureArgIndex = (unsigned)argElmt.argIdx;
+    if (SILInstruction *inst = argElmt.inst.unbridged()) {
+      mangler.setArgumentClosureProp(closureArgIndex, inst);
     } else {
-      auto *TTTFI = cast<ThinToThickFunctionInst>(closureArg);
-      mangler.setArgumentClosureProp(closureArgIndex,
-                                     const_cast<ThinToThickFunctionInst *>(TTTFI));
+      mangler.setArgumentClosurePropPreviousArg(closureArgIndex, argElmt.otherArgIdx);
     }
   }