text: //===- 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;
  
  //===----------------------------------------------------------------------===//
  // AsmParserState::Impl
  //===----------------------------------------------------------------------===//
  
  struct AsmParserState::Impl {
    /// A map from a SymbolRefAttr to a range of uses.
    using SymbolUseMap =
        DenseMap<Attribute, SmallVector<SmallVector<SMRange>, 0>>;
  
    struct PartialOpDef {
      explicit PartialOpDef(const OperationName &opName) {
        if (opName.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 in the IR.
    void resolveSymbolUses();
  
    /// 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<SMLoc>> placeholderValueUses;
  
    /// The symbol table operations within the IR.
    SmallVector<std::pair<Operation *, std::unique_ptr<SymbolUseMap>>>
        symbolTableOperations;
  
    /// 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() {
    SmallVector<Operation *> symbolOps;
    for (auto &opAndUseMapIt : symbolTableOperations) {
      for (auto &it : *opAndUseMapIt.second) {
        symbolOps.clear();
        if (failed(symbolTable.lookupSymbolIn(
                opAndUseMapIt.first, it.first.cast<SymbolRefAttr>(), symbolOps)))
          continue;
  
        for (ArrayRef<SMRange> useRange : it.second) {
          for (const auto &symIt : llvm::zip(symbolOps, useRange)) {
            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() = default;
  AsmParserState &AsmParserState::operator=(AsmParserState &&other) {
    impl = std::move(other.impl);
    return *this;
  }
  
  //===----------------------------------------------------------------------===//
  // Access State
  
  auto AsmParserState::getBlockDefs() const -> iterator_range<BlockDefIterator> {
    return llvm::make_pointee_range(llvm::makeArrayRef(impl->blocks));
  }
  
  auto AsmParserState::getBlockDef(Block *block) const
      -> const BlockDefinition * {
    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));
  }
  
  auto AsmParserState::getOpDef(Operation *op) const
      -> const OperationDefinition * {
    auto it = impl->operationToIdx.find(op);
    return it == impl->operationToIdx.end() ? nullptr
                                            : &*impl->operations[it->second];
  }
  
  SMRange AsmParserState::convertIdLocToRange(SMLoc loc) {
    if (!loc.isValid())
      return 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 (*curPtr && isIdentifierChar(*(++curPtr)))
      continue;
    return SMRange(loc, SMLoc::getFromPointer(curPtr));
  }
  
  //===----------------------------------------------------------------------===//
  // Populate State
  
  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->symbolTableOperations.emplace_back(
          topLevelOp, std::move(partialOpDef.symbolTable));
    }
    impl->resolveSymbolUses();
  }
  
  void AsmParserState::startOperationDefinition(const OperationName &opName) {
    impl->partialOperations.emplace_back(opName);
  }
  
  void AsmParserState::finalizeOperationDefinition(
      Operation *op, SMRange nameLoc, SMLoc endLoc,
      ArrayRef<std::pair<unsigned, 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, nameLoc, endLoc);
    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->symbolTableOperations.emplace_back(
          op, std::move(partialOpDef.symbolTable));
    }
  }
  
  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, 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)));
      return;
    }
  
    // If an entry already exists, this was a forward declaration that now has a
    // proper definition.
    impl->blocks[it->second]->definition.loc = convertIdLocToRange(location);
  }
  
  void AsmParserState::addDefinition(BlockArgument blockArg,
                                     SMLoc location) {
    auto it = impl->blocksToIdx.find(blockArg.getOwner());
    assert(it != impl->blocksToIdx.end() &&
           "expected owner block to have an entry");
    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::addUses(Value value, ArrayRef<SMLoc> locations) {
    // Handle the case where the value is an operation result.
    if (OpResult result = value.dyn_cast<OpResult>()) {
      // 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 (existingIt == impl->operationToIdx.end()) {
        impl->placeholderValueUses[value].append(locations.begin(),
                                                 locations.end());
        return;
      }
  
      // If a definition does exist, locate the value's result group and add the
      // use. The result groups are ordered by increasing start index, so we just
      // need to find the last group that has a smaller/equal start index.
      unsigned resultNo = result.getResultNumber();
      OperationDefinition &def = *impl->operations[existingIt->second];
      for (auto &resultGroup : llvm::reverse(def.resultGroups)) {
        if (resultNo >= resultGroup.startIndex) {
          for (SMLoc loc : locations)
            resultGroup.definition.uses.push_back(convertIdLocToRange(loc));
          return;
        }
      }
      llvm_unreachable("expected valid result group for value use");
    }
  
    // Otherwise, this is a block argument.
    BlockArgument arg = value.cast<BlockArgument>();
    auto existingIt = impl->blocksToIdx.find(arg.getOwner());
    assert(existingIt != impl->blocksToIdx.end() &&
           "expected valid block definition for block argument");
    BlockDefinition &blockDef = *impl->blocks[existingIt->second];
    SMDefinition &argDef = blockDef.arguments[arg.getArgNumber()];
    for (SMLoc loc : locations)
      argDef.uses.emplace_back(convertIdLocToRange(loc));
  }
  
  void AsmParserState::addUses(Block *block, ArrayRef<SMLoc> locations) {
    auto it = impl->blocksToIdx.find(block);
    if (it == impl->blocksToIdx.end()) {
      it = impl->blocksToIdx.try_emplace(block, impl->blocks.size()).first;
      impl->blocks.emplace_back(std::make_unique<BlockDefinition>(block));
    }
  
    BlockDefinition &def = *impl->blocks[it->second];
    for (SMLoc loc : locations)
      def.definition.uses.push_back(convertIdLocToRange(loc));
  }
  
  void AsmParserState::addUses(SymbolRefAttr refAttr,
                               ArrayRef<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].emplace_back(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);
  }
  