fingerprint() performance improvements (#2175)

Cherry-picked from #2174

When used on top of #2173, this accelerates second-stage witness
generation for the main machine from ~35s to ~10s for the example
mentioned in #2173.

std/protocols/bus.asm

@@ -82,21 +82,25 @@ let bus_interaction: expr, expr[], expr -> () = constr |id, tuple, multiplicity|
 /// 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: 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_next = sub_ext(eval_ext(beta), fingerprint_with_id(eval(id'), array::eval(array::next(tuple)), alpha));
 
-    let m_ext = from_base(multiplicity);
-    let m_ext_next = next_ext(m_ext);
+    let m_next = eval(multiplicity');
+    let m_ext_next = from_base(m_next);
 
     let is_first_next = eval(is_first');
     let current_acc = if is_first_next == 1 {from_base(0)} else {eval_ext(acc)};
 
     // acc' = current_acc + multiplicity' / folded'
-    let res = add_ext(
-        current_acc,
-        mul_ext(eval_ext(m_ext_next), inv_ext(folded_next))
-    );
+    let res = if m_next == 0 {
+        current_acc
+    } else {
+        // Implemented as: folded = (beta - fingerprint(id, tuple...));
+        // `multiplicity / (beta - fingerprint(id, tuple...))` to `acc`
+        let folded_next = sub_ext(eval_ext(beta), fingerprint_with_id(eval(id'), array::eval(array::next(tuple)), alpha));
+        add_ext(
+            current_acc,
+            mul_ext(m_ext_next, inv_ext(folded_next))
+        )
+    };
 
     unpack_ext_array(res)
 };

std/protocols/fingerprint.asm

@@ -11,15 +11,21 @@ use std::check::assert;
 /// 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$
 /// To generate an expression that computes the fingerprint, use `fingerprint_inter` instead.
 /// Note that alpha is passed as an expressions, so that it is only evaluated if needed (i.e., if len(expr_array) > 1).
-let fingerprint: fe[], Fp2<expr> -> Fp2<fe> = query |expr_array, alpha| if len(expr_array) == 1 {
+let fingerprint: fe[], Fp2<expr> -> Fp2<fe> = query |expr_array, alpha| if array::len(expr_array) == 1 {
+    // No need to evaluate `alpha` (which would be removed by the optimizer).
+    from_base(expr_array[0])
+} else {
+    fingerprint_impl(expr_array, eval_ext(alpha), len(expr_array))
+};
+
+let fingerprint_impl: fe[], Fp2<fe>, int -> Fp2<fe> = query |expr_array, alpha, l| if l == 1 {
     // Base case
     from_base(expr_array[0])
 } else {
-    assert(len(expr_array) > 1, || "fingerprint requires at least one element");
 
     // Recursively compute the fingerprint as fingerprint(expr_array[:-1], alpha) * alpha + expr_array[-1]
-    let intermediate_fingerprint = fingerprint(array::sub_array(expr_array, 0, len(expr_array) - 1), alpha);
-    add_ext(mul_ext(eval_ext(alpha), intermediate_fingerprint), from_base(expr_array[len(expr_array) - 1]))
+    let intermediate_fingerprint = fingerprint_impl(expr_array, alpha, l - 1);
+    add_ext(mul_ext(alpha, intermediate_fingerprint), from_base(expr_array[l - 1]))
 };
 
 /// Like `fingerprint`, but "materializes" the intermediate results as intermediate columns.