Build generic qm31 constants.

Author: leo-starkware

crates/circuits/src/finalize_constants.rs

@@ -103,6 +103,9 @@ fn finalize_constants_with_min_base(context: &mut Context<impl IValue>, min_base
         m31_base,
         qm31_basis,
     );
+    // Build the remaining QM31 constants.
+    decompose_qm31_constants(context, &mut qm31_constants, &mut m31_cache, m31_base, qm31_basis);
+    assert!(qm31_constants.is_empty());
 }
 
 /// Finds the largest integer N such that all values in [0, N] are present as constants.
@@ -252,3 +255,33 @@ fn decompose_broadcast_constants(
         false
     });
 }
+
+fn decompose_qm31_constants(
+    context: &mut Context<impl IValue>,
+    qm31_constants: &mut IndexMap<QM31, Var>,
+    m31_cache: &mut HashMap<M31, usize>,
+    base: M31,
+    qm_basis: [Var; 3],
+) {
+    let [i_var, u_var, iu_var] = qm_basis;
+
+    for (qm31_value, qm31_var) in qm31_constants.drain(..) {
+        let limbs = qm31_value.to_m31_array();
+        let [a, b, c, d]: [Var; 4] = std::array::from_fn(|j| {
+            let m31_val = limbs[j];
+            if !m31_cache.contains_key(&m31_val) {
+                build_m31_from_base(context, m31_cache, &mut IndexMap::new(), base, m31_val);
+            }
+            Var { idx: *m31_cache.get(&m31_val).unwrap() }
+        });
+
+        // a + b*i + c*u + d*iu → output to reserved idx
+        let first_half = eval!(context, (a) + ((b) * (i_var)));
+        let second_half = eval!(context, ((c) * (u_var)) + ((d) * (iu_var)));
+        context.circuit.add.push(Add {
+            in0: first_half.idx,
+            in1: second_half.idx,
+            out: qm31_var.idx,
+        });
+    }
+}

crates/circuits/src/finalize_constants_test.rs

@@ -151,3 +151,77 @@ fn test_broadcast_decomposition() {
     context.circuit.check_yields();
     context.validate_circuit();
 }
+
+#[test]
+fn test_mixed_m31_and_qm31_constants_small() {
+    let mut context = TraceContext::default();
+    // Add `u`.
+    // TODO(Leo): remove this once `u` is added to the default constants.
+    context.constant(qm31_from_u32s(0, 0, 1, 0));
+    // General (non-broadcast, non-base-field) QM31 constant. All limbs (1, 2, 3, 4) live in the
+    // chain, so no base decomposition is triggered.
+    context.constant(qm31_from_u32s(1, 2, 3, 4));
+    finalize_constants_with_min_base(&mut context, 5);
+
+    // The plus-one chain populates [5]..=[8] for values 2..=5. The QM31 basis allocates [9] = u*u,
+    // [10] = u² - 2 = i, [11] = i*u = iu. The ones vector is built as ([1] + [10]) + ([2] + [11])
+    // yielding wires [12], [13], [14] (unused here, since (1, 2, 3, 4) isn't a broadcast). The
+    // general QM31 constant is then assembled as a + b*i + c*u + d*iu:
+    //   [15] = [5] * [10] (2 * i),   [16] = [1] + [15]      (1 + 2*i)
+    //   [17] = [6] * [2]  (3 * u),   [18] = [7] * [11]      (4 * iu),   [19] = [17] + [18]
+    //   [3] = [16] + [19]            (finally constrain the constant).
+    expect![[r#"
+        [0] = [0] + [0]
+        [2] = [2] + [0]
+        [1] = [1] + [0]
+        [5] = [1] + [1]
+        [6] = [5] + [1]
+        [7] = [6] + [1]
+        [8] = [7] + [1]
+        [12] = [1] + [10]
+        [13] = [2] + [11]
+        [14] = [12] + [13]
+        [16] = [1] + [15]
+        [19] = [17] + [18]
+        [3] = [16] + [19]
+        [10] = [9] - [5]
+        [4] = [2] * [1]
+        [9] = [2] * [2]
+        [11] = [10] * [2]
+        [15] = [5] * [10]
+        [17] = [6] * [2]
+        [18] = [7] * [11]
+        [4] = [2]
+        output [2]
+    "#]]
+    .assert_eq(&format!("{:?}", context.circuit));
+
+    // The reserved Var carries the assembled QM31 value; the partial sums hold the two halves.
+    assert_eq!(context.get(Var { idx: 3 }), qm31_from_u32s(1, 2, 3, 4));
+    assert_eq!(context.get(Var { idx: 16 }), qm31_from_u32s(1, 2, 0, 0));
+    assert_eq!(context.get(Var { idx: 19 }), qm31_from_u32s(0, 0, 3, 4));
+    context.circuit.check_yields();
+    context.validate_circuit();
+}
+
+#[test]
+fn test_mixed_m31_and_qm31_constants_large() {
+    let mut context = TraceContext::default();
+    // Add `u`.
+    // TODO(Leo): remove this once `u` is added to the default constants.
+    context.constant(qm31_from_u32s(0, 0, 1, 0));
+    // Add constants of various types.
+    context.constant(qm31_from_u32s(1000, 2000, 3000, 4000));
+    context.constant(qm31_from_u32s(1, 1, 1, 1));
+    context.constant(qm31_from_u32s(2, 2, 2, 2));
+    context.constant(qm31_from_u32s(666, 666, 666, 666));
+    context.constant(qm31_from_u32s(3456, 0, 0, 0));
+    context.constant(qm31_from_u32s(7890, 0, 0, 0));
+    context.constant(qm31_from_u32s(1234, 2, 3, 4));
+    context.constant(qm31_from_u32s(0, 1234, 0, 0));
+    context.constant(qm31_from_u32s(0, 0, 1234, 0));
+    context.constant(qm31_from_u32s(0, 0, 0, 1234));
+    finalize_constants(&mut context);
+    context.circuit.check_yields();
+    context.validate_circuit();
+}