From 623669af847e592c08ef892ed5151b2330c1e518 Mon Sep 17 00:00:00 2001
From: Chao Ma <mc8305@gmail.com>
Date: Mon, 9 Sep 2024 10:32:51 +0700
Subject: [PATCH] add lookup records for mockprover

---
 ceno_zkvm/src/instructions/riscv/blt.rs | 10 ++++--
 ceno_zkvm/src/scheme/mock_prover.rs     | 41 +++++++++++++++++++++++++
 ceno_zkvm/src/uint/arithmetic.rs        |  2 +-
 3 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/ceno_zkvm/src/instructions/riscv/blt.rs b/ceno_zkvm/src/instructions/riscv/blt.rs
index fc0181d4b..323c7726b 100644
--- a/ceno_zkvm/src/instructions/riscv/blt.rs
+++ b/ceno_zkvm/src/instructions/riscv/blt.rs
@@ -235,9 +235,12 @@ impl<E: ExtensionField> Instruction<E> for BltInstruction {
 mod test {
     use super::*;
     use ceno_emul::StepRecord;
+    use ff::Field;
     use goldilocks::GoldilocksExt2;
     use itertools::Itertools;
     use multilinear_extensions::mle::IntoMLEs;
+    use std::array;
+    use rand::rngs::OsRng;
 
     use crate::{circuit_builder::ConstraintSystem, scheme::mock_prover::MockProver};
 
@@ -257,6 +260,8 @@ mod test {
         )
         .unwrap();
 
+        let rng = OsRng;
+        let challenges: [GoldilocksExt2;2] = array::from_fn(|_| GoldilocksExt2::random(rng));
         MockProver::run(
             &mut circuit_builder,
             &raw_witin
@@ -265,9 +270,8 @@ mod test {
                 .into_iter()
                 .map(|v| v.into())
                 .collect_vec(),
-            None,
-        )
-        .expect_err("lookup will fail");
+            Some(challenges),
+        ).unwrap();
         Ok(())
     }
 
diff --git a/ceno_zkvm/src/scheme/mock_prover.rs b/ceno_zkvm/src/scheme/mock_prover.rs
index f116afb01..096f93425 100644
--- a/ceno_zkvm/src/scheme/mock_prover.rs
+++ b/ceno_zkvm/src/scheme/mock_prover.rs
@@ -266,6 +266,8 @@ impl<'a, E: ExtensionField> MockProver<E> {
         let mut table_vec = vec![];
         load_u5_table(&mut table_vec, cb, challenge);
         load_u16_table(&mut table_vec, cb, challenge);
+        load_and_table(&mut table_vec, cb, challenge);
+        load_ltu_table(&mut table_vec, cb, challenge);
 
         // Lookup expressions
         for (expr, name) in cb
@@ -350,6 +352,45 @@ pub fn load_u16_table<E: ExtensionField>(
     }
 }
 
+
+pub fn load_and_table<E: ExtensionField>(
+    t_vec: &mut Vec<E>,
+    cb: &CircuitBuilder<E>,
+    challenge: [E; 2],
+) {
+    for i in 0..(1 << 16) {
+        let a = i >> 8;
+        let b = i & 0b0111_1111;
+        let c = a & b;
+        let rlc_record = cb.rlc_chip_record(vec![
+            Expression::Constant(E::BaseField::from(ROMType::And as u64)),
+            i.into(),
+            c.into(),
+        ]);
+        let rlc_record = eval_by_expr(&[], &challenge, &rlc_record);
+        t_vec.push(rlc_record);
+    }
+}
+
+pub fn load_ltu_table<E: ExtensionField>(
+    t_vec: &mut Vec<E>,
+    cb: &CircuitBuilder<E>,
+    challenge: [E; 2],
+) {
+    for i in 0..(1 << 16) {
+        let a = i >> 8;
+        let b = i & 0xFF;
+        let c = (a < b) as usize;
+        let rlc_record = cb.rlc_chip_record(vec![
+            Expression::Constant(E::BaseField::from(ROMType::Ltu as u64)),
+            i.into(),
+            c.into(),
+        ]);
+        let rlc_record = eval_by_expr(&[], &challenge, &rlc_record);
+        t_vec.push(rlc_record);
+    }
+}
+
 #[allow(unused_imports)]
 #[cfg(test)]
 mod tests {
diff --git a/ceno_zkvm/src/uint/arithmetic.rs b/ceno_zkvm/src/uint/arithmetic.rs
index caeb59273..fe6dc88ef 100644
--- a/ceno_zkvm/src/uint/arithmetic.rs
+++ b/ceno_zkvm/src/uint/arithmetic.rs
@@ -254,7 +254,7 @@ impl<const M: usize, E: ExtensionField> UInt<M, 8, E> {
         circuit_builder.lookup_and_byte(
             high_limb_no_msb.expr(),
             high_limb.clone(),
-            Expression::from(1 << 7),
+            Expression::from(0xFF),
         )?;
 
         let inv_128 = F::from(128).invert().unwrap();