[mlir] Allow SymbolTable to update existing symbols

[tpopp]

Summary

Previous behavior would fail if inserting an operation that already
existed. Now SymbolTable::insert can also be used as a way to make a
symbol's name unique even after insertion.

Further TODOs have been left over naming and consistent behavior
considerations.

[tpopp]

Correct typo and add comment

[rriddle]

Nice!

mlir/lib/IR/SymbolTable.cpp

153–154

Can you update the doc here and in the header? It should probably say that it only inserts into the symbol table operation if it isn't already a child.

157–158

All of this can be wrapped into the if given that it is only useful when inserting the operation.

mlir/include/mlir/IR/SymbolTable.h

template <typename T> T lookup(StringRef name) const {
     return dyn_cast_or_null<T>(lookup(name));
   }
 
   /// Erase the given symbol from the table.
   void erase(Operation *symbol);
 
   /// Insert a new symbol into the table, and rename it as necessary to avoid
   /// collisions. Also insert at the specified location in the body of the
-  /// associated operation.
+  /// associated operation if it is not already there. It is asserted that the
+  /// symbol is not inside another operation.
   void insert(Operation *symbol, Block::iterator insertPt = {});
 
   /// Return the name of the attribute used for symbol names.
   static StringRef getSymbolAttrName() { return "sym_name"; }
 
   /// Returns the associated operation.
   Operation *getOp() const { return symbolTableOp; }
 

mlir/lib/IR/SymbolTable.cpp

assert(symbol->getParentOp() == symbolTableOp &&
          "SymbolTable");
 
   auto it = symbolTable.find(*name);
   if (it != symbolTable.end() && it->second == symbol) {
     symbolTable.erase(it);
     symbol->erase();
   }
 }
 
-/// Insert a new symbol into the table and associated operation, and rename it
-/// as necessary to avoid collisions.
+// TODO: Consider if this should be renamed to something like insertOrUpdate
+/// Insert a new symbol into the table and associated operation if not already
+/// there and rename it as necessary to avoid collisions.
 void SymbolTable::insert(Operation *symbol, Block::iterator insertPt) {
+  // The symbol cannot be the child of another op and must be the child of the
+  // symbolTableOp after this.
+  //
+  // TODO: consider if SymbolTable's constructor should behave the same.
+  if (!symbol->getParentOp()) {
     auto &body = symbolTableOp->getRegion(0).front();
     if (insertPt == Block::iterator() || insertPt == body.end())
       insertPt = Block::iterator(body.getTerminator());
 
     assert(insertPt->getParentOp() == symbolTableOp &&
            "expected insertPt to be in the associated module operation");
 
     body.getOperations().insert(insertPt, symbol);
+  }
+  assert(symbol->getParentOp() == symbolTableOp &&
+         "symbol is already inserted in another op");
 
   // Add this symbol to the symbol table, uniquing the name if a conflict is
   // detected.
   StringRef name = getSymbolName(symbol);
   if (symbolTable.insert({name, symbol}).second)
     return;
+  // If the symbol was already in the table, also return.
+  if (symbolTable.lookup(name) == symbol)
+    return;
   // If a conflict was detected, then the symbol will not have been added to
   // the symbol table. Try suffixes until we get to a unique name that works.
   SmallString<128> nameBuffer(name);
   unsigned originalLength = nameBuffer.size();
 
   // Iteratively try suffixes until we find one that isn't used.
   do {
     nameBuffer.resize(originalLength);