SymbolHelper.h

Go to the documentation of this file.

//===-- SymbolHelper.h ------------------------------------------*- 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
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "llzk/Util/SymbolLookup.h"
 
#include <mlir/Interfaces/CallInterfaces.h>
 
#include <cassert>
#include <optional>
#include <ranges>
 
namespace llzk {
 
namespace component {
class StructType;
class StructDefOp;
class MemberDefOp;
} // namespace component
 
namespace function {
class FuncDefOp;
} // namespace function
namespace polymorphic {
class TemplateOp;
} // namespace polymorphic
 
llvm::SmallVector<mlir::StringRef> getNames(mlir::SymbolRefAttr ref);
llvm::SmallVector<mlir::FlatSymbolRefAttr> getPieces(mlir::SymbolRefAttr ref);

inline mlir::FlatSymbolRefAttr
getFlatSymbolRefAttr(mlir::MLIRContext *context, const mlir::Twine &twine) {
  return mlir::FlatSymbolRefAttr::get(mlir::StringAttr::get(context, twine));
}

inline mlir::SymbolRefAttr asSymbolRefAttr(mlir::StringAttr root, mlir::SymbolRefAttr tail) {
  return mlir::SymbolRefAttr::get(root, getPieces(tail));
}

inline mlir::SymbolRefAttr asSymbolRefAttr(llvm::ArrayRef<mlir::FlatSymbolRefAttr> path) {
  return mlir::SymbolRefAttr::get(path.front().getAttr(), path.drop_front());
}

inline mlir::SymbolRefAttr asSymbolRefAttr(const std::vector<mlir::FlatSymbolRefAttr> &path) {
  return asSymbolRefAttr(llvm::ArrayRef<mlir::FlatSymbolRefAttr>(path));
}

inline mlir::SymbolRefAttr getTailAsSymbolRefAttr(mlir::SymbolRefAttr symbol) {
  return asSymbolRefAttr(symbol.getNestedReferences());
}

inline mlir::SymbolRefAttr getPrefixAsSymbolRefAttr(mlir::SymbolRefAttr symbol) {
  return mlir::SymbolRefAttr::get(
      symbol.getRootReference(), symbol.getNestedReferences().drop_back()
  );
}

mlir::SymbolRefAttr replaceLeaf(mlir::SymbolRefAttr orig, mlir::FlatSymbolRefAttr newLeaf);
inline mlir::SymbolRefAttr replaceLeaf(mlir::SymbolRefAttr orig, mlir::StringAttr newLeaf) {
  return replaceLeaf(orig, mlir::FlatSymbolRefAttr::get(newLeaf));
}
inline mlir::SymbolRefAttr replaceLeaf(mlir::SymbolRefAttr orig, const mlir::Twine &newLeaf) {
  return replaceLeaf(orig, mlir::StringAttr::get(orig.getContext(), newLeaf));
}

mlir::SymbolRefAttr appendLeaf(mlir::SymbolRefAttr orig, mlir::FlatSymbolRefAttr newLeaf);
inline mlir::SymbolRefAttr appendLeaf(mlir::SymbolRefAttr orig, mlir::StringAttr newLeaf) {
  return appendLeaf(orig, mlir::FlatSymbolRefAttr::get(newLeaf));
}
inline mlir::SymbolRefAttr appendLeaf(mlir::SymbolRefAttr orig, const mlir::Twine &newLeaf) {
  return appendLeaf(orig, mlir::StringAttr::get(orig.getContext(), newLeaf));
}

mlir::SymbolRefAttr appendLeafName(mlir::SymbolRefAttr orig, const mlir::Twine &newLeafSuffix);

mlir::FailureOr<mlir::ModuleOp> getRootModule(mlir::Operation *from);
mlir::FailureOr<mlir::SymbolRefAttr>
getPathFromRoot(mlir::SymbolOpInterface to, mlir::ModuleOp *foundRoot = nullptr);
mlir::FailureOr<mlir::SymbolRefAttr>
getPathFromRoot(component::StructDefOp &to, mlir::ModuleOp *foundRoot = nullptr);
mlir::FailureOr<mlir::SymbolRefAttr>
getPathFromRoot(component::MemberDefOp &to, mlir::ModuleOp *foundRoot = nullptr);
mlir::FailureOr<mlir::SymbolRefAttr>
getPathFromRoot(function::FuncDefOp &to, mlir::ModuleOp *foundRoot = nullptr);

inline mlir::SymbolRefAttr
getFullyQualifiedName(mlir::SymbolOpInterface symbol, bool requireParent = true) {
  if (!requireParent && symbol.getOperation()->getParentOp() == nullptr) {
    return mlir::SymbolRefAttr::get(symbol.getOperation());
  }
  mlir::FailureOr<mlir::SymbolRefAttr> res = getPathFromRoot(symbol);
  assert(mlir::succeeded(res));
  return res.value();
}

mlir::FailureOr<mlir::ModuleOp> getTopRootModule(mlir::Operation *from);
mlir::FailureOr<mlir::SymbolRefAttr>
getPathFromTopRoot(mlir::SymbolOpInterface to, mlir::ModuleOp *foundRoot = nullptr);
mlir::FailureOr<mlir::SymbolRefAttr>
getPathFromTopRoot(component::StructDefOp &to, mlir::ModuleOp *foundRoot = nullptr);
mlir::FailureOr<mlir::SymbolRefAttr>
getPathFromTopRoot(component::MemberDefOp &to, mlir::ModuleOp *foundRoot = nullptr);
mlir::FailureOr<mlir::SymbolRefAttr>
getPathFromTopRoot(function::FuncDefOp &to, mlir::ModuleOp *foundRoot = nullptr);

mlir::FailureOr<llzk::component::StructType> getMainInstanceType(mlir::Operation *lookupFrom);

mlir::FailureOr<SymbolLookupResult<llzk::component::StructDefOp>>
getMainInstanceDef(mlir::SymbolTableCollection &symbolTable, mlir::Operation *lookupFrom);

template <typename T>
inline mlir::FailureOr<SymbolLookupResult<T>>
resolveCallable(mlir::SymbolTableCollection &symbolTable, mlir::CallOpInterface call) {
  mlir::CallInterfaceCallable callable = call.getCallableForCallee();
  if (auto symbolVal = llvm::dyn_cast<mlir::Value>(callable)) {
    return SymbolLookupResult<T>(symbolVal.getDefiningOp());
  }
 
  // If the callable isn't a value, lookup the symbol reference.
  // We first try to resolve in the nearest symbol table, as per the default
  // MLIR behavior. If the resulting operation is not found, we will then
  // use the LLZK lookup helpers.
  auto symbolRef = llvm::cast<mlir::SymbolRefAttr>(callable);
  mlir::Operation *op = symbolTable.lookupNearestSymbolFrom(call.getOperation(), symbolRef);
 
  if (op) {
    return SymbolLookupResult<T>(std::move(op));
  }
  // Otherwise, use the top-level lookup.
  return lookupTopLevelSymbol<T>(symbolTable, symbolRef, call.getOperation());
}

template <typename T>
inline mlir::FailureOr<SymbolLookupResult<T>>
resolveCallableSilently(mlir::SymbolTableCollection &symbolTable, mlir::CallOpInterface call) {
  mlir::CallInterfaceCallable callable = call.getCallableForCallee();
  if (auto symbolVal = llvm::dyn_cast<mlir::Value>(callable)) {
    SymbolLookupResult<T> result(symbolVal.getDefiningOp());
    if (!result) {
      return mlir::failure();
    }
    return result;
  }
 
  auto symbolRef = llvm::cast<mlir::SymbolRefAttr>(callable);
  if (mlir::Operation *op = symbolTable.lookupNearestSymbolFrom(call.getOperation(), symbolRef)) {
    SymbolLookupResult<T> result(op);
    if (!result) {
      return mlir::failure();
    }
    return result;
  }
  return lookupTopLevelSymbol<T>(
      symbolTable, symbolRef, call.getOperation(), /*reportMissing=*/false
  );
}
 
template <typename T>
inline mlir::FailureOr<SymbolLookupResult<T>> resolveCallable(mlir::CallOpInterface call) {
  mlir::SymbolTableCollection symbolTable;
  return resolveCallable<T>(symbolTable, call);
}

mlir::FailureOr<polymorphic::TemplateOp>
getConstResolutionTemplate(mlir::SymbolTableCollection &tables, mlir::Operation *origin);

mlir::LogicalResult verifyParamOfType(
    mlir::SymbolTableCollection &tables, mlir::SymbolRefAttr param, mlir::Type structOrArrayType,
    mlir::Operation *origin, std::optional<mlir::Type> requiredParamType = std::nullopt
);

mlir::LogicalResult verifyParamsOfType(
    mlir::SymbolTableCollection &tables, mlir::ArrayRef<mlir::Attribute> tyParams,
    mlir::Type structOrArrayType, mlir::Operation *origin,
    std::optional<mlir::Type> requiredParamType = std::nullopt
);

mlir::FailureOr<component::StructDefOp> verifyStructTypeResolution(
    mlir::SymbolTableCollection &tables, component::StructType ty, mlir::Operation *origin
);

mlir::LogicalResult
verifyTypeResolution(mlir::SymbolTableCollection &tables, mlir::Operation *origin, mlir::Type type);

template <std::ranges::input_range Range>
mlir::LogicalResult verifyTypeResolution(
    mlir::SymbolTableCollection &tables, mlir::Operation *origin, const Range &types
) {
  // Check all before returning to present all applicable type errors in one compilation.
  bool failed = false;
  for (const auto &t : types) {
    failed |= mlir::failed(verifyTypeResolution(tables, origin, t));
  }
  return mlir::LogicalResult::failure(failed);
}
 
} // namespace llzk