Skip to content

Commit 9304ce9

Browse files
authored
Merge pull request #85707 from eeckstein/embedded-witness-method-specialization
embedded: change the function representation of directly called witness methods
2 parents 62a1396 + 64dd574 commit 9304ce9

File tree

27 files changed

+452
-69
lines changed

27 files changed

+452
-69
lines changed

SwiftCompilerSources/Sources/AST/Type.swift

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,25 @@ extension TypeProperties {
227227
GenericSignature(bridged: rawType.bridged.getInvocationGenericSignatureOfFunctionType())
228228
}
229229

230+
public var functionTypeRepresentation: FunctionTypeRepresentation {
231+
switch rawType.bridged.getFunctionTypeRepresentation() {
232+
case .Thick: return .thick
233+
case .Block: return .block
234+
case .Thin: return .thin
235+
case .CFunctionPointer: return .cFunctionPointer
236+
case .Method: return .method
237+
case .ObjCMethod: return .objCMethod
238+
case .WitnessMethod: return .witnessMethod
239+
case .Closure: return .closure
240+
case .CXXMethod: return .cxxMethod
241+
case .KeyPathAccessorGetter: return .keyPathAccessorGetter
242+
case .KeyPathAccessorSetter: return .keyPathAccessorSetter
243+
case .KeyPathAccessorEquals: return .keyPathAccessorEquals
244+
case .KeyPathAccessorHash: return .keyPathAccessorHash
245+
default: fatalError()
246+
}
247+
}
248+
230249
//===--------------------------------------------------------------------===//
231250
// Type properties
232251
//===--------------------------------------------------------------------===//
@@ -307,6 +326,59 @@ extension TypeProperties {
307326
}
308327
}
309328

