Add new generalized cloner.

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -311,10 +311,9 @@ private struct FunctionSpecializations {
       // This can happen if a previous run of the pass already created this specialization.
       return
     }
-    let cloner = SpecializationCloner(emptySpecializedFunction: specializedFunc, context)
-    cloner.cloneFunctionBody(from: original)
-
     context.buildSpecializedFunction(specializedFunction: specializedFunc) { (specializedFunc, specContext) in
+      cloneFunction(from: original, toEmpty: specializedFunc, specContext)
+
       replaceBoxWithStackArguments(in: specializedFunc, original: original, specContext)
     }
     context.notifyNewFunction(function: specializedFunc, derivedFrom: original)

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -247,9 +247,10 @@ private func getOrCreateSpecializedFunction(
   context.buildSpecializedFunction(
     specializedFunction: specializedFunction,
     buildFn: { (emptySpecializedFunction, functionPassContext) in
-      let closureSpecCloner = SpecializationCloner(
-        emptySpecializedFunction: emptySpecializedFunction, functionPassContext)
+      var closureSpecCloner = Cloner(
+        cloneToEmptyFunction: emptySpecializedFunction, functionPassContext)
       closureSpecCloner.cloneAndSpecializeFunctionBody(using: pullbackClosureInfo)
+      closureSpecCloner.deinitialize()
     })
 
   return (specializedFunction, false)
@@ -731,15 +732,14 @@ private func markConvertedAndReabstractedClosuresAsUsed(
   }
 }
 
