From d4eb57eee03e2fc48e77fd3b620e3cd45f39ae6d Mon Sep 17 00:00:00 2001
From: "sm.wu" <hero78119@gmail.com>
Date: Thu, 5 Sep 2024 12:17:32 +0800
Subject: [PATCH] cache limbs to avoid conversion multiple time

---
 ceno_zkvm/src/instructions/riscv/addsub.rs | 11 +++----
 ceno_zkvm/src/uint.rs                      | 37 ++++++++++++++--------
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index 2fea09a06..e3c5df248 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -156,17 +156,17 @@ impl<E: ExtensionField> Instruction<E> for AddInstruction {
         // TODO use fields from step
         set_val!(instance, config.pc, 1);
         set_val!(instance, config.ts, 2);
-        let addend_0 = UIntValue(step.rs1().unwrap().value);
-        let addend_1 = UIntValue(step.rs2().unwrap().value);
+        let addend_0 = UIntValue::new(step.rs1().unwrap().value);
+        let addend_1 = UIntValue::new(step.rs2().unwrap().value);
         config
             .prev_rd_value
             .assign_limbs(instance, [0, 0].iter().map(E::BaseField::from).collect());
         config
             .addend_0
-            .assign_limbs(instance, addend_0.as_u16_fields());
+            .assign_limbs(instance, addend_0.u16_fields());
         config
             .addend_1
-            .assign_limbs(instance, addend_1.as_u16_fields());
+            .assign_limbs(instance, addend_1.u16_fields());
         let carries = addend_0.add_u16_carries(&addend_1);
         config.outcome.assign_carries(
             instance,
@@ -234,9 +234,6 @@ impl<E: ExtensionField> Instruction<E> for SubInstruction {
 
 #[cfg(test)]
 mod test {
-
-    use std::u32;
-
     use ceno_emul::{ReadOp, StepRecord};
     use goldilocks::GoldilocksExt2;
     use itertools::Itertools;
diff --git a/ceno_zkvm/src/uint.rs b/ceno_zkvm/src/uint.rs
index 971f3da89..8c9d3bfe9 100644
--- a/ceno_zkvm/src/uint.rs
+++ b/ceno_zkvm/src/uint.rs
@@ -459,10 +459,22 @@ impl<E: ExtensionField, const M: usize, const C: usize> ToExpr<E> for UInt<M, C,
     }
 }
 
-pub struct UIntValue<T: Into<u64> + Copy>(pub T);
+pub struct UIntValue<T: Into<u64> + Copy> {
+    #[allow(dead_code)]
+    val: T,
+    pub limbs: Vec<u16>,
+}
 
 // TODO generalize to support non 16 bit limbs
+// TODO optimize api with fixed size array
 impl<T: Into<u64> + Copy> UIntValue<T> {
+    pub fn new(val: T) -> Self {
+        UIntValue::<T> {
+            val,
+            limbs: Self::split_to_u16(val),
+        }
+    }
+
     fn split_to_u16(value: T) -> Vec<u16> {
         let mut limbs = Vec::new();
         let value: u64 = value.into(); // Convert to u64 for generality
@@ -477,23 +489,19 @@ impl<T: Into<u64> + Copy> UIntValue<T> {
         limbs
     }
 
-    pub fn as_u16_limbs(&self) -> Vec<u16> {
-        Self::split_to_u16(self.0)
+    pub fn as_u16_limbs(&self) -> &[u16] {
+        &self.limbs
     }
 
-    pub fn as_u16_fields<F: SmallField>(&self) -> Vec<F> {
-        Self::split_to_u16(self.0)
-            .into_iter()
-            .map(|v| F::from(v as u64))
-            .collect_vec()
+    pub fn u16_fields<F: SmallField>(&self) -> Vec<F> {
+        self.limbs.iter().map(|v| F::from(*v as u64)).collect_vec()
     }
 
     pub fn add_u16_carries(&self, rhs: &Self) -> Vec<bool> {
-        self.as_u16_limbs()
-            .into_iter()
-            .zip(rhs.as_u16_limbs())
-            .fold(vec![], |mut acc, (a_limb, b_limb)| {
-                let (a, b) = a_limb.overflowing_add(b_limb);
+        self.as_u16_limbs().iter().zip(rhs.as_u16_limbs()).fold(
+            vec![],
+            |mut acc, (a_limb, b_limb)| {
+                let (a, b) = a_limb.overflowing_add(*b_limb);
                 if let Some(prev_carry) = acc.last() {
                     let (_, d) = a.overflowing_add(*prev_carry as u16);
                     acc.push(b || d);
@@ -501,7 +509,8 @@ impl<T: Into<u64> + Copy> UIntValue<T> {
                     acc.push(b);
                 }
                 acc
-            })
+            },
+        )
     }
 }