SourceRef.h

Go to the documentation of this file.

//===-- SourceRef.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/Analysis/AbstractLatticeValue.h"
#include "llzk/Dialect/Felt/IR/Ops.h"
#include "llzk/Dialect/LLZK/IR/AttributeHelper.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Util/DynamicAPIntHelper.h"
#include "llzk/Util/ErrorHelper.h"
#include "llzk/Util/Hash.h"
 
#include <mlir/Analysis/DataFlowFramework.h>
#include <mlir/Dialect/Arith/IR/Arith.h>
#include <mlir/Pass/AnalysisManager.h>
 
#include <llvm/ADT/DynamicAPInt.h>
#include <llvm/ADT/EquivalenceClasses.h>
 
#include <unordered_set>
#include <vector>
 
namespace llzk {

class SourceRefIndex {
  using IndexRange = std::pair<llvm::DynamicAPInt, llvm::DynamicAPInt>;
 
public:
  explicit SourceRefIndex(component::MemberDefOp f) : index(f) {}
  explicit SourceRefIndex(SymbolLookupResult<component::MemberDefOp> f) : index(f) {}
  explicit SourceRefIndex(const llvm::DynamicAPInt &i) : index(i) {}
  explicit SourceRefIndex(const llvm::APInt &i) : index(toDynamicAPInt(i)) {}
  explicit SourceRefIndex(int64_t i) : index(llvm::DynamicAPInt(i)) {}
  SourceRefIndex(const llvm::APInt &low, const llvm::APInt &high)
      : index(IndexRange {toDynamicAPInt(low), toDynamicAPInt(high)}) {}
  explicit SourceRefIndex(IndexRange r) : index(r) {}
 
  bool isMember() const {
    return std::holds_alternative<SymbolLookupResult<component::MemberDefOp>>(index) ||
           std::holds_alternative<component::MemberDefOp>(index);
  }
  component::MemberDefOp getMember() const {
    ensure(isMember(), "SourceRefIndex: member requested but not contained");
    if (std::holds_alternative<component::MemberDefOp>(index)) {
      return std::get<component::MemberDefOp>(index);
    }
    return std::get<SymbolLookupResult<component::MemberDefOp>>(index).get();
  }
 
  bool isIndex() const { return std::holds_alternative<llvm::DynamicAPInt>(index); }
  llvm::DynamicAPInt getIndex() const {
    ensure(isIndex(), "SourceRefIndex: index requested but not contained");
    return std::get<llvm::DynamicAPInt>(index);
  }
 
  bool isIndexRange() const { return std::holds_alternative<IndexRange>(index); }
  IndexRange getIndexRange() const {
    ensure(isIndexRange(), "SourceRefIndex: index range requested but not contained");
    return std::get<IndexRange>(index);
  }
 
  inline void dump() const { print(llvm::errs()); }
  void print(mlir::raw_ostream &os) const;
 
  inline bool operator==(const SourceRefIndex &rhs) const {
    if (isMember() && rhs.isMember()) {
      // We compare the underlying members, since the member could be in a symbol
      // lookup or not.
      return getMember() == rhs.getMember();
    }
    if (isIndex() && rhs.isIndex()) {
      return getIndex() == rhs.getIndex();
    }
    return index == rhs.index;
  }
 
  bool operator<(const SourceRefIndex &rhs) const;
 
  bool operator>(const SourceRefIndex &rhs) const { return rhs < *this; }
 
  struct Hash {
    size_t operator()(const SourceRefIndex &c) const;
  };
 
  inline size_t getHash() const { return Hash {}(*this); }
 
private:
  std::variant<
      component::MemberDefOp, SymbolLookupResult<component::MemberDefOp>, llvm::DynamicAPInt,
      IndexRange>
      index;
};
 
static inline mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const SourceRefIndex &rhs) {
  rhs.print(os);
  return os;
}

