[SILOptimizer] Generalize optimization of static keypaths

We have an optimization in SILCombiner that "inlines" the use of compile-time constant key paths by performing the property access directly instead of calling a runtime function (leading to huge performance gains e.g. for heavy use of @dynamicMemberLookup). However, this optimization previously only supported key paths which solely access stored properties, so computed properties, optional chaining, etc. still had to call a runtime function. This commit generalizes the optimization to support all types of key paths.

Author: NobodyNada

include/swift/SILOptimizer/Utils/KeyPathProjector.h

@@ -0,0 +1,81 @@
+//===-- KeyPathProjector.h - Project a static key path ----------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2019 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// Utility class to project a statically known key path
+/// expression to a direct property access sequence.
+///
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SILOPTIMIZER_UTILS_KEYPATHPROJECTOR_H
+#define SWIFT_SILOPTIMIZER_UTILS_KEYPATHPROJECTOR_H
+
+#include "swift/SIL/SILBuilder.h"
+#include <memory>
+
+namespace swift {
+
+/// Projects a statically known key path expression to
+/// a direct property access.
+class KeyPathProjector {
+public:
+  /// The type of a key path access.
+  enum class AccessType {
+    /// A get-only access (i.e. swift_getAtKeyPath).
+    Get,
+    
+    /// A set-only access (i.e. swift_setAtWritableKeyPath).
+    Set,
+    
+    /// A modification (i.e. swift_modifyAtWritableKeyPath).
+    Modify
+  };
+  
+  /// Creates a key path projector for a key path.
+  ///
+  /// Returns nullptr if \p keyPath is not a keypath instruction or if there is
+  /// any other reason why the optimization cannot be done.
+  ///
+  /// \param keyPath The key path to project. Must be the result of either
+  ///    a keypath instruction or an upcast of a key path instruction.
+  /// \param root The address of the object the key path is applied to.
+  /// \param loc The location of the key path application.
+  /// \param builder The SILBuilder to use.
+  static std::unique_ptr<KeyPathProjector>
+  create(SILValue keyPath, SILValue root, SILLocation loc, SILBuilder &builder);
+  
+  /// Projects the key path to an address. Sets up the projection,
+  /// invokes the callback, then tears down the projection.
+  /// \param accessType The access type of the projected address.
+  /// \param callback A callback to invoke with the projected adddress.
+  ///     The projected address is only valid from within \p callback.
+  virtual void project(AccessType accessType,
+               std::function<void(SILValue addr)> callback) = 0;
+    
+  virtual ~KeyPathProjector() {};
+  
+  /// Whether this projection returns a struct.
+  virtual bool isStruct() = 0;
+protected:
+  KeyPathProjector(SILLocation loc, SILBuilder &builder)
+      : loc(loc), builder(builder) {}
+  
+  /// The location of the key path application.
+  SILLocation loc;
+  
+  /// The SILBuilder to use.
+  SILBuilder &builder;
+};
+
+}  // end namespace swift
+
+#endif /* KeyPathProjector_h */

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -28,11 +28,13 @@
 #include "swift/SILOptimizer/Analysis/ValueTracking.h"
 #include "swift/SILOptimizer/Utils/CFGOptUtils.h"
 #include "swift/SILOptimizer/Utils/Existential.h"
+#include "swift/SILOptimizer/Utils/KeyPathProjector.h"
 #include "swift/SILOptimizer/Utils/ValueLifetime.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
+#include <utility>
 
 using namespace swift;
 using namespace swift::PatternMatch;
@@ -199,92 +201,6 @@ SILCombiner::optimizeApplyOfConvertFunctionInst(FullApplySite AI,
   return NAI;
 }
 
