ConstrainRef.h

Go to the documentation of this file.

//===-- ConstrainRef.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/APIntHelper.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/EquivalenceClasses.h>
 
#include <unordered_set>
#include <vector>
 
namespace llzk {

class ConstrainRefIndex {
  using IndexRange = std::pair<mlir::APInt, mlir::APInt>;
 
public:
  explicit ConstrainRefIndex(component::FieldDefOp f) : index(f) {}
  explicit ConstrainRefIndex(SymbolLookupResult<component::FieldDefOp> f) : index(f) {}
  explicit ConstrainRefIndex(mlir::APInt i) : index(i) {}
  explicit ConstrainRefIndex(int64_t i) : index(toAPInt(i)) {}
  ConstrainRefIndex(mlir::APInt low, mlir::APInt high) : index(IndexRange {low, high}) {}
  explicit ConstrainRefIndex(IndexRange r) : index(r) {}
 
  bool isField() const {
    return std::holds_alternative<SymbolLookupResult<component::FieldDefOp>>(index) ||
           std::holds_alternative<component::FieldDefOp>(index);
  }
  component::FieldDefOp getField() const {
    ensure(isField(), "ConstrainRefIndex: field requested but not contained");
    if (std::holds_alternative<component::FieldDefOp>(index)) {
      return std::get<component::FieldDefOp>(index);
    }
    return std::get<SymbolLookupResult<component::FieldDefOp>>(index).get();
  }
 
  bool isIndex() const { return std::holds_alternative<mlir::APInt>(index); }
  mlir::APInt getIndex() const {
    ensure(isIndex(), "ConstrainRefIndex: index requested but not contained");
    return std::get<mlir::APInt>(index);
  }
 
  bool isIndexRange() const { return std::holds_alternative<IndexRange>(index); }
  IndexRange getIndexRange() const {
    ensure(isIndexRange(), "ConstrainRefIndex: 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 ConstrainRefIndex &rhs) const {
    if (isField() && rhs.isField()) {
      // We compare the underlying fields, since the field could be in a symbol
      // lookup or not.
      return getField() == rhs.getField();
    }
    if (isIndex() && rhs.isIndex()) {
      return safeEq(mlir::APSInt(getIndex()), mlir::APSInt(rhs.getIndex()));
    }
    return index == rhs.index;
  }
 
  bool operator<(const ConstrainRefIndex &rhs) const;
 
  bool operator>(const ConstrainRefIndex &rhs) const { return rhs < *this; }
 
  struct Hash {
    size_t operator()(const ConstrainRefIndex &c) const;
  };
 
  inline size_t getHash() const { return Hash {}(*this); }
 
private:
  std::variant<
      component::FieldDefOp, SymbolLookupResult<component::FieldDefOp>, mlir::APInt, IndexRange>
      index;
};
 
static inline mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const ConstrainRefIndex &rhs) {
  rhs.print(os);
  return os;
}

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

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

  static std::vector<ConstrainRef>
  getAllConstrainRefs(component::StructDefOp structDef, component::FieldDefOp fieldDef);
 
  explicit ConstrainRef(mlir::BlockArgument b) : root(b), fieldRefs(), constantVal(std::nullopt) {}
  ConstrainRef(mlir::BlockArgument b, std::vector<ConstrainRefIndex> f)
      : root(b), fieldRefs(std::move(f)), constantVal(std::nullopt) {}
 
  explicit ConstrainRef(component::CreateStructOp createOp)
      : root(createOp), fieldRefs(), constantVal(std::nullopt) {}
  ConstrainRef(component::CreateStructOp createOp, std::vector<ConstrainRefIndex> f)
      : root(createOp), fieldRefs(std::move(f)), constantVal(std::nullopt) {}
 
