[mlir][lsp] Add LSP support for attribute and type aliases

[Mogball]

Summary

This wires in attribute and type aliases into the MLIR LSP server. This
will allow goto definition and find references on attribute and type
references, which should make debugging locations and other metadata
easier.

Depends on D158781

[Mogball]

fix typo

[rriddle]

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

660–661

You could also wrap in an mlir markdown block, which would give syntax highlighting

671–672

Same here.

[Mogball]

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

660–661

Pretty cool. I'll give that a try.

[Mogball]

markdown format

[Mogball]

rebase

mlir/include/mlir/AsmParser/AsmParserState.h

 //===- AsmParserState.h -----------------------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ASMPARSER_ASMPARSERSTATE_H
 #define MLIR_ASMPARSER_ASMPARSERSTATE_H
 
-#include "mlir/Support/LLVM.h"
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringRef.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Types.h"
 #include "llvm/Support/SMLoc.h"
 #include <cstddef>
 
 namespace mlir {
 class Block;
 class BlockArgument;
 class FileLineColLoc;
 class Operation;
struct BlockDefinition {
 
     /// Source definitions for any arguments of this block.
     SmallVector<SMDefinition> arguments;
   };
 
   /// This class represents the information for an attribute alias definition
   /// within the input file.
   struct AttributeAliasDefinition {
-    AttributeAliasDefinition(StringRef name, SMRange loc = {})
-        : name(name), definition(loc) {}
+    AttributeAliasDefinition(StringRef name, SMRange loc = {},
+                             Attribute value = {})
+        : name(name), definition(loc), value(value) {}
 
     /// The name of the attribute alias.
     StringRef name;
 
     /// The source location for the alias.
     SMDefinition definition;
+
+    /// The value of the alias.
+    Attribute value;
   };
 
   /// This class represents the information for type definition within the input
   /// file.
   struct TypeAliasDefinition {
-    TypeAliasDefinition(StringRef name, SMRange loc)
-        : name(name), definition(loc) {}
+    TypeAliasDefinition(StringRef name, SMRange loc, Type value)
+        : name(name), definition(loc), value(value) {}
 
     /// The name of the attribute alias.
     StringRef name;
 
     /// The source location for the alias.
     SMDefinition definition;
+
+    /// The value of the alias.
+    Type value;
   };
 
   AsmParserState();
   ~AsmParserState();
   AsmParserState &operator=(AsmParserState &&other);
 
   //===--------------------------------------------------------------------===//
   // Access State
   //===--------------------------------------------------------------------===//
 
   using BlockDefIterator = llvm::pointee_iterator<
       ArrayRef<std::unique_ptr<BlockDefinition>>::iterator>;
   using OperationDefIterator = llvm::pointee_iterator<
       ArrayRef<std::unique_ptr<OperationDefinition>>::iterator>;
+  using AttributeDefIterator = llvm::pointee_iterator<
+      ArrayRef<std::unique_ptr<AttributeAliasDefinition>>::iterator>;
+  using TypeDefIterator = llvm::pointee_iterator<
+      ArrayRef<std::unique_ptr<TypeAliasDefinition>>::iterator>;
 
   /// Return a range of the BlockDefinitions held by the current parser state.
   iterator_range<BlockDefIterator> getBlockDefs() const;
 
   /// 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;
 
   /// Return the definition for the given operation, or nullptr if the given
   /// operation does not have a definition.
   const OperationDefinition *getOpDef(Operation *op) const;
 
+  /// Return a range of the AttributeAliasDefinitions held by the current parser
+  /// state.
+  iterator_range<AttributeDefIterator> getAttributeAliasDefs() const;
+
+  /// Return the definition for the given attribute alias, or nullptr if the
+  /// given alias does not have a definition.
+  const AttributeAliasDefinition *getAttributeAliasDef(StringRef name) const;
+
+  /// Return a range of the TypeAliasDefinitions held by the current parser
+  /// state.
+  iterator_range<TypeDefIterator> getTypeAliasDefs() const;
+
+  /// Return the definition for the given type alias, or nullptr if the given
+  /// alias does not have a definition.
+  const TypeAliasDefinition *getTypeAliasDef(StringRef name) const;
+
   /// Returns (heuristically) the range of an identifier given a SMLoc
   /// corresponding to the start of an identifier location.
   static SMRange convertIdLocToRange(SMLoc loc);
 
   //===--------------------------------------------------------------------===//
   // Populate State
   //===--------------------------------------------------------------------===//
 
public:
   void startRegionDefinition();
 
   /// Finalize the most recently started region definition.
   void finalizeRegionDefinition();
 
   /// Add a definition of the given entity.
   void addDefinition(Block *block, SMLoc location);
   void addDefinition(BlockArgument blockArg, SMLoc location);
-  void addAttrAliasDefinition(StringRef name, SMRange location);
-  void addTypeAliasDefinition(StringRef name, SMRange location);
+  void addAttrAliasDefinition(StringRef name, SMRange location,
+                              Attribute value);
+  void addTypeAliasDefinition(StringRef name, SMRange location, Type value);
 
   /// Add a source uses of the given value.
   void addUses(Value value, ArrayRef<SMLoc> locations);
   void addUses(Block *block, ArrayRef<SMLoc> locations);
   void addAttrAliasUses(StringRef name, SMRange locations);
   void addTypeAliasUses(StringRef name, SMRange locations);
 
   /// Add source uses for all the references nested under `refAttr`. The

mlir/lib/AsmParser/AsmParserState.cpp

 
 auto AsmParserState::getOpDef(Operation *op) const
     -> const OperationDefinition * {
   auto it = impl->operationToIdx.find(op);
   return it == impl->operationToIdx.end() ? nullptr
                                           : &*impl->operations[it->second];
 }
 
+auto AsmParserState::getAttributeAliasDefs() const
+    -> iterator_range<AttributeDefIterator> {
+  return llvm::make_pointee_range(ArrayRef(impl->attrAliases));
+}
+
+auto AsmParserState::getAttributeAliasDef(StringRef name) const
+    -> const AttributeAliasDefinition * {
+  auto it = impl->attrAliasToIdx.find(name);
+  return it == impl->attrAliasToIdx.end() ? nullptr
+                                          : &*impl->attrAliases[it->second];
+}
+
+auto AsmParserState::getTypeAliasDefs() const
+    -> iterator_range<TypeDefIterator> {
+  return llvm::make_pointee_range(ArrayRef(impl->typeAliases));
+}
+
+auto AsmParserState::getTypeAliasDef(StringRef name) const
+    -> const TypeAliasDefinition * {
+  auto it = impl->typeAliasToIdx.find(name);
+  return it == impl->typeAliasToIdx.end() ? nullptr
+                                          : &*impl->typeAliases[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;
void AsmParserState::addDefinition(BlockArgument blockArg, SMLoc location) {
   BlockDefinition &def = *impl->blocks[it->second];
   unsigned argIdx = blockArg.getArgNumber();
 
   if (def.arguments.size() <= argIdx)
     def.arguments.resize(argIdx + 1);
   def.arguments[argIdx] = SMDefinition(convertIdLocToRange(location));
 }
 
-void AsmParserState::addAttrAliasDefinition(StringRef name, SMRange location) {
+void AsmParserState::addAttrAliasDefinition(StringRef name, SMRange location,
+                                            Attribute value) {
   auto [it, inserted] =
       impl->attrAliasToIdx.try_emplace(name, impl->attrAliases.size());
   // Location aliases may be referenced before they are defined.
   if (inserted) {
     impl->attrAliases.push_back(
-        std::make_unique<AttributeAliasDefinition>(name, location));
+        std::make_unique<AttributeAliasDefinition>(name, location, value));
   } else {
-    impl->attrAliases[it->second]->definition.loc = location;
+    AttributeAliasDefinition &attr = *impl->attrAliases[it->second];
+    attr.definition.loc = location;
+    attr.value = value;
   }
 }
 
-void AsmParserState::addTypeAliasDefinition(StringRef name, SMRange location) {
+void AsmParserState::addTypeAliasDefinition(StringRef name, SMRange location,
+                                            Type value) {
   auto [it, inserted] =
       impl->typeAliasToIdx.try_emplace(name, impl->typeAliases.size());
   assert(inserted && "unexpected attribute alias redefinition");
   impl->typeAliases.push_back(
-      std::make_unique<TypeAliasDefinition>(name, location));
+      std::make_unique<TypeAliasDefinition>(name, location, value));
 }
 
 void AsmParserState::addUses(Value value, ArrayRef<SMLoc> locations) {
   // Handle the case where the value is an operation result.
   if (OpResult result = dyn_cast<OpResult>(value)) {
     // Check to see if a definition for the parent operation has been recorded.
     // If one hasn't, we treat the provided value as a placeholder value that
     // will be refined further later.

mlir/lib/AsmParser/Parser.cpp

ParseResult TopLevelOperationParser::parseAttributeAliasDef() {
 
   // Parse the attribute value.
   Attribute attr = parseAttribute();
   if (!attr)
     return failure();
 
   // Register this alias with the parser state.
   if (state.asmState)
-    state.asmState->addAttrAliasDefinition(aliasName, location);
+    state.asmState->addAttrAliasDefinition(aliasName, location, attr);
   state.symbols.attributeAliasDefinitions[aliasName] = attr;
   return success();
 }
 
 ParseResult TopLevelOperationParser::parseTypeAliasDef() {
   assert(getToken().is(Token::exclamation_identifier));
   StringRef aliasName = getTokenSpelling().drop_front();
 
ParseResult TopLevelOperationParser::parseTypeAliasDef() {
 
   // Parse the type.
   Type aliasedType = parseType();
   if (!aliasedType)
     return failure();
 
   // Register this alias with the parser state.
   if (state.asmState)
-    state.asmState->addTypeAliasDefinition(aliasName, location);
+    state.asmState->addTypeAliasDefinition(aliasName, location, aliasedType);
   state.symbols.typeAliasDefinitions.try_emplace(aliasName, aliasedType);
   return success();
 }
 
 ParseResult TopLevelOperationParser::parseFileMetadataDictionary() {
   consumeToken(Token::file_metadata_begin);
   return parseCommaSeparatedListUntil(
       Token::file_metadata_end, [&]() -> ParseResult {

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

lsp::Hover buildHoverForOperationResult(SMRange hoverRange, Operation *op,
                                           unsigned resultStart,
                                           unsigned resultEnd, SMLoc posLoc);
   lsp::Hover buildHoverForBlock(SMRange hoverRange,
                                 const AsmParserState::BlockDefinition &block);
   lsp::Hover
   buildHoverForBlockArgument(SMRange hoverRange, BlockArgument arg,
                              const AsmParserState::BlockDefinition &block);
 
+  lsp::Hover buildHoverForAttributeAlias(
+      SMRange hoverRange, const AsmParserState::AttributeAliasDefinition &attr);
+  lsp::Hover
+  buildHoverForTypeAlias(SMRange hoverRange,
+                         const AsmParserState::TypeAliasDefinition &type);
+
   //===--------------------------------------------------------------------===//
   // Document Symbols
   //===--------------------------------------------------------------------===//
 
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
   void findDocumentSymbols(Operation *op,
                            std::vector<lsp::DocumentSymbol> &symbols);
 
void MLIRDocument::getLocationsOf(const lsp::URIForFile &uri,
   // Check all definitions related to blocks.
   for (const AsmParserState::BlockDefinition &block : asmState.getBlockDefs()) {
     if (containsPosition(block.definition))
       return;
     for (const AsmParserState::SMDefinition &arg : block.arguments)
       if (containsPosition(arg))
         return;
   }
+
+  // Check all alias definitions.
+  for (const AsmParserState::AttributeAliasDefinition &attr :
+       asmState.getAttributeAliasDefs()) {
+    if (containsPosition(attr.definition))
+      return;
+  }
+  for (const AsmParserState::TypeAliasDefinition &type :
+       asmState.getTypeAliasDefs()) {
+    if (containsPosition(type.definition))
+      return;
+  }
 }
 
 void MLIRDocument::findReferencesOf(const lsp::URIForFile &uri,
                                     const lsp::Position &pos,
                                     std::vector<lsp::Location> &references) {
   // Functor used to append all of the definitions/uses of the given SM
   // definition to the reference list.
   auto appendSMDef = [&](const AsmParserState::SMDefinition &def) {
void MLIRDocument::findReferencesOf(const lsp::URIForFile &uri,
   for (const AsmParserState::BlockDefinition &block : asmState.getBlockDefs()) {
     if (isDefOrUse(block.definition, posLoc))
       return appendSMDef(block.definition);
 
     for (const AsmParserState::SMDefinition &arg : block.arguments)
       if (isDefOrUse(arg, posLoc))
         return appendSMDef(arg);
   }
+
+  // Check all alias definitions.
+  for (const AsmParserState::AttributeAliasDefinition &attr :
+       asmState.getAttributeAliasDefs()) {
+    if (isDefOrUse(attr.definition, posLoc))
+      return appendSMDef(attr.definition);
+  }
+  for (const AsmParserState::TypeAliasDefinition &type :
+       asmState.getTypeAliasDefs()) {
+    if (isDefOrUse(type.definition, posLoc))
+      return appendSMDef(type.definition);
+  }
 }
 
 //===----------------------------------------------------------------------===//
 // MLIRDocument: Hover
 //===----------------------------------------------------------------------===//
 
 std::optional<lsp::Hover>
 MLIRDocument::findHover(const lsp::URIForFile &uri,
for (const AsmParserState::BlockDefinition &block : asmState.getBlockDefs()) {
     for (const auto &arg : llvm::enumerate(block.arguments)) {
       if (!isDefOrUse(arg.value(), posLoc, &hoverRange))
         continue;
 
       return buildHoverForBlockArgument(
           hoverRange, block.block->getArgument(arg.index()), block);
     }
   }
+
+  // Check to see if the hover is over an alias.
+  for (const AsmParserState::AttributeAliasDefinition &attr :
+       asmState.getAttributeAliasDefs()) {
+    if (isDefOrUse(attr.definition, posLoc, &hoverRange))
+      return buildHoverForAttributeAlias(hoverRange, attr);
+  }
+  for (const AsmParserState::TypeAliasDefinition &type :
+       asmState.getTypeAliasDefs()) {
+    if (isDefOrUse(type.definition, posLoc, &hoverRange))
+      return buildHoverForTypeAlias(hoverRange, type);
+  }
+
   return std::nullopt;
 }
 
 std::optional<lsp::Hover> MLIRDocument::buildHoverForOperation(
     SMRange hoverRange, const AsmParserState::OperationDefinition &op) {
   lsp::Hover hover(lsp::Range(sourceMgr, hoverRange));
   llvm::raw_string_ostream os(hover.contents.value);
 
os << "Operation: \"" << block.block->getParentOp()->getName() << "\"\n\n"
      << "Block: ";
   printDefBlockName(os, block);
   os << "\n\nArgument #" << arg.getArgNumber() << "\n\n"
      << "Type: `" << arg.getType() << "`\n\n";
 
   return hover;
 }
 
+lsp::Hover MLIRDocument::buildHoverForAttributeAlias(
+    SMRange hoverRange, const AsmParserState::AttributeAliasDefinition &attr) {
+  lsp::Hover hover(lsp::Range(sourceMgr, hoverRange));
+  llvm::raw_string_ostream os(hover.contents.value);
+
+  os << "Attribute Alias: \"" << attr.name << "\n\n";
+  os << "Value: " << attr.value << "\n\n";
+
+  return hover;
+}
+
+lsp::Hover MLIRDocument::buildHoverForTypeAlias(
+    SMRange hoverRange, const AsmParserState::TypeAliasDefinition &type) {
+  lsp::Hover hover(lsp::Range(sourceMgr, hoverRange));
+  llvm::raw_string_ostream os(hover.contents.value);
+
+  os << "TYpe Alias: \"" << type.name << "\n\n";
+  os << "Value: " << type.value << "\n\n";
+
+  return hover;
+}
+
 //===----------------------------------------------------------------------===//
 // MLIRDocument: Document Symbols
 //===----------------------------------------------------------------------===//
 
 void MLIRDocument::findDocumentSymbols(
     std::vector<lsp::DocumentSymbol> &symbols) {
   for (Operation &op : parsedIR)
     findDocumentSymbols(&op, symbols);

mlir/test/mlir-lsp-server/definition.test

 // RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"mlir","capabilities":{},"trace":"off"}}
 // -----
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
   "uri":"test:///foo.mlir",
   "languageId":"mlir",
   "version":1,
-  "text":"func.func @foo() -> i1 {\n%value = arith.constant true\nreturn %value : i1\n}"
+  "text":"#attr = 1 : index\n!type = index\nfunc.func @foo(%arg0: !type) -> i1 attributes {attr = #attr} {\n%value = arith.constant true loc(#loc)\nreturn %value : i1\n}\n#loc = loc(\"foo.mlir\":1:2)"
 }}}
 // -----
 {"jsonrpc":"2.0","id":1,"method":"textDocument/definition","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
-  "position":{"line":2,"character":12}
+  "position":{"line":4,"character":12}
 }}
 //      CHECK:  "id": 1
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": [
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 6,
+// CHECK-NEXT:          "line": 3
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 0,
+// CHECK-NEXT:          "line": 3
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    }
+// -----
+{"jsonrpc":"2.0","id":2,"method":"textDocument/definition","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":2,"character":12}
+}}
+//      CHECK:  "id": 2
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 9,
+// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 0,
+// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    }
+// -----
+{"jsonrpc":"2.0","id":2,"method":"textDocument/definition","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":2,"character":25}
+}}
+//      CHECK:  "id": 2
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 5,
 // CHECK-NEXT:          "line": 1
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 0,
 // CHECK-NEXT:          "line": 1
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    }
 // -----
 {"jsonrpc":"2.0","id":2,"method":"textDocument/definition","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
-  "position":{"line":0,"character":12}
+  "position":{"line":2,"character":57}
 }}
 //      CHECK:  "id": 2
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": [
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
-// CHECK-NEXT:          "character": 9,
+// CHECK-NEXT:          "character": 5,
 // CHECK-NEXT:          "line": 0
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 0,
 // CHECK-NEXT:          "line": 0
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    }
 // -----
