Add method to invalidate cached symbol table

Author: mscuttari

mlir/include/mlir/IR/SymbolTable.h

@@ -316,6 +316,13 @@ class SymbolTableCollection {
   /// Lookup, or create, a symbol table for an operation.
   SymbolTable &getSymbolTable(Operation *op);
 
+  /// Invalidate the cached symbol table for an operation.
+  /// This is important when doing IR modifications that erase and also create
+  /// operations having the 'OpTrait::SymbolTable' trait. If a symbol table of
+  /// an erased operation is not invalidated, a new operation sharing the same
+  /// address would be associated with outdated, and wrong, information.
+  void invalidateSymbolTable(Operation *op);
+
 private:
   friend class LockedSymbolTableCollection;
 

mlir/lib/IR/SymbolTable.cpp

@@ -998,6 +998,13 @@ SymbolTable &SymbolTableCollection::getSymbolTable(Operation *op) {
   return *it.first->second;
 }
 
+void SymbolTableCollection::invalidateSymbolTable(Operation *op) {
+  auto it = symbolTables.find(op);
+
+  if (it != symbolTables.end())
+    symbolTables.erase(it);
+}
+
 //===----------------------------------------------------------------------===//
 // LockedSymbolTableCollection
 //===----------------------------------------------------------------------===//