[mlir] Fixed double-free bug in SymbolUserMap

`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.

Differential Revision: https://reviews.llvm.org/D129345
This commit is contained in:
Nandor Licker 2022-07-08 08:57:37 +03:00
parent a246eb6814
commit f92d319c70
1 changed files with 10 additions and 8 deletions

View File

@ -1066,10 +1066,9 @@ void SymbolUserMap::replaceAllUsesWith(Operation *symbol,
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
@ -1077,13 +1076,16 @@ void SymbolUserMap::replaceAllUsesWith(Operation *symbol,
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);
}
}