329+
public enum FunctionTypeRepresentation {
330+
/// A freestanding thick function.
331+
case thick
332+
333+
/// A thick function that is represented as an Objective-C block.
334+
case block
335+
336+
/// A freestanding thin function that needs no context.
337+
case thin
338+
339+
/// A C function pointer, which is thin and also uses the C calling convention.
340+
case cFunctionPointer
341+
342+
/// A Swift instance method.
343+
case method
344+
345+
/// An Objective-C method.
346+
case objCMethod
347+
348+
/// A Swift protocol witness.
349+
case witnessMethod
350+
351+
/// A closure invocation function that has not been bound to a context.
352+
case closure
353+
354+
/// A C++ method that takes a "this" argument (not a static C++ method or constructor).
355+
/// Except for handling the "this" argument, has the same behavior as "CFunctionPointer".
356+
case cxxMethod
357+
358+
/// A KeyPath accessor function, which is thin and also uses the variadic length generic
359+
/// components serialization in trailing buffer. Each representation has a different convention
360+
/// for which parameters have serialized generic type info.
361+
case keyPathAccessorGetter, keyPathAccessorSetter, keyPathAccessorEquals, keyPathAccessorHash
362+
363+
public var bridged: BridgedASTType.FunctionTypeRepresentation {
364+
switch self {
365+
case .thick: return .Thick
366+
case .block: return .Block
367+
case .thin: return .Thin
368+
case .cFunctionPointer: return .CFunctionPointer
369+
case .method: return .Method
370+
case .objCMethod: return .ObjCMethod
371+
case .witnessMethod: return .WitnessMethod
372+
case .closure: return .Closure
373+
case .cxxMethod: return .CXXMethod
374+
case .keyPathAccessorGetter: return .KeyPathAccessorGetter
375+
case .keyPathAccessorSetter: return .KeyPathAccessorSetter
376+
case .keyPathAccessorEquals: return .KeyPathAccessorEquals
377+
case .keyPathAccessorHash: return .KeyPathAccessorHash
378+
}
379+
}
380+
}
381+
310382
public struct TypeArray : RandomAccessCollection, CustomReflectable {
311383
public let bridged: BridgedASTTypeArray
312384

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

Lines changed: 2 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -257,31 +257,8 @@ private struct FunctionSpecializations {
257257
return argOp.value
258258
}
259259
}
260-
let specializedCallee = builder.createFunctionRef(originalToSpecialized[callee]!)
261-
262-
switch apply {
263-
case let applyInst as ApplyInst:
264-
let newApply = builder.createApply(function: specializedCallee, applyInst.substitutionMap, arguments: newArgs, isNonThrowing: applyInst.isNonThrowing)
265-
applyInst.replace(with: newApply, context)
266-
case let partialAp as PartialApplyInst:
267-
let newApply = builder.createPartialApply(function: specializedCallee, substitutionMap:
268-
partialAp.substitutionMap,
269-
capturedArguments: newArgs,
270-
calleeConvention: partialAp.calleeConvention,
271-
hasUnknownResultIsolation: partialAp.hasUnknownResultIsolation,
272-
isOnStack: partialAp.isOnStack)
273-
partialAp.replace(with: newApply, context)
274-
case let tryApply as TryApplyInst:
275-
builder.createTryApply(function: specializedCallee, tryApply.substitutionMap, arguments: newArgs,
276-
normalBlock: tryApply.normalBlock, errorBlock: tryApply.errorBlock)
277-
context.erase(instruction: tryApply)
278-
case let beginApply as BeginApplyInst:
279-
let newApply = builder.createBeginApply(function: specializedCallee, beginApply.substitutionMap,
280-
arguments: newArgs)
281-
beginApply.replace(with: newApply, context)
282-
default:
283-
fatalError("unknown apply")
284-
}
260+
apply.replace(withCallTo: originalToSpecialized[callee]!, arguments: newArgs, context)
261+
285262
// It is important to delete the dead `function_ref`. Otherwise it will still reference the original
286263
// function which prevents deleting it in the mandatory-allocbox-to-stack pass.
287264
if fri.uses.isEmpty {

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ swift_compiler_sources(Optimizer
2222
DeinitDevirtualizer.swift
2323
DestroyHoisting.swift
2424
DiagnoseInfiniteRecursion.swift
25+
EmbeddedWitnessCallSpecialization.swift
2526
InitializeStaticGlobals.swift
2627
LetPropertyLowering.swift
2728
LifetimeDependenceDiagnostics.swift

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/ClosureSpecialization.swift

Lines changed: 2 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -410,7 +410,7 @@ private struct SpecializationInfo {
410410
// The specialized function is always a thin function. This is important because we add additional
411411
// parameters after the Self parameter of witness methods. In this case the new function is not a
412412
// method anymore.
413-
makeThin: true, makeBare: true)
413+
withRepresentation: .thin, makeBare: true)
414414

415415
context.buildSpecializedFunction(
416416
specializedFunction: specializedFunction,
@@ -538,38 +538,7 @@ private struct SpecializationInfo {
538538

539539
let newApplyArgs = getNewApplyArguments(context)
540540

541-
let builder = Builder(before: apply, context)
542-
let funcRef = builder.createFunctionRef(specializedFunction)
543-
544-
switch apply {
545-
case let oldPartialApply as PartialApplyInst:
546-
let newPartialApply = builder.createPartialApply(
547-
function: funcRef, substitutionMap: SubstitutionMap(),
548-
capturedArguments: newApplyArgs, calleeConvention: oldPartialApply.calleeConvention,
549-
hasUnknownResultIsolation: oldPartialApply.hasUnknownResultIsolation,
550-
isOnStack: oldPartialApply.isOnStack)
551-
oldPartialApply.replace(with: newPartialApply, context)
552-
553-
case let oldApply as ApplyInst:
554-
let newApply = builder.createApply(function: funcRef, SubstitutionMap(), arguments: newApplyArgs,
555-
isNonThrowing: oldApply.isNonThrowing,
556-
isNonAsync: oldApply.isNonAsync)
557-
oldApply.replace(with: newApply, context)
558-
559-
case let oldTryApply as TryApplyInst:
560-
builder.createTryApply(function: funcRef, SubstitutionMap(), arguments: newApplyArgs,
561-
normalBlock: oldTryApply.normalBlock, errorBlock: oldTryApply.errorBlock,
562-
isNonAsync: oldTryApply.isNonAsync)
563-
context.erase(instruction: oldTryApply)
564-
565-
case let oldBeginApply as BeginApplyInst:
566-
let newApply = builder.createBeginApply(function: funcRef, SubstitutionMap(), arguments: newApplyArgs,
567-
isNonThrowing: oldBeginApply.isNonThrowing,
568-
isNonAsync: oldBeginApply.isNonAsync)
569-
oldBeginApply.replace(with: newApply, context)
570-
default:
571-
fatalError("unknown apply")
572-
}
541+
apply.replace(withCallTo: specializedFunction, arguments: newApplyArgs, context)
573542
}
574543

575544
private func insertCompensatingDestroysForOwnedClosureArguments(_ context: FunctionPassContext) {
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
//===--- EmbeddedWitnessCallSpecialization.swift ---------------------------==//
2+
//
3+
// This source file is part of the Swift.org open source project
4+
//
5+
// Copyright (c) 2014 - 2025 Apple Inc. and the Swift project authors
6+
// Licensed under Apache License v2.0 with Runtime Library Exception
7+
//
8+
// See https://swift.org/LICENSE.txt for license information
9+
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
10+
//
11+
//===----------------------------------------------------------------------===//
12+
13+
import SIL
14+
15+
/// Changes the function representation of directly called witness methods in Embedded Swift.
16+
///
17+
/// If a function with `witness_method` convention is directly called, the function is specialized
18+
/// by changing the convention to `method` and the call is replaced by a call to the specialized
19+
/// function:
20+
///
21+
/// ```
22+
/// %1 = function_ref @callee : $@convention(witness_method: P) (@guaranteed C) -> ()
23+
/// %2 = apply %1(%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
24+
/// ...
25+
/// sil [ossa] @callee : $@convention(witness_method: P) (@guaranteed C) -> () {
26+
/// ...
27+
/// }
28+
/// ```
29+
/// ->
30+
/// ```
31+
/// %1 = function_ref @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> ()
32+
/// %2 = apply %1(%0) : $@convention(method) (@guaranteed C) -> ()
33+
/// ...
34+
/// // specialized callee
35+
/// sil shared [ossa] @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> () {
36+
/// ...
37+
/// }
38+
/// ```
39+
///
40+
/// This is needed in Embedded Swift because the `witness_method` convention requires passing the
41+
/// witness table to the callee. However, the witness table is not necessarily available.
42+
/// A witness table is only generated if an existential value of a protocol is created.
43+
///
44+
/// This is a rare situation because only witness thunks have `witness_method` convention and those
45+
/// thunks are created as "transparent" functions, which means they are always inlined (after de-
46+
/// virtualization of a witness method call). However, inlining - even of transparent functions -
47+
/// can fail for some reasons.
48+
///
49+
let embeddedWitnessCallSpecialization = FunctionPass(name: "embedded-witness-call-specialization") {
50+
(function: Function, context: FunctionPassContext) in
51+
52+
guard context.options.enableEmbeddedSwift,
53+
!function.isGeneric
54+
else {
55+
return
56+
}
57+
for inst in function.instructions {
58+
if let apply = inst as? FullApplySite {
59+
specializeDirectWitnessMethodCall(apply: apply, context)
60+
}
61+
}
62+
}
63+
64+
private func specializeDirectWitnessMethodCall(apply: FullApplySite, _ context: FunctionPassContext) {
65+
guard apply.callee.type.functionTypeRepresentation == .witnessMethod,
66+
let callee = apply.referencedFunction,
67+
callee.isDefinition
68+
else {
69+
return
70+
}
71+
72+
let specializedFunctionName = context.mangle(withChangedRepresentation: callee)
73+
74+
let specializedFunction: Function
75+
76+
if let existingSpecializedFunction = context.lookupFunction(name: specializedFunctionName) {
77+
specializedFunction = existingSpecializedFunction
78+
79+
} else {
80+
specializedFunction = context.createSpecializedFunctionDeclaration(
81+
from: callee, withName: specializedFunctionName,
82+
withParams: Array(callee.convention.parameters),
83+
withRepresentation: .method)
84+
85+
context.buildSpecializedFunction(
86+
specializedFunction: specializedFunction,
87+
buildFn: { (specializedFunction, specializedContext) in
88+
cloneFunction(from: callee, toEmpty: specializedFunction, specializedContext)
89+
})
90+
91+
context.notifyNewFunction(function: specializedFunction, derivedFrom: callee)
92+
}
93+
94+
apply.replace(withCallTo: specializedFunction, arguments: Array(apply.arguments), context)
95+
}

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/PackSpecialization.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -553,7 +553,7 @@ private struct PackExplodedFunction {
553553
withParams: newParameters,
554554
withResults: newResults,
555555
// If a method has a dynamic self parameter, it cannot be converted into a thin function (non-method).
556-
makeThin: !original.mayBindDynamicSelf)
556+
withRepresentation: original.mayBindDynamicSelf ? nil : .thin)
557557

558558
self.buildSpecializedFunction(context)
559559
}

SwiftCompilerSources/Sources/Optimizer/PassManager/FunctionPassContext.swift

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -164,16 +164,21 @@ struct FunctionPassContext : MutatingContext {
164164
}
165165
}
166166

