WitgenLowering.cpp

Go to the documentation of this file.

//===-- LLZKWitgenLoweringPass.cpp -----------------------------*- 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
//
//===----------------------------------------------------------------------===//
 
#include "WitgenLowering.h"
 
#include "WitgenDriver.h"
#include "WitgenUtils.h"
#include "WitnessSelection.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/LLZK/IR/Ops.h"
#include "llzk/Dialect/POD/IR/Attrs.h"
#include "llzk/Dialect/POD/IR/Ops.h"
#include "llzk/Dialect/Polymorphic/Transforms/TransformationPasses.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Transforms/LLZKTransformationPasses.h"
#include "llzk/Util/Compare.h"
#include "llzk/Util/Constants.h"
#include "llzk/Util/DynamicAPIntHelper.h"
#include "llzk/Util/Field.h"
#include "llzk/Util/SymbolHelper.h"
#include "llzk/Util/Walk.h"
 
#include <mlir/Conversion/AffineToStandard/AffineToStandard.h>
#include <mlir/Dialect/Arith/IR/Arith.h>
#include <mlir/Dialect/ControlFlow/IR/ControlFlowOps.h>
#include <mlir/Dialect/Func/IR/FuncOps.h>
#include <mlir/Dialect/LLVMIR/LLVMDialect.h>
#include <mlir/Dialect/MemRef/IR/MemRef.h>
#include <mlir/Dialect/SCF/IR/SCF.h>
#include <mlir/Dialect/Utils/IndexingUtils.h>
#include <mlir/IR/Builders.h>
#include <mlir/IR/BuiltinAttributes.h>
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/SymbolTable.h>
#include <mlir/Pass/PassManager.h>
#include <mlir/Transforms/Passes.h>
 
#include <llvm/ADT/APInt.h>
#include <llvm/ADT/STLExtras.h>
#include <llvm/ADT/SmallString.h>
#include <llvm/ADT/StringMap.h>
#include <llvm/ADT/TypeSwitch.h>
#include <llvm/Support/MathExtras.h>
 
#include <limits>
 
using namespace mlir;
 
namespace llzk::witgen {
namespace {

struct LoweredValue {
  Type sourceType;
  llvm::SmallVector<Value> leaves;
};

static FailureOr<std::reference_wrapper<const Field>> getModuleField(ModuleOp moduleOp) {
  FieldSet fields;
  if (failed(collectFields(moduleOp.getOperation(), fields, false))) {
    moduleOp.emitError("failed to collect fields for llzk-witgen lowering");
    return failure();
  }
  if (fields.size() != 1) {
    moduleOp.emitError("llzk-witgen execution-engine lowering requires exactly one field");
    return failure();
  }
  return *fields.begin();
}

static std::string mangleFunctionName(function::FuncDefOp funcOp) {
  auto symbolRef = funcOp.getFullyQualifiedName(false);
  llvm::SmallString<128> result("__llzk_witgen_");
  for (StringRef piece : getNames(symbolRef)) {
    if (!result.empty() && result.back() != '_') {
      result += "__";
    }
    for (char c : piece) {
      result += llvm::isAlnum(c) ? c : '_';
    }
  }
  return std::string(result);
}

static Value makeIndexConstant(OpBuilder &builder, Location loc, int64_t value) {
  return builder.create<arith::ConstantIndexOp>(loc, value).getResult();
}

static Value makeOneFelt(OpBuilder &builder, Location loc, const Field &field) {
  return builder.create<arith::ConstantOp>(
      loc, IntegerAttr::get(IntegerType::get(builder.getContext(), field.bitWidth()), 1)
  );
}

static FailureOr<Type> lowerScalarType(MLIRContext *context, Type type, const Field &field) {
  if (isa<felt::FeltType>(type)) {
    return IntegerType::get(context, field.bitWidth());
  }
  if (isa<IndexType>(type)) {
    return type;
  }
  if (auto intType = dyn_cast<IntegerType>(type)) {
    if (intType.getWidth() == 1) {
      return intType;
    }
  }
  return failure();
}

static bool isScalarType(Type type) {
  return isa<felt::FeltType, IndexType>(type) ||
         (isa<IntegerType>(type) && llvm::cast<IntegerType>(type).getWidth() == 1);
}

static LogicalResult flattenTypeLeaves(
    Type type, SymbolTableCollection &tables, Operation *origin, const Field &field,
    SmallVectorImpl<Type> &out, llvm::ArrayRef<int64_t> prefixShape = {}, bool storage = false
) {
  auto emitScalarLeaf = [&](Type leafType) {
    auto lowered = lowerScalarType(origin->getContext(), leafType, field);
    if (failed(lowered)) {
      return failure();
    }
    if (!storage && prefixShape.empty()) {
      out.push_back(*lowered);
      return success();
    }
    llvm::SmallVector<int64_t> shape(prefixShape.begin(), prefixShape.end());
    if (shape.empty()) {
      shape.push_back(1);
    }
    out.push_back(MemRefType::get(shape, *lowered));
    return success();
  };
 
  if (isScalarType(type)) {
    return emitScalarLeaf(type);
  }
 
  if (auto arrayType = dyn_cast<array::ArrayType>(type)) {
    llvm::SmallVector<int64_t> newPrefix(prefixShape.begin(), prefixShape.end());
    newPrefix.append(arrayType.getShape().begin(), arrayType.getShape().end());
    return flattenTypeLeaves(
        arrayType.getElementType(), tables, origin, field, out, newPrefix, true
    );
  }
 
  if (auto podType = dyn_cast<pod::PodType>(type)) {
    for (pod::RecordAttr record : podType.getRecords()) {
      if (failed(
              flattenTypeLeaves(record.getType(), tables, origin, field, out, prefixShape, true)
          )) {
        return failure();
      }
    }
    return success();
  }
 
  if (auto structType = dyn_cast<component::StructType>(type)) {
    auto def = structType.getDefinition(tables, origin);
    if (failed(def)) {
      origin->emitError("could not resolve struct type during witgen lowering");
      return failure();
    }
    for (component::MemberDefOp member : def->get().getMemberDefs()) {
      if (failed(
              flattenTypeLeaves(member.getType(), tables, origin, field, out, prefixShape, true)
          )) {
        return failure();
      }
    }
    return success();
  }
 
  origin->emitError("unsupported type in llzk-witgen lowering: ") << type;
  return failure();
}

static MemRefType
getStridedMemRefType(MLIRContext *context, ArrayRef<int64_t> shape, Type elementType) {
  SmallVector<int64_t> strides(shape.size(), ShapedType::kDynamic);
  return MemRefType::get(
      shape, elementType, StridedLayoutAttr::get(context, ShapedType::kDynamic, strides)
  );
}

static LogicalResult flattenABILeafTypes(
    Type type, SymbolTableCollection &tables, Operation *origin, const Field &field,
    SmallVectorImpl<Type> &out, size_t prefixRank = 0, bool aggregateStorage = false
) {
  auto emitScalarLeaf = [&](Type leafType) {
    auto lowered = lowerScalarType(origin->getContext(), leafType, field);
    if (failed(lowered)) {
      return failure();
    }
    if (!aggregateStorage && prefixRank == 0) {
      out.push_back(*lowered);
      return success();
    }
    SmallVector<int64_t> shape;
    if (prefixRank == 0) {
      shape.push_back(1);
    } else {
      shape.assign(prefixRank, ShapedType::kDynamic);
    }
    out.push_back(getStridedMemRefType(origin->getContext(), shape, *lowered));
    return success();
  };
 
  if (isScalarType(type)) {
    return emitScalarLeaf(type);
  }
 
  if (auto arrayType = dyn_cast<array::ArrayType>(type)) {
    return flattenABILeafTypes(
        arrayType.getElementType(), tables, origin, field, out, prefixRank + arrayType.getRank(),
        true
    );
  }
 
  if (auto podType = dyn_cast<pod::PodType>(type)) {
    for (pod::RecordAttr record : podType.getRecords()) {
      if (failed(
              flattenABILeafTypes(record.getType(), tables, origin, field, out, prefixRank, true)
          )) {
        return failure();
      }
    }
    return success();
  }
 
  if (auto structType = dyn_cast<component::StructType>(type)) {
    auto def = structType.getDefinition(tables, origin);
    if (failed(def)) {
      origin->emitError("could not resolve struct type during witgen lowering");
      return failure();
    }
    for (component::MemberDefOp member : def->get().getMemberDefs()) {
      if (failed(
              flattenABILeafTypes(member.getType(), tables, origin, field, out, prefixRank, true)
          )) {
        return failure();
      }
    }
    return success();
  }
 
  origin->emitError("unsupported type in llzk-witgen lowering: ") << type;
  return failure();
}

static FailureOr<size_t>
getLeafCount(Type type, SymbolTableCollection &tables, Operation *origin, const Field &field) {
  SmallVector<Type> leaves;
  if (failed(flattenTypeLeaves(type, tables, origin, field, leaves))) {
    return failure();
  }
  return leaves.size();
}

static FailureOr<SmallVector<Type>>
getLeafTypes(Type type, SymbolTableCollection &tables, Operation *origin, const Field &field) {
  SmallVector<Type> leaves;
  if (failed(flattenTypeLeaves(type, tables, origin, field, leaves))) {
    return failure();
  }
  return leaves;
}

static FailureOr<SmallVector<Type>>
getABILeafTypes(Type type, SymbolTableCollection &tables, Operation *origin, const Field &field) {
  SmallVector<Type> leaves;
  if (failed(flattenABILeafTypes(type, tables, origin, field, leaves))) {
    return failure();
  }
  return leaves;
}

static FailureOr<std::pair<size_t, size_t>> getNamedLeafSpan(
    Type ownerType, StringRef name, SymbolTableCollection &tables, Operation *origin,
    const Field &field
) {
  if (auto podType = dyn_cast<pod::PodType>(ownerType)) {
    size_t running = 0;
    for (pod::RecordAttr record : podType.getRecords()) {
      auto count = getLeafCount(record.getType(), tables, origin, field);
      if (failed(count)) {
        return failure();
      }
      if (record.getName().getValue() == name) {
        return std::pair<size_t, size_t> {running, *count};
      }
      running += *count;
    }
  }
 
  if (auto structType = dyn_cast<component::StructType>(ownerType)) {
    auto def = structType.getDefinition(tables, origin);
    if (failed(def)) {
      origin->emitError("could not resolve struct type during witgen lowering");
      return failure();
    }
    size_t running = 0;
    for (component::MemberDefOp member : def->get().getMemberDefs()) {
      auto count = getLeafCount(member.getType(), tables, origin, field);
      if (failed(count)) {
        return failure();
      }
      if (member.getSymName() == name) {
        return std::pair<size_t, size_t> {running, *count};
      }
      running += *count;
    }
  }
 
  origin->emitError("could not resolve aggregate member/record @") << name;
  return failure();
}

static FailureOr<Type>
getNamedSubType(Type ownerType, StringRef name, SymbolTableCollection &tables, Operation *origin) {
  if (auto podType = dyn_cast<pod::PodType>(ownerType)) {
    for (pod::RecordAttr record : podType.getRecords()) {
      if (record.getName().getValue() == name) {
        return record.getType();
      }
    }
  }
  if (auto structType = dyn_cast<component::StructType>(ownerType)) {
    auto def = structType.getDefinition(tables, origin);
    if (failed(def)) {
      origin->emitError("could not resolve struct type during witgen lowering");
      return failure();
    }
    for (component::MemberDefOp member : def->get().getMemberDefs()) {
      if (member.getSymName() == name) {
        return member.getType();
      }
    }
  }
  origin->emitError("could not resolve aggregate member/record @") << name;
  return failure();
}

static FailureOr<Value> createZeroMemRef(OpBuilder &builder, Location loc, MemRefType memrefType) {
  auto elementCount = getStaticElementCount(memrefType, "witgen zero memref");
  if (!elementCount) {
    emitError(loc) << llvm::toString(elementCount.takeError());
    return failure();
  }
  Value alloc = builder.create<memref::AllocOp>(loc, memrefType);
  auto elementType = memrefType.getElementType();
  Value zero;
  if (isa<IndexType>(elementType)) {
    zero = builder.create<arith::ConstantIndexOp>(loc, 0);
  } else {
    zero = builder.create<arith::ConstantOp>(
        loc, IntegerAttr::get(llvm::cast<IntegerType>(elementType), 0)
    );
  }
  auto strides = mlir::computeStrides(memrefType.getShape());
  for (size_t flat = 0; flat < *elementCount; ++flat) {
    auto flatSigned = checkedCast<int64_t>(flat);
    if (!flatSigned) {
      emitError(loc) << llvm::toString(flatSigned.takeError());
      return failure();
    }
    SmallVector<Value> indices;
    for (int64_t index : mlir::delinearize(*flatSigned, strides)) {
      indices.push_back(makeIndexConstant(builder, loc, index));
    }
    builder.create<memref::StoreOp>(loc, zero, alloc, indices);
  }
  return alloc;
}

static FailureOr<Value> createRandomMemRef(
    OpBuilder &builder, Location loc, MemRefType memrefType, const Field &field,
    std::mt19937_64 &rng
) {
  auto elementCount = getStaticElementCount(memrefType, "witgen random memref");
  if (!elementCount) {
    emitError(loc) << llvm::toString(elementCount.takeError());
    return failure();
  }
  Value alloc = builder.create<memref::AllocOp>(loc, memrefType);
  auto elementType = memrefType.getElementType();
  auto strides = mlir::computeStrides(memrefType.getShape());
  for (size_t flat = 0; flat < *elementCount; ++flat) {
    auto flatSigned = checkedCast<int64_t>(flat);
    if (!flatSigned) {
      emitError(loc) << llvm::toString(flatSigned.takeError());
      return failure();
    }
    SmallVector<Value> indices;
    for (int64_t index : mlir::delinearize(*flatSigned, strides)) {
      indices.push_back(makeIndexConstant(builder, loc, index));
    }
    if (isa<IndexType>(elementType)) {
      auto value = randomIndexValue(rng);
      builder.create<memref::StoreOp>(
          loc, builder.create<arith::ConstantIndexOp>(loc, value), alloc, indices
      );
      continue;
    }
    auto intType = llvm::cast<IntegerType>(elementType);
    if (intType.getWidth() == 1) {
      builder.create<memref::StoreOp>(
          loc,
          builder.create<arith::ConstantOp>(
              loc, IntegerAttr::get(intType, APInt(1, randomBoolValue(rng)))
          ),
          alloc, indices
      );
      continue;
    }
    auto candidate = randomFieldElement(rng, field);
    builder.create<memref::StoreOp>(
        loc,
        builder.create<arith::ConstantOp>(
            loc, IntegerAttr::get(intType, llzk::toExactWidthAPInt(candidate, intType.getWidth()))
        ),
        alloc, indices
    );
  }
  return alloc;
}

static FailureOr<LoweredValue> createDefaultValue(
    OpBuilder &builder, Location loc, Type type, SymbolTableCollection &tables, Operation *origin,
    const Field &field, UninitializedBehavior behavior, std::mt19937_64 &rng
) {
  LoweredValue lowered {type, {}};
  auto leafTypes = getLeafTypes(type, tables, origin, field);
  if (failed(leafTypes)) {
    return failure();
  }
  for (Type leafType : *leafTypes) {
    if (behavior == UninitializedBehavior::Fail) {
      origin->emitError(
          "fail-mode default materialization is unsupported in witgen lowering because it would "
          "hide uninitialized reads"
      );
      return failure();
    }
    if (behavior == UninitializedBehavior::Random) {
      if (auto memrefType = dyn_cast<MemRefType>(leafType)) {
        auto randomMemRef = createRandomMemRef(builder, loc, memrefType, field, rng);
        if (failed(randomMemRef)) {
          return failure();
        }
        lowered.leaves.push_back(*randomMemRef);
        continue;
      }
      if (isa<IndexType>(leafType)) {
        lowered.leaves.push_back(
            builder.create<arith::ConstantIndexOp>(loc, randomIndexValue(rng))
        );
        continue;
      }
      auto intType = llvm::cast<IntegerType>(leafType);
      if (intType.getWidth() == 1) {
        lowered.leaves.push_back(builder.create<arith::ConstantOp>(
            loc, IntegerAttr::get(intType, APInt(1, randomBoolValue(rng)))
        ));
        continue;
      }
      auto candidate = randomFieldElement(rng, field);
      lowered.leaves.push_back(builder.create<arith::ConstantOp>(
          loc, IntegerAttr::get(intType, llzk::toExactWidthAPInt(candidate, intType.getWidth()))
      ));
      continue;
    }
    if (auto memrefType = dyn_cast<MemRefType>(leafType)) {
      auto zeroMemRef = createZeroMemRef(builder, loc, memrefType);
      if (failed(zeroMemRef)) {
        return failure();
      }
      lowered.leaves.push_back(*zeroMemRef);
      continue;
    }
    if (isa<IndexType>(leafType)) {
      lowered.leaves.push_back(builder.create<arith::ConstantIndexOp>(loc, 0));
      continue;
    }
    lowered.leaves.push_back(builder.create<arith::ConstantOp>(
        loc, IntegerAttr::get(llvm::cast<IntegerType>(leafType), 0)
    ));
  }
  return lowered;
}

static Value normalizeWideValue(
    OpBuilder &builder, Location loc, Value wideValue, unsigned dstWidth, const Field &field
) {
  auto wideType = llvm::cast<IntegerType>(wideValue.getType());
  Value modulus = builder.create<arith::ConstantOp>(
      loc, field.getPrimeAttr(builder.getContext(), wideType.getWidth())
  );
  Value reduced = builder.create<arith::RemUIOp>(loc, wideValue, modulus);
  return builder.create<arith::TruncIOp>(
      loc, IntegerType::get(builder.getContext(), dstWidth), reduced
  );
}

static Value normalizeSignedWideValue(
    OpBuilder &builder, Location loc, Value wideValue, unsigned dstWidth, const Field &field
) {
  auto wideType = llvm::cast<IntegerType>(wideValue.getType());
  Value modulus = builder.create<arith::ConstantOp>(
      loc, field.getPrimeAttr(builder.getContext(), wideType.getWidth())
  );
  Value reduced = builder.create<arith::RemSIOp>(loc, wideValue, modulus);
  Value zero = builder.create<arith::ConstantOp>(loc, IntegerAttr::get(wideType, 0));
  Value isNegative = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, reduced, zero);
  Value adjusted = builder.create<arith::AddIOp>(loc, reduced, modulus);
  Value canonical = builder.create<arith::SelectOp>(loc, isNegative, adjusted, reduced);
  return builder.create<arith::TruncIOp>(
      loc, IntegerType::get(builder.getContext(), dstWidth), canonical
  );
}

static Value
lowerFeltToSignedWide(OpBuilder &builder, Location loc, Value operand, const Field &field) {
  unsigned width = field.bitWidth();
  unsigned wideWidth = width + 1;
  auto feltType = IntegerType::get(builder.getContext(), width);
  auto wideType = IntegerType::get(builder.getContext(), wideWidth);
  Value operandWide = builder.create<arith::ExtUIOp>(loc, wideType, operand);
  Value prime =
      builder.create<arith::ConstantOp>(loc, field.getPrimeAttr(builder.getContext(), wideWidth));
  Value half = builder.create<arith::ConstantOp>(
      loc, IntegerAttr::get(feltType, toExactWidthAPInt(field.half(), width))
  );
  Value isNegative = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::uge, operand, half);
  Value signedOperand = builder.create<arith::SubIOp>(loc, operandWide, prime);
  return builder.create<arith::SelectOp>(loc, isNegative, signedOperand, operandWide);
}

static void assertNonZeroFelt(OpBuilder &builder, Location loc, Value operand, StringRef message) {
  auto operandType = llvm::cast<IntegerType>(operand.getType());
  Value zero = builder.create<arith::ConstantOp>(loc, IntegerAttr::get(operandType, 0));
  Value isNonZero = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, operand, zero);
  builder.create<cf::AssertOp>(loc, isNonZero, message);
}

static Value
lowerFeltAdd(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  unsigned width = field.bitWidth();
  unsigned wideWidth = width + 1;
  auto wideType = IntegerType::get(builder.getContext(), wideWidth);
  Value lhsWide = builder.create<arith::ExtUIOp>(loc, wideType, lhs);
  Value rhsWide = builder.create<arith::ExtUIOp>(loc, wideType, rhs);
  Value sum = builder.create<arith::AddIOp>(loc, lhsWide, rhsWide);
  return normalizeWideValue(builder, loc, sum, width, field);
}

static Value
lowerFeltSub(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  unsigned width = field.bitWidth();
  unsigned wideWidth = width + 1;
  auto wideType = IntegerType::get(builder.getContext(), wideWidth);
  Value lhsWide = builder.create<arith::ExtUIOp>(loc, wideType, lhs);
  Value rhsWide = builder.create<arith::ExtUIOp>(loc, wideType, rhs);
  Value modulus =
      builder.create<arith::ConstantOp>(loc, field.getPrimeAttr(builder.getContext(), wideWidth));
  Value lhsPlusMod = builder.create<arith::AddIOp>(loc, lhsWide, modulus);
  Value diff = builder.create<arith::SubIOp>(loc, lhsPlusMod, rhsWide);
  return normalizeWideValue(builder, loc, diff, width, field);
}

static Value lowerFeltNeg(OpBuilder &builder, Location loc, Value operand, const Field &field) {
  unsigned width = field.bitWidth();
  unsigned wideWidth = width + 1;
  auto wideType = IntegerType::get(builder.getContext(), wideWidth);
  Value operandWide = builder.create<arith::ExtUIOp>(loc, wideType, operand);
  Value modulus =
      builder.create<arith::ConstantOp>(loc, field.getPrimeAttr(builder.getContext(), wideWidth));
  Value diff = builder.create<arith::SubIOp>(loc, modulus, operandWide);
  return normalizeWideValue(builder, loc, diff, width, field);
}

static Value
lowerFeltMul(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  unsigned width = field.bitWidth();
  unsigned wideWidth = width * 2;
  auto wideType = IntegerType::get(builder.getContext(), wideWidth);
  Value lhsWide = builder.create<arith::ExtUIOp>(loc, wideType, lhs);
  Value rhsWide = builder.create<arith::ExtUIOp>(loc, wideType, rhs);
  Value product = builder.create<arith::MulIOp>(loc, lhsWide, rhsWide);
  return normalizeWideValue(builder, loc, product, width, field);
}

static Value lowerFeltInv(OpBuilder &builder, Location loc, Value operand, const Field &field) {
  llvm::APInt exponent = toExactWidthAPInt(field.prime() - 2, field.bitWidth());
  Value result = makeOneFelt(builder, loc, field);
  Value base = operand;
  for (unsigned bit = 0; bit < exponent.getBitWidth(); ++bit) {
    if (exponent[bit]) {
      result = lowerFeltMul(builder, loc, result, base, field);
    }
    if (bit + 1 < exponent.getBitWidth()) {
      base = lowerFeltMul(builder, loc, base, base, field);
    }
  }
  return result;
}

static Value
lowerFeltDiv(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  return lowerFeltMul(builder, loc, lhs, lowerFeltInv(builder, loc, rhs, field), field);
}

static Value
lowerFeltPow(OpBuilder &builder, Location loc, Value base, Value exponent, const Field &field) {
  auto feltType = IntegerType::get(builder.getContext(), field.bitWidth());
  Value zero = builder.create<arith::ConstantOp>(loc, IntegerAttr::get(feltType, 0));
  Value one = builder.create<arith::ConstantOp>(loc, IntegerAttr::get(feltType, 1));
  Value result = makeOneFelt(builder, loc, field);
  Value currentBase = base;
  for (unsigned bit = 0; bit < field.bitWidth(); ++bit) {
    Value bitIndex = builder.create<arith::ConstantOp>(
        loc, IntegerAttr::get(feltType, llvm::APInt(field.bitWidth(), bit))
    );
    Value shifted = builder.create<arith::ShRUIOp>(loc, exponent, bitIndex);
    Value masked = builder.create<arith::AndIOp>(loc, shifted, one);
    Value bitIsSet = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, masked, zero);
    auto ifOp = builder.create<scf::IfOp>(loc, TypeRange {feltType}, bitIsSet, true);
    {
      OpBuilder::InsertionGuard guard(builder);
      builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
      Value multiplied = lowerFeltMul(builder, loc, result, currentBase, field);
      builder.create<scf::YieldOp>(loc, multiplied);
    }
    {
      OpBuilder::InsertionGuard guard(builder);
      builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
      builder.create<scf::YieldOp>(loc, result);
    }
    result = ifOp.getResult(0);
    if (bit + 1 < field.bitWidth()) {
      currentBase = lowerFeltMul(builder, loc, currentBase, currentBase, field);
    }
  }
  return result;
}

static Value
lowerFeltShl(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  auto feltType = IntegerType::get(builder.getContext(), field.bitWidth());
  Value two = builder.create<arith::ConstantOp>(loc, IntegerAttr::get(feltType, 2));
  return lowerFeltMul(builder, loc, lhs, lowerFeltPow(builder, loc, two, rhs, field), field);
}

static Value
lowerFeltOr(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  unsigned width = field.bitWidth();
  auto wideType = IntegerType::get(builder.getContext(), width + 1);
  Value orValue = builder.create<arith::OrIOp>(loc, lhs, rhs);
  Value orWide = builder.create<arith::ExtUIOp>(loc, wideType, orValue);
  return normalizeWideValue(builder, loc, orWide, width, field);
}

static Value
lowerFeltXor(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  unsigned width = field.bitWidth();
  auto wideType = IntegerType::get(builder.getContext(), width + 1);
  Value xorValue = builder.create<arith::XOrIOp>(loc, lhs, rhs);
  Value xorWide = builder.create<arith::ExtUIOp>(loc, wideType, xorValue);
  return normalizeWideValue(builder, loc, xorWide, width, field);
}

static Value lowerFeltUnsignedDiv(OpBuilder &builder, Location loc, Value lhs, Value rhs) {
  return builder.create<arith::DivUIOp>(loc, lhs, rhs);
}

static Value
lowerFeltSignedDiv(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  unsigned width = field.bitWidth();
  Value lhsSigned = lowerFeltToSignedWide(builder, loc, lhs, field);
  Value rhsSigned = lowerFeltToSignedWide(builder, loc, rhs, field);
  Value quotient = builder.create<arith::DivSIOp>(loc, lhsSigned, rhsSigned);
  return normalizeSignedWideValue(builder, loc, quotient, width, field);
}

static Value lowerFeltUnsignedMod(OpBuilder &builder, Location loc, Value lhs, Value rhs) {
  return builder.create<arith::RemUIOp>(loc, lhs, rhs);
}

static Value
lowerFeltSignedMod(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  unsigned width = field.bitWidth();
  Value lhsSigned = lowerFeltToSignedWide(builder, loc, lhs, field);
  Value rhsSigned = lowerFeltToSignedWide(builder, loc, rhs, field);
  Value remainder = builder.create<arith::RemSIOp>(loc, lhsSigned, rhsSigned);
  return normalizeSignedWideValue(builder, loc, remainder, width, field);
}

static Value
lowerFeltShr(OpBuilder &builder, Location loc, Value lhs, Value rhs, const Field &field) {
  auto feltType = IntegerType::get(builder.getContext(), field.bitWidth());
  Value width = builder.create<arith::ConstantOp>(
      loc, IntegerAttr::get(feltType, llvm::APInt(field.bitWidth(), field.bitWidth()))
  );
  Value shiftTooLarge = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::uge, rhs, width);
  Value zero = builder.create<arith::ConstantOp>(loc, IntegerAttr::get(feltType, 0));
  Value maxValidShift = builder.create<arith::ConstantOp>(
      loc, IntegerAttr::get(feltType, llvm::APInt(field.bitWidth(), field.bitWidth() - 1))
  );
  Value clampedShift = builder.create<arith::MinUIOp>(loc, rhs, maxValidShift);
  Value shifted = builder.create<arith::ShRUIOp>(loc, lhs, clampedShift);
  return builder.create<arith::SelectOp>(loc, shiftTooLarge, zero, shifted);
}

static Value lowerFeltNot(OpBuilder &builder, Location loc, Value operand, const Field &field) {
  unsigned width = field.bitWidth();
  auto feltType = IntegerType::get(builder.getContext(), width);
  auto wideType = IntegerType::get(builder.getContext(), width + 1);
  Value maxMask = builder.create<arith::ConstantOp>(
      loc, IntegerAttr::get(feltType, llvm::APInt::getAllOnes(width))
  );
  Value complement = builder.create<arith::XOrIOp>(loc, operand, maxMask);
  Value complementWide = builder.create<arith::ExtUIOp>(loc, wideType, complement);
  return normalizeWideValue(builder, loc, complementWide, width, field);
}

static Value loadStorageScalar(OpBuilder &builder, Location loc, Value storageLeaf) {
  auto memrefType = llvm::cast<MemRefType>(storageLeaf.getType());
  SmallVector<Value> indices;
  indices.reserve(memrefType.getRank());
  for (int64_t dim = 0; dim < memrefType.getRank(); ++dim) {
    indices.push_back(makeIndexConstant(builder, loc, 0));
  }
  return builder.create<memref::LoadOp>(loc, storageLeaf, indices);
}

static void storeStorageScalar(OpBuilder &builder, Location loc, Value scalar, Value storageLeaf) {
  auto memrefType = llvm::cast<MemRefType>(storageLeaf.getType());
  SmallVector<Value> indices;
  indices.reserve(memrefType.getRank());
  for (int64_t dim = 0; dim < memrefType.getRank(); ++dim) {
    indices.push_back(makeIndexConstant(builder, loc, 0));
  }
  builder.create<memref::StoreOp>(loc, scalar, storageLeaf, indices);
}

static LogicalResult copyIntoStorage(
    OpBuilder &builder, Location loc, Type sourceType, ArrayRef<Value> destLeaves,
    ArrayRef<Value> sourceLeaves, SymbolTableCollection &tables, Operation *origin,
    const Field &field
) {
  auto leafTypes = getLeafTypes(sourceType, tables, origin, field);
  if (failed(leafTypes)) {
    return failure();
  }
  if (destLeaves.size() != sourceLeaves.size() || destLeaves.size() != leafTypes->size()) {
    origin->emitError("flattened leaf mismatch while copying aggregate storage");
    return failure();
  }
  for (auto [leafType, destLeaf, srcLeaf] : llvm::zip(*leafTypes, destLeaves, sourceLeaves)) {
    if (isa<MemRefType>(leafType)) {
      builder.create<memref::CopyOp>(loc, srcLeaf, destLeaf);
      continue;
    }
    storeStorageScalar(builder, loc, srcLeaf, destLeaf);
  }
  return success();
}

static FailureOr<LoweredValue> readNamedAggregateValue(
    OpBuilder &builder, Location loc, Type ownerType, StringRef name, const LoweredValue &owner,
    SymbolTableCollection &tables, Operation *origin, const Field &field
) {
  auto subType = getNamedSubType(ownerType, name, tables, origin);
  if (failed(subType)) {
    return failure();
  }
  auto span = getNamedLeafSpan(ownerType, name, tables, origin, field);
  if (failed(span)) {
    return failure();
  }
  LoweredValue result {*subType, {}};
  auto leafTypes = getLeafTypes(*subType, tables, origin, field);
  if (failed(leafTypes)) {
    return failure();
  }
  auto leaves = ArrayRef<Value>(owner.leaves).slice(span->first, span->second);
  for (auto [leafType, leafValue] : llvm::zip(*leafTypes, leaves)) {
    if (isa<MemRefType>(leafType)) {
      result.leaves.push_back(leafValue);
    } else {
      result.leaves.push_back(loadStorageScalar(builder, loc, leafValue));
    }
  }
  return result;
}

static LogicalResult writeNamedAggregateValue(
    OpBuilder &builder, Location loc, Type ownerType, StringRef name, LoweredValue &owner,
    const LoweredValue &value, SymbolTableCollection &tables, Operation *origin, const Field &field
) {
  auto subType = getNamedSubType(ownerType, name, tables, origin);
  if (failed(subType)) {
    return failure();
  }
  auto span = getNamedLeafSpan(ownerType, name, tables, origin, field);
  if (failed(span)) {
    return failure();
  }
  return copyIntoStorage(
      builder, loc, *subType, ArrayRef<Value>(owner.leaves).slice(span->first, span->second),
      value.leaves, tables, origin, field
  );
}

static FailureOr<Value>
createElementSubview(OpBuilder &builder, Location loc, Value source, ValueRange outerIndices) {
  auto sourceType = llvm::cast<MemRefType>(source.getType());
  SmallVector<OpFoldResult> mixedOffsets;
  SmallVector<OpFoldResult> mixedSizes;
  SmallVector<OpFoldResult> mixedStrides;
  auto indexedRank = checkedCast<int64_t>(outerIndices.size());
  if (!indexedRank) {
    emitError(loc) << llvm::toString(indexedRank.takeError());
    return failure();
  }
  mixedOffsets.reserve(sourceType.getRank());
  mixedSizes.reserve(sourceType.getRank());
  mixedStrides.reserve(sourceType.getRank());
  for (Value index : outerIndices) {
    mixedOffsets.push_back(index);
  }
  for (int64_t dim = *indexedRank; dim < sourceType.getRank(); ++dim) {
    mixedOffsets.push_back(builder.getIndexAttr(0));
  }
  for (int64_t dim = 0; dim < *indexedRank; ++dim) {
    mixedSizes.push_back(builder.getIndexAttr(1));
  }
  for (int64_t dim = *indexedRank; dim < sourceType.getRank(); ++dim) {
    mixedSizes.push_back(memref::getMixedSize(builder, loc, source, dim));
  }
  for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) {
    mixedStrides.push_back(builder.getIndexAttr(1));
  }
  SmallVector<int64_t> desiredShape;
  auto reserveSize = checkedCast<size_t>(sourceType.getRank() - *indexedRank);
  if (!reserveSize) {
    emitError(loc) << llvm::toString(reserveSize.takeError());
    return failure();
  }
  desiredShape.reserve(*reserveSize);
  for (int64_t dim = *indexedRank; dim < sourceType.getRank(); ++dim) {
    auto dimIndex = checkedCast<size_t>(dim);
    if (!dimIndex) {
      emitError(loc) << llvm::toString(dimIndex.takeError());
      return failure();
    }
    if (auto attr = llvm::dyn_cast<Attribute>(mixedSizes[*dimIndex])) {
      desiredShape.push_back(llvm::cast<IntegerAttr>(attr).getInt());
    } else {
      desiredShape.push_back(ShapedType::kDynamic);
    }
  }
  if (desiredShape.empty()) {
    desiredShape.push_back(1);
  }
  auto resultType = llvm::cast<MemRefType>(memref::SubViewOp::inferRankReducedResultType(
      desiredShape, sourceType, mixedOffsets, mixedSizes, mixedStrides
  ));
  auto op = builder.create<memref::SubViewOp>(
      loc, resultType, source, mixedOffsets, mixedSizes, mixedStrides
  );
  return success(op.getResult());
}

static FailureOr<LoweredValue> readArrayElement(
    OpBuilder &builder, Location loc, array::ArrayType arrayType, const LoweredValue &arrayValue,
    ArrayRef<Value> indices
) {
  Type elementType = arrayType.getElementType();
  LoweredValue result {elementType, {}};
  if (isScalarType(elementType)) {
    result.leaves.push_back(
        builder.create<memref::LoadOp>(loc, arrayValue.leaves.front(), indices)
    );
    return result;
  }
 
  for (Value sourceLeaf : arrayValue.leaves) {
    auto subview = createElementSubview(builder, loc, sourceLeaf, indices);
    if (failed(subview)) {
      return failure();
    }
    result.leaves.push_back(*subview);
  }
  return result;
}

static LogicalResult writeArrayElement(
    OpBuilder &builder, Location loc, array::ArrayType arrayType, LoweredValue &arrayValue,
    ArrayRef<Value> indices, const LoweredValue &elementValue
) {
  Type elementType = arrayType.getElementType();
  if (isScalarType(elementType)) {
    builder.create<memref::StoreOp>(
        loc, elementValue.leaves.front(), arrayValue.leaves.front(), indices
    );
    return success();
  }
 
  for (auto [destLeaf, srcLeaf] : llvm::zip(arrayValue.leaves, elementValue.leaves)) {
    auto subview = createElementSubview(builder, loc, destLeaf, indices);
    if (failed(subview)) {
      return failure();
    }
    builder.create<memref::CopyOp>(loc, srcLeaf, *subview);
  }
  return success();
}

static LogicalResult appendFlatLeavesToTypes(
    OpBuilder &builder, Location loc, const LoweredValue &value, ArrayRef<Type> targetLeafTypes,
    SmallVectorImpl<Value> &out, Operation *origin
) {
  if (targetLeafTypes.size() != value.leaves.size()) {
    origin->emitError("flattened leaf mismatch during call lowering");
    return failure();
  }
  for (auto [leafValue, leafType] : llvm::zip(value.leaves, targetLeafTypes)) {
    if (leafValue.getType() == leafType) {
      out.push_back(leafValue);
      continue;
    }
    if (isa<MemRefType>(leafValue.getType()) && isa<MemRefType>(leafType)) {
      out.push_back(builder.create<memref::CastOp>(loc, leafType, leafValue));
      continue;
    }
    origin->emitError("lowered leaf type mismatch during call lowering");
    return failure();
  }
  return success();
}

class BodyLowerer {
public:
  BodyLowerer(
      ModuleOp mod, SymbolTableCollection &symbolTables, const Field &moduleField,
      const WitgenOptions &options
  )
      : moduleOp(mod), tables(symbolTables), field(moduleField),
        uninitializedBehavior(options.uninitializedBehavior), rng(makeDefaultValueRng(options)) {}

  FailureOr<func::FuncOp> lowerFunction(function::FuncDefOp funcOp) {
    if (funcOp.isExternal()) {
      funcOp.emitError("execution-engine backend does not lower extern functions");
      return failure();
    }
    if (!funcOp.getBody().hasOneBlock()) {
      funcOp.emitError("execution-engine backend only supports single-block functions");
      return failure();
    }
 
    SmallVector<Type> loweredArgTypes;
    for (Type argType : funcOp.getArgumentTypes()) {
      if (failed(
              flattenABILeafTypes(argType, tables, funcOp.getOperation(), field, loweredArgTypes)
          )) {
        return failure();
      }
    }
    SmallVector<Type> loweredResultTypes;
    for (Type resultType : funcOp.getResultTypes()) {
      if (failed(flattenABILeafTypes(
              resultType, tables, funcOp.getOperation(), field, loweredResultTypes
          ))) {
        return failure();
      }
    }
 
    OpBuilder moduleBuilder(moduleOp.getContext());
    moduleBuilder.setInsertionPointToEnd(moduleOp.getBody());
    auto loweredFunc = moduleBuilder.create<func::FuncOp>(
        funcOp.getLoc(), mangleFunctionName(funcOp),
        moduleBuilder.getFunctionType(loweredArgTypes, loweredResultTypes)
    );
    Block *entry = loweredFunc.addEntryBlock();
    OpBuilder builder(entry, entry->begin());
 
    DenseMap<Value, LoweredValue> valueMap;
    unsigned cursor = 0;
    for (auto [arg, argType] :
         llvm::zip(funcOp.getBody().front().getArguments(), funcOp.getArgumentTypes())) {
      auto leafCount = getLeafCount(argType, tables, funcOp.getOperation(), field);
      if (failed(leafCount)) {
        loweredFunc.erase();
        return failure();
      }
      LoweredValue lowered {argType, {}};
      lowered.leaves.append(
          entry->getArguments().begin() + cursor,
          entry->getArguments().begin() + cursor + *leafCount
      );
      cursor += *leafCount;
      valueMap[arg] = std::move(lowered);
    }
 
    if (failed(lowerBlock(builder, funcOp.getBody().front(), valueMap))) {
      loweredFunc.erase();
      return failure();
    }
    return loweredFunc;
  }
 
private:
  ModuleOp moduleOp;
  SymbolTableCollection &tables;
  const Field &field;
  UninitializedBehavior uninitializedBehavior;
  std::mt19937_64 rng;

  FailureOr<LoweredValue>
  lookup(Value value, DenseMap<Value, LoweredValue> &valueMap, Operation *origin) {
    auto it = valueMap.find(value);
    if (it == valueMap.end()) {
      origin->emitError("failed to find lowered SSA value");
      return failure();
    }
    return it->second;
  }

  FailureOr<Value>
  lookupScalar(Value value, DenseMap<Value, LoweredValue> &valueMap, Operation *origin) {
    auto lowered = lookup(value, valueMap, origin);
    if (failed(lowered) || lowered->leaves.size() != 1 ||
        isa<MemRefType>(lowered->leaves.front().getType())) {
      origin->emitError("expected scalar lowered value");
      return failure();
    }
    return lowered->leaves.front();
  }

  LogicalResult
  lowerBlock(OpBuilder &builder, Block &block, DenseMap<Value, LoweredValue> &valueMap) {
    for (Operation &op : block) {
      if (failed(lowerOperation(builder, op, valueMap))) {
        return failure();
      }
    }
    return success();
  }

  FailureOr<Value>
  lowerFeltCmp(OpBuilder &builder, Location loc, boolean::CmpOp cmpOp, Value lhs, Value rhs) {
    arith::CmpIPredicate predicate;
    switch (cmpOp.getPredicate()) {
    case boolean::FeltCmpPredicate::EQ:
      predicate = arith::CmpIPredicate::eq;
      break;
    case boolean::FeltCmpPredicate::NE:
      predicate = arith::CmpIPredicate::ne;
      break;
    case boolean::FeltCmpPredicate::LT:
      predicate = arith::CmpIPredicate::ult;
      break;
    case boolean::FeltCmpPredicate::LE:
      predicate = arith::CmpIPredicate::ule;
      break;
    case boolean::FeltCmpPredicate::GT:
      predicate = arith::CmpIPredicate::ugt;
      break;
    case boolean::FeltCmpPredicate::GE:
      predicate = arith::CmpIPredicate::uge;
      break;
    }
    return builder.create<arith::CmpIOp>(loc, predicate, lhs, rhs).getResult();
  }

