ConstraintDependencyGraph.cpp

Go to the documentation of this file.

//===-- ConstraintDependencyGraph.cpp ---------------------------*- 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/Array/IR/Ops.h"
#include "llzk/Dialect/Constrain/IR/Ops.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Util/Hash.h"
#include "llzk/Util/SymbolHelper.h"
#include "llzk/Util/TypeHelper.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-cdg"
 
using namespace mlir;
 
namespace llzk {
 
using namespace array;
using namespace component;
using namespace constrain;
using namespace function;
 
/* ConstrainRefAnalysis */
 
void ConstrainRefAnalysis::visitCallControlFlowTransfer(
    mlir::CallOpInterface call, dataflow::CallControlFlowAction action,
    const ConstrainRefLattice &before, ConstrainRefLattice *after
) {
  LLVM_DEBUG(
      llvm::dbgs() << "ConstrainRefAnalysis::visitCallControlFlowTransfer: " << call << '\n'
  );
  auto fnOpRes = resolveCallable<FuncDefOp>(tables, call);
  ensure(succeeded(fnOpRes), "could not resolve called function");
 
  LLVM_DEBUG({
    llvm::dbgs().indent(4) << "parent op is ";
    if (auto s = call->getParentOfType<StructDefOp>()) {
      llvm::dbgs() << s.getName();
    } else if (auto p = call->getParentOfType<FuncDefOp>()) {
      llvm::dbgs() << p.getName();
    } else {
      llvm::dbgs() << "<UNKNOWN PARENT TYPE>";
    }
    llvm::dbgs() << '\n';
  });

  if (action == dataflow::CallControlFlowAction::EnterCallee) {
    // We skip updating the incoming lattice for function calls,
    // as ConstrainRefs are relative to the containing function/struct, so we don't need to pollute
    // the callee with the callers values.
    // This also avoids a non-convergence scenario, as calling a
    // function from other contexts can cause the lattice values to oscillate and constantly
    // change (thus looping infinitely).
 
    setToEntryState(after);
  }
  else if (action == dataflow::CallControlFlowAction::ExitCallee) {
    // Get the argument values of the lattice by getting the state as it would
    // have been for the callsite.
    const ConstrainRefLattice *beforeCall = nullptr;
    if (auto *prev = call->getPrevNode()) {
      beforeCall = static_cast<const ConstrainRefLattice *>(getLattice(prev));
    } else {
      beforeCall = static_cast<const ConstrainRefLattice *>(getLattice(call->getBlock()));
    }
    ensure(beforeCall, "could not get prior lattice");
 
    // Translate argument values based on the operands given at the call site.
    std::unordered_map<ConstrainRef, ConstrainRefLatticeValue, ConstrainRef::Hash> translation;
    auto funcOpRes = resolveCallable<FuncDefOp>(tables, call);
    ensure(mlir::succeeded(funcOpRes), "could not lookup called function");
    auto funcOp = funcOpRes->get();
 
    auto callOp = mlir::dyn_cast<CallOp>(call.getOperation());
    ensure(callOp, "call is not a llzk::CallOp");
 
    for (unsigned i = 0; i < funcOp.getNumArguments(); i++) {
      auto key = ConstrainRef(funcOp.getArgument(i));
      auto val = beforeCall->getOrDefault(callOp.getOperand(i));
      translation[key] = val;
    }
 
    // The lattice at the return is:
    // - the lattice before the call, plus
    // - the translated, return values, plus
    // - any translated internal values (so we can see where values are used)
    mlir::ChangeResult updated = after->join(*beforeCall);
    for (unsigned i = 0; i < callOp.getNumResults(); i++) {
      auto retVal = before.getReturnValue(i);
      auto [translatedVal, _] = retVal.translate(translation);
      updated |= after->setValue(callOp->getResult(i), translatedVal);
    }
    for (const auto &[val, refVal] : before.getMap()) {
      auto [translatedVal, _] = refVal.translate(translation);
      updated |= after->setValue(val, translatedVal);
    }
    propagateIfChanged(after, updated);
  }
  else if (action == mlir::dataflow::CallControlFlowAction::ExternalCallee) {
    // For external calls, we propagate what information we already have from
    // before the call to after the call, since the external call won't invalidate
    // any of that information. It also, conservatively, makes no assumptions about
    // external calls and their computation, so CDG edges will not be computed over
    // input arguments to external functions.
    join(after, before);
  }
}
 
void ConstrainRefAnalysis::visitOperation(
    mlir::Operation *op, const ConstrainRefLattice &before, ConstrainRefLattice *after
) {
  LLVM_DEBUG(llvm::dbgs() << "ConstrainRefAnalysis::visitOperation: " << *op << '\n');
  // Collect the references that are made by the operands to `op`.
  ConstrainRefLattice::ValueMap operandVals;
  for (mlir::OpOperand &operand : op->getOpOperands()) {
    operandVals[operand.get()] = before.getOrDefault(operand.get());
  }
 
  // Propagate existing state.
  join(after, before);
 
  // Add operand values, if not already added. Ensures that the default value
  // of a ConstrainRef (the source of the ref) is visible in the lattice.
  ChangeResult res = after->setValues(operandVals);
 
  // We will now join the the operand refs based on the type of operand.
  if (auto fieldRefOp = mlir::dyn_cast<FieldRefOpInterface>(op)) {
    // The operand is indexed into by the FieldDefOp.
    auto fieldOpRes = fieldRefOp.getFieldDefOp(tables);
    ensure(mlir::succeeded(fieldOpRes), "could not find field read");
 
    ConstrainRefLattice::ValueTy fieldRefRes;
    if (fieldRefOp.isRead()) {
      fieldRefRes = fieldRefOp.getVal();
    } else {
      fieldRefRes = fieldRefOp;
    }
 
    const auto &ops = operandVals.at(fieldRefOp.getComponent());
    auto [fieldVals, _] = ops.referenceField(fieldOpRes.value());
 
    res |= after->setValue(fieldRefRes, fieldVals);
  } else if (auto arrayAccessOp = mlir::dyn_cast<ArrayAccessOpInterface>(op)) {
    // Covers read/write/extract/insert array ops
    arraySubdivisionOpUpdate(arrayAccessOp, operandVals, before, after);
  } else if (auto createArray = mlir::dyn_cast<CreateArrayOp>(op)) {
    // Create an array using the operand values, if they exist.
    // Currently, the new array must either be fully initialized or uninitialized.
 
    auto newArrayVal = ConstrainRefLatticeValue(createArray.getType().getShape());
    // If the array is statically initialized, iterate through all operands and initialize the array
    // value.
    const auto &elements = createArray.getElements();
    if (!elements.empty()) {
      for (unsigned i = 0; i < elements.size(); i++) {
        auto currentOp = elements[i];
        auto &opVals = operandVals[currentOp];
        (void)newArrayVal.getElemFlatIdx(i).setValue(opVals);
      }
    }
 
    auto createArrayRes = createArray.getResult();
 
    res |= after->setValue(createArrayRes, newArrayVal);
  } else if (auto structNewOp = mlir::dyn_cast<CreateStructOp>(op)) {
    auto newOpRes = structNewOp.getResult();
    auto newStructValue = before.getOrDefault(newOpRes);
    res |= after->setValue(newOpRes, newStructValue);
  } else {
    // Standard union of operands into the results value.
    // TODO: Could perform constant computation/propagation here for, e.g., arithmetic
    // over constants, but such analysis may be better suited for a dedicated pass.
    res |= fallbackOpUpdate(op, operandVals, before, after);
  }
 
  propagateIfChanged(after, res);
}
 
// Perform a standard union of operands into the results value.
mlir::ChangeResult ConstrainRefAnalysis::fallbackOpUpdate(
    mlir::Operation *op, const ConstrainRefLattice::ValueMap &operandVals,
    const ConstrainRefLattice &before, ConstrainRefLattice *after
) {
  auto updated = mlir::ChangeResult::NoChange;
  for (auto res : op->getResults()) {
    auto cur = before.getOrDefault(res);
 
    for (auto &[_, opVal] : operandVals) {
      (void)cur.update(opVal);
    }
    updated |= after->setValue(res, cur);
  }
  return updated;
}
 
// Perform the update for either a readarr op or an extractarr op, which
// operate very similarly: index into the first operand using a variable number
// of provided indices.
void ConstrainRefAnalysis::arraySubdivisionOpUpdate(
    ArrayAccessOpInterface arrayAccessOp, const ConstrainRefLattice::ValueMap &operandVals,
    const ConstrainRefLattice &before, ConstrainRefLattice *after
) {
  // We index the first operand by all remaining indices.
  ConstrainRefLattice::ValueTy res;
  if (mlir::isa<ReadArrayOp, ExtractArrayOp>(arrayAccessOp)) {
    res = arrayAccessOp->getResult(0);
  } else {
    res = arrayAccessOp;
  }
 
  auto array = arrayAccessOp.getArrRef();
  auto it = operandVals.find(array);
  ensure(it != operandVals.end(), "improperly constructed operandVals map");
  auto currVals = it->second;
 
  std::vector<ConstrainRefIndex> indices;
 
  for (unsigned i = 0; i < arrayAccessOp.getIndices().size(); ++i) {
    auto idxOperand = arrayAccessOp.getIndices()[i];
    auto idxIt = operandVals.find(idxOperand);
    ensure(idxIt != operandVals.end(), "improperly constructed operandVals map");
    auto &idxVals = idxIt->second;
 
    // Note: we allow constant values regardless of if they are felt or index,
    // as if they were felt, there would need to be a cast to index, and if it
    // was missing, there would be a semantic check failure. So we accept either
    // so we don't have to track the cast ourselves.
    if (idxVals.isSingleValue() && idxVals.getSingleValue().isConstant()) {
      ConstrainRefIndex idx(idxVals.getSingleValue().getConstantValue());
      indices.push_back(idx);
    } else {
      // Otherwise, assume any range is valid.
      auto arrayType = mlir::dyn_cast<ArrayType>(array.getType());
      auto lower = mlir::APInt::getZero(64);
      mlir::APInt upper(64, arrayType.getDimSize(i));
      auto idxRange = ConstrainRefIndex(lower, upper);
      indices.push_back(idxRange);
    }
  }
 
  auto [newVals, _] = currVals.extract(indices);
 
  if (mlir::isa<ReadArrayOp, WriteArrayOp>(arrayAccessOp)) {
    ensure(newVals.isScalar(), "array read/write must produce a scalar value");
  }
  // an extract operation may yield a "scalar" value if not all dimensions of
  // the source array are instantiated; for example, if extracting an array from
  // an input arg, the current value is a "scalar" with an array type, and extracting
  // from that yields another single value with indices. For example: extracting [0][1]
  // from { arg1 } yields { arg1[0][1] }.
 
  propagateIfChanged(after, after->setValue(res, newVals));
}
 
/* ConstraintDependencyGraph */
 
mlir::FailureOr<ConstraintDependencyGraph> ConstraintDependencyGraph::compute(
    mlir::ModuleOp m, StructDefOp s, mlir::DataFlowSolver &solver, mlir::AnalysisManager &am,
    bool runIntraprocedural
) {
  ConstraintDependencyGraph cdg(m, s, runIntraprocedural);
  if (cdg.computeConstraints(solver, am).failed()) {
    return mlir::failure();
  }
  return cdg;
}
 
void ConstraintDependencyGraph::dump() const { print(llvm::errs()); }

void ConstraintDependencyGraph::print(llvm::raw_ostream &os) const {
  // the EquivalenceClasses::iterator is sorted, but the EquivalenceClasses::member_iterator is
  // not guaranteed to be sorted. So, we will sort members before printing them.
  // We also want to add the constant values into the printing.
  std::set<std::set<ConstrainRef>> sortedSets;
  for (auto it = signalSets.begin(); it != signalSets.end(); it++) {
    if (!it->isLeader()) {
      continue;
    }
 
    std::set<ConstrainRef> sortedMembers;
    for (auto mit = signalSets.member_begin(it); mit != signalSets.member_end(); mit++) {
      sortedMembers.insert(*mit);
    }
 
    // We only want to print sets with a size > 1, because size == 1 means the
    // signal is not in a constraint.
    if (sortedMembers.size() > 1) {
      sortedSets.insert(sortedMembers);
    }
  }
  // Add the constants in separately.
  for (auto &[ref, constSet] : constantSets) {
    if (constSet.empty()) {
      continue;
    }
    std::set<ConstrainRef> sortedMembers(constSet.begin(), constSet.end());
    sortedMembers.insert(ref);
    sortedSets.insert(sortedMembers);
  }
 
  os << "ConstraintDependencyGraph { ";
 
  for (auto it = sortedSets.begin(); it != sortedSets.end();) {
    os << "\n    { ";
    for (auto mit = it->begin(); mit != it->end();) {
      os << *mit;
      mit++;
      if (mit != it->end()) {
        os << ", ";
      }
    }
 
    it++;
    if (it == sortedSets.end()) {
      os << " }\n";
    } else {
      os << " },";
    }
  }
 
  os << "}\n";
}
 
mlir::LogicalResult ConstraintDependencyGraph::computeConstraints(
    mlir::DataFlowSolver &solver, mlir::AnalysisManager &am
) {
  // Fetch the constrain function. This is a required feature for all LLZK structs.
  FuncDefOp constrainFnOp = structDef.getConstrainFuncOp();
  ensure(
      constrainFnOp,
      "malformed struct " + mlir::Twine(structDef.getName()) + " must define a constrain function"
  );

 
  // - Union all constraints from the analysis
  // This requires iterating over all of the emit operations
  constrainFnOp.walk([this, &solver](Operation *op) {
    const auto *refLattice = solver.lookupState<ConstrainRefLattice>(op);
    // aggregate the ref2Val map across operations, as some may have nested
    // regions and blocks that aren't propagated to the function terminator
    if (refLattice) {
      for (auto &[ref, vals] : refLattice->getRef2Val()) {
        ref2Val[ref].insert(vals.begin(), vals.end());
      }
    }
    if (isa<EmitEqualityOp, EmitContainmentOp>(op)) {
      this->walkConstrainOp(solver, op);
    }
  });

  auto fnCallWalker = [this, &solver, &am](CallOp fnCall) mutable {
    auto res = resolveCallable<FuncDefOp>(tables, fnCall);
    ensure(mlir::succeeded(res), "could not resolve constrain call");
 
    auto fn = res->get();
    if (!fn.isStructConstrain()) {
      return;
    }
    // Nested
    auto calledStruct = fn.getOperation()->getParentOfType<StructDefOp>();
    ConstrainRefRemappings translations;
 
    auto lattice = solver.lookupState<ConstrainRefLattice>(fnCall.getOperation());
    ensure(lattice, "could not find lattice for call operation");
 
    // Map fn parameters to args in the call op
    for (unsigned i = 0; i < fn.getNumArguments(); i++) {
      auto prefix = ConstrainRef(fn.getArgument(i));
      auto val = lattice->getOrDefault(fnCall.getOperand(i));
      translations.push_back({prefix, val});
    }
    auto &childAnalysis =
        am.getChildAnalysis<ConstraintDependencyGraphStructAnalysis>(calledStruct);
    if (!childAnalysis.constructed()) {
      ensure(
          mlir::succeeded(childAnalysis.runAnalysis(solver, am, /* runIntraprocedural */ false)),
          "could not construct CDG for child struct"
      );
    }
    auto translatedCDG = childAnalysis.getResult().translate(translations);
 
    // Now, union sets based on the translation
    // We should be able to just merge what is in the translatedCDG to the current CDG
    auto &tSets = translatedCDG.signalSets;
    for (auto lit = tSets.begin(); lit != tSets.end(); lit++) {
      if (!lit->isLeader()) {
        continue;
      }
      auto leader = lit->getData();
      for (auto mit = tSets.member_begin(lit); mit != tSets.member_end(); mit++) {
        signalSets.unionSets(leader, *mit);
      }
    }
    // And update the constant sets
    for (auto &[ref, constSet] : translatedCDG.constantSets) {
      constantSets[ref].insert(constSet.begin(), constSet.end());
    }
  };
  if (!runIntraprocedural) {
    constrainFnOp.walk(fnCallWalker);
  }
 
  return mlir::success();
}
 
void ConstraintDependencyGraph::walkConstrainOp(
    mlir::DataFlowSolver &solver, mlir::Operation *emitOp
) {
  std::vector<ConstrainRef> signalUsages, constUsages;
 
  const ConstrainRefLattice *refLattice = solver.lookupState<ConstrainRefLattice>(emitOp);
  ensure(refLattice, "missing lattice for constrain op");
 
  for (auto operand : emitOp->getOperands()) {
    auto latticeVal = refLattice->getOrDefault(operand);
    for (auto &ref : latticeVal.foldToScalar()) {
      if (ref.isConstant()) {
        constUsages.push_back(ref);
      } else {
        signalUsages.push_back(ref);
      }
    }
  }
 
  // Compute a transitive closure over the signals.
  if (!signalUsages.empty()) {
    auto it = signalUsages.begin();
    auto leader = signalSets.getOrInsertLeaderValue(*it);
    for (it++; it != signalUsages.end(); it++) {
      signalSets.unionSets(leader, *it);
    }
  }
  // Also update constant references for each value.
  for (auto &sig : signalUsages) {
    constantSets[sig].insert(constUsages.begin(), constUsages.end());
  }
}
 
ConstraintDependencyGraph ConstraintDependencyGraph::translate(ConstrainRefRemappings translation
) const {
  ConstraintDependencyGraph res(mod, structDef);
  auto translate = [&translation](const ConstrainRef &elem
                   ) -> mlir::FailureOr<std::vector<ConstrainRef>> {
    std::vector<ConstrainRef> refs;
    for (auto &[prefix, vals] : translation) {
      if (!elem.isValidPrefix(prefix)) {
        continue;
      }
 
      if (vals.isArray()) {
        // Try to index into the array
        auto suffix = elem.getSuffix(prefix);
        ensure(
            mlir::succeeded(suffix), "failure is nonsensical, we already checked for valid prefix"
        );
 
        auto [resolvedVals, _] = vals.extract(suffix.value());
        auto folded = resolvedVals.foldToScalar();
        refs.insert(refs.end(), folded.begin(), folded.end());
      } else {
        for (auto &replacement : vals.getScalarValue()) {
          auto translated = elem.translate(prefix, replacement);
          if (mlir::succeeded(translated)) {
            refs.push_back(translated.value());
          }
        }
      }
    }
    if (refs.empty()) {
      return mlir::failure();
    }
    return refs;
  };
 
  for (auto leaderIt = signalSets.begin(); leaderIt != signalSets.end(); leaderIt++) {
    if (!leaderIt->isLeader()) {
      continue;
    }
    // translate everything in this set first
    std::vector<ConstrainRef> translatedSignals, translatedConsts;
    for (auto mit = signalSets.member_begin(leaderIt); mit != signalSets.member_end(); mit++) {
      auto member = translate(*mit);
      if (mlir::failed(member)) {
        continue;
      }
      for (auto &ref : *member) {
        if (ref.isConstant()) {
          translatedConsts.push_back(ref);
        } else {
          translatedSignals.push_back(ref);
        }
      }
      // Also add the constants from the original CDG
      if (auto it = constantSets.find(*mit); it != constantSets.end()) {
        auto &origConstSet = it->second;
        translatedConsts.insert(translatedConsts.end(), origConstSet.begin(), origConstSet.end());
      }
    }
 
    if (translatedSignals.empty()) {
      continue;
    }
 
    // Now we can insert the translated signals
    auto it = translatedSignals.begin();
    auto leader = *it;
    res.signalSets.insert(leader);
    for (it++; it != translatedSignals.end(); it++) {
      res.signalSets.insert(*it);
      res.signalSets.unionSets(leader, *it);
    }
 
    // And update the constant references
    for (auto &ref : translatedSignals) {
      res.constantSets[ref].insert(translatedConsts.begin(), translatedConsts.end());
    }
  }
 
  // Translate ref2Val as well
  for (auto &[ref, vals] : ref2Val) {
    auto translationRes = translate(ref);
    if (succeeded(translationRes)) {
      for (const auto &translatedRef : *translationRes) {
        res.ref2Val[translatedRef].insert(vals.begin(), vals.end());
      }
    }
  }
 
  return res;
}
 
ConstrainRefSet ConstraintDependencyGraph::getConstrainingValues(const ConstrainRef &ref) const {
  ConstrainRefSet res;
  auto currRef = mlir::FailureOr<ConstrainRef>(ref);
  while (mlir::succeeded(currRef)) {
    // Add signals
    for (auto it = signalSets.findLeader(*currRef); it != signalSets.member_end(); it++) {
      if (currRef.value() != *it) {
        res.insert(*it);
      }
    }
    // Add constants
    auto constIt = constantSets.find(*currRef);
    if (constIt != constantSets.end()) {
      res.insert(constIt->second.begin(), constIt->second.end());
    }
    // Go to parent
    currRef = currRef->getParentPrefix();
  }
  return res;
}
 
/* ConstraintDependencyGraphStructAnalysis */
 
mlir::LogicalResult ConstraintDependencyGraphStructAnalysis::runAnalysis(
    mlir::DataFlowSolver &solver, mlir::AnalysisManager &moduleAnalysisManager,
    bool runIntraprocedural
) {
  auto result = ConstraintDependencyGraph::compute(
      getModule(), getStruct(), solver, moduleAnalysisManager, runIntraprocedural
  );
  if (mlir::failed(result)) {
    return mlir::failure();
  }
  setResult(std::move(*result));
  return mlir::success();
}
 
} // namespace llzk