From 68c333f2a823021705845266404df8ea6bae50be Mon Sep 17 00:00:00 2001
From: Jaremy Creechley <creechley@gmail.com>
Date: Mon, 11 Dec 2023 13:13:59 -0700
Subject: [PATCH] rework funcs

---
 src/ffi.rs            |  3 +-
 src/storage_proofs.rs | 65 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 66 insertions(+), 2 deletions(-)

diff --git a/src/ffi.rs b/src/ffi.rs
index bea22c0..cecef83 100644
--- a/src/ffi.rs
+++ b/src/ffi.rs
@@ -234,13 +234,14 @@ mod tests {
 
         // Serialize the value types to an array pointer
         write_value(&mut buf, &data).unwrap();
-        let mut rd = &buf[..];
+        let mut rd: &[u8] = &buf[..];
         
         let args = read_value(&mut rd).unwrap();
 
         assert!(Value::is_map(&args));
         assert!(Value::is_array(&args["chunks"]));
         assert!(Value::is_array(&args["chunks"][0]));
+
         let mut arg_chunks: Vec<Vec<U256>> = Vec::new();
 
         // deserialize the data back into u256's
diff --git a/src/storage_proofs.rs b/src/storage_proofs.rs
index 603995e..8033fea 100644
--- a/src/storage_proofs.rs
+++ b/src/storage_proofs.rs
@@ -10,6 +10,30 @@ use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Read};
 use ark_std::rand::rngs::ThreadRng;
 use ruint::aliases::U256;
 
+use rmpv;
+use rmpv::decode::read_value;
+
+type CircomBuilderParams = CircomBuilder<ark_ec::bn::Bn<ark_bn254::Parameters>>;
+
+const EXT_ID_U256_LE: i8 = 50;
+const EXT_ID_U256_BE: i8 = 51;
+
+fn decode_u256(val: &rmpv::Value) -> Result<U256, String> {
+    match val {
+        rmpv::Value::Ext(id, val) => {
+            match *id {
+                EXT_ID_U256_LE =>
+                    match U256::try_from_le_slice(val) {
+                        Some(i) => Ok(i),
+                        None => Err("error parsing 256".to_string()),
+                    }
+                num => return Err(format!("unhandled ext id {}", num)),
+            }
+        },
+        _ => return Err("expected ext mpack kind".to_string()),
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct StorageProofs {
     builder: CircomBuilder<Bn254>,
@@ -83,9 +107,48 @@ impl StorageProofs {
 
     pub fn proof_build_inputs(
         &mut self,
+        mut inputs: &[u8]
     ) -> Result<CircomCircuit<ark_ec::bn::Bn<ark_bn254::Parameters>>, String> {
 
-        let mut builder = self.builder.clone();
+        let values: rmpv::Value = read_value(&mut inputs).map_err(|e| e.to_string())?;
+        let args: &Vec<(rmpv::Value, rmpv::Value)> = match values.as_map() {
+            Some(args) => args,
+            None => return Err("args must be a map of string to arrays".to_string()),
+        };
+
+        let mut builder: CircomBuilderParams = self.builder.clone();
+        for (key, val) in args {
+            let name = match key.as_str() {
+                Some(n) => n,
+                None => return Err(format!("expected string value")),
+            };
+            match val {
+                rmpv::Value::Array(vals) => {
+                    // add a (name, Vec<u256>) or (name, Vev<Vec<u256>>) arrays
+                    if vals.len() > 0 && vals[0].is_array() {
+                        for inner_val in vals {
+                            match inner_val.as_array() {
+                                Some(inner_vals) => {
+                                    for val in inner_vals {
+                                        builder.push_input(name, decode_u256(val)?);
+                                    }
+                                },
+                                None => todo!(),
+                            }
+                        }
+                    } else {
+                        for val in vals {
+                            builder.push_input(name, decode_u256(val)?);
+                        }
+                    }
+                },
+                rmpv::Value::Ext(_, _) => {
+                    // directly add a (name,u256) arg pair 
+                    builder.push_input(name, decode_u256(val)?);
+                },
+                _ => return Err("unhandled argument kind".to_string()),
+            }
+        }
 
         // vec of vecs is flattened, since wasm expects a contiguous array in memory