SIL: Don't DFE functions referenced from KeyPath patterns.

Fixes rdar://problem/31776015.

Author: jckarter

include/swift/SIL/SILInstruction.h

@@ -1599,6 +1599,9 @@ class KeyPathPatternComponent {
                                    getter, setter, indices, ty);
   }
 
+  void incrementRefCounts() const;
+  void decrementRefCounts() const;
+  
   void Profile(llvm::FoldingSetNodeID &ID);
 };
 
@@ -1696,6 +1699,7 @@ class KeyPathInst final
 
 public:
   KeyPathPattern *getPattern() const;
+  bool hasPattern() const { return (bool)Pattern; }
 
   ArrayRef<Operand> getAllOperands() const {
     // TODO: Subscript keypaths will have operands.
@@ -1712,6 +1716,8 @@ class KeyPathInst final
     return V->getKind() == ValueKind::KeyPathInst;
   }
 
+  void dropReferencedPattern();
+  
   ~KeyPathInst();
 };
 

lib/SIL/SILInstruction.cpp

@@ -152,11 +152,22 @@ void SILInstruction::dropAllReferences() {
 
   // If we have a function ref inst, we need to especially drop its function
   // argument so that it gets a proper ref decrement.
-  auto *FRI = dyn_cast<FunctionRefInst>(this);
-  if (!FRI || !FRI->getReferencedFunction())
+  if (auto *FRI = dyn_cast<FunctionRefInst>(this)) {
+    if (!FRI->getReferencedFunction())
+      return;
+    FRI->dropReferencedFunction();
     return;
+  }
 
-  FRI->dropReferencedFunction();
+  // If we have a KeyPathInst, drop its pattern reference so that we can
+  // decrement refcounts on referenced functions.
+  if (auto *KPI = dyn_cast<KeyPathInst>(this)) {
+    if (!KPI->hasPattern())
+      return;
+    
+    KPI->dropReferencedPattern();
+    return;
+  }
 }
 
 void SILInstruction::replaceAllUsesWithUndef() {

lib/SIL/SILInstructions.cpp

@@ -1978,6 +1978,40 @@ bool KeyPathPatternComponent::isComputedSettablePropertyMutating() const {
   }
 }
 
