[mlir:LSP] Improve conversion between SourceMgr and LSP locations

[rriddle]

Summary

SourceMgr generally uses 1-based locations, whereas the LSP is zero based.
This commit corrects this conversion and also enhances the conversion from SMLoc
to SMRange to support string tokens.

Depends on D124583

[jpienaar]

mlir/lib/Parser/AsmParserState.cpp

125

Function comment?

[rriddle]

mlir/lib/Parser/AsmParserState.cpp

125

Ooops, thanks.

mlir/lib/Parser/AsmParserState.cpp

 //===- AsmParserState.cpp -------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Parser/AsmParserState.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/StringExtras.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // AsmParserState::Impl
 //===----------------------------------------------------------------------===//
 
 struct AsmParserState::Impl {

 
 auto AsmParserState::getOpDef(Operation *op) const
     -> const OperationDefinition * {
   auto it = impl->operationToIdx.find(op);
   return it == impl->operationToIdx.end() ? nullptr
                                           : &*impl->operations[it->second];
 }
 
+/// Lex a string token whose contents start at the given `curPtr`. Returns the
+/// position at the end of the string, after a terminal or invalid character
+/// (e.g. `"` or `\0`).
+static const char *lexLocStringTok(const char *curPtr) {
+  while (char c = *curPtr++) {
+    // Check for various terminal characters.
+    if (StringRef("\"\n\v\f").contains(c))
+      return curPtr;
+
+    // Check for escape sequences.
+    if (c == '\\') {
+      // Check a few known escapes and \xx hex digits.
+      if (*curPtr == '"' || *curPtr == '\\' || *curPtr == 'n' || *curPtr == 't')
+        ++curPtr;
+      else if (llvm::isHexDigit(*curPtr) && llvm::isHexDigit(curPtr[1]))
+        curPtr += 2;
+      else
+        return curPtr;
+    }
+  }
+
+  // If we hit this point, we've reached the end of the buffer. Update the end
+  // pointer to not point past the buffer.
+  return curPtr - 1;
+}
+
 SMRange AsmParserState::convertIdLocToRange(SMLoc loc) {
   if (!loc.isValid())
     return SMRange();
+  const char *curPtr = loc.getPointer();
+
+  // Check if this is a string token.
+  if (*curPtr == '"') {
+    curPtr = lexLocStringTok(curPtr + 1);
 
+    // Otherwise, default to handling an identifier.
+  } else {
     // Return if the given character is a valid identifier character.
     auto isIdentifierChar = [](char c) {
       return isalnum(c) || c == '$' || c == '.' || c == '_' || c == '-';
     };
 
-  const char *curPtr = loc.getPointer();
     while (*curPtr && isIdentifierChar(*(++curPtr)))
       continue;
+  }
+
   return SMRange(loc, SMLoc::getFromPointer(curPtr));
 }
 
 //===----------------------------------------------------------------------===//
 // Populate State
 
 void AsmParserState::initialize(Operation *topLevelOp) {
   startOperationDefinition(topLevelOp->getName());

mlir/lib/Tools/lsp-server-support/Protocol.h

friend bool operator<=(const Position &lhs, const Position &rhs) {
     return std::tie(lhs.line, lhs.character) <=
            std::tie(rhs.line, rhs.character);
   }
 
   /// Convert this position into a source location in the main file of the given
   /// source manager.
   SMLoc getAsSMLoc(llvm::SourceMgr &mgr) const {
     return mgr.FindLocForLineAndColumn(mgr.getMainFileID(), line + 1,
-                                       character);
+                                       character + 1);
   }
 };
 
 /// Add support for JSON serialization.
 bool fromJSON(const llvm::json::Value &value, Position &result,
               llvm::json::Path path);
 llvm::json::Value toJSON(const Position &value);
 raw_ostream &operator<<(raw_ostream &os, const Position &value);

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

if (!sourceURI) {
     lsp::Logger::error("Failed to create URI for file `{0}`: {1}",
                        loc.getFilename(),
                        llvm::toString(sourceURI.takeError()));
     return llvm::None;
   }
 
   lsp::Position position;
   position.line = loc.getLine() - 1;
