[Prototype] Implement public KeyPath elimination

Author: kateinoigakukun

include/swift/SIL/SILInstruction.h

@@ -2953,7 +2953,10 @@ class KeyPathPatternComponent {
   void visitReferencedFunctionsAndMethods(
       std::function<void (SILFunction *)> functionCallBack,
       std::function<void (SILDeclRef)> methodCallBack) const;
-    
+
+  void clearReferencedFunctions_if(
+      llvm::function_ref<bool(SILFunction *)> predicate);
+
   void incrementRefCounts() const;
   void decrementRefCounts() const;
 

include/swift/SIL/SILProperty.h

@@ -67,7 +67,14 @@ class SILProperty : public llvm::ilist_node<SILProperty>,
   const Optional<KeyPathPatternComponent> &getComponent() const {
     return Component;
   }
-  
+
+  void clearReferencedFunctions_if(
+      llvm::function_ref<bool(SILFunction *)> predicate) {
+    if (Component) {
+      Component->clearReferencedFunctions_if(predicate);
+    }
+  }
+
   void print(SILPrintContext &Ctx) const;
   void dump() const;
 

lib/IRGen/GenKeyPath.cpp

@@ -119,7 +119,11 @@ getAccessorForComputedComponent(IRGenModule &IGM,
     accessor = component.getSubscriptIndexHash();
     break;
   }
-  
+
+  if (!accessor) {
+    return nullptr;
+  }
+
   // If the accessor is not generic, and locally available, we can use it as is.
   // If it's only externally available, we need a local thunk to relative-
   // reference.
@@ -510,10 +514,18 @@ getWitnessTableForComputedComponent(IRGenModule &IGM,
     fields.addNullPointer(IGM.FunctionPtrTy);
   fields.addSignedPointer(copy, schemaKeyPath,
                           PointerAuthEntity::Special::KeyPathCopy);
-  fields.addSignedPointer(equals, schemaKeyPath,
-                          PointerAuthEntity::Special::KeyPathEquals);
-  fields.addSignedPointer(hash, schemaKeyPath,
-                          PointerAuthEntity::Special::KeyPathHash);
+  if (equals)
+    fields.addSignedPointer(equals, schemaKeyPath,
+                            PointerAuthEntity::Special::KeyPathEquals);
+  else
+    fields.addNullPointer(IGM.FunctionPtrTy);
+
+  if (hash)
+    fields.addSignedPointer(hash, schemaKeyPath,
+                            PointerAuthEntity::Special::KeyPathHash);
+  else
+    fields.addNullPointer(IGM.FunctionPtrTy);
+
   return fields.finishAndCreateGlobal(
       "keypath_witnesses", IGM.getPointerAlignment(), /*constant*/ true,
       llvm::GlobalVariable::PrivateLinkage);
@@ -914,6 +926,11 @@ emitKeyPathComponent(IRGenModule &IGM,
     switch (id.getKind()) {
     case KeyPathPatternComponent::ComputedPropertyId::Function: {
       idKind = KeyPathComponentHeader::Pointer;
+      if (!id.getFunction()) {
+        idValue = nullptr;
+        idResolution = KeyPathComponentHeader::Resolved;
+        break;
+      }
       // FIXME: Does this need to be signed?
       auto idRef = IGM.getAddrOfLLVMVariableOrGOTEquivalent(
         LinkEntity::forSILFunction(id.getFunction(), false));
@@ -1066,7 +1083,7 @@ emitKeyPathComponent(IRGenModule &IGM,
     switch (idKind) {
     case KeyPathComponentHeader::Pointer:
       // Use a relative offset to the referent.
-      fields.addRelativeAddress(idValue);
+      fields.addRelativeAddressOrNull(idValue);
       break;
 
     case KeyPathComponentHeader::VTableOffset:
@@ -1077,15 +1094,12 @@ emitKeyPathComponent(IRGenModule &IGM,
     }
 
     // Push the accessors, possibly thunked to marshal generic environment.
-    fields.addRelativeAddress(
-      getAccessorForComputedComponent(IGM, component, Getter,
-                                      genericEnv, requirements,
-                                      hasSubscriptIndices));
+    fields.addRelativeAddressOrNull(getAccessorForComputedComponent(
+        IGM, component, Getter, genericEnv, requirements, hasSubscriptIndices));
     if (settable)
-      fields.addRelativeAddress(
-        getAccessorForComputedComponent(IGM, component, Setter,
-                                        genericEnv, requirements,
-                                        hasSubscriptIndices));
+      fields.addRelativeAddressOrNull(
+          getAccessorForComputedComponent(IGM, component, Setter, genericEnv,
+                                          requirements, hasSubscriptIndices));
 
     if (!isInstantiableOnce) {
       // If there's generic context or subscript indexes, embed as

lib/SIL/IR/SILInstructions.cpp

@@ -2530,6 +2530,57 @@ forEachRefcountableReference(const KeyPathPatternComponent &component,
   }
 }
 
+void KeyPathPatternComponent::clearReferencedFunctions_if(
+    llvm::function_ref<bool(SILFunction *)> predicate) {
+  switch (getKind()) {
+  case KeyPathPatternComponent::Kind::StoredProperty:
+  case KeyPathPatternComponent::Kind::OptionalChain:
+  case KeyPathPatternComponent::Kind::OptionalWrap:
+  case KeyPathPatternComponent::Kind::OptionalForce:
+  case KeyPathPatternComponent::Kind::TupleElement:
+    return;
+  case KeyPathPatternComponent::Kind::SettableProperty: {
+    auto setter = getComputedPropertySetter();
+    if (predicate(setter)) {
+      SetterAndIdKind.setPointer(nullptr);
+    }
+    LLVM_FALLTHROUGH;
+  }
+  case KeyPathPatternComponent::Kind::GettableProperty:
+    auto getter = getComputedPropertyGetter();
+    if (predicate(getter)) {
+      ValueAndKind.setPointer(nullptr);
+    }
+    switch (getComputedPropertyId().getKind()) {
+    case KeyPathPatternComponent::ComputedPropertyId::DeclRef:
+      // Mark the vtable entry as used somehow?
+      break;
+    case KeyPathPatternComponent::ComputedPropertyId::Function: {
+      auto idFn = getComputedPropertyId().getFunction();
+      if (predicate(idFn)) {
+        IdValue = ComputedPropertyId::ValueType();
+      }
+      break;
+    }
+    case KeyPathPatternComponent::ComputedPropertyId::Property:
+      break;
+    }
+
+    if (auto equals = getSubscriptIndexEquals()) {
+      if (predicate(equals)) {
+        IndexEquality.Equal = nullptr;
+      }
+    }
+
+    if (auto hash = getSubscriptIndexHash()) {
+      if (predicate(hash)) {
+        IndexEquality.Hash = nullptr;
+      }
+    }
+    return;
+  }
+}
+
 void KeyPathPatternComponent::incrementRefCounts() const {
   forEachRefcountableReference(*this,
     [&](SILFunction *f) { f->incrementRefCount(); });

lib/SILOptimizer/IPO/CrossDeadFunctionElimination.cpp

@@ -113,6 +113,17 @@ class SILCrossDeadFuncElimination : public SILModuleTransform {
     }
   }
 
+  void cleanupSILProperty(SILModule &M,
+                          const std::vector<SILFunction *> &DeadFunctions) {
+    std::set<SILFunction *> DeadFunctionsSet(DeadFunctions.begin(),
+                                             DeadFunctions.end());
+    for (SILProperty &P : M.getPropertyList()) {
+      P.clearReferencedFunctions_if([&](SILFunction *f) {
+        return DeadFunctionsSet.find(f) != DeadFunctionsSet.end();
+      });
+    }
+  }
+
   void run() override {
     LLVM_DEBUG(llvm::dbgs() << "Running CrossDeadFuncElimination\n");
     auto &Context = getModule()->getASTContext();
@@ -134,6 +145,7 @@ class SILCrossDeadFuncElimination : public SILModuleTransform {
     this->eliminateDeadEntriesFromTables(M);
     std::vector<SILFunction *> DeadFunctions;
     this->eliminateDeadFunctions(M, DeadFunctions);
+    this->cleanupSILProperty(M, DeadFunctions);
 
     while (!DeadFunctions.empty()) {
       SILFunction *F = DeadFunctions.back();