From b6a32d947524dc07a7871956a1e94165446b0174 Mon Sep 17 00:00:00 2001 From: River Riddle Date: Fri, 9 Sep 2022 02:03:47 -0700 Subject: [PATCH] [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 --- mlir/include/mlir/IR/SymbolTable.h | 5 ++++- mlir/lib/IR/SymbolTable.cpp | 16 +++++++++------- 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/mlir/include/mlir/IR/SymbolTable.h b/mlir/include/mlir/IR/SymbolTable.h index 86ca3b1..3cc4800 100644 --- a/mlir/include/mlir/IR/SymbolTable.h +++ b/mlir/include/mlir/IR/SymbolTable.h @@ -41,7 +41,10 @@ public: return dyn_cast_or_null(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 792bf42..669f538 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 -- 2.7.4