[mlir-lsp] Report range of potential identifier starting at location of diagnostic

Currently the diagnostics reports the file:line:col, but some LSP
frontends require a non-empty range. Report either the range of an
identifier that starts at location, or a range of 1. Expose the id
location to range helper and reuse here.

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

mlir/include/mlir/Parser/AsmParserState.h

@@ -104,6 +104,10 @@ public:
   /// state.
   iterator_range<OperationDefIterator> getOpDefs() const;
 
+  /// Returns (heuristically) the range of an identifier given a SMLoc
+  /// corresponding to the start of an identifier location.
+  static llvm::SMRange convertIdLocToRange(llvm::SMLoc loc);
+
   //===--------------------------------------------------------------------===//
   // Populate State
   //===--------------------------------------------------------------------===//

mlir/lib/Parser/AsmParserState.cpp

@@ -11,23 +11,6 @@
 
 using namespace mlir;
 
-/// Given a SMLoc corresponding to an identifier location, return a location
-/// representing the full range of the identifier.
-static llvm::SMRange convertIdLocToRange(llvm::SMLoc loc) {
-  if (!loc.isValid())
-    return llvm::SMRange();
-
-  // 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 (isIdentifierChar(*(++curPtr)))
-    continue;
-  return llvm::SMRange(loc, llvm::SMLoc::getFromPointer(curPtr));
-}
-
 //===----------------------------------------------------------------------===//
 // AsmParserState::Impl
 //===----------------------------------------------------------------------===//
@@ -74,6 +57,23 @@ auto AsmParserState::getOpDefs() const -> iterator_range<OperationDefIterator> {
   return llvm::make_pointee_range(llvm::makeArrayRef(impl->operations));
 }
 
+/// Returns (heuristically) the range of an identifier given a SMLoc
+/// corresponding to the start of an identifier location.
+llvm::SMRange AsmParserState::convertIdLocToRange(llvm::SMLoc loc) {
+  if (!loc.isValid())
+    return llvm::SMRange();
+
+  // 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 (isIdentifierChar(*(++curPtr)))
+    continue;
+  return llvm::SMRange(loc, llvm::SMLoc::getFromPointer(curPtr));
+}
+
 //===----------------------------------------------------------------------===//
 // Populate State
 

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

@@ -67,7 +67,8 @@ static Optional<lsp::Location> getLocationFromLoc(FileLineColLoc loc) {
 /// 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(Location loc, const lsp::URIForFile *uri = nullptr) {
+getLocationFromLoc(llvm::SourceMgr &sourceMgr, Location loc,
+                   const lsp::URIForFile *uri = nullptr) {
   Optional<lsp::Location> location;
   loc->walk([&](Location nestedLoc) {
     FileLineColLoc fileLoc = nestedLoc.dyn_cast<FileLineColLoc>();
@@ -77,6 +78,17 @@ getLocationFromLoc(Location loc, const lsp::URIForFile *uri = nullptr) {
     Optional<lsp::Location> sourceLoc = getLocationFromLoc(fileLoc);
     if (sourceLoc && (!uri || sourceLoc->uri == *uri)) {
       location = *sourceLoc;
+      llvm::SMLoc loc = sourceMgr.FindLocForLineAndColumn(
+          sourceMgr.getMainFileID(), fileLoc.getLine(), fileLoc.getColumn());
+
+      // Use range of potential identifier starting at location, else length 1
+      // range.
+      location->range.end.character += 1;
+      if (Optional<llvm::SMRange> range =
+              AsmParserState::convertIdLocToRange(loc)) {
+        auto lineCol = sourceMgr.getLineAndColumn(range->End);
+        location->range.end.character = lineCol.second - 1;
+      }
       return WalkResult::interrupt();
     }
     return WalkResult::advance();
@@ -195,7 +207,8 @@ static void printDefBlockName(raw_ostream &os,
 }
 
 /// Convert the given MLIR diagnostic to the LSP form.
-static lsp::Diagnostic getLspDiagnoticFromDiag(Diagnostic &diag,
+static lsp::Diagnostic getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr,
+                                               Diagnostic &diag,
                                                const lsp::URIForFile &uri) {
   lsp::Diagnostic lspDiag;
   lspDiag.source = "mlir";
@@ -208,7 +221,7 @@ static lsp::Diagnostic getLspDiagnoticFromDiag(Diagnostic &diag,
   // TODO: For simplicity, we just grab the first one. It may be likely that we
   // will need a more interesting heuristic here.'
   Optional<lsp::Location> lspLocation =
-      getLocationFromLoc(diag.getLocation(), &uri);
+      getLocationFromLoc(sourceMgr, diag.getLocation(), &uri);
   if (lspLocation)
     lspDiag.range = lspLocation->range;
 
@@ -232,7 +245,8 @@ static lsp::Diagnostic getLspDiagnoticFromDiag(Diagnostic &diag,
   std::vector<lsp::DiagnosticRelatedInformation> relatedDiags;
   for (Diagnostic &note : diag.getNotes()) {
     lsp::Location noteLoc;
-    if (Optional<lsp::Location> loc = getLocationFromLoc(note.getLocation()))
+    if (Optional<lsp::Location> loc =
+            getLocationFromLoc(sourceMgr, note.getLocation()))
       noteLoc = *loc;
     else
       noteLoc.uri = uri;
@@ -306,7 +320,7 @@ MLIRDocument::MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
     : context(registry) {
   context.allowUnregisteredDialects();
   ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) {
-    diagnostics.push_back(getLspDiagnoticFromDiag(diag, uri));
+    diagnostics.push_back(getLspDiagnoticFromDiag(sourceMgr, diag, uri));
   });
 
   // Try to parsed the given IR string.

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

@@ -15,7 +15,7 @@
 // CHECK-NEXT:         "message": "custom op 'func' expected valid '@'-identifier for symbol name",
 // CHECK-NEXT:         "range": {
 // CHECK-NEXT:           "end": {
-// CHECK-NEXT:             "character": 6,
+// CHECK-NEXT:             "character": 7,
 // CHECK-NEXT:             "line": 0
 // CHECK-NEXT:           },
 // CHECK-NEXT:           "start": {