ConstraintScalarFacts.h

Go to the documentation of this file.

//===-- ConstraintScalarFacts.h ---------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
 
#pragma once
 
#include "ScalarEvaluation.h"
 
#include <mlir/IR/Matchers.h>
#include <mlir/Interfaces/SideEffectInterfaces.h>
 
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/DenseSet.h>
 
namespace llzk::polymorphic::detail {
class ConstraintScalarFacts {
  llvm::DenseMap<mlir::Value, mlir::Attribute> constants;

  void fold(mlir::Operation *producer) {
    llvm::SmallVector<mlir::Attribute> operands;
    for (mlir::Value operand : producer->getOperands()) {
      operands.push_back(constants.lookup(operand));
    }
    llvm::SmallVector<mlir::OpFoldResult> results;
    if (mlir::succeeded(evaluateScalarOperation(producer, operands, results))) {
      for (auto [result, folded] : llvm::zip_equal(producer->getResults(), results)) {
        constants[result] = llvm::isa<mlir::Attribute>(folded)
                                ? llvm::cast<mlir::Attribute>(folded)
                                : constants.lookup(llvm::cast<mlir::Value>(folded));
      }
    }
  }
 
public:
  ConstraintScalarFacts() = default;
  explicit ConstraintScalarFacts(mlir::MLIRContext * /*context*/) {}

  mlir::Attribute known(mlir::Value value) {
    llvm::SmallVector<std::pair<mlir::Value, bool>> pending {{value, false}};
    llvm::DenseSet<mlir::Value> examined;
    while (!pending.empty()) {
      auto [current, ready] = pending.pop_back_val();
      if (constants.contains(current)) {
        continue;
      }
      mlir::Attribute literal;
      if (mlir::matchPattern(current, mlir::m_Constant(&literal))) {
        constants[current] = literal;
        continue;
      }
      mlir::Operation *producer = current.getDefiningOp();
      // Assert casts may have effects, but their fold hooks succeed only when the
      // range check is satisfied. Other effectful operations remain unknown.
      if (!producer || producer->getNumRegions() ||
          (!mlir::isMemoryEffectFree(producer) &&
           !llvm::isa<cast::FeltToIndexOp, cast::IntToFeltOp>(producer))) {
        constants.try_emplace(current, mlir::Attribute());
        continue;
      }
      if (!ready) {
        if (!examined.insert(current).second) {
          continue;
        }
        pending.push_back({current, true});
        for (mlir::Value operand : producer->getOperands()) {
          pending.push_back({operand, false});
        }
        continue;
      }
      fold(producer);
      constants.try_emplace(current, mlir::Attribute());
    }
    return constants.lookup(value);
  }

  std::optional<int64_t> integer(mlir::Value value) {
    auto attr = llvm::dyn_cast_or_null<mlir::IntegerAttr>(known(value));
    if (!attr || !attr.getValue().isSignedIntN(64)) {
      return std::nullopt;
    }
    return attr.getInt();
  }
};

inline bool isInvariantTrueCondition(mlir::Value condition) {
  ConstraintScalarFacts facts(condition.getContext());
  auto value = facts.integer(condition);
  return value && *value != 0;
}
 
} // namespace llzk::polymorphic::detail