+{"jsonrpc":"2.0","id":2,"method":"textDocument/definition","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":3,"character":37}
+}}
+//      CHECK:  "id": 2
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 4,
+// CHECK-NEXT:          "line": 6
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 0,
+// CHECK-NEXT:          "line": 6
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    }
+// -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}

mlir/test/mlir-lsp-server/references.test

 // RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"mlir","capabilities":{},"trace":"off"}}
 // -----
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
   "uri":"test:///foo.mlir",
   "languageId":"mlir",
   "version":1,
-  "text":"func.func @foo() -> i1 {\n%value = arith.constant true\n%result = call @foo() : () -> i1\nreturn %value : i1\n}"
+  "text":"#attr = 1 : index\n!type = index\nfunc.func @foo(%arg0: !type) -> i1 attributes {attr = #attr} {\n%value = arith.constant true\n%result = call @foo(%arg0) : (!type) -> i1\nreturn %value : i1 loc(#loc)\n}\n#loc = loc(\"foo.mlir\":1:2)"
 }}}
 // -----
 {"jsonrpc":"2.0","id":1,"method":"textDocument/references","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
-  "position":{"line":1,"character":2},
+  "position":{"line":3,"character":2},
   "context":{"includeDeclaration": false}
 }}
 //      CHECK:  "id": 1
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": [
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 6,
-// CHECK-NEXT:          "line": 1
+// CHECK-NEXT:          "line": 3
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 0,
-// CHECK-NEXT:          "line": 1
+// CHECK-NEXT:          "line": 3
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 13,
-// CHECK-NEXT:          "line": 3
+// CHECK-NEXT:          "line": 5
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 7,
-// CHECK-NEXT:          "line": 3
+// CHECK-NEXT:          "line": 5
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    }
 // CHECK-NEXT:  ]
 // -----
 {"jsonrpc":"2.0","id":2,"method":"textDocument/references","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
-  "position":{"line":0,"character":12},
+  "position":{"line":2,"character":12},
   "context":{"includeDeclaration": false}
 }}
 //      CHECK:  "id": 2
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": [
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 14,
-// CHECK-NEXT:          "line": 0
+// CHECK-NEXT:          "line": 2
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 10,
-// CHECK-NEXT:          "line": 0
+// CHECK-NEXT:          "line": 2
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 19,
-// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:          "line": 4
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 15,
+// CHECK-NEXT:          "line": 4
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ]
+// -----
+{"jsonrpc":"2.0","id":2,"method":"textDocument/references","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":0,"character":3},
+  "context":{"includeDeclaration": false}
+}}
+//      CHECK:  "id": 2
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 5,
+// CHECK-NEXT:          "line": 0
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 0,
+// CHECK-NEXT:          "line": 0
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 59,
+// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 54,
 // CHECK-NEXT:          "line": 2
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    }
 // CHECK-NEXT:  ]
 // -----
+{"jsonrpc":"2.0","id":2,"method":"textDocument/references","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":1,"character":3},
+  "context":{"includeDeclaration": false}
+}}
+//      CHECK:  "id": 2
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 5,
+// CHECK-NEXT:          "line": 1
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 0,
+// CHECK-NEXT:          "line": 1
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 27,
+// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 22,
+// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 35,
+// CHECK-NEXT:          "line": 4
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 30,
+// CHECK-NEXT:          "line": 4
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ]
+// -----
+{"jsonrpc":"2.0","id":2,"method":"textDocument/references","params":{
+  "textDocument":{"uri":"test:///foo.mlir"},
+  "position":{"line":7,"character":3},
+  "context":{"includeDeclaration": false}
+}}
+//      CHECK:  "id": 2
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 4,
+// CHECK-NEXT:          "line": 7
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 0,
+// CHECK-NEXT:          "line": 7
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 27,
+// CHECK-NEXT:          "line": 5
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 23,
+// CHECK-NEXT:          "line": 5
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ]
+// -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}