ConstrainRefLattice.cpp

Go to the documentation of this file.

//===-- ConstrainRefLattice.cpp - ConstrainRef lattice & utils --*- 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
//
//===----------------------------------------------------------------------===//
 
#include "llzk/Analysis/ConstrainRefLattice.h"
#include "llzk/Analysis/ConstraintDependencyGraph.h"
#include "llzk/Analysis/DenseAnalysis.h"
#include "llzk/Dialect/Felt/IR/Ops.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Util/Hash.h"
#include "llzk/Util/SymbolHelper.h"
 
#include <mlir/Analysis/DataFlow/DeadCodeAnalysis.h>
#include <mlir/IR/Value.h>
 
#include <llvm/Support/Debug.h>
 
#include <numeric>
#include <unordered_set>
 
#define DEBUG_TYPE "llzk-constrain-ref-lattice"
 
using namespace mlir;
 
namespace llzk {
 
using namespace component;
using namespace felt;
using namespace polymorphic;
 
/* ConstrainRefLatticeValue */
 
mlir::ChangeResult ConstrainRefLatticeValue::insert(const ConstrainRef &rhs) {
  auto rhsVal = ConstrainRefLatticeValue(rhs);
  if (isScalar()) {
    return updateScalar(rhsVal.getScalarValue());
  } else {
    return foldAndUpdate(rhsVal);
  }
}
 
std::pair<ConstrainRefLatticeValue, mlir::ChangeResult>
ConstrainRefLatticeValue::translate(const TranslationMap &translation) const {
  auto newVal = *this;
  auto res = mlir::ChangeResult::NoChange;
  if (newVal.isScalar()) {
    res = newVal.translateScalar(translation);
  } else {
    for (auto &elem : newVal.getArrayValue()) {
      auto [newElem, elemRes] = elem->translate(translation);
      (*elem) = newElem;
      res |= elemRes;
    }
  }
  return {newVal, res};
}
 
std::pair<ConstrainRefLatticeValue, mlir::ChangeResult>
ConstrainRefLatticeValue::referenceField(SymbolLookupResult<FieldDefOp> fieldRef) const {
  ConstrainRefIndex idx(fieldRef);
  auto transform = [&idx](const ConstrainRef &r) -> ConstrainRef { return r.createChild(idx); };
  return elementwiseTransform(transform);
}
 
std::pair<ConstrainRefLatticeValue, mlir::ChangeResult>
ConstrainRefLatticeValue::extract(const std::vector<ConstrainRefIndex> &indices) const {
  if (isArray()) {
    ensure(indices.size() <= getNumArrayDims(), "invalid extract array operands");
 
    // First, compute what chunk(s) to index
    std::vector<size_t> currIdxs {0};
    for (unsigned i = 0; i < indices.size(); i++) {
      auto &idx = indices[i];
      auto currDim = getArrayDim(i);
 
      std::vector<size_t> newIdxs;
      ensure(idx.isIndex() || idx.isIndexRange(), "wrong type of index for array");
      if (idx.isIndex()) {
        auto idxVal = fromAPInt(idx.getIndex());
        std::transform(
            currIdxs.begin(), currIdxs.end(), std::back_inserter(newIdxs),
            [&currDim, &idxVal](size_t j) { return j * currDim + idxVal; }
        );
      } else {
        auto [low, high] = idx.getIndexRange();
        for (auto idxVal = fromAPInt(low); idxVal < fromAPInt(high); idxVal++) {
          std::transform(
              currIdxs.begin(), currIdxs.end(), std::back_inserter(newIdxs),
              [&currDim, &idxVal](size_t j) { return j * currDim + idxVal; }
          );
        }
      }
 
      currIdxs = newIdxs;
    }
    std::vector<int64_t> newArrayDims;
    size_t chunkSz = 1;
    for (unsigned i = indices.size(); i < getNumArrayDims(); i++) {
      auto dim = getArrayDim(i);
      newArrayDims.push_back(dim);
      chunkSz *= dim;
    }
    if (newArrayDims.empty()) {
      // read case, where the return value is a scalar (single element)
      ConstrainRefLatticeValue extractedVal;
      for (auto idx : currIdxs) {
        (void)extractedVal.update(getElemFlatIdx(idx));
      }
      return {extractedVal, mlir::ChangeResult::Change};
    } else {
      // extract case, where the return value is an array of fewer dimensions.
      ConstrainRefLatticeValue extractedVal(newArrayDims);
      for (auto chunkStart : currIdxs) {
        for (size_t i = 0; i < chunkSz; i++) {
          (void)extractedVal.getElemFlatIdx(i).update(getElemFlatIdx(chunkStart + i));
        }
      }
      return {extractedVal, mlir::ChangeResult::Change};
    }
  } else {
    auto currVal = *this;
    auto res = mlir::ChangeResult::NoChange;
    for (auto &idx : indices) {
      auto transform = [&idx](const ConstrainRef &r) -> ConstrainRef { return r.createChild(idx); };
      auto [newVal, transformRes] = currVal.elementwiseTransform(transform);
      currVal = std::move(newVal);
      res |= transformRes;
    }
    return {currVal, res};
  }
}
 
mlir::ChangeResult ConstrainRefLatticeValue::translateScalar(const TranslationMap &translation) {
  auto res = mlir::ChangeResult::NoChange;
  // copy the current value
  auto currVal = getScalarValue();
  // reset this value
  getValue() = ScalarTy();
  // For each current element, see if the translation map contains a valid prefix.
  // If so, translate the current element with all replacement prefixes indicated
  // by the translation value.
  for (const ConstrainRef &currRef : currVal) {
    for (auto &[prefix, replacementVal] : translation) {
      if (currRef.isValidPrefix(prefix)) {
        for (const ConstrainRef &replacementPrefix : replacementVal.foldToScalar()) {
          auto translatedRefRes = currRef.translate(prefix, replacementPrefix);
          if (succeeded(translatedRefRes)) {
            res |= insert(*translatedRefRes);
          }
        }
      }
    }
  }
  return res;
}
 
std::pair<ConstrainRefLatticeValue, mlir::ChangeResult>
ConstrainRefLatticeValue::elementwiseTransform(
    llvm::function_ref<ConstrainRef(const ConstrainRef &)> transform
) const {
  auto newVal = *this;
  auto res = mlir::ChangeResult::NoChange;
  if (newVal.isScalar()) {
    ScalarTy indexed;
    for (auto &ref : newVal.getScalarValue()) {
      auto [_, inserted] = indexed.insert(transform(ref));
      if (inserted) {
        res |= mlir::ChangeResult::Change;
      }
    }
    newVal.getScalarValue() = indexed;
  } else {
    for (auto &elem : newVal.getArrayValue()) {
      auto [newElem, elemRes] = elem->elementwiseTransform(transform);
      (*elem) = newElem;
      res |= elemRes;
    }
  }
  return {newVal, res};
}
 
mlir::raw_ostream &operator<<(mlir::raw_ostream &os, const ConstrainRefLatticeValue &v) {
  v.print(os);
  return os;
}
 
/* ConstrainRefLattice */
 
mlir::FailureOr<ConstrainRef> ConstrainRefLattice::getSourceRef(mlir::Value val) {
  if (auto blockArg = mlir::dyn_cast<mlir::BlockArgument>(val)) {
    return ConstrainRef(blockArg);
  } else if (auto defOp = val.getDefiningOp()) {
    if (auto feltConst = mlir::dyn_cast<FeltConstantOp>(defOp)) {
      return ConstrainRef(feltConst);
    } else if (auto constIdx = mlir::dyn_cast<mlir::arith::ConstantIndexOp>(defOp)) {
      return ConstrainRef(constIdx);
    } else if (auto readConst = mlir::dyn_cast<ConstReadOp>(defOp)) {
      return ConstrainRef(readConst);
    } else if (auto structNew = mlir::dyn_cast<CreateStructOp>(defOp)) {
      return ConstrainRef(structNew);
    }
  }
  return mlir::failure();
}
 
void ConstrainRefLattice::print(mlir::raw_ostream &os) const {
  os << "ConstrainRefLattice { ";
  for (auto mit = valMap.begin(); mit != valMap.end();) {
    auto &[val, latticeVal] = *mit;
    os << "\n    (";
    if (val.is<Value>()) {
      os << val.get<Value>();
    } else if (val.is<Operation *>()) {
      os << *val.get<Operation *>();
    } else {
      llvm_unreachable("unhandled ValueTy print case");
    }
    os << ") => " << latticeVal;
    mit++;
    if (mit != valMap.end()) {
      os << ',';
    } else {
      os << '\n';
    }
  }
  os << "}\n";
}
 
mlir::ChangeResult ConstrainRefLattice::setValues(const ValueMap &rhs) {
  auto res = mlir::ChangeResult::NoChange;
 
  for (auto &[v, s] : rhs) {
    res |= setValue(v, s);
  }
  return res;
}
 
mlir::ChangeResult ConstrainRefLattice::setValue(ValueTy v, const ConstrainRefLatticeValue &rhs) {
  for (const ConstrainRef &ref : rhs.foldToScalar()) {
    refMap[ref].insert(v);
  }
  return valMap[v].setValue(rhs);
}
 
mlir::ChangeResult ConstrainRefLattice::setValue(ValueTy v, const ConstrainRef &ref) {
  refMap[ref].insert(v);
  return valMap[v].setValue(ConstrainRefLatticeValue(ref));
}
 
ConstrainRefLatticeValue ConstrainRefLattice::getOrDefault(ConstrainRefLattice::ValueTy v) const {
  auto it = valMap.find(v);
  if (it != valMap.end()) {
    return it->second;
  }
 
  if (v.is<Value>()) {
    auto sourceRef = getSourceRef(v.get<Value>());
    if (mlir::succeeded(sourceRef)) {
      return ConstrainRefLatticeValue(sourceRef.value());
    }
  }
  return ConstrainRefLatticeValue();
}
 
ConstrainRefLatticeValue ConstrainRefLattice::getReturnValue(unsigned i) const {
  auto op = this->getPoint().get<mlir::Operation *>();
  if (auto retOp = mlir::dyn_cast<function::ReturnOp>(op)) {
    if (i >= retOp.getNumOperands()) {
      llvm::report_fatal_error("return value requested is out of range");
    }
    return this->getOrDefault(retOp.getOperand(i));
  }
  return ConstrainRefLatticeValue();
}
 
ConstrainRefLattice::ValueSet ConstrainRefLattice::lookupValues(const ConstrainRef &ref) const {
  if (auto it = refMap.find(ref); it != refMap.end()) {
    return it->second;
  }
  return ConstrainRefLattice::ValueSet();
}
 
llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const ConstrainRefLattice &lattice) {
  lattice.print(os);
  return os;
}
 
} // namespace llzk
 
namespace llvm {
 
raw_ostream &operator<<(raw_ostream &os, llvm::PointerUnion<mlir::Value, mlir::Operation *> ptr) {
  if (ptr.is<Value>()) {
    os << ptr.get<Value>();
  } else {
    Operation *op = ptr.get<Operation *>();
    if (op) {
      os << *op;
    } else {
      os << "<null operation>";
    }
  }
  return os;
}
} // namespace llvm