EmbeddedSwiftDiagnostics: check apply diagnostics also for begin_apply and partial_apply

SwiftCompilerSources/Sources/Optimizer/ModulePasses/EmbeddedSwiftDiagnostics.swift

@@ -131,20 +131,26 @@ private struct FunctionChecker {
                               at: instruction.location)
       }
 
-    case is BeginApplyInst:
-      if context.options.noAllocations {
-        throw Diagnostic(.embedded_swift_allocating_coroutine, at: instruction.location)
-      }
-
     case is ThunkInst:
       if context.options.noAllocations {
         throw Diagnostic(.embedded_swift_allocating, at: instruction.location)
       }
 
+    case let ba as BeginApplyInst:
+      if context.options.noAllocations {
+        throw Diagnostic(.embedded_swift_allocating_coroutine, at: instruction.location)
+      }
+      try checkApply(apply: ba)
+
     case let pai as PartialApplyInst:
       if context.options.noAllocations && !pai.isOnStack {
         throw Diagnostic(.embedded_swift_allocating_closure, at: instruction.location)
       }
+      try checkApply(apply: pai)
+
+    // Remaining apply instructions
+    case let apply as ApplySite:
+      try checkApply(apply: apply)
 
     case let bi as BuiltinInst:
       switch bi.id {
@@ -162,42 +168,43 @@ private struct FunctionChecker {
         break
       }
 
-    case let apply as ApplySite:
-      if context.options.noAllocations && apply.isAsync {
-        throw Diagnostic(.embedded_swift_allocating_type, at: instruction.location)
-      }
-
-      if !apply.callee.type.hasValidSignatureForEmbedded,
-         // Some runtime functions have generic parameters in SIL, which are not used in IRGen.
-         // Therefore exclude runtime functions at all.
-         !apply.callsEmbeddedRuntimeFunction
-      {
-        switch apply.callee {
-        case let cmi as ClassMethodInst:
-          throw Diagnostic(.embedded_cannot_specialize_class_method, cmi.member, at: instruction.location)
-        case let wmi as WitnessMethodInst:
-          throw Diagnostic(.embedded_cannot_specialize_witness_method, wmi.member, at: instruction.location)
-        default:
-          throw Diagnostic(.embedded_call_generic_function, at: instruction.location)
-        }
-      }
-
-      // Although all (non-generic) functions are initially put into the worklist there are two reasons
-      // to call `checkFunction` recursively:
-      // * To get a better caller info in the diagnostics.
-      // * When passing an opened existential to a generic function, it's valid in Embedded swift even if the
-      //   generic is not specialized. We need to check such generic functions, too.
-      if let callee = apply.referencedFunction {
-        callStack.push(CallSite(apply: apply, callee: callee))
-        try checkFunction(callee)
-        _ = callStack.pop()
-      }
-
     default:
       break
     }
   }
 
+  mutating func checkApply(apply: ApplySite) throws {
+    if context.options.noAllocations && apply.isAsync {
+      throw Diagnostic(.embedded_swift_allocating_type, at: apply.location)
+    }
+
+    if !apply.callee.type.hasValidSignatureForEmbedded,
+       // Some runtime functions have generic parameters in SIL, which are not used in IRGen.
+       // Therefore exclude runtime functions at all.
+       !apply.callsEmbeddedRuntimeFunction
+    {
+      switch apply.callee {
+      case let cmi as ClassMethodInst:
+        throw Diagnostic(.embedded_cannot_specialize_class_method, cmi.member, at: apply.location)
+      case let wmi as WitnessMethodInst:
+        throw Diagnostic(.embedded_cannot_specialize_witness_method, wmi.member, at: apply.location)
+      default:
+        throw Diagnostic(.embedded_call_generic_function, at: apply.location)
+      }
+    }
+
+    // Although all (non-generic) functions are initially put into the worklist there are two reasons
+    // to call `checkFunction` recursively:
+    // * To get a better caller info in the diagnostics.
+    // * When passing an opened existential to a generic function, it's valid in Embedded swift even if the
+    //   generic is not specialized. We need to check such generic functions, too.
+    if let callee = apply.referencedFunction {
+      callStack.push(CallSite(apply: apply, callee: callee))
+      try checkFunction(callee)
+      _ = callStack.pop()
+    }
+  }
+
   // Check for any violations in witness tables for existentials.
   mutating func checkConformance(_ conformance: Conformance, location: Location) throws {
     guard conformance.isConcrete,