Merge pull request #63748 from eeckstein/fix-keypath-inst

SIL: add type-dependent operands to the `keypath` instruction

Author: eeckstein

include/swift/SIL/SILInstruction.h

@@ -3870,7 +3870,8 @@ class KeyPathInst final
   friend TrailingObjects;
 
   KeyPathPattern *Pattern;
-  unsigned NumOperands;
+  unsigned numPatternOperands;
+  unsigned numTypeDependentOperands;
   SubstitutionMap Substitutions;
 
   static KeyPathInst *create(SILDebugLocation Loc,
@@ -3883,11 +3884,12 @@ class KeyPathInst final
   KeyPathInst(SILDebugLocation Loc,
               KeyPathPattern *Pattern,
               SubstitutionMap Subs,
-              ArrayRef<SILValue> Args,
+              ArrayRef<SILValue> allOperands,
+              unsigned numPatternOperands,
               SILType Ty);
 
   size_t numTrailingObjects(OverloadToken<Operand>) const {
-    return NumOperands;
+    return numPatternOperands + numTypeDependentOperands;
   }
 
 public:
@@ -3899,6 +3901,23 @@ class KeyPathInst final
   }
   MutableArrayRef<Operand> getAllOperands();
 
+  ArrayRef<Operand> getPatternOperands() const {
+    return getAllOperands().slice(0, numPatternOperands);
+  }
+
+  MutableArrayRef<Operand> getPatternOperands() {
+    return getAllOperands().slice(0, numPatternOperands);
+  }
+
+
+  ArrayRef<Operand> getTypeDependentOperands() const {
+    return getAllOperands().slice(numPatternOperands);
+  }
+
+  MutableArrayRef<Operand> getTypeDependentOperands() {
+    return getAllOperands().slice(numPatternOperands);
+  }
+
   SubstitutionMap getSubstitutions() const { return Substitutions; }
 
   void dropReferencedPattern();

lib/SIL/IR/SILInstructions.cpp

@@ -2899,24 +2899,30 @@ KeyPathInst::create(SILDebugLocation Loc,
   assert(Args.size() == Pattern->getNumOperands()
          && "number of key path args doesn't match pattern");
 
-  auto totalSize = totalSizeToAlloc<Operand>(Args.size());
+  SmallVector<SILValue, 8> allOperands(Args.begin(), Args.end());
+  collectTypeDependentOperands(allOperands, F, Ty);
+
+  auto totalSize = totalSizeToAlloc<Operand>(allOperands.size());
   void *mem = F.getModule().allocateInst(totalSize, alignof(KeyPathInst));
-  return ::new (mem) KeyPathInst(Loc, Pattern, Subs, Args, Ty);
+  return ::new (mem) KeyPathInst(Loc, Pattern, Subs, allOperands, Args.size(), Ty);
 }
 
 KeyPathInst::KeyPathInst(SILDebugLocation Loc,
                          KeyPathPattern *Pattern,
                          SubstitutionMap Subs,
-                         ArrayRef<SILValue> Args,
+                         ArrayRef<SILValue> allOperands,
+                         unsigned numPatternOperands,
                          SILType Ty)
   : InstructionBase(Loc, Ty),
     Pattern(Pattern),
