IntervalAnalysis.h

Go to the documentation of this file.

//===-- IntervalAnalysis.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/Analysis/AnalysisWrappers.h"
#include "llzk/Analysis/ConstraintDependencyGraph.h"
#include "llzk/Analysis/DenseAnalysis.h"
#include "llzk/Analysis/Field.h"
#include "llzk/Analysis/Intervals.h"
#include "llzk/Dialect/Array/IR/Ops.h"
#include "llzk/Dialect/Bool/IR/Ops.h"
#include "llzk/Dialect/Cast/IR/Ops.h"
#include "llzk/Dialect/Constrain/IR/Ops.h"
#include "llzk/Dialect/Felt/IR/Ops.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/Global/IR/Ops.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Util/APIntHelper.h"
#include "llzk/Util/Compare.h"
 
#include <mlir/IR/BuiltinOps.h>
#include <mlir/Pass/AnalysisManager.h>
#include <mlir/Support/LLVM.h>
 
#include <llvm/ADT/MapVector.h>
#include <llvm/Support/SMTAPI.h>
 
#include <array>
#include <mutex>
 
namespace llzk {
 
/* ExpressionValue */

class ExpressionValue {
public:
  /* Must be default initializable to be a ScalarLatticeValue. */
  ExpressionValue() : i(), expr(nullptr) {}
 
  explicit ExpressionValue(const Field &f, llvm::SMTExprRef exprRef)
      : i(Interval::Entire(f)), expr(exprRef) {}
 
  ExpressionValue(const Field &f, llvm::SMTExprRef exprRef, llvm::APSInt singleVal)
      : i(Interval::Degenerate(f, singleVal)), expr(exprRef) {}
 
  ExpressionValue(llvm::SMTExprRef exprRef, Interval interval) : i(interval), expr(exprRef) {}
 
  llvm::SMTExprRef getExpr() const { return expr; }
 
  const Interval &getInterval() const { return i; }
 
  const Field &getField() const { return i.getField(); }

  ExpressionValue withInterval(const Interval &newInterval) const {
    return ExpressionValue(expr, newInterval);
  }
 
  /* Required to be a ScalarLatticeValue. */
  ExpressionValue &join(const ExpressionValue &rhs) {
    i = Interval::Entire(getField());
    return *this;
  }
 
  bool operator==(const ExpressionValue &rhs) const;

  friend ExpressionValue
  intersection(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);

