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

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.

Differential Revision: https://reviews.llvm.org/D124584

Author: River707

mlir/lib/Parser/AsmParserState.cpp

@@ -9,6 +9,7 @@
 #include "mlir/Parser/AsmParserState.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/StringExtras.h"
 
 using namespace mlir;
 
@@ -121,18 +122,52 @@ auto AsmParserState::getOpDef(Operation *op) const
                                           : &*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();
 
-  // Return if the given character is a valid identifier character.
-  auto isIdentifierChar = [](char c) {
-    return isalnum(c) || c == '$' || c == '.' || c == '_' || c == '-';
-  };
+  // 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 == '-';
+    };
+
+    while (*curPtr && isIdentifierChar(*(++curPtr)))
+      continue;
+  }
 
-  const char *curPtr = loc.getPointer();
-  while (*curPtr && isIdentifierChar(*(++curPtr)))
-    continue;
   return SMRange(loc, SMLoc::getFromPointer(curPtr));
 }
 

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

@@ -279,7 +279,7 @@ struct Position {
   /// source manager.
   SMLoc getAsSMLoc(llvm::SourceMgr &mgr) const {
     return mgr.FindLocForLineAndColumn(mgr.getMainFileID(), line + 1,
-                                       character);
+                                       character + 1);
   }
 };
 

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

@@ -30,7 +30,7 @@ static Optional<lsp::Location> getLocationFromLoc(FileLineColLoc loc) {
 
   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)};
 }
 

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

@@ -983,9 +983,6 @@ PDLDocument::getCodeCompletion(const lsp::URIForFile &uri,
   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;
@@ -1132,9 +1129,6 @@ lsp::SignatureHelp PDLDocument::getSignatureHelp(const lsp::URIForFile &uri,
   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;

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

@@ -19,7 +19,37 @@
 // 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:         },

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

@@ -13,7 +13,7 @@
 // -----
 {"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",