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

[jpienaar]

Summary

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.

[rriddle]

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

70

I'd rather base this on something real(i.e. heuristically matching a token) than using a magic number.

[rriddle]

It would be nice of we could reuse some of the tokenizing logic from the parser to recover the full range. Maybe we could use the ParserAsmState for this injection, and add a method to transform positions into ranges.

[jpienaar]

Expose and reuse convertIdLocToRange heuristic instead of fixed range. Now the range is either that of an identifier that starts at that location or a range of one character if not a valid identifier.

[rriddle]

mlir/include/mlir/Parser/AsmParserState.h

109

If this is intended to be just for identifiers, can you rephrase it to indicate that? e.g. ^

mlir/lib/Parser/AsmParserState.cpp

60

nit: None of the other methods here have the description copied from the header.

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

83

nit: Add a newline before this comment.

[jpienaar]

Fixing the test (updated incorrectly)

mlir/include/mlir/Parser/AsmParserState.h

public:
   /// Return the definition for the given block, or nullptr if the given
   /// block does not have a definition.
   const BlockDefinition *getBlockDef(Block *block) const;
 
   /// Return a range of the OperationDefinitions held by the current parser
   /// 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
   //===--------------------------------------------------------------------===//
 
   /// Add a definition of the given operation.
   void addDefinition(
       Operation *op, llvm::SMRange location,
       ArrayRef<std::pair<unsigned, llvm::SMLoc>> resultGroups = llvm::None);

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"
 
 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
 //===----------------------------------------------------------------------===//
 
 struct AsmParserState::Impl {
   /// A mapping from operations in the input source file to their parser state.
   SmallVector<std::unique_ptr<OperationDefinition>> operations;
   DenseMap<Operation *, unsigned> operationToIdx;
auto AsmParserState::getBlockDef(Block *block) const
   auto it = impl->blocksToIdx.find(block);
   return it == impl->blocksToIdx.end() ? nullptr : &*impl->blocks[it->second];
 }
 
 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
 
 void AsmParserState::addDefinition(
     Operation *op, llvm::SMRange location,
     ArrayRef<std::pair<unsigned, llvm::SMLoc>> resultGroups) {
   std::unique_ptr<OperationDefinition> def =
       std::make_unique<OperationDefinition>(op, location);

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

static Optional<lsp::Location> getLocationFromLoc(FileLineColLoc loc) {
   position.character = loc.getColumn();
   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(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>();
     if (!fileLoc)
       return WalkResult::advance();
 
     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();
   });
   return location;
 }
 
 /// Collect all of the locations from the given MLIR location that are not
static void printDefBlockName(raw_ostream &os, Block *block,
   os << "<Block #" << getBlockNumber(block) << ">";
 }
 static void printDefBlockName(raw_ostream &os,
                               const AsmParserState::BlockDefinition &def) {
   printDefBlockName(os, def.block, def.definition.loc);
 }
 
 /// 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";
 
   // Note: Right now all of the diagnostics are treated as parser issues, but
   // some are parser and some are verifier.
   lspDiag.category = "Parse Error";
 
   // Try to grab a file location for this diagnostic.
   // 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;
 
   // Convert the severity for the diagnostic.
   switch (diag.getSeverity()) {
   case DiagnosticSeverity::Note:
     llvm_unreachable("expected notes to be handled separately");
   case DiagnosticSeverity::Warning:
     lspDiag.severity = lsp::DiagnosticSeverity::Warning;
     break;
   case DiagnosticSeverity::Error:
     lspDiag.severity = lsp::DiagnosticSeverity::Error;
     break;
   case DiagnosticSeverity::Remark:
     lspDiag.severity = lsp::DiagnosticSeverity::Information;
     break;
   }
   lspDiag.message = diag.str();
 
   // Attach any notes to the main diagnostic as related information.
   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;
     relatedDiags.emplace_back(noteLoc, note.str());
   }
   if (!relatedDiags.empty())
     lspDiag.relatedInformation = std::move(relatedDiags);
 

 } // namespace
 
 MLIRDocument::MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
                            DialectRegistry &registry,
                            std::vector<lsp::Diagnostic> &diagnostics)
     : 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.
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
   if (!memBuffer) {
     lsp::Logger::error("Failed to create memory buffer for file", uri.file());
     return;
   }

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

 // CHECK: "method": "textDocument/publishDiagnostics",
 // CHECK-NEXT: "params": {
 // CHECK-NEXT:     "diagnostics": [
 // CHECK-NEXT:       {
 // CHECK-NEXT:         "category": "Parse Error",
 // 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": {
 // CHECK-NEXT:             "character": 6,
 // CHECK-NEXT:             "line": 0
 // CHECK-NEXT:           }
 // CHECK-NEXT:         },
 // CHECK-NEXT:         "severity": 1,