Naive implementation of support for VTable expansion in extensions

Author: Brian King

lib/SILGen/SILGen.cpp

@@ -1466,6 +1466,8 @@ SILModule::constructSIL(ModuleDecl *mod, SILOptions &options, FileUnit *SF,
     if (hasSIB)
       M->getSILLoader()->getAllForModule(mod->getName(), nullptr);
   }
+  
+  SGM.emitVTables();
 
   // Emit external definitions used by this module.
   for (size_t i = 0, e = mod->getASTContext().LastCheckedExternalDefinition;

lib/SILGen/SILGen.h

@@ -63,6 +63,12 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   llvm::DenseMap<SILDeclRef, SILFunction*> emittedFunctions;
   /// Mapping from ProtocolConformances to emitted SILWitnessTables.
   llvm::DenseMap<ProtocolConformance*, SILWitnessTable*> emittedWitnessTables;
+  
+  /// VTable entry storage that is built as the module is generated
+  llvm::DenseMap<ClassDecl *, std::vector<SILVTable::Entry> *> VTableEntryMap;
+  
+  /// Set of classes for which vtables need to be emitted.
+  llvm::DenseSet<ClassDecl *> VTablesPending;
 
   struct DelayedFunction {
     /// Insert the entity after the given function when it's emitted.
@@ -292,6 +298,9 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
   /// Add a global variable to the SILModule.
   void addGlobalVariable(VarDecl *global);
 
+  /// Emit SIL related to the processed VTables
+  void emitVTables();
+  
   /// Emit SIL related to a Clang-imported declaration.
   void emitExternalDefinition(Decl *d);
 

lib/SILGen/SILGenDecl.cpp

@@ -1222,6 +1222,14 @@ void SILGenModule::emitExternalWitnessTable(ProtocolConformance *c) {
   lastEmittedConformance = root;
 }
 
+void SILGenModule::emitVTables() {
+  for (auto *c : VTablesPending) {
+    auto vtableEntries = VTableEntryMap[c];
+    ClassDecl *theClass = c;
+    SILVTable::create(M, theClass, *vtableEntries);
+  }
+}
+
 void SILGenModule::emitExternalDefinition(Decl *d) {
   switch (d->getKind()) {
   case DeclKind::Func: {

lib/SILGen/SILGenType.cpp

@@ -151,20 +151,27 @@ class SILGenVTable : public Lowering::ASTVisitor<SILGenVTable> {
 public:
   SILGenModule &SGM;
   ClassDecl *theClass;
-  std::vector<SILVTable::Entry> vtableEntries;
+  
+  std::vector<SILVTable::Entry> &getVTableEntries() {
+    auto *entries = SGM.VTableEntryMap.lookup(theClass);
+    if (!entries) {
+      entries = new std::vector<SILVTable::Entry>();
+      SGM.VTableEntryMap.insert({theClass, entries});
+    }
+    return *entries;
+  }
 
   SILGenVTable(SILGenModule &SGM, ClassDecl *theClass)
-    : SGM(SGM), theClass(theClass)
-  {
+    : SGM(SGM), theClass(theClass) {}
+
+  void visitAncestors() {
     // Populate the superclass members, if any.
     Type super = theClass->getSuperclass();
     if (super && super->getClassOrBoundGenericClass())
       visitAncestor(super->getClassOrBoundGenericClass());
-  }
-
-  ~SILGenVTable() {
-    // Create the vtable.
-    SILVTable::create(SGM.M, theClass, vtableEntries);
+      
+    // Store an iterable class
+    SGM.VTablesPending.insert(theClass);
   }
 
   void visitAncestor(ClassDecl *ancestor) {
@@ -204,7 +211,7 @@ class SILGenVTable : public Lowering::ASTVisitor<SILGenVTable> {
     // NB: Mutates vtableEntries in-place
     // FIXME: O(n^2)
     if (auto overridden = member.getNextOverriddenVTableEntry()) {
-      for (SILVTable::Entry &entry : vtableEntries) {
+      for (SILVTable::Entry &entry : getVTableEntries()) {
         SILDeclRef ref = overridden;
 
         do {
@@ -230,7 +237,7 @@ class SILGenVTable : public Lowering::ASTVisitor<SILGenVTable> {
       return;
 
     // Otherwise, introduce a new vtable entry.
-    vtableEntries.push_back(getVtableEntry(member));
+    getVTableEntries().push_back(getVtableEntry(member));
   }
 
   // Default for members that don't require vtable entries.
@@ -325,8 +332,10 @@ class SILGenType : public TypeMemberVisitor<SILGenType> {
   /// Emit SIL functions for all the members of the type.
   void emitType() {
     // Start building a vtable if this is a class.
-    if (auto theClass = dyn_cast<ClassDecl>(theType))
+    if (auto theClass = dyn_cast<ClassDecl>(theType)) {
       genVTable.emplace(SGM, theClass);
+      genVTable->visitAncestors();
+    }
 
     for (Decl *member : theType->getMembers()) {
       if (genVTable)
@@ -430,14 +439,24 @@ void SILGenModule::visitNominalTypeDecl(NominalTypeDecl *ntd) {
 class SILGenExtension : public TypeMemberVisitor<SILGenExtension> {
 public:
   SILGenModule &SGM;
+  Optional<SILGenVTable> genVTable;
 
   SILGenExtension(SILGenModule &SGM)
     : SGM(SGM) {}
 
   /// Emit SIL functions for all the members of the extension.
   void emitExtension(ExtensionDecl *e) {
-    for (Decl *member : e->getMembers())
+    
+    if (auto theClass = e->getAsClassOrClassExtensionContext())
+      genVTable.emplace(SGM, theClass);
+
+    for (Decl *member : e->getMembers()) {
+      if (genVTable) {
+        genVTable->visit(member);
+      }
+    
       visit(member);
+    }
 
     if (!e->getExtendedType()->isExistentialType()) {
       // Emit witness tables for protocol conformances introduced by the

lib/Sema/TypeCheckDecl.cpp

@@ -6046,14 +6046,14 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
 
     // Non-Objective-C declarations in extensions cannot override or
     // be overridden.
-    if ((base->getDeclContext()->isExtensionContext() ||
-         override->getDeclContext()->isExtensionContext()) &&
-        !base->isObjC() && !isKnownObjC) {
-      TC.diagnose(override, diag::override_decl_extension,
-                  !override->getDeclContext()->isExtensionContext());
-      TC.diagnose(base, diag::overridden_here);
-      return true;
-    }
+//    if ((base->getDeclContext()->isExtensionContext() ||
+//         override->getDeclContext()->isExtensionContext()) &&
+//        !base->isObjC() && !isKnownObjC) {
+//      TC.diagnose(override, diag::override_decl_extension,
+//                  !override->getDeclContext()->isExtensionContext());
+//      TC.diagnose(base, diag::overridden_here);
+//      return true;
+//    }
 
     // If the overriding declaration does not have the 'override' modifier on
     // it, complain.
@@ -6067,11 +6067,30 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
             .fixItInsert(override->getStartLoc(), "override ");
       else
         TC.diagnose(override, diag::missing_override);
+
       TC.diagnose(base, diag::overridden_here);
       override->getAttrs().add(
           new (TC.Context) OverrideAttr(SourceLoc()));
     }
 
+    // If the overridden method is declared in a Swift Class Declaration,
+    // dispatch will use table dispatch. If the override is in an extension
+    // warn, since it is not added to the class vtable.
+    //
+    // FIXME: Only warn if the extension is in another module, and if
+    // it is in the same module, update the vtable.
+//    if (auto *baseDecl = dyn_cast<ClassDecl>(base->getDeclContext())) {
+//      if (baseDecl->hasKnownSwiftImplementation() && 
+//          !base->isDynamic() &&
+//          override->getDeclContext()->isExtensionContext()) {
+//        // For compatibility, only generate a warning in Swift 3
+//        TC.diagnose(override, (TC.Context.isSwiftVersion3()
+//          ? diag::override_class_declaration_in_extension_warning
+//          : diag::override_class_declaration_in_extension));
+//        TC.diagnose(base, diag::overridden_here);
+//      }
+//    }
+    
     // If the overriding declaration is 'throws' but the base is not,
     // complain.
     if (auto overrideFn = dyn_cast<AbstractFunctionDecl>(override)) {

test/SILGen/vtables.swift

@@ -9,11 +9,6 @@ import gizmo
 
 // Test for compilation order independence
 class C : B {
-  // foo inherited from B
-  override func bar() {}
-  // bas inherited from A
-  override func qux() {}
-
   // zim inherited from B
   override func zang() {}
 
@@ -22,6 +17,13 @@ class C : B {
   func flopsy() {}
   func mopsy() {}
 }
+
+extension C {
+  // foo inherited from B
+  override func bar() {}
+  // bas inherited from A
+  override func qux() {}
+}
 // CHECK: sil_vtable C {
 // CHECK:   #A.foo!1: _TFC7vtables1B3foo
 // CHECK:   #A.bar!1: _TFC7vtables1C3bar