ClosureSpecialization: support for OSSA and a big overhaul

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ConstantCapturePropagation.swift

@@ -112,7 +112,7 @@ private func constantPropagateCaptures(of partialApply: PartialApplyInst, _ cont
     // Escaping closures consume their arguments. Therefore we need to destroy the removed argument values.
     addCompensatingDestroys(for: constArgs, context)
   }
-  let newArguments = nonConstArgs.map { $0.value }
+  let newArguments = Array(nonConstArgs.values)
   rewritePartialApply(partialApply, withSpecialized: specializedCallee, arguments: newArguments, context)
 }
 
@@ -195,8 +195,7 @@ private func cloneArgument(_ argumentOp: Operand,
   if partialApply.calleeArgumentConventions[calleeArgIdx].isGuaranteed {
     // If the original argument was passed as guaranteed, i.e. is _not_ destroyed in the closure, we have
     // to destroy the cloned argument at function exits.
-    for block in targetFunction.blocks where block.terminator.isFunctionExiting {
-      let builder = Builder(before: block.terminator, context)
+    Builder.insertCleanupAtFunctionExits(of: targetFunction, context) { builder in
       builder.emitDestroy(of: clonedArg)
     }
   }

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/InitializeStaticGlobals.swift

@@ -404,7 +404,7 @@ private extension Value {
       case is LoadInst:
         return true
       case is StructInst, is TupleInst:
-        worklist.pushIfNotVisited(contentsOf: (v as! Instruction).operands.lazy.map { $0.value })
+        worklist.pushIfNotVisited(contentsOf: (v as! Instruction).operands.values)
       case let ei as EnumInst:
         if let payload = ei.payload {
           worklist.pushIfNotVisited(payload)

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyCondBranch.swift

@@ -27,9 +27,9 @@ private extension CondBranchInst {
     }
     let builder = Builder(before: self, context)
     if conditionValue == 0 {
-      builder.createBranch(to: falseBlock, arguments: Array(falseOperands.map { $0.value }))
+      builder.createBranch(to: falseBlock, arguments: Array(falseOperands.values))
     } else {
-      builder.createBranch(to: trueBlock, arguments: Array(trueOperands.map { $0.value }))
+      builder.createBranch(to: trueBlock, arguments: Array(trueOperands.values))
     }
     context.erase(instruction: self)
   }

SwiftCompilerSources/Sources/Optimizer/ModulePasses/StackProtection.swift

@@ -339,13 +339,9 @@ private struct StackProtectionOptimization {
 
     builder.createCopyAddr(from: argument, to: temporary, takeSource: true, initializeDest: true)
 
-    for block in function.blocks {
-      let terminator = block.terminator
-      if terminator.isFunctionExiting {
-        let exitBuilder = Builder(before: terminator, location: terminator.location.asAutoGenerated, context)
-        exitBuilder.createCopyAddr(from: temporary, to: argument, takeSource: true, initializeDest: true)
-        exitBuilder.createDeallocStack(temporary)
-      }
+    Builder.insertCleanupAtFunctionExits(of: function, context) { builder in
+      builder.createCopyAddr(from: temporary, to: argument, takeSource: true, initializeDest: true)
+      builder.createDeallocStack(temporary)
     }
     log("move addr protection in \(function.name): \(argument)")
 

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))
     }
   }

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -105,7 +105,7 @@ private func registerSwiftPasses() {
   registerPass(tempRValueElimination, { tempRValueElimination.run($0) })
   registerPass(mandatoryTempRValueElimination, { mandatoryTempRValueElimination.run($0) })
   registerPass(tempLValueElimination, { tempLValueElimination.run($0) })
-  registerPass(generalClosureSpecialization, { generalClosureSpecialization.run($0) })
+  registerPass(closureSpecialization, { closureSpecialization.run($0) })
   registerPass(autodiffClosureSpecialization, { autodiffClosureSpecialization.run($0) })
   registerPass(loopInvariantCodeMotionPass, { loopInvariantCodeMotionPass.run($0) })
 

SwiftCompilerSources/Sources/Optimizer/Utilities/FunctionTest.swift

@@ -43,7 +43,6 @@ public func registerOptimizerTests() {
   registerFunctionTests(
     addressOwnershipLiveRangeTest,
     argumentConventionsTest,
-    getPullbackClosureInfoTest,
     interiorLivenessTest,
     lifetimeDependenceRootTest,
     lifetimeDependenceScopeTest,
@@ -52,8 +51,6 @@ public func registerOptimizerTests() {
     localVariableReachableUsesTest,
     localVariableReachingAssignmentsTest,
     rangeOverlapsPathTest,
-    rewrittenCallerBodyTest,
-    specializedFunctionSignatureAndBodyTest,
     variableIntroducerTest
   )
 

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -71,7 +71,7 @@ extension Value {
       }
       switch v {
       case let fw as ForwardingInstruction:
-        worklist.pushIfNotVisited(contentsOf: fw.definedOperands.lazy.map { $0.value })
+        worklist.pushIfNotVisited(contentsOf: fw.definedOperands.values)
       case let bf as BorrowedFromInst:
         worklist.pushIfNotVisited(bf.borrowedValue)
       case let bb as BeginBorrowInst:
@@ -82,7 +82,7 @@ extension Value {
         } else if let termResult = TerminatorResult(arg),
                let fw = termResult.terminator as? ForwardingInstruction
         {
-          worklist.pushIfNotVisited(contentsOf: fw.definedOperands.lazy.map { $0.value })
+          worklist.pushIfNotVisited(contentsOf: fw.definedOperands.values)
         }
       default:
         continue
@@ -262,6 +262,18 @@ extension Builder {
     }
   }
 
+  static func insertCleanupAtFunctionExits(
+    of function: Function,
+    _ context: some MutatingContext,
+    insert: (Builder) -> ()
+  ) {
+    for exitBlock in function.blocks where exitBlock.terminator.isFunctionExiting {
+      let terminator = exitBlock.terminator
+      let builder = Builder(before: terminator, location: terminator.location.asCleanup, context)
+      insert(builder)
+    }
+  }
+
   func destroyCapturedArgs(for paiOnStack: PartialApplyInst) {
     precondition(paiOnStack.isOnStack, "Function must only be called for `partial_apply`s on stack!")
     self.bridged.destroyCapturedArgs(paiOnStack.bridged)

SwiftCompilerSources/Sources/SIL/Argument.swift

@@ -170,7 +170,7 @@ public struct Phi {
   }
 
   public var incomingValues: LazyMapSequence<LazyMapSequence<PredecessorList, Operand>, Value> {
-    incomingOperands.lazy.map { $0.value }
+    incomingOperands.values
   }
 
   public var isReborrow: Bool { value.isReborrow }

SwiftCompilerSources/Sources/SIL/Builder.swift

@@ -289,6 +289,11 @@ public struct Builder {
     return notifyNew(cast.getAs(UncheckedAddrCastInst.self))
   }
 
+  public func createUncheckedValueCast(from value: Value, to type: Type) -> UncheckedValueCastInst {
+    let cast = bridged.createUncheckedValueCast(value.bridged, type.bridged)
+    return notifyNew(cast.getAs(UncheckedValueCastInst.self))
+  }
+
   public func createUpcast(from value: Value, to type: Type) -> UpcastInst {
     let cast = bridged.createUpcast(value.bridged, type.bridged)
     return notifyNew(cast.getAs(UpcastInst.self))

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -788,7 +788,9 @@ final public class UncheckedTrivialBitCastInst : SingleValueInstruction, UnaryIn
 }
 
 final public class UncheckedBitwiseCastInst : SingleValueInstruction, UnaryInstruction {}
-final public class UncheckedValueCastInst : SingleValueInstruction, UnaryInstruction {}
+final public class UncheckedValueCastInst : SingleValueInstruction, UnaryInstruction {
+  public var fromValue: Value { operand.value }
+}
 
 final public class RefToRawPointerInst : SingleValueInstruction, UnaryInstruction {}
 final public class RefToUnmanagedInst : SingleValueInstruction, UnaryInstruction {}
@@ -1719,6 +1721,9 @@ final public class BeginApplyInst : MultipleValueInstruction, FullApplySite {
   public var yieldedValues: Results {
     Results(inst: self, numResults: resultCount - (isCalleeAllocated ? 2 : 1))
   }
+
+  public var isNonThrowing: Bool { bridged.BeginApplyInst_getNonThrowing() }
+  public var isNonAsync: Bool { bridged.BeginApplyInst_getNonAsync() }
 }
 
 final public class EndApplyInst : SingleValueInstruction, UnaryInstruction {

SwiftCompilerSources/Sources/SIL/Operand.swift

@@ -104,10 +104,6 @@ public struct OperandArray : RandomAccessCollection, CustomReflectable {
       base: OptionalBridgedOperand(op: base.advancedBy(bounds.lowerBound).op),
       count: bounds.upperBound - bounds.lowerBound)
   }
-
-  public var values: LazyMapSequence<LazySequence<OperandArray>.Elements, Value> {
-    self.lazy.map { $0.value }
-  }
 }
 
 public struct UseList : CollectionLikeSequence {
@@ -143,6 +139,10 @@ public struct UseList : CollectionLikeSequence {
 }
 
 extension Sequence where Element == Operand {
+  public var values: LazyMapSequence<Self, Value> {
+    self.lazy.map { $0.value }
+  }
+
   public var singleUse: Operand? {
     var result: Operand? = nil
     for op in self {

SwiftCompilerSources/Sources/SIL/Utilities/BorrowUtils.swift

@@ -534,7 +534,7 @@ public final class EnclosingValueIterator : IteratorProtocol {
         } else if let forwardingInst = value.forwardingInstruction {
           // Recurse through guaranteed forwarding non-phi instructions.
           let ops = forwardingInst.forwardedOperands
-          worklist.pushIfNotVisited(contentsOf: ops.lazy.map { $0.value })
+          worklist.pushIfNotVisited(contentsOf: ops.values)
         } else if value.isGuaranteedApplyResult {
           let selfArgument = (value as! ApplyInst).arguments.last!
           worklist.pushIfNotVisited(selfArgument)

SwiftCompilerSources/Sources/SIL/Utilities/Cloner.swift

@@ -67,7 +67,9 @@ public struct Cloner<Context: MutatingContext> {
     if let entryBlock = targetFunction.blocks.first {
       return entryBlock
     }
-    return targetFunction.appendNewBlock(context)
+    let entryBlock = targetFunction.appendNewBlock(context)
+    bridged.setInsertionBlockIfNotSet(entryBlock.bridged)
+    return entryBlock
   }
 
   public func cloneFunctionBody(from originalFunction: Function, entryBlockArguments: [Value]) {

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/AST/SILOptions.h

@@ -303,10 +303,6 @@ class SILOptions {
   /// Are we building in embedded Swift + -no-allocations?
   bool NoAllocations = false;
 
-  /// Should we use the experimental Swift based closure-specialization
-  /// optimization pass instead of the existing C++ one.
-  bool EnableExperimentalSwiftBasedClosureSpecialization = false;
-
   /// The name of the file to which the backend should save optimization
   /// records.
   std::string OptRecordFile;

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/Option/FrontendOptions.td

@@ -393,10 +393,6 @@ def public_autolink_library :
 // HIDDEN FLAGS
 let Flags = [FrontendOption, NoDriverOption, HelpHidden] in {
 
-def enable_experimental_swift_based_closure_specialization :
-  Flag<["-"], "experimental-swift-based-closure-specialization">,
-  HelpText<"Use the experimental Swift based closure-specialization optimization pass instead of the existing C++ one">;
-
 def checked_async_objc_bridging : Joined<["-"], "checked-async-objc-bridging=">,
   HelpText<"Control whether checked continuations are used when bridging "
            "async calls from Swift to ObjC: 'on', 'off' ">;

include/swift/SIL/SILBridging.h

@@ -243,8 +243,6 @@ BridgedParameterInfoArray SILFunctionType_getParameters(BridgedCanType);
 
 BRIDGED_INLINE bool SILFunctionType_hasSelfParam(BridgedCanType);
 
-BRIDGED_INLINE bool SILFunctionType_isTrivialNoescape(BridgedCanType);
-
 SWIFT_IMPORT_UNSAFE BRIDGED_INLINE
 BridgedYieldInfoArray SILFunctionType_getYields(BridgedCanType);
 
@@ -816,6 +814,8 @@ struct BridgedInstruction {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedGenericSpecializationInformation ApplyInst_getSpecializationInfo() const;
   BRIDGED_INLINE bool TryApplyInst_getNonAsync() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedGenericSpecializationInformation TryApplyInst_getSpecializationInfo() const;
+  BRIDGED_INLINE bool BeginApplyInst_getNonThrowing() const;
+  BRIDGED_INLINE bool BeginApplyInst_getNonAsync() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclRef ClassMethodInst_getMember() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclRef WitnessMethodInst_getMember() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedCanType WitnessMethodInst_getLookupType() const;
@@ -1229,6 +1229,8 @@ struct BridgedBuilder{
                                                                                BridgedType type) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUncheckedAddrCast(BridgedValue op,
                                                                                 BridgedType type) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUncheckedValueCast(BridgedValue op,
+                                                                                 BridgedType type) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUpcast(BridgedValue op, BridgedType type) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createCheckedCastAddrBranch(
       BridgedValue source, BridgedCanType sourceFormalType,
@@ -1545,6 +1547,7 @@ struct BridgedCloner {
   void recordClonedInstruction(BridgedInstruction origInst, BridgedInstruction clonedInst) const;
   void recordFoldedValue(BridgedValue orig, BridgedValue mapped) const;
   BridgedInstruction clone(BridgedInstruction inst) const;
+  void setInsertionBlockIfNotSet(BridgedBasicBlock block) const;
 };
 
 struct BridgedTypeSubstCloner {

include/swift/SIL/SILBridgingImpl.h

@@ -260,10 +260,6 @@ bool SILFunctionType_hasSelfParam(BridgedCanType funcTy) {
   return funcTy.unbridged()->castTo<swift::SILFunctionType>()->hasSelfParam();
 }
 
-bool SILFunctionType_isTrivialNoescape(BridgedCanType funcTy) {
-  return funcTy.unbridged()->castTo<swift::SILFunctionType>()->isTrivialNoEscape();
-}
-
 BridgedYieldInfoArray SILFunctionType_getYields(BridgedCanType funcTy) {
   return {funcTy.unbridged()->castTo<swift::SILFunctionType>()->getYields()};
 }
@@ -1349,6 +1345,14 @@ BridgedGenericSpecializationInformation BridgedInstruction::TryApplyInst_getSpec
   return {getAs<swift::TryApplyInst>()->getSpecializationInfo()};
 }
 
+bool BridgedInstruction::BeginApplyInst_getNonThrowing() const {
+  return getAs<swift::BeginApplyInst>()->isNonThrowing();
+}
+
+bool BridgedInstruction::BeginApplyInst_getNonAsync() const {
+  return getAs<swift::BeginApplyInst>()->isNonAsync();
+}
+
 BridgedDeclRef BridgedInstruction::ClassMethodInst_getMember() const {
   return getAs<swift::ClassMethodInst>()->getMember();
 }
@@ -2301,6 +2305,11 @@ BridgedInstruction BridgedBuilder::createUncheckedAddrCast(BridgedValue op, Brid
                                               type.unbridged())};
 }
 
+BridgedInstruction BridgedBuilder::createUncheckedValueCast(BridgedValue op, BridgedType type) const {
+  return {unbridged().createUncheckedValueCast(regularLoc(), op.getSILValue(),
+                                              type.unbridged())};
+}
+
 BridgedInstruction BridgedBuilder::createUpcast(BridgedValue op, BridgedType type) const {
   return {unbridged().createUpcast(regularLoc(), op.getSILValue(),
                                    type.unbridged())};