+static void
+forEachRefcountableReference(const KeyPathPatternComponent &component,
+                         llvm::function_ref<void (SILFunction*)> forFunction) {
+  switch (component.getKind()) {
+  case KeyPathPatternComponent::Kind::StoredProperty:
+    return;
+  case KeyPathPatternComponent::Kind::SettableProperty:
+    forFunction(component.getComputedPropertySetter());
+    LLVM_FALLTHROUGH;
+  case KeyPathPatternComponent::Kind::GettableProperty:
+    forFunction(component.getComputedPropertyGetter());
+    
+    switch (component.getComputedPropertyId().getKind()) {
+    case KeyPathPatternComponent::ComputedPropertyId::DeclRef:
+      // Mark the vtable entry as used somehow?
+      return;
+    case KeyPathPatternComponent::ComputedPropertyId::Function:
+      forFunction(component.getComputedPropertyId().getFunction());
+      return;
+    case KeyPathPatternComponent::ComputedPropertyId::Property:
+      return;
+    }
+  }
+}
+
+void KeyPathPatternComponent::incrementRefCounts() const {
+  forEachRefcountableReference(*this,
+    [&](SILFunction *f) { f->incrementRefCount(); });
+}
+void KeyPathPatternComponent::decrementRefCounts() const {
+  forEachRefcountableReference(*this,
+                               [&](SILFunction *f) { f->decrementRefCount(); });
+}
+
 KeyPathPattern *
 KeyPathPattern::get(SILModule &M, CanGenericSignature signature,
                     CanType rootType, CanType valueType,
@@ -2116,6 +2150,11 @@ KeyPathInst::KeyPathInst(SILDebugLocation Loc,
 {
   auto *subsBuf = getTrailingObjects<Substitution>();
   std::uninitialized_copy(Subs.begin(), Subs.end(), subsBuf);
+  
+  // Increment the use of any functions referenced from the keypath pattern.
+  for (auto component : Pattern->getComponents()) {
+    component.incrementRefCounts();
+  }
 }
 
 MutableArrayRef<Substitution>
@@ -2130,9 +2169,24 @@ KeyPathInst::getAllOperands() {
 }
 
 KeyPathInst::~KeyPathInst() {
+  if (!Pattern)
+    return;
+
+  // Decrement the use of any functions referenced from the keypath pattern.
+  for (auto component : Pattern->getComponents()) {
+    component.decrementRefCounts();
+  }
   // TODO: destroy operands
 }
 
 KeyPathPattern *KeyPathInst::getPattern() const {
+  assert(Pattern && "pattern was reset!");
   return Pattern;
 }
+
+void KeyPathInst::dropReferencedPattern() {
+  for (auto component : Pattern->getComponents()) {
+    component.decrementRefCounts();
+  }
+  Pattern = nullptr;
+}

lib/SILOptimizer/IPO/DeadFunctionElimination.cpp

@@ -192,6 +192,46 @@ class FunctionLivenessComputation {
 
   }
 
+  /// Marks the declarations referenced by a key path pattern as alive if they
+  /// aren't yet.
+  void ensureKeyPathComponentsAreAlive(KeyPathPattern *KP) {
+    for (auto &component : KP->getComponents()) {
+      switch (component.getKind()) {
+      case KeyPathPatternComponent::Kind::SettableProperty:
+        ensureAlive(component.getComputedPropertySetter());
+        LLVM_FALLTHROUGH;
+      case KeyPathPatternComponent::Kind::GettableProperty: {
+        ensureAlive(component.getComputedPropertyGetter());
+        auto id = component.getComputedPropertyId();
+        switch (id.getKind()) {
+        case KeyPathPatternComponent::ComputedPropertyId::DeclRef: {
+          auto decl = cast<AbstractFunctionDecl>(id.getDeclRef().getDecl());
+          if (auto clas = dyn_cast<ClassDecl>(decl->getDeclContext())) {
+            ensureAliveClassMethod(getMethodInfo(decl, /*witness*/ false),
+                                   dyn_cast<FuncDecl>(decl),
+                                   clas);
+          } else if (auto proto =
+                       dyn_cast<ProtocolDecl>(decl->getDeclContext())) {
+            ensureAliveProtocolMethod(getMethodInfo(decl, /*witness*/ true));
+          } else {
+            llvm_unreachable("key path keyed by a non-class, non-protocol method");
+          }
+          break;
+        }
+        case KeyPathPatternComponent::ComputedPropertyId::Function:
+          ensureAlive(id.getFunction());
+          break;
+        case KeyPathPatternComponent::ComputedPropertyId::Property:
+          break;
+        }
+        continue;
+      }
+      case KeyPathPatternComponent::Kind::StoredProperty:
+        continue;
+      }
+    }
+  }
+  
   /// Marks a function as alive if it is not alive yet.
   void ensureAlive(SILFunction *F) {
     if (!isAlive(F))
@@ -313,6 +353,8 @@ class FunctionLivenessComputation {
           ensureAliveClassMethod(mi, dyn_cast<FuncDecl>(funcDecl), MethodCl);
         } else if (auto *FRI = dyn_cast<FunctionRefInst>(&I)) {
           ensureAlive(FRI->getReferencedFunction());
+        } else if (auto *KPI = dyn_cast<KeyPathInst>(&I)) {
+          ensureKeyPathComponentsAreAlive(KPI->getPattern());
         }
       }
     }

test/stdlib/KeyPath.swift

@@ -4,9 +4,6 @@
 // REQUIRES: executable_test
 // REQUIRES: PTRSIZE=64
 
-// Disabled for now
-// REQUIRES: rdar://problem/31776015
-
 import StdlibUnittest
 
 var keyPath = TestSuite("key paths")

test/stdlib/KeyPathImplementation.swift

@@ -4,10 +4,6 @@
 // REQUIRES: executable_test
 // REQUIRES: PTRSIZE=64
 
-// Disabled for now
-// REQUIRES: rdar://problem/31776015
-
-
 import StdlibUnittest
 
 var keyPathImpl = TestSuite("key path implementation")

test/stdlib/KeyPathObjC.swift

@@ -5,9 +5,6 @@
 // REQUIRES: PTRSIZE=64
 // REQUIRES: objc_interop
 
-// Disabled for now
-// REQUIRES: rdar://problem/31776015
-
 import StdlibUnittest
 import Foundation