Simplify permutation using prover functions similiar to lookups (#1826)

Similar to #1802, we can also simplify permutation in PIL using prover
functions.

Author: onurinanc

pipeline/tests/powdr_std.rs

@@ -119,7 +119,6 @@ fn permutation_via_challenges_bn() {
 }
 
 #[test]
-#[should_panic = "Error reducing expression to constraint:\nExpression: std::protocols::permutation::permutation([main::z], main::alpha, main::beta, main::permutation_constraint)\nError: FailedAssertion(\"The field is too small and needs to move to the extension field. Pass two elements instead!\")"]
 fn permutation_via_challenges_gl() {
     let f = "std/permutation_via_challenges.asm";
     make_simple_prepared_pipeline::<GoldilocksField>(f);

std/protocols/permutation.asm

@@ -1,5 +1,4 @@
 use std::array::map;
-use std::array::len;
 use std::check::assert;
 use std::check::panic;
 use std::math::fp2::Fp2;
@@ -14,6 +13,8 @@ use std::math::fp2::eval_ext;
 use std::math::fp2::from_base;
 use std::math::fp2::fp2_from_array;
 use std::math::fp2::constrain_eq_ext;
+use std::math::fp2::required_extension_size;
+use std::math::fp2::needs_extension;
 use std::protocols::fingerprint::fingerprint;
 use std::utils::unwrap_or_else;
 
@@ -53,16 +54,9 @@ let compute_next_z: Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr -> fe[] = query |acc
 };
 
 /// Returns constraints that enforce that lhs is a permutation of rhs
-///
-/// # Arguments:
-/// - acc: A phase-2 witness column to be used as the accumulator. If 2 are provided, computations
-///        are done on the F_{p^2} extension field.
-/// - alpha: A challenge used to compress the LHS and RHS values
-/// - beta: A challenge used to update the accumulator
-/// - permutation_constraint: The permutation constraint
-///
-/// # Returns:
-/// - Constraints to be added to enforce the permutation
+/// WARNING: This function can currently not be used multiple times since
+/// the used challenges would overlap
+/// TODO: Implement this for an array of constraints
 ///
 /// # Implementation:
 /// This function implements a permutation argument described e.g. in
@@ -79,7 +73,12 @@ let compute_next_z: Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr -> fe[] = query |acc
 /// the wrapping behavior: The first accumulator is constrained to be 1, and the last
 /// accumulator is the same as the first one, because of wrapping.
 /// For small fields, this computation should happen in the extension field.
-let permutation: expr[], Fp2<expr>, Fp2<expr>, Constr -> () = constr |acc, alpha, beta, permutation_constraint| {
+let permutation: Constr -> () = constr |permutation_constraint| {
+    std::check::assert(required_extension_size() <= 2, || "Invalid extension size");
+    // Alpha is used to compress the LHS and RHS arrays
+    let alpha = fp2_from_array(std::array::new(required_extension_size(), |i| challenge(0, i + 1)));
+    // Beta is used to update the accumulator
+    let beta = fp2_from_array(std::array::new(required_extension_size(), |i| challenge(0, i + 3)));
 
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_permutation_constraint(permutation_constraint);
 
@@ -88,6 +87,8 @@ let permutation: expr[], Fp2<expr>, Fp2<expr>, Constr -> () = constr |acc, alpha
     // Implemented as: folded = selector * (beta - fingerprint(values) - 1) + 1;
     let lhs_folded = selected_or_one(lhs_selector, sub_ext(beta, fingerprint(lhs, alpha)));
     let rhs_folded = selected_or_one(rhs_selector, sub_ext(beta, fingerprint(rhs, alpha)));
+
+    let acc = std::array::new(required_extension_size(), |i| std::prover::new_witness_col_at_stage("acc", 1));
     let acc_ext = fp2_from_array(acc);
     let next_acc = next_ext(acc_ext);
 
@@ -108,4 +109,21 @@ let permutation: expr[], Fp2<expr>, Fp2<expr>, Constr -> () = constr |acc, alpha
     is_first * (acc_1 - 1) = 0;
     is_first * acc_2 = 0;
     constrain_eq_ext(update_expr, from_base(0));
+
+    // In the extension field, we need a prover function for the accumulator.
+    if needs_extension() {
+        // TODO: Helper columns, because we can't access the previous row in hints
+        let acc_next_col = std::array::map(acc, |_| std::prover::new_witness_col_at_stage("acc_next", 1));
+        query |i| {
+            let _ = std::array::zip(
+                acc_next_col,
+                compute_next_z(acc_ext, alpha, beta, permutation_constraint),
+                |acc_next, hint_val| std::prover::provide_value(acc_next, i, hint_val)
+            );
+        };
+        std::array::zip(acc, acc_next_col, |acc_col, acc_next| {
+            acc_col' = acc_next
+        });
+    } else {
+    }
 };

test_data/std/permutation_via_challenges.asm

@@ -6,9 +6,6 @@ use std::prover::challenge;
 
 machine Main with degree: 8 {
 
-    let alpha = from_base(challenge(0, 1));
-    let beta = from_base(challenge(0, 2));
-
     col fixed first_four = [1, 1, 1, 1, 0, 0, 0, 0];
 
     // Two pairs of witness columns, claimed to be permutations of one another
@@ -23,8 +20,5 @@ machine Main with degree: 8 {
 
     let permutation_constraint = first_four $ [a1, a2] is (1 - first_four) $ [b1, b2];
 
-    // TODO: Functions currently cannot add witness columns at later stages,
-    // so we have to manually create it here and pass it to permutation(). 
-    col witness stage(1) z;
-    permutation([z], alpha, beta, permutation_constraint);
+    permutation(permutation_constraint);
 }

test_data/std/permutation_via_challenges_ext.asm

@@ -7,13 +7,6 @@ use std::prover::challenge;
 
 machine Main with degree: 8 {
 
-    let alpha1: expr = challenge(0, 1);
-    let alpha2: expr = challenge(0, 2);
-    let beta1: expr = challenge(0, 3);
-    let beta2: expr = challenge(0, 4);
-    let alpha = Fp2::Fp2(alpha1, alpha2);
-    let beta = Fp2::Fp2(beta1, beta2);
-
     col fixed first_four = [1, 1, 1, 1, 0, 0, 0, 0];
 
     // Two pairs of witness columns, claimed to be permutations of one another
@@ -27,24 +20,6 @@ machine Main with degree: 8 {
     };
 
     let permutation_constraint = first_four $ [a1, a2] is (1 - first_four) $ [b1, b2];
-
-    // TODO: Functions currently cannot add witness columns at later stages,
-    // so we have to manually create it here and pass it to permutation(). 
-    col witness stage(1) z1;
-    col witness stage(1) z2;
-    let z = Fp2::Fp2(z1, z2);
-
-    permutation([z1, z2], alpha, beta, permutation_constraint);
-
-    // TODO: Helper columns, because we can't access the previous row in hints
-    col witness stage(1) z1_next;
-    col witness stage(1) z2_next;
-    query |i| {
-        let hint = compute_next_z(z, alpha, beta, permutation_constraint);
-        std::prover::provide_value(z1_next, i, hint[0]);
-        std::prover::provide_value(z2_next, i, hint[1]);
-    };
-
-    z1' = z1_next;
-    z2' = z2_next;
+    
+    permutation(permutation_constraint);
 }