-extension SpecializationCloner {
+extension Cloner where Context == FunctionPassContext {
   fileprivate func cloneAndSpecializeFunctionBody(using pullbackClosureInfo: PullbackClosureInfo) {
     self.cloneEntryBlockArgsWithoutOrigClosures(usingOrigCalleeAt: pullbackClosureInfo)
 
     let (allSpecializedEntryBlockArgs, closureArgIndexToAllClonedReleasableClosures) =
       cloneAllClosures(at: pullbackClosureInfo)
 
-    self.cloneFunctionBody(
-      from: pullbackClosureInfo.pullbackFn, entryBlockArguments: allSpecializedEntryBlockArgs)
+    self.cloneFunctionBody(from: pullbackClosureInfo.pullbackFn, entryBlockArguments: allSpecializedEntryBlockArgs)
 
     self.insertCleanupCodeForClonedReleasableClosures(
       from: pullbackClosureInfo,
@@ -750,8 +750,8 @@ extension SpecializationCloner {
     usingOrigCalleeAt pullbackClosureInfo: PullbackClosureInfo
   ) {
     let originalEntryBlock = pullbackClosureInfo.pullbackFn.entryBlock
-    let clonedFunction = self.cloned
-    let clonedEntryBlock = self.entryBlock
+    let clonedFunction = self.targetFunction
+    let clonedEntryBlock = self.getOrCreateEntryBlock()
 
     originalEntryBlock.arguments
       .enumerated()
@@ -782,7 +782,8 @@ extension SpecializationCloner {
     )
   {
     func entryBlockArgsWithOrigClosuresSkipped() -> [Value?] {
-      var clonedNonClosureEntryBlockArgs = self.entryBlock.arguments.makeIterator()
+      let clonedEntryBlock = self.getOrCreateEntryBlock()
+      var clonedNonClosureEntryBlockArgs = clonedEntryBlock.arguments.makeIterator()
 
       return pullbackClosureInfo.pullbackFn
         .entryBlock
@@ -823,8 +824,8 @@ extension SpecializationCloner {
   {
     let (origToClonedValueMap, capturedArgRange) = self.addEntryBlockArgs(
       forValuesCapturedBy: closureArgDesc)
-    let clonedFunction = self.cloned
-    let clonedEntryBlock = self.entryBlock
+    let clonedFunction = self.targetFunction
+    let clonedEntryBlock = self.getOrCreateEntryBlock()
     let clonedClosureArgs = Array(clonedEntryBlock.arguments[capturedArgRange])
 
     let builder =
@@ -853,8 +854,8 @@ extension SpecializationCloner {
     -> (origToClonedValueMap: [HashableValue: Value], capturedArgRange: Range<Int>)
   {
     var origToClonedValueMap: [HashableValue: Value] = [:]
-    let clonedFunction = self.cloned
-    let clonedEntryBlock = self.entryBlock
+    let clonedFunction = self.targetFunction
+    let clonedEntryBlock = self.getOrCreateEntryBlock()
 
     let capturedArgRangeStart = clonedEntryBlock.arguments.count
 
@@ -908,8 +909,8 @@ extension SpecializationCloner {
       }
     }
 
-    if self.context.needFixStackNesting {
-      self.context.fixStackNesting(in: self.cloned)
+    if (self.context.needFixStackNesting) {
+      self.context.fixStackNesting(in: targetFunction)
     }
   }
 }
@@ -1425,12 +1426,10 @@ private struct PullbackClosureInfo {
   }
 
   func specializedCalleeName(_ context: FunctionPassContext) -> String {
-    let closureArgs = Array(self.closureArgDescriptors.map { $0.closure })
-    let closureIndices = Array(self.closureArgDescriptors.map { $0.closureArgIndex })
-
-    return context.mangle(
-      withClosureArguments: closureArgs, closureArgIndices: closureIndices,
-      from: pullbackFn)
+    let closureArgs = Array(self.closureArgDescriptors.map {
+      (argumentIndex: $0.closureArgIndex, argumentValue: $0.closure)
+    })
+    return context.mangle(withClosureArguments: closureArgs, from: pullbackFn)
   }
 }
 

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/InitializeStaticGlobals.swift

@@ -271,7 +271,7 @@ private indirect enum GlobalInitValue {
 
   /// Creates SIL for this global init value in the initializer of the `global`.
   func materialize(into global: GlobalVariable, from function: Function, _ context: FunctionPassContext) {
-    var cloner = StaticInitCloner(cloneTo: global, context)
+    var cloner = Cloner(cloneToGlobal: global, context)
     defer { cloner.deinitialize() }
     let builder = Builder(staticInitializerOf: global, context)
 
@@ -280,7 +280,7 @@ private indirect enum GlobalInitValue {
 
   private func materializeRecursively(
     type: Type,
-    _ cloner: inout StaticInitCloner<FunctionPassContext>,
+    _ cloner: inout Cloner<FunctionPassContext>,
     _ builder: Builder,
     _ function: Function
   ) -> Value {
@@ -289,7 +289,7 @@ private indirect enum GlobalInitValue {
       fatalError("cannot materialize undefined init value")
 
     case .constant(let value):
-      return cloner.clone(value)
+      return cloner.cloneRecursivelyToGlobal(value: value)
 
     case .aggregate(let fields):
       if type.isStruct {

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ObjectOutliner.swift

@@ -364,13 +364,13 @@ private func constructObject(of allocRef: AllocRefInstBase,
                              inInitializerOf global: GlobalVariable,
                              _ storesToClassFields: [StoreInst], _ storesToTailElements: [StoreInst],
                              _ context: FunctionPassContext) {
-  var cloner = StaticInitCloner(cloneTo: global, context)
+  var cloner = Cloner(cloneToGlobal: global, context)
   defer { cloner.deinitialize() }
 
   // Create the initializers for the fields
   var objectArgs = [Value]()
   for store in storesToClassFields {
-    objectArgs.append(cloner.clone(store.source as! SingleValueInstruction))
+    objectArgs.append(cloner.cloneRecursivelyToGlobal(value: store.source as! SingleValueInstruction))
   }
   let globalBuilder = Builder(staticInitializerOf: global, context)
 
@@ -382,7 +382,7 @@ private func constructObject(of allocRef: AllocRefInstBase,
       for elementIdx in 0..<allocRef.numTailElements! {
         let tupleElems = (0..<numTailTupleElems).map { tupleIdx in
             let store = storesToTailElements[elementIdx * numTailTupleElems + tupleIdx]
-            return cloner.clone(store.source as! SingleValueInstruction)
+            return cloner.cloneRecursivelyToGlobal(value: store.source as! SingleValueInstruction)
         }
         let tuple = globalBuilder.createTuple(type: allocRef.tailAllocatedTypes[0], elements: tupleElems)
         objectArgs.append(tuple)
@@ -390,7 +390,7 @@ private func constructObject(of allocRef: AllocRefInstBase,
     } else {
       // The non-tuple element case.
       for store in storesToTailElements {
-        objectArgs.append(cloner.clone(store.source as! SingleValueInstruction))
+        objectArgs.append(cloner.cloneRecursivelyToGlobal(value: store.source as! SingleValueInstruction))
       }
     }
   }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyLoad.swift

@@ -104,10 +104,10 @@ extension LoadInst : OnoneSimplifiable, SILCombineSimplifiable {
     if !globalInitVal.canBeCopied(into: parentFunction, context) {
       return false
     }
-    var cloner = StaticInitCloner(cloneBefore: self, context)
+    var cloner = Cloner(cloneBefore: self, context)
     defer { cloner.deinitialize() }
 
-    let initVal = cloner.clone(globalInitVal)
+    let initVal = cloner.cloneRecursivelyToGlobal(value: globalInitVal)
 
     uses.replaceAll(with: initVal, context)
     // Also erases a builtin "once" on which the global_addr depends on. This is fine
@@ -323,7 +323,7 @@ private extension Value {
         return false
       }
       if let fri = value as? FunctionRefInst {
-        if function.isAnySerialized, 
+        if function.isAnySerialized,
            !fri.referencedFunction.hasValidLinkageForFragileRef(function.serializedKind)
         {
           return false
@@ -447,4 +447,3 @@ private func getGlobalInitialization(
   }
   return nil
 }
-

SwiftCompilerSources/Sources/Optimizer/PassManager/ContextCommon.swift

@@ -152,3 +152,28 @@ extension MutatingContext {
     bridgedPassContext.notifyDependencyOnBodyOf(otherFunction.bridged)
   }
 }
+extension Cloner where Context == FunctionPassContext {
+  func getOrCreateEntryBlock() -> BasicBlock {
+    if let entryBlock = targetFunction.blocks.first {
+      return entryBlock
+    }
+    return targetFunction.appendNewBlock(context)
+  }
+
+  func cloneFunctionBody(from originalFunction: Function, entryBlockArguments: [Value]) {
+    entryBlockArguments.withBridgedValues { bridgedEntryBlockArgs in
+      let entryBlock = getOrCreateEntryBlock()
+      bridged.cloneFunctionBody(originalFunction.bridged, entryBlock.bridged, bridgedEntryBlockArgs)
+    }
+  }
+
+  func cloneFunctionBody(from originalFunction: Function) {
+    bridged.cloneFunctionBody(originalFunction.bridged)
+  }
+}
+
+func cloneFunction(from originalFunction: Function, toEmpty targetFunction: Function, _ context: FunctionPassContext) {
+  var cloner = Cloner(cloneToEmptyFunction: targetFunction, context)
+  defer { cloner.deinitialize() }
+  cloner.cloneFunctionBody(from: originalFunction)
+}

SwiftCompilerSources/Sources/Optimizer/PassManager/FunctionPassContext.swift

@@ -92,15 +92,11 @@ struct FunctionPassContext : MutatingContext {
     return String(taking: bridgedPassContext.mangleOutlinedVariable(function.bridged))
   }
 
-  func mangle(withClosureArguments closureArgs: [Value], closureArgIndices: [Int], from applySiteCallee: Function) -> String {
-    closureArgs.withBridgedValues { bridgedClosureArgsRef in
-      closureArgIndices.withBridgedArrayRef{bridgedClosureArgIndicesRef in
-        String(taking: bridgedPassContext.mangleWithClosureArgs(
-          bridgedClosureArgsRef,
-          bridgedClosureArgIndicesRef,
-          applySiteCallee.bridged
-        ))
-      }
+  func mangle(withClosureArguments closureArgs: [(argumentIndex: Int, argumentValue: Value)],
+              from applySiteCallee: Function
+  ) -> String {
+    closureArgs.withBridgedArrayRef{ bridgedClosureArgs in
+      String(taking: bridgedPassContext.mangleWithClosureArgs(bridgedClosureArgs, applySiteCallee.bridged))
     }
   }
 

SwiftCompilerSources/Sources/Optimizer/Utilities/CMakeLists.txt

@@ -8,7 +8,6 @@
 
 swift_compiler_sources(Optimizer
   AddressUtils.swift
-  SpecializationCloner.swift
   Devirtualization.swift
   EscapeUtils.swift
   FunctionSignatureTransforms.swift
@@ -18,5 +17,4 @@ swift_compiler_sources(Optimizer
   LocalVariableUtils.swift
   OptUtils.swift
   OwnershipLiveness.swift
-  StaticInitCloner.swift
 )

SwiftCompilerSources/Sources/Optimizer/Utilities/SpecializationCloner.swift

@@ -1,52 +0,0 @@
-//===--- SpecializationCloner.swift --------------------------------------------==//
-//
-// This source file is part of the Swift.org open source project
-//
-// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
-// Licensed under Apache License v2.0 with Runtime Library Exception
-//
-// See https://swift.org/LICENSE.txt for license information
-// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
-//
-//===----------------------------------------------------------------------===//
-
-import OptimizerBridging
-import SIL
-
-/// Utility cloner type that can be used by optimizations that generate new functions or specialized versions of
-/// existing functions. 
-struct SpecializationCloner {
-  private let bridged: BridgedSpecializationCloner
-  let context: FunctionPassContext
-
-  init(emptySpecializedFunction: Function, _ context: FunctionPassContext) {
-    self.bridged = BridgedSpecializationCloner(emptySpecializedFunction.bridged)
-    self.context = context
-  }
-  
-  var cloned: Function {
-    bridged.getCloned().function
-  }
-
-  var entryBlock: BasicBlock {
-    if cloned.blocks.isEmpty {
-      return cloned.appendNewBlock(context)
-    } else {
-      return cloned.entryBlock
-    }
-  }
-
-  func getClonedBlock(for originalBlock: BasicBlock) -> BasicBlock {
-    bridged.getClonedBasicBlock(originalBlock.bridged).block
-  }
-
-  func cloneFunctionBody(from originalFunction: Function, entryBlockArguments: [Value]) {
-    entryBlockArguments.withBridgedValues { bridgedEntryBlockArgs in
-      bridged.cloneFunctionBody(originalFunction.bridged, self.entryBlock.bridged, bridgedEntryBlockArgs)
-    }
-  }
-
-  func cloneFunctionBody(from originalFunction: Function) {
-    bridged.cloneFunctionBody(originalFunction.bridged)
-  }
-}

SwiftCompilerSources/Sources/Optimizer/Utilities/StaticInitCloner.swift

@@ -1,78 +0,0 @@
-//===--- StaticInitCloner.swift --------------------------------------------==//
-//
-// This source file is part of the Swift.org open source project
-//
-// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
-// Licensed under Apache License v2.0 with Runtime Library Exception
-//
-// See https://swift.org/LICENSE.txt for license information
-// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
-//
-//===----------------------------------------------------------------------===//
-
-import SIL
-import OptimizerBridging
-
-/// Clones the initializer value of a GlobalVariable.
-///
-/// Used to transitively clone "constant" instructions, including their operands,
-/// from or to the static initializer value of a GlobalVariable.
-///
-struct StaticInitCloner<Context: MutatingContext> {
-  private var bridged: BridgedCloner
-  private let context: Context
-  private let cloningIntoFunction: Bool
-
-  init(cloneTo global: GlobalVariable, _ context: Context) {
-    self.bridged = BridgedCloner(global.bridged, context._bridged)
-    self.context = context
-    self.cloningIntoFunction = false
-  }
-
-  init(cloneBefore inst: Instruction, _ context: Context) {
-    self.bridged = BridgedCloner(inst.bridged, context._bridged)
-    self.context = context
-    self.cloningIntoFunction = true
-  }
-
-  mutating func deinitialize() {
-    bridged.destroy(context._bridged)
-  }
-
-  /// Transitively clones `value` including its defining instruction's operands.
-  mutating func clone(_ value: Value) -> Value {
-
-    if isCloned(value: value) {
-      return getClonedValue(of: value)
-    }
-
-    if let beginAccess = value as? BeginAccessInst {
-      // Skip access instructions, which might be generated for UnsafePointer globals which point to other globals.
-      let clonedOperand = clone(beginAccess.address)
-      bridged.recordFoldedValue(beginAccess.bridged, clonedOperand.bridged)
-      return clonedOperand
-    }
-
-    let inst = value.definingInstruction!
-    assert(!(inst is ScopedInstruction), "global init value must not contain a scoped instruction")
-
-    for op in inst.operands {
-      _ = clone(op.value)
-    }
-
-    bridged.clone(inst.bridged)
-    let clonedValue = getClonedValue(of: value)
-    if cloningIntoFunction {
-      context.notifyInstructionChanged(clonedValue.definingInstruction!)
-    }
-    return clonedValue
-  }
-
-  mutating func getClonedValue(of value: Value) -> Value {
-    bridged.getClonedValue(value.bridged).value
-  }
-
-  func isCloned(value: Value) -> Bool {
-    bridged.isValueCloned(value.bridged)
-  }
-}

SwiftCompilerSources/Sources/SIL/Utilities/CMakeLists.txt

@@ -9,6 +9,7 @@
 swift_compiler_sources(SIL
   AccessUtils.swift
   BorrowUtils.swift
+  Cloner.swift
   ForwardingUtils.swift
   PhiUpdater.swift
   SequenceUtilities.swift

SwiftCompilerSources/Sources/SIL/Utilities/Cloner.swift

@@ -0,0 +1,141 @@
+//===--- Cloner.swift ------------------------------------------------------==//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import OptimizerBridging
+import SILBridging
+
+/// Clones the initializer value of a GlobalVariable.
+///
+/// Used to transitively clone "constant" instructions, including their operands,
+/// from or to the static initializer value of a GlobalVariable.
+///
+public struct Cloner<Context: MutatingContext> {
+  public var bridged: BridgedCloner
+  public let context: Context
+
+  public enum GetClonedResult {
+    case defaultValue
+    case customValue(Value)
+    case stopCloning
+  }
+  
+  public enum Target {
+    case function(Function)
+    case global(GlobalVariable)
+  }
+  public let target: Target
+
+  public init(cloneToGlobal: GlobalVariable, _ context: Context) {
+    self.bridged = BridgedCloner(cloneToGlobal.bridged, context._bridged)
+    self.context = context
+    self.target = .global(cloneToGlobal)
+  }
+
+  public init(cloneBefore inst: Instruction, _ context: Context) {
+    self.bridged = BridgedCloner(inst.bridged, context._bridged)
+    self.context = context
+    self.target = .function(inst.parentFunction)
+  }
+  
+  public init(cloneToEmptyFunction: Function, _ context: Context) {
+    self.bridged = BridgedCloner(cloneToEmptyFunction.bridged, context._bridged)
+    self.context = context
+    self.target = .function(cloneToEmptyFunction)
+  }
+
+  public mutating func deinitialize() {
+    bridged.destroy(context._bridged)
+  }
+
+  public var targetFunction: Function {
+    guard case .function(let function) = target else {
+      fatalError("expected cloning into a function")
+    }
+    return function
+  }
+
+  public mutating func clone(instruction: Instruction) -> Instruction {
+    let cloned = bridged.clone(instruction.bridged).instruction
+    if case .function = target {
+      context.notifyInstructionChanged(cloned)
+      context.notifyInstructionsChanged()
+    }
+    return cloned
+  }
+  
+  public mutating func cloneRecursivelyToGlobal(value: Value) -> Value {
+    guard let cloned = cloneRecursively(value: value, customGetCloned: { value, cloner in
+      guard let beginAccess = value as? BeginAccessInst else {
+        return .defaultValue
+      }
+      
+      // Skip access instructions, which might be generated for UnsafePointer globals which point to other globals.
+      let clonedOperand = cloner.cloneRecursivelyToGlobal(value: beginAccess.address)
+      cloner.recordFoldedValue(beginAccess, mappedTo: clonedOperand)
+      return .customValue(clonedOperand)
+    }) else {
+      fatalError("Clone recursively to global shouldn't bail.")
+    }
+    
+    return cloned
+  }
+
+  /// Transitively clones `value` including its defining instruction's operands.
+  public mutating func cloneRecursively(value: Value, customGetCloned: (Value, inout Cloner) -> GetClonedResult) -> Value? {
+    if isCloned(value: value) {
+      return getClonedValue(of: value)
+    }
+    
+    switch customGetCloned(value, &self) {
+    case .customValue(let base):
+      return base
+    case .stopCloning:
+      return nil
+    case .defaultValue:
+      break
+    }
+
+    guard let inst = value.definingInstruction else {
+      fatalError("expected instruction to clone or already cloned value")
+    }
+
+    for op in inst.operands {
+      if cloneRecursively(value: op.value, customGetCloned: customGetCloned) == nil {
+        return nil
+      }
+    }
+
+    let cloned = clone(instruction: inst)
+    if let svi = cloned as? SingleValueInstruction {
+      return svi
+    } else if let originalMvi = value as? MultipleValueInstructionResult {
+      return cloned.results[originalMvi.index]
+    }
+    fatalError("unexpected instruction kind")
+  }
+
+  public mutating func getClonedValue(of originalValue: Value) -> Value {
+    bridged.getClonedValue(originalValue.bridged).value
+  }
+
+  public func isCloned(value: Value) -> Bool {
+    bridged.isValueCloned(value.bridged)
+  }
+
+  public func getClonedBlock(for originalBlock: BasicBlock) -> BasicBlock {
+    bridged.getClonedBasicBlock(originalBlock.bridged).block
+  }
+  
+  public func recordFoldedValue(_ origValue: Value, mappedTo mappedValue: Value) {
+    bridged.recordFoldedValue(origValue.bridged, mappedValue.bridged)
+  }
+}

include/swift/SIL/SILBridging.h

@@ -60,6 +60,7 @@ class SymbolicValueBumpAllocator;
 class ConstExprEvaluator;
 class SILWitnessTable;
 class SILDefaultWitnessTable;
+class BridgedClonerImpl;
 class SILDebugLocation;
 class NominalTypeDecl;
 class VarDecl;
@@ -1279,7 +1280,7 @@ struct BridgedBuilder{
                                                                            BridgedArgumentConvention calleeConvention,
                                                                            BridgedSubstitutionMap bridgedSubstitutionMap = BridgedSubstitutionMap(),
                                                                            bool hasUnknownIsolation = true,
-                                                                           bool isOnStack = false) const;                                                                                  
+                                                                           bool isOnStack = false) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createBranch(BridgedBasicBlock destBlock,
                                                                      BridgedValueArray arguments) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUnreachable() const;
@@ -1376,18 +1377,6 @@ struct BridgedOperandSet {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedFunction getFunction() const;
 };
 
-struct BridgedCloner {
-  swift::ClonerWithFixedLocation * _Nonnull cloner;
-
-  BridgedCloner(BridgedGlobalVar var, BridgedContext context);
-  BridgedCloner(BridgedInstruction inst, BridgedContext context);
-  void destroy(BridgedContext context);
-  SWIFT_IMPORT_UNSAFE BridgedValue getClonedValue(BridgedValue v);
-  bool isValueCloned(BridgedValue v) const;
-  void clone(BridgedInstruction inst);
-  void recordFoldedValue(BridgedValue origValue, BridgedValue mappedValue);
-};
-
 struct BridgedContext {
   swift::SILContext * _Nonnull context;
 
@@ -1501,6 +1490,25 @@ struct BridgedContext {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE Slab freeSlab(Slab slab) const;
 };
 
+struct BridgedCloner {
+  swift::BridgedClonerImpl * _Nonnull cloner;
+
+  BridgedCloner(BridgedGlobalVar var, BridgedContext context);
+  BridgedCloner(BridgedInstruction inst, BridgedContext context);
+  BridgedCloner(BridgedFunction emptyFunction, BridgedContext context);
+  void destroy(BridgedContext context);
+  SWIFT_IMPORT_UNSAFE BridgedFunction getCloned() const;
+  SWIFT_IMPORT_UNSAFE BridgedBasicBlock getClonedBasicBlock(BridgedBasicBlock originalBasicBlock) const;
+  void cloneFunctionBody(BridgedFunction originalFunction, BridgedBasicBlock clonedEntryBlock,
+                         BridgedValueArray clonedEntryBlockArgs) const;
+  void cloneFunctionBody(BridgedFunction originalFunction) const;
+  SWIFT_IMPORT_UNSAFE BridgedValue getClonedValue(BridgedValue v);
+  bool isValueCloned(BridgedValue v) const;
+  void recordClonedInstruction(BridgedInstruction origInst, BridgedInstruction clonedInst) const;
+  void recordFoldedValue(BridgedValue orig, BridgedValue mapped) const;
+  BridgedInstruction clone(BridgedInstruction inst);
+};
+
 struct BridgedVerifier {
   typedef void (* _Nonnull VerifyFunctionFn)(BridgedContext, BridgedFunction);
 

lib/SIL/Utils/SILBridging.cpp

@@ -546,59 +546,96 @@ BridgedInstruction BridgedBuilder::createSwitchEnumAddrInst(BridgedValue enumAdd
 }
 
 //===----------------------------------------------------------------------===//
-//                             BridgedCloner
+//                               BridgedCloner
 //===----------------------------------------------------------------------===//
 
 // Need to put ClonerWithFixedLocation into namespace swift to forward reference
-// it in SILBridging.h.
+// it in OptimizerBridging.h.
 namespace swift {
 
-class ClonerWithFixedLocation : public SILCloner<ClonerWithFixedLocation> {
-  friend class SILInstructionVisitor<ClonerWithFixedLocation>;
-  friend class SILCloner<ClonerWithFixedLocation>;
+class BridgedClonerImpl : public SILCloner<BridgedClonerImpl> {
+  friend class SILInstructionVisitor<BridgedClonerImpl>;
+  friend class SILCloner<BridgedClonerImpl>;
 
-  SILDebugLocation insertLoc;
+  bool hasFixedLocation;
+  union {
+    SILDebugLocation fixedLocation;
+    ScopeCloner scopeCloner;
+  };
+
+  SILInstruction *result = nullptr;
 
 public:
-  ClonerWithFixedLocation(SILGlobalVariable *gVar)
-  : SILCloner<ClonerWithFixedLocation>(gVar),
-  insertLoc(ArtificialUnreachableLocation(), nullptr) {}
+  BridgedClonerImpl(SILGlobalVariable *gVar)
+    : SILCloner<BridgedClonerImpl>(gVar),
+      hasFixedLocation(true),
+      fixedLocation(ArtificialUnreachableLocation(), nullptr) {}
 
-  ClonerWithFixedLocation(SILInstruction *insertionPoint)
-  : SILCloner<ClonerWithFixedLocation>(*insertionPoint->getFunction()),
-  insertLoc(insertionPoint->getDebugLocation()) {
+  BridgedClonerImpl(SILInstruction *insertionPoint)
+    : SILCloner<BridgedClonerImpl>(*insertionPoint->getFunction()),
+      hasFixedLocation(true),
+      fixedLocation(insertionPoint->getDebugLocation()) {
     Builder.setInsertionPoint(insertionPoint);
   }
 
+  BridgedClonerImpl(SILFunction &emptyFunction)
+    : SILCloner<BridgedClonerImpl>(emptyFunction),
+      hasFixedLocation(false),
+      scopeCloner(ScopeCloner(emptyFunction)) {}
+
+  ~BridgedClonerImpl() {
+    if (hasFixedLocation) {
+      fixedLocation.~SILDebugLocation();
+    } else {
+      scopeCloner.~ScopeCloner();
+    }
+  }
+
   SILValue getClonedValue(SILValue v) {
     return getMappedValue(v);
   }
 
-  void cloneInst(SILInstruction *inst) {
+  SILInstruction *cloneInst(SILInstruction *inst) {
+    result = nullptr;
     visit(inst);
+    ASSERT(result && "instruction not cloned");
+    return result;
   }
 
-protected:
-
   SILLocation remapLocation(SILLocation loc) {
-    return insertLoc.getLocation();
+    if (hasFixedLocation)
+      return fixedLocation.getLocation();
+    return loc;
   }
 
   const SILDebugScope *remapScope(const SILDebugScope *DS) {
-    return insertLoc.getScope();
+    if (hasFixedLocation)
+      return fixedLocation.getScope();
+    return scopeCloner.getOrCreateClonedScope(DS);
   }
+
+  void postProcess(SILInstruction *Orig, SILInstruction *Cloned) {
+    result = Cloned;
+    SILCloner<BridgedClonerImpl>::postProcess(Orig, Cloned);
+  }
+
 };
 
 } // namespace swift
 
 BridgedCloner::BridgedCloner(BridgedGlobalVar var, BridgedContext context)
-  : cloner(new ClonerWithFixedLocation(var.getGlobal())) {
+  : cloner(new BridgedClonerImpl(var.getGlobal())) {
   context.context->notifyNewCloner();
 }
 
 BridgedCloner::BridgedCloner(BridgedInstruction inst,
                              BridgedContext context)
-    : cloner(new ClonerWithFixedLocation(inst.unbridged())) {
+    : cloner(new BridgedClonerImpl(inst.unbridged())) {
+  context.context->notifyNewCloner();
+}
+
+BridgedCloner::BridgedCloner(BridgedFunction emptyFunction, BridgedContext context)
+  : cloner(new BridgedClonerImpl(*emptyFunction.getFunction())) {
   context.context->notifyNewCloner();
 }
 
@@ -608,6 +645,10 @@ void BridgedCloner::destroy(BridgedContext context) {
   context.context->notifyClonerDestroyed();
 }
 
+BridgedFunction BridgedCloner::getCloned() const {
+  return { &cloner->getBuilder().getFunction() };
+}
+
 BridgedValue BridgedCloner::getClonedValue(BridgedValue v) {
   return {cloner->getClonedValue(v.getSILValue())};
 }
@@ -616,12 +657,32 @@ bool BridgedCloner::isValueCloned(BridgedValue v) const {
   return cloner->isValueCloned(v.getSILValue());
 }
 
-void BridgedCloner::clone(BridgedInstruction inst) {
-  cloner->cloneInst(inst.unbridged());
+void BridgedCloner::recordClonedInstruction(BridgedInstruction origInst, BridgedInstruction clonedInst) const {
+  cloner->recordClonedInstruction(origInst.unbridged(), clonedInst.unbridged());
 }
 
-void BridgedCloner::recordFoldedValue(BridgedValue origValue, BridgedValue mappedValue) {
-  cloner->recordFoldedValue(origValue.getSILValue(), mappedValue.getSILValue());
+void BridgedCloner::recordFoldedValue(BridgedValue orig, BridgedValue mapped) const {
+  cloner->recordFoldedValue(orig.getSILValue(), mapped.getSILValue());
+}
+
+BridgedInstruction BridgedCloner::clone(BridgedInstruction inst) {
+  return {cloner->cloneInst(inst.unbridged())->asSILNode()};
+}
+
+BridgedBasicBlock BridgedCloner::getClonedBasicBlock(BridgedBasicBlock originalBasicBlock) const {
+  return { cloner->getOpBasicBlock(originalBasicBlock.unbridged()) };
+}
+
+void BridgedCloner::cloneFunctionBody(BridgedFunction originalFunction,
+                                      BridgedBasicBlock clonedEntryBlock,
+                                      BridgedValueArray clonedEntryBlockArgs) const {
+  llvm::SmallVector<swift::SILValue, 16> clonedEntryBlockArgsStorage;
+  auto clonedEntryBlockArgsArrayRef = clonedEntryBlockArgs.getValues(clonedEntryBlockArgsStorage);
+  cloner->cloneFunctionBody(originalFunction.getFunction(), clonedEntryBlock.unbridged(), clonedEntryBlockArgsArrayRef);
+}
+
+void BridgedCloner::cloneFunctionBody(BridgedFunction originalFunction) const {
+  cloner->cloneFunction(originalFunction.getFunction());
 }
 
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Utils/OptimizerBridging.cpp

@@ -272,32 +272,31 @@ BridgedOwnedString BridgedPassContext::mangleWithDeadArgs(BridgedArrayRef bridge
 }
 
 BridgedOwnedString BridgedPassContext::mangleWithClosureArgs(
-  BridgedValueArray bridgedClosureArgs,
-  BridgedArrayRef bridgedClosureArgIndices,
-  BridgedFunction applySiteCallee
+  BridgedArrayRef bridgedClosureArgs, 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());
 
-  llvm::SmallVector<swift::SILValue, 16> closureArgsStorage;
-  auto closureArgs = bridgedClosureArgs.getValues(closureArgsStorage);
-  auto closureArgIndices = bridgedClosureArgIndices.unbridged<SwiftInt>();
+  auto closureArgs = bridgedClosureArgs.unbridged<ClosureArgElement>();
 
-  assert(closureArgs.size() == closureArgIndices.size() &&
-         "Number of closures arguments and number of closure indices do not match!");
-
-  for (size_t i = 0; i < closureArgs.size(); i++) {
-    auto closureArg = closureArgs[i];
-    auto closureArgIndex = closureArgIndices[i];
+  for (ClosureArgElement argElmt : closureArgs) {
+    auto closureArg = argElmt.argValue.value.getSILValue();
+    auto closureArgIndex = argElmt.argIdx;
 
     if (auto *PAI = dyn_cast<PartialApplyInst>(closureArg)) {
       mangler.setArgumentClosureProp(closureArgIndex,
                                      const_cast<PartialApplyInst *>(PAI));
     } else {
       auto *TTTFI = cast<ThinToThickFunctionInst>(closureArg);
-      mangler.setArgumentClosureProp(closureArgIndex, 
+      mangler.setArgumentClosureProp(closureArgIndex,
                                      const_cast<ThinToThickFunctionInst *>(TTTFI));
     }
   }
@@ -459,43 +458,12 @@ BridgedCalleeAnalysis::CalleeList BridgedCalleeAnalysis::getDestructors(BridgedT
   return ca->getDestructors(type.unbridged(), isExactType);
 }
 
-namespace swift {
-  class SpecializationCloner: public SILClonerWithScopes<SpecializationCloner> {
-    friend class SILInstructionVisitor<SpecializationCloner>;
-    friend class SILCloner<SpecializationCloner>;
-  public:
-    using SuperTy = SILClonerWithScopes<SpecializationCloner>;
-    SpecializationCloner(SILFunction &emptySpecializedFunction): SuperTy(emptySpecializedFunction) {}
-  };
-} // namespace swift
-
-BridgedSpecializationCloner::BridgedSpecializationCloner(BridgedFunction emptySpecializedFunction): 
-  cloner(new SpecializationCloner(*emptySpecializedFunction.getFunction())) {}
-
-BridgedFunction BridgedSpecializationCloner::getCloned() const {
-  return { &cloner->getBuilder().getFunction() };
-}
-
-BridgedBasicBlock BridgedSpecializationCloner::getClonedBasicBlock(BridgedBasicBlock originalBasicBlock) const {
-  return { cloner->getOpBasicBlock(originalBasicBlock.unbridged()) };
-}
-
-void BridgedSpecializationCloner::cloneFunctionBody(BridgedFunction originalFunction, BridgedBasicBlock clonedEntryBlock, BridgedValueArray clonedEntryBlockArgs) const {
-  llvm::SmallVector<swift::SILValue, 16> clonedEntryBlockArgsStorage;
-  auto clonedEntryBlockArgsArrayRef = clonedEntryBlockArgs.getValues(clonedEntryBlockArgsStorage);
-  cloner->cloneFunctionBody(originalFunction.getFunction(), clonedEntryBlock.unbridged(), clonedEntryBlockArgsArrayRef);
-}
-
-void BridgedSpecializationCloner::cloneFunctionBody(BridgedFunction originalFunction) const {
-  cloner->cloneFunction(originalFunction.getFunction());
-}
-
 void BridgedBuilder::destroyCapturedArgs(BridgedInstruction partialApply) const {
   if (auto *pai = llvm::dyn_cast<PartialApplyInst>(partialApply.unbridged()); pai->isOnStack()) {
     auto b = unbridged();
-    return swift::insertDestroyOfCapturedArguments(pai, b); 
+    return swift::insertDestroyOfCapturedArguments(pai, b);
   } else {
-    assert(false && "`destroyCapturedArgs` must only be called on a `partial_apply` on stack!");   
+    assert(false && "`destroyCapturedArgs` must only be called on a `partial_apply` on stack!");
   }
 }