Merge pull request #5053 from shajrawi/vtable_entry_cache

Add a sidetable cache in the module for VTable entries, replacing linear lookup

Author: Joe Shajrawi

include/swift/SIL/SILModule.h

@@ -132,6 +132,10 @@ class SILModule {
   /// The list of SILVTables in the module.
   VTableListType vtables;
 
+  /// This is a cache of vtable entries for quick look-up
+  llvm::DenseMap<std::pair<const SILVTable *, SILDeclRef>, SILFunction *>
+      VTableEntryCache;
+
   /// Lookup table for SIL witness tables from conformances.
   llvm::DenseMap<const NormalProtocolConformance *, SILWitnessTable *>
   WitnessTableMap;

include/swift/SIL/SILVTable.h

@@ -86,6 +86,7 @@ class SILVTable : public llvm::ilist_node<SILVTable>,
                                [&](Pair &entry) -> bool {
       if (predicate(entry)) {
         entry.second->decrementRefCount();
+        removeFromVTableCache(entry);
         return true;
       }
       return false;
@@ -99,6 +100,9 @@ class SILVTable : public llvm::ilist_node<SILVTable>,
   /// Print the vtable.
   void print(llvm::raw_ostream &OS, bool Verbose = false) const;
   void dump() const;
+
+private:
+  void removeFromVTableCache(Pair &entry);
 };
 
 } // end swift namespace

lib/SIL/SILVTable.cpp

@@ -33,27 +33,30 @@ SILVTable *SILVTable::create(SILModule &M, ClassDecl *Class,
   SILVTable *vt = ::new (buf) SILVTable(Class, Entries);
   M.vtables.push_back(vt);
   M.VTableMap[Class] = vt;
+  // Update the Module's cache with new vtable + vtable entries:
+  for (auto &entry : Entries) {
+    M.VTableEntryCache.insert({{vt, entry.first}, entry.second});
+  }
   return vt;
 }
 
 SILFunction *
 SILVTable::getImplementation(SILModule &M, SILDeclRef method) const {
-  // FIXME: We should build a sidetable cache in the module. Linear lookup here
-  // is lame.
-  
-  for (auto &entry : getEntries()) {
-    // Check whether this mapping matches either the given decl directly or
-    // one of its overridden decl.
-    SILDeclRef m = method;
-    do {
-      if (entry.first == m)
-        return entry.second;
-    } while ((m = m.getOverridden()));
-  }
-  
+  SILDeclRef m = method;
+  do {
+    auto entry = M.VTableEntryCache.find({this, m});
+    if (entry != M.VTableEntryCache.end()) {
+      return (*entry).second;
+    }
+  } while ((m = m.getOverridden()));
   return nullptr;
 }
 
+void SILVTable::removeFromVTableCache(Pair &entry) {
+  SILModule &M = entry.second->getModule();
+  M.VTableEntryCache.erase({this, entry.first});
+}
+
 SILVTable::SILVTable(ClassDecl *c, ArrayRef<Pair> entries)
   : Class(c), NumEntries(entries.size())
 {

lib/SIL/SILVerifier.cpp

@@ -3546,15 +3546,27 @@ void SILModule::verify() const {
     g.verify();
   }
 
-  // Check all vtables.
+  // Check all vtables and the vtable cache.
   llvm::DenseSet<ClassDecl*> vtableClasses;
+  unsigned EntriesSZ = 0;
   for (const SILVTable &vt : getVTables()) {
     if (!vtableClasses.insert(vt.getClass()).second) {
       llvm::errs() << "Vtable redefined: " << vt.getClass()->getName() << "!\n";
       assert(false && "triggering standard assertion failure routine");
     }
     vt.verify(*this);
+    // Check if there is a cache entry for each vtable entry
+    for (auto entry : vt.getEntries()) {
+      if (VTableEntryCache.find({&vt, entry.first}) == VTableEntryCache.end()) {
+        llvm::errs() << "Vtable entry for function: " << entry.second->getName()
+                     << "not in cache!\n";
+        assert(false && "triggering standard assertion failure routine");
+      }
+      EntriesSZ++;
+    }
   }
+  assert(EntriesSZ == VTableEntryCache.size() &&
+         "Cache size is not equal to true number of VTable entries");
 
   // Check all witness tables.
   DEBUG(llvm::dbgs() << "*** Checking witness tables for duplicates ***\n");