[Mlir-commits] [mlir] f92d319 - [mlir] Fixed double-free bug in SymbolUserMap

[Nandor Licker]

Author: Nandor Licker
Date: 2022-07-08T20:06:35+03:00
New Revision: f92d319c70b5893deb33a996af4b5a74e193d66b

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

LOG: [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

Added: 
    

Modified: 
    mlir/lib/IR/SymbolTable.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp
index 961f6c159705b..b55487c84d910 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -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);
   }
 }