[mlir] Fix leak in case of failed parse of Block

[jpienaar]

Summary

A Block is optionally allocated and can leak in failed parse. Change Block
parsing to ensure Block gets freed unless parse is successful. defineBlockNamed
had only 1 call site, inline and track whether allocation was performed.

[rriddle]

mlir/lib/Parser/Parser.cpp

1956

Can you add a comment for this?

mlir/lib/Parser/Parser.cpp

public:
   /// Parse a (possibly empty) list of block arguments.
   ParseResult parseOptionalBlockArgList(Block *owner);
 
   /// Get the block with the specified name, creating it if it doesn't
   /// already exist.  The location specified is the point of use, which allows
   /// us to diagnose references to blocks that are not defined precisely.
   Block *getBlockNamed(StringRef name, SMLoc loc);
 
-  /// Define the block with the specified name. Returns the Block* or nullptr in
-  /// the case of redefinition.
-  Block *defineBlockNamed(StringRef name, SMLoc loc, Block *existing);
-
 private:
   /// This class represents a definition of a Block.
   struct BlockDefinition {
     /// A pointer to the defined Block.
     Block *block;
     /// The location that the Block was defined at.
     SMLoc loc;
   };
ParseResult OperationParser::parseBlock(Block *&block) {
   if (block && getToken().isNot(Token::caret_identifier))
     return parseBlockBody(block);
 
   SMLoc nameLoc = getToken().getLoc();
   auto name = getTokenSpelling();
   if (parseToken(Token::caret_identifier, "expected block name"))
     return failure();
 
-  block = defineBlockNamed(name, nameLoc, 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;
+
+  // 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();
+    }
 
-  // Fail if the block was already defined.
-  if (!block)
+    // 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 << "'";
+  }
+
+  // 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();
 
-  return parseBlockBody(block);
+  ParseResult res = parseBlockBody(block);
+  if (succeeded(res))
+    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.
   while (getToken().isNot(Token::caret_identifier, Token::r_brace))
Block *OperationParser::getBlockNamed(StringRef name, SMLoc loc) {
 
   // Populate the high level assembly state if necessary.
   if (state.asmState)
     state.asmState->addUses(blockDef.block, loc);
 
   return blockDef.block;
 }
 
-/// Define the block with the specified name. Returns the Block* or nullptr in
-/// the case of redefinition.
-Block *OperationParser::defineBlockNamed(StringRef name, SMLoc loc,
-                                         Block *existing) {
-  auto &blockAndLoc = getBlockInfoByName(name);
-  blockAndLoc.loc = loc;
-
-  // 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) {
-    blockAndLoc.block = existing ? existing : new Block();
-
-    // 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.
-  } else if (!eraseForwardRef(blockAndLoc.block)) {
-    return nullptr;
-  }
-
-  // Populate the high level assembly state if necessary.
-  if (state.asmState)
-    state.asmState->addDefinition(blockAndLoc.block, loc);
-
-  return blockAndLoc.block;
-}
-
 /// Parse a (possibly empty) list of SSA operands with types as block arguments
 /// enclosed in parentheses.
 ///
 ///   value-id-and-type-list ::= value-id-and-type (`,` ssa-id-and-type)*
 ///   block-arg-list ::= `(` value-id-and-type-list? `)`
 ///
 ParseResult OperationParser::parseOptionalBlockArgList(Block *owner) {
   if (getToken().is(Token::r_brace))

mlir/test/IR/invalid-ops.mlir

 
 // -----
 
 func @invalid_splat(%v : f32) { // expected-note {{prior use here}}
   vector.splat %v : vector<8xf64>
   // expected-error@-1 {{expects different type than prior uses}}
   return
 }
+
+// -----
+
+// Case that resulted in leak previously.
+
+// expected-error@+1 {{expected ':' after block name}}
+"g"()({^a:^b })