[SILOptimizer] fix crash in key path projector

When a computed property returns a generic, the accessor's function
type may involve a type parameter that needs to be resolved using
the key path instruction's substitution map.

Author: NobodyNada

lib/SILOptimizer/Utils/KeyPathProjector.cpp

@@ -208,7 +208,9 @@ class GettablePropertyProjector : public ComponentProjector {
       // The callback expects a memory address it can read from,
       // so allocate a buffer.
       auto &function = builder.getFunction();
-      SILType type = function.getLoweredType(component.getComponentType());
+      auto substType = component.getComponentType().subst(keyPath->getSubstitutions(),
+                                                           None);
+      SILType type = function.getLoweredType(substType);
       auto addr = builder.createAllocStack(loc, type);
 
       assertHasNoContext();
@@ -283,7 +285,9 @@ class SettablePropertyProjector : public GettablePropertyProjector {
           // The callback expects a memory address it can write to,
           // so allocate a writeback buffer.
           auto &function = builder.getFunction();
-          SILType type = function.getLoweredType(component.getComponentType());
+          auto substType = component.getComponentType().subst(keyPath->getSubstitutions(),
+                                                               None);
+          SILType type = function.getLoweredType(substType);
           auto addr = builder.createAllocStack(loc, type);
 
           assertHasNoContext();
@@ -315,10 +319,12 @@ class SettablePropertyProjector : public GettablePropertyProjector {
 
 class OptionalWrapProjector : public ComponentProjector {
 public:
-  OptionalWrapProjector(const KeyPathPatternComponent &component,
+  OptionalWrapProjector(KeyPathInst *kpInst,
+                        const KeyPathPatternComponent &component,
                         std::unique_ptr<KeyPathProjector> parent,
                         SILLocation loc, SILBuilder &builder)
-        : ComponentProjector(component, std::move(parent), loc, builder) {}
+        : ComponentProjector(component, std::move(parent), loc, builder),
+          keyPath(kpInst) {}
 
   void project(AccessType accessType,
                std::function<void(SILValue addr)> callback) override {
@@ -328,7 +334,9 @@ class OptionalWrapProjector : public ComponentProjector {
 
     parent->project(AccessType::Get, [&](SILValue parentValue) {
       auto &function = builder.getFunction();
-      SILType optType = function.getLoweredType(component.getComponentType());
+      auto substType = component.getComponentType().subst(keyPath->getSubstitutions(),
+                                                           None);
+      SILType optType = function.getLoweredType(substType);
       SILType objType = optType.getOptionalObjectType().getAddressType();
 
       assert(objType && "optional wrap must return an optional");
@@ -352,6 +360,9 @@ class OptionalWrapProjector : public ComponentProjector {
       builder.createDeallocStack(loc, optAddr);
     });
   }
+
+private:
+  KeyPathInst *keyPath;
 };
 
 class OptionalForceProjector : public ComponentProjector {
@@ -531,7 +542,8 @@ class CompleteKeyPathProjector : public KeyPathProjector {
       // If we're reading an optional chain, create an optional result.
       auto resultCanType = components.back().getComponentType();
       auto &function = builder.getFunction();
-      auto optType = function.getLoweredType(resultCanType);
+      auto substType = resultCanType.subst(keyPath->getSubstitutions(), None);
+      auto optType = function.getLoweredType(substType);
 
       assert(optType.getOptionalObjectType() &&
              "Optional-chained key path should result in an optional");
@@ -608,7 +620,7 @@ class CompleteKeyPathProjector : public KeyPathProjector {
         break;
       case KeyPathPatternComponent::Kind::OptionalWrap:
         projector = std::make_unique<OptionalWrapProjector>
-            (comp, std::move(parent), loc, builder);
+            (keyPath, comp, std::move(parent), loc, builder);
         break;
       case KeyPathPatternComponent::Kind::OptionalForce:
         projector = std::make_unique<OptionalForceProjector>

test/SILOptimizer/optimize_keypath.swift

@@ -16,6 +16,8 @@ protocol P {
 struct GenStruct<T : P> : P {
   var st: T
   var computed: Int { get { st.computed } set { st.computed = newValue } }
+
+  var computedGeneric: T { get { st} set { st = newValue} }
 
   init(_ st: T) { self.st = st }
 
@@ -581,6 +583,20 @@ func testGetComplex(_ s: SimpleClass) -> Int {
   return s[keyPath: kp]
 }
 
+// allow exactly one unoptimzedkey path instruction, in this function
+// CHECK-ALL: sil {{.*}}makeKeyPathInGenericContext
+// CHECK-ALL: = keypath
+func makeKeyPathInGenericContext<T: P>(of: T.Type) -> WritableKeyPath<GenStruct<T>, T> {
+  \.computedGeneric
+}
+
+// CHECK-ALL-NOT: = keypath
+
+func testGenericResult(_ s: inout GenStruct<SimpleStruct>) {
+    let kp = makeKeyPathInGenericContext(of: SimpleStruct.self)
+    s[keyPath: kp].i += 1
+}
+
 // CHECK-LABEL: sil {{.*}}testit
 func testit() {
   // CHECK-OUTPUT: GenStructRead: 27
@@ -680,6 +696,11 @@ func testit() {
 
   // CHECK-OUTPUT: testGetComplex: 1
   print("testGetComplex: \(testGetComplex(c4))")
+  
+  // CHECK-OUTPUT: testGenericResult: 2
+  var s5 = GenStruct(SimpleStruct(i: 1))
+  testGenericResult(&s5)
+  print("testGenericResult: \(s5.st.i)")
 }
 
 testit()