[Mlir-commits] [mlir] [mlir] Change end of OperationDefinition. (PR #77273)

[Jacques Pienaar]

https://github.com/jpienaar created https://github.com/llvm/llvm-project/pull/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 https://gist.github.com/jpienaar/9598339b504157b189c3a3c38314a703 for example of effect of change).

Discovered while working on a little textual post processing tool.

>From 6a9eb928fc63694774399b4ce5bbcb4f969d926d Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Sun, 7 Jan 2024 21:35:18 -0800
Subject: [PATCH] [mlir] Change end of OperationDefinition.

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 return (especially true in the case of comments, see
https://gist.github.com/jpienaar/9598339b504157b189c3a3c38314a703 for example
of effect of change).
---
 mlir/lib/AsmParser/Parser.cpp    | 12 +++++++-----
 mlir/lib/AsmParser/Parser.h      |  5 +++++
 mlir/lib/AsmParser/ParserState.h |  7 +++++--
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/AsmParser/Parser.cpp b/mlir/lib/AsmParser/Parser.cpp
index 3aa9adcbe1c5df..00f2b0c0c2f12f 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 01c55f97a08c2c..b959e67b8e2583 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 1428ea3a82cee9..159058a18fa4e1 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;