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

[rriddle]

Summary

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

mlir/lib/AsmParser/Parser.cpp

ParseResult OperationParser::parseBlock(Block *&block) {
   // Define the block with the specified name.
   auto &blockAndLoc = getBlockInfoByName(name);
   blockAndLoc.loc = nameLoc;
 
   // Use a unique pointer for in-flight block being parsed. Release ownership
   // 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.
   if (!blockAndLoc.block) {
     if (block) {
       blockAndLoc.block = block;
     } else {
       inflightBlock = std::make_unique<Block>();
       blockAndLoc.block = inflightBlock.get();
     }
 
     // Otherwise, the block has a forward declaration. Forward declarations are
     // removed once defined, so if we are defining a existing block and it is
     // not a forward declaration, then it is a redeclaration. Fail if the 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;
 }
 
 ParseResult OperationParser::parseBlockBody(Block *block) {
   // Set the insertion point to the end of the block to parse.
   opBuilder.setInsertionPointToEnd(block);
 
   // Parse the list of operations that make up the body of the block.

mlir/test/IR/invalid.mlir

 // 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:
   cf.br ^missing  // expected-error {{reference to an undefined block}}
 }
 
 // -----