Optimize getter keypath patterns

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -212,6 +212,85 @@ static void insertEndAccess(BeginAccessInst *&beginAccess, bool isModify,
   }
 }
 
+static SILValue createKeypathStoredPropertyProjections(
+    SILValue addr, SILLocation loc, BeginAccessInst *&beginAccess,
+    SILBuilder &builder, const KeyPathPatternComponent &comp) {
+  assert(comp.getKind() == KeyPathPatternComponent::Kind::StoredProperty);
+  VarDecl *storedProperty = comp.getStoredPropertyDecl();
+  SILValue elementAddr;
+  if (addr->getType().getStructOrBoundGenericStruct()) {
+    addr = builder.createStructElementAddr(loc, addr, storedProperty);
+  } else if (addr->getType().getClassOrBoundGenericClass()) {
+    SingleValueInstruction *Ref =
+        builder.createLoad(loc, addr, LoadOwnershipQualifier::Unqualified);
+    insertEndAccess(beginAccess, /*isModify*/ false, builder);
+
+    // Handle the case where the storedProperty is in a super class.
+    while (Ref->getType().getClassOrBoundGenericClass() !=
+           storedProperty->getDeclContext()) {
+      SILType superCl = Ref->getType().getSuperclass();
+      if (!superCl) {
+        // This should never happen, because the property should be in the
+        // decl or in a superclass of it. Just handle this to be on the safe
+        // side.
+        return SILValue();
+      }
+      Ref = builder.createUpcast(loc, Ref, superCl);
+    }
+
+    addr = builder.createRefElementAddr(loc, Ref, storedProperty);
+
+    // Class members need access enforcement.
+    if (builder.getModule().getOptions().EnforceExclusivityDynamic) {
+      beginAccess = builder.createBeginAccess(loc, addr, SILAccessKind::Read,
+                                              SILAccessEnforcement::Dynamic,
+                                              /*noNestedConflict*/ false,
+                                              /*fromBuiltin*/ false);
+      addr = beginAccess;
+    }
+  } else {
+    // This should never happen, as a stored-property pattern can only be
+    // applied to classes and structs. But to be safe - and future prove -
+    // let's handle this case and bail.
+    insertEndAccess(beginAccess, /*isModify*/ false, builder);
+    return SILValue();
+  }
+
+  return addr;
+}
+
+static SILValue createKeypathGettablePropertyProjections(
+    KeyPathInst *keyPath, SILValue addr, SILLocation loc,
+    BeginAccessInst *&beginAccess, SILBuilder &builder,
+    const KeyPathPatternComponent &comp) {
+  assert(comp.getKind() == KeyPathPatternComponent::Kind::GettableProperty);
+
+  if (!addr->getType().getStructOrBoundGenericStruct() &&
+      !addr->getType().getClassOrBoundGenericClass()) {
+    // This should never happen, as a stored-property pattern can only be
+    // applied to classes and structs. But to be safe - and future prove -
+    // let's handle this case and bail.
+    insertEndAccess(beginAccess, /*isModify*/ false, builder);
+    return SILValue();
+  }
+
+  SingleValueInstruction *ref =
+      builder.createLoad(loc, addr, LoadOwnershipQualifier::Unqualified);
+
+  insertEndAccess(beginAccess, /*isModify*/ false, builder);
+
+  SmallVector<SILValue, 4> args;
+  args.reserve(keyPath->getNumOperands());
+  for (auto &op : keyPath->getAllOperands()) {
+    args.push_back(op.get());
+  }
+  args.push_back(ref);
+
+  auto *getter = comp.getComputedPropertyId().getFunction();
+  auto *getterFuncRef = builder.createFunctionRef(loc, getter);
+  return builder.createApply(loc, getterFuncRef, SubstitutionMap{}, args);
+}
+
 /// Creates the projection pattern for a keypath instruction.
 ///
 /// Currently only the StoredProperty pattern is handled.
