[mlir][lsp] Add parser support for tracking alias definitions and uses

[Mogball]

Summary

This adds fields to AsmParserState to track attribute and type alias
definitions and uses and teachers the parser to inform the
AsmParserState about them. This will be used to add LSP support for goto
definition and find references for aliases.

Attribute aliases are tolerant to use before def, because certain
location aliases may be deferred.

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 "llvm/Support/SMLoc.h"
 #include <cstddef>
 
 namespace mlir {
 class Block;
 class BlockArgument;
 class FileLineColLoc;
 class Operation;
struct BlockDefinition {
     /// The source location for the block, i.e. the location of its name, and
     /// any uses it has.
     SMDefinition definition;
 
     /// 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) {}
+
+    /// The name of the attribute alias.
+    StringRef name;
+
+    /// The source location for the alias.
+    SMDefinition definition;
+  };
+
+  /// This class represents the information for type definition within the input
+  /// file.
+  struct TypeAliasDefinition {
+    TypeAliasDefinition(StringRef name, SMRange loc)
+        : name(name), definition(loc) {}
+
+    /// The name of the attribute alias.
+    StringRef name;
+
+    /// The source location for the alias.
+    SMDefinition definition;
+  };
+
   AsmParserState();
   ~AsmParserState();
   AsmParserState &operator=(AsmParserState &&other);
 
   //===--------------------------------------------------------------------===//
   // Access 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);
 
   /// 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
   /// 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<SMRange> refLocations);
 
   /// Refine the `oldValue` to the `newValue`. This is used to indicate that

mlir/lib/AsmParser/AsmParserState.cpp

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;
 
   /// A mapping from blocks in the input source file to their parser state.
   SmallVector<std::unique_ptr<BlockDefinition>> blocks;
   DenseMap<Block *, unsigned> blocksToIdx;
 
+  /// A mapping from aliases in the input source file to their parser state.
+  SmallVector<std::unique_ptr<AttributeAliasDefinition>> attrAliases;
+  SmallVector<std::unique_ptr<TypeAliasDefinition>> typeAliases;
+  llvm::StringMap<unsigned> attrAliasToIdx;
+  llvm::StringMap<unsigned> typeAliasToIdx;
+
   /// A set of value definitions that are placeholders for forward references.
   /// This map should be empty if the parser finishes successfully.
   DenseMap<Value, SmallVector<SMLoc>> placeholderValueUses;
 
   /// The symbol table operations within the IR.
   SmallVector<std::pair<Operation *, std::unique_ptr<SymbolUseMap>>>
       symbolTableOperations;
 
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) {
+  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));
+  } else {
+    impl->attrAliases[it->second]->definition.loc = location;
+  }
+}
+
+void AsmParserState::addTypeAliasDefinition(StringRef name, SMRange location) {
+  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));
+}
+
 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.
     Operation *parentOp = result.getOwner();
     auto existingIt = impl->operationToIdx.find(parentOp);
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].emplace_back(locations.begin(),
                                                         locations.end());
 }
 
+void AsmParserState::addAttrAliasUses(StringRef name, SMRange location) {
+  auto it = impl->attrAliasToIdx.find(name);
+  // Location aliases may be referenced before they are defined.
+  if (it == impl->attrAliasToIdx.end()) {
+    it = impl->attrAliasToIdx.try_emplace(name, impl->attrAliases.size()).first;
+    impl->attrAliases.push_back(
+        std::make_unique<AttributeAliasDefinition>(name));
+  }
+  AttributeAliasDefinition &def = *impl->attrAliases[it->second];
+  def.definition.uses.push_back(location);
+}
+
+void AsmParserState::addTypeAliasUses(StringRef name, SMRange location) {
+  auto it = impl->typeAliasToIdx.find(name);
+  // Location aliases may be referenced before they are defined.
+  assert(it != impl->typeAliasToIdx.end() &&
+         "expected valid type alias definition");
+  TypeAliasDefinition &def = *impl->typeAliases[it->second];
+  def.definition.uses.push_back(location);
+}
+
 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/AsmParser/DialectSymbolParser.cpp

ParseResult Parser::parseDialectSymbolBody(StringRef &body,
 
   unsigned length = curPtr - body.begin();
   body = StringRef(body.data(), length);
   return success();
 }
 
 /// Parse an extended dialect symbol.
 template <typename Symbol, typename SymbolAliasMap, typename CreateFn>
