[mlir] Fixed double-free bug in SymbolUserMap

[nand]

Summary

SymbolUserMap relied on try_emplace and std::move to relocate an entry to another key. However, if this triggered the resizing of the DenseMap, the value was destroyed before it could be moved to the new storage location, leading to a dangling users reference to be inserted into the map. On destruction, since a new entry was created from one that was already freed, a double-free error occurred.

Fixed issue by re-fetching the iterator after the mutation of the container.

[rriddle]

We could alternatively try to move the users first and then update later, but this is fine as well.

mlir/lib/IR/SymbolTable.cpp

SymbolTable::walkSymbolTables(symbolTableOp, /*allSymUsesVisible=*/false,
                                 walkFn);
 }
 
 void SymbolUserMap::replaceAllUsesWith(Operation *symbol,
                                        StringAttr newSymbolName) {
   auto it = symbolToUsers.find(symbol);
   if (it == symbolToUsers.end())
     return;
-  SetVector<Operation *> &users = it->second;
 
   // Replace the uses within the users of `symbol`.
-  for (Operation *user : users)
+  for (Operation *user : it->second)
     (void)SymbolTable::replaceAllSymbolUses(symbol, newSymbolName, user);
 
   // Move the current users of `symbol` to the new symbol if it is in the
   // symbol table.
   Operation *newSymbol =
       symbolTable.lookupSymbolIn(symbol->getParentOp(), newSymbolName);
   if (newSymbol != symbol) {
-    // Transfer over the users to the new symbol.
-    auto newIt = symbolToUsers.find(newSymbol);
-    if (newIt == symbolToUsers.end())
-      symbolToUsers.try_emplace(newSymbol, std::move(users));
+    // Transfer over the users to the new symbol.  The reference to the old one
+    // is fetched again as the iterator is invalidated during the insertion.
+    auto newIt = symbolToUsers.try_emplace(newSymbol, SetVector<Operation *>{});
+    auto oldIt = symbolToUsers.find(symbol);
+    assert(oldIt != symbolToUsers.end() && "missing old users list");
+    if (newIt.second)
+      newIt.first->second = std::move(oldIt->second);
     else
-      newIt->second.set_union(users);
-    symbolToUsers.erase(symbol);
+      newIt.first->second.set_union(oldIt->second);
+    symbolToUsers.erase(oldIt);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Visibility parsing implementation.
 //===----------------------------------------------------------------------===//
 
 ParseResult impl::parseOptionalVisibilityKeyword(OpAsmParser &parser,