[mlir-lsp-server] Add support for tracking the use/def chains of symbols

[rriddle]

Summary

This revision adds assembly state tracking for uses of symbols, allowing for go-to-definition and references support for SymbolRefAttrs.

[jpienaar]

Looks good thanks

mlir/include/mlir/Parser/AsmParserState.h

69

What is the range of a symbol use? E.g., just the name of it so that hover would report it is the op?

mlir/lib/Parser/AsmParserState.cpp

41

When would this one be needed?

42

Do we need this to be implicitly convertable?

71

Would all these all be nested?

mlir/lib/Parser/Parser.cpp

1996

Nit: move this to line 1992 (it feels closer linked to the conditional there and then also clearer why you just consume vs consumeIf before)

2001

Nice lint warning :) I can't see it off hand though, is this due to parseBlock ?

mlir/test/mlir-lsp-server/hover.test

8

Could we have this as a comment with line & column markers?

Would make it easier to check the tests manually:

  0    5    a
0  func @foo(%arg: i1) {
  0         9
1  %value = constant true
2  br ^bb2

[rriddle]

Update

[rriddle]

mlir/include/mlir/Parser/AsmParserState.h

69

The range is the full at_identifier: i.e. @foo. I haven't added hover behavior, because I'm not entirely sure what we would actually want to display in this case.

mlir/lib/Parser/AsmParserState.cpp

71

Whether they are or not isn't really that important, the expectation of nesting here is that the most recently started operation is stopped first.

mlir/lib/Parser/Parser.cpp

2001

Yeah, that is the only reasonable thing it could be caused by. (I don't plan on changing anything though, because well... parsing is generally a recursive play land).

mlir/test/mlir-lsp-server/hover.test

8

I can look at this in a followup, as this applies to all of the tests.

mlir/include/mlir/Parser/AsmParserState.h

 #include "llvm/Support/SMLoc.h"
 #include <cstddef>
 
 namespace mlir {
 class Block;
 class BlockArgument;
 class FileLineColLoc;
 class Operation;
+class OperationName;
+class SymbolRefAttr;
 class Value;
 
 /// This class represents state from a parsed MLIR textual format string. It is
 /// useful for building additional analysis and language utilities on top of
 /// textual MLIR. This should generally not be used for traditional compilation.
 class AsmParserState {
 public:
   /// This class represents a definition within the source manager, containing
struct OperationDefinition {
     /// The operation representing this definition.
     Operation *op;
 
     /// The source location for the operation, i.e. the location of its name.
     llvm::SMRange loc;
 
     /// Source definitions for any result groups of this operation.
     SmallVector<std::pair<unsigned, SMDefinition>> resultGroups;
+
+    /// If this operation is a symbol operation, this vector contains symbol
+    /// uses of this operation.
+    SmallVector<llvm::SMRange> symbolUses;
   };
 
   /// This class represents the information for a block definition within the
   /// input file.
   struct BlockDefinition {
     BlockDefinition(Block *block, llvm::SMRange loc = {})
         : block(block), definition(loc) {}
 
public:
   /// Return a range of the OperationDefinitions held by the current parser
   /// state.
   iterator_range<OperationDefIterator> getOpDefs() const;
 
   //===--------------------------------------------------------------------===//
   // Populate State
   //===--------------------------------------------------------------------===//
 
-  /// Add a definition of the given operation.
-  void addDefinition(
-      Operation *op, llvm::SMRange location,
+  /// Initialize the state in preparation for populating more parser state under
+  /// the given top-level operation.
+  void initialize(Operation *topLevelOp);
+
+  /// Finalize any in-progress parser state under the given top-level operation.
+  void finalize(Operation *topLevelOp);
+
+  /// Start a definition for an operation with the given name.
+  void startOperationDefinition(const OperationName &opName);
+
+  /// Finalize the most recently started operation definition.
+  void finalizeOperationDefinition(
+      Operation *op, llvm::SMRange nameLoc,
       ArrayRef<std::pair<unsigned, llvm::SMLoc>> resultGroups = llvm::None);
+
+  /// Start a definition for a region nested under the current operation.
+  void startRegionDefinition();
+
+  /// Finalize the most recently started region definition.
+  void finalizeRegionDefinition();
+
+  /// Add a definition of the given entity.
   void addDefinition(Block *block, llvm::SMLoc location);
   void addDefinition(BlockArgument blockArg, llvm::SMLoc location);
 
   /// Add a source uses of the given value.
   void addUses(Value value, ArrayRef<llvm::SMLoc> locations);
   void addUses(Block *block, ArrayRef<llvm::SMLoc> locations);
 
+  /// Add source uses for all the references nested under `refAttr`. The
+  /// provided `locations` should match 1-1 with the number of references in
+  /// `refAttr`, i.e.:
+  ///   nestedReferences.size() + /*leafReference=*/1 == refLocations.size()
+  void addUses(SymbolRefAttr refAttr, ArrayRef<llvm::SMRange> refLocations);
+
   /// Refine the `oldValue` to the `newValue`. This is used to indicate that
   /// `oldValue` was a placeholder, and the uses of it should really refer to
   /// `newValue`.
   void refineDefinition(Value oldValue, Value newValue);
 
 private:
   struct Impl;
 
   /// A pointer to the internal implementation of this class.
   std::unique_ptr<Impl> impl;
 };
 
 } // end namespace mlir
 
 #endif // MLIR_PARSER_ASMPARSERSTATE_H

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"
 
 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();
static llvm::SMRange convertIdLocToRange(llvm::SMLoc loc) {
   return llvm::SMRange(loc, llvm::SMLoc::getFromPointer(curPtr));
 }
 
 //===----------------------------------------------------------------------===//
 // AsmParserState::Impl
 //===----------------------------------------------------------------------===//
 
 struct AsmParserState::Impl {
+  /// A map from a SymbolRefAttr to a range of uses.
+  using SymbolUseMap = DenseMap<Attribute, SmallVector<llvm::SMRange>>;
+
+  struct PartialOpDef {
+    PartialOpDef() : symbolTable(std::make_unique<SymbolUseMap>()) {}
+    PartialOpDef(const OperationName &opName) {
+      const auto *abstractOp = opName.getAbstractOperation();
+      if (abstractOp && abstractOp->hasTrait<OpTrait::SymbolTable>())
+        symbolTable = std::make_unique<SymbolUseMap>();
+    }
+
+    /// Return if this operation is a symbol table.
+    bool isSymbolTable() const { return symbolTable.get(); }
+
+    /// If this operation is a symbol table, the following contains symbol uses
+    /// within this operation.
+    std::unique_ptr<SymbolUseMap> symbolTable;
+  };
+
+  /// Resolve any symbol table uses under the given partial operation.
+  void resolveSymbolUses(Operation *op, PartialOpDef &opDef);
+
   /// A mapping from operations in the input source file to their parser state.
   SmallVector<std::unique_ptr<OperationDefinition>> operations;
   DenseMap<Operation *, unsigned> operationToIdx;
 
   /// A mapping from blocks in the input source file to their parser state.
   SmallVector<std::unique_ptr<BlockDefinition>> blocks;
   DenseMap<Block *, unsigned> blocksToIdx;
 
   /// A set of value definitions that are placeholders for forward references.
   /// This map should be empty if the parser finishes successfully.
   DenseMap<Value, SmallVector<llvm::SMLoc>> placeholderValueUses;
+
+  /// A stack of partial operation definitions that have been started but not
+  /// yet finalized.
+  SmallVector<PartialOpDef> partialOperations;
+
+  /// A stack of symbol use scopes. This is used when collecting symbol table
+  /// uses during parsing.
+  SmallVector<SymbolUseMap *> symbolUseScopes;
+
+  /// A symbol table containing all of the symbol table operations in the IR.
+  SymbolTableCollection symbolTable;
 };
 
+void AsmParserState::Impl::resolveSymbolUses(Operation *op,
+                                             PartialOpDef &opDef) {
+  assert(opDef.isSymbolTable() && "expected op to be a symbol table");
+
+  SmallVector<Operation *> symbolOps;
+  for (auto &it : *opDef.symbolTable) {
+    symbolOps.clear();
+    if (failed(symbolTable.lookupSymbolIn(op, it.first.cast<SymbolRefAttr>(),
+                                          symbolOps)))
+      continue;
+
+    for (const auto &symIt : llvm::zip(symbolOps, it.second)) {
+      auto opIt = operationToIdx.find(std::get<0>(symIt));
+      if (opIt != operationToIdx.end())
+        operations[opIt->second]->symbolUses.push_back(std::get<1>(symIt));
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // AsmParserState
 //===----------------------------------------------------------------------===//
 
 AsmParserState::AsmParserState() : impl(std::make_unique<Impl>()) {}
 AsmParserState::~AsmParserState() {}
 AsmParserState &AsmParserState::operator=(AsmParserState &&other) {
   impl = std::move(other.impl);

 
 auto AsmParserState::getOpDefs() const -> iterator_range<OperationDefIterator> {
   return llvm::make_pointee_range(llvm::makeArrayRef(impl->operations));
 }
 
 //===----------------------------------------------------------------------===//
 // Populate State
 
-void AsmParserState::addDefinition(
-    Operation *op, llvm::SMRange location,
+void AsmParserState::initialize(Operation *topLevelOp) {
+  startOperationDefinition(topLevelOp->getName());
+
+  // If the top-level operation is a symbol table, push a new symbol scope.
+  Impl::PartialOpDef &partialOpDef = impl->partialOperations.back();
+  if (partialOpDef.isSymbolTable())
+    impl->symbolUseScopes.push_back(partialOpDef.symbolTable.get());
+}
+
+void AsmParserState::finalize(Operation *topLevelOp) {
+  assert(!impl->partialOperations.empty() &&
+         "expected valid partial operation definition");
+  Impl::PartialOpDef partialOpDef = impl->partialOperations.pop_back_val();
+
+  // If this operation is a symbol table, resolve any symbol uses.
+  if (partialOpDef.isSymbolTable())
+    impl->resolveSymbolUses(topLevelOp, partialOpDef);
+}
+
+void AsmParserState::startOperationDefinition(const OperationName &opName) {
+  impl->partialOperations.push_back(opName);
+}
+
+void AsmParserState::finalizeOperationDefinition(
+    Operation *op, llvm::SMRange nameLoc,
     ArrayRef<std::pair<unsigned, llvm::SMLoc>> resultGroups) {
+  assert(!impl->partialOperations.empty() &&
+         "expected valid partial operation definition");
+  Impl::PartialOpDef partialOpDef = impl->partialOperations.pop_back_val();
+
+  // Build the full operation definition.
   std::unique_ptr<OperationDefinition> def =
-      std::make_unique<OperationDefinition>(op, location);
+      std::make_unique<OperationDefinition>(op, nameLoc);
   for (auto &resultGroup : resultGroups)
     def->resultGroups.emplace_back(resultGroup.first,
                                    convertIdLocToRange(resultGroup.second));
-
   impl->operationToIdx.try_emplace(op, impl->operations.size());
   impl->operations.emplace_back(std::move(def));
+
+  // If this operation is a symbol table, resolve any symbol uses.
+  if (partialOpDef.isSymbolTable())
+    impl->resolveSymbolUses(op, partialOpDef);
+}
+
+void AsmParserState::startRegionDefinition() {
+  assert(!impl->partialOperations.empty() &&
+         "expected valid partial operation definition");
+
+  // If the parent operation of this region is a symbol table, we also push a
+  // new symbol scope.
+  Impl::PartialOpDef &partialOpDef = impl->partialOperations.back();
+  if (partialOpDef.isSymbolTable())
+    impl->symbolUseScopes.push_back(partialOpDef.symbolTable.get());
+}
+
+void AsmParserState::finalizeRegionDefinition() {
+  assert(!impl->partialOperations.empty() &&
+         "expected valid partial operation definition");
+
+  // If the parent operation of this region is a symbol table, pop the symbol
+  // scope for this region.
+  Impl::PartialOpDef &partialOpDef = impl->partialOperations.back();
+  if (partialOpDef.isSymbolTable())
+    impl->symbolUseScopes.pop_back();
 }
 
 void AsmParserState::addDefinition(Block *block, llvm::SMLoc location) {
   auto it = impl->blocksToIdx.find(block);
   if (it == impl->blocksToIdx.end()) {
     impl->blocksToIdx.try_emplace(block, impl->blocks.size());
     impl->blocks.emplace_back(std::make_unique<BlockDefinition>(
         block, convertIdLocToRange(location)));
if (it == impl->blocksToIdx.end()) {
     impl->blocks.emplace_back(std::make_unique<BlockDefinition>(block));
   }
 
   BlockDefinition &def = *impl->blocks[it->second];
   for (llvm::SMLoc loc : locations)
     def.definition.uses.push_back(convertIdLocToRange(loc));
 }
 
+void AsmParserState::addUses(SymbolRefAttr refAttr,
+                             ArrayRef<llvm::SMRange> locations) {
+  // Ignore this symbol if no scopes are active.
+  if (impl->symbolUseScopes.empty())
+    return;
+
+  assert((refAttr.getNestedReferences().size() + 1) == locations.size() &&
+         "expected the same number of references as provided locations");
+  (*impl->symbolUseScopes.back())[refAttr].append(locations.begin(),
+                                                  locations.end());
+}
+
 void AsmParserState::refineDefinition(Value oldValue, Value newValue) {
   auto it = impl->placeholderValueUses.find(oldValue);
   assert(it != impl->placeholderValueUses.end() &&
          "expected `oldValue` to be a placeholder");
   addUses(newValue, it->second);
   impl->placeholderValueUses.erase(oldValue);
 }

mlir/lib/Parser/AttributeParser.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "Parser.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/IntegerSet.h"
+#include "mlir/Parser/AsmParserState.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Endian.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 /// Parse an arbitrary attribute.
 ///
if (!type && consumeIf(Token::colon) && !(type = parseType()))
       return Attribute();
 
     return type ? StringAttr::get(val, type)
                 : StringAttr::get(getContext(), val);
   }
 
   // Parse a symbol reference attribute.
   case Token::at_identifier: {
+    // When populating the parser state, this is a list of locations for all of
+    // the nested references.
+    SmallVector<llvm::SMRange> referenceLocations;
+    if (state.asmState)
+      referenceLocations.push_back(getToken().getLocRange());
+
+    // Parse the top-level reference.
     std::string nameStr = getToken().getSymbolReference();
     consumeToken(Token::at_identifier);
 
     // Parse any nested references.
     std::vector<FlatSymbolRefAttr> nestedRefs;
     while (getToken().is(Token::colon)) {
       // Check for the '::' prefix.
       const char *curPointer = getToken().getLoc().getPointer();
       consumeToken(Token::colon);
       if (!consumeIf(Token::colon)) {
         state.lex.resetPointer(curPointer);
         consumeToken();
         break;
       }
       // Parse the reference itself.
       auto curLoc = getToken().getLoc();
       if (getToken().isNot(Token::at_identifier)) {
         emitError(curLoc, "expected nested symbol reference identifier");
         return Attribute();
       }
 
+      // If we are populating the assembly state, add the location for this
+      // reference.
+      if (state.asmState)
+        referenceLocations.push_back(getToken().getLocRange());
+
       std::string nameStr = getToken().getSymbolReference();
       consumeToken(Token::at_identifier);
       nestedRefs.push_back(SymbolRefAttr::get(getContext(), nameStr));
     }
+    SymbolRefAttr symbolRefAttr = builder.getSymbolRefAttr(nameStr, nestedRefs);
 
-    return builder.getSymbolRefAttr(nameStr, nestedRefs);
+    // If we are populating the assembly state, record this symbol reference.
+    if (state.asmState)
+      state.asmState->addUses(symbolRefAttr, referenceLocations);
+    return symbolRefAttr;
   }
 
   // Parse a 'unit' attribute.
   case Token::kw_unit:
     consumeToken(Token::kw_unit);
     return builder.getUnitAttr();
 
   default:

mlir/lib/Parser/Parser.cpp

 // OperationParser
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class provides support for parsing operations and regions of
 /// operations.
 class OperationParser : public Parser {
 public:
-  OperationParser(ParserState &state, Operation *topLevelOp)
-      : Parser(state), opBuilder(topLevelOp->getRegion(0)),
-        topLevelOp(topLevelOp) {
-    // The top level operation starts a new name scope.
-    pushSSANameScope(/*isIsolated=*/true);
-  }
-
+  OperationParser(ParserState &state, Operation *topLevelOp);
   ~OperationParser();
 
   /// After parsing is finished, this function must be called to see if there
   /// are any remaining issues.
-  ParseResult finalize();
+  ParseResult finalize(Operation *topLevelOp);
 
   //===--------------------------------------------------------------------===//
   // SSA Value Handling
   //===--------------------------------------------------------------------===//
 
   /// This represents a use of an SSA value in the program.  The first two
   /// entries in the tuple are the name and result number of a reference.  The
   /// third is the location of the reference, which is used in case this ends
public:
   /// Parse a region into 'region' with the provided entry block arguments.
   /// 'isIsolatedNameScope' indicates if the naming scope of this region is
   /// isolated from those above.
   ParseResult parseRegion(Region &region,
                           ArrayRef<std::pair<SSAUseInfo, Type>> entryArguments,
                           bool isIsolatedNameScope = false);
 
   /// Parse a region body into 'region'.
-  ParseResult parseRegionBody(Region &region);
+  ParseResult
+  parseRegionBody(Region &region, llvm::SMLoc startLoc,
+                  ArrayRef<std::pair<SSAUseInfo, Type>> entryArguments,
+                  bool isIsolatedNameScope);
 
   //===--------------------------------------------------------------------===//
   // Block Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a new block into 'block'.
   ParseResult parseBlock(Block *&block);
 
private:
   /// The builder used when creating parsed operation instances.
   OpBuilder opBuilder;
 
   /// The top level operation that holds all of the parsed operations.
   Operation *topLevelOp;
 };
 } // end anonymous namespace
 
+OperationParser::OperationParser(ParserState &state, Operation *topLevelOp)
+    : Parser(state), opBuilder(topLevelOp->getRegion(0)),
+      topLevelOp(topLevelOp) {
+  // The top level operation starts a new name scope.
+  pushSSANameScope(/*isIsolated=*/true);
+
+  // If we are populating the parser state, prepare it for parsing.
+  if (state.asmState)
+    state.asmState->initialize(topLevelOp);
+}
+
 OperationParser::~OperationParser() {
   for (auto &fwd : forwardRefPlaceholders) {
     // Drop all uses of undefined forward declared reference and destroy
     // defining operation.
     fwd.first.dropAllUses();
     fwd.first.getDefiningOp()->destroy();
   }
   for (const auto &scope : forwardRef) {
     for (const auto &fwd : scope) {
       // Delete all blocks that were created as forward references but never
       // included into a region.
       fwd.first->dropAllUses();
       delete fwd.first;
     }
   }
 }
 
 /// After parsing is finished, this function must be called to see if there are
 /// any remaining issues.
-ParseResult OperationParser::finalize() {
+ParseResult OperationParser::finalize(Operation *topLevelOp) {
   // Check for any forward references that are left.  If we find any, error
   // out.
   if (!forwardRefPlaceholders.empty()) {
     SmallVector<const char *, 4> errors;
     // Iteration over the map isn't deterministic, so sort by source location.
     for (auto entry : forwardRefPlaceholders)
       errors.push_back(entry.second.getPointer());
     llvm::array_pod_sort(errors.begin(), errors.end());
if (!locAttr)
              << "expected location, but found '" << attr << "'";
     auto opOrArgument = it.first;
     if (auto *op = opOrArgument.dyn_cast<Operation *>())
       op->setLoc(locAttr);
     else
       opOrArgument.get<BlockArgument>().setLoc(locAttr);
   }
 
+  // If we are populating the parser state, finalize the top-level operation.
+  if (state.asmState)
+    state.asmState->finalize(topLevelOp);
+
   // Pop the top level name scope.
   return popSSANameScope();
 }
 
 //===----------------------------------------------------------------------===//
 // SSA Value Handling
 //===----------------------------------------------------------------------===//
 

 ///                        trailing-location?
 ///  generic-operation ::= string-literal `(` ssa-use-list? `)`
 ///                        successor-list? (`(` region-list `)`)?
 ///                        attribute-dict? `:` function-type
 ///  custom-operation  ::= bare-id custom-operation-format
 ///  op-result-list    ::= op-result (`,` op-result)* `=`
 ///  op-result         ::= ssa-id (`:` integer-literal)
 ///
 ParseResult OperationParser::parseOperation() {
   auto loc = getToken().getLoc();
   SmallVector<ResultRecord, 1> resultIDs;
   size_t numExpectedResults = 0;
   if (getToken().is(Token::percent_identifier)) {
     // Parse the group of result ids.
     auto parseNextResult = [&]() -> ParseResult {
       // Parse the next result id.
       if (!getToken().is(Token::percent_identifier))
if (!resultIDs.empty()) {
     if (state.asmState) {
       unsigned resultIt = 0;
       SmallVector<std::pair<unsigned, llvm::SMLoc>> asmResultGroups;
       asmResultGroups.reserve(resultIDs.size());
       for (ResultRecord &record : resultIDs) {
         asmResultGroups.emplace_back(resultIt, std::get<2>(record));
         resultIt += std::get<1>(record);
       }
-      state.asmState->addDefinition(op, nameTok.getLocRange(), asmResultGroups);
+      state.asmState->finalizeOperationDefinition(op, nameTok.getLocRange(),
+                                                  asmResultGroups);
     }
 
     // Add definitions for each of the result groups.
     unsigned opResI = 0;
     for (ResultRecord &resIt : resultIDs) {
       for (unsigned subRes : llvm::seq<unsigned>(0, std::get<1>(resIt))) {
         if (addDefinition({std::get<0>(resIt), subRes, std::get<2>(resIt)},
                           op->getResult(opResI++)))
           return failure();
       }
     }
 
     // Add this operation to the assembly state if it was provided to populate.
   } else if (state.asmState) {
-    state.asmState->addDefinition(op, nameTok.getLocRange());
+    state.asmState->finalizeOperationDefinition(op, nameTok.getLocRange());
   }
 
   return success();
 }
 
 /// Parse a single operation successor.
 ///
 ///   successor ::= block-id
Operation *OperationParser::parseGenericOperation() {
   if (!result.name.getAbstractOperation()) {
     StringRef dialectName = StringRef(name).split('.').first;
     if (!getContext()->getLoadedDialect(dialectName) &&
         getContext()->getOrLoadDialect(dialectName)) {
       result.name = OperationName(name, getContext());
     }
   }
 
+  // If we are populating the parser state, start a new operation definition.
+  if (state.asmState)
+    state.asmState->startOperationDefinition(result.name);
+
   // Parse the operand list.
   SmallVector<SSAUseInfo, 8> operandInfos;
   if (parseToken(Token::l_paren, "expected '(' to start operand list") ||
       parseOptionalSSAUseList(operandInfos) ||
       parseToken(Token::r_paren, "expected ')' to end operand list")) {
     return nullptr;
   }
 
Operation *OperationParser::parseGenericOperation() {
   Operation *op = opBuilder.createOperation(result);
   if (parseTrailingLocationSpecifier(op))
     return nullptr;
   return op;
 }
 
 Operation *OperationParser::parseGenericOperation(Block *insertBlock,
                                                   Block::iterator insertPt) {
+  Token nameToken = getToken();
+
   OpBuilder::InsertionGuard restoreInsertionPoint(opBuilder);
   opBuilder.setInsertionPoint(insertBlock, insertPt);
-  return parseGenericOperation();
+  Operation *op = parseGenericOperation();
+  if (!op)
+    return nullptr;
+
+  // If we are populating the parser asm state, finalize this operation
+  // definition.
+  if (state.asmState)
+    state.asmState->finalizeOperationDefinition(op, nameToken.getLocRange());
+  return op;
 }
 
 namespace {
 class CustomOpAsmParser : public OpAsmParser {
 public:
   CustomOpAsmParser(
       SMLoc nameLoc, ArrayRef<OperationParser::ResultRecord> resultIDs,
       function_ref<ParseResult(OpAsmParser &, OperationState &)> parseAssembly,
ParseResult parseOptionalSymbolName(StringAttr &result, StringRef attrName,
                                       NamedAttrList &attrs) override {
     Token atToken = parser.getToken();
     if (atToken.isNot(Token::at_identifier))
       return failure();
 
     result = getBuilder().getStringAttr(atToken.getSymbolReference());
     attrs.push_back(getBuilder().getNamedAttr(attrName, result));
     parser.consumeToken();
+
+    // If we are populating the assembly parser state, record this as a symbol
+    // reference.
+    if (parser.getState().asmState) {
+      parser.getState().asmState->addUses(
+          getBuilder().getSymbolRefAttr(result.getValue()),
+          atToken.getLocRange());
+    }
     return success();
   }
 
   /// Parse a loc(...) specifier if present, filling in result if so.
   ParseResult
   parseOptionalLocationSpecifier(Optional<Location> &result) override {
     // If there is a 'loc' we parse a trailing location.
     if (!parser.consumeIf(Token::kw_loc))
OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
   // If the custom op parser crashes, produce some indication to help
   // debugging.
   std::string opNameStr = opName.str();
   llvm::PrettyStackTraceFormat fmt("MLIR Parser: custom op parser '%s'",
                                    opNameStr.c_str());
 
   // Get location information for the operation.
   auto srcLocation = getEncodedSourceLocation(opLoc);
+  OperationState opState(srcLocation, opName);
+
+  // If we are populating the parser state, start a new operation definition.
+  if (state.asmState)
+    state.asmState->startOperationDefinition(opState.name);
 
   // Have the op implementation take a crack and parsing this.
-  OperationState opState(srcLocation, opName);
   CleanupOpStateRegions guard{opState};
   CustomOpAsmParser opAsmParser(opLoc, resultIDs, parseAssemblyFn,
                                 isIsolatedFromAbove, opName, *this);
   if (opAsmParser.parseOperation(opState))
     return nullptr;
 
   // If it emitted an error, we failed.
   if (opAsmParser.didEmitError())
OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Region Parsing
 //===----------------------------------------------------------------------===//
 
-/// Region.
-///
-///   region ::= '{' region-body
-///
 ParseResult OperationParser::parseRegion(
     Region &region,
     ArrayRef<std::pair<OperationParser::SSAUseInfo, Type>> entryArguments,
     bool isIsolatedNameScope) {
   // Parse the '{'.
   Token lBraceTok = getToken();
   if (parseToken(Token::l_brace, "expected '{' to begin a region"))
     return failure();
 
-  // Check for an empty region.
-  if (entryArguments.empty() && consumeIf(Token::r_brace))
-    return success();
+  // If we are populating the parser state, start a new region definition.
+  if (state.asmState)
+    state.asmState->startRegionDefinition();
+
+  // Parse the region body.
+  if ((!entryArguments.empty() || getToken().isNot(Token::r_brace)) &&
+      parseRegionBody(region, lBraceTok.getLoc(), entryArguments,
+                      isIsolatedNameScope)) {
+    return failure();
+  }
+
+  // If we are populating the parser state, finalize this region.
+  if (state.asmState)
+    state.asmState->finalizeRegionDefinition();
+  consumeToken(Token::r_brace);
+
+}
+
+ParseResult OperationParser::parseRegionBody(
+    Region &region, llvm::SMLoc startLoc,
+    ArrayRef<std::pair<OperationParser::SSAUseInfo, Type>> entryArguments,
+    bool isIsolatedNameScope) {
   auto currentPt = opBuilder.saveInsertionPoint();
 
   // Push a new named value scope.
   pushSSANameScope(isIsolatedNameScope);
 
   // Parse the first block directly to allow for it to be unnamed.
   auto owning_block = std::make_unique<Block>();
   Block *block = owning_block.get();
 
   // If this block is not defined in the source file, add a definition for it
   // now in the assembly state. Blocks with a name will be defined when the name
   // is parsed.
   if (state.asmState && getToken().isNot(Token::caret_identifier))
-    state.asmState->addDefinition(block, lBraceTok.getLoc());
+    state.asmState->addDefinition(block, startLoc);
 
   // Add arguments to the entry block.
   if (!entryArguments.empty()) {
     for (auto &placeholderArgPair : entryArguments) {
       auto &argInfo = placeholderArgPair.first;
 
       // Ensure that the argument was not already defined.
       if (auto defLoc = getReferenceLoc(argInfo.name, argInfo.number)) {
ParseResult OperationParser::parseRegionBody(
   // Verify that no other arguments were parsed.
   if (!entryArguments.empty() &&
       block->getNumArguments() > entryArguments.size()) {
     return emitError("entry block arguments were already defined");
   }
 
   // Parse the rest of the region.
   region.push_back(owning_block.release());
-  if (parseRegionBody(region))
+  while (getToken().isNot(Token::r_brace)) {
+    Block *newBlock = nullptr;
+    if (parseBlock(newBlock))
       return failure();
+    region.push_back(newBlock);
+  }
 
   // Pop the SSA value scope for this region.
   if (popSSANameScope())
     return failure();
 
   // Reset the original insertion point.
   opBuilder.restoreInsertionPoint(currentPt);
   return success();
 }
 
-/// Region.
-///
-///   region-body ::= block* '}'
-///
-ParseResult OperationParser::parseRegionBody(Region &region) {
-  // Parse the list of blocks.
-  while (!consumeIf(Token::r_brace)) {
-    Block *newBlock = nullptr;
-    if (parseBlock(newBlock))
-      return failure();
-    region.push_back(newBlock);
-  }
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // Block Parsing
 //===----------------------------------------------------------------------===//
 
 /// Block declaration.
 ///
 ///   block ::= block-label? operation*
 ///   block-label    ::= block-id block-arg-list? `:`
while (true) {
     default:
       // Parse a top-level operation.
       if (opParser.parseOperation())
         return failure();
       break;
 
     // If we got to the end of the file, then we're done.
     case Token::eof: {
-      if (opParser.finalize())
+      if (opParser.finalize(topLevelOp.get()))
         return failure();
 
       // Verify that the parsed operations are valid.
       if (failed(verify(topLevelOp.get())))
         return failure();
 
       // Splice the blocks of the parsed operation over to the provided
       // top-level block.

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

 
 using namespace mlir;
 
 /// Returns a language server position for the given source location.
 static lsp::Position getPosFromLoc(llvm::SourceMgr &mgr, llvm::SMLoc loc) {
   std::pair<unsigned, unsigned> lineAndCol = mgr.getLineAndColumn(loc);
   lsp::Position pos;
   pos.line = lineAndCol.first - 1;
-  pos.character = lineAndCol.second;
+  pos.character = lineAndCol.second - 1;
   return pos;
 }
 
 /// Returns a source location from the given language server position.
 static llvm::SMLoc getPosFromLoc(llvm::SourceMgr &mgr, lsp::Position pos) {
   return mgr.FindLocForLineAndColumn(mgr.getMainFileID(), pos.line + 1,
                                      pos.character);
 }
 
 /// Returns a language server range for the given source range.
 static lsp::Range getRangeFromLoc(llvm::SourceMgr &mgr, llvm::SMRange range) {
-  // lsp::Range is an inclusive range, SMRange is half-open.
-  llvm::SMLoc inclusiveEnd =
-      llvm::SMLoc::getFromPointer(range.End.getPointer() - 1);
-  return {getPosFromLoc(mgr, range.Start), getPosFromLoc(mgr, inclusiveEnd)};
+  return {getPosFromLoc(mgr, range.Start), getPosFromLoc(mgr, range.End)};
 }
 
 /// Returns a language server location from the given source range.
 static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr,
                                         llvm::SMRange range,
                                         const lsp::URIForFile &uri) {
   return lsp::Location{uri, getRangeFromLoc(mgr, range)};
 }
void MLIRDocument::getLocationsOf(const lsp::URIForFile &uri,
 
   // Check all definitions related to operations.
   for (const AsmParserState::OperationDefinition &op : asmState.getOpDefs()) {
     if (contains(op.loc, posLoc))
       return collectLocationsFromLoc(op.op->getLoc(), locations, uri);
     for (const auto &result : op.resultGroups)
       if (containsPosition(result.second))
         return collectLocationsFromLoc(op.op->getLoc(), locations, uri);
+    for (const auto &symUse : op.symbolUses) {
+      if (contains(symUse, posLoc)) {
+        locations.push_back(getLocationFromLoc(sourceMgr, op.loc, uri));
+        return collectLocationsFromLoc(op.op->getLoc(), locations, 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))
void MLIRDocument::findReferencesOf(const lsp::URIForFile &uri,
 
   llvm::SMLoc posLoc = getPosFromLoc(sourceMgr, pos);
 
   // Check all definitions related to operations.
   for (const AsmParserState::OperationDefinition &op : asmState.getOpDefs()) {
     if (contains(op.loc, posLoc)) {
       for (const auto &result : op.resultGroups)
         appendSMDef(result.second);
+      for (const auto &symUse : op.symbolUses)
+        if (contains(symUse, posLoc))
+          references.push_back(getLocationFromLoc(sourceMgr, symUse, uri));
       return;
     }
     for (const auto &result : op.resultGroups)
       if (isDefOrUse(result.second, posLoc))
         return appendSMDef(result.second);
+    for (const auto &symUse : op.symbolUses) {
+      if (!contains(symUse, posLoc))
+        continue;
+      for (const auto &symUse : op.symbolUses)
+        references.push_back(getLocationFromLoc(sourceMgr, symUse, uri));
+      return;
+    }
   }
 
   // Check all definitions related to blocks.
   for (const AsmParserState::BlockDefinition &block : asmState.getBlockDefs()) {
     if (isDefOrUse(block.definition, posLoc))
       return appendSMDef(block.definition);
 
     for (const AsmParserState::SMDefinition &arg : block.arguments)

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

 // 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": 1,
+// CHECK-NEXT:          "character": 0,
 // CHECK-NEXT:          "line": 3
 // 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/definition.test

 // CHECK-NEXT:  "result": [
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 6,
 // CHECK-NEXT:          "line": 1
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
-// CHECK-NEXT:          "character": 1,
+// 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":7}
+}}
+//      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": 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":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}

mlir/test/mlir-lsp-server/hover.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 @foo(%arg: i1) {\n%value = constant true\nbr ^bb2\n^bb2:\nreturn\n}"
 }}}
 // -----
 // Hover on an operation.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":1,"character":12}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "```mlir\n%true = \"std.constant\"() {value = true} : () -> i1\n```\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 17,
 // CHECK-NEXT:        "line": 1
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
-// CHECK-NEXT:        "character": 10,
+// CHECK-NEXT:        "character": 9,
 // CHECK-NEXT:        "line": 1
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on an operation result.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":1,"character":2}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "Operation: \"std.constant\"\n\nResult #0\n\nType: `i1`\n\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 6,
 // CHECK-NEXT:        "line": 1
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
-// CHECK-NEXT:        "character": 1,
+// CHECK-NEXT:        "character": 0,
 // CHECK-NEXT:        "line": 1
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a Block.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":3,"character":2}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "Operation: \"func\"\n\nBlock #1\n\nPredecessors: <Block #0>\n\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 4,
 // CHECK-NEXT:        "line": 3
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
-// CHECK-NEXT:        "character": 1,
+// CHECK-NEXT:        "character": 0,
 // CHECK-NEXT:        "line": 3
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a Block argument.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":0,"character":12}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "Operation: \"func\"\n\nBlock: <Block #0>\n\nArgument #0\n\nType: `i1`\n\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 14,
 // 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:  }
 // -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}

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

 // CHECK-NEXT:  "result": [
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 6,
 // CHECK-NEXT:          "line": 1
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
-// CHECK-NEXT:          "character": 1,
+// 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": 13,
 // CHECK-NEXT:          "line": 2
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
-// CHECK-NEXT:          "character": 8,
+// CHECK-NEXT:          "character": 7,
 // 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":0,"character":7},
+  "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": 9,
+// CHECK-NEXT:          "line": 0
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 5,
+// CHECK-NEXT:          "line": 0
+// 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"}