SIL: add better support for specialized witness tables.

Store specialize witness tables in a separate lookup table in the module. This allows that for a normal conformance there can exist the original _and_ a specialized witness table.
Also, add a boolean property `isSpecialized` to `WitnessTable` which indicates whether the witness table is specialized or not.

Author: eeckstein

SwiftCompilerSources/Sources/SIL/WitnessTable.swift

@@ -111,6 +111,9 @@ public struct WitnessTable : CustomStringConvertible, NoReflectionChildren {
 
   public var isDefinition: Bool { !bridged.isDeclaration() }
 
+  // True, if this is a specialized witness table (currently only used in embedded mode).
+  public var isSpecialized: Bool { bridged.isSpecialized() }
+
   public var description: String {
     return String(taking: bridged.getDebugDescription())
   }

include/swift/SIL/SILBridging.h

@@ -1089,6 +1089,7 @@ struct BridgedWitnessTable {
   BRIDGED_INLINE SwiftInt getNumEntries() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedWitnessTableEntry getEntry(SwiftInt index) const;
   BRIDGED_INLINE bool isDeclaration() const;
+  BRIDGED_INLINE bool isSpecialized() const;
 };
 
 struct OptionalBridgedWitnessTable {

include/swift/SIL/SILBridgingImpl.h

@@ -1978,6 +1978,10 @@ bool BridgedWitnessTable::isDeclaration() const {
   return table->isDeclaration();
 }
 
+bool BridgedWitnessTable::isSpecialized() const {
+  return table->isSpecialized();
+}
+
 SwiftInt BridgedDefaultWitnessTable::getNumEntries() const {
   return SwiftInt(table->getEntries().size());
 }

include/swift/SIL/SILModule.h

@@ -255,9 +255,14 @@ class SILModule {
       VTableEntryCache;
 
   /// Lookup table for SIL witness tables from conformances.
-  llvm::DenseMap<const ProtocolConformance *, SILWitnessTable *>
+  llvm::DenseMap<const RootProtocolConformance *, SILWitnessTable *>
   WitnessTableMap;
 
+  /// Lookup table for specialized witness tables from conformances.
+  /// Currently only used in embedded mode.
+  llvm::DenseMap<const ProtocolConformance *, SILWitnessTable *>
+  specializedWitnessTableMap;
+
   /// The list of SILWitnessTables in the module.
   WitnessTableListType witnessTables;
 
@@ -866,16 +871,17 @@ class SILModule {
   /// i.e. it can be linked by linkFunction.
   bool hasFunction(StringRef Name);
 
-  /// Look up the SILWitnessTable representing the lowering of a protocol
-  /// conformance, and collect the substitutions to apply to the referenced
-  /// witnesses, if any.
-  ///
-  /// \arg C The protocol conformance mapped key to use to lookup the witness
-  ///        table.
-  /// \arg deserializeLazily If we cannot find the witness table should we
-  ///                        attempt to lazily deserialize it.
+  /// Look up the SILWitnessTable representing the lowering of a conformance `C`.
+  /// If a specialized witness table exists for the conformance, return the specialized table.
+  /// Specialized witness tables are currently only used in embedded mode.
   SILWitnessTable *lookUpWitnessTable(const ProtocolConformance *C);
 
+  /// Look up the SILWitnessTable representing the lowering of a conformance `C`.
+  /// The `isSpecialized` flag specifies if only non-specialized or specialized witness
+  /// tables are looked up.
+  /// Specialized witness tables are currently only used in embedded mode.
+  SILWitnessTable *lookUpWitnessTable(const ProtocolConformance *C, bool isSpecialized);
+
   /// Attempt to lookup \p Member in the witness table for \p C.
   ///
   /// Also, deserialize all referenced functions according to the \p linkgingMode.

include/swift/SIL/SILWitnessTable.h

@@ -181,7 +181,9 @@ class SILWitnessTable : public llvm::ilist_node<SILWitnessTable>,
   /// whether or not entries is empty since you can have an empty witness table
   /// that is not a declaration.
   bool IsDeclaration;
- 
+
+  bool specialized;
+
   /// Whether or not this witness table is serialized, which allows
   /// devirtualization from another module.
   unsigned SerializedKind : 2;
@@ -190,11 +192,12 @@ class SILWitnessTable : public llvm::ilist_node<SILWitnessTable>,
   SILWitnessTable(SILModule &M, SILLinkage Linkage, SerializedKind_t Serialized,
                   StringRef name, ProtocolConformance *conformance,
                   ArrayRef<Entry> entries,
-                  ArrayRef<ProtocolConformanceRef> conditionalConformances);
+                  ArrayRef<ProtocolConformanceRef> conditionalConformances,
+                  bool specialized);
 
   /// Private constructor for making SILWitnessTable declarations.
   SILWitnessTable(SILModule &M, SILLinkage Linkage, StringRef Name,
-                  ProtocolConformance *conformance);
+                  ProtocolConformance *conformance, bool specialized);
 
   void addWitnessTable();
 
@@ -203,11 +206,13 @@ class SILWitnessTable : public llvm::ilist_node<SILWitnessTable>,
   static SILWitnessTable *
   create(SILModule &M, SILLinkage Linkage, SerializedKind_t SerializedKind,
          ProtocolConformance *conformance, ArrayRef<Entry> entries,
-         ArrayRef<ProtocolConformanceRef> conditionalConformances);
+         ArrayRef<ProtocolConformanceRef> conditionalConformances,
+         bool specialized);
 
   /// Create a new SILWitnessTable declaration.
   static SILWitnessTable *create(SILModule &M, SILLinkage Linkage,
-                                 ProtocolConformance *conformance);
+                                 ProtocolConformance *conformance,
+                                 bool specialized);
 
   ~SILWitnessTable();
 
