[Mlir-commits] [mlir] c35a4f5 - [mlir][Parser] Fix memory leak when failing to parse a forward declared block

[River Riddle]

Author: River Riddle
Date: 2022-07-25T17:29:49-07:00
New Revision: c35a4f58045ad39faf40b27b57bf52e9a267c019

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

LOG: [mlir][Parser] Fix memory leak when failing to parse a forward declared block

This commit fixes a failure edge case where we accidentally drop forward
declared blocks in the error case. This allows for running the
invalid.mlir test in asan mode now.

Fixes #51387

Differential Revision: https://reviews.llvm.org/D130132

Added: 
    

Modified: 
    mlir/lib/AsmParser/Parser.cpp
    mlir/test/IR/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp
index 7b17125befec0..9934ca529f7aa 100644
--- a/mlir/lib/AsmParser/Parser.cpp
+++ b/mlir/lib/AsmParser/Parser.cpp
@@ -2090,6 +2090,10 @@ ParseResult OperationParser::parseBlock(Block *&block) {
   // only in the case of a successful parse. This ensures that the Block
   // allocated is released if the parse fails and control returns early.
   std::unique_ptr<Block> inflightBlock;
+  auto cleanupOnFailure = llvm::make_scope_exit([&] {
+    if (inflightBlock)
+      inflightBlock->dropAllDefinedValueUses();
+  });
 
   // If a block has yet to be set, this is a new definition. If the caller
   // provided a block, use it. Otherwise create a new one.
@@ -2107,25 +2111,31 @@ ParseResult OperationParser::parseBlock(Block *&block) {
     // was already defined.
   } else if (!eraseForwardRef(blockAndLoc.block)) {
     return emitError(nameLoc, "redefinition of block '") << name << "'";
+  } else {
+    // This was a forward reference block that is now floating. Keep track of it
+    // as inflight in case of error, so that it gets cleaned up properly.
+    inflightBlock.reset(blockAndLoc.block);
   }
 
   // Populate the high level assembly state if necessary.
   if (state.asmState)
     state.asmState->addDefinition(blockAndLoc.block, nameLoc);
-
   block = blockAndLoc.block;
 
   // If an argument list is present, parse it.
   if (getToken().is(Token::l_paren))
     if (parseOptionalBlockArgList(block))
       return failure();
-
   if (parseToken(Token::colon, "expected ':' after block name"))
     return failure();
 
+  // Parse the body of the block.
   ParseResult res = parseBlockBody(block);
+
+  // If parsing was successful, drop the inflight block. We relinquish ownership
+  // back up to the caller.
   if (succeeded(res))
-    inflightBlock.release();
+    (void)inflightBlock.release();
   return res;
 }
 

diff  --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index 0bd39f62a03b0..064b2ed6f749f 100644
--- a/mlir/test/IR/invalid.mlir
+++ b/mlir/test/IR/invalid.mlir
@@ -1,10 +1,6 @@
 // RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -verify-diagnostics
 
-// See http://llvm.org/pr52045
-// UNSUPPORTED: asan
-
 // Check 
diff erent error cases.
-// -----
 
 func.func @bad_branch() {
 ^bb12: