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.
This commit is contained in:
onurinanc
2024-07-12 14:38:14 +03:00
committed by GitHub
parent fbf2b4074d
commit 6bd82ff12e
7 changed files with 71 additions and 74 deletions

View File

@@ -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;

View File

@@ -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))
);

View File

@@ -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...

View File

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

View File

@@ -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))
};

View File

@@ -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`.

View File

@@ -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 };