Implement Basic Bus (#1566)

Related to the issue of implementing basic bus (#1497), I have
implemented basic bus together with an example
(`permutation_via_bus.asm`) as specified inside the issue.

Currently, `test_data/std/bus_permutation_via_challenges.asm` works as
intended (To make it sound, stage(1) witness columns need to be exposed
publicly and verifier needs to check such as `out_z1 + out_z2 = 0`) We
can now check using RUST_LOG=trace and adding the final z1 and z2 is
equal to 0.

However, `test_data/std/bus_permutation_via_challenges_ext.asm` is not
working correctly as intended. This will be fixed with the following
commits.

pipeline/tests/powdr_std.rs

@@ -156,6 +156,18 @@ fn lookup_via_challenges_ext_simple() {
         .unwrap();
 }
 
+#[test]
+fn bus_permutation_via_challenges_bn() {
+    let f = "std/bus_permutation_via_challenges.asm";
+    test_halo2(f, Default::default());
+}
+
+#[test]
+fn bus_permutation_via_challenges_ext_bn() {
+    let f = "std/bus_permutation_via_challenges_ext.asm";
+    test_halo2(f, Default::default());
+}
+
 #[test]
 fn write_once_memory_test() {
     let f = "std/write_once_memory_test.asm";

std/protocols/bus.asm

@@ -0,0 +1,104 @@
+use std::check::assert;
+use std::check::panic;
+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::inv_ext;
+use std::math::fp2::eval_ext;
+use std::math::fp2::unpack_ext;
+use std::math::fp2::unpack_ext_array;
+use std::math::fp2::next_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_with_id;
+use std::prover::eval;
+
+/// Sends the tuple (id, tuple...) to the bus by adding
+/// `multiplicity / (beta - fingerprint(id, tuple...))` to `acc`
+/// It is the callers responsibility to properly constrain the multiplicity (e.g. constrain
+/// it to be boolean) if needed.
+///
+/// # Arguments:
+///
+/// - is_first: A column that is 1 for the first row and 0 for the rest
+/// - id: Interaction Id
+/// - tuple: An array of columns to be sent to the bus
+/// - multiplicity: The multiplicity which shows how many times a column will be sent
+/// - 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 id and tuple
+/// - beta: A challenge used to update the accumulator
+///
+/// # Returns:
+///
+/// - Constraints to be added to enforce the bus
+let bus_interaction: expr, expr, expr[], expr, expr[], Fp2<expr>, Fp2<expr> -> Constr[] = |is_first, id, tuple, multiplicity, acc, alpha, beta| {
+
+    // Implemented as: folded = (beta - fingerprint(id, tuple...));
+    let folded = sub_ext(beta, fingerprint_with_id(id, tuple, alpha));
+    let folded_next = next_ext(folded);
+
+    let m_ext = from_base(multiplicity);
+    let m_ext_next = next_ext(m_ext);
+
+    let acc_ext = fp2_from_array(acc);
+    let next_acc = next_ext(acc_ext);
+
+    let is_first_next = from_base(is_first');
+
+    // Update rule:
+    // acc' =  acc * (1 - is_first') + multiplicity' / folded'
+    // or equivalently:
+    // folded' * (acc' - acc * (1 - is_first')) - multiplicity' = 0
+    let update_expr = sub_ext(
+        mul_ext(folded_next, sub_ext(next_acc, mul_ext(acc_ext, sub_ext(from_base(1), is_first_next)))), m_ext_next
+    );
+
+    constrain_eq_ext(update_expr, from_base(0))
+};
+
+/// Compute acc' = acc * (1 - is_first') + multiplicity' / fingerprint_with_id(id, (a1', a2')),
+/// using extension field arithmetic.
+/// This is intended to be used as a hint in the extension field case; for the base case
+/// automatic witgen is smart enough to figure out the value of the accumulator.
+let compute_next_z_send: expr, expr, expr[], expr, Fp2<expr>, Fp2<expr>, Fp2<expr> -> fe[] = query |is_first, id, tuple, multiplicity, acc, alpha, beta| {
+    // Implemented as: folded = (beta - fingerprint(id, tuple...));
+    // `multiplicity / (beta - fingerprint(id, tuple...))` to `acc`
+    let folded = sub_ext(beta, fingerprint_with_id(id, tuple, alpha));
+    let folded_next = next_ext(folded);
+
+    let m_ext = from_base(multiplicity);
+    let m_ext_next = next_ext(m_ext);
+
+    let is_first_next = eval(is_first');
+    let current_acc = if is_first_next == 1 {from_base(0)} else {eval_ext(acc)};
+
+    // acc' = acc * (1 - is_first') + multiplicity / fingerprint_with_id(id, (a1', a2'))
+    let res = add_ext(
+        current_acc,
+        mul_ext(eval_ext(m_ext_next), inv_ext(eval_ext(folded_next)))
+    );
+
+    unpack_ext_array(res)
+};
+
+/// Compute acc' = acc * (1 - is_first') - multiplicity' / fingerprint_with_id(id, (a1', a2')),
+/// using extension field arithmetic.
+/// This is intended to be used as a hint in the extension field case; for the base case
+/// automatic witgen is smart enough to figure out the value of the accumulator.
+let compute_next_z_receive: expr, expr, expr[], expr, Fp2<expr>, Fp2<expr>, Fp2<expr> -> fe[] = query |is_first, id, tuple, multiplicity, acc, alpha, beta| 
+    compute_next_z_send(is_first, id, tuple, -multiplicity, acc, alpha, beta);
+
+/// Convenience function for bus interaction to send columns
+let bus_send: expr, expr, expr[], expr, expr[], Fp2<expr>, Fp2<expr> -> Constr[] = |is_first, id, tuple, multiplicity, acc, alpha, beta| {
+    bus_interaction(is_first, id, tuple, multiplicity, acc, alpha, beta)
+};
+
+/// Convenience function for bus interaction to receive columns
+let bus_receive: expr, expr, expr[], expr, expr[], Fp2<expr>, Fp2<expr> -> Constr[] = |is_first, id, tuple, multiplicity, acc, alpha, beta| {
+    bus_interaction(is_first, id, tuple, -1 * multiplicity, acc, alpha, beta)
+};

std/protocols/fingerprint.asm

@@ -9,4 +9,7 @@ let<T: Add + Mul + FromLiteral> fingerprint: T[], Fp2<T> -> Fp2<T> = |expr_array
     expr_array,
     from_base(0),
     |sum_acc, el| add_ext(mul_ext(alpha, sum_acc), from_base(el))
-);
+);
+
+/// Maps [id, 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_with_id: T, T[], Fp2<T> -> Fp2<T> = |id, expr_array, alpha| fingerprint([id] + expr_array, alpha);

std/protocols/lookup.asm

@@ -28,7 +28,7 @@ let unpack_lookup_constraint: Constr -> (expr, expr[], expr, expr[]) = |lookup_c
     _ => panic("Expected lookup constraint")
 };
 
-// Compute z' = z + 1/(beta-a_i) * lhs_selector - m_i/(beta-b_i) * rhs_selector, using extension field arithmetic
+/// 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);
 
@@ -51,15 +51,15 @@ let compute_next_z: Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr, expr -> fe[] = quer
     unpack_ext_array(res)
 };
 
-// Adds constraints that enforce that rhs is the lookup for lhs
-// Arguments:
-// - is_first: A column that is 1 for the first row and 0 for the rest
-// - alpha: A challenge used to compress the LHS and RHS values
-// - beta: A challenge used to update the accumulator
-// - 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                  
+/// Adds constraints that enforce that rhs is the lookup for lhs
+/// Arguments:
+/// - is_first: A column that is 1 for the first row and 0 for the rest
+/// - 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
+/// - lookup_constraint: The lookup constraint
+/// - multiplicities: The multiplicities which shows how many times each RHS value appears in the LHS                  
 let lookup: expr, expr[], Fp2<expr>, Fp2<expr>, Constr, expr -> Constr[] = |is_first, acc, alpha, beta, lookup_constraint, multiplicities| {
 
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_lookup_constraint(lookup_constraint);

std/protocols/mod.asm

@@ -1,3 +1,5 @@
+mod bus;
 mod fingerprint;
 mod lookup;
-mod permutation;
+mod permutation;
+mod permutation_via_bus;

std/protocols/permutation.asm

@@ -55,8 +55,11 @@ 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:
+/// - is_first: A column that is 1 for the first row and 0 for the rest
 /// - 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:

std/protocols/permutation_via_bus.asm

@@ -0,0 +1,26 @@
+use std::array::len;
+use std::check::assert;
+use std::protocols::bus::bus_send;
+use std::protocols::bus::bus_receive;
+use std::protocols::bus::compute_next_z_send;
+use std::protocols::bus::compute_next_z_receive;
+use std::protocols::permutation::unpack_permutation_constraint;
+use std::math::fp2::Fp2;
+
+// Example usage of the bus: Implement a permutation constraint
+// To make this sound, the last values of `acc_lhs` and `acc_rhs` need to be
+// exposed as publics, and the verifier needs to assert that they sum to 0.
+let permutation: expr, expr, expr[], expr[], Fp2<expr>, Fp2<expr>, Constr -> Constr[] = |is_first, id, acc_lhs, acc_rhs, alpha, beta, permutation_constraint| {
+    let (lhs_selector, lhs, rhs_selector, rhs) = unpack_permutation_constraint(permutation_constraint);
+    bus_send(is_first, id, lhs, lhs_selector, acc_lhs, alpha, beta) + bus_receive(is_first, id, rhs, rhs_selector, acc_rhs, alpha, beta)
+};
+
+let compute_next_z_send_permutation: expr, expr, Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr -> fe[] = query |is_first, id, acc, alpha, beta, permutation_constraint| {
+    let (lhs_selector, lhs, rhs_selector, rhs) = unpack_permutation_constraint(permutation_constraint);  
+    compute_next_z_send(is_first, id, lhs, lhs_selector, acc, alpha, beta)
+};
+
+let compute_next_z_receive_permutation: expr, expr, Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr -> fe[] = query |is_first, id, acc, alpha, beta, permutation_constraint| {
+    let (lhs_selector, lhs, rhs_selector, rhs) = unpack_permutation_constraint(permutation_constraint);  
+    compute_next_z_receive(is_first, id, rhs, rhs_selector, acc, alpha, beta)
+};

test_data/std/bus_permutation_via_challenges.asm

@@ -0,0 +1,34 @@
+use std::prover::Query;
+use std::convert::fe;
+use std::protocols::permutation_via_bus::permutation;
+use std::math::fp2::from_base;
+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
+    // (when selected by first_four and (1 - first_four), respectively)
+    col witness a1(i) query Query::Hint(fe(i));
+    col witness a2(i) query Query::Hint(fe(i + 42));
+    col witness b1(i) query Query::Hint(fe(7 - i));
+    col witness b2(i) query Query::Hint(fe(7 - i + 42));
+
+    let permutation_constraint = Constr::Permutation(
+        (Option::Some(first_four), Option::Some(1 - first_four)),
+        [(a1, b1), (a2, 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;
+    col witness stage(1) u;
+
+    let is_first: col = std::well_known::is_first;
+    permutation(is_first, 1, [z], [u], alpha, beta, permutation_constraint);
+
+    is_first' * (z + u) = 0;
+}

test_data/std/bus_permutation_via_challenges_ext.asm

@@ -0,0 +1,61 @@
+use std::prover::Query;
+use std::convert::fe;
+use std::protocols::permutation_via_bus::permutation;
+use std::protocols::permutation_via_bus::compute_next_z_send_permutation;
+use std::protocols::permutation_via_bus::compute_next_z_receive_permutation;
+use std::math::fp2::Fp2;
+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
+    // (when selected by first_four and (1 - first_four), respectively)
+    col witness a1(i) query Query::Hint(fe(i));
+    col witness a2(i) query Query::Hint(fe(i + 42));
+    col witness b1(i) query Query::Hint(fe(7 - i));
+    col witness b2(i) query Query::Hint(fe(7 - i + 42));
+
+    let permutation_constraint = Constr::Permutation(
+        (Option::Some(first_four), Option::Some(1 - first_four)),
+        [(a1, b1), (a2, 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);
+
+    col witness stage(1) u1;
+    col witness stage(1) u2;
+    let u = Fp2::Fp2(u1, u2);
+
+    let is_first: col = std::well_known::is_first;
+    permutation(is_first, 1, [z1,z2], [u1, u2], alpha, beta, permutation_constraint);
+
+    let hint_send = query |i| Query::Hint(compute_next_z_send_permutation(is_first, 1, z, alpha, beta, permutation_constraint)[i]);
+    col witness stage(1) z1_next(i) query hint_send(0);
+    col witness stage(1) z2_next(i) query hint_send(1);
+
+    z1' = z1_next;
+    z2' = z2_next;
+
+    let hint_receive = query |i| Query::Hint(compute_next_z_receive_permutation(is_first, 1, u, alpha, beta, permutation_constraint)[i]);
+    col witness stage(1) u1_next(i) query hint_receive(0);
+    col witness stage(1) u2_next(i) query hint_receive(1);
+
+    u1' = u1_next;
+    u2' = u2_next;
+
+    is_first' * (z1 + u1) = 0;
+    is_first' * (z2 + u2) = 0;
+}