Embedded: specialized witness tables, part2: support remaining cases of class existentials

SwiftCompilerSources/Sources/AST/Conformance.swift

@@ -47,6 +47,16 @@ public struct Conformance: CustomStringConvertible, NoReflectionChildren {
     return bridged.getGenericConformance().conformance
   }
 
+  public var isInherited: Bool {
+    assert(isConcrete)
+    return bridged.isInheritedConformance()
+  }
+
+  public var inheritedConformance: Conformance {
+    assert(isInherited)
+    return bridged.getInheritedConformance().conformance
+  }
+
   public var specializedSubstitutions: SubstitutionMap {
     assert(isSpecialized)
     return SubstitutionMap(bridged: bridged.getSpecializedSubstitutions())

SwiftCompilerSources/Sources/AST/SubstitutionMap.swift

@@ -18,7 +18,7 @@ import ASTBridging
 ///
 /// Substitution maps are primarily used when performing substitutions into any entity that
 /// can reference type parameters and conformances.
-public struct SubstitutionMap {
+public struct SubstitutionMap: CustomStringConvertible {
   public let bridged: BridgedSubstitutionMap
 
   public init(bridged: BridgedSubstitutionMap) {
@@ -29,6 +29,10 @@ public struct SubstitutionMap {
     self.bridged = BridgedSubstitutionMap()
   }
 
+  public var description: String {
+    return String(taking: bridged.getDebugDescription())
+  }
+
   public var isEmpty: Bool { bridged.isEmpty() }
 
   public var hasAnySubstitutableParams: Bool { bridged.hasAnySubstitutableParams() }

SwiftCompilerSources/Sources/AST/Type.swift

@@ -37,6 +37,10 @@ public struct Type: CustomStringConvertible, NoReflectionChildren {
   public var isEscapable: Bool { bridged.isEscapable() }
   public var isNoEscape: Bool { bridged.isNoEscape() }
   public var isInteger: Bool { bridged.isInteger() }
+
+  public func subst(with substitutionMap: SubstitutionMap) -> Type {
+    return Type(bridged: bridged.subst(substitutionMap.bridged))
+  }
 }
 
 /// A Type that is statically known to be canonical.
@@ -55,4 +59,8 @@ public struct CanonicalType: CustomStringConvertible, NoReflectionChildren {
   public var isEscapable: Bool { type.isEscapable }
   public var isNoEscape: Bool { type.isNoEscape }
   public var isInteger: Bool { type.isInteger }
+
+  public func subst(with substitutionMap: SubstitutionMap) -> CanonicalType {
+    return type.subst(with: substitutionMap).canonical
+  }
 }

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -116,6 +116,10 @@ private func optimize(function: Function, _ context: FunctionPassContext, _ modu
         if context.options.enableEmbeddedSwift {
           _ = context.specializeClassMethodInst(classMethod)
         }
+      case let witnessMethod as WitnessMethodInst:
+        if context.options.enableEmbeddedSwift {
+          _ = context.specializeWitnessMethodInst(witnessMethod)
+        }
 
       case let initExRef as InitExistentialRefInst:
         if context.options.enableEmbeddedSwift {
@@ -267,12 +271,14 @@ private func shouldInline(apply: FullApplySite, callee: Function, alreadyInlined
 private func specializeWitnessTables(for initExRef: InitExistentialRefInst, _ context: ModulePassContext,
                                      _ worklist: inout FunctionWorklist)
 {
-  for conformance in initExRef.conformances where conformance.isConcrete {
+  for c in initExRef.conformances where c.isConcrete {
+    let conformance = c.isInherited ? c.inheritedConformance : c
     let origWitnessTable = context.lookupWitnessTable(for: conformance)
     if conformance.isSpecialized {
       if origWitnessTable == nil {
-        let wt = specializeWitnessTable(forConformance: conformance, errorLocation: initExRef.location, context)
-        worklist.addWitnessMethods(of: wt)
+        specializeWitnessTable(forConformance: conformance, errorLocation: initExRef.location, context) {
+          worklist.addWitnessMethods(of: $0)
+        }
       }
     } else if let origWitnessTable {
       checkForGenericMethods(in: origWitnessTable, errorLocation: initExRef.location, context)

SwiftCompilerSources/Sources/Optimizer/PassManager/Context.swift

@@ -74,6 +74,14 @@ extension Context {
     return _bridged.lookupSpecializedVTable(classType.bridged).vTable
   }
 
+  func getSpecializedConformance(of genericConformance: Conformance,
+                                 for type: AST.`Type`,
+                                 substitutions: SubstitutionMap) -> Conformance
+  {
+    let c = _bridged.getSpecializedConformance(genericConformance.bridged, type.bridged, substitutions.bridged)
+    return Conformance(bridged: c)
+  }
+
   func notifyNewFunction(function: Function, derivedFrom: Function) {
     _bridged.addFunctionToPassManagerWorklist(function.bridged, derivedFrom.bridged)
   }
@@ -349,6 +357,15 @@ struct FunctionPassContext : MutatingContext {
     return false
   }
 
+  func specializeWitnessMethodInst(_ wm: WitnessMethodInst) -> Bool {
+    if _bridged.specializeWitnessMethodInst(wm.bridged) {
+      notifyInstructionsChanged()
+      notifyCallsChanged()
+      return true
+    }
+    return false
+  }
+
   func specializeApplies(in function: Function, isMandatory: Bool) -> Bool {
     if _bridged.specializeAppliesInFunction(function.bridged, isMandatory) {
       notifyInstructionsChanged()

SwiftCompilerSources/Sources/Optimizer/Utilities/GenericSpecialization.swift

@@ -78,7 +78,8 @@ func specializeVTablesOfSuperclasses(_ moduleContext: ModulePassContext) {
 
 func specializeWitnessTable(forConformance conformance: Conformance,
                             errorLocation: Location,
-                            _ context: ModulePassContext) -> WitnessTable
+                            _ context: ModulePassContext,
+                            _ notifyNewWitnessTable: (WitnessTable) -> ())
 {
   let genericConformance = conformance.genericConformance
   guard let witnessTable = context.lookupWitnessTable(for: genericConformance) else {
@@ -88,6 +89,8 @@ func specializeWitnessTable(forConformance conformance: Conformance,
 
   let newEntries = witnessTable.entries.map { origEntry in
     switch origEntry {
+    case .invalid:
+      return WitnessTable.Entry.invalid
     case .method(let requirement, let witness):
       guard let origMethod = witness else {
         return origEntry
@@ -101,11 +104,23 @@ func specializeWitnessTable(forConformance conformance: Conformance,
         return origEntry
       }
       return .method(requirement: requirement, witness: specializedMethod)
-    default:
-      // TODO: handle other witness table entry kinds
-      fatalError("unsupported witness table etnry")
+    case .baseProtocol(let requirement, let witness):
+      let baseConf = context.getSpecializedConformance(of: witness,
+                                                       for: conformance.type,
+                                                       substitutions: conformance.specializedSubstitutions)
+      specializeWitnessTable(forConformance: baseConf, errorLocation: errorLocation, context, notifyNewWitnessTable)
+      return .baseProtocol(requirement: requirement, witness: baseConf)
+    case .associatedType(let requirement, let witness):
+      let substType = witness.subst(with: conformance.specializedSubstitutions)
+      return .associatedType(requirement: requirement, witness: substType)
+    case .associatedConformance(let requirement, let proto, let witness):
+      if witness.isSpecialized {
+        specializeWitnessTable(forConformance: witness, errorLocation: errorLocation, context, notifyNewWitnessTable)
+      }
+      return .associatedConformance(requirement: requirement, protocol: proto, witness: witness)
     }
-    return origEntry
   }
-  return context.createWitnessTable(entries: newEntries, conformance: conformance, linkage: .shared, serialized: false)
+  let newWT = context.createWitnessTable(entries: newEntries,conformance: conformance,
+                                         linkage: .shared, serialized: false)
+  notifyNewWitnessTable(newWT)
 }

docs/EmbeddedSwift/ABI.md

@@ -20,14 +20,31 @@ The compiler respects the ABIs and calling conventions of C and C++ when interop
 
 ## Metadata ABI of Embedded Swift
 
-Embedded Swift eliminates almost all metadata compared to full Swift. However, class metadata is still used, because those serve as vtables for dynamic dispatch of methods to implement runtime polymorphism. The layout of Embedded Swift's class metadata is *different* from full Swift:
+Embedded Swift eliminates almost all metadata compared to full Swift. However, class and existential metadata are still used, because those serve as vtables and witness tables for dynamic dispatch of methods to implement runtime polymorphism with classes and existentials.
+
+### Class Metadata ABI
+
+The layout of Embedded Swift's class metadata is *different* from full Swift:
 
 - The **super pointer** pointing to the class metadata record for the superclass is stored at **offset 0**. If the class is a root class, it is null.
 - The **destructor pointer** is stored at **offset 1**. This function is invoked by Swift's deallocator when the class instance is destroyed.
 - The **ivar destroyer** is stored at **offset 2**. This function is invoked to destroy instance members when creation of the object is cancelled (e.g. in a failable initializer).
 - Lastly, the **vtable** is stored at **offset 3**: For each Swift class in the class's inheritance hierarchy, in order starting
   from the root class and working down to the most derived class, the function pointers to the implementation of every method of the class in declaration order in stored.
 
+### Witness Tables ABI
+
+The layout of Embedded Swift's witness tables is *different* from full Swift:
+
+- The first word is always a null pointer (TODO: it can be eliminated)
+- The following words are witness table entries which can be one of the following:
+  - A method witness: a pointer to the witness function.
+  - An associated conformance witness: a pointer to the witness table of the associated conformance
+
+Note that witness tables in Embedded Swift do not contain associated type entries.
+
+Witness functions are always specialized for concrete types. This also means that parameters and return values are passed directly (if possible).
+
 ## Heap object layout in Embedded Swift
 
 Heap objects have the following layout in Embedded Swift:

include/swift/AST/ASTBridging.h

@@ -2004,6 +2004,7 @@ struct BridgedASTType {
   BRIDGED_INLINE bool isEscapable() const;
   BRIDGED_INLINE bool isNoEscape() const;
   BRIDGED_INLINE bool isInteger() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType subst(BridgedSubstitutionMap substMap) const;
 };
 
 class BridgedCanType {
@@ -2031,8 +2032,10 @@ struct BridgedConformance {
   BRIDGED_INLINE bool isConcrete() const;
   BRIDGED_INLINE bool isValid() const;
   BRIDGED_INLINE bool isSpecializedConformance() const;
+  BRIDGED_INLINE bool isInheritedConformance() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedASTType getType() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedConformance getGenericConformance() const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedConformance getInheritedConformance() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedSubstitutionMap getSpecializedSubstitutions() const;
 };
 
@@ -2060,6 +2063,7 @@ struct BridgedSubstitutionMap {
   BRIDGED_INLINE BridgedSubstitutionMap(swift::SubstitutionMap map);
   BRIDGED_INLINE swift::SubstitutionMap unbridged() const;
   BRIDGED_INLINE BridgedSubstitutionMap();
+  BridgedOwnedString getDebugDescription() const;
   BRIDGED_INLINE bool isEmpty() const;
   BRIDGED_INLINE bool hasAnySubstitutableParams() const;
   BRIDGED_INLINE SwiftInt getNumConformances() const;

include/swift/AST/ASTBridgingImpl.h

@@ -126,6 +126,10 @@ bool BridgedASTType::isInteger() const {
   return unbridged()->is<swift::IntegerType>();
 }
 
+BridgedASTType BridgedASTType::subst(BridgedSubstitutionMap substMap) const {
+  return {unbridged().subst(substMap.unbridged()).getPointer()};
+}
+
 //===----------------------------------------------------------------------===//
 // MARK: BridgedCanType
 //===----------------------------------------------------------------------===//
@@ -159,6 +163,10 @@ bool BridgedConformance::isSpecializedConformance() const {
   return swift::isa<swift::SpecializedProtocolConformance>(unbridged().getConcrete());
 }
 
+bool BridgedConformance::isInheritedConformance() const {
+  return swift::isa<swift::InheritedProtocolConformance>(unbridged().getConcrete());
+}
+
 BridgedASTType BridgedConformance::getType() const {
   return {unbridged().getConcrete()->getType().getPointer()};
 }
@@ -168,6 +176,11 @@ BridgedConformance BridgedConformance::getGenericConformance() const {
   return {swift::ProtocolConformanceRef(specPC->getGenericConformance())};
 }
 
+BridgedConformance BridgedConformance::getInheritedConformance() const {
+  auto *inheritedConf = swift::cast<swift::InheritedProtocolConformance>(unbridged().getConcrete());
+  return {swift::ProtocolConformanceRef(inheritedConf->getInheritedConformance())};
+}
+
 BridgedSubstitutionMap BridgedConformance::getSpecializedSubstitutions() const {
   auto *specPC = swift::cast<swift::SpecializedProtocolConformance>(unbridged().getConcrete());
   return {specPC->getSubstitutionMap()};

include/swift/AST/IRGenOptions.h

@@ -120,6 +120,8 @@ struct PointerAuthOptions : clang::PointerAuthOptions {
 
   /// Swift protocol witness table associated conformance witness table
   /// access functions.
+  /// In Embedded Swift used for associated conformance witness table
+  /// pointers.
   PointerAuthSchema ProtocolAssociatedTypeWitnessTableAccessFunctions;
 
   /// Swift class v-table functions.

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/SILCloner.h

@@ -2682,7 +2682,7 @@ SILCloner<ImplClass>::visitWitnessMethodInst(WitnessMethodInst *Inst) {
   recordClonedInstruction(Inst,
                           getBuilder().createWitnessMethod(
                               getOpLocation(Inst->getLoc()), newLookupType,
-                              conformance, Inst->getMember(), Inst->getType()));
+                              conformance, Inst->getMember(), getOpType(Inst->getType())));
 }
 
 template<typename ImplClass>

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

include/swift/SILOptimizer/OptimizerBridging.h

@@ -256,6 +256,7 @@ struct BridgedPassContext {
                                                                  BridgedSubstitutionMap substitutions) const;
   void deserializeAllCallees(BridgedFunction function, bool deserializeAll) const;
   bool specializeClassMethodInst(BridgedInstruction cm) const;
+  bool specializeWitnessMethodInst(BridgedInstruction wm) const;
   bool specializeAppliesInFunction(BridgedFunction function, bool isMandatory) const;
   BridgedOwnedString mangleOutlinedVariable(BridgedFunction function) const;
   BridgedOwnedString mangleAsyncRemoved(BridgedFunction function) const;
@@ -332,6 +333,10 @@ struct BridgedPassContext {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE
   OptionalBridgedVTable lookupSpecializedVTable(BridgedType classType) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE
+  BridgedConformance getSpecializedConformance(BridgedConformance genericConformance,
+                                                       BridgedASTType type,
+                                                       BridgedSubstitutionMap substitutions) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE
   OptionalBridgedWitnessTable lookupWitnessTable(BridgedConformance conformance) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedWitnessTable createWitnessTable(BridgedLinkage linkage,
                                                                             bool serialized,

include/swift/SILOptimizer/OptimizerBridgingImpl.h

@@ -417,6 +417,16 @@ OptionalBridgedVTable BridgedPassContext::lookupSpecializedVTable(BridgedType cl
   return {mod->lookUpSpecializedVTable(classType.unbridged())};
 }
 
+BridgedConformance BridgedPassContext::getSpecializedConformance(
+                                                     BridgedConformance genericConformance,
+                                                     BridgedASTType type,
+                                                     BridgedSubstitutionMap substitutions) const {
+  auto &ctxt = invocation->getPassManager()->getModule()->getASTContext();
+  auto *genConf = llvm::cast<swift::NormalProtocolConformance>(genericConformance.unbridged().getConcrete());
+  auto *c = ctxt.getSpecializedConformance(type.unbridged(), genConf, substitutions.unbridged());
+  return swift::ProtocolConformanceRef(c);
+}
+
 OptionalBridgedWitnessTable BridgedPassContext::lookupWitnessTable(BridgedConformance conformance) const {
   swift::ProtocolConformanceRef ref = conformance.unbridged();
   if (!ref.isConcrete()) {

include/swift/SILOptimizer/Utils/InstOptUtils.h

@@ -607,6 +607,7 @@ bool optimizeMemoryAccesses(SILFunction *fn);
 bool eliminateDeadAllocations(SILFunction *fn, DominanceInfo *domInfo);
 
 bool specializeClassMethodInst(ClassMethodInst *cm);
+bool specializeWitnessMethodInst(WitnessMethodInst *wm);
 
 bool specializeAppliesInFunction(SILFunction &F,
                                  SILTransform *transform,

lib/AST/ASTBridging.cpp

@@ -2744,16 +2744,31 @@ void BridgedTypeRepr_dump(void *type) { static_cast<TypeRepr *>(type)->dump(); }
 
 #pragma clang diagnostic pop
 
+//===----------------------------------------------------------------------===//
+// MARK: Conformance
+//===----------------------------------------------------------------------===//
+
 BridgedOwnedString BridgedConformance::getDebugDescription() const {
   std::string str;
   llvm::raw_string_ostream os(str);
   unbridged().print(os);
   return str;
 }
 
+//===----------------------------------------------------------------------===//
+// MARK: SubstitutionMap
+//===----------------------------------------------------------------------===//
+
 static_assert(sizeof(BridgedSubstitutionMap) >= sizeof(swift::SubstitutionMap),
               "BridgedSubstitutionMap has wrong size");
 
+BridgedOwnedString BridgedSubstitutionMap::getDebugDescription() const {
+  std::string str;
+  llvm::raw_string_ostream os(str);
+  unbridged().dump(os);
+  return str;
+}
+
 //===----------------------------------------------------------------------===//
 // MARK: Plugins
 //===----------------------------------------------------------------------===//