feat(zkvms_host_io): Make Output type contain public inputs

2 files changed, 19 insertions, 1 deletions

diff --git a/zkvms_host_io/input_macros/Cargo.toml b/zkvms_host_io/input_macros/Cargo.toml
index 88cfb35..a8cb035 100644
--- a/zkvms_host_io/input_macros/Cargo.toml
+++ b/zkvms_host_io/input_macros/Cargo.toml
@@ -5,3 +5,6 @@ edition = "2021"
 
 [lib]
 proc_macro = true
+
+[dependencies]
+toml = "0.8.19"
diff --git a/zkvms_host_io/input_macros/src/lib.rs b/zkvms_host_io/input_macros/src/lib.rs
index f313a1e..5184192 100644
--- a/zkvms_host_io/input_macros/src/lib.rs
+++ b/zkvms_host_io/input_macros/src/lib.rs
@@ -14,7 +14,22 @@ fn get_types() -> (TokenStream, TokenStream) {
 #[proc_macro]
 pub fn generate_output_type_input_struct(_: TokenStream) -> TokenStream {
     let (args, ret) = get_types();
-    let output_type = format!("pub type Output = {};", ret).to_string();
+    let (patterns, types) = args_divide(&args);
+
+    let mut patterns = patterns
+        .iter()
+        .map(|x| x.to_string());
+    let public_inputs = toml::from_str::<toml::Table>(
+            include_str!(concat!(env!("INPUTS_DIR"), "/default_public_input.toml"))
+        )
+        .unwrap();
+    let mut commitment = String::new();
+    for input in public_inputs.keys() {
+        if let Some(index) = patterns.clone().position(|x| x == *input) {
+            commitment += &format!("{}, ", types[index]);
+        }
+    }
+    let output_type = format!("pub type Output = ({} {});", commitment, ret).to_string();
 
     let all_args = args_split(&args);