  friend ExpressionValue
  join(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  // arithmetic ops
 
  friend ExpressionValue
  add(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  sub(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  mul(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  div(llvm::SMTSolverRef solver, felt::DivFeltOp op, const ExpressionValue &lhs,
      const ExpressionValue &rhs);
 
  friend ExpressionValue
  mod(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  bitAnd(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  shiftLeft(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  shiftRight(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  cmp(llvm::SMTSolverRef solver, boolean::CmpOp op, const ExpressionValue &lhs,
      const ExpressionValue &rhs);
 
  friend ExpressionValue
  boolAnd(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  boolOr(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);
 
  friend ExpressionValue
  boolXor(llvm::SMTSolverRef solver, const ExpressionValue &lhs, const ExpressionValue &rhs);

  friend ExpressionValue fallbackBinaryOp(
      llvm::SMTSolverRef solver, mlir::Operation *op, const ExpressionValue &lhs,
      const ExpressionValue &rhs
  );
 
  friend ExpressionValue neg(llvm::SMTSolverRef solver, const ExpressionValue &val);
 
  friend ExpressionValue notOp(llvm::SMTSolverRef solver, const ExpressionValue &val);
 
  friend ExpressionValue boolNot(llvm::SMTSolverRef solver, const ExpressionValue &val);
 
  friend ExpressionValue
  fallbackUnaryOp(llvm::SMTSolverRef solver, mlir::Operation *op, const ExpressionValue &val);
 
  /* Utility */
 
  void print(mlir::raw_ostream &os) const;
 
  friend mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const ExpressionValue &e) {
    e.print(os);
    return os;
  }
 
  struct Hash {
    unsigned operator()(const ExpressionValue &e) const {
      return Interval::Hash {}(e.i) ^ llvm::hash_value(e.expr);
    }
  };
 
private:
  Interval i;
  llvm::SMTExprRef expr;
};
 
/* IntervalAnalysisLatticeValue */
 
class IntervalAnalysisLatticeValue
    : public dataflow::AbstractLatticeValue<IntervalAnalysisLatticeValue, ExpressionValue> {
public:
  using AbstractLatticeValue::AbstractLatticeValue;
};
 
/* IntervalAnalysisLattice */
 
class IntervalDataFlowAnalysis;

class IntervalAnalysisLattice : public dataflow::AbstractDenseLattice {
public:
  using LatticeValue = IntervalAnalysisLatticeValue;
  // Map mlir::Values to LatticeValues
  using ValueMap = mlir::DenseMap<mlir::Value, LatticeValue>;
  // Map field references to LatticeValues. Used for field reads and writes.
  // Structure is component value -> field attribute -> latticeValue
  using FieldMap = mlir::DenseMap<mlir::Value, mlir::DenseMap<mlir::StringAttr, LatticeValue>>;
  // Expression to interval map for convenience.
  using ExpressionIntervals = mlir::DenseMap<llvm::SMTExprRef, Interval>;
  // Tracks all constraints and assignments in insertion order
  using ConstraintSet = llvm::SetVector<ExpressionValue>;
 
  using AbstractDenseLattice::AbstractDenseLattice;
 
  mlir::ChangeResult join(const AbstractDenseLattice &other) override;
 
  mlir::ChangeResult meet(const AbstractDenseLattice &rhs) override {
    llvm::report_fatal_error("IntervalDataFlowAnalysis::meet : unsupported");
    return mlir::ChangeResult::NoChange;
  }
 
  void print(mlir::raw_ostream &os) const override;
 
  mlir::FailureOr<LatticeValue> getValue(mlir::Value v) const;
  mlir::FailureOr<LatticeValue> getValue(mlir::Value v, mlir::StringAttr f) const;
 
  mlir::ChangeResult setValue(mlir::Value v, ExpressionValue e);
  mlir::ChangeResult setValue(mlir::Value v, mlir::StringAttr f, ExpressionValue e);
 
  mlir::ChangeResult addSolverConstraint(ExpressionValue e);
 
  friend mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const IntervalAnalysisLattice &l) {
    l.print(os);
    return os;
  }
 
  const ConstraintSet &getConstraints() const { return constraints; }
 
  mlir::FailureOr<Interval> findInterval(llvm::SMTExprRef expr) const;
  mlir::ChangeResult setInterval(llvm::SMTExprRef expr, Interval i);
 
  size_t size() const { return valMap.size(); }
 
  const ValueMap &getMap() const { return valMap; }
 
  ValueMap::iterator begin() { return valMap.begin(); }
  ValueMap::iterator end() { return valMap.end(); }
  ValueMap::const_iterator begin() const { return valMap.begin(); }
  ValueMap::const_iterator end() const { return valMap.end(); }
 
private:
  ValueMap valMap;
  FieldMap fieldMap;
  ConstraintSet constraints;
  ExpressionIntervals intervals;
};
 
/* IntervalDataFlowAnalysis */
 
class IntervalDataFlowAnalysis
    : public dataflow::DenseForwardDataFlowAnalysis<IntervalAnalysisLattice> {
  using Base = dataflow::DenseForwardDataFlowAnalysis<IntervalAnalysisLattice>;
  using Lattice = IntervalAnalysisLattice;
  using LatticeValue = IntervalAnalysisLattice::LatticeValue;
 
  // Map fields to their symbols
  using SymbolMap = mlir::DenseMap<ConstrainRef, llvm::SMTExprRef>;
 
public:
  explicit IntervalDataFlowAnalysis(
      mlir::DataFlowSolver &dataflowSolver, llvm::SMTSolverRef smt, const Field &f,
      bool propInputConstraints
  )
      : Base::DenseForwardDataFlowAnalysis(dataflowSolver), _dataflowSolver(dataflowSolver),
        smtSolver(smt), field(f), propagateInputConstraints(propInputConstraints) {}
 
  void visitCallControlFlowTransfer(
      mlir::CallOpInterface call, dataflow::CallControlFlowAction action, const Lattice &before,
      Lattice *after
  ) override;
 
  void visitOperation(mlir::Operation *op, const Lattice &before, Lattice *after) override;

  llvm::SMTExprRef getOrCreateSymbol(const ConstrainRef &r);
 
private:
  mlir::DataFlowSolver &_dataflowSolver;
  llvm::SMTSolverRef smtSolver;
  SymbolMap refSymbols;
  std::reference_wrapper<const Field> field;
  bool propagateInputConstraints;
  mlir::SymbolTableCollection tables;
 
  void setToEntryState(Lattice *lattice) override {
    // initial state should be empty, so do nothing here
  }
 
  llvm::SMTExprRef createFeltSymbol(const ConstrainRef &r) const;
 
  llvm::SMTExprRef createFeltSymbol(mlir::Value val) const;
 
  llvm::SMTExprRef createFeltSymbol(const char *name) const;
 
  bool isConstOp(mlir::Operation *op) const {
    return mlir::isa<
        felt::FeltConstantOp, mlir::arith::ConstantIndexOp, mlir::arith::ConstantIntOp>(op);
  }
 
  llvm::APSInt getConst(mlir::Operation *op) const;
 
  llvm::SMTExprRef createConstBitvectorExpr(llvm::APSInt v) const {
    return smtSolver->mkBitvector(v, field.get().bitWidth());
  }
 
  llvm::SMTExprRef createConstBoolExpr(bool v) const {
    return smtSolver->mkBitvector(mlir::APSInt((int)v), field.get().bitWidth());
  }
 
  bool isArithmeticOp(mlir::Operation *op) const {
    return mlir::isa<
        felt::AddFeltOp, felt::SubFeltOp, felt::MulFeltOp, felt::DivFeltOp, felt::ModFeltOp,
        felt::NegFeltOp, felt::InvFeltOp, felt::AndFeltOp, felt::OrFeltOp, felt::XorFeltOp,
        felt::NotFeltOp, felt::ShlFeltOp, felt::ShrFeltOp, boolean::CmpOp, boolean::AndBoolOp,
        boolean::OrBoolOp, boolean::XorBoolOp, boolean::NotBoolOp>(op);
  }
 
  ExpressionValue
  performBinaryArithmetic(mlir::Operation *op, const LatticeValue &a, const LatticeValue &b);
 
  ExpressionValue performUnaryArithmetic(mlir::Operation *op, const LatticeValue &a);

  mlir::ChangeResult
  applyInterval(mlir::Operation *originalOp, Lattice *after, mlir::Value val, Interval newInterval);

  mlir::FailureOr<std::pair<llvm::DenseSet<mlir::Value>, Interval>> getGeneralizedDecompInterval(
      const ConstrainRefLattice *constrainRefLattice, mlir::Value lhs, mlir::Value rhs
  );
 
  bool isBoolOp(mlir::Operation *op) const {
    return mlir::isa<boolean::AndBoolOp, boolean::OrBoolOp, boolean::XorBoolOp, boolean::NotBoolOp>(
        op
    );
  }
 
  bool isConversionOp(mlir::Operation *op) const {
    return mlir::isa<cast::IntToFeltOp, cast::FeltToIndexOp>(op);
  }
 
  bool isApplyMapOp(mlir::Operation *op) const { return mlir::isa<polymorphic::ApplyMapOp>(op); }
 
  bool isAssertOp(mlir::Operation *op) const { return mlir::isa<boolean::AssertOp>(op); }
 
  bool isReadOp(mlir::Operation *op) const {
    return mlir::isa<component::FieldReadOp, polymorphic::ConstReadOp, array::ReadArrayOp>(op);
  }
 
  bool isWriteOp(mlir::Operation *op) const {
    return mlir::isa<component::FieldWriteOp, array::WriteArrayOp, array::InsertArrayOp>(op);
  }
 
  bool isArrayLengthOp(mlir::Operation *op) const { return mlir::isa<array::ArrayLengthOp>(op); }
 
  bool isEmitOp(mlir::Operation *op) const {
    return mlir::isa<constrain::EmitEqualityOp, constrain::EmitContainmentOp>(op);
  }
 
  bool isCreateOp(mlir::Operation *op) const {
    return mlir::isa<component::CreateStructOp, array::CreateArrayOp>(op);
  }
 
  bool isExtractArrayOp(mlir::Operation *op) const { return mlir::isa<array::ExtractArrayOp>(op); }
 
  bool isDefinitionOp(mlir::Operation *op) const {
    return mlir::isa<
        component::StructDefOp, function::FuncDefOp, component::FieldDefOp, global::GlobalDefOp,
        mlir::ModuleOp>(op);
  }
 
  bool isCallOp(mlir::Operation *op) const { return mlir::isa<function::CallOp>(op); }
 
  bool isReturnOp(mlir::Operation *op) const { return mlir::isa<function::ReturnOp>(op); }
};
 
/* StructIntervals */

struct IntervalAnalysisContext {
  IntervalDataFlowAnalysis *intervalDFA;
  llvm::SMTSolverRef smtSolver;
  std::reference_wrapper<const Field> field;
  bool propagateInputConstraints;
 
  llvm::SMTExprRef getSymbol(const ConstrainRef &r) { return intervalDFA->getOrCreateSymbol(r); }
  const Field &getField() const { return field.get(); }
  bool doInputConstraintPropagation() const { return propagateInputConstraints; }
};
 
class StructIntervals {
public:
  static mlir::FailureOr<StructIntervals> compute(
      mlir::ModuleOp mod, component::StructDefOp s, mlir::DataFlowSolver &solver,
      IntervalAnalysisContext &ctx
  ) {
    StructIntervals si(mod, s);
    if (si.computeIntervals(solver, ctx).failed()) {
      return mlir::failure();
    }
    return si;
  }
 
  mlir::LogicalResult computeIntervals(mlir::DataFlowSolver &solver, IntervalAnalysisContext &ctx);
 
  void print(mlir::raw_ostream &os, bool withConstraints = false, bool printCompute = false) const;
 
  const llvm::MapVector<ConstrainRef, Interval> &getConstrainIntervals() const {
    return constrainFieldRanges;
  }
 
  const llvm::SetVector<ExpressionValue> getConstrainSolverConstraints() const {
    return constrainSolverConstraints;
  }
 
  const llvm::MapVector<ConstrainRef, Interval> &getComputeIntervals() const {
    return computeFieldRanges;
  }
 
  const llvm::SetVector<ExpressionValue> getComputeSolverConstraints() const {
    return computeSolverConstraints;
  }
 
  friend mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const StructIntervals &si) {
    si.print(os);
    return os;
  }
 
private:
  mlir::ModuleOp mod;
  component::StructDefOp structDef;
  llvm::SMTSolverRef smtSolver;
  // llvm::MapVector keeps insertion order for consistent iteration
  llvm::MapVector<ConstrainRef, Interval> constrainFieldRanges, computeFieldRanges;
  // llvm::SetVector for the same reasons as above
  llvm::SetVector<ExpressionValue> constrainSolverConstraints, computeSolverConstraints;
 
  StructIntervals(mlir::ModuleOp m, component::StructDefOp s) : mod(m), structDef(s) {}
};
 
/* StructIntervalAnalysis */
 
class ModuleIntervalAnalysis;
 
class StructIntervalAnalysis : public StructAnalysis<StructIntervals, IntervalAnalysisContext> {
public:
  using StructAnalysis::StructAnalysis;
  virtual ~StructIntervalAnalysis() = default;
 
  mlir::LogicalResult runAnalysis(
      mlir::DataFlowSolver &solver, mlir::AnalysisManager &_, IntervalAnalysisContext &ctx
  ) override {
    auto computeRes = StructIntervals::compute(getModule(), getStruct(), solver, ctx);
    if (mlir::failed(computeRes)) {
      return mlir::failure();
    }
    setResult(std::move(*computeRes));
    return mlir::success();
  }
};
 
/* ModuleIntervalAnalysis */
 
class ModuleIntervalAnalysis
    : public ModuleAnalysis<StructIntervals, IntervalAnalysisContext, StructIntervalAnalysis> {
 
public:
  ModuleIntervalAnalysis(mlir::Operation *op)
      : ModuleAnalysis(op), smtSolver(llvm::CreateZ3Solver()), field(std::nullopt) {}
  virtual ~ModuleIntervalAnalysis() = default;
 
  void setField(const Field &f) { field = f; }
  void setPropagateInputConstraints(bool prop) { propagateInputConstraints = prop; }
 
protected:
  void initializeSolver() override {
    ensure(field.has_value(), "field not set, could not generate analysis context");
    (void)solver.load<ConstrainRefAnalysis>();
    auto smtSolverRef = smtSolver;
    bool prop = propagateInputConstraints;
    intervalDFA = solver.load<IntervalDataFlowAnalysis, llvm::SMTSolverRef, const Field &, bool>(
        std::move(smtSolverRef), field.value(), std::move(prop)
    );
  }
 
  IntervalAnalysisContext getContext() override {
    ensure(field.has_value(), "field not set, could not generate analysis context");
    return {
        .intervalDFA = intervalDFA,
        .smtSolver = smtSolver,
        .field = field.value(),
        .propagateInputConstraints = propagateInputConstraints,
    };
  }
 
private:
  llvm::SMTSolverRef smtSolver;
  IntervalDataFlowAnalysis *intervalDFA;
  std::optional<std::reference_wrapper<const Field>> field;
  bool propagateInputConstraints;
};
 
} // namespace llzk
 
namespace llvm {
 
template <> struct DenseMapInfo<llzk::ExpressionValue> {
 
  static SMTExprRef getEmptyExpr() {
    static auto emptyPtr = reinterpret_cast<SMTExprRef>(1);
    return emptyPtr;
  }
  static SMTExprRef getTombstoneExpr() {
    static auto tombstonePtr = reinterpret_cast<SMTExprRef>(2);
    return tombstonePtr;
  }
 
  static llzk::ExpressionValue getEmptyKey() {
    return llzk::ExpressionValue(llzk::Field::getField("bn128"), getEmptyExpr());
  }
  static inline llzk::ExpressionValue getTombstoneKey() {
    return llzk::ExpressionValue(llzk::Field::getField("bn128"), getTombstoneExpr());
  }
  static unsigned getHashValue(const llzk::ExpressionValue &e) {
    return llzk::ExpressionValue::Hash {}(e);
  }
  static bool isEqual(const llzk::ExpressionValue &lhs, const llzk::ExpressionValue &rhs) {
    if (lhs.getExpr() == getEmptyExpr() || lhs.getExpr() == getTombstoneExpr() ||
        rhs.getExpr() == getEmptyExpr() || rhs.getExpr() == getTombstoneExpr()) {
      return lhs.getExpr() == rhs.getExpr();
    }
    return lhs == rhs;
  }
};
 
} // namespace llvm