diff --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp index 3aa9adcbe1c5d..00f2b0c0c2f12 100644 --- a/mlir/lib/AsmParser/Parser.cpp +++ b/mlir/lib/AsmParser/Parser.cpp @@ -1209,7 +1209,7 @@ ParseResult OperationParser::parseOperation() { resultIt += std::get<1>(record); } state.asmState->finalizeOperationDefinition( - op, nameTok.getLocRange(), /*endLoc=*/getToken().getLoc(), + op, nameTok.getLocRange(), /*endLoc=*/getLastToken().getEndLoc(), asmResultGroups); } @@ -1225,8 +1225,9 @@ ParseResult OperationParser::parseOperation() { // Add this operation to the assembly state if it was provided to populate. } else if (state.asmState) { - state.asmState->finalizeOperationDefinition(op, nameTok.getLocRange(), - /*endLoc=*/getToken().getLoc()); + state.asmState->finalizeOperationDefinition( + op, nameTok.getLocRange(), + /*endLoc=*/getLastToken().getEndLoc()); } return success(); @@ -1500,8 +1501,9 @@ Operation *OperationParser::parseGenericOperation(Block *insertBlock, // If we are populating the parser asm state, finalize this operation // definition. if (state.asmState) - state.asmState->finalizeOperationDefinition(op, nameToken.getLocRange(), - /*endLoc=*/getToken().getLoc()); + state.asmState->finalizeOperationDefinition( + op, nameToken.getLocRange(), + /*endLoc=*/getLastToken().getEndLoc()); return op; } diff --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h index 01c55f97a08c2..b959e67b8e258 100644 --- a/mlir/lib/AsmParser/Parser.h +++ b/mlir/lib/AsmParser/Parser.h @@ -102,6 +102,9 @@ class Parser { const Token &getToken() const { return state.curToken; } StringRef getTokenSpelling() const { return state.curToken.getSpelling(); } + /// Return the last parsed token. + const Token &getLastToken() const { return state.lastToken; } + /// If the current token has the specified kind, consume it and return true. /// If not, return false. bool consumeIf(Token::Kind kind) { @@ -115,6 +118,7 @@ class Parser { void consumeToken() { assert(state.curToken.isNot(Token::eof, Token::error) && "shouldn't advance past EOF or errors"); + state.lastToken = state.curToken; state.curToken = state.lex.lexToken(); } @@ -129,6 +133,7 @@ class Parser { /// Reset the parser to the given lexer position. void resetToken(const char *tokPos) { state.lex.resetPointer(tokPos); + state.lastToken = state.curToken; state.curToken = state.lex.lexToken(); } diff --git a/mlir/lib/AsmParser/ParserState.h b/mlir/lib/AsmParser/ParserState.h index 1428ea3a82cee..159058a18fa4e 100644 --- a/mlir/lib/AsmParser/ParserState.h +++ b/mlir/lib/AsmParser/ParserState.h @@ -54,8 +54,8 @@ struct ParserState { AsmParserCodeCompleteContext *codeCompleteContext) : config(config), lex(sourceMgr, config.getContext(), codeCompleteContext), - curToken(lex.lexToken()), symbols(symbols), asmState(asmState), - codeCompleteContext(codeCompleteContext) {} + curToken(lex.lexToken()), lastToken(Token::error, ""), symbols(symbols), + asmState(asmState), codeCompleteContext(codeCompleteContext) {} ParserState(const ParserState &) = delete; void operator=(const ParserState &) = delete; @@ -68,6 +68,9 @@ struct ParserState { /// This is the next token that hasn't been consumed yet. Token curToken; + /// This is the last token that has been consumed. + Token lastToken; + /// The current state for symbol parsing. SymbolState &symbols;