class SourceRef {
 
public:
  static std::vector<SourceRef>
  getAllSourceRefs(mlir::SymbolTableCollection &tables, mlir::ModuleOp mod, SourceRef root);

  static std::vector<SourceRef>
  getAllSourceRefs(component::StructDefOp structDef, function::FuncDefOp fnOp);

  static std::vector<SourceRef>
  getAllSourceRefs(component::StructDefOp structDef, component::MemberDefOp memberDef);
 
  explicit SourceRef(mlir::BlockArgument b) : root(b), memberRefs(), constantVal(std::nullopt) {}
  SourceRef(mlir::BlockArgument b, std::vector<SourceRefIndex> f)
      : root(b), memberRefs(std::move(f)), constantVal(std::nullopt) {}
 
  explicit SourceRef(component::CreateStructOp createOp)
      : root(createOp), memberRefs(), constantVal(std::nullopt) {}
  SourceRef(component::CreateStructOp createOp, std::vector<SourceRefIndex> f)
      : root(createOp), memberRefs(std::move(f)), constantVal(std::nullopt) {}
 
  explicit SourceRef(felt::FeltConstantOp c) : root(std::nullopt), memberRefs(), constantVal(c) {}
  explicit SourceRef(mlir::arith::ConstantIndexOp c)
      : root(std::nullopt), memberRefs(), constantVal(c) {}
  explicit SourceRef(polymorphic::ConstReadOp c)
      : root(std::nullopt), memberRefs(), constantVal(c) {}
 
  mlir::Type getType() const;
 
  bool isConstantFelt() const {
    return constantVal.has_value() && std::holds_alternative<felt::FeltConstantOp>(*constantVal);
  }
  bool isConstantIndex() const {
    return constantVal.has_value() &&
           std::holds_alternative<mlir::arith::ConstantIndexOp>(*constantVal);
  }
  bool isTemplateConstant() const {
    return constantVal.has_value() &&
           std::holds_alternative<polymorphic::ConstReadOp>(*constantVal);
  }
  bool isConstant() const { return constantVal.has_value(); }
  bool isConstantInt() const { return isConstantFelt() || isConstantIndex(); }
 
  bool isFeltVal() const { return llvm::isa<felt::FeltType>(getType()); }
  bool isIndexVal() const { return llvm::isa<mlir::IndexType>(getType()); }
  bool isIntegerVal() const { return llvm::isa<mlir::IntegerType>(getType()); }
  bool isTypeVarVal() const { return llvm::isa<polymorphic::TypeVarType>(getType()); }
  bool isScalar() const {
    return isConstant() || isFeltVal() || isIndexVal() || isIntegerVal() || isTypeVarVal();
  }
 
  bool isBlockArgument() const {
    return root.has_value() && std::holds_alternative<mlir::BlockArgument>(*root);
  }
  mlir::BlockArgument getBlockArgument() const {
    ensure(isBlockArgument(), "is not a block argument");
    return std::get<mlir::BlockArgument>(*root);
  }
  unsigned getInputNum() const { return getBlockArgument().getArgNumber(); }
 
  bool isCreateStructOp() const {
    return root.has_value() && std::holds_alternative<component::CreateStructOp>(*root);
  }
  component::CreateStructOp getCreateStructOp() const {
    ensure(isCreateStructOp(), "is not a create struct op");
    return std::get<component::CreateStructOp>(*root);
  }
 
  llvm::DynamicAPInt getConstantFeltValue() const {
    ensure(
        isConstantFelt(), mlir::Twine(mlir::StringRef(__FUNCTION__), " requires a constant felt!")
    );
    llvm::APInt i = std::get<felt::FeltConstantOp>(*constantVal).getValue();
    return toDynamicAPInt(i);
  }
  llvm::DynamicAPInt getConstantIndexValue() const {
    ensure(
        isConstantIndex(), mlir::Twine(mlir::StringRef(__FUNCTION__), " requires a constant index!")
    );
    return llvm::DynamicAPInt(std::get<mlir::arith::ConstantIndexOp>(*constantVal).value());
  }
  llvm::DynamicAPInt getConstantValue() const {
    ensure(
        isConstantFelt() || isConstantIndex(),
        mlir::Twine(mlir::StringRef(__FUNCTION__), " requires a constant int type!")
    );
    return isConstantFelt() ? getConstantFeltValue() : getConstantIndexValue();
  }

