SIL: support specialized witness_method instructions

In Embedded Swift, witness method lookup is done from specialized witness tables.
For this to work, the type of witness_method must be specialized as well.
Otherwise the method call would be done with wrong parameter conventions (indirect instead of direct).

Author: eeckstein

include/swift/SIL/InstructionUtils.h

@@ -15,7 +15,7 @@
 
 #include "swift/SIL/InstWrappers.h"
 #include "swift/SIL/RuntimeEffect.h"
-#include "swift/SIL/SILInstruction.h"
+#include "swift/SIL/SILModule.h"
 
 namespace swift {
 
@@ -233,6 +233,9 @@ bool visitExplodedTupleType(SILType type,
 bool visitExplodedTupleValue(SILValue value,
                              llvm::function_ref<SILValue(SILValue, std::optional<unsigned>)> callback);
 
+std::pair<SILFunction *, SILWitnessTable *>
+lookUpFunctionInWitnessTable(WitnessMethodInst *wmi, SILModule::LinkingMode linkingMode);
+
 } // end namespace swift
 
 #endif

include/swift/SIL/SILInstruction.h

@@ -7715,6 +7715,13 @@ class WitnessMethodInst final
     return getMember().getDecl()->getDeclContext()->getSelfProtocolDecl();
   }
 
+  // Returns true if it's expected that the witness method is looked up up from
+  // a specialized witness table.
+  // This is the case in Embedded Swift.
+  bool isSpecialized() const {
+    return !getType().castTo<SILFunctionType>()->isPolymorphic();
+  }
+
   ProtocolConformanceRef getConformance() const { return Conformance; }
 };
 

include/swift/SIL/SILModule.h

@@ -853,6 +853,7 @@ class SILModule {
   std::pair<SILFunction *, SILWitnessTable *>
   lookUpFunctionInWitnessTable(ProtocolConformanceRef C,
                                SILDeclRef Requirement,
+                               bool lookupInSpecializedWitnessTable,
                                SILModule::LinkingMode linkingMode);
 
   /// Look up the SILDefaultWitnessTable representing the default witnesses

lib/SIL/IR/SILModule.cpp

@@ -531,6 +531,7 @@ SerializedSILLoader *SILModule::getSILLoader() {
 std::pair<SILFunction *, SILWitnessTable *>
 SILModule::lookUpFunctionInWitnessTable(ProtocolConformanceRef C,
                                         SILDeclRef Requirement,
+                                        bool lookupInSpecializedWitnessTable,
                                         SILModule::LinkingMode linkingMode) {
   if (!C.isConcrete())
     return {nullptr, nullptr};
@@ -539,7 +540,12 @@ SILModule::lookUpFunctionInWitnessTable(ProtocolConformanceRef C,
     SILLinkerVisitor linker(*this, linkingMode);
     linker.processConformance(C);
   }
-  SILWitnessTable *wt = lookUpWitnessTable(C.getConcrete()->getRootConformance());
+  ProtocolConformance *conf = C.getConcrete();
+  if (!isa<SpecializedProtocolConformance>(conf) || !lookupInSpecializedWitnessTable) {
+    conf = conf->getRootConformance();
+  }
+
+  SILWitnessTable *wt = lookUpWitnessTable(conf);
 
   if (!wt) {
     LLVM_DEBUG(llvm::dbgs() << "        Failed speculative lookup of "

lib/SIL/Utils/CalleeCache.cpp

@@ -12,6 +12,7 @@
 
 #include "swift/SIL/CalleeCache.h"
 #include "swift/SIL/SILModule.h"
+#include "swift/SIL/InstructionUtils.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/Basic/Assertions.h"
 #include "llvm/Support/Compiler.h"
@@ -256,8 +257,7 @@ CalleeCache::getSingleCalleeForWitnessMethod(WitnessMethodInst *WMI) const {
   SILWitnessTable *WT;
 
   // Attempt to find a specific callee for the given conformance and member.
-  std::tie(CalleeFn, WT) = WMI->getModule().lookUpFunctionInWitnessTable(
-      WMI->getConformance(), WMI->getMember(), SILModule::LinkingMode::LinkNormal);
+  std::tie(CalleeFn, WT) = lookUpFunctionInWitnessTable(WMI, SILModule::LinkingMode::LinkNormal);
 
   return CalleeFn;
 }

lib/SIL/Utils/InstructionUtils.cpp

@@ -1389,3 +1389,11 @@ bool swift::visitExplodedTupleValue(
 
   return true;
 }
+
+std::pair<SILFunction *, SILWitnessTable *>
+swift::lookUpFunctionInWitnessTable(WitnessMethodInst *wmi,
+                                    SILModule::LinkingMode linkingMode) {
+  SILModule &mod = wmi->getModule();
+  return mod.lookUpFunctionInWitnessTable(wmi->getConformance(), wmi->getMember(),
+                                          wmi->isSpecialized(), linkingMode);
+}

lib/SIL/Verifier/SILVerifier.cpp

@@ -4114,29 +4114,31 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
               == F.getModule().Types.getProtocolWitnessRepresentation(protocol),
             "result of witness_method must have correct representation for protocol");
 
-    require(methodType->isPolymorphic(),
-            "result of witness_method must be polymorphic");
-
-    auto genericSig = methodType->getInvocationGenericSignature();
-
-    auto selfGenericParam = genericSig.getGenericParams()[0];
-    require(selfGenericParam->getDepth() == 0
-            && selfGenericParam->getIndex() == 0,
-            "method should be polymorphic on Self parameter at depth 0 index 0");
-    std::optional<Requirement> selfRequirement;
-    for (auto req : genericSig.getRequirements()) {
-      if (req.getKind() != RequirementKind::SameType) {
-        selfRequirement = req;
-        break;
+    if (methodType->isPolymorphic()) {
+      require(methodType->isPolymorphic(),
+              "result of witness_method must be polymorphic");
+
+      auto genericSig = methodType->getInvocationGenericSignature();
+
+      auto selfGenericParam = genericSig.getGenericParams()[0];
+      require(selfGenericParam->getDepth() == 0
+              && selfGenericParam->getIndex() == 0,
+              "method should be polymorphic on Self parameter at depth 0 index 0");
+      std::optional<Requirement> selfRequirement;
+      for (auto req : genericSig.getRequirements()) {
+        if (req.getKind() != RequirementKind::SameType) {
+          selfRequirement = req;
+          break;
+        }
       }
-    }
 
-    require(selfRequirement &&
-            selfRequirement->getKind() == RequirementKind::Conformance,
-            "first non-same-typerequirement should be conformance requirement");
-    const auto protos = genericSig->getRequiredProtocols(selfGenericParam);
-    require(std::find(protos.begin(), protos.end(), protocol) != protos.end(),
-            "requirement Self parameter must conform to called protocol");
+      require(selfRequirement &&
+              selfRequirement->getKind() == RequirementKind::Conformance,
+              "first non-same-typerequirement should be conformance requirement");
+      const auto protos = genericSig->getRequiredProtocols(selfGenericParam);
+      require(std::find(protos.begin(), protos.end(), protocol) != protos.end(),
+              "requirement Self parameter must conform to called protocol");
+    }
 
     auto lookupType = AMI->getLookupType();
     if (getLocalArchetypeOf(lookupType) || lookupType->hasDynamicSelfType()) {

lib/SILOptimizer/Mandatory/DiagnoseInfiniteRecursion.cpp

@@ -122,8 +122,7 @@ static bool isRecursiveCall(FullApplySite applySite) {
   }
 
   if (auto *WMI = dyn_cast<WitnessMethodInst>(callee)) {
-    auto funcAndTable = parentFunc->getModule().lookUpFunctionInWitnessTable(
-        WMI->getConformance(), WMI->getMember(), SILModule::LinkingMode::LinkNormal);
+    auto funcAndTable = lookUpFunctionInWitnessTable(WMI, SILModule::LinkingMode::LinkNormal);
     return funcAndTable.first == parentFunc;
   }
   return false;

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -1134,8 +1134,6 @@ SILInstruction *SILCombiner::createApplyWithConcreteType(
     FullApplySite Apply,
     const llvm::SmallDenseMap<unsigned, ConcreteOpenedExistentialInfo> &COAIs,
     SILBuilderContext &BuilderCtx) {
-  // Ensure that the callee is polymorphic.
-  assert(Apply.getOrigCalleeType()->isPolymorphic());
 
   // Create the new set of arguments to apply including their substitutions.
   SubstitutionMap NewCallSubs = Apply.getSubstitutionMap();

lib/SILOptimizer/Utils/ConstExpr.cpp

@@ -433,7 +433,7 @@ SymbolicValue ConstExprFunctionState::computeConstantValue(SILValue value) {
                         UnknownReason::UnknownWitnessMethodConformance);
     auto &module = wmi->getModule();
     SILFunction *fn =
-        module.lookUpFunctionInWitnessTable(conf, wmi->getMember(),
+        module.lookUpFunctionInWitnessTable(conf, wmi->getMember(), wmi->isSpecialized(),
             SILModule::LinkingMode::LinkAll).first;
     // If we were able to resolve it, then we can proceed.
     if (fn)