From c85fe44a171ce6200c5d52c5b9142be7cf60572c Mon Sep 17 00:00:00 2001
From: Roshan <19766713+rpalakkal@users.noreply.github.com>
Date: Tue, 20 Feb 2024 15:24:48 -0500
Subject: [PATCH] fix: fix len vec input (#11)

* fix: fix len vec input

* fix: use F::from instead of F::from_u128
---
 circuit/src/input/flatten.rs   | 24 ++++++++++++++----------
 circuit/src/input/raw_input.rs | 30 ++++++++++++++++++++++++++----
 sdk/src/api.rs                 | 17 +++++++++++++++++
 3 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/circuit/src/input/flatten.rs b/circuit/src/input/flatten.rs
index 8320579..0ab83e8 100644
--- a/circuit/src/input/flatten.rs
+++ b/circuit/src/input/flatten.rs
@@ -11,26 +11,30 @@ pub trait InputFlatten<T: Copy>: Sized {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
-pub struct FixLenVec<T: Copy, const N: usize> {
-    pub vec: Vec<T>,
+pub struct FixLenVec<T: Copy + Default, const N: usize>(pub Vec<T>);
+
+impl<T: Copy + Default, const N: usize> Default for FixLenVec<T, N> {
+    fn default() -> Self {
+        Self(vec![T::default(); N])
+    }
 }
 
-impl<T: Copy, const N: usize> FixLenVec<T, N> {
+impl<T: Copy + Default, const N: usize> FixLenVec<T, N> {
     pub fn new(vec: Vec<T>) -> anyhow::Result<Self> {
         if vec.len() != N {
             anyhow::bail!("Invalid input length: {} != {}", vec.len(), N);
         }
-        Ok(FixLenVec { vec })
+        Ok(FixLenVec(vec))
     }
 
     pub fn into_inner(self) -> Vec<T> {
-        self.vec
+        self.0
     }
 }
 
-impl<T: Copy, const N: usize> From<Vec<T>> for FixLenVec<T, N> {
+impl<T: Copy + Default, const N: usize> From<Vec<T>> for FixLenVec<T, N> {
     fn from(vec: Vec<T>) -> Self {
-        Self { vec }
+        Self(vec)
     }
 }
 
@@ -46,14 +50,14 @@ macro_rules! check_input_length {
     };
 }
 
-impl<T: Copy, const N: usize> InputFlatten<T> for FixLenVec<T, N> {
+impl<T: Copy + Default, const N: usize> InputFlatten<T> for FixLenVec<T, N> {
     const NUM_FE: usize = N;
     fn flatten_vec(&self) -> Vec<T> {
-        self.vec.clone()
+        self.0.clone()
     }
     fn unflatten(vec: Vec<T>) -> Result<Self> {
         check_input_length!(vec);
-        Ok(FixLenVec { vec })
+        Ok(FixLenVec(vec))
     }
 }
 
diff --git a/circuit/src/input/raw_input.rs b/circuit/src/input/raw_input.rs
index 5f16fff..62484d1 100644
--- a/circuit/src/input/raw_input.rs
+++ b/circuit/src/input/raw_input.rs
@@ -56,18 +56,40 @@ impl<F: Field, const N: usize> RawInput<F> for [u8; N] {
     fn convert(&self) -> Self::FEType<F> {
         let mut res = [F::ZERO; N];
         for i in 0..N {
-            res[i] = F::from_u128(self[i] as u128);
+            res[i] = F::from(self[i] as u64);
+        }
+        res
+    }
+}
+
+impl<F: Field, const N: usize> RawInput<F> for [u64; N] {
+    type FEType<T: Copy> = [T; N];
+    fn convert(&self) -> Self::FEType<F> {
+        let mut res = [F::ZERO; N];
+        for i in 0..N {
+            res[i] = F::from(self[i]);
+        }
+        res
+    }
+}
+
+impl<F: Field, const N: usize> RawInput<F> for [usize; N] {
+    type FEType<T: Copy> = [T; N];
+    fn convert(&self) -> Self::FEType<F> {
+        let mut res = [F::ZERO; N];
+        for i in 0..N {
+            res[i] = F::from(self[i] as u64);
         }
         res
     }
 }
 
 impl<F: Field, const N: usize> RawInput<F> for FixLenVec<usize, N> {
-    type FEType<T: Copy> = Vec<T>;
+    type FEType<T: Copy> = [T; N];
     fn convert(&self) -> Self::FEType<F> {
-        let mut res = Vec::new();
+        let mut res = [F::ZERO; N];
         for i in 0..N {
-            res.push(F::from_u128(self.vec[i] as u128));
+            res[i] = F::from(self.0[i] as u64);
         }
         res
     }
diff --git a/sdk/src/api.rs b/sdk/src/api.rs
index d4736b2..09f2d70 100644
--- a/sdk/src/api.rs
+++ b/sdk/src/api.rs
@@ -6,6 +6,7 @@ use axiom_circuit::{
         halo2_base::{gates::RangeChip, safe_types::SafeTypeChip, AssignedValue, Context},
         keccak::promise::{KeccakFixLenCall, KeccakVarLenCall},
         rlc::circuit::builder::RlcCircuitBuilder,
+        utils::uint_to_bytes_be,
     },
     subquery::caller::SubqueryCaller,
     utils::{from_hi_lo, to_hi_lo},
@@ -78,6 +79,22 @@ impl<'a> AxiomAPI<'a> {
         to_hi_lo(ctx, self.range, val)
     }
 
+    /// Decomposes a `AssignedValue<Fr>` into bytes, in big-endian, and returns the bytes.
+    ///
+    /// * `uint` - The `AssignedValue<Fr>` object to convert.
+    /// * `num_bytes` - The number of bytes in `uint`.
+    pub fn to_bytes_be(
+        &mut self,
+        uint: AssignedValue<Fr>,
+        num_bytes: usize,
+    ) -> Vec<AssignedValue<Fr>> {
+        let ctx = self.builder.base.main(0);
+        uint_to_bytes_be(ctx, self.range, &uint, num_bytes)
+            .iter()
+            .map(|x| *x.as_ref())
+            .collect()
+    }
+
     /// Returns an [Account] builder given block number and address.
     ///
     /// * `block_number` - The block number as an `AssignedValue<Fr>`.