implement lookup arguments from logarithmic derivatives in pil (#1477)

This PR solves #1374

There are 3 examples included in the PR:

1. `lookup_via_challenges_ext_simple.asm` is the most basic example
which implements {a} in {b} using extension field without using
selectors
2. `lookup_via_challenges_ext_no_selector.asm` implements {a1, a2} in
{b1, b2} using extension field without using selectors
3. `lookup_via_challenges_ext.asm` is a more complex example than others
which implements {a1, a2, a3} in {b1, b2, b3} using extension field
(which also handles tuples using Reed-Solomon fingerprinting). It also
use different lhs and rhs selectors.

---------

Co-authored-by: Georg Wiese <georgwiese@gmail.com>

pipeline/tests/powdr_std.rs

@@ -115,6 +115,36 @@ fn permutation_via_challenges_ext() {
         .unwrap();
 }
 
+#[test]
+fn lookup_via_challenges_bn() {
+    let f = "std/lookup_via_challenges.asm";
+    test_halo2(f, Default::default());
+}
+
+#[test]
+fn lookup_via_challenges_ext() {
+    let f = "std/lookup_via_challenges_ext.asm";
+    test_halo2(f, Default::default());
+    // Note that this does not actually run the second-phase witness generation, because no
+    // Goldilocks backend support challenges yet.
+    Pipeline::<GoldilocksField>::default()
+        .from_file(resolve_test_file(f))
+        .compute_witness()
+        .unwrap();
+}
+
+#[test]
+fn lookup_via_challenges_ext_simple() {
+    let f = "std/lookup_via_challenges_ext_simple.asm";
+    test_halo2(f, Default::default());
+    // Note that this does not actually run the second-phase witness generation, because no
+    // Goldilocks backend support challenges yet.
+    Pipeline::<GoldilocksField>::default()
+        .from_file(resolve_test_file(f))
+        .compute_witness()
+        .unwrap();
+}
+
 #[test]
 fn write_once_memory_test() {
     let f = "std/write_once_memory_test.asm";

std/protocols/lookup.asm

@@ -0,0 +1,160 @@
+use std::prover::challenge;
+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;
+use std::math::fp2::mul_ext;
+use std::math::fp2::unpack_ext;
+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::constrain_eq_ext;
+
+
+let is_first: col = |i| if i == 0 { 1 } else { 0 };
+
+// challenges to be used in polynomial evaluation and folding different columns
+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 unpack_lookup_constraint: Constr -> (expr, expr[], expr, expr[]) = |lookup_constraint| match lookup_constraint {
+    Constr::Lookup((lhs_selector, rhs_selector), values) => (
+        unwrap_or_else(lhs_selector, || 1),
+        map(values, |(lhs, _)| lhs),
+        unwrap_or_else(rhs_selector, || 1),
+        map(values, |(_, rhs)| rhs)
+    ),
+    _ => 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>, Constr, expr -> fe[] = query |acc, lookup_constraint, multiplicities| {
+    let (lhs_selector, lhs, rhs_selector, rhs) = unpack_lookup_constraint(lookup_constraint);
+    let alpha = if len(lhs) > 1 {
+        Fp2::Fp2(alpha1, alpha2)
+    } else {
+        // The optimizer will have removed alpha, but the compression function
+        // still accesses it (to multiply by 0 in this case)
+        from_base(0)
+    };
+    let beta = Fp2::Fp2(beta1, beta2);
+    
+    let lhs_denom = sub_ext(beta, compress_expression_array(lhs, alpha));
+    let rhs_denom = sub_ext(beta, compress_expression_array(rhs, alpha));
+    let m_ext = from_base(multiplicities);
+    
+    // acc' = acc + 1/(beta-a_i) * lhs_selector - m_i/(beta-b_i) * rhs_selector
+    let res = add_ext(
+        eval_ext(acc),
+        sub_ext(
+            mul_ext(
+                inv_ext(eval_ext(lhs_denom)), 
+                eval_ext(from_base(lhs_selector))),
+            mul_ext(
+                mul_ext(eval_ext(m_ext), inv_ext(eval_ext(rhs_denom))),
+                eval_ext(from_base(rhs_selector))
+        )
+    ));
+    match res {
+        Fp2::Fp2(a0_fe, a1_fe) => [a0_fe, a1_fe]
+    }
+};
+    
+// Adds constraints that enforce that rhs is the lookup for lhs
+// 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.
+// - lookup_constraint: The lookup constraint
+// - multiplicities: The multiplicities which shows how many times each RHS value appears in the LHS                  
+let lookup: expr[], Constr, expr -> Constr[] = |acc, lookup_constraint, multiplicities| {
+
+    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 {
+        assert(!needs_extension(), || "The Goldilocks field is too small and needs to move to the extension field. Pass two accumulators instead!")
+    } 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 alpha = fp2_from_array([alpha1, alpha2]);
+    let beta = fp2_from_array([beta1, beta2]);
+
+    let lhs_denom = sub_ext(beta, compress_expression_array(lhs, alpha));
+    let rhs_denom = sub_ext(beta, compress_expression_array(rhs, alpha));
+    let m_ext = from_base(multiplicities);
+
+    let next_acc = if with_extension {
+        next_ext(acc_ext)
+    } else {
+        // The second component is 0, but the next operator is not defined on it...
+        from_base(acc[0]')
+    };
+
+    // Update rule:
+    // acc' * (beta - A) * (beta - B)  + m * rhs_selector * (beta - A) = acc * (beta - A) * (beta - B) + lhs_selector * (beta - B)
+    // => (acc' - acc) * lhs_denom * rhs_denom
+    //    + m * rhs_selector * lhs_denom
+    //    - lhs_selector * rhs_denom = 0
+
+    let update_expr = sub_ext(
+        add_ext(
+            mul_ext(
+                mul_ext(lhs_denom, rhs_denom),
+                sub_ext(next_acc, acc_ext)
+            ),
+            mul_ext(
+                mul_ext(m_ext, from_base(rhs_selector)),
+                lhs_denom
+            )
+        ),
+        mul_ext(
+            from_base(lhs_selector),
+            rhs_denom
+        )
+    );
+
+    let (acc_1, acc_2) = unpack_ext(acc_ext);
+
+    [
+        is_first * acc_1 = 0,
+
+        is_first * acc_2 = 0
+    ] + constrain_eq_ext(update_expr, from_base(0))
+};

std/protocols/mod.asm

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

test_data/std/lookup_via_challenges.asm

@@ -0,0 +1,27 @@
+use std::prover::Query;
+use std::convert::fe;
+use std::protocols::lookup::lookup;
+
+machine Main with degree: 8 {
+    col fixed random_six = [1, 1, 1, 0, 1, 1, 1, 0];
+    col fixed first_seven = [1, 1, 1, 1, 1, 1, 1, 0];
+
+    col fixed a1 = [1, 2, 4, 3, 1, 1, 4, 1];
+    col fixed a2 = [1, 2, 4, 1, 1, 1, 4, 1];
+    col fixed a3 = [1, 2, 4, 1, 1, 1, 4, 3];
+    col witness b1(i) query Query::Hint(fe(i+1));
+    col witness b2(i) query Query::Hint(fe(i+1));
+    col witness b3(i) query Query::Hint(fe(i+1));
+    col fixed m = [3, 1, 0, 2, 0, 0, 0, 0];
+
+    let lookup_constraint = Constr::Lookup(
+        (Option::Some(random_six), Option::Some(first_seven)),
+        [(a1, b1), (a2, b2), (a3, b3)]
+    );
+
+    // 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;
+    lookup([z], lookup_constraint, m);
+    
+}

test_data/std/lookup_via_challenges_ext.asm

@@ -0,0 +1,38 @@
+use std::prover::Query;
+use std::convert::fe;
+use std::protocols::lookup::lookup;
+use std::protocols::lookup::compute_next_z;
+use std::math::fp2::Fp2;
+
+machine Main with degree: 8 {
+    col fixed a_sel = [0, 1, 1, 1, 0, 1, 0, 0];
+    col fixed b_sel = [1, 1, 0, 1, 1, 1, 1, 0];
+
+    col fixed a1 = [16, 20, 22, 17, 16, 20, 4, 1];
+    col fixed a2 = [12, 5, 7, 2, 5, 5, 4, 1];
+    col fixed a3 = [20, 36, 38, 33, 36, 36, 4, 3];
+    col witness b1(i) query Query::Hint(fe(i+16));
+    col witness b2(i) query Query::Hint(fe(i+1));
+    col witness b3(i) query Query::Hint(fe(i+32));
+    col fixed m = [0, 1, 0, 0, 2, 0, 1, 0];
+
+    let lookup_constraint = Constr::Lookup(
+        (Option::Some(a_sel), Option::Some(b_sel)),
+        [(a1, b1), (a2, b2), (a3, b3)]
+    );
+
+    // TODO: Functions currently cannot add witness columns at later stages,
+    // so we have to manually create it here and pass it to lookup(). 
+    col witness stage(1) z1;
+    col witness stage(1) z2;
+
+    lookup([z1, z2], lookup_constraint, m);
+
+    // TODO: Helper columns, because we can't access the previous row in hints
+    let hint = query |i| Query::Hint(compute_next_z(Fp2::Fp2(z1, z2), lookup_constraint, m)[i]); 
+    col witness stage(1) z1_next(i) query hint(0);
+    col witness stage(1) z2_next(i) query hint(1);
+
+    z1' = z1_next;
+    z2' = z2_next;
+}

test_data/std/lookup_via_challenges_ext_simple.asm

@@ -0,0 +1,31 @@
+use std::prover::Query;
+use std::convert::fe;
+use std::protocols::lookup::lookup;
+use std::protocols::lookup::compute_next_z;
+use std::math::fp2::Fp2;
+
+machine Main with degree: 8 {
+    col fixed a = [1, 1, 4, 1, 1, 2, 1, 1];
+    col witness b(i) query Query::Hint(fe(i+1));
+    col fixed m = [6, 1, 0, 1, 0, 0, 0, 0];
+
+    let lookup_constraint = Constr::Lookup(
+        (Option::None, Option::None),
+        [(a, b)]
+    );
+
+    // TODO: Functions currently cannot add witness columns at later stages,
+    // so we have to manually create it here and pass it to lookup(). 
+    col witness stage(1) z1;
+    col witness stage(1) z2;
+
+    lookup([z1, z2], lookup_constraint, m);
+
+    // TODO: Helper columns, because we can't access the previous row in hints
+    let hint = query |i| Query::Hint(compute_next_z(Fp2::Fp2(z1, z2), lookup_constraint, m)[i]); 
+    col witness stage(1) z1_next(i) query hint(0);
+    col witness stage(1) z2_next(i) query hint(1);
+
+    z1' = z1_next;
+    z2' = z2_next;
+}