optimize keypath instructions

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -226,6 +226,7 @@ class SILCombiner :
   SILInstruction *visitRetainValueAddrInst(RetainValueAddrInst *CI);
   SILInstruction *visitPartialApplyInst(PartialApplyInst *AI);
   SILInstruction *visitApplyInst(ApplyInst *AI);
+  SILInstruction *visitBeginApplyInst(BeginApplyInst *BAI);
   SILInstruction *visitTryApplyInst(TryApplyInst *AI);
   SILInstruction *optimizeStringObject(BuiltinInst *BI);
   SILInstruction *visitBuiltinInst(BuiltinInst *BI);
@@ -299,6 +300,10 @@ class SILCombiner :
 
   SILInstruction *optimizeApplyOfConvertFunctionInst(FullApplySite AI,
                                                      ConvertFunctionInst *CFI);
+
+  bool tryOptimizeKeypath(ApplyInst *AI);
+  bool tryOptimizeInoutKeypath(BeginApplyInst *AI);
+
   // Optimize concatenation of string literals.
   // Constant-fold concatenation of string literals known at compile-time.
   SILInstruction *optimizeConcatenationOfStringLiterals(ApplyInst *AI);

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -30,10 +30,13 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
 
 using namespace swift;
 using namespace swift::PatternMatch;
 
+STATISTIC(NumOptimizedKeypaths, "Number of optimized keypath instructions");
+
 /// Remove pointless reabstraction thunk closures.
 ///   partial_apply %reabstraction_thunk_typeAtoB(
 ///      partial_apply %reabstraction_thunk_typeBtoA %closure_typeB))