  explicit ConstrainRef(felt::FeltConstantOp c) : root(std::nullopt), fieldRefs(), constantVal(c) {}
  explicit ConstrainRef(mlir::arith::ConstantIndexOp c)
      : root(std::nullopt), fieldRefs(), constantVal(c) {}
  explicit ConstrainRef(polymorphic::ConstReadOp c)
      : root(std::nullopt), fieldRefs(), 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 mlir::isa<felt::FeltType>(getType()); }
  bool isIndexVal() const { return mlir::isa<mlir::IndexType>(getType()); }
  bool isIntegerVal() const { return mlir::isa<mlir::IntegerType>(getType()); }
  bool isTypeVarVal() const { return mlir::isa<polymorphic::TypeVarType>(getType()); }
  bool isScalar() const {
    return isConstant() || isFeltVal() || isIndexVal() || isIntegerVal() || isTypeVarVal();
  }
  bool isSignal() const { return isSignalType(getType()); }
 
  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);
  }
 
  mlir::APInt getConstantFeltValue() const {
    ensure(isConstantFelt(), __FUNCTION__ + mlir::Twine(" requires a constant felt!"));
    return std::get<felt::FeltConstantOp>(*constantVal).getValueAttr().getValue();
  }
  mlir::APInt getConstantIndexValue() const {
    ensure(isConstantIndex(), __FUNCTION__ + mlir::Twine(" requires a constant index!"));
    return toAPInt(std::get<mlir::arith::ConstantIndexOp>(*constantVal).value());
  }
  mlir::APInt getConstantValue() const {
    ensure(
        isConstantFelt() || isConstantIndex(),
        __FUNCTION__ + mlir::Twine(" requires a constant int type!")
    );
    return isConstantFelt() ? getConstantFeltValue() : getConstantIndexValue();
  }

  bool isValidPrefix(const ConstrainRef &prefix) const;

  mlir::FailureOr<std::vector<ConstrainRefIndex>> getSuffix(const ConstrainRef &prefix) const;

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

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

  std::vector<ConstrainRef>
  getAllChildren(mlir::SymbolTableCollection &tables, mlir::ModuleOp mod) const;
 
  ConstrainRef createChild(ConstrainRefIndex r) const {
    auto copy = *this;
    copy.fieldRefs.push_back(r);
    return copy;
  }
 
  ConstrainRef createChild(ConstrainRef other) const {
    assert(other.isConstantIndex());
    return createChild(ConstrainRefIndex(other.getConstantIndexValue()));
  }
 
  const std::vector<ConstrainRefIndex> &getPieces() const { return fieldRefs; }
 
  void print(mlir::raw_ostream &os) const;
  void dump() const { print(llvm::errs()); }
 
  bool operator==(const ConstrainRef &rhs) const;
 
  bool operator!=(const ConstrainRef &rhs) const { return !(*this == rhs); }
 
  // required for EquivalenceClasses usage
  bool operator<(const ConstrainRef &rhs) const;
 
  bool operator>(const ConstrainRef &rhs) const { return rhs < *this; }
 
  struct Hash {
    size_t operator()(const ConstrainRef &val) const;
  };
 
private:
  std::optional<std::variant<mlir::BlockArgument, component::CreateStructOp>> root;
 
  std::vector<ConstrainRefIndex> fieldRefs;
  // 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 ConstrainRef &rhs);
 
/* ConstrainRefSet */
 
class ConstrainRefSet : public std::unordered_set<ConstrainRef, ConstrainRef::Hash> {
  using Base = std::unordered_set<ConstrainRef, ConstrainRef::Hash>;
 
public:
  using Base::Base;
 
  ConstrainRefSet &join(const ConstrainRefSet &rhs);
 
  friend mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const ConstrainRefSet &rhs);
};
 
static_assert(
    dataflow::ScalarLatticeValue<ConstrainRefSet>,
    "ConstrainRefSet must satisfy the ScalarLatticeValue requirements"
);
 
} // namespace llzk
 
namespace llvm {
 
template <> struct DenseMapInfo<llzk::ConstrainRef> {
  static llzk::ConstrainRef getEmptyKey() {
    return llzk::ConstrainRef(mlir::BlockArgument(reinterpret_cast<mlir::detail::ValueImpl *>(1)));
  }
  static inline llzk::ConstrainRef getTombstoneKey() {
    return llzk::ConstrainRef(mlir::BlockArgument(reinterpret_cast<mlir::detail::ValueImpl *>(2)));
  }
  static unsigned getHashValue(const llzk::ConstrainRef &ref) {
    if (ref == getEmptyKey() || ref == getTombstoneKey()) {
      return llvm::hash_value(ref.getBlockArgument().getAsOpaquePointer());
    }
    return llzk::ConstrainRef::Hash {}(ref);
  }
  static bool isEqual(const llzk::ConstrainRef &lhs, const llzk::ConstrainRef &rhs) {
    return lhs == rhs;
  }
};
 
} // namespace llvm