From c35a4f58045ad39faf40b27b57bf52e9a267c019 Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Tue, 19 Jul 2022 16:12:02 -0700
Subject: [PATCH] [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
---
 mlir/lib/AsmParser/Parser.cpp | 16 +++++++++++++---
 mlir/test/IR/invalid.mlir     |  4 ----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp
index 7b17125..9934ca5 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 0bd39f6..064b2ed 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 different error cases.
-// -----
 
 func.func @bad_branch() {
 ^bb12:
-- 
2.7.4