  bool isValidPrefix(const SourceRef &prefix) const;

  mlir::FailureOr<std::vector<SourceRefIndex>> getSuffix(const SourceRef &prefix) const;

  mlir::FailureOr<SourceRef> translate(const SourceRef &prefix, const SourceRef &other) const;

  mlir::FailureOr<SourceRef> getParentPrefix() const {
    if (isConstantFelt() || memberRefs.empty()) {
      return mlir::failure();
    }
    auto copy = *this;
    copy.memberRefs.pop_back();
    return copy;
  }

  std::vector<SourceRef>
  getAllChildren(mlir::SymbolTableCollection &tables, mlir::ModuleOp mod) const;
 
  SourceRef createChild(SourceRefIndex r) const {
    auto copy = *this;
    copy.memberRefs.push_back(r);
    return copy;
  }
 
  SourceRef createChild(SourceRef other) const {
    assert(other.isConstantIndex());
    return createChild(SourceRefIndex(other.getConstantIndexValue()));
  }
 
  const std::vector<SourceRefIndex> &getPieces() const { return memberRefs; }
 
  void print(mlir::raw_ostream &os) const;
  void dump() const { print(llvm::errs()); }
 
  bool operator==(const SourceRef &rhs) const;
 
  bool operator!=(const SourceRef &rhs) const { return !(*this == rhs); }
 
  // required for EquivalenceClasses usage
  bool operator<(const SourceRef &rhs) const;
 
  bool operator>(const SourceRef &rhs) const { return rhs < *this; }
 
  struct Hash {
    size_t operator()(const SourceRef &val) const;
  };
 
private:
  std::optional<std::variant<mlir::BlockArgument, component::CreateStructOp>> root;
 
  std::vector<SourceRefIndex> memberRefs;
  // using mutable to reduce constant casts for certain get* functions.
  mutable std::optional<
      std::variant<felt::FeltConstantOp, mlir::arith::ConstantIndexOp, polymorphic::ConstReadOp>>
      constantVal;
};
 
mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const SourceRef &rhs);
 
/* SourceRefSet */
 
class SourceRefSet : public std::unordered_set<SourceRef, SourceRef::Hash> {
  using Base = std::unordered_set<SourceRef, SourceRef::Hash>;
 
public:
  using Base::Base;
 
  SourceRefSet &join(const SourceRefSet &rhs);
 
  friend mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const SourceRefSet &rhs);
};
 
static_assert(
    dataflow::ScalarLatticeValue<SourceRefSet>,
    "SourceRefSet must satisfy the ScalarLatticeValue requirements"
);
 
} // namespace llzk
 
namespace llvm {
 
template <> struct DenseMapInfo<llzk::SourceRef> {
  static llzk::SourceRef getEmptyKey() {
    return llzk::SourceRef(mlir::BlockArgument(reinterpret_cast<mlir::detail::ValueImpl *>(1)));
  }
  static inline llzk::SourceRef getTombstoneKey() {
    return llzk::SourceRef(mlir::BlockArgument(reinterpret_cast<mlir::detail::ValueImpl *>(2)));
  }
  static unsigned getHashValue(const llzk::SourceRef &ref) {
    if (ref == getEmptyKey() || ref == getTombstoneKey()) {
      return llvm::hash_value(ref.getBlockArgument().getAsOpaquePointer());
    }
    return llzk::SourceRef::Hash {}(ref);
  }
  static bool isEqual(const llzk::SourceRef &lhs, const llzk::SourceRef &rhs) { return lhs == rhs; }
};
 
} // namespace llvm