@@ -238,6 +243,9 @@ class SILWitnessTable : public llvm::ilist_node<SILWitnessTable>,
   /// Returns true if this witness table is a definition.
   bool isDefinition() const { return !isDeclaration(); }
 
+  // Returns true, if this is a specialized witness table (currently only used in embedded mode).
+  bool isSpecialized() const { return specialized; }
+
   /// Returns true if this witness table is going to be (or was) serialized.
   bool isSerialized() const {
     return SerializedKind_t(SerializedKind) == IsSerialized;

include/swift/SILOptimizer/OptimizerBridgingImpl.h

@@ -453,7 +453,7 @@ BridgedWitnessTable BridgedPassContext::createWitnessTable(BridgedLinkage linkag
   return {swift::SILWitnessTable::create(*mod, (swift::SILLinkage)linkage,
                                          serialized ? swift::IsSerialized : swift::IsNotSerialized,
                                          conformance.unbridged().getConcrete(),
-                                         entries, {})};
+                                         entries, {}, /*specialized=*/true)};
 }
 
 BridgedVTable BridgedPassContext::createSpecializedVTable(BridgedType classType,

lib/AST/ASTMangler.cpp

@@ -299,7 +299,7 @@ std::string ASTMangler::mangleWitnessTable(const ProtocolConformance *C) {
   if (auto *sc = dyn_cast<SpecializedProtocolConformance>(C)) {
     appendProtocolConformance(sc);
     appendOperator("WP");
-  } else if (isa<NormalProtocolConformance>(C)) {
+  } else if (isa<NormalProtocolConformance>(C) || isa<InheritedProtocolConformance>(C)) {
     appendProtocolConformance(C);
     appendOperator("WP");
   } else if (isa<SelfProtocolConformance>(C)) {

lib/IRGen/GenProto.cpp

@@ -1387,9 +1387,9 @@ class DirectConformanceInfo : public ConformanceInfo {
 class SpecializedConformanceInfo : public ConformanceInfo {
   friend ProtocolInfo;
 
-  const SpecializedProtocolConformance *Conformance;
+  const ProtocolConformance *Conformance;
 public:
-  SpecializedConformanceInfo(const SpecializedProtocolConformance *C)
+  SpecializedConformanceInfo(const ProtocolConformance *C)
       : Conformance(C) {}
 
   llvm::Value *getTable(IRGenFunction &IGF,
@@ -2608,8 +2608,15 @@ IRGenModule::getConformanceInfo(const ProtocolDecl *protocol,
   const ConformanceInfo *info;
 
   auto *specConf = conformance;
-  if (auto *inheritedC = dyn_cast<InheritedProtocolConformance>(conformance))
+  if (auto *inheritedC = dyn_cast<InheritedProtocolConformance>(conformance)) {
+    SILWitnessTable *wt = getSILModule().lookUpWitnessTable(inheritedC);
+    if (wt && wt->getConformance() == inheritedC) {
+      info = new SpecializedConformanceInfo(inheritedC);
+      Conformances.try_emplace(conformance, info);
+      return *info;
+    }
     specConf = inheritedC->getInheritedConformance();
+  }
 
   // If there is a specialized SILWitnessTable for the specialized conformance,
   // directly use it.

lib/SIL/IR/Linker.cpp

@@ -299,7 +299,8 @@ void SILLinkerVisitor::visitProtocolConformance(
 
       SILLinkage linkage = getLinkageForProtocolConformance(rootC, NotForDefinition);
       WT = SILWitnessTable::create(Mod, linkage,
-                                   const_cast<RootProtocolConformance *>(rootC));
+                                   const_cast<RootProtocolConformance *>(rootC),
+                                   /*specialized=*/ false);
     }
     // If the module is at or past the Lowered stage, then we can't do any
     // further deserialization, since pre-IRGen SIL lowering changes the types

lib/SIL/IR/SILModule.cpp

@@ -240,14 +240,28 @@ void SILModule::flushDeletedInsts() {
 
 SILWitnessTable *
 SILModule::lookUpWitnessTable(const ProtocolConformance *C) {
-  assert(C && "null conformance passed to lookUpWitnessTable");
+  // First try to lookup a specialized witness table for that conformance.
+  if (auto *wt = lookUpWitnessTable(C, /*isSpecialized=*/true)) {
+    return wt;
+  }
+  return lookUpWitnessTable(C, /*isSpecialized=*/false);
+}
 
-  // Attempt to lookup the witness table from the table.
-  auto found = WitnessTableMap.find(C);
-  if (found == WitnessTableMap.end())
-    return nullptr;
+SILWitnessTable *
+SILModule::lookUpWitnessTable(const ProtocolConformance *C, bool isSpecialized) {
+  assert(C && "null conformance passed to lookUpWitnessTable");
 
-  return found->second;
+  if (isSpecialized) {
+    // First try to lookup a specialized witness table for that conformance.
+    auto foundSpec = specializedWitnessTableMap.find(C);
+    if (foundSpec != specializedWitnessTableMap.end())
+      return foundSpec->second;
+  } else if (auto *rootConf = dyn_cast<RootProtocolConformance>(C)) {
+    auto found = WitnessTableMap.find(rootConf);
+    if (found != WitnessTableMap.end())
+      return found->second;
+  }
+  return nullptr;
 }
 
 SILDefaultWitnessTable *
@@ -287,7 +301,9 @@ void SILModule::deleteWitnessTable(SILWitnessTable *Wt) {
   auto Conf = Wt->getConformance();
   assert(lookUpWitnessTable(Conf) == Wt);
   getSILLoader()->invalidateWitnessTable(Wt);
-  WitnessTableMap.erase(Conf);
+  specializedWitnessTableMap.erase(Conf);
+  if (auto *rootConf = dyn_cast<RootProtocolConformance>(Conf))
+    WitnessTableMap.erase(rootConf);
   witnessTables.erase(Wt);
 }
 
@@ -579,15 +595,24 @@ SILModule::lookUpFunctionInWitnessTable(ProtocolConformanceRef C,
     linker.processConformance(C);
   }
   ProtocolConformance *conf = C.getConcrete();
-  if (auto *inheritedC = dyn_cast<InheritedProtocolConformance>(conf))
-    conf = inheritedC->getInheritedConformance();
-
-  if (!isa<SpecializedProtocolConformance>(conf) || !lookupInSpecializedWitnessTable) {
-    conf = conf->getRootConformance();
+  SILWitnessTable *wt = nullptr;
+
+  if (lookupInSpecializedWitnessTable) {
+    wt = lookUpWitnessTable(conf);
+    if (!wt) {
+      if (auto *inheritedC = dyn_cast<InheritedProtocolConformance>(conf)) {
+        conf = inheritedC->getInheritedConformance();
+        wt = lookUpWitnessTable(conf);
+      }
+      if (!wt && !isa<SpecializedProtocolConformance>(conf)) {
+        conf = conf->getRootConformance();
+        wt = lookUpWitnessTable(conf);
+      }
+    }
+  } else {
+    wt = lookUpWitnessTable(conf->getRootConformance());
   }
 
-  SILWitnessTable *wt = lookUpWitnessTable(conf);
-
   if (!wt) {
     LLVM_DEBUG(llvm::dbgs() << "        Failed speculative lookup of "
                "witness for: ";