Merge common functions in permutation and lookup (#1518)

This PR aims to merge common functions in `permutation.asm` and
`lookup.asm` as much as possible.

std/math/fp2.asm

@@ -1,7 +1,13 @@
+use std::array::len;
+use std::array::fold;
+use std::check::assert;
+use std::check::panic;
 use std::convert::fe;
 use std::convert::int;
 use std::convert::expr;
 use std::field::modulus;
+use std::field::known_field;
+use std::field::KnownField;
 use std::prover::eval;
 
 /// An element of the extension field over the implied base field (which has to be either
@@ -90,6 +96,23 @@ let<T> unpack_ext: Fp2<T> -> (T, T) = |a| match a {
     Fp2::Fp2(a0, a1) => (a0, a1)
 };
 
+/// Whether we need to operate on the F_{p^2} extension field (because the current field is too small).
+let needs_extension: -> bool = || match known_field() {
+    Option::Some(KnownField::Goldilocks) => true,
+    Option::Some(KnownField::BN254) => false,
+    None => panic("The permutation/lookup argument is not implemented for the current field!")
+};
+
+/// Matches whether the length of a given array is correct to operate on the extension field
+let is_extension = |arr| match len(arr) {
+        1 => false,
+        2 => true,
+        _ => panic("Expected 1 or 2 accumulator columns!")
+};
+
+/// Constructs an extension field element `a0 + a1 * X` from either `[a0, a1]` or `[a0]` (setting `a1`to zero in that case)
+let fp2_from_array = |arr| if is_extension(arr) { Fp2::Fp2(arr[0], arr[1]) } else { from_base(arr[0]) };
+
 mod test {
     use super::Fp2;
     use super::from_base;

std/protocols/fingerprint.asm

@@ -0,0 +1,12 @@
+use std::array::fold;
+use std::math::fp2::Fp2;
+use std::math::fp2::add_ext;
+use std::math::fp2::mul_ext;
+use std::math::fp2::from_base;
+
+/// Maps [x_1, x_2, ..., x_n] to its Read-Solomon fingerprint, using a challenge alpha: $\sum_{i=1}^n alpha**{(n - i)} * x_i$
+let<T: Add + Mul + FromLiteral> fingerprint: T[], Fp2<T> -> Fp2<T> = |expr_array, alpha| fold(
+    expr_array,
+    from_base(0),
+    |sum_acc, el| add_ext(mul_ext(alpha, sum_acc), from_base(el))
+);

std/protocols/lookup.asm

@@ -1,11 +1,8 @@
 use std::array::fold;
-use std::utils::unwrap_or_else;
 use std::array::len;
 use std::array::map;
 use std::check::assert;
 use std::check::panic;
-use std::field::known_field;
-use std::field::KnownField;
 use std::math::fp2::Fp2;
 use std::math::fp2::add_ext;
 use std::math::fp2::sub_ext;
@@ -15,7 +12,12 @@ use std::math::fp2::next_ext;
 use std::math::fp2::inv_ext;
 use std::math::fp2::eval_ext;
 use std::math::fp2::from_base;
+use std::math::fp2::is_extension;
+use std::math::fp2::fp2_from_array;
+use std::math::fp2::needs_extension;
 use std::math::fp2::constrain_eq_ext;
+use std::protocols::fingerprint::fingerprint;
+use std::utils::unwrap_or_else;
 
 let unpack_lookup_constraint: Constr -> (expr, expr[], expr, expr[]) = |lookup_constraint| match lookup_constraint {
     Constr::Lookup((lhs_selector, rhs_selector), values) => (
@@ -27,27 +29,12 @@ let unpack_lookup_constraint: Constr -> (expr, expr[], expr, expr[]) = |lookup_c
     _ => panic("Expected lookup constraint")
 };
 
-/// Whether we need to operate on the F_{p^2} extension field (because the current field is too small).
-let needs_extension: -> bool = || match known_field() {
-    Option::Some(KnownField::Goldilocks) => true,
-    Option::Some(KnownField::BN254) => false,
-    None => panic("The lookup argument is not implemented for the current field!")
-};
-
-//* Generic for both permutation and lookup arguments
-/// Maps [x_1, x_2, ..., x_n] to alpha**(n - 1) * x_1 + alpha ** (n - 2) * x_2 + ... + x_n
-let<T: Add + Mul + FromLiteral> compress_expression_array: T[], Fp2<T> -> Fp2<T> = |expr_array, alpha| fold(
-    expr_array,
-    from_base(0),
-    |sum_acc, el| add_ext(mul_ext(alpha, sum_acc), from_base(el))
-);
-
 // Compute z' = z + 1/(beta-a_i) * lhs_selector - m_i/(beta-b_i) * rhs_selector, using extension field arithmetic
 let compute_next_z: Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr, expr -> fe[] = query |acc, alpha, beta, lookup_constraint, multiplicities| {
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_lookup_constraint(lookup_constraint);
 
-    let lhs_denom = sub_ext(beta, compress_expression_array(lhs, alpha));
-    let rhs_denom = sub_ext(beta, compress_expression_array(rhs, alpha));
+    let lhs_denom = sub_ext(beta, fingerprint(lhs, alpha));
+    let rhs_denom = sub_ext(beta, fingerprint(rhs, alpha));
     let m_ext = from_base(multiplicities);
 
     // acc' = acc + 1/(beta-a_i) * lhs_selector - m_i/(beta-b_i) * rhs_selector
@@ -81,28 +68,20 @@ let lookup: expr, expr[], Fp2<expr>, Fp2<expr>, Constr, expr -> Constr[] = |is_f
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_lookup_constraint(lookup_constraint);
 
     let _ = assert(len(lhs) == len(rhs), || "LHS and RHS should have equal length");
-
-    let with_extension = match len(acc) {
-        1 => false,
-        2 => true,
-        _ => panic("Expected 1 or 2 accumulator columns!")
-    };
-
-    let _ = if !with_extension {
+    let _ = if !is_extension(acc) {
         assert(!needs_extension(), || "The Goldilocks field is too small and needs to move to the extension field. Pass two accumulators instead!")
-    } else { () };
+    } else { };
 
     // On the extension field, we'll need two field elements to represent the challenge.
     // If we don't need an extension field, we can simply set the second component to 0,
     // in which case the operations below effectively only operate on the first component.
-    let fp2_from_array = |arr| if with_extension { Fp2::Fp2(arr[0], arr[1]) } else { from_base(arr[0]) };
     let acc_ext = fp2_from_array(acc);
 
-    let lhs_denom = sub_ext(beta, compress_expression_array(lhs, alpha));
-    let rhs_denom = sub_ext(beta, compress_expression_array(rhs, alpha));
+    let lhs_denom = sub_ext(beta, fingerprint(lhs, alpha));
+    let rhs_denom = sub_ext(beta, fingerprint(rhs, alpha));
     let m_ext = from_base(multiplicities);
 
-    let next_acc = if with_extension {
+    let next_acc = if is_extension(acc) {
         next_ext(acc_ext)
     } else {
         // The second component is 0, but the next operator is not defined on it...

std/protocols/mod.asm

@@ -1,2 +1,3 @@
-mod permutation;
-mod lookup;
+mod fingerprint;
+mod lookup;
+mod permutation;

std/protocols/permutation.asm

@@ -1,11 +1,7 @@
-use std::array::fold;
 use std::array::map;
-use std::utils::unwrap_or_else;
 use std::array::len;
 use std::check::assert;
 use std::check::panic;
-use std::field::known_field;
-use std::field::KnownField;
 use std::math::fp2::Fp2;
 use std::math::fp2::add_ext;
 use std::math::fp2::sub_ext;
@@ -15,7 +11,12 @@ use std::math::fp2::next_ext;
 use std::math::fp2::inv_ext;
 use std::math::fp2::eval_ext;
 use std::math::fp2::from_base;
+use std::math::fp2::needs_extension;
+use std::math::fp2::is_extension;
+use std::math::fp2::fp2_from_array;
 use std::math::fp2::constrain_eq_ext;
+use std::protocols::fingerprint::fingerprint;
+use std::utils::unwrap_or_else;
 
 let unpack_permutation_constraint: Constr -> (expr, expr[], expr, expr[]) = |permutation_constraint| match permutation_constraint {
     Constr::Permutation((lhs_selector, rhs_selector), values) => (
@@ -27,20 +28,6 @@ let unpack_permutation_constraint: Constr -> (expr, expr[], expr, expr[]) = |per
     _ => panic("Expected permutation constraint")
 };
 
-/// Whether we need to operate on the F_{p^2} extension field (because the current field is too small).
-let needs_extension: -> bool = || match known_field() {
-    Option::Some(KnownField::Goldilocks) => true,
-    Option::Some(KnownField::BN254) => false,
-    None => panic("The permutation argument is not implemented for the current field!")
-};
-
-/// Maps [x_1, x_2, ..., x_n] to its Read-Solomon fingerprint, using challenge alpha: $\sum_{i=1}^n alpha**{(n - i)} * x_i$
-let<T: Add + Mul + FromLiteral> compress_expression_array: T[], Fp2<T> -> Fp2<T> = |expr_array, alpha| fold(
-    expr_array,
-    from_base(0),
-    |sum_acc, el| add_ext(mul_ext(alpha, sum_acc), from_base(el))
-);
-
 /// Takes a boolean selector (0/1) and a value, returns equivalent of `if selector { value } else { 1 }`
 /// Implemented as: selector * (value - 1) + 1
 let<T: Add + Mul + Sub + FromLiteral> selected_or_one: T, Fp2<T> -> Fp2<T> = |selector, value| add_ext(mul_ext(from_base(selector), sub_ext(value, from_base(1))), from_base(1));
@@ -54,8 +41,8 @@ let compute_next_z: Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr -> fe[] = query |acc
 
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_permutation_constraint(permutation_constraint);
 
-    let lhs_folded = selected_or_one(lhs_selector, sub_ext(beta, compress_expression_array(lhs, alpha)));
-    let rhs_folded = selected_or_one(rhs_selector, sub_ext(beta, compress_expression_array(rhs, alpha)));
+    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)));
 
     // acc' = acc * lhs_folded / rhs_folded
     let res = mul_ext(
@@ -84,7 +71,7 @@ let compute_next_z: Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr -> fe[] = query |acc
 /// page 99, paragraph "Multiset equality checking (a.k.a. permutation checking)
 /// via fingerprinting". In short:
 /// 1. The LHS and RHS are Reed-Solomon fingerprinted using challenge $\alpha$
-///    (see `compress_expression_array`).
+///    (see `std::fingerprint::fingerprint`).
 /// 2. If the selector is one, the accumulator is updated as:
 ///    `acc' = acc * (beta - lhs) / (beta - rhs)`.
 /// This iteratively evaluates the fraction of polynomials $\prod_i (X - lhs_i)$
@@ -98,29 +85,23 @@ let permutation: expr, expr[], Fp2<expr>, Fp2<expr>, Constr -> Constr[] = |is_fi
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_permutation_constraint(permutation_constraint);
 
     let _ = assert(len(lhs) == len(rhs), || "LHS and RHS should have equal length");
-    let with_extension = match len(acc) {
-        1 => false,
-        2 => true,
-        _ => panic("Expected 1 or 2 accumulator columns!")
-    };
-
-    let _ = if !with_extension {
+    let _ = if !is_extension(acc) {
         assert(!needs_extension(), || "The Goldilocks field is too small and needs to move to the extension field. Pass two accumulators instead!")
-    } else { () };
-    
+    } else { };
+
     // On the extension field, we'll need two field elements to represent the challenge.
     // If we don't need an extension field, we can simply set the second component to 0,
     // in which case the operations below effectively only operate on the first component.
-    let fp2_from_array = |arr| if with_extension { Fp2::Fp2(arr[0], arr[1]) } else { from_base(arr[0]) };
+    let fp2_from_array = |arr| if is_extension(acc) { Fp2::Fp2(arr[0], arr[1]) } else { from_base(arr[0]) };
     let acc_ext = fp2_from_array(acc);
 
-    // If the selector is 1, contribute a factor of `beta - compress_expression_array(lhs)` to accumulator.
+    // If the selector is 1, contribute a factor of `beta - fingerprint(lhs)` to accumulator.
     // If the selector is 0, contribute a factor of 1 to the accumulator.
-    // Implemented as: folded = selector * (beta - compress_expression_array(values) - 1) + 1;
-    let lhs_folded = selected_or_one(lhs_selector, sub_ext(beta, compress_expression_array(lhs, alpha)));
-    let rhs_folded = selected_or_one(rhs_selector, sub_ext(beta, compress_expression_array(rhs, 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 next_acc = if with_extension {
+    let next_acc = if is_extension(acc) {
         next_ext(acc_ext)
     } else {
         // The second component is 0, but the next operator is not defined on it...
@@ -130,7 +111,7 @@ let permutation: expr, expr[], Fp2<expr>, Fp2<expr>, Constr -> Constr[] = |is_fi
     // Update rule:
     // acc' = acc * lhs_folded / rhs_folded
     // => rhs_folded * acc' - lhs_folded * acc = 0
-    let diff_from_expected = sub_ext(
+    let update_expr = sub_ext(
         mul_ext(rhs_folded, next_acc),
         mul_ext(lhs_folded, acc_ext)
     );
@@ -145,5 +126,5 @@ let permutation: expr, expr[], Fp2<expr>, Fp2<expr>, Constr -> Constr[] = |is_fi
         // Note that if with_extension is false, this generates 0 = 0 and is removed
         // by the optimizer.
         is_first * acc_2 = 0
-    ] + constrain_eq_ext(diff_from_expected, from_base(0))
+    ] + constrain_eq_ext(update_expr, from_base(0))
 };

std/utils.asm

@@ -1,4 +1,3 @@
-
 /// Evaluates to folder(...folder(folder(initial, f(0)), f(1)) ..., f(length - 1)),
 /// i.e. calls f(0), f(1), ..., f(length - 1) and combines the results
 /// using the function `folder`, starting with the value `initial`.

std/well_known.asm

@@ -1 +1,3 @@
+/// Evaluates to 1 on the first row and 0 on all other rows.
+/// Useful to define a fixed column of that property.
 let is_first: int -> int = |i| if i == 0 { 1 } else { 0 };