-    NumOperands(Pattern->getNumOperands()),
+    numPatternOperands(numPatternOperands),
+    numTypeDependentOperands(allOperands.size() - numPatternOperands),
     Substitutions(Subs)
 {
+  assert(allOperands.size() >= numPatternOperands);
   auto *operandsBuf = getTrailingObjects<Operand>();
-  for (unsigned i = 0; i < Args.size(); ++i) {
-    ::new ((void*)&operandsBuf[i]) Operand(this, Args[i]);
+  for (unsigned i = 0; i < allOperands.size(); ++i) {
+    ::new ((void*)&operandsBuf[i]) Operand(this, allOperands[i]);
   }
 
   // Increment the use of any functions referenced from the keypath pattern.
@@ -2927,7 +2933,7 @@ KeyPathInst::KeyPathInst(SILDebugLocation Loc,
 
 MutableArrayRef<Operand>
 KeyPathInst::getAllOperands() {
-  return {getTrailingObjects<Operand>(), NumOperands};
+  return {getTrailingObjects<Operand>(), numPatternOperands + numTypeDependentOperands};
 }
 
 KeyPathInst::~KeyPathInst() {

lib/SIL/IR/SILPrinter.cpp

@@ -2673,10 +2673,10 @@ class SILPrinter : public SILInstructionVisitor<SILPrinter> {
       *this << ' ';
       printSubstitutions(KPI->getSubstitutions());
     }
-    if (!KPI->getAllOperands().empty()) {
+    if (!KPI->getPatternOperands().empty()) {
       *this << " (";
 
-      interleave(KPI->getAllOperands(),
+      interleave(KPI->getPatternOperands(),
         [&](const Operand &operand) {
           *this << Ctx.getID(operand.get());
         }, [&]{

lib/SIL/Verifier/SILVerifier.cpp

@@ -5240,7 +5240,7 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
           baseTy,
           leafTy,
           component,
-          KPI->getAllOperands(),
+          KPI->getPatternOperands(),
           KPI->getPattern()->getGenericSignature(),
           KPI->getSubstitutions(),
           /*property descriptor*/false,

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -73,7 +73,7 @@ static SILInstruction *getConstant(SILValue V) {
     // mangling scheme.
     // But currently it's not worth it because we do not optimize subscript
     // keypaths in SILCombine.
-    if (kp->getNumOperands() != 0)
+    if (kp->getPatternOperands().size() != 0)
       return nullptr;
     if (!kp->hasPattern())
       return nullptr;

lib/SILOptimizer/Transforms/DeadObjectElimination.cpp

@@ -882,7 +882,7 @@ bool DeadObjectElimination::processKeyPath(KeyPathInst *KPI) {
 
   // For simplicity just bail if the keypath has a non-trivial operands.
   // TODO: don't bail but insert compensating destroys for such operands.
-  for (const Operand &Op : KPI->getAllOperands()) {
+  for (const Operand &Op : KPI->getPatternOperands()) {
     if (!Op.get()->getType().isTrivial(*KPI->getFunction()))
       return false;
   }

lib/Serialization/SerializeSIL.cpp

@@ -2490,7 +2490,7 @@ void SILSerializer::writeSILInstruction(const SILInstruction &SI) {
       writeKeyPathPatternComponent(component, ListOfValues);
     }
 
-    for (auto &operand : KPI->getAllOperands()) {
+    for (auto &operand : KPI->getPatternOperands()) {
       auto value = operand.get();
       ListOfValues.push_back(addValueRef(value));
       ListOfValues.push_back(S.addTypeRef(value->getType().getASTType()));

test/SILOptimizer/inline_keypath.swift

@@ -0,0 +1,18 @@
+// RUN: %target-swift-frontend -O -emit-sil %s | %FileCheck %s
+
+public protocol P {}
+
+// CHECK-LABEL: sil @$s14inline_keypath6testitySbAA1P_pXp_s14PartialKeyPathCyxGtSlRzlF
+// CHECK:         [[T:%[0-9]+]] = open_existential_metatype %0
+// CHECK:         = keypath {{.*}}@opened{{.*}} type-defs: [[T]]
+// CHECK:       } // end sil function '$s14inline_keypath6testitySbAA1P_pXp_s14PartialKeyPathCyxGtSlRzlF'
+public func testit<C: Collection>(_ t: any P.Type, _ kp: PartialKeyPath<C>) -> Bool {
+  return tyFunc(t, kp)
+}
+
+@inline(__always)
+func tyFunc<C: Collection, H: P>(_ h: H.Type, _ kp: PartialKeyPath<C>) -> Bool {
+  return kp == \Array<H>.count
+}
+
+