Validate poly lowering split fixes

backends/r1cs/lib/r1cs/Transforms/R1CSLoweringPass.cpp

@@ -340,7 +340,7 @@ class PassImpl : public r1cs::impl::R1CSLoweringPassBase<PassImpl> {
         if (degLhs == 2 && degRhs == 2) {
           builder.setInsertionPoint(op);
           std::string auxName = R1CS_AUXILIARY_MEMBER_PREFIX + std::to_string(auxCounter++);
-          MemberDefOp auxMember = addAuxMember(structDef, auxName);
+          MemberDefOp auxMember = addAuxMember(structDef, auxName, val.getType());
           Value aux = builder.create<MemberReadOp>(
               val.getLoc(), val.getType(), constrainFunc.getSelfValueFromConstrain(),
               auxMember.getNameAttr()
@@ -532,7 +532,7 @@ class PassImpl : public r1cs::impl::R1CSLoweringPassBase<PassImpl> {
       // Entire linear combination was zero
       result = builder.create<r1cs::ConstOp>(
           loc, r1cs::LinearType::get(builder.getContext()),
-          r1cs::FeltAttr::get(builder.getContext(), 0)
+          r1cs::FeltAttr::get(builder.getContext(), toAPSInt(lc.constant))
       );
     }
 
@@ -673,7 +673,7 @@ class PassImpl : public r1cs::impl::R1CSLoweringPassBase<PassImpl> {
         if (degLhs == 2 && degRhs == 2) {
           builder.setInsertionPoint(eqOp);
           std::string auxName = R1CS_AUXILIARY_MEMBER_PREFIX + std::to_string(auxCounter++);
-          MemberDefOp auxMember = addAuxMember(structDef, auxName);
+          MemberDefOp auxMember = addAuxMember(structDef, auxName, lhs.getType());
           Value aux = builder.create<MemberReadOp>(
               eqOp.getLoc(), lhs.getType(), constrainFunc.getSelfValueFromConstrain(),
               auxMember.getNameAttr()

changelogs/unreleased/poly-aux-write-order.yaml

@@ -0,0 +1,2 @@
+fixed:
+  - Topologically order polynomial-lowering auxiliary writes in compute functions

changelogs/unreleased/poly-compute-straight-line.yaml

@@ -0,0 +1,2 @@
+fixed:
+  - Reject non-straight-line compute bodies before polynomial lowering inserts auxiliary writes

changelogs/unreleased/poly-member-read-access.yaml

@@ -0,0 +1,2 @@
+fixed:
+  - Preserve member-read component operands, access attributes, and index map operands when polynomial lowering rebuilds auxiliary producers

changelogs/unreleased/poly-rewrite-dominance.yaml

@@ -0,0 +1,2 @@
+fixed:
+  - Avoid reusing polynomial-lowering replacement values at uses they do not dominate

changelogs/unreleased/poly-typed-aux-members.yaml

@@ -0,0 +1,4 @@
+fixed:
+  - Emit polynomial-lowering auxiliary members with the exact type of the materialized expression
+  - Emit all R1CS-lowering auxiliary members with the exact type of the materialized expression
+  - Emit synthesized zero R1CS linear-combination constants with a printable integer width

include/llzk/Transforms/LLZKLoweringUtils.h

@@ -12,6 +12,7 @@
 
 #include <mlir/IR/Builders.h>
 #include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/Types.h>
 #include <mlir/IR/Value.h>
 #include <mlir/Support/LogicalResult.h>
 
@@ -32,13 +33,22 @@ mlir::Value rebuildExprInCompute(
 mlir::LogicalResult
 checkForAuxMemberConflicts(component::StructDefOp structDef, llvm::StringRef auxPrefix);
 
-/// Rejects control flow under `constrainFunc`; polynomial and R1CS auxiliary
-/// materialization assumes control flow has already been flattened or otherwise
-/// lowered away. The region check catches multi-block function bodies before
-/// the operation walk rejects nested regions or successor-bearing operations.
+/// Rejects control flow under `func`; auxiliary materialization assumes control
+/// flow has already been flattened or otherwise lowered away. The region check
+/// catches multi-block function bodies before the operation walk rejects nested
+/// regions or successor-bearing operations.
+mlir::LogicalResult checkFuncBodyIsStraightLine(
+    function::FuncDefOp func, llvm::StringRef passName, llvm::StringRef funcName
+);
+
+/// Rejects control flow under `constrainFunc`; kept as a convenience wrapper
+/// for passes whose public diagnostic already names constrain bodies.
 mlir::LogicalResult
 checkConstrainBodyIsStraightLine(function::FuncDefOp constrainFunc, llvm::StringRef passName);
 
+component::MemberDefOp
+addAuxMember(component::StructDefOp structDef, llvm::StringRef name, mlir::Type type);
+
 component::MemberDefOp addAuxMember(component::StructDefOp structDef, llvm::StringRef name);
 
 unsigned getFeltDegree(mlir::Value val, llvm::DenseMap<mlir::Value, unsigned> &memo);

include/llzk/Transforms/LLZKTransformationPasses.td

@@ -74,10 +74,10 @@ def PolyLoweringPass : LLZKPass<"llzk-poly-lowering-pass"> {
     through degree-neutral roots such as `felt.add`, `felt.sub`, and `felt.neg`, so composite
     equality operands and struct constrain call arguments satisfy their degree bounds.
 
-    The pass expects already-flattened, straight-line `constrain` bodies. Run `llzk-flatten` or
-    another control-flow lowering pass before this pass; nested-region, successor-bearing, or
-    multi-block constraints are rejected instead of being rewritten with component-lifetime
-    auxiliary members.
+    The pass expects already-flattened, straight-line `compute` and `constrain` bodies. Run
+    `llzk-flatten` or another control-flow lowering pass before this pass; nested-region,
+    successor-bearing, or multi-block functions are rejected instead of being rewritten with
+    component-lifetime auxiliary members.
 
     This pass is best used on already-flattened input as part of the `-llzk-full-poly-lowering`
     pipeline, which includes additional cleanup passes to ensure correctness and optimal

lib/Transforms/LLZKLoweringUtils.cpp

@@ -11,6 +11,7 @@
 #include <mlir/IR/Block.h>
 #include <mlir/IR/Builders.h>
 #include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/IRMapping.h>
 #include <mlir/IR/Operation.h>
 #include <mlir/Support/LogicalResult.h>
 
@@ -26,6 +27,33 @@ using namespace llzk::constrain;
 
 namespace llzk {
 
+namespace {
+
+Value mapBlockArgumentInCompute(BlockArgument barg, FuncDefOp computeFunc) {
+  // constrain(%self, args...) corresponds to compute(args...).  When rebuilding
+  // expressions in compute(), the constrain self argument maps to the compute
+  // self value, and every other constrain argument maps one position earlier in
+  // compute().
+  if (barg.getArgNumber() == 0) {
+    return computeFunc.getSelfValueFromCompute();
+  }
+  return computeFunc.getArgument(barg.getArgNumber() - 1);
+}
+
+Value mapValueIntoCompute(
+    Value val, FuncDefOp computeFunc, OpBuilder &builder, DenseMap<Value, Value> &memo
+) {
+  if (auto it = memo.find(val); it != memo.end()) {
+    return it->second;
+  }
+  if (auto barg = llvm::dyn_cast<BlockArgument>(val)) {
+    return memo[val] = mapBlockArgumentInCompute(barg, computeFunc);
+  }
+  return rebuildExprInCompute(val, computeFunc, builder, memo);
+}
+
+} // namespace
+
 Value rebuildExprInCompute(
     Value val, FuncDefOp computeFunc, OpBuilder &builder, DenseMap<Value, Value> &memo
 ) {
@@ -34,17 +62,36 @@ Value rebuildExprInCompute(
   }
 
   if (auto barg = llvm::dyn_cast<BlockArgument>(val)) {
-    unsigned index = barg.getArgNumber();
-    Value mapped = computeFunc.getArgument(index - 1);
-    return memo[val] = mapped;
+    return memo[val] = mapBlockArgumentInCompute(barg, computeFunc);
   }
 
   if (auto readOp = val.getDefiningOp<MemberReadOp>()) {
-    Value self = computeFunc.getSelfValueFromCompute();
-    Value rebuilt = builder.create<MemberReadOp>(
-        readOp.getLoc(), readOp.getType(), self, readOp.getMemberNameAttr().getAttr()
+    IRMapping mapper;
+    for (Value operand : readOp->getOperands()) {
+      mapper.map(operand, mapValueIntoCompute(operand, computeFunc, builder, memo));
+    }
+
+    Operation *rebuiltOp = builder.clone(*readOp.getOperation(), mapper);
+    assert(rebuiltOp->getNumResults() == 1 && "member reads have exactly one result");
+    return memo[val] = rebuiltOp->getResult(0);
+  }
+
+  if (val.getType().isIndex()) {
+    Operation *defOp = val.getDefiningOp();
+    assert(defOp && "index block arguments should already be mapped");
+
+    IRMapping mapper;
+    for (Value operand : defOp->getOperands()) {
+      mapper.map(operand, mapValueIntoCompute(operand, computeFunc, builder, memo));
+    }
+
+    Operation *rebuiltOp = builder.clone(*defOp, mapper);
+    assert(
+        rebuiltOp->getNumResults() == defOp->getNumResults() &&
+        "cloned index op should preserve result count"
     );
-    return memo[val] = rebuilt;
+    unsigned resultNumber = llvm::cast<OpResult>(val).getResultNumber();
+    return memo[val] = rebuiltOp->getResult(resultNumber);
   }
 
   if (auto add = val.getDefiningOp<AddFeltOp>()) {
@@ -99,16 +146,15 @@ LogicalResult checkForAuxMemberConflicts(StructDefOp structDef, StringRef prefix
   return failure(conflictFound);
 }
 
-LogicalResult checkConstrainBodyIsStraightLine(FuncDefOp constrainFunc, StringRef passName) {
-  auto emitStraightLineError = [passName](Operation *op) {
-    op->emitError() << passName
-                    << " expects a straight-line constrain body; run `llzk-flatten` or another "
-                       "control-flow lowering pass first";
+LogicalResult checkFuncBodyIsStraightLine(FuncDefOp func, StringRef passName, StringRef funcName) {
+  auto emitStraightLineError = [passName, funcName](Operation *op) {
+    op->emitError() << passName << " expects a straight-line " << funcName
+                    << " body; run `llzk-flatten` or another control-flow lowering pass first";
   };
 
-  Region &body = constrainFunc.getBody();
+  Region &body = func.getBody();
   if (!body.hasOneBlock()) {
-    emitStraightLineError(constrainFunc.getOperation());
+    emitStraightLineError(func.getOperation());
     return failure();
   }
 
@@ -129,6 +175,10 @@ LogicalResult checkConstrainBodyIsStraightLine(FuncDefOp constrainFunc, StringRe
   return failure();
 }
 
+LogicalResult checkConstrainBodyIsStraightLine(FuncDefOp constrainFunc, StringRef passName) {
+  return checkFuncBodyIsStraightLine(constrainFunc, passName, "constrain");
+}
+
 void replaceSubsequentUsesWith(Value oldVal, Value newVal, Operation *afterOp) {
   assert(afterOp && "afterOp must be a valid Operation*");
 
@@ -148,12 +198,17 @@ void replaceSubsequentUsesWith(Value oldVal, Value newVal, Operation *afterOp) {
   }
 }
 
-MemberDefOp addAuxMember(StructDefOp structDef, StringRef name) {
+MemberDefOp addAuxMember(StructDefOp structDef, StringRef name, Type type) {
+  assert(type && "auxiliary member type must be non-null");
+
   OpBuilder builder(structDef);
   builder.setInsertionPointToEnd(structDef.getBody());
-  return builder.create<MemberDefOp>(
-      structDef.getLoc(), builder.getStringAttr(name), builder.getType<FeltType>()
-  );
+  return builder.create<MemberDefOp>(structDef.getLoc(), builder.getStringAttr(name), type);
+}
+
+MemberDefOp addAuxMember(StructDefOp structDef, StringRef name) {
+  OpBuilder builder(structDef);
+  return addAuxMember(structDef, name, builder.getType<FeltType>());
 }
 
 unsigned getFeltDegree(Value val, DenseMap<Value, unsigned> &memo) {