[embedded] Compile-time (literal) KeyPaths for Embedded Swift

Enable KeyPath/AnyKeyPath/PartialKeyPath/WritableKeyPath in Embedded Swift, but
for compile-time use only:

- Add keypath optimizations into the mandatory optimizations pipeline
- Allow keypath optimizations to look through begin_borrow, to make them work
  even in OSSA.
- If a use of a KeyPath doesn't optimize away, diagnose in PerformanceDiagnostics
- Make UnsafePointer.pointer(to:) transparent to allow the keypath optimization
  to happen in the callers of UnsafePointer.pointer(to:).

Author: kubamracek

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/CMakeLists.txt

@@ -21,6 +21,7 @@ swift_compiler_sources(Optimizer
   SimplifyDestructure.swift
   SimplifyGlobalValue.swift
   SimplifyInitEnumDataAddr.swift
+  SimplifyKeyPath.swift
   SimplifyLoad.swift
   SimplifyPartialApply.swift
   SimplifyPointerToAddress.swift

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyApply.swift

@@ -17,6 +17,10 @@ extension ApplyInst : OnoneSimplifyable {
     if tryTransformThickToThinCallee(of: self, context) {
       return
     }
+    if context.tryOptimizeKeypath(apply: self) {
+      context.erase(instruction: self)
+      return
+    }
     _ = context.tryDevirtualize(apply: self, isMandatory: false)
   }
 }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyKeyPath.swift

@@ -0,0 +1,35 @@
+//===--- SimplifyKeyPath.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
+
+extension KeyPathInst : OnoneSimplifyable {
+  func simplify(_ context: SimplifyContext) {
+    if allUsesRemovable(instruction: self) {
+      context.erase(instructionIncludingAllUsers: self)
+    }
+  }
+}
+
+fileprivate func allUsesRemovable(instruction: Instruction) -> Bool {
+  for result in instruction.results {
+    for use in result.uses {
+      if !(use.instruction is UpcastInst || use.instruction is DestroyValueInst || use.instruction is BeginBorrowInst || use.instruction is EndBorrowInst) {
+        return false
+      }
+      if !allUsesRemovable(instruction: use.instruction) {
+        return false
+      }
+    }
+  }
+  return true
+}

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -170,6 +170,10 @@ extension MutatingContext {
     }
     return nil
   }
+  
+  func tryOptimizeKeypath(apply: FullApplySite) -> Bool {
+    return _bridged.tryOptimizeKeypath(apply.bridged)
+  }
 
   func inlineFunction(apply: FullApplySite, mandatoryInline: Bool) {
     // This is only a best-effort attempt to notity the new cloned instructions as changed.

include/swift/AST/DiagnosticsSIL.def

@@ -379,6 +379,8 @@ ERROR(embedded_swift_metatype_type,none,
       "cannot use metatype of type %0 in embedded Swift", (Type))
 ERROR(embedded_swift_metatype,none,
       "cannot use metatype in embedded Swift", ())
+ERROR(embedded_swift_keypath,none,
+      "cannot use key path in embedded Swift", ())
 ERROR(embedded_swift_allocating_type,none,
       "cannot use allocating type %0 in -no-allocations mode", (Type))
 ERROR(embedded_swift_allocating,none,

include/swift/SILOptimizer/OptimizerBridging.h

@@ -214,6 +214,7 @@ struct BridgedPassContext {
   bool tryOptimizeApplyOfPartialApply(BridgedInstruction closure) const;
   bool tryDeleteDeadClosure(BridgedInstruction closure, bool needKeepArgsAlive) const;
   SWIFT_IMPORT_UNSAFE DevirtResult tryDevirtualizeApply(BridgedInstruction apply, bool isMandatory) const;
+  bool tryOptimizeKeypath(BridgedInstruction apply) const;
   SWIFT_IMPORT_UNSAFE OptionalBridgedValue constantFoldBuiltin(BridgedInstruction builtin) const;
   SWIFT_IMPORT_UNSAFE swift::SILVTable * _Nullable specializeVTableForType(BridgedType type,
                                                                            BridgedFunction function) const;

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -599,6 +599,13 @@ bool specializeAppliesInFunction(SILFunction &F,
                                  SILTransform *transform,
                                  bool isMandatory);
 
+bool tryOptimizeKeypath(ApplyInst *AI, SILBuilder Builder);
+bool tryOptimizeKeypathApplication(ApplyInst *AI, SILFunction *callee, SILBuilder Builder);
+bool tryOptimizeKeypathOffsetOf(ApplyInst *AI, FuncDecl *calleeFn,
+                                KeyPathInst *kp, SILBuilder Builder);
+bool tryOptimizeKeypathKVCString(ApplyInst *AI, FuncDecl *calleeFn,
+                                KeyPathInst *kp, SILBuilder Builder);
+
 /// Instantiate the specified type by recursively tupling and structing the
 /// unique instances of the empty types and undef "instances" of the non-empty
 /// types aggregated together at each level.

lib/SILOptimizer/Mandatory/PerformanceDiagnostics.cpp

@@ -519,6 +519,10 @@ bool PerformanceDiagnostics::visitInst(SILInstruction *inst,
         diagnose(loc, diag::embedded_swift_value_deinit, impactType.getASTType());
         return true;
       }
+      if (isa<KeyPathInst>(inst)) {
+        diagnose(loc, diag::embedded_swift_keypath);
+        return true;
+      }
       if (!allowedMetadataUseInEmbeddedSwift(inst)) {
         PrettyStackTracePerformanceDiagnostics stackTrace("metatype", inst);
         if (impactType) {

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1600,6 +1600,11 @@ BridgedPassContext::DevirtResult BridgedPassContext::tryDevirtualizeApply(Bridge
   return {{nullptr}, false};
 }
 
+bool BridgedPassContext::tryOptimizeKeypath(BridgedInstruction apply) const {
+  SILBuilder builder(apply.unbridged());
+  return ::tryOptimizeKeypath(apply.getAs<ApplyInst>(), builder);
+}
+
 OptionalBridgedValue BridgedPassContext::constantFoldBuiltin(BridgedInstruction builtin) const {
   auto bi = builtin.getAs<BuiltinInst>();
   std::optional<bool> resultsInError;

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -330,13 +330,7 @@ class SILCombiner :
   SILInstruction *optimizeApplyOfConvertFunctionInst(FullApplySite AI,
                                                      ConvertFunctionInst *CFI);
 
-  bool tryOptimizeKeypath(ApplyInst *AI);
   bool tryOptimizeInoutKeypath(BeginApplyInst *AI);
-  bool tryOptimizeKeypathApplication(ApplyInst *AI, SILFunction *callee);
-  bool tryOptimizeKeypathOffsetOf(ApplyInst *AI, FuncDecl *calleeFn,
-                                  KeyPathInst *kp);
-  bool tryOptimizeKeypathKVCString(ApplyInst *AI, FuncDecl *calleeFn,
-                                  KeyPathInst *kp);
 
   /// Sinks owned forwarding instructions to their uses if they do not have
   /// non-debug non-consuming uses. Deletes any debug_values and destroy_values