embedded: change the function representation of directly called witness methods

This is needed in Embedded Swift because the `witness_method` convention requires passing the witness table to the callee.
However, the witness table is not necessarily available.
A witness table is only generated if an existential value of a protocol is created.

This is a rare situation because only witness thunks have `witness_method` convention and those thunks are created as "transparent" functions, which means they are always inlined (after de-virtualization of a witness method call).
However, inlining - even of transparent functions - can fail for some reasons.

This change adds a new EmbeddedWitnessCallSpecialization pass:
If a function with `witness_method` convention is directly called, the function is specialized by changing the convention to `method` and the call is replaced by a call to the specialized function:

```
  %1 = function_ref @callee : $@convention(witness_method: P) (@guaranteed C) -> ()
  %2 = apply %1(%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
...
sil [ossa] @callee : $@convention(witness_method: P) (@guaranteed C) -> () {
  ...
}
```
->
```
  %1 = function_ref @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> ()
  %2 = apply %1(%0) : $@convention(method) (@guaranteed C) -> ()
...
// specialized callee
sil shared [ossa] @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> () {
  ...
}
```

Fixes a compiler crash
rdar://165184147

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/CMakeLists.txt

@@ -21,6 +21,7 @@ swift_compiler_sources(Optimizer
   DeinitDevirtualizer.swift
   DestroyHoisting.swift
   DiagnoseInfiniteRecursion.swift
+  EmbeddedWitnessCallSpecialization.swift
   InitializeStaticGlobals.swift
   LetPropertyLowering.swift
   LifetimeDependenceDiagnostics.swift

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/EmbeddedWitnessCallSpecialization.swift

@@ -0,0 +1,95 @@
+//===--- EmbeddedWitnessCallSpecialization.swift ---------------------------==//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2025 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
+
+/// Changes the function representation of directly called witness methods in Embedded Swift.
+///
+/// If a function with `witness_method` convention is directly called, the function is specialized
+/// by changing the convention to `method` and the call is replaced by a call to the specialized
+/// function:
+///
+/// ```
+///   %1 = function_ref @callee : $@convention(witness_method: P) (@guaranteed C) -> ()
+///   %2 = apply %1(%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
+/// ...
+/// sil [ossa] @callee : $@convention(witness_method: P) (@guaranteed C) -> () {
+///   ...
+/// }
+/// ```
+/// ->
+/// ```
+///   %1 = function_ref @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> ()
+///   %2 = apply %1(%0) : $@convention(method) (@guaranteed C) -> ()
+/// ...
+/// // specialized callee
+/// sil shared [ossa] @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> () {
+///   ...
+/// }
+/// ```
+///
+/// This is needed in Embedded Swift because the `witness_method` convention requires passing the
+/// witness table to the callee. However, the witness table is not necessarily available.
+/// A witness table is only generated if an existential value of a protocol is created.
+///
+/// This is a rare situation because only witness thunks have `witness_method` convention and those
+/// thunks are created as "transparent" functions, which means they are always inlined (after de-
+/// virtualization of a witness method call). However, inlining - even of transparent functions -
+/// can fail for some reasons.
+///
+let embeddedWitnessCallSpecialization = FunctionPass(name: "embedded-witness-call-specialization") {
+  (function: Function, context: FunctionPassContext) in
+
+  guard context.options.enableEmbeddedSwift,
+        !function.isGeneric
+  else {
+    return
+  }
+  for inst in function.instructions {
+    if let apply = inst as? FullApplySite {
+      specializeDirectWitnessMethodCall(apply: apply, context)
+    }
+  }
+}
+
+private func specializeDirectWitnessMethodCall(apply: FullApplySite, _ context: FunctionPassContext) {
+  guard apply.callee.type.functionTypeRepresentation == .witnessMethod,
+        let callee = apply.referencedFunction,
+        callee.isDefinition
+  else {
+    return
+  }
+
+  let specializedFunctionName = context.mangle(withChangedRepresentation: callee)
+
+  let specializedFunction: Function
+
+  if let existingSpecializedFunction = context.lookupFunction(name: specializedFunctionName) {
+    specializedFunction = existingSpecializedFunction
+
+  } else {
+    specializedFunction = context.createSpecializedFunctionDeclaration(
+      from: callee, withName: specializedFunctionName,
+      withParams: Array(callee.convention.parameters),
+      withRepresentation: .method)
+
+    context.buildSpecializedFunction(
+      specializedFunction: specializedFunction,
+      buildFn: { (specializedFunction, specializedContext) in
+        cloneFunction(from: callee, toEmpty: specializedFunction, specializedContext)
+      })
+
+    context.notifyNewFunction(function: specializedFunction, derivedFrom: callee)
+  }
+
+  apply.replace(withCallTo: specializedFunction, arguments: Array(apply.arguments), context)
+}

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -73,6 +73,7 @@ private func registerSwiftPasses() {
   registerPass(allocBoxToStack, { allocBoxToStack.run($0) })
   registerPass(asyncDemotion, { asyncDemotion.run($0) })
   registerPass(booleanLiteralFolding, { booleanLiteralFolding.run($0) })
+  registerPass(embeddedWitnessCallSpecialization, { embeddedWitnessCallSpecialization.run($0) })
   registerPass(letPropertyLowering, { letPropertyLowering.run($0) })
   registerPass(mergeCondFailsPass, { mergeCondFailsPass.run($0) })
   registerPass(constantCapturePropagation, { constantCapturePropagation.run($0) })

include/swift/SILOptimizer/PassManager/Passes.def

@@ -75,6 +75,8 @@ PASS(MandatoryDestroyHoisting, "mandatory-destroy-hoisting",
      "Hoist destroy_value instructions for non-lexical values")
 PASS(DeadEndBlockDumper, "dump-deadendblocks",
      "Tests the DeadEndBlocks utility")
+PASS(EmbeddedWitnessCallSpecialization, "embedded-witness-call-specialization",
+     "Mandatory witness method call specialization")
 PASS(EscapeInfoDumper, "dump-escape-info",
      "Dumps escape information")
 PASS(AddressEscapeInfoDumper, "dump-addr-escape-info",

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -265,6 +265,7 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   P.addMandatoryPerformanceOptimizations();
   P.addOnoneSimplification();
   P.addInitializeStaticGlobals();
+  P.addEmbeddedWitnessCallSpecialization();
 
   P.addMandatoryDestroyHoisting();
 
@@ -893,6 +894,10 @@ static void addLastChanceOptPassPipeline(SILPassPipelinePlan &P) {
     P.addAssumeSingleThreaded();
   }
 
+  // Needs to run again at the end of the pipeline (after all de-virtualizations
+  // are done) in case an optimization pass de-virtualizes a witness method call.
+  P.addEmbeddedWitnessCallSpecialization();
+
   // Emits remarks on all functions with @_assemblyVision attribute.
   P.addAssemblyVisionRemarkGenerator();
 

test/SILOptimizer/witness-call-specialization.sil

@@ -0,0 +1,98 @@
+// RUN: %target-sil-opt -enable-experimental-feature Embedded %s -embedded-witness-call-specialization | %FileCheck %s
+// RUN: %target-sil-opt %s -embedded-witness-call-specialization | %FileCheck --check-prefix=CHECKNE %s
+
+// REQUIRES: swift_feature_Embedded
+
+sil_stage canonical
+
+import Builtin
+
+protocol P {}
+
+class C: P {}
+
+// CHECK-LABEL: sil [ossa] @tst_simple :
+// CHECK:         [[F:%.*]] = function_ref @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> ()
+// CHECK:         apply [[F]](%0) : $@convention(method) (@guaranteed C) -> ()
+// CHECK:       } // end sil function 'test_simple'
+// CHECKNE-LABEL: sil [ossa] @test_simple :
+// CHECKNE:         [[F:%.*]] = function_ref @callee : $@convention(witness_method: P) (@guaranteed C) -> ()
+// CHECKNE:         apply [[F]](%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
+// CHECKNE:       } // end sil function 'test_simple'
+sil [ossa] @test_simple : $@convention(thin) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  %1 = function_ref @callee : $@convention(witness_method: P) (@guaranteed C) -> ()
+  %2 = apply %1(%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
+  %3 = tuple ()
+  return %3
+}
+
+// CHECK-LABEL: sil [ossa] @second_use :
+// CHECK:         [[F:%.*]] = function_ref @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> ()
+// CHECK:         apply [[F]](%0) : $@convention(method) (@guaranteed C) -> ()
+// CHECK:       } // end sil function 'second_use'
+sil [ossa] @second_use : $@convention(thin) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  %1 = function_ref @callee : $@convention(witness_method: P) (@guaranteed C) -> ()
+  %2 = apply %1(%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
+  %3 = tuple ()
+  return %3
+}
+
+// CHECK-LABEL: sil [ossa] @dont_handle_generic_functions :
+// CHECK:         [[F:%.*]] = function_ref @callee : $@convention(witness_method: P) (@guaranteed C) -> ()
+// CHECK:         apply [[F]](%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
+// CHECK:       } // end sil function 'dont_handle_generic_functions'
+sil [ossa] @dont_handle_generic_functions : $@convention(thin) <T> (@guaranteed C, @in_guaranteed T) -> () {
+bb0(%0 : @guaranteed $C, %1 : $*T):
+  %2 = function_ref @callee : $@convention(witness_method: P) (@guaranteed C) -> ()
+  %3 = apply %2(%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
+  %4 = tuple ()
+  return %4
+}
+
+// CHECK-LABEL: sil [ossa] @only_specialize_witness_methods :
+// CHECK:         [[F:%.*]] = function_ref @thin_function : $@convention(thin) (@guaranteed C) -> ()
+// CHECK:         apply [[F]](%0) : $@convention(thin) (@guaranteed C) -> ()
+// CHECK:       } // end sil function 'only_specialize_witness_methods'
+sil [ossa] @only_specialize_witness_methods : $@convention(thin) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  %1 = function_ref @thin_function : $@convention(thin) (@guaranteed C) -> ()
+  %2 = apply %1(%0) : $@convention(thin) (@guaranteed C) -> ()
+  %3 = tuple ()
+  return %3
+}
+
+// CHECK-LABEL: sil [ossa] @callee_must_have_body :
+// CHECK:         [[F:%.*]] = function_ref @no_body : $@convention(witness_method: P) (@guaranteed C) -> ()
+// CHECK:         apply [[F]](%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
+// CHECK:       } // end sil function 'callee_must_have_body'
+sil [ossa] @callee_must_have_body : $@convention(thin) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  %1 = function_ref @no_body : $@convention(witness_method: P) (@guaranteed C) -> ()
+  %2 = apply %1(%0) : $@convention(witness_method: P) (@guaranteed C) -> ()
+  %3 = tuple ()
+  return %3
+}
+
+
+// CHECK-LABEL: sil shared [ossa] @$e6calleeTfr9 : $@convention(method) (@guaranteed C) -> () {
+// CHECK:         fix_lifetime %0
+// CHECK-NEXT:    %2 = tuple ()
+// CHECK-NEXT:    return %2
+// CHECK:       } // end sil function '$e6calleeTfr9'
+sil [ossa] @callee : $@convention(witness_method: P) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  fix_lifetime %0
+  %2 = tuple ()
+  return %2
+}
+
+sil [ossa] @thin_function : $@convention(thin) (@guaranteed C) -> () {
+bb0(%0 : @guaranteed $C):
+  fix_lifetime %0
+  %2 = tuple ()
+  return %2
+}
+
+sil [ossa] @no_body : $@convention(witness_method: P) (@guaranteed C) -> ()

test/embedded/accessors.swift

@@ -21,6 +21,25 @@ public class C {
   var y: Int = 27
 }
 
+public protocol P {
+  var d: [Int : WrappedBool] { get set }
+}
+
+extension P {
+  mutating func set(key: Int) {
+    d[key]?.b = true
+  }
+}
+
+public struct WrappedBool {
+  public var b: Bool = true
+}
+
+public class S: P {
+  public var d: [Int : WrappedBool] = [:]
+  public func foo() {}
+}
+
 @main
 struct Main {
   static func main() {
@@ -33,5 +52,11 @@ struct Main {
 
     print("2") // CHECK: 2
     print("")
+
+    var handler = S()
+    handler.d[27] = WrappedBool(b: false)
+    handler.set(key: 27)
+    // CHECK: true
+    print(handler.d[27]!.b ? "true" : "false") 
   }
 }