167+
func mangle(withChangedRepresentation original: Function) -> String {
168+
String(taking: bridgedPassContext.mangleWithChangedRepresentation(original.bridged))
169+
}
170+
167171
func createSpecializedFunctionDeclaration(
168172
from original: Function, withName specializedFunctionName: String,
169173
withParams specializedParameters: [ParameterInfo],
170174
withResults specializedResults: [ResultInfo]? = nil,
171-
makeThin: Bool = false,
175+
withRepresentation: FunctionTypeRepresentation? = nil,
172176
makeBare: Bool = false,
173177
preserveGenericSignature: Bool = true
174178
) -> Function {
175179
return specializedFunctionName._withBridgedStringRef { nameRef in
176180
let bridgedParamInfos = specializedParameters.map { $0._bridged }
181+
let repr = withRepresentation ?? original.loweredFunctionType.functionTypeRepresentation
177182

178183
return bridgedParamInfos.withUnsafeBufferPointer { paramBuf in
179184

@@ -183,15 +188,15 @@ struct FunctionPassContext : MutatingContext {
183188
return bridgedPassContext.createSpecializedFunctionDeclaration(
184189
nameRef, paramBuf.baseAddress, paramBuf.count,
185190
resultBuf.baseAddress, resultBuf.count,
186-
original.bridged, makeThin, makeBare,
191+
original.bridged, repr.bridged, makeBare,
187192
preserveGenericSignature
188193
).function
189194
}
190195
} else {
191196
return bridgedPassContext.createSpecializedFunctionDeclaration(
192197
nameRef, paramBuf.baseAddress, paramBuf.count,
193198
nil, 0,
194-
original.bridged, makeThin, makeBare,
199+
original.bridged, repr.bridged, makeBare,
195200
preserveGenericSignature
196201
).function
197202
}

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ private func registerSwiftPasses() {
7373
registerPass(allocBoxToStack, { allocBoxToStack.run($0) })
7474
registerPass(asyncDemotion, { asyncDemotion.run($0) })
7575
registerPass(booleanLiteralFolding, { booleanLiteralFolding.run($0) })
76+
registerPass(embeddedWitnessCallSpecialization, { embeddedWitnessCallSpecialization.run($0) })
7677
registerPass(letPropertyLowering, { letPropertyLowering.run($0) })
7778
registerPass(mergeCondFailsPass, { mergeCondFailsPass.run($0) })
7879
registerPass(constantCapturePropagation, { constantCapturePropagation.run($0) })

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,47 @@ extension FullApplySite {
241241
}
242242
}
243243

244+
extension ApplySite {
245+
func replace(withCallTo callee: Function, arguments newArguments: [Value], _ context: some MutatingContext) {
246+
let builder = Builder(before: self, context)
247+
let calleeRef = builder.createFunctionRef(callee)
248+
249+
switch self {
250+
case let applyInst as ApplyInst:
251+
let newApply = builder.createApply(function: calleeRef,
252+
applyInst.substitutionMap,
253+
arguments: newArguments,
254+
isNonThrowing: applyInst.isNonThrowing)
255+
applyInst.replace(with: newApply, context)
256+
257+
case let partialAp as PartialApplyInst:
258+
let newApply = builder.createPartialApply(function: calleeRef,
259+
substitutionMap: partialAp.substitutionMap,
260+
capturedArguments: newArguments,
261+
calleeConvention: partialAp.calleeConvention,
262+
hasUnknownResultIsolation: partialAp.hasUnknownResultIsolation,
263+
isOnStack: partialAp.isOnStack)
264+
partialAp.replace(with: newApply, context)
265+
266+
case let tryApply as TryApplyInst:
267+
builder.createTryApply(function: calleeRef,
268+
tryApply.substitutionMap,
269+
arguments: newArguments,
270+
normalBlock: tryApply.normalBlock, errorBlock: tryApply.errorBlock)
271+
context.erase(instruction: tryApply)
272+
273+
case let beginApply as BeginApplyInst:
274+
let newApply = builder.createBeginApply(function: calleeRef,
275+
beginApply.substitutionMap,
276+
arguments: newArguments)
277+
beginApply.replace(with: newApply, context)
278+
279+
default:
280+
fatalError("unknown apply")
281+
}
282+
}
283+
}
284+
244285
extension Builder {
245286
static func insert(after inst: Instruction, _ context: some MutatingContext, insertFunc: (Builder) -> ()) {
246287
Builder.insert(after: inst, location: inst.location, context, insertFunc: insertFunc)

0 commit comments

Comments
 (0)