Optimizer: extract replacing an apply into a utility function ApplySite.replace(withCallTo:)

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -257,31 +257,8 @@ private struct FunctionSpecializations {
         return argOp.value
       }
     }
-    let specializedCallee = builder.createFunctionRef(originalToSpecialized[callee]!)
-
-    switch apply {
-    case let applyInst as ApplyInst:
-      let newApply = builder.createApply(function: specializedCallee, applyInst.substitutionMap, arguments: newArgs, isNonThrowing: applyInst.isNonThrowing)
-      applyInst.replace(with: newApply, context)
-    case let partialAp as PartialApplyInst:
-      let newApply = builder.createPartialApply(function: specializedCallee, substitutionMap:
-                                                partialAp.substitutionMap,
-                                                capturedArguments: newArgs,
-                                                calleeConvention: partialAp.calleeConvention,
-                                                hasUnknownResultIsolation: partialAp.hasUnknownResultIsolation,
-                                                isOnStack: partialAp.isOnStack)
-      partialAp.replace(with: newApply, context)
-    case let tryApply as TryApplyInst:
-      builder.createTryApply(function: specializedCallee, tryApply.substitutionMap, arguments: newArgs,
-                             normalBlock: tryApply.normalBlock, errorBlock: tryApply.errorBlock)
-      context.erase(instruction: tryApply)
-    case let beginApply as BeginApplyInst:
-      let newApply = builder.createBeginApply(function: specializedCallee, beginApply.substitutionMap,
-                                              arguments: newArgs)
-      beginApply.replace(with: newApply, context)
-    default:
-      fatalError("unknown apply")
-    }
+    apply.replace(withCallTo: originalToSpecialized[callee]!, arguments: newArgs, context)
+
     // It is important to delete the dead `function_ref`. Otherwise it will still reference the original
     // function which prevents deleting it in the mandatory-allocbox-to-stack pass.
     if fri.uses.isEmpty {

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

@@ -538,38 +538,7 @@ private struct SpecializationInfo {
 
     let newApplyArgs = getNewApplyArguments(context)
 
-    let builder = Builder(before: apply, context)
-    let funcRef = builder.createFunctionRef(specializedFunction)
-
-    switch apply {
-    case let oldPartialApply as PartialApplyInst:
-      let newPartialApply = builder.createPartialApply(
-        function: funcRef, substitutionMap: SubstitutionMap(),
-        capturedArguments: newApplyArgs, calleeConvention: oldPartialApply.calleeConvention,
-        hasUnknownResultIsolation: oldPartialApply.hasUnknownResultIsolation,
-        isOnStack: oldPartialApply.isOnStack)
-      oldPartialApply.replace(with: newPartialApply, context)
-
-    case let oldApply as ApplyInst:
-      let newApply = builder.createApply(function: funcRef, SubstitutionMap(), arguments: newApplyArgs,
-                                         isNonThrowing: oldApply.isNonThrowing,
-                                         isNonAsync: oldApply.isNonAsync)
-      oldApply.replace(with: newApply, context)
-
-    case let oldTryApply as TryApplyInst:
-      builder.createTryApply(function: funcRef, SubstitutionMap(), arguments: newApplyArgs,
-                             normalBlock: oldTryApply.normalBlock, errorBlock: oldTryApply.errorBlock,
-                             isNonAsync: oldTryApply.isNonAsync)
-      context.erase(instruction: oldTryApply)
-
-    case let oldBeginApply as BeginApplyInst:
-      let newApply = builder.createBeginApply(function: funcRef, SubstitutionMap(), arguments: newApplyArgs,
-                                              isNonThrowing: oldBeginApply.isNonThrowing,
-                                              isNonAsync: oldBeginApply.isNonAsync)
-      oldBeginApply.replace(with: newApply, context)
-    default:
-      fatalError("unknown apply")
-    }
+    apply.replace(withCallTo: specializedFunction, arguments: newApplyArgs, context)
   }
 
   private func insertCompensatingDestroysForOwnedClosureArguments(_ context: FunctionPassContext) {

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -241,6 +241,47 @@ extension FullApplySite {
   }
 }
 
+extension ApplySite {
+  func replace(withCallTo callee: Function, arguments newArguments: [Value], _ context: some MutatingContext) {
+    let builder = Builder(before: self, context)
+    let calleeRef = builder.createFunctionRef(callee)
+
+    switch self {
+    case let applyInst as ApplyInst:
+      let newApply = builder.createApply(function: calleeRef,
+                                         applyInst.substitutionMap,
+                                         arguments: newArguments,
+                                         isNonThrowing: applyInst.isNonThrowing)
+      applyInst.replace(with: newApply, context)
+
+    case let partialAp as PartialApplyInst:
+      let newApply = builder.createPartialApply(function: calleeRef,
+                                                substitutionMap: partialAp.substitutionMap,
+                                                capturedArguments: newArguments,
+                                                calleeConvention: partialAp.calleeConvention,
+                                                hasUnknownResultIsolation: partialAp.hasUnknownResultIsolation,
+                                                isOnStack: partialAp.isOnStack)
+      partialAp.replace(with: newApply, context)
+
+    case let tryApply as TryApplyInst:
+      builder.createTryApply(function: calleeRef,
+                             tryApply.substitutionMap,
+                             arguments: newArguments,
+                             normalBlock: tryApply.normalBlock, errorBlock: tryApply.errorBlock)
+      context.erase(instruction: tryApply)
+
+    case let beginApply as BeginApplyInst:
+      let newApply = builder.createBeginApply(function: calleeRef,
+                                              beginApply.substitutionMap,
+                                              arguments: newArguments)
+      beginApply.replace(with: newApply, context)
+
+    default:
+      fatalError("unknown apply")
+    }
+  }
+}
+
 extension Builder {
   static func insert(after inst: Instruction, _ context: some MutatingContext, insertFunc: (Builder) -> ()) {
     Builder.insert(after: inst, location: inst.location, context, insertFunc: insertFunc)