@@ -548,6 +551,189 @@ 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()) {
+      LoadInst *Ref = builder.createLoad(loc, addr,
+                                         LoadOwnershipQualifier::Unqualified);
+      insertEndAccess(beginAccess, /*isModify*/ false, builder);
+      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):
+///   %kp = keypath ...
+///   apply %keypath_runtime_function(%addr, %kp, %root_object)
+/// with:
+///   %addr = struct_element_addr/ref_element_addr %root_object
+///   ...
+///   load/store %addr
+bool SILCombiner::tryOptimizeKeypath(ApplyInst *AI) {
+  SILFunction *callee = AI->getReferencedFunction();
+  if (!callee)
+    return false;
+
+  if (AI->getNumArguments() != 3)
+    return false;
+
+  SILValue keyPath, rootAddr, valueAddr;
+  bool isModify = false;
+  if (callee->getName() == "swift_setAtWritableKeyPath" ||
+      callee->getName() == "swift_setAtReferenceWritableKeyPath") {
+    keyPath = AI->getArgument(1);
+    rootAddr = AI->getArgument(0);
+    valueAddr = AI->getArgument(2);
+    isModify = true;
+  } else if (callee->getName() == "swift_getAtKeyPath") {
+    keyPath = AI->getArgument(2);
+    rootAddr = AI->getArgument(1);
+    valueAddr = AI->getArgument(0);
+  } else {
+    return false;
+  }
+
+  BeginAccessInst *beginAccess = nullptr;
+  SILValue projectedAddr = createKeypathProjections(keyPath, rootAddr,
+                                                    AI->getLoc(), beginAccess,
+                                                    Builder);
+  if (!projectedAddr)
+    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);
+  eraseInstFromFunction(*AI);
+  ++NumOptimizedKeypaths;
+  return true;
+}
+
+/// Try to optimize a keypath application with an apply instruction.
+///
+/// Replaces (simplified SIL):
+///   %kp = keypath ...
+///   %inout_addr = begin_apply %keypath_runtime_function(%kp, %root_object)
+///   // use %inout_addr
+///   end_apply
+/// with:
+///   %addr = struct_element_addr/ref_element_addr %root_object
+///   // use %inout_addr
+bool SILCombiner::tryOptimizeInoutKeypath(BeginApplyInst *AI) {
+  SILFunction *callee = AI->getReferencedFunction();
+  if (!callee)
+    return false;
+
+  if (AI->getNumArguments() != 2)
+    return false;
+
+  SILValue keyPath = AI->getArgument(1);
+  SILValue rootAddr = AI->getArgument(0);
+  bool isModify = false;
+  if (callee->getName() == "swift_modifyAtWritableKeyPath" ||
+      callee->getName() == "swift_modifyAtReferenceWritableKeyPath") {
+    isModify = true;
+  } else if (callee->getName() != "swift_readAtKeyPath") {
+    return false;
+  }
+
+  SILInstructionResultArray yields = AI->getYieldedValues();
+  if (yields.size() != 1)
+    return false;
+
+  SILValue valueAddr = yields[0];
+  Operand *AIUse = AI->getTokenResult()->getSingleUse();
+  if (!AIUse)
+    return false;
+  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)
+    return false;
+
+  // Replace the projected address.
+  valueAddr->replaceAllUsesWith(projectedAddr);
+
+  Builder.setInsertionPoint(endApply);
+  insertEndAccess(beginAccess, isModify, Builder);
+  eraseInstFromFunction(*endApply);
+  eraseInstFromFunction(*AI);
+  ++NumOptimizedKeypaths;
+  return true;
+}
+
 bool
 SILCombiner::recursivelyCollectARCUsers(UserListTy &Uses, ValueBase *Value) {
   // FIXME: We could probably optimize this case too
@@ -1327,6 +1513,9 @@ SILInstruction *SILCombiner::visitApplyInst(ApplyInst *AI) {
   if (auto *CFI = dyn_cast<ConvertFunctionInst>(AI->getCallee()))
     return optimizeApplyOfConvertFunctionInst(AI, CFI);
 
+  if (tryOptimizeKeypath(AI))
+    return nullptr;
+
   // Optimize readonly functions with no meaningful users.
   SILFunction *SF = AI->getReferencedFunction();
   if (SF && SF->getEffectsKind() < EffectsKind::ReleaseNone) {
@@ -1393,6 +1582,12 @@ SILInstruction *SILCombiner::visitApplyInst(ApplyInst *AI) {
   return nullptr;
 }
 
+SILInstruction *SILCombiner::visitBeginApplyInst(BeginApplyInst *BAI) {
+  if (tryOptimizeInoutKeypath(BAI))
+    return nullptr;
+  return nullptr;
+}
+
 bool SILCombiner::
 isTryApplyResultNotUsed(UserListTy &AcceptedUses, TryApplyInst *TAI) {
   SILBasicBlock *NormalBB = TAI->getNormalBB();

lib/SILOptimizer/Transforms/DeadObjectElimination.cpp

@@ -52,6 +52,9 @@ STATISTIC(DeadAllocRefEliminated,
 STATISTIC(DeadAllocStackEliminated,
           "number of AllocStack instructions removed");
 
+STATISTIC(DeadKeyPathEliminated,
+          "number of keypath instructions removed");
+
 STATISTIC(DeadAllocApplyEliminated,
           "number of allocating Apply instructions removed");
 
@@ -208,6 +211,9 @@ static bool canZapInstruction(SILInstruction *Inst, bool acceptRefCountInsts,
   if (isa<InjectEnumAddrInst>(Inst))
     return true;
 
+  if (isa<KeyPathInst>(Inst))
+    return true;
+
   // We know that the destructor has no side effects so we can remove the
   // deallocation instruction too.
   if (isa<DeallocationInst>(Inst) || isa<AllocationInst>(Inst))
@@ -643,17 +649,20 @@ class DeadObjectElimination : public SILFunctionTransform {
   llvm::SmallVector<SILInstruction*, 16> Allocations;
 
   void collectAllocations(SILFunction &Fn) {
-    for (auto &BB : Fn)
+    for (auto &BB : Fn) {
       for (auto &II : BB) {
-        if (isa<AllocationInst>(&II))
-          Allocations.push_back(&II);
-        else if (isAllocatingApply(&II))
+        if (isa<AllocationInst>(&II) ||
+            isAllocatingApply(&II) ||
+            isa<KeyPathInst>(&II)) {
           Allocations.push_back(&II);
+        }
       }
+    }
   }
 
   bool processAllocRef(AllocRefInst *ARI);
   bool processAllocStack(AllocStackInst *ASI);
+  bool processKeyPath(KeyPathInst *KPI);
   bool processAllocBox(AllocBoxInst *ABI){ return false;}
   bool processAllocApply(ApplyInst *AI, DeadEndBlocks &DEBlocks);
 
@@ -668,6 +677,8 @@ class DeadObjectElimination : public SILFunctionTransform {
         Changed |= processAllocRef(A);
       else if (auto *A = dyn_cast<AllocStackInst>(II))
         Changed |= processAllocStack(A);
+      else if (auto *KPI = dyn_cast<KeyPathInst>(II))
+        Changed |= processKeyPath(KPI);
       else if (auto *A = dyn_cast<AllocBoxInst>(II))
         Changed |= processAllocBox(A);
       else if (auto *A = dyn_cast<ApplyInst>(II))
@@ -749,6 +760,30 @@ bool DeadObjectElimination::processAllocStack(AllocStackInst *ASI) {
   return true;
 }
 
+bool DeadObjectElimination::processKeyPath(KeyPathInst *KPI) {
+  UserList UsersToRemove;
+  if (hasUnremovableUsers(KPI, UsersToRemove, /*acceptRefCountInsts=*/ true,
+      /*onlyAcceptTrivialStores*/ false)) {
+    LLVM_DEBUG(llvm::dbgs() << "    Found a use that cannot be zapped...\n");
+    return false;
+  }
+
+  // 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()) {
+    if (!Op.get()->getType().isTrivial(*KPI->getFunction()))
+      return false;
+  }
+
+  // Remove the keypath and all of its users.
+  removeInstructions(
+    ArrayRef<SILInstruction*>(UsersToRemove.begin(), UsersToRemove.end()));
+  LLVM_DEBUG(llvm::dbgs() << "    Success! Eliminating keypath.\n");
+
+  ++DeadKeyPathEliminated;
+  return true;
+}
+
 /// If AI is the version of an initializer where we pass in either an apply or
 /// an alloc_ref to initialize in place, validate that we are able to continue
 /// optimizing and return To