-static Symbol parseExtendedSymbol(Parser &p, SymbolAliasMap &aliases,
+static Symbol parseExtendedSymbol(Parser &p, AsmParserState *asmState,
+                                  SymbolAliasMap &aliases,
                                   CreateFn &&createSymbol) {
   Token tok = p.getToken();
 
   // Handle code completion of the extended symbol.
   StringRef identifier = tok.getSpelling().drop_front();
   if (tok.isCodeCompletion() && identifier.empty())
     return p.codeCompleteDialectSymbol(aliases);
 
   // Parse the dialect namespace.
+  SMRange range = p.getToken().getLocRange();
   SMLoc loc = p.getToken().getLoc();
   p.consumeToken();
 
   // Check to see if this is a pretty name.
   auto [dialectName, symbolData] = identifier.split('.');
   bool isPrettyName = !symbolData.empty() || identifier.back() == '.';
 
   // Check to see if the symbol has trailing data, i.e. has an immediately
   // following '<'.
   bool hasTrailingData =
       p.getToken().is(Token::less) &&
       identifier.bytes_end() == p.getTokenSpelling().bytes_begin();
 
   // If there is no '<' token following this, and if the typename contains no
   // dot, then we are parsing a symbol alias.
   if (!hasTrailingData && !isPrettyName) {
     // Check for an alias for this type.
     auto aliasIt = aliases.find(identifier);
     if (aliasIt == aliases.end())
       return (p.emitWrongTokenError("undefined symbol alias id '" + identifier +
                                     "'"),
               nullptr);
+    if (asmState) {
+      if constexpr (std::is_same_v<Symbol, Type>)
+        asmState->addTypeAliasUses(identifier, range);
+      else
+        asmState->addAttrAliasUses(identifier, range);
+    }
     return aliasIt->second;
   }
 
   // If this isn't an alias, we are parsing a dialect-specific symbol. If the
   // name contains a dot, then this is the "pretty" form. If not, it is the
   // verbose form that looks like <...>.
   if (!isPrettyName) {
     // Point the symbol data to the end of the dialect name to start.

 ///   dialect-attribute  ::= `#` dialect-namespace `<` attr-data `>`
 ///                          (`:` type)?
 ///                        | `#` alias-name pretty-dialect-sym-body? (`:` type)?
 ///   attribute-alias    ::= `#` alias-name
 ///
 Attribute Parser::parseExtendedAttr(Type type) {
   MLIRContext *ctx = getContext();
   Attribute attr = parseExtendedSymbol<Attribute>(
-      *this, state.symbols.attributeAliasDefinitions,
+      *this, state.asmState, state.symbols.attributeAliasDefinitions,
       [&](StringRef dialectName, StringRef symbolData, SMLoc loc) -> Attribute {
         // Parse an optional trailing colon type.
         Type attrType = type;
         if (consumeIf(Token::colon) && !(attrType = parseType()))
           return Attribute();
 
         // If we found a registered dialect, then ask it to parse the attribute.
         if (Dialect *dialect =

 ///   extended-type ::= (dialect-type | type-alias)
 ///   dialect-type  ::= `!` dialect-namespace `<` `"` type-data `"` `>`
 ///   dialect-type  ::= `!` alias-name pretty-dialect-attribute-body?
 ///   type-alias    ::= `!` alias-name
 ///
 Type Parser::parseExtendedType() {
   MLIRContext *ctx = getContext();
   return parseExtendedSymbol<Type>(
-      *this, state.symbols.typeAliasDefinitions,
+      *this, state.asmState, state.symbols.typeAliasDefinitions,
       [&](StringRef dialectName, StringRef symbolData, SMLoc loc) -> Type {
         // If we found a registered dialect, then ask it to parse the type.
         if (auto *dialect = ctx->getOrLoadDialect(dialectName)) {
           // Temporarily reset the lexer to let the dialect parse the type.
           const char *curLexerPos = getToken().getLoc().getPointer();
           resetToken(symbolData.data());
 
           // Parse the type.

mlir/lib/AsmParser/Parser.cpp

ParseResult OperationParser::parseLocationAlias(LocationAttr &loc) {
   Token tok = getToken();
   consumeToken(Token::hash_identifier);
   StringRef identifier = tok.getSpelling().drop_front();
   if (identifier.contains('.')) {
     return emitError(tok.getLoc())
            << "expected location, but found dialect attribute: '#" << identifier
            << "'";
   }
+  if (state.asmState)
+    state.asmState->addAttrAliasUses(identifier, tok.getLocRange());
 
   // If this alias can be resolved, do it now.
   Attribute attr = state.symbols.attributeAliasDefinitions.lookup(identifier);
   if (attr) {
     if (!(loc = dyn_cast<LocationAttr>(attr)))
       return emitError(tok.getLoc())
              << "expected location, but found '" << attr << "'";
   } else {
ParseResult TopLevelOperationParser::parseAttributeAliasDef() {
   if (state.symbols.attributeAliasDefinitions.count(aliasName) > 0)
     return emitError("redefinition of attribute alias id '" + aliasName + "'");
 
   // Make sure this isn't invading the dialect attribute namespace.
   if (aliasName.contains('.'))
     return emitError("attribute names with a '.' are reserved for "
                      "dialect-defined names");
 
+  SMRange location = getToken().getLocRange();
   consumeToken(Token::hash_identifier);
 
   // Parse the '='.
   if (parseToken(Token::equal, "expected '=' in attribute alias definition"))
     return failure();
 
   // 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.symbols.attributeAliasDefinitions[aliasName] = attr;
   return success();
 }
 
 ParseResult TopLevelOperationParser::parseTypeAliasDef() {
   assert(getToken().is(Token::exclamation_identifier));
   StringRef aliasName = getTokenSpelling().drop_front();
 
   // Check for redefinitions.
   if (state.symbols.typeAliasDefinitions.count(aliasName) > 0)
     return emitError("redefinition of type alias id '" + aliasName + "'");
 
   // Make sure this isn't invading the dialect type namespace.
   if (aliasName.contains('.'))
     return emitError("type names with a '.' are reserved for "
                      "dialect-defined names");
+
+  SMRange location = getToken().getLocRange();
   consumeToken(Token::exclamation_identifier);
 
   // Parse the '='.
   if (parseToken(Token::equal, "expected '=' in type alias definition"))
     return failure();
 
   // 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.symbols.typeAliasDefinitions.try_emplace(aliasName, aliasedType);
   return success();
 }
 
 ParseResult TopLevelOperationParser::parseFileMetadataDictionary() {
   consumeToken(Token::file_metadata_begin);
   return parseCommaSeparatedListUntil(
       Token::file_metadata_end, [&]() -> ParseResult {