  LogicalResult
  lowerOperation(OpBuilder &builder, Operation &op, DenseMap<Value, LoweredValue> &valueMap) {
    Location loc = op.getLoc();
 
    auto bind = [&](Value result, LoweredValue lowered) {
      valueMap[result] = std::move(lowered);
      return success();
    };
 
    if (auto returnOp = dyn_cast<function::ReturnOp>(op)) {
      SmallVector<Value> results;
      for (Value operand : returnOp.getOperands()) {
        auto lowered = lookup(operand, valueMap, returnOp.getOperation());
        auto leafTypes = getABILeafTypes(operand.getType(), tables, returnOp.getOperation(), field);
        if (failed(lowered) || failed(leafTypes) ||
            failed(appendFlatLeavesToTypes(
                builder, loc, *lowered, *leafTypes, results, returnOp.getOperation()
            ))) {
          return failure();
        }
      }
      builder.create<func::ReturnOp>(loc, results);
      return success();
    }
 
    if (auto yieldOp = dyn_cast<scf::YieldOp>(op)) {
      SmallVector<Value> results;
      for (Value operand : yieldOp.getOperands()) {
        auto lowered = lookup(operand, valueMap, yieldOp.getOperation());
        auto leafTypes = getABILeafTypes(operand.getType(), tables, yieldOp.getOperation(), field);
        if (failed(lowered) || failed(leafTypes) ||
            failed(appendFlatLeavesToTypes(
                builder, loc, *lowered, *leafTypes, results, yieldOp.getOperation()
            ))) {
          return failure();
        }
      }
      builder.create<scf::YieldOp>(loc, results);
      return success();
    }
    if (auto conditionOp = dyn_cast<scf::ConditionOp>(op)) {
      auto condition =
          lookupScalar(conditionOp.getCondition(), valueMap, conditionOp.getOperation());
      if (failed(condition)) {
        return failure();
      }
      SmallVector<Value> results;
      for (Value operand : conditionOp.getArgs()) {
        auto lowered = lookup(operand, valueMap, conditionOp.getOperation());
        auto leafTypes =
            getABILeafTypes(operand.getType(), tables, conditionOp.getOperation(), field);
        if (failed(lowered) || failed(leafTypes) ||
            failed(appendFlatLeavesToTypes(
                builder, loc, *lowered, *leafTypes, results, conditionOp.getOperation()
            ))) {
          return failure();
        }
      }
      builder.create<scf::ConditionOp>(loc, *condition, results);
      return success();
    }
 
    if (auto constantOp = dyn_cast<arith::ConstantOp>(op)) {
      Operation *clone = builder.clone(op);
      return bind(
          constantOp.getResult(), LoweredValue {constantOp.getType(), {clone->getResult(0)}}
      );
    }
 
    if (auto feltConst = dyn_cast<felt::FeltConstantOp>(op)) {
      auto intType = IntegerType::get(builder.getContext(), field.bitWidth());
      // Reduce into the field first, then build an APInt with the exact storage width.
      auto constVal = toDynamicAPInt(feltConst.getValue().getValue());
      auto modVal = constVal % field.prime();
      auto intVal = llzk::toExactWidthAPInt(modVal, field.bitWidth());
      Value lowered = builder.create<arith::ConstantOp>(loc, IntegerAttr::get(intType, intVal));
      return bind(feltConst.getResult(), LoweredValue {feltConst.getType(), {lowered}});
    }
 
    if (auto nondetOp = dyn_cast<llzk::NonDetOp>(op)) {
      auto lowered = createDefaultValue(
          builder, loc, nondetOp.getType(), tables, nondetOp.getOperation(), field,
          uninitializedBehavior, rng
      );
      if (failed(lowered)) {
        return failure();
      }
      return bind(nondetOp.getResult(), std::move(*lowered));
    }
 
    if (auto addOp = dyn_cast<felt::AddFeltOp>(op)) {
      auto lhs = lookupScalar(addOp.getLhs(), valueMap, addOp.getOperation());
      auto rhs = lookupScalar(addOp.getRhs(), valueMap, addOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          addOp.getResult(),
          LoweredValue {addOp.getType(), {lowerFeltAdd(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto powOp = dyn_cast<felt::PowFeltOp>(op)) {
      auto lhs = lookupScalar(powOp.getLhs(), valueMap, powOp.getOperation());
      auto rhs = lookupScalar(powOp.getRhs(), valueMap, powOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          powOp.getResult(),
          LoweredValue {powOp.getType(), {lowerFeltPow(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto andOp = dyn_cast<felt::AndFeltOp>(op)) {
      auto lhs = lookupScalar(andOp.getLhs(), valueMap, andOp.getOperation());
      auto rhs = lookupScalar(andOp.getRhs(), valueMap, andOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          andOp.getResult(),
          LoweredValue {andOp.getType(), {builder.create<arith::AndIOp>(loc, *lhs, *rhs)}}
      );
    }
    if (auto orOp = dyn_cast<felt::OrFeltOp>(op)) {
      auto lhs = lookupScalar(orOp.getLhs(), valueMap, orOp.getOperation());
      auto rhs = lookupScalar(orOp.getRhs(), valueMap, orOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          orOp.getResult(),
          LoweredValue {orOp.getType(), {lowerFeltOr(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto xorOp = dyn_cast<felt::XorFeltOp>(op)) {
      auto lhs = lookupScalar(xorOp.getLhs(), valueMap, xorOp.getOperation());
      auto rhs = lookupScalar(xorOp.getRhs(), valueMap, xorOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          xorOp.getResult(),
          LoweredValue {xorOp.getType(), {lowerFeltXor(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto subOp = dyn_cast<felt::SubFeltOp>(op)) {
      auto lhs = lookupScalar(subOp.getLhs(), valueMap, subOp.getOperation());
      auto rhs = lookupScalar(subOp.getRhs(), valueMap, subOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          subOp.getResult(),
          LoweredValue {subOp.getType(), {lowerFeltSub(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto mulOp = dyn_cast<felt::MulFeltOp>(op)) {
      auto lhs = lookupScalar(mulOp.getLhs(), valueMap, mulOp.getOperation());
      auto rhs = lookupScalar(mulOp.getRhs(), valueMap, mulOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          mulOp.getResult(),
          LoweredValue {mulOp.getType(), {lowerFeltMul(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto negOp = dyn_cast<felt::NegFeltOp>(op)) {
      auto operand = lookupScalar(negOp.getOperand(), valueMap, negOp.getOperation());
      if (failed(operand)) {
        return failure();
      }
      return bind(
          negOp.getResult(),
          LoweredValue {negOp.getType(), {lowerFeltNeg(builder, loc, *operand, field)}}
      );
    }
    if (auto invOp = dyn_cast<felt::InvFeltOp>(op)) {
      auto operand = lookupScalar(invOp.getOperand(), valueMap, invOp.getOperation());
      if (failed(operand)) {
        return failure();
      }
      return bind(
          invOp.getResult(),
          LoweredValue {invOp.getType(), {lowerFeltInv(builder, loc, *operand, field)}}
      );
    }
    if (auto divOp = dyn_cast<felt::DivFeltOp>(op)) {
      auto lhs = lookupScalar(divOp.getLhs(), valueMap, divOp.getOperation());
      auto rhs = lookupScalar(divOp.getRhs(), valueMap, divOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          divOp.getResult(),
          LoweredValue {divOp.getType(), {lowerFeltDiv(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto uintDivOp = dyn_cast<felt::UnsignedIntDivFeltOp>(op)) {
      auto lhs = lookupScalar(uintDivOp.getLhs(), valueMap, uintDivOp.getOperation());
      auto rhs = lookupScalar(uintDivOp.getRhs(), valueMap, uintDivOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      assertNonZeroFelt(builder, loc, *rhs, "felt.uintdiv divisor must be non-zero");
      return bind(
          uintDivOp.getResult(),
          LoweredValue {uintDivOp.getType(), {lowerFeltUnsignedDiv(builder, loc, *lhs, *rhs)}}
      );
    }
    if (auto sintDivOp = dyn_cast<felt::SignedIntDivFeltOp>(op)) {
      auto lhs = lookupScalar(sintDivOp.getLhs(), valueMap, sintDivOp.getOperation());
      auto rhs = lookupScalar(sintDivOp.getRhs(), valueMap, sintDivOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      assertNonZeroFelt(builder, loc, *rhs, "felt.sintdiv divisor must be non-zero");
      return bind(
          sintDivOp.getResult(),
          LoweredValue {sintDivOp.getType(), {lowerFeltSignedDiv(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto umodOp = dyn_cast<felt::UnsignedModFeltOp>(op)) {
      auto lhs = lookupScalar(umodOp.getLhs(), valueMap, umodOp.getOperation());
      auto rhs = lookupScalar(umodOp.getRhs(), valueMap, umodOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      assertNonZeroFelt(builder, loc, *rhs, "felt.umod divisor must be non-zero");
      return bind(
          umodOp.getResult(),
          LoweredValue {umodOp.getType(), {lowerFeltUnsignedMod(builder, loc, *lhs, *rhs)}}
      );
    }
    if (auto smodOp = dyn_cast<felt::SignedModFeltOp>(op)) {
      auto lhs = lookupScalar(smodOp.getLhs(), valueMap, smodOp.getOperation());
      auto rhs = lookupScalar(smodOp.getRhs(), valueMap, smodOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      assertNonZeroFelt(builder, loc, *rhs, "felt.smod divisor must be non-zero");
      return bind(
          smodOp.getResult(),
          LoweredValue {smodOp.getType(), {lowerFeltSignedMod(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto shrOp = dyn_cast<felt::ShrFeltOp>(op)) {
      auto lhs = lookupScalar(shrOp.getLhs(), valueMap, shrOp.getOperation());
      auto rhs = lookupScalar(shrOp.getRhs(), valueMap, shrOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          shrOp.getResult(),
          LoweredValue {shrOp.getType(), {lowerFeltShr(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto shlOp = dyn_cast<felt::ShlFeltOp>(op)) {
      auto lhs = lookupScalar(shlOp.getLhs(), valueMap, shlOp.getOperation());
      auto rhs = lookupScalar(shlOp.getRhs(), valueMap, shlOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          shlOp.getResult(),
          LoweredValue {shlOp.getType(), {lowerFeltShl(builder, loc, *lhs, *rhs, field)}}
      );
    }
    if (auto notOp = dyn_cast<felt::NotFeltOp>(op)) {
      auto operand = lookupScalar(notOp.getOperand(), valueMap, notOp.getOperation());
      if (failed(operand)) {
        return failure();
      }
      return bind(
          notOp.getResult(),
          LoweredValue {notOp.getType(), {lowerFeltNot(builder, loc, *operand, field)}}
      );
    }
 
    if (auto cmpOp = dyn_cast<boolean::CmpOp>(op)) {
      auto lhs = lookupScalar(cmpOp.getLhs(), valueMap, cmpOp.getOperation());
      auto rhs = lookupScalar(cmpOp.getRhs(), valueMap, cmpOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      auto lowered = lowerFeltCmp(builder, loc, cmpOp, *lhs, *rhs);
      if (failed(lowered)) {
        return failure();
      }
      return bind(cmpOp.getResult(), LoweredValue {cmpOp.getType(), {*lowered}});
    }
    if (auto assertOp = dyn_cast<boolean::AssertOp>(op)) {
      auto condition = lookupScalar(assertOp.getCondition(), valueMap, assertOp.getOperation());
      if (failed(condition)) {
        return failure();
      }
      builder.create<cf::AssertOp>(
          loc, *condition, assertOp.getMsg() ? assertOp.getMsg()->str() : "bool.assert failed"
      );
      return success();
    }
    if (auto andOp = dyn_cast<boolean::AndBoolOp>(op)) {
      auto lhs = lookupScalar(andOp.getLhs(), valueMap, andOp.getOperation());
      auto rhs = lookupScalar(andOp.getRhs(), valueMap, andOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          andOp.getResult(),
          LoweredValue {andOp.getType(), {builder.create<arith::AndIOp>(loc, *lhs, *rhs)}}
      );
    }
    if (auto orOp = dyn_cast<boolean::OrBoolOp>(op)) {
      auto lhs = lookupScalar(orOp.getLhs(), valueMap, orOp.getOperation());
      auto rhs = lookupScalar(orOp.getRhs(), valueMap, orOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          orOp.getResult(),
          LoweredValue {orOp.getType(), {builder.create<arith::OrIOp>(loc, *lhs, *rhs)}}
      );
    }
    if (auto xorOp = dyn_cast<boolean::XorBoolOp>(op)) {
      auto lhs = lookupScalar(xorOp.getLhs(), valueMap, xorOp.getOperation());
      auto rhs = lookupScalar(xorOp.getRhs(), valueMap, xorOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          xorOp.getResult(),
          LoweredValue {xorOp.getType(), {builder.create<arith::XOrIOp>(loc, *lhs, *rhs)}}
      );
    }
    if (auto notOp = dyn_cast<boolean::NotBoolOp>(op)) {
      auto operand = lookupScalar(notOp.getOperand(), valueMap, notOp.getOperation());
      if (failed(operand)) {
        return failure();
      }
      Value one = builder.create<arith::ConstantOp>(
          loc, IntegerAttr::get(IntegerType::get(builder.getContext(), 1), 1)
      );
      return bind(
          notOp.getResult(),
          LoweredValue {notOp.getType(), {builder.create<arith::XOrIOp>(loc, *operand, one)}}
      );
    }
 
    if (auto intToFelt = dyn_cast<cast::IntToFeltOp>(op)) {
      auto operand = lookupScalar(intToFelt.getValue(), valueMap, intToFelt.getOperation());
      if (failed(operand)) {
        return failure();
      }
      auto dstType = IntegerType::get(builder.getContext(), field.bitWidth());
      Value lowered;
      if (isa<IndexType>((*operand).getType())) {
        lowered = builder.create<arith::IndexCastUIOp>(loc, dstType, *operand);
      } else {
        auto intType = llvm::cast<IntegerType>((*operand).getType());
        if (intType.getWidth() < dstType.getWidth()) {
          lowered = builder.create<arith::ExtUIOp>(loc, dstType, *operand);
        } else if (intType.getWidth() > dstType.getWidth()) {
          lowered = normalizeWideValue(builder, loc, *operand, dstType.getWidth(), field);
        } else {
          lowered = *operand;
        }
      }
      return bind(intToFelt.getResult(), LoweredValue {intToFelt.getType(), {lowered}});
    }
    if (auto feltToIndex = dyn_cast<cast::FeltToIndexOp>(op)) {
      auto operand = lookupScalar(feltToIndex.getValue(), valueMap, feltToIndex.getOperation());
      if (failed(operand)) {
        return failure();
      }
      return bind(
          feltToIndex.getResult(),
          LoweredValue {
              feltToIndex.getType(),
              {builder.create<arith::IndexCastUIOp>(loc, builder.getIndexType(), *operand)}
          }
      );
    }
 
    if (auto structNewOp = dyn_cast<component::CreateStructOp>(op)) {
      auto lowered = createDefaultValue(
          builder, loc, structNewOp.getType(), tables, structNewOp.getOperation(), field,
          uninitializedBehavior, rng
      );
      if (failed(lowered)) {
        return failure();
      }
      return bind(structNewOp.getResult(), std::move(*lowered));
    }
    if (auto readMemberOp = dyn_cast<component::MemberReadOp>(op)) {
      auto componentValue =
          lookup(readMemberOp.getComponent(), valueMap, readMemberOp.getOperation());
      if (failed(componentValue)) {
        return failure();
      }
      auto lowered = readNamedAggregateValue(
          builder, loc, readMemberOp.getComponent().getType(), readMemberOp.getMemberName(),
          *componentValue, tables, readMemberOp.getOperation(), field
      );
      if (failed(lowered)) {
        return failure();
      }
      return bind(readMemberOp.getResult(), std::move(*lowered));
    }
    if (auto writeMemberOp = dyn_cast<component::MemberWriteOp>(op)) {
      auto componentValue =
          lookup(writeMemberOp.getComponent(), valueMap, writeMemberOp.getOperation());
      auto memberValue = lookup(writeMemberOp.getVal(), valueMap, writeMemberOp.getOperation());
      if (failed(componentValue) || failed(memberValue)) {
        return failure();
      }
      return writeNamedAggregateValue(
          builder, loc, writeMemberOp.getComponent().getType(), writeMemberOp.getMemberName(),
          valueMap[writeMemberOp.getComponent()], *memberValue, tables,
          writeMemberOp.getOperation(), field
      );
    }
 
    if (auto newPodOp = dyn_cast<pod::NewPodOp>(op)) {
      auto lowered = createDefaultValue(
          builder, loc, newPodOp.getType(), tables, newPodOp.getOperation(), field,
          uninitializedBehavior, rng
      );
      if (failed(lowered)) {
        return failure();
      }
      for (pod::RecordValue init : newPodOp.getInitializedRecordValues()) {
        auto value = lookup(init.value, valueMap, newPodOp.getOperation());
        if (failed(value) || failed(writeNamedAggregateValue(
                                 builder, loc, newPodOp.getType(), init.name, *lowered, *value,
                                 tables, newPodOp.getOperation(), field
                             ))) {
          return failure();
        }
      }
      return bind(newPodOp.getResult(), std::move(*lowered));
    }
    if (auto readPodOp = dyn_cast<pod::ReadPodOp>(op)) {
      auto podValue = lookup(readPodOp.getPodRef(), valueMap, readPodOp.getOperation());
      if (failed(podValue)) {
        return failure();
      }
      auto lowered = readNamedAggregateValue(
          builder, loc, readPodOp.getPodRef().getType(), readPodOp.getRecordName(), *podValue,
          tables, readPodOp.getOperation(), field
      );
      if (failed(lowered)) {
        return failure();
      }
      return bind(readPodOp.getResult(), std::move(*lowered));
    }
    if (auto writePodOp = dyn_cast<pod::WritePodOp>(op)) {
      auto recordValue = lookup(writePodOp.getValue(), valueMap, writePodOp.getOperation());
      if (failed(recordValue)) {
        return failure();
      }
      return writeNamedAggregateValue(
          builder, loc, writePodOp.getPodRef().getType(), writePodOp.getRecordName(),
          valueMap[writePodOp.getPodRef()], *recordValue, tables, writePodOp.getOperation(), field
      );
    }
 
    if (auto arrayNewOp = dyn_cast<array::CreateArrayOp>(op)) {
      auto lowered = createDefaultValue(
          builder, loc, arrayNewOp.getType(), tables, arrayNewOp.getOperation(), field,
          uninitializedBehavior, rng
      );
      if (failed(lowered)) {
        return failure();
      }
      if (!arrayNewOp.getElements().empty()) {
        auto elementCount = checkedCast<size_t>(arrayNewOp.getType().getNumElements());
        if (!elementCount) {
          arrayNewOp.emitError() << llvm::toString(elementCount.takeError());
          return failure();
        }
        if (arrayNewOp.getElements().size() != *elementCount) {
          arrayNewOp.emitError("expected one explicit element per array slot in witgen lowering");
          return failure();
        }
        auto shape = arrayNewOp.getType().getShape();
        for (auto [flatIndex, operand] : llvm::enumerate(arrayNewOp.getElements())) {
          auto elementValue = lookup(operand, valueMap, arrayNewOp.getOperation());
          if (failed(elementValue)) {
            return failure();
          }
          SmallVector<Value> indices;
          auto strides = mlir::computeStrides(shape);
          auto flatSigned = checkedCast<int64_t>(flatIndex);
          if (!flatSigned) {
            arrayNewOp.emitError() << llvm::toString(flatSigned.takeError());
            return failure();
          }
          for (int64_t index : mlir::delinearize(*flatSigned, strides)) {
            indices.push_back(makeIndexConstant(builder, loc, index));
          }
          if (failed(writeArrayElement(
                  builder, loc, arrayNewOp.getType(), *lowered, indices, *elementValue
              ))) {
            return failure();
          }
        }
      }
      return bind(arrayNewOp.getResult(), std::move(*lowered));
    }
    if (auto readArrayOp = dyn_cast<array::ReadArrayOp>(op)) {
      SmallVector<Value> indices;
      for (Value indexValue : readArrayOp.getIndices()) {
        auto loweredIndex = lookupScalar(indexValue, valueMap, readArrayOp.getOperation());
        if (failed(loweredIndex)) {
          return failure();
        }
        indices.push_back(*loweredIndex);
      }
      auto arrayValue = lookup(readArrayOp.getArrRef(), valueMap, readArrayOp.getOperation());
      if (failed(arrayValue)) {
        return failure();
      }
      auto lowered = readArrayElement(
          builder, loc, llvm::cast<array::ArrayType>(readArrayOp.getArrRef().getType()),
          *arrayValue, indices
      );
      if (failed(lowered)) {
        return failure();
      }
      return bind(readArrayOp.getResult(), std::move(*lowered));
    }
    if (auto writeArrayOp = dyn_cast<array::WriteArrayOp>(op)) {
      SmallVector<Value> indices;
      for (Value indexValue : writeArrayOp.getIndices()) {
        auto loweredIndex = lookupScalar(indexValue, valueMap, writeArrayOp.getOperation());
        if (failed(loweredIndex)) {
          return failure();
        }
        indices.push_back(*loweredIndex);
      }
      auto elementValue = lookup(writeArrayOp.getRvalue(), valueMap, writeArrayOp.getOperation());
      if (failed(elementValue)) {
        return failure();
      }
      return writeArrayElement(
          builder, loc, llvm::cast<array::ArrayType>(writeArrayOp.getArrRef().getType()),
          valueMap[writeArrayOp.getArrRef()], indices, *elementValue
      );
    }
 
    if (auto cmpiOp = dyn_cast<arith::CmpIOp>(op)) {
      auto lhs = lookupScalar(cmpiOp.getLhs(), valueMap, cmpiOp.getOperation());
      auto rhs = lookupScalar(cmpiOp.getRhs(), valueMap, cmpiOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          cmpiOp.getResult(),
          LoweredValue {
              cmpiOp.getType(),
              {builder.create<arith::CmpIOp>(loc, cmpiOp.getPredicate(), *lhs, *rhs)}
          }
      );
    }
    if (auto selectOp = dyn_cast<arith::SelectOp>(op)) {
      auto cond = lookupScalar(selectOp.getCondition(), valueMap, selectOp.getOperation());
      auto trueValue = lookupScalar(selectOp.getTrueValue(), valueMap, selectOp.getOperation());
      auto falseValue = lookupScalar(selectOp.getFalseValue(), valueMap, selectOp.getOperation());
      if (failed(cond) || failed(trueValue) || failed(falseValue)) {
        return failure();
      }
      return bind(
          selectOp.getResult(),
          LoweredValue {
              selectOp.getType(),
              {builder.create<arith::SelectOp>(loc, *cond, *trueValue, *falseValue)}
          }
      );
    }
    if (auto addiOp = dyn_cast<arith::AddIOp>(op)) {
      auto lhs = lookupScalar(addiOp.getLhs(), valueMap, addiOp.getOperation());
      auto rhs = lookupScalar(addiOp.getRhs(), valueMap, addiOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          addiOp.getResult(),
          LoweredValue {addiOp.getType(), {builder.create<arith::AddIOp>(loc, *lhs, *rhs)}}
      );
    }
    if (auto subiOp = dyn_cast<arith::SubIOp>(op)) {
      auto lhs = lookupScalar(subiOp.getLhs(), valueMap, subiOp.getOperation());
      auto rhs = lookupScalar(subiOp.getRhs(), valueMap, subiOp.getOperation());
      if (failed(lhs) || failed(rhs)) {
        return failure();
      }
      return bind(
          subiOp.getResult(),
          LoweredValue {subiOp.getType(), {builder.create<arith::SubIOp>(loc, *lhs, *rhs)}}
      );
    }
 
    if (auto callOp = dyn_cast<function::CallOp>(op)) {
      if (callOp.getTemplateParams() || !callOp.getMapOperands().empty()) {
        callOp.emitError("execution-engine backend encountered an unflattened function.call");
        return failure();
      }
      auto *callable = callOp.resolveCallableInTable(&tables);
      auto callee = dyn_cast_or_null<function::FuncDefOp>(callable);
      if (!callee) {
        callOp.emitError("failed to resolve callee during execution-engine lowering");
        return failure();
      }
      SmallVector<Type> resultTypes;
      for (Type resultType : callOp.getResultTypes()) {
        if (failed(
                flattenABILeafTypes(resultType, tables, callOp.getOperation(), field, resultTypes)
            )) {
          return failure();
        }
      }
      SmallVector<Value> flatArgs;
      for (Value operand : callOp.getArgOperands()) {
        auto lowered = lookup(operand, valueMap, callOp.getOperation());
        auto leafTypes = getABILeafTypes(operand.getType(), tables, callOp.getOperation(), field);
        if (failed(lowered) || failed(leafTypes) ||
            failed(appendFlatLeavesToTypes(
                builder, loc, *lowered, *leafTypes, flatArgs, callOp.getOperation()
            ))) {
          return failure();
        }
      }
      auto loweredCall =
          builder.create<func::CallOp>(loc, mangleFunctionName(callee), resultTypes, flatArgs);
      auto loweredCallResults = loweredCall.getResults();
      size_t totalResults = loweredCallResults.size();
      size_t cursor = 0;
      for (auto [oldResult, oldType] : llvm::zip(callOp.getResults(), callOp.getResultTypes())) {
        auto leafCount = getLeafCount(oldType, tables, callOp.getOperation(), field);
        if (failed(leafCount)) {
          return failure();
        }
        bool overflow = false;
        size_t nextCursor = llvm::SaturatingAdd(cursor, *leafCount, &overflow);
        if (overflow || nextCursor > totalResults) {
          callOp.emitError("leaf count overflow while lowering function call results");
          return failure();
        }
        LoweredValue lowered {oldType, {}};
        lowered.leaves.append(
            loweredCallResults.begin() + static_cast<ptrdiff_t>(cursor),
            loweredCallResults.begin() + static_cast<ptrdiff_t>(nextCursor)
        );
        valueMap[oldResult] = std::move(lowered);
        cursor = nextCursor;
      }
      return success();
    }
 
    if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
      SmallVector<Value> initArgs;
      SmallVector<size_t> beforeLeafCounts;
      for (auto [init, initType] : llvm::zip(whileOp.getInits(), whileOp.getOperandTypes())) {
        auto lowered = lookup(init, valueMap, whileOp.getOperation());
        auto leafTypes = getABILeafTypes(initType, tables, whileOp.getOperation(), field);
        if (failed(lowered) || failed(leafTypes) ||
            failed(appendFlatLeavesToTypes(
                builder, loc, *lowered, *leafTypes, initArgs, whileOp.getOperation()
            ))) {
          return failure();
        }
        auto count = getLeafCount(initType, tables, whileOp.getOperation(), field);
        if (failed(count)) {
          return failure();
        }
        beforeLeafCounts.push_back(*count);
      }
 
      SmallVector<size_t> resultLeafCounts;
      SmallVector<Type> loweredResultTypes;
      for (Type resultType : whileOp.getResultTypes()) {
        auto leafTypes = getABILeafTypes(resultType, tables, whileOp.getOperation(), field);
        auto count = getLeafCount(resultType, tables, whileOp.getOperation(), field);
        if (failed(leafTypes) || failed(count)) {
          return failure();
        }
        loweredResultTypes.append(leafTypes->begin(), leafTypes->end());
        resultLeafCounts.push_back(*count);
      }
 
      auto mapRegionArguments = [&](auto oldArgs, auto oldTypes, auto leafCounts, auto newArgs,
                                    StringRef overflowMessage,
                                    DenseMap<Value, LoweredValue> &regionMap) -> LogicalResult {
        size_t totalArgs = newArgs.size();
        size_t cursor = 0;
        for (auto [oldArg, oldType, leafCount] : llvm::zip(oldArgs, oldTypes, leafCounts)) {
          bool overflow = false;
          size_t nextCursor = llvm::SaturatingAdd(cursor, leafCount, &overflow);
          if (overflow || nextCursor > totalArgs) {
            whileOp.emitError(overflowMessage);
            return failure();
          }
          LoweredValue lowered {oldType, {}};
          lowered.leaves.append(
              newArgs.begin() + llzk::checkedCast<ptrdiff_t>(cursor),
              newArgs.begin() + llzk::checkedCast<ptrdiff_t>(nextCursor)
          );
          regionMap[oldArg] = std::move(lowered);
          cursor = nextCursor;
        }
        return success();
      };
 
      LogicalResult whileLoweringStatus = success();
      auto newWhile = builder.create<scf::WhileOp>(
          loc, loweredResultTypes, initArgs,
          [&](OpBuilder &regionBuilder, Location /*regionLoc*/, ValueRange beforeArgs) {
        DenseMap<Value, LoweredValue> beforeMap(valueMap.begin(), valueMap.end());
        if (failed(mapRegionArguments(
                whileOp.getBeforeArguments(), whileOp.getOperandTypes(), beforeLeafCounts,
                beforeArgs, "leaf count overflow while lowering while-loop before-region args",
                beforeMap
            )) ||
            failed(lowerBlock(regionBuilder, whileOp.getBefore().front(), beforeMap))) {
          whileLoweringStatus = failure();
        }
      }, [&](OpBuilder &regionBuilder, Location /*regionLoc*/, ValueRange afterArgs) {
        DenseMap<Value, LoweredValue> afterMap(valueMap.begin(), valueMap.end());
        if (failed(mapRegionArguments(
                whileOp.getAfterArguments(), whileOp.getResultTypes(), resultLeafCounts, afterArgs,
                "leaf count overflow while lowering while-loop after-region args", afterMap
            )) ||
            failed(lowerBlock(regionBuilder, whileOp.getAfter().front(), afterMap))) {
          whileLoweringStatus = failure();
        }
      }
      );
      if (failed(whileLoweringStatus)) {
        newWhile.erase();
        return failure();
      }
 
      auto newWhileResults = newWhile.getResults();
      size_t totalResults = newWhileResults.size();
      size_t cursor = 0;
      for (auto [oldResult, oldType, leafCount] :
           llvm::zip(whileOp.getResults(), whileOp.getResultTypes(), resultLeafCounts)) {
        bool overflow = false;
        size_t nextCursor = llvm::SaturatingAdd(cursor, leafCount, &overflow);
        if (overflow || nextCursor > totalResults) {
          whileOp.emitError("leaf count overflow while lowering while-loop results");
          return failure();
        }
        LoweredValue lowered {oldType, {}};
        lowered.leaves.append(
            newWhileResults.begin() + llzk::checkedCast<ptrdiff_t>(cursor),
            newWhileResults.begin() + llzk::checkedCast<ptrdiff_t>(nextCursor)
        );
        valueMap[oldResult] = std::move(lowered);
        cursor = nextCursor;
      }
      return success();
    }
 
    if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
      auto condition = lookupScalar(ifOp.getCondition(), valueMap, ifOp.getOperation());
      if (failed(condition)) {
        return failure();
      }
 
      SmallVector<size_t> resultLeafCounts;
      SmallVector<Type> loweredResultTypes;
      for (Type resultType : ifOp.getResultTypes()) {
        auto leafTypes = getABILeafTypes(resultType, tables, ifOp.getOperation(), field);
        auto count = getLeafCount(resultType, tables, ifOp.getOperation(), field);
        if (failed(leafTypes) || failed(count)) {
          return failure();
        }
        loweredResultTypes.append(leafTypes->begin(), leafTypes->end());
        resultLeafCounts.push_back(*count);
      }
 
      auto newIf = builder.create<scf::IfOp>(
          loc, loweredResultTypes, *condition, true, !ifOp.getElseRegion().empty()
      );
 
      {
        OpBuilder thenBuilder = OpBuilder::atBlockBegin(&newIf.getThenRegion().front());
        DenseMap<Value, LoweredValue> thenMap(valueMap.begin(), valueMap.end());
        if (failed(lowerBlock(thenBuilder, ifOp.getThenRegion().front(), thenMap))) {
          newIf.erase();
          return failure();
        }
      }
      if (!ifOp.getElseRegion().empty()) {
        OpBuilder elseBuilder = OpBuilder::atBlockBegin(&newIf.getElseRegion().front());
        DenseMap<Value, LoweredValue> elseMap(valueMap.begin(), valueMap.end());
        if (failed(lowerBlock(elseBuilder, ifOp.getElseRegion().front(), elseMap))) {
          newIf.erase();
          return failure();
        }
      }
      auto newIfResults = newIf.getResults();
      size_t totalResults = newIfResults.size();
      size_t cursor = 0;
      for (auto [oldResult, oldType, leafCount] :
           llvm::zip(ifOp.getResults(), ifOp.getResultTypes(), resultLeafCounts)) {
        bool overflow = false;
        size_t nextCursor = llvm::SaturatingAdd(cursor, leafCount, &overflow);
        if (overflow || nextCursor > totalResults) {
          ifOp.emitError("leaf count overflow while lowering if-op results");
          return failure();
        }
        LoweredValue lowered {oldType, {}};
        lowered.leaves.append(
            newIfResults.begin() + llzk::checkedCast<ptrdiff_t>(cursor),
            newIfResults.begin() + llzk::checkedCast<ptrdiff_t>(nextCursor)
        );
        valueMap[oldResult] = std::move(lowered);
        cursor = nextCursor;
      }
      return success();
    }
 
    if (auto forOp = dyn_cast<scf::ForOp>(op)) {
      auto lb = lookupScalar(forOp.getLowerBound(), valueMap, forOp.getOperation());
      auto ub = lookupScalar(forOp.getUpperBound(), valueMap, forOp.getOperation());
      auto step = lookupScalar(forOp.getStep(), valueMap, forOp.getOperation());
      if (failed(lb) || failed(ub) || failed(step)) {
        return failure();
      }
 
      SmallVector<Value> initArgs;
      SmallVector<size_t> initLeafCounts;
      for (auto [init, resultType] : llvm::zip(forOp.getInitArgs(), forOp.getResultTypes())) {
        auto lowered = lookup(init, valueMap, forOp.getOperation());
        auto leafTypes = getABILeafTypes(resultType, tables, forOp.getOperation(), field);
        if (failed(lowered) || failed(leafTypes) ||
            failed(appendFlatLeavesToTypes(
                builder, loc, *lowered, *leafTypes, initArgs, forOp.getOperation()
            ))) {
          return failure();
        }
        auto count = getLeafCount(resultType, tables, forOp.getOperation(), field);
        if (failed(count)) {
          return failure();
        }
        initLeafCounts.push_back(*count);
      }
 
      auto newFor = builder.create<scf::ForOp>(loc, *lb, *ub, *step, initArgs);
      if (Attribute unsignedCmpAttr = forOp->getAttr("unsignedCmp")) {
        newFor->setAttr("unsignedCmp", unsignedCmpAttr);
      }
      DenseMap<Value, LoweredValue> bodyMap(valueMap.begin(), valueMap.end());
      bodyMap[forOp.getInductionVar()] =
          LoweredValue {forOp.getInductionVar().getType(), {newFor.getInductionVar()}};
      {
        auto newForIterArgs = newFor.getRegionIterArgs();
        size_t totalIterArgs = newForIterArgs.size();
        size_t cursor = 0;
        for (auto [oldIterArg, oldType, leafCount] :
             llvm::zip(forOp.getRegionIterArgs(), forOp.getResultTypes(), initLeafCounts)) {
          bool overflow = false;
          size_t nextCursor = llvm::SaturatingAdd(cursor, leafCount, &overflow);
          if (overflow || nextCursor > totalIterArgs) {
            forOp.emitError("leaf count overflow while lowering for-loop region iter args");
            return failure();
          }
          LoweredValue lowered {oldType, {}};
          lowered.leaves.append(
              newForIterArgs.begin() + static_cast<ptrdiff_t>(cursor),
              newForIterArgs.begin() + static_cast<ptrdiff_t>(nextCursor)
          );
          bodyMap[oldIterArg] = std::move(lowered);
          cursor = nextCursor;
        }
      }
 
      newFor.getBody()->clear();
      OpBuilder bodyBuilder = OpBuilder::atBlockBegin(newFor.getBody());
      if (failed(lowerBlock(bodyBuilder, *forOp.getBody(), bodyMap))) {
        return failure();
      }
 
      {
        auto newForResults = newFor.getResults();
        size_t totalForResults = newForResults.size();
        size_t cursor = 0;
        for (auto [oldResult, oldType, leafCount] :
             llvm::zip(forOp.getResults(), forOp.getResultTypes(), initLeafCounts)) {
          bool overflow = false;
          size_t nextCursor = llvm::SaturatingAdd(cursor, leafCount, &overflow);
          if (overflow || nextCursor > totalForResults) {
            forOp.emitError("leaf count overflow while lowering for-loop results");
            return failure();
          }
          LoweredValue lowered {oldType, {}};
          lowered.leaves.append(
              newForResults.begin() + static_cast<ptrdiff_t>(cursor),
              newForResults.begin() + static_cast<ptrdiff_t>(nextCursor)
          );
          valueMap[oldResult] = std::move(lowered);
          cursor = nextCursor;
        }
      }
      return success();
    }
 
    op.emitError("unsupported operation in execution-engine lowering: ") << op.getName();
    return failure();
  }
};

class LowerComputeToCorePass : public PassWrapper<LowerComputeToCorePass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerComputeToCorePass)
 
  explicit LowerComputeToCorePass(const WitgenOptions &opts) : options(opts) {}

  StringRef getArgument() const final { return "llzk-lower-compute-to-core"; }

  StringRef getDescription() const final {
    return "Lower LLZK compute IR to func/arith/cf/scf/memref";
  }

  StringRef getName() const override { return "LowerComputeToCorePass"; }

  void runOnOperation() override {
    ModuleOp moduleOp = getOperation();
    auto field = getModuleField(moduleOp);
    if (failed(field)) {
      signalPassFailure();
      return;
    }
 
    SymbolTableCollection tables;
    BodyLowerer lowerer(moduleOp, tables, field->get(), options);
    auto funcs = walkCollect<function::FuncDefOp>(moduleOp, [](auto funcOp) {
      return !funcOp.nameIsConstrain();
    });
    for (function::FuncDefOp funcOp : funcs) {
      if (failed(lowerer.lowerFunction(funcOp))) {
        signalPassFailure();
        return;
      }
    }
  }
 
private:
  WitgenOptions options;
};

class CreateWitgenEntryPass : public PassWrapper<CreateWitgenEntryPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CreateWitgenEntryPass)
 
  
  explicit CreateWitgenEntryPass(OutputScope newOutputScope) : outputScope(newOutputScope) {}

  StringRef getArgument() const final { return "llzk-create-witgen-entry"; }

  StringRef getDescription() const final {
    return "Create the llzk-witgen execution-engine entry wrapper";
  }

  StringRef getName() const override { return "CreateWitgenEntryPass"; }

  void runOnOperation() override {
    ModuleOp moduleOp = getOperation();
    auto field = getModuleField(moduleOp);
    if (failed(field)) {
      signalPassFailure();
      return;
    }
 
    SymbolTableCollection tables;
    auto mainDef = getMainInstanceDef(tables, moduleOp.getOperation());
    if (failed(mainDef) || !mainDef.value()) {
      moduleOp.emitError("module is missing a concrete llzk.main struct");
      signalPassFailure();
      return;
    }
    function::FuncDefOp computeFunc = mainDef->get().getComputeFuncOp();
    if (!computeFunc) {
      moduleOp.emitError("main struct is missing @compute");
      signalPassFailure();
      return;
    }
 
    auto outputs =
        collectOutputBindings(mainDef->get(), tables, computeFunc.getOperation(), outputScope);
    if (failed(outputs)) {
      signalPassFailure();
      return;
    }
 
    OpBuilder builder(moduleOp.getContext());
    builder.setInsertionPointToEnd(moduleOp.getBody());
 
    SmallVector<Type> wrapperArgs;
    for (Type argType : computeFunc.getArgumentTypes()) {
      SmallVector<Type> loweredLeafTypes;
      if (failed(flattenTypeLeaves(
              argType, tables, computeFunc.getOperation(), field->get(), loweredLeafTypes, {}, true
          ))) {
        signalPassFailure();
        return;
      }
      if (loweredLeafTypes.size() != 1 || !isa<MemRefType>(loweredLeafTypes.front())) {
        computeFunc.emitError(
            "execution-engine wrapper only supports felt and array<...xfelt> inputs"
        );
        signalPassFailure();
        return;
      }
      wrapperArgs.push_back(loweredLeafTypes.front());
    }
    for (const OutputBinding &output : *outputs) {
      SmallVector<Type> loweredLeafTypes;
      if (failed(flattenTypeLeaves(
              output.type, tables, computeFunc.getOperation(), field->get(), loweredLeafTypes, {},
              true
          ))) {
        signalPassFailure();
        return;
      }
      if (loweredLeafTypes.size() != 1 || !isa<MemRefType>(loweredLeafTypes.front())) {
        computeFunc.emitError(
            "execution-engine wrapper only supports felt and array<...xfelt> outputs"
        );
        signalPassFailure();
        return;
      }
      wrapperArgs.push_back(loweredLeafTypes.front());
    }
 
    auto wrapper = builder.create<func::FuncOp>(
        computeFunc.getLoc(), "__llzk_witgen_main",
        builder.getFunctionType(wrapperArgs, TypeRange {})
    );
    wrapper->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(), builder.getUnitAttr());
    Block *entry = wrapper.addEntryBlock();
    builder.setInsertionPointToStart(entry);
 
    SmallVector<Type> loweredMainResultTypes;
    for (Type resultType : computeFunc.getResultTypes()) {
      if (failed(flattenABILeafTypes(
              resultType, tables, computeFunc.getOperation(), field->get(), loweredMainResultTypes
          ))) {
        signalPassFailure();
        return;
      }
    }
 
    SmallVector<Value> mainArgs;
    for (auto [argType, wrapperArg] : llvm::zip(
             computeFunc.getArgumentTypes(),
             entry->getArguments().take_front(computeFunc.getNumArguments())
         )) {
      if (isScalarType(argType)) {
        mainArgs.push_back(loadStorageScalar(builder, computeFunc.getLoc(), wrapperArg));
      } else {
        auto abiLeafTypes =
            getABILeafTypes(argType, tables, computeFunc.getOperation(), field->get());
        if (failed(abiLeafTypes) || abiLeafTypes->size() != 1 ||
            !isa<MemRefType>(abiLeafTypes->front())) {
          computeFunc.emitError("failed to derive execution-engine ABI type for main input");
          signalPassFailure();
          return;
        }
        if (wrapperArg.getType() == abiLeafTypes->front()) {
          mainArgs.push_back(wrapperArg);
        } else {
          mainArgs.push_back(builder.create<memref::CastOp>(
              computeFunc.getLoc(), abiLeafTypes->front(), wrapperArg
          ));
        }
      }
    }
    auto loweredMain = builder.create<func::CallOp>(
        computeFunc.getLoc(), mangleFunctionName(computeFunc), loweredMainResultTypes, mainArgs
    );
 
    LoweredValue mainResultValue {
        computeFunc.getResultTypes().front(),
        llvm::SmallVector<Value>(loweredMain.getResults().begin(), loweredMain.getResults().end())
    };
 
    auto extractOutputSlice = [&](ArrayRef<std::string> path, Type currentType,
                                  ArrayRef<Value> leaves,
                                  auto &self) -> FailureOr<SmallVector<Value>> {
      if (path.empty()) {
        return SmallVector<Value>(leaves.begin(), leaves.end());
      }
      if (auto structType = dyn_cast<component::StructType>(currentType)) {
        auto defLookup = structType.getDefinition(tables, computeFunc.getOperation());
        if (failed(defLookup)) {
          return failure();
        }
        unsigned localCursor = 0;
        for (component::MemberDefOp member : defLookup->get().getMemberDefs()) {
          auto leafCount =
              getLeafCount(member.getType(), tables, member.getOperation(), field->get());
          if (failed(leafCount)) {
            return failure();
          }
          ArrayRef<Value> slice = ArrayRef<Value>(leaves).slice(localCursor, *leafCount);
          localCursor += *leafCount;
          if (member.getSymName() == path.front()) {
            return self(path.drop_front(), member.getType(), slice, self);
          }
        }
        computeFunc.emitError("failed to find struct member while wiring witgen outputs");
        return failure();
      }
      if (auto podType = dyn_cast<pod::PodType>(currentType)) {
        unsigned localCursor = 0;
        for (pod::RecordAttr record : podType.getRecords()) {
          auto leafCount =
              getLeafCount(record.getType(), tables, computeFunc.getOperation(), field->get());
          if (failed(leafCount)) {
            return failure();
          }
          ArrayRef<Value> slice = ArrayRef<Value>(leaves).slice(localCursor, *leafCount);
          localCursor += *leafCount;
          if (record.getName().getValue() == path.front()) {
            return self(path.drop_front(), record.getType(), slice, self);
          }
        }
        computeFunc.emitError("failed to find POD record while wiring witgen outputs");
        return failure();
      }
      computeFunc.emitError("extra witness path components for non-aggregate output");
      return failure();
    };
 
    auto outputArgs = entry->getArguments().drop_front(computeFunc.getNumArguments());
    for (auto [output, outputMemRef] : llvm::zip(*outputs, outputArgs)) {
      auto slice = extractOutputSlice(
          output.path, mainResultValue.sourceType, mainResultValue.leaves, extractOutputSlice
      );
      if (failed(slice) || slice->empty()) {
        wrapper.emitError("missing selected witness output slice while building witgen entry");
        signalPassFailure();
        return;
      }
      if (isScalarType(output.type)) {
        storeStorageScalar(
            builder, computeFunc.getLoc(),
            loadStorageScalar(builder, computeFunc.getLoc(), slice->front()), outputMemRef
        );
      } else {
        builder.create<memref::CopyOp>(computeFunc.getLoc(), slice->front(), outputMemRef);
      }
    }
    builder.create<func::ReturnOp>(computeFunc.getLoc());
 
    // Remove `llzk.main` attribute because the main struct is deleted below.
    moduleOp->removeAttr(MAIN_ATTR_NAME);
 
    SmallVector<Operation *> toErase;
    for (Operation &op : moduleOp.getBody()->getOperations()) {
      if (!isa<func::FuncOp>(op)) {
        toErase.push_back(&op);
      }
    }
    for (Operation *op : toErase) {
      op->erase();
    }
  }
 
private:
  OutputScope outputScope;
};
 
} // namespace
 
void addWitgenPreparePipeline(OpPassManager &pm, const WitgenOptions &) {
  using namespace llzk::polymorphic;
  pm.addPass(createFlatteningPass(
      FlatteningPassOptions {.cleanupMode = FlatteningCleanupMode::ConcreteAsRoot}
  ));
  pm.addPass(mlir::createLowerAffinePass());
  // TODO: simplify lowering with `llzk-inline-structs` and `llzk-pod-to-scalar` when both are
  // available and support PODs.
  pm.addPass(mlir::createCanonicalizerPass());
  pm.addPass(mlir::createCSEPass());
}
 
std::unique_ptr<Pass> createLowerComputeToCorePass(const WitgenOptions &options) {
  return std::make_unique<LowerComputeToCorePass>(options);
}
 
std::unique_ptr<Pass> createCreateWitgenEntryPass(OutputScope outputScope) {
  return std::make_unique<CreateWitgenEntryPass>(outputScope);
}
 
} // namespace llzk::witgen