@@ -233,55 +312,23 @@ static SILValue createKeypathProjections(SILValue keyPath, SILValue root,
 
   auto components = kpInst->getPattern()->getComponents();
 
-  // Check if the keypath only contains patterns which we support.
-  for (const KeyPathPatternComponent &comp : components) {
-    if (comp.getKind() != KeyPathPatternComponent::Kind::StoredProperty)
-      return SILValue();
-  }
-
   SILValue addr = root;
+  // Check if the keypath only contains patterns which we support.
   for (const KeyPathPatternComponent &comp : components) {
-    assert(comp.getKind() == KeyPathPatternComponent::Kind::StoredProperty);
-    VarDecl *storedProperty = comp.getStoredPropertyDecl();
-    SILValue elementAddr;
-    if (addr->getType().getStructOrBoundGenericStruct()) {
-      addr = builder.createStructElementAddr(loc, addr, storedProperty);
-    } else if (addr->getType().getClassOrBoundGenericClass()) {
-      SingleValueInstruction *Ref = builder.createLoad(loc, addr,
-                                         LoadOwnershipQualifier::Unqualified);
-      insertEndAccess(beginAccess, /*isModify*/ false, builder);
-
-      // Handle the case where the storedProperty is in a super class.
-      while (Ref->getType().getClassOrBoundGenericClass() !=
-             storedProperty->getDeclContext()) {
-        SILType superCl = Ref->getType().getSuperclass();
-        if (!superCl) {
-          // This should never happen, because the property should be in the
-          // decl or in a superclass of it. Just handle this to be on the safe
-          // side.
-          return SILValue();
-        }
-        Ref = builder.createUpcast(loc, Ref, superCl);
-      }
-
-      addr = builder.createRefElementAddr(loc, Ref, storedProperty);
-
-      // Class members need access enforcement.
-      if (builder.getModule().getOptions().EnforceExclusivityDynamic) {
-        beginAccess = builder.createBeginAccess(loc, addr, SILAccessKind::Read,
-                                                SILAccessEnforcement::Dynamic,
-                                                /*noNestedConflict*/ false,
-                                                /*fromBuiltin*/ false);
-        addr = beginAccess;
-      }
-    } else {
-      // This should never happen, as a stored-property pattern can only be
-      // applied to classes and structs. But to be safe - and future prove -
-      // let's handle this case and bail.
-      insertEndAccess(beginAccess, /*isModify*/ false, builder);
+    switch (comp.getKind()) {
+    case KeyPathPatternComponent::Kind::StoredProperty:
+      addr = createKeypathStoredPropertyProjections(addr, loc, beginAccess,
+                                                    builder, comp);
+      break;
+    case KeyPathPatternComponent::Kind::GettableProperty:
+      addr = createKeypathGettablePropertyProjections(
+          cast<KeyPathInst>(keyPath), addr, loc, beginAccess, builder, comp);
+      break;
+    default:
       return SILValue();
     }
   }
+
   return addr;
 }
 
@@ -325,13 +372,20 @@ bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
   if (!projectedAddr)
     return false;
 
-  if (isModify) {
-    Builder.createCopyAddr(AI->getLoc(), valueAddr, projectedAddr,
-                           IsTake, IsNotInitialization);
+  if (projectedAddr->getType().isAddress()) {
+    if (isModify) {
+      Builder.createCopyAddr(AI->getLoc(), valueAddr, projectedAddr, IsTake,
+                             IsNotInitialization);
+    } else {
+      Builder.createCopyAddr(AI->getLoc(), projectedAddr, valueAddr, IsNotTake,
+                             IsInitialization);
+    }
   } else {
-    Builder.createCopyAddr(AI->getLoc(), projectedAddr, valueAddr,
-                           IsNotTake, IsInitialization);
+    assert(valueAddr->getType().isAddress());
+    Builder.createStore(AI->getLoc(), projectedAddr, valueAddr,
+                        StoreOwnershipQualifier::Unqualified);
   }
+
   insertEndAccess(beginAccess, isModify, Builder);
   eraseInstFromFunction(*AI);
   ++NumOptimizedKeypaths;
@@ -379,9 +433,8 @@ bool SILCombiner::tryOptimizeInoutKeypath(BeginApplyInst *AI) {
     return false;
 
   BeginAccessInst *beginAccess = nullptr;
-  SILValue projectedAddr = createKeypathProjections(keyPath, rootAddr,
-                                                    AI->getLoc(), beginAccess,
-                                                    Builder);
+  SILValue projectedAddr = createKeypathProjections(
+      keyPath, rootAddr, AI->getLoc(), beginAccess, Builder);
   if (!projectedAddr)
     return false;
 

test/SILOptimizer/optimize_keypath.swift

@@ -309,4 +309,45 @@ print("SimpleClass obj count: \(SimpleClass.numObjs)")
 // CHECK-OUTPUT: GenClass obj count: 0
 print("GenClass obj count: \(numGenClassObjs)")
 
+// Optimize subscript getter
+
+struct Sub {
+  public let x : Int
+  public subscript(y: Int, z: Int) -> Int { return x + y + z }
+  public subscript(y: Int) -> Int { return x + y }
+  public subscript() -> Int { return x }
+}
+
+// CHECK-LABEL: sil {{.*}}test_subscript_1
+// CHECK: bb0:
+// CHECK-NEXT: integer_literal $Builtin.Int64, 42
+// CHECK-NEXT: struct
+// CHECK-NEXT: return
+public func test_subscript_1() -> Int {
+  let v = Sub(x: 40)
+  let xKeyPath = \Sub.[1, 1]
+  return v[keyPath: xKeyPath]
+}
+
+// CHECK-LABEL: sil {{.*}}test_subscript_2
+// CHECK: bb0:
+// CHECK-NEXT: integer_literal $Builtin.Int64, 42
+// CHECK-NEXT: struct
+// CHECK-NEXT: return
+public func test_subscript_2() -> Int {
+  let v = Sub(x: 40)
+  let xKeyPath = \Sub.[2]
+  return v[keyPath: xKeyPath]
+}
+
+// CHECK-LABEL: sil {{.*}}test_subscript_3
+// CHECK: bb0:
+// CHECK-NEXT: integer_literal $Builtin.Int64, 42
+// CHECK-NEXT: struct
+// CHECK-NEXT: return
+public func test_subscript_3() -> Int {
+  let v = Sub(x: 42)
+  let xKeyPath = \Sub.[]
+  return v[keyPath: xKeyPath]
+}