-/// Ends the begin_access "scope" if a begin_access was inserted for optimizing
-/// a keypath pattern.
-static void insertEndAccess(BeginAccessInst *&beginAccess, bool isModify,
-                            SILBuilder &builder) {
-  if (beginAccess) {
-    builder.createEndAccess(beginAccess->getLoc(), beginAccess,
-                            /*aborted*/ false);
-    if (isModify)
-      beginAccess->setAccessKind(SILAccessKind::Modify);
-    beginAccess = nullptr;
-  }
-}
-
-/// Creates the projection pattern for a keypath instruction.
-///
-/// Currently only the StoredProperty pattern is handled.
-/// TODO: handle other patterns, like getters/setters, optional chaining, etc.
-///
-/// Returns false if \p keyPath is not a keypath instruction or if there is any
-/// other reason why the optimization cannot be done.
-static SILValue createKeypathProjections(SILValue keyPath, SILValue root,
-                                         SILLocation loc,
-                                         BeginAccessInst *&beginAccess,
-                                         SILBuilder &builder) {
-  if (auto *upCast = dyn_cast<UpcastInst>(keyPath))
-    keyPath = upCast->getOperand();
-
-  // Is it a keypath instruction at all?
-  auto *kpInst = dyn_cast<KeyPathInst>(keyPath);
-  if (!kpInst || !kpInst->hasPattern())
-    return SILValue();
-
-  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;
-  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);
-      return SILValue();
-    }
-  }
-  return addr;
-}
-
 /// Try to optimize a keypath application with an apply instruction.
 ///
 /// Replaces (simplified SIL):
@@ -317,22 +233,26 @@ bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
   } else {
     return false;
   }
-
-  BeginAccessInst *beginAccess = nullptr;
-  SILValue projectedAddr = createKeypathProjections(keyPath, rootAddr,
-                                                    AI->getLoc(), beginAccess,
-                                                    Builder);
-  if (!projectedAddr)
+  
+  auto projector = KeyPathProjector::create(keyPath, rootAddr,
+                                            AI->getLoc(), Builder);
+  if (!projector)
     return false;
-
-  if (isModify) {
-    Builder.createCopyAddr(AI->getLoc(), valueAddr, projectedAddr,
-                           IsTake, IsNotInitialization);
-  } else {
-    Builder.createCopyAddr(AI->getLoc(), projectedAddr, valueAddr,
-                           IsNotTake, IsInitialization);
-  }
-  insertEndAccess(beginAccess, isModify, Builder);
+  
+  KeyPathProjector::AccessType accessType;
+  if (isModify) accessType = KeyPathProjector::AccessType::Set;
+  else accessType = KeyPathProjector::AccessType::Get;
+  
+  projector->project(accessType, [&](SILValue projectedAddr) {
+    if (isModify) {
+      Builder.createCopyAddr(AI->getLoc(), valueAddr, projectedAddr,
+                             IsTake, IsNotInitialization);
+    } else {
+      Builder.createCopyAddr(AI->getLoc(), projectedAddr, valueAddr,
+                             IsNotTake, IsInitialization);
+    }
+  });
+  
   eraseInstFromFunction(*AI);
   ++NumOptimizedKeypaths;
   return true;
@@ -377,19 +297,24 @@ bool SILCombiner::tryOptimizeInoutKeypath(BeginApplyInst *AI) {
   EndApplyInst *endApply = dyn_cast<EndApplyInst>(AIUse->getUser());
   if (!endApply)
     return false;
-
-  BeginAccessInst *beginAccess = nullptr;
-  SILValue projectedAddr = createKeypathProjections(keyPath, rootAddr,
-                                                    AI->getLoc(), beginAccess,
-                                                    Builder);
-  if (!projectedAddr)
+  
+  auto projector = KeyPathProjector::create(keyPath, rootAddr,
+                                            AI->getLoc(), Builder);
+  if (!projector)
     return false;
+    
+  KeyPathProjector::AccessType accessType;
+  if (isModify) accessType = KeyPathProjector::AccessType::Modify;
+  else accessType = KeyPathProjector::AccessType::Get;
+  
+  projector->project(accessType, [&](SILValue projectedAddr) {
+    // Replace the projected address.
+    valueAddr->replaceAllUsesWith(projectedAddr);
+    
+    // Skip to the end of the key path application before cleaning up.
+    Builder.setInsertionPoint(endApply);
+  });
 
-  // Replace the projected address.
-  valueAddr->replaceAllUsesWith(projectedAddr);
-
-  Builder.setInsertionPoint(endApply);
-  insertEndAccess(beginAccess, isModify, Builder);
   eraseInstFromFunction(*endApply);
   eraseInstFromFunction(*AI);
   ++NumOptimizedKeypaths;

lib/SILOptimizer/Utils/CMakeLists.txt

@@ -11,6 +11,7 @@ silopt_register_sources(
   GenericCloner.cpp
   Generics.cpp
   InstOptUtils.cpp
+  KeyPathProjector.cpp
   LoadStoreOptUtils.cpp
   LoopUtils.cpp
   OptimizerStatsUtils.cpp