SymbolLookup.h

Go to the documentation of this file.

//===-- SymbolLookup.h - Symbol Lookup Functions ----------------*- 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/Constants.h"
 
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/Operation.h>
#include <mlir/IR/OwningOpRef.h>
 
#include <llvm/ADT/ArrayRef.h>
#include <llvm/ADT/StringRef.h>
 
#include <variant>
#include <vector>
 
namespace llzk {
 
template <typename T> class SymbolLookupResult;
 
using ManagedResources =
    std::shared_ptr<std::pair<mlir::OwningOpRef<mlir::ModuleOp>, mlir::SymbolTableCollection>>;
 
class SymbolLookupResultUntyped {
public:
  SymbolLookupResultUntyped() : op(nullptr) {}
  SymbolLookupResultUntyped(mlir::Operation *opPtr) : op(opPtr) {}
 
  SymbolLookupResultUntyped(const SymbolLookupResultUntyped &other)
      : op(other.op), managedResources(other.managedResources),
        includeSymNameStack(other.includeSymNameStack), namespaceStack(other.namespaceStack) {}
  template <typename T> SymbolLookupResultUntyped(const SymbolLookupResult<T> &other);
 
  SymbolLookupResultUntyped &operator=(const SymbolLookupResultUntyped &other) {
    if (this == &other) {
      return *this;
    }
    this->op = other.op;
    this->managedResources = other.managedResources;
    this->includeSymNameStack = other.includeSymNameStack;
    this->namespaceStack = other.namespaceStack;
    return *this;
  }
  template <typename T> SymbolLookupResultUntyped &operator=(const SymbolLookupResult<T> &other);
 
  SymbolLookupResultUntyped(SymbolLookupResultUntyped &&other) noexcept
      : op(other.op), managedResources(std::move(other.managedResources)),
        includeSymNameStack(std::move(other.includeSymNameStack)),
        namespaceStack(std::move(other.namespaceStack)) {
    other.op = nullptr;
  }
  template <typename T> SymbolLookupResultUntyped(SymbolLookupResult<T> &&other);
 
  SymbolLookupResultUntyped &operator=(SymbolLookupResultUntyped &&other) noexcept {
    if (this != &other) {
      this->op = other.op;
      other.op = nullptr;
      this->managedResources = std::move(other.managedResources);
      this->includeSymNameStack = std::move(other.includeSymNameStack);
      this->namespaceStack = std::move(other.namespaceStack);
    }
    return *this;
  }
  template <typename T> SymbolLookupResultUntyped &operator=(SymbolLookupResult<T> &&other);

  mlir::Operation *operator->();
  mlir::Operation &operator*();
  mlir::Operation &operator*() const;
  mlir::Operation *get();
  mlir::Operation *get() const;

  operator bool() const;

  std::vector<llvm::StringRef> getIncludeSymNames() const { return includeSymNameStack; }

  llvm::ArrayRef<llvm::StringRef> getNamespace() const { return namespaceStack; }

  bool viaInclude() const { return !includeSymNameStack.empty(); }
 
  mlir::SymbolTableCollection *getSymbolTableCache() {
    if (managedResources) {
      return &managedResources->second;
    } else {
      return nullptr;
    }
  }

  bool isManaged() const { return managedResources != nullptr; }

  void manage(mlir::OwningOpRef<mlir::ModuleOp> &&ptr, mlir::SymbolTableCollection &&tables);

  void trackIncludeAsName(llvm::StringRef includeOpSymName);

  void pushNamespace(llvm::StringRef symName);

  void prependNamespace(llvm::ArrayRef<llvm::StringRef> ns);
 
  bool operator==(const SymbolLookupResultUntyped &rhs) const { return op == rhs.op; }
 
private:
  mlir::Operation *op;
  ManagedResources managedResources;
  std::vector<llvm::StringRef> includeSymNameStack;
  std::vector<llvm::StringRef> namespaceStack;
 
