[mlir] Change end of OperationDefinition. (#77273)

Store the last token parsed in the parser state so that the range parsed
can utilize its end rather than the start of the token after parsed.
This results in a tighter range (especially true in the case of
comments, see

```mlir
|%c4 = arith.constant 4 : index

  // Foo

  |
```

vs

```mlir
|%c4 = arith.constant 4 : index|
```

).

Discovered while working on a little textual post processing tool.

Author: jpienaar

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;
 }
 

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();
   }
 

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;