SIL: add an API to replace all entries of a VTable

* add `ModulePassContext.replaceVTableEntries()`
* add `ModulePassContext.notifyFunctionTablesChanged()`

SwiftCompilerSources/Sources/Optimizer/PassManager/ModulePassContext.swift

@@ -168,6 +168,14 @@ struct ModulePassContext : Context, CustomStringConvertible {
     return VTable(bridged: bridgedVTable)
   }
 
+  func replaceVTableEntries(of vTable: VTable, with entries: [VTable.Entry]) {
+    let bridgedEntries = entries.map { $0.bridged }
+    bridgedEntries.withBridgedArrayRef {
+      vTable.bridged.replaceEntries($0)
+    }
+    notifyFunctionTablesChanged()
+  }
+
   func createEmptyFunction(
     name: String,
     parameters: [ParameterInfo],
@@ -201,6 +209,10 @@ struct ModulePassContext : Context, CustomStringConvertible {
       }
     }
   }
+
+  func notifyFunctionTablesChanged() {
+    _bridged.asNotificationHandler().notifyChanges(.functionTablesChanged)
+  }
 }
 
 extension GlobalVariable {

SwiftCompilerSources/Sources/SIL/VTable.swift

@@ -14,7 +14,7 @@ import AST
 import SILBridging
 
 public struct VTable : CustomStringConvertible, NoReflectionChildren {
-  let bridged: BridgedVTable
+  public let bridged: BridgedVTable
 
   public init(bridged: BridgedVTable) { self.bridged = bridged }
 

include/swift/SIL/SILBridging.h

@@ -1065,17 +1065,18 @@ struct BridgedVTableEntryArray {
 };
 
 struct BridgedVTable {
-  const swift::SILVTable * _Nonnull vTable;
+  swift::SILVTable * _Nonnull vTable;
 
   BridgedOwnedString getDebugDescription() const;
   BRIDGED_INLINE SwiftInt getNumEntries() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedVTableEntry getEntry(SwiftInt index) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclObj getClass() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedType getSpecializedClassType() const;
+  BRIDGED_INLINE void replaceEntries(BridgedArrayRef bridgedEntries) const;
 };
 
 struct OptionalBridgedVTable {
-  const swift::SILVTable * _Nullable table;
+  swift::SILVTable * _Nullable table;
 };
 
 struct BridgedWitnessTableEntry {
@@ -1328,7 +1329,8 @@ struct BridgedChangeNotificationHandler {
     instructionsChanged,
     callsChanged,
     branchesChanged,
-    effectsChanged
+    effectsChanged,
+    functionTablesChanged
   };
 
   void notifyChanges(Kind changeKind) const;

include/swift/SIL/SILBridgingImpl.h

@@ -1715,6 +1715,14 @@ BridgedType BridgedVTable::getSpecializedClassType() const {
   return {vTable->getClassType()};
 }
 
+void BridgedVTable::replaceEntries(BridgedArrayRef bridgedEntries) const {
+  llvm::SmallVector<swift::SILVTableEntry, 8> entries;
+  for (const BridgedVTableEntry &e : bridgedEntries.unbridged<BridgedVTableEntry>()) {
+    entries.push_back(e.unbridged());
+  }
+  vTable->replaceEntries(entries);
+}
+
 //===----------------------------------------------------------------------===//
 //               BridgedWitnessTable, BridgedDefaultWitnessTable
 //===----------------------------------------------------------------------===//

include/swift/SIL/SILVTable.h

@@ -205,6 +205,8 @@ public:
     NumEntries = std::distance(Entries.begin(), end);
   }
 
+  void replaceEntries(ArrayRef<Entry> newEntries);
+
   /// Verify that the vtable is well-formed for the given class.
   void verify(const SILModule &M) const;
 

include/swift/SILOptimizer/PassManager/PassManager.h

@@ -69,6 +69,8 @@ class SwiftPassInvocation {
   SILAnalysis::InvalidationKind changeNotifications =
       SILAnalysis::InvalidationKind::Nothing;
 
+  bool functionTablesChanged = false;
+
   /// All slabs, allocated by the pass.
   SILModule::SlabList allocatedSlabs;
 
@@ -141,6 +143,8 @@ public:
   /// Called by the pass when changes are made to the SIL.
   void notifyChanges(SILAnalysis::InvalidationKind invalidationKind);
 
+  void notifyFunctionTablesChanged();
+
   /// Called by the pass manager before the pass starts running.
   void startModulePassRun(SILModuleTransform *transform);
 
@@ -513,6 +517,9 @@ notifyChanges(SILAnalysis::InvalidationKind invalidationKind) {
     changeNotifications = (SILAnalysis::InvalidationKind)
         (changeNotifications | invalidationKind);
 }
+inline void SwiftPassInvocation::notifyFunctionTablesChanged() {
+  functionTablesChanged = true;
+}
 
 } // end namespace swift
 

lib/SIL/IR/SILVTable.cpp

@@ -75,6 +75,22 @@ void SILVTable::updateVTableCache(const Entry &entry) {
   M.VTableEntryCache[{this, entry.getMethod()}] = entry;
 }
 
+void SILVTable::replaceEntries(ArrayRef<Entry> newEntries) {
+  auto entries = getMutableEntries();
+  ASSERT(newEntries.size() <= entries.size());
+  for (unsigned i = 0; i < entries.size(); ++i) {
+    entries[i].getImplementation()->decrementRefCount();
+    if (i < newEntries.size()) {
+      entries[i] = newEntries[i];
+      entries[i].getImplementation()->incrementRefCount();
+      updateVTableCache(entries[i]);
+    } else {
+      removeFromVTableCache(entries[i]);
+    }
+  }
+  NumEntries = newEntries.size();
+}
+
 SILVTable::SILVTable(ClassDecl *c, SILType classType,
                      SerializedKind_t serialized, ArrayRef<Entry> entries)
     : Class(c), classType(classType), SerializedKind(serialized),

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -1516,6 +1516,7 @@ void SwiftPassInvocation::finishedModulePassRun() {
   endPass();
   assert(!function && transform && "not running a pass");
   assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing
+         && !functionTablesChanged
          && "unhandled change notifications at end of module pass");
   transform = nullptr;
 }
@@ -1551,6 +1552,7 @@ void SwiftPassInvocation::endPass() {
 void SwiftPassInvocation::beginTransformFunction(SILFunction *function) {
   assert(!this->function && transform && "not running a pass");
   assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing
+         && !functionTablesChanged
          && "change notifications not cleared");
   this->function = function;
 }
@@ -1561,6 +1563,10 @@ void SwiftPassInvocation::endTransformFunction() {
     passManager->invalidateAnalysis(function, changeNotifications);
     changeNotifications = SILAnalysis::InvalidationKind::Nothing;
   }
+  if (functionTablesChanged) {
+    passManager->invalidateFunctionTables();
+    functionTablesChanged = false;
+  }
   function = nullptr;
   assert(numBlockSetsAllocated == 0 && "Not all BasicBlockSets deallocated");
   assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
@@ -1580,6 +1586,7 @@ void SwiftPassInvocation::endVerifyFunction() {
   assert(function);
   if (!transform) {
     assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing &&
+           !functionTablesChanged &&
            "verifyication must not change the SIL of a function");
     assert(numBlockSetsAllocated == 0 && "Not all BasicBlockSets deallocated");
     assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
@@ -1634,6 +1641,9 @@ void BridgedChangeNotificationHandler::notifyChanges(Kind changeKind) const {
   case Kind::effectsChanged:
     invocation->notifyChanges(SILAnalysis::InvalidationKind::Effects);
     break;
+  case Kind::functionTablesChanged:
+    invocation->notifyFunctionTablesChanged();
+    break;
   }
 }