[Mlir-commits] [mlir] b6a32d9 - [mlir:SymbolTable] Add "remove" method that drops a symbol without erasing it

[River Riddle]

Author: River Riddle
Date: 2022-09-09T15:09:57-07:00
New Revision: b6a32d947524dc07a7871956a1e94165446b0174

URL: https://github.com/llvm/llvm-project/commit/b6a32d947524dc07a7871956a1e94165446b0174
DIFF: https://github.com/llvm/llvm-project/commit/b6a32d947524dc07a7871956a1e94165446b0174.diff

LOG: [mlir:SymbolTable] Add "remove" method that drops a symbol without erasing it

There are various use cases where we don't want to immediately erase
an operation when removing it from the symbol table. This commit adds
a "remove" method to support that.

Differential Revision: https://reviews.llvm.org/D133564

Added: 
    

Modified: 
    mlir/include/mlir/IR/SymbolTable.h
    mlir/lib/IR/SymbolTable.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/SymbolTable.h b/mlir/include/mlir/IR/SymbolTable.h
index 86ca3b173b372..3cc48000e1944 100644
--- a/mlir/include/mlir/IR/SymbolTable.h
+++ b/mlir/include/mlir/IR/SymbolTable.h
@@ -41,7 +41,10 @@ class SymbolTable {
     return dyn_cast_or_null<T>(lookup(name));
   }
 
-  /// Erase the given symbol from the table.
+  /// Remove the given symbol from the table, without deleting it.
+  void remove(Operation *op);
+
+  /// Erase the given symbol from the table and delete the operation.
   void erase(Operation *symbol);
 
   /// Insert a new symbol into the table, and rename it as necessary to avoid

diff  --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 792bf42488372..669f53867d740 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -146,19 +146,21 @@ Operation *SymbolTable::lookup(StringAttr name) const {
   return symbolTable.lookup(name);
 }
 
-/// Erase the given symbol from the table.
-void SymbolTable::erase(Operation *symbol) {
-  StringAttr name = getNameIfSymbol(symbol);
+void SymbolTable::remove(Operation *op) {
+  StringAttr name = getNameIfSymbol(op);
   assert(name && "expected valid 'name' attribute");
-  assert(symbol->getParentOp() == symbolTableOp &&
+  assert(op->getParentOp() == symbolTableOp &&
          "expected this operation to be inside of the operation with this "
          "SymbolTable");
 
   auto it = symbolTable.find(name);
-  if (it != symbolTable.end() && it->second == symbol) {
+  if (it != symbolTable.end() && it->second == op)
     symbolTable.erase(it);
-    symbol->erase();
-  }
+}
+
+void SymbolTable::erase(Operation *symbol) {
+  remove(symbol);
+  symbol->erase();
 }
 
 // TODO: Consider if this should be renamed to something like insertOrUpdate