-  position.character = loc.getColumn();
+  position.character = loc.getColumn() ? loc.getColumn() - 1 : 0;
   return lsp::Location{*sourceURI, lsp::Range(position)};
 }
 
 /// Returns a language server location from the given MLIR location, or None if
 /// one couldn't be created. `uri` is an optional additional filter that, when
 /// present, is used to filter sub locations that do not share the same uri.
 static Optional<lsp::Location>
 getLocationFromLoc(llvm::SourceMgr &sourceMgr, Location loc,

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

 
 lsp::CompletionList
 PDLDocument::getCodeCompletion(const lsp::URIForFile &uri,
                                const lsp::Position &completePos) {
   SMLoc posLoc = completePos.getAsSMLoc(sourceMgr);
   if (!posLoc.isValid())
     return lsp::CompletionList();
 
-  // Adjust the position one further to after the completion trigger token.
-  posLoc = SMLoc::getFromPointer(posLoc.getPointer() + 1);
-
   // To perform code completion, we run another parse of the module with the
   // code completion context provided.
   ods::Context tmpODSContext;
   lsp::CompletionList completionList;
   LSPCodeCompleteContext lspCompleteContext(
       posLoc, completionList, tmpODSContext, sourceMgr.getIncludeDirs());
 
   ast::Context tmpContext(tmpODSContext);

 } // namespace
 
 lsp::SignatureHelp PDLDocument::getSignatureHelp(const lsp::URIForFile &uri,
                                                  const lsp::Position &helpPos) {
   SMLoc posLoc = helpPos.getAsSMLoc(sourceMgr);
   if (!posLoc.isValid())
     return lsp::SignatureHelp();
 
-  // Adjust the position one further to after the completion trigger token.
-  posLoc = SMLoc::getFromPointer(posLoc.getPointer() + 1);
-
   // To perform code completion, we run another parse of the module with the
   // code completion context provided.
   ods::Context tmpODSContext;
   lsp::SignatureHelp signatureHelp;
   LSPSignatureHelpContext completeContext(posLoc, signatureHelp, tmpODSContext);
 
   ast::Context tmpContext(tmpODSContext);
   (void)parsePDLAST(tmpContext, sourceMgr, &completeContext);

mlir/test/mlir-lsp-server/diagnostics.test

 // CHECK-NEXT:         "category": "Parse Error",
 // CHECK-NEXT:         "message": "custom op 'func.func' expected valid '@'-identifier for symbol name",
 // CHECK-NEXT:         "range": {
 // CHECK-NEXT:           "end": {
 // CHECK-NEXT:             "character": 12,
 // CHECK-NEXT:             "line": 0
 // CHECK-NEXT:           },
 // CHECK-NEXT:           "start": {
-// CHECK-NEXT:             "character": 11,
+// CHECK-NEXT:             "character": 10,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           }
+// CHECK-NEXT:         },
+// CHECK-NEXT:         "severity": 1,
+// CHECK-NEXT:         "source": "mlir"
+// CHECK-NEXT:       }
+// CHECK-NEXT:     ],
+// CHECK-NEXT:     "uri": "test:///foo.mlir",
+// CHECK-NEXT:     "version": 1
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.mlir",
+  "languageId":"mlir",
+  "version":1,
+  "text":"\"\""
+}}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [
+// CHECK-NEXT:       {
+// CHECK-NEXT:         "category": "Parse Error",
+// CHECK-NEXT:         "message": "empty operation name is invalid",
+// CHECK-NEXT:         "range": {
+// CHECK-NEXT:           "end": {
+// CHECK-NEXT:             "character": 2,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           },
+// CHECK-NEXT:           "start": {
+// CHECK-NEXT:             "character": 0,
 // CHECK-NEXT:             "line": 0
 // CHECK-NEXT:           }
 // CHECK-NEXT:         },
 // CHECK-NEXT:         "severity": 1,
 // CHECK-NEXT:         "source": "mlir"
 // CHECK-NEXT:       }
 // CHECK-NEXT:     ],
 // CHECK-NEXT:     "uri": "test:///foo.mlir",
 // CHECK-NEXT:     "version": 1
 // CHECK-NEXT:   }
 // -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}

mlir/test/mlir-pdll-lsp-server/definition-split-file.test

 // RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck %s
 // This test checks support for split files by attempting to find the definition
 // of a symbol in a split file. The interesting part of this test is that the
 // file chunk before the one we are looking for the definition in has an error.
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"pdll","capabilities":{},"trace":"off"}}
 // -----
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
   "uri":"test:///foo.pdll",
   "languageId":"pdll",
   "version":1,
   "text":"Pattern Foo {\n// -----\nPattern {\n  erase root: Op<toy.test>;\n }"
 }}}
 // -----
 {"jsonrpc":"2.0","id":1,"method":"textDocument/definition","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":3,"character":12}
+  "position":{"line":3,"character":11}
 }}
 //      CHECK:  "id": 1
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": [
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 12,