  friend class Within;
};
 
template <typename T> class SymbolLookupResult {
public:
  SymbolLookupResult(SymbolLookupResultUntyped &&innerRes) : inner(std::move(innerRes)) {}

  T operator->() { return llvm::dyn_cast<T>(*inner); }
  T operator*() { return llvm::dyn_cast<T>(*inner); }
  const T operator*() const { return llvm::dyn_cast<T>(*inner); }
  T get() { return llvm::dyn_cast<T>(inner.get()); }
  T get() const { return llvm::dyn_cast<T>(inner.get()); }
 
  operator bool() const { return inner && llvm::isa<T>(*inner); }

  std::vector<llvm::StringRef> getIncludeSymNames() const { return inner.getIncludeSymNames(); }

  llvm::ArrayRef<llvm::StringRef> getNamespace() const { return inner.getNamespace(); }

  void prependNamespace(llvm::ArrayRef<llvm::StringRef> ns) { inner.prependNamespace(ns); }

  bool viaInclude() const { return inner.viaInclude(); }
 
  bool operator==(const SymbolLookupResult<T> &rhs) const { return inner == rhs.inner; }

  bool isManaged() const { return inner.isManaged(); }
 
private:
  SymbolLookupResultUntyped inner;
 
  friend class Within;
  friend class SymbolLookupResultUntyped;
};
 
// These methods' definitions need to be here, after the declaration of SymbolLookupResult<T>
 
template <typename T>
SymbolLookupResultUntyped::SymbolLookupResultUntyped(const SymbolLookupResult<T> &other)
    : SymbolLookupResultUntyped(other.inner) {}
 
template <typename T>
SymbolLookupResultUntyped &
SymbolLookupResultUntyped::operator=(const SymbolLookupResult<T> &other) {
  *this = other.inner;
  return *this;
}
 
template <typename T>
SymbolLookupResultUntyped::SymbolLookupResultUntyped(SymbolLookupResult<T> &&other)
    : SymbolLookupResultUntyped(std::move(other.inner)) {}
 
template <typename T>
SymbolLookupResultUntyped &SymbolLookupResultUntyped::operator=(SymbolLookupResult<T> &&other) {
  *this = std::move(other.inner);
  return *this;
}
 
class Within {
public:
  Within() : from(nullptr) {}
  Within(mlir::Operation *op) : from(op) { assert(op && "cannot lookup within nullptr"); }
  Within(SymbolLookupResultUntyped &&res) : from(std::move(res)) {}
  template <typename T> Within(SymbolLookupResult<T> &&res) : Within(std::move(res.inner)) {}
 
  Within(const Within &) = delete;
  Within(Within &&other) noexcept : from(std::move(other.from)) {}
  Within &operator=(const Within &) = delete;
  Within &operator=(Within &&) noexcept;
 
  inline static Within root() { return Within(); }
 
  mlir::FailureOr<SymbolLookupResultUntyped> lookup(
      mlir::SymbolTableCollection &tables, mlir::SymbolRefAttr symbol, mlir::Operation *origin,
      bool reportMissing = true
  ) &&;
 
private:
  std::variant<mlir::Operation *, SymbolLookupResultUntyped> from;
};
 
inline mlir::FailureOr<SymbolLookupResultUntyped> lookupSymbolIn(
    mlir::SymbolTableCollection &tables, mlir::SymbolRefAttr symbol, Within &&lookupWithin,
    mlir::Operation *origin, bool reportMissing = true
) {
  return std::move(lookupWithin).lookup(tables, symbol, origin, reportMissing);
}
 
inline mlir::FailureOr<SymbolLookupResultUntyped> lookupTopLevelSymbol(
    mlir::SymbolTableCollection &tables, mlir::SymbolRefAttr symbol, mlir::Operation *origin,
    bool reportMissing = true
) {
  return Within().lookup(tables, symbol, origin, reportMissing);
}
 
template <typename T>
inline mlir::FailureOr<SymbolLookupResult<T>> lookupSymbolIn(
    mlir::SymbolTableCollection &tables, mlir::SymbolRefAttr symbol, Within &&lookupWithin,
    mlir::Operation *origin, bool reportMissing = true
) {
  auto found = lookupSymbolIn(tables, symbol, std::move(lookupWithin), origin, reportMissing);
  if (mlir::failed(found)) {
    return mlir::failure(); // lookupSymbolIn() already emits a sufficient error message
  }
  // Keep a copy of the op ptr in case we need it for displaying diagnostics
  mlir::Operation *op = found->get();
  // ... since the untyped result gets moved here into a typed result.
  SymbolLookupResult<T> ret(std::move(*found));
  if (!ret) {
    if (reportMissing) {
      return origin->emitError() << "symbol \"" << symbol << "\" references a '" << op->getName()
                                 << "' but expected a '" << T::getOperationName() << '\'';
    } else {
      return mlir::failure();
    }
  }
  return ret;
}
 
template <typename T>
inline mlir::FailureOr<SymbolLookupResult<T>> lookupTopLevelSymbol(
    mlir::SymbolTableCollection &tables, mlir::SymbolRefAttr symbol, mlir::Operation *origin,
    bool reportMissing = true
) {
  return lookupSymbolIn<T>(tables, symbol, Within(), origin, reportMissing);
}
 
} // namespace llzk