babybear test for addi/logici

Author: hero78119

ceno_zkvm/src/instructions/riscv/arith_imm.rs

@@ -26,34 +26,40 @@ mod test {
         structs::ProgramParams,
     };
     use ceno_emul::{Change, InsnKind, PC_STEP_SIZE, StepRecord, encode_rv32};
-    use ff_ext::GoldilocksExt2;
+    use ff_ext::{BabyBearExt4, ExtensionField, GoldilocksExt2};
     use gkr_iop::circuit_builder::DebugIndex;
 
     #[test]
-    fn test_opcode_addi_v1() {
-        test_opcode_addi(1000, 1003, 3);
-        test_opcode_addi(1000, 997, -3);
+    fn test_opcode_addi() {
+        let cases = vec![
+            (1000, 1003, 3), // positive immediate
+            (1000, 997, -3), // negative immediate
+        ];
+
+        for &(rs1, expected, imm) in &cases {
+            test_opcode_addi_internal::<GoldilocksExt2>(rs1, expected, imm);
+            #[cfg(feature = "u16limb_circuit")]
+            test_opcode_addi_internal::<BabyBearExt4>(rs1, expected, imm);
+        }
     }
 
-    fn test_opcode_addi(rs1: u32, rd: u32, imm: i32) {
-        let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
+    fn test_opcode_addi_internal<E: ExtensionField>(rs1: u32, rd: u32, imm: i32) {
+        let mut cs = ConstraintSystem::<E>::new(|| "riscv");
         let mut cb = CircuitBuilder::new(&mut cs);
         let config = cb
             .namespace(
                 || "addi",
                 |cb| {
-                    let config = AddiInstruction::<GoldilocksExt2>::construct_circuit(
-                        cb,
-                        &ProgramParams::default(),
-                    );
+                    let config =
+                        AddiInstruction::<E>::construct_circuit(cb, &ProgramParams::default());
                     Ok(config)
                 },
             )
             .unwrap()
             .unwrap();
 
         let insn_code = encode_rv32(InsnKind::ADDI, 2, 0, 4, imm);
-        let (raw_witin, lkm) = AddiInstruction::<GoldilocksExt2>::assign_instances(
+        let (raw_witin, lkm) = AddiInstruction::<E>::assign_instances(
             &config,
             cb.cs.num_witin as usize,
             cb.cs.num_structural_witin as usize,

ceno_zkvm/src/instructions/riscv/logic_imm/logic_imm_circuit.rs

@@ -124,7 +124,7 @@ impl<E: ExtensionField> LogicConfig<E> {
 #[cfg(test)]
 mod test {
     use ceno_emul::{Change, InsnKind, PC_STEP_SIZE, StepRecord, encode_rv32u};
-    use ff_ext::GoldilocksExt2;
+    use ff_ext::{BabyBearExt4, ExtensionField, GoldilocksExt2};
     use gkr_iop::circuit_builder::DebugIndex;
 
     use crate::{
@@ -150,26 +150,55 @@ mod test {
 
     #[test]
     fn test_opcode_andi() {
-        verify::<AndiOp>("basic", 0x0000_0011, 3, 0x0000_0011 & 3);
-        verify::<AndiOp>("zero result", 0x0000_0100, 3, 0x0000_0100 & 3);
-        verify::<AndiOp>("negative imm", TEST, NEG, TEST & NEG);
+        let cases = vec![
+            ("basic", 0x0000_0011, 3, 0x0000_0011 & 3),
+            ("zero result", 0x0000_0100, 3, 0x0000_0100 & 3),
+            ("negative imm", TEST, NEG, TEST & NEG),
+        ];
+
+        for &(name, rs1, imm, expected) in &cases {
+            verify::<AndiOp, GoldilocksExt2>(name, rs1, imm, expected);
+            #[cfg(feature = "u16limb_circuit")]
+            verify::<AndiOp, BabyBearExt4>(name, rs1, imm, expected);
+        }
     }
 
     #[test]
     fn test_opcode_ori() {
-        verify::<OriOp>("basic", 0x0000_0011, 3, 0x0000_0011 | 3);
-        verify::<OriOp>("basic2", 0x0000_0100, 3, 0x0000_0100 | 3);
-        verify::<OriOp>("negative imm", TEST, NEG, TEST | NEG);
+        let cases = vec![
+            ("basic",       0x0000_0011, 3, 0x0000_0011 | 3),
+            ("basic2",      0x0000_0100, 3, 0x0000_0100 | 3),
+            ("negative imm", TEST,       NEG, TEST | NEG),
+        ];
+
+        for &(name, rs1, imm, expected) in &cases {
+            verify::<OriOp, GoldilocksExt2>(name, rs1, imm, expected);
+            #[cfg(feature = "u16limb_circuit")]
+            verify::<OriOp, BabyBearExt4>(name, rs1, imm, expected);
+        }
     }
 
     #[test]
     fn test_opcode_xori() {
-        verify::<XoriOp>("basic", 0x0000_0011, 3, 0x0000_0011 ^ 3);
-        verify::<XoriOp>("non-overlap", 0x0000_0100, 3, 0x0000_0100 ^ 3);
-        verify::<XoriOp>("negative imm", TEST, NEG, TEST ^ NEG);
+        let cases = vec![
+            ("basic",        0x0000_0011, 3, 0x0000_0011 ^ 3),
+            ("non-overlap",  0x0000_0100, 3, 0x0000_0100 ^ 3),
+            ("negative imm", TEST,       NEG, TEST ^ NEG),
+        ];
+
+        for &(name, rs1, imm, expected) in &cases {
+            verify::<XoriOp, GoldilocksExt2>(name, rs1, imm, expected);
+            #[cfg(feature = "u16limb_circuit")]
+            verify::<XoriOp, BabyBearExt4>(name, rs1, imm, expected);
+        }
     }
 
-    fn verify<I: LogicOp>(name: &'static str, rs1_read: u32, imm: u32, expected_rd_written: u32) {
+    fn verify<I: LogicOp, E: ExtensionField>(
+        name: &'static str,
+        rs1_read: u32,
+        imm: u32,
+        expected_rd_written: u32,
+    ) {
         let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
         let mut cb = CircuitBuilder::new(&mut cs);