SymbolTableLLZK.cpp

Go to the documentation of this file.

//===-- SymbolTableLLZK.cpp -------------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2025 Veridise Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Adapted from the LLVM Project's mlir/lib/IR/SymbolTable.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 "llzk/Util/SymbolTableLLZK.h"
 
#include <llvm/ADT/SmallPtrSet.h>
 
using namespace mlir;
 
//===----------------------------------------------------------------------===//
// Symbol Use Lists
//===----------------------------------------------------------------------===//
 
namespace {

static bool isPotentiallyUnknownSymbolTable(Operation *op) {
  return op->getNumRegions() == 1 && !op->getDialect();
}

static StringAttr getNameIfSymbol(Operation *op, StringAttr symbolAttrNameId) {
  return op->getAttrOfType<StringAttr>(symbolAttrNameId);
}

static LogicalResult collectValidReferencesFor(
    Operation *symbol, StringAttr symbolName, Operation *within,
    SmallVectorImpl<SymbolRefAttr> &results
) {
  assert(within->isAncestor(symbol) && "expected 'within' to be an ancestor");
  MLIRContext *ctx = symbol->getContext();
 
  auto leafRef = FlatSymbolRefAttr::get(symbolName);
  results.push_back(leafRef);
 
  // Early exit for when 'within' is the parent of 'symbol'.
  Operation *symbolTableOp = symbol->getParentOp();
  if (within == symbolTableOp) {
    return success();
  }
 
  // Collect references until 'symbolTableOp' reaches 'within'.
  SmallVector<FlatSymbolRefAttr, 1> nestedRefs(1, leafRef);
  StringAttr symbolNameId = StringAttr::get(ctx, SymbolTable::getSymbolAttrName());
  do {
    // Each parent of 'symbol' should define a symbol table.
    if (!symbolTableOp->hasTrait<OpTrait::SymbolTable>()) {
      return failure();
    }
    // Each parent of 'symbol' should also be a symbol.
    StringAttr symbolTableName = getNameIfSymbol(symbolTableOp, symbolNameId);
    if (!symbolTableName) {
      return failure();
    }
    results.push_back(SymbolRefAttr::get(symbolTableName, nestedRefs));
 
    symbolTableOp = symbolTableOp->getParentOp();
    if (symbolTableOp == within) {
      break;
    }
    nestedRefs.insert(nestedRefs.begin(), FlatSymbolRefAttr::get(symbolTableName));
  } while (true);
  return success();
}

static std::optional<WalkResult> walkSymbolTable(
    MutableArrayRef<Region> regions, function_ref<std::optional<WalkResult>(Operation *)> callback
) {
  SmallVector<Region *, 1> worklist(llvm::make_pointer_range(regions));
  while (!worklist.empty()) {
    for (Operation &op : worklist.pop_back_val()->getOps()) {
      std::optional<WalkResult> result = callback(&op);
      if (result != WalkResult::advance()) {
        return result;
      }
 
      // If this op defines a new symbol table scope, we can't traverse. Any
      // symbol references nested within 'op' are different semantically.
      if (!op.hasTrait<OpTrait::SymbolTable>()) {
        for (Region &region : op.getRegions()) {
          worklist.push_back(&region);
        }
      }
    }
  }
  return WalkResult::advance();
}

static WalkResult
walkSymbolRefs(Operation *op, function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {
  // This is modified for LLZK.
  auto walkFn = [&op, &callback](SymbolRefAttr symbolRef) {
    if (callback({op, symbolRef}).wasInterrupted()) {
      return WalkResult::interrupt();
    }
    return WalkResult::skip(); // Don't walk nested references.
  };
  for (Type t : op->getOperandTypes()) {
    if (t.walk<WalkOrder::PreOrder>(walkFn).wasInterrupted()) {
      return WalkResult::interrupt();
    }
  }
  for (Type t : op->getResultTypes()) {
    if (t.walk<WalkOrder::PreOrder>(walkFn).wasInterrupted()) {
      return WalkResult::interrupt();
    }
  }
  return op->getAttrDictionary().walk<WalkOrder::PreOrder>(walkFn);
}

static std::optional<WalkResult> walkSymbolUses(
    MutableArrayRef<Region> regions, function_ref<WalkResult(SymbolTable::SymbolUse)> callback
) {
  return walkSymbolTable(regions, [&](Operation *op) -> std::optional<WalkResult> {
    // Check that this isn't a potentially unknown symbol table.
    if (isPotentiallyUnknownSymbolTable(op)) {
      return std::nullopt;
    }
    return walkSymbolRefs(op, callback);
  });
}

static std::optional<WalkResult>
walkSymbolUses(Operation *from, function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {
  // If this operation has regions, and it, as well as its dialect, isn't
  // registered then conservatively fail. The operation may define a
  // symbol table, so we can't opaquely know if we should traverse to find
  // nested uses.
  if (isPotentiallyUnknownSymbolTable(from)) {
    return std::nullopt;
  }
 
  // Walk the uses on this operation.
  if (walkSymbolRefs(from, callback).wasInterrupted()) {
    return WalkResult::interrupt();
  }
 
  // Only recurse if this operation is not a symbol table. A symbol table
  // defines a new scope, so we can't walk the attributes from within the symbol
  // table op.
  if (!from->hasTrait<OpTrait::SymbolTable>()) {
    return walkSymbolUses(from->getRegions(), callback);
  }
  return WalkResult::advance();
}

struct SymbolScope {
  template <
      typename CallbackT,
      std::enable_if_t<!std::is_same<
          typename llvm::function_traits<CallbackT>::result_t, void>::value> * = nullptr>
  std::optional<WalkResult> walk(CallbackT cback) {
    if (Region *region = llvm::dyn_cast_if_present<Region *>(limit)) {
      return walkSymbolUses(*region, cback);
    }
    return walkSymbolUses(llvm::cast<Operation *>(limit), cback);
  }
  template <
      typename CallbackT,
      std::enable_if_t<std::is_same<
          typename llvm::function_traits<CallbackT>::result_t, void>::value> * = nullptr>
  std::optional<WalkResult> walk(CallbackT cback) {
    return walk([=](SymbolTable::SymbolUse use) { return cback(use), WalkResult::advance(); });
  }

  template <typename CallbackT>
  [[maybe_unused]]
  std::optional<WalkResult> walkSymbolTable(CallbackT &&cback) {
    if (Region *region = llvm::dyn_cast_if_present<Region *>(limit)) {
      return ::walkSymbolTable(*region, cback);
    }
    return ::walkSymbolTable(llvm::cast<Operation *>(limit), cback);
  }

  SymbolRefAttr symbol;

  llvm::PointerUnion<Operation *, Region *> limit;
};

static SmallVector<SymbolScope, 2> collectSymbolScopes(Operation *symbol, Operation *limit) {
  StringAttr symName = SymbolTable::getSymbolName(symbol);
  assert(!symbol->hasTrait<OpTrait::SymbolTable>() || symbol != limit);
 
  // Compute the ancestors of 'limit'.
  SetVector<Operation *, SmallVector<Operation *, 4>, SmallPtrSet<Operation *, 4>> limitAncestors;
  Operation *limitAncestor = limit;
  do {
    // Check to see if 'symbol' is an ancestor of 'limit'.
    if (limitAncestor == symbol) {
      // Check that the nearest symbol table is 'symbol's parent. SymbolRefAttr
      // doesn't support parent references.
      if (SymbolTable::getNearestSymbolTable(limit->getParentOp()) == symbol->getParentOp()) {
        return {{SymbolRefAttr::get(symName), limit}};
      }
      return {};
    }
 
    limitAncestors.insert(limitAncestor);
  } while ((limitAncestor = limitAncestor->getParentOp()));
 
  // Try to find the first ancestor of 'symbol' that is an ancestor of 'limit'.
  Operation *commonAncestor = symbol->getParentOp();
  do {
    if (limitAncestors.count(commonAncestor)) {
      break;
    }
  } while ((commonAncestor = commonAncestor->getParentOp()));
  assert(commonAncestor && "'limit' and 'symbol' have no common ancestor");
 
  // Compute the set of valid nested references for 'symbol' as far up to the
  // common ancestor as possible.
  SmallVector<SymbolRefAttr, 2> references;
  bool collectedAllReferences =
      succeeded(collectValidReferencesFor(symbol, symName, commonAncestor, references));
 
  // Handle the case where the common ancestor is 'limit'.
  if (commonAncestor == limit) {
    SmallVector<SymbolScope, 2> scopes;
 
    // Walk each of the ancestors of 'symbol', calling the compute function for
    // each one.
    Operation *limitIt = symbol->getParentOp();
    for (size_t i = 0, e = references.size(); i != e; ++i, limitIt = limitIt->getParentOp()) {
      assert(limitIt->hasTrait<OpTrait::SymbolTable>());
      scopes.push_back({references[i], &limitIt->getRegion(0)});
    }
    return scopes;
  }
 
  // Otherwise, we just need the symbol reference for 'symbol' that will be
  // used within 'limit'. This is the last reference in the list we computed
  // above if we were able to collect all references.
  if (!collectedAllReferences) {
    return {};
  }
  return {{references.back(), limit}};
}
 
static SmallVector<SymbolScope, 2> collectSymbolScopes(Operation *symbol, Region *limit) {
  auto scopes = collectSymbolScopes(symbol, limit->getParentOp());
 
  // If we collected some scopes to walk, make sure to constrain the one for
  // limit to the specific region requested.
  if (!scopes.empty()) {
    scopes.back().limit = limit;
  }
  return scopes;
}
 
static SmallVector<SymbolScope, 1> collectSymbolScopes(StringAttr symbol, Region *limit) {
  return {{SymbolRefAttr::get(symbol), limit}};
}
 
static SmallVector<SymbolScope, 1> collectSymbolScopes(StringAttr symbol, Operation *limit) {
  SmallVector<SymbolScope, 1> scopes;
  auto symbolRef = SymbolRefAttr::get(symbol);
  for (auto &region : limit->getRegions()) {
    scopes.push_back({symbolRef, &region});
  }
  return scopes;
}

static bool isReferencePrefixOf(SymbolRefAttr subRef, SymbolRefAttr ref) {
  if (ref == subRef) {
    return true;
  }
 
  // If the references are not pointer equal, check to see if `subRef` is a
  // prefix of `ref`.
  if (llvm::isa<FlatSymbolRefAttr>(ref) || ref.getRootReference() != subRef.getRootReference()) {
    return false;
  }
 
  auto refLeafs = ref.getNestedReferences();
  auto subRefLeafs = subRef.getNestedReferences();
  return subRefLeafs.size() < refLeafs.size() &&
         subRefLeafs == refLeafs.take_front(subRefLeafs.size());
}
 
} // namespace
 
//===----------------------------------------------------------------------===//
// llzk::getSymbolUses
 
namespace {

template <typename FromT>
static std::optional<SymbolTable::UseRange> getSymbolUsesImpl(FromT from) {
  std::vector<SymbolTable::SymbolUse> uses;
  auto walkFn = [&](SymbolTable::SymbolUse symbolUse) {
    uses.push_back(symbolUse);
    return WalkResult::advance();
  };
  auto result = walkSymbolUses(from, walkFn);
  return result ? std::optional<SymbolTable::UseRange>(std::move(uses)) : std::nullopt;
}
 
} // namespace

std::optional<SymbolTable::UseRange> llzk::getSymbolUses(Operation *from) {
  return getSymbolUsesImpl(from);
}
std::optional<SymbolTable::UseRange> llzk::getSymbolUses(Region *from) {
  return getSymbolUsesImpl(MutableArrayRef<Region>(*from));
}
 
//===----------------------------------------------------------------------===//
// llzk::getSymbolUses
 
namespace {

template <typename SymbolT, typename IRUnitT>
static std::optional<SymbolTable::UseRange> getSymbolUsesImpl(SymbolT symbol, IRUnitT *limit) {
  std::vector<SymbolTable::SymbolUse> uses;
  for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {
    if (!scope.walk([&](SymbolTable::SymbolUse symbolUse) {
      if (isReferencePrefixOf(scope.symbol, symbolUse.getSymbolRef())) {
        uses.push_back(symbolUse);
      }
    })) {
      return std::nullopt;
    }
  }
  return SymbolTable::UseRange(std::move(uses));
}
 
} // namespace

std::optional<SymbolTable::UseRange> llzk::getSymbolUses(StringAttr symbol, Operation *from) {
  return getSymbolUsesImpl(symbol, from);
}
std::optional<SymbolTable::UseRange> llzk::getSymbolUses(Operation *symbol, Operation *from) {
  return getSymbolUsesImpl(symbol, from);
}
std::optional<SymbolTable::UseRange> llzk::getSymbolUses(StringAttr symbol, Region *from) {
  return getSymbolUsesImpl(symbol, from);
}
std::optional<SymbolTable::UseRange> llzk::getSymbolUses(Operation *symbol, Region *from) {
  return getSymbolUsesImpl(symbol, from);
}
 
//===----------------------------------------------------------------------===//
// llzk::getSymbolsUsedIn
 
void llzk::getSymbolsUsedIn(Type t, llvm::SmallDenseSet<SymbolRefAttr> &symbolsUsed) {
  t.walk([&symbolsUsed](SymbolRefAttr symbolRef) { symbolsUsed.insert(symbolRef); });
}
 
void llzk::getSymbolsUsedIn(ArrayRef<Type> types, llvm::SmallDenseSet<SymbolRefAttr> &symbolsUsed) {
  for (Type t : types) {
    getSymbolsUsedIn(t, symbolsUsed);
  }
}
 
llvm::SmallDenseSet<SymbolRefAttr> llzk::getSymbolsUsedIn(Type t) {
  llvm::SmallDenseSet<SymbolRefAttr> symbolsUsed;
  getSymbolsUsedIn(t, symbolsUsed);
  return symbolsUsed;
}
 
llvm::SmallDenseSet<SymbolRefAttr> llzk::getSymbolsUsedIn(ArrayRef<Type> types) {
  llvm::SmallDenseSet<SymbolRefAttr> symbolsUsed;
  getSymbolsUsedIn(types, symbolsUsed);
  return symbolsUsed;
}
 
//===----------------------------------------------------------------------===//
// llzk::symbolKnownUseEmpty
 
namespace {

template <typename SymbolT, typename IRUnitT>
static bool symbolKnownUseEmptyImpl(SymbolT symbol, IRUnitT *limit) {
  for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {
    // Walk all of the symbol uses looking for a reference to 'symbol'.
    if (scope.walk([&](SymbolTable::SymbolUse symbolUse) {
      return isReferencePrefixOf(scope.symbol, symbolUse.getSymbolRef()) ? WalkResult::interrupt()
                                                                         : WalkResult::advance();
    }) != WalkResult::advance()) {
      return false;
    }
  }
  return true;
}
 
} // namespace

bool llzk::symbolKnownUseEmpty(StringAttr symbol, Operation *from) {
  return symbolKnownUseEmptyImpl(symbol, from);
}
bool llzk::symbolKnownUseEmpty(Operation *symbol, Operation *from) {
  return symbolKnownUseEmptyImpl(symbol, from);
}
bool llzk::symbolKnownUseEmpty(StringAttr symbol, Region *from) {
  return symbolKnownUseEmptyImpl(symbol, from);
}
bool llzk::symbolKnownUseEmpty(Operation *symbol, Region *from) {
  return symbolKnownUseEmptyImpl(symbol, from);
}
 
//===----------------------------------------------------------------------===//
// llzk::getSymbolName
 
StringAttr llzk::getSymbolName(Operation *op) {
  // This is modified for LLZK.
  // `SymbolTable::getSymbolName(Operation*)` asserts if there is no name (ex: in the case of
  // ModuleOp where the symbol name is optional) and there's no other way to check if the name
  // exists so this fully involved retrieval method must be used to return `nullptr` if no name.
  return op->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName());
}