From 6e3ef85fd7b9e2e97583105a28a5a863527cfe59 Mon Sep 17 00:00:00 2001
From: Kamen Mladenov <kamen@syndamia.com>
Date: Tue, 28 Jan 2025 13:04:02 +0200
Subject: feat(zkvms_host_io): Generate input type from string and read input
 data from TOML file

---
 zkvms_host_io/Cargo.toml              |  3 +++
 zkvms_host_io/input_macros/Cargo.toml |  7 +++++++
 zkvms_host_io/input_macros/src/lib.rs | 36 +++++++++++++++++++++++++++++++++++
 zkvms_host_io/src/lib.rs              | 15 ++++++++++-----
 4 files changed, 56 insertions(+), 5 deletions(-)
 create mode 100644 zkvms_host_io/input_macros/Cargo.toml
 create mode 100644 zkvms_host_io/input_macros/src/lib.rs

diff --git a/zkvms_host_io/Cargo.toml b/zkvms_host_io/Cargo.toml
index b4d6e8e..371758e 100644
--- a/zkvms_host_io/Cargo.toml
+++ b/zkvms_host_io/Cargo.toml
@@ -4,5 +4,8 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+input_macros = { path = "./input_macros" }
 clap = { version = "4.5.26", features = ["derive"] }
 num-traits = "0.2.19"
+serde = { version = "1.0.217", features = ["derive"] }
+toml = "0.8.19"
diff --git a/zkvms_host_io/input_macros/Cargo.toml b/zkvms_host_io/input_macros/Cargo.toml
new file mode 100644
index 0000000..88cfb35
--- /dev/null
+++ b/zkvms_host_io/input_macros/Cargo.toml
@@ -0,0 +1,7 @@
+[package]
+name = "input_macros"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+proc_macro = true
diff --git a/zkvms_host_io/input_macros/src/lib.rs b/zkvms_host_io/input_macros/src/lib.rs
new file mode 100644
index 0000000..51fd613
--- /dev/null
+++ b/zkvms_host_io/input_macros/src/lib.rs
@@ -0,0 +1,36 @@
+use proc_macro::TokenStream;
+
+#[path = "../../../guests_macro/src/parse_fn.rs"]
+mod parse_fn;
+use crate::parse_fn::{ args_split, args_divide, group_streams };
+
+
+fn get_args() -> TokenStream {
+    "(graph: Vec<Vec<bool>>, colors: u32, coloring: Vec<Vec<u32>>,)".parse::<TokenStream>().unwrap()
+}
+
+#[proc_macro]
+pub fn generate_input_struct(_: TokenStream) -> TokenStream {
+    let args = &get_args();
+    let all_args = args_split(&args);
+
+    let mut struct_def = "#[derive(Debug, Serialize, Deserialize)] pub struct Input {".to_string();
+    for arg in all_args {
+        struct_def += &format!("pub {arg},");
+    }
+
+    let (patterns, types) = args_divide(&args);
+    let types = group_streams(&types);
+    struct_def += &format!("}}
+        impl From<Input> for {types} {{
+            fn from(input: Input) -> {types} {{
+                (
+    ");
+
+    for field in patterns {
+        struct_def += &format!("input.{field},");
+    }
+    struct_def += ") } }";
+
+    struct_def.parse::<TokenStream>().unwrap()
+}
diff --git a/zkvms_host_io/src/lib.rs b/zkvms_host_io/src/lib.rs
index 29723db..74c3992 100644
--- a/zkvms_host_io/src/lib.rs
+++ b/zkvms_host_io/src/lib.rs
@@ -1,11 +1,16 @@
 use clap::{Parser, ValueEnum};
 use num_traits::NumCast;
+use serde::{ Serialize, Deserialize };
+pub use input_macros::foreach_input_field;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 struct Cli {
     /// What the ZKVM is going to do
     run_type: RunType,
+
+    #[arg(default_value = "./public_input.toml")]
+    public_input: String,
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
@@ -21,15 +26,15 @@ pub struct RunWith<T> {
     pub input: T,
 }
 
-fn read_input() -> (Vec<Vec<bool>>, u32, Vec<Vec<u32>>) {
-    include!(env!("INPUTS"))
-}
+input_macros::generate_input_struct!();
 
-pub fn read_args() -> RunWith<(Vec<Vec<bool>>, u32, Vec<Vec<u32>>)> {
+pub fn read_args() -> RunWith<Input> {
     let cli = Cli::parse();
 
+    let val: Input = toml::from_str(&std::fs::read_to_string(cli.public_input).unwrap()).unwrap();
+
     RunWith {
         run_type: cli.run_type,
-        input: read_input(),
+        input: val,
     }
 }
-- 
cgit v1.2.3