feat: Add output type to zkvms_host_io macro

3 files changed, 15 insertions, 8 deletions

diff --git a/guests_macro/src/lib.rs b/guests_macro/src/lib.rs
index f87e8a5..0ab7ec2 100644
--- a/guests_macro/src/lib.rs
+++ b/guests_macro/src/lib.rs
@@ -8,7 +8,8 @@ pub fn proving_entrypoint(_: TokenStream, mut item: TokenStream) -> TokenStream
 
     // Put the file in zkVMs-benchmarks/guests/
     let mut output = File::create("../type.txt").unwrap();
-    write!(output, "{}", args);
+    writeln!(output, "{}", args);
+    write!(output, "{}", ret);
 
     item.extend(format!("#[macro_export]
         macro_rules! entrypoint_expr {{
diff --git a/zkvms_host_io/input_macros/src/lib.rs b/zkvms_host_io/input_macros/src/lib.rs
index 84c2792..f313a1e 100644
--- a/zkvms_host_io/input_macros/src/lib.rs
+++ b/zkvms_host_io/input_macros/src/lib.rs
@@ -4,13 +4,18 @@ use proc_macro::TokenStream;
 mod parse_fn;
 use crate::parse_fn::{ args_split, args_divide, group_streams };
 
-fn get_args() -> TokenStream {
-    include_str!("../../../guests/type.txt").parse::<TokenStream>().unwrap()
+fn get_types() -> (TokenStream, TokenStream) {
+    let types: Vec<&str> = include_str!("../../../guests/type.txt")
+        .split('\n')
+        .collect();
+    (types[0].parse::<TokenStream>().unwrap(), types[1].parse::<TokenStream>().unwrap())
 }
 
 #[proc_macro]
-pub fn generate_input_struct(_: TokenStream) -> TokenStream {
-    let args = &get_args();
+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 all_args = args_split(&args);
 
     let mut struct_def = "#[derive(Debug, Serialize, Deserialize)] pub struct Input {".to_string();
@@ -31,12 +36,13 @@ pub fn generate_input_struct(_: TokenStream) -> TokenStream {
     }
     struct_def += ") } }";
 
-    struct_def.parse::<TokenStream>().unwrap()
+    (output_type + &struct_def).parse::<TokenStream>().unwrap()
 }
 
 #[proc_macro]
 pub fn foreach_input_field(item: TokenStream) -> TokenStream {
-    let arg_patterns = args_divide(&get_args()).0;
+    let (args, _) = get_types();
+    let arg_patterns = args_divide(&args).0;
 
     let expr = format!("{}", item);
     let mut out = String::new();
diff --git a/zkvms_host_io/src/lib.rs b/zkvms_host_io/src/lib.rs
index c029975..d00a1b1 100644
--- a/zkvms_host_io/src/lib.rs
+++ b/zkvms_host_io/src/lib.rs
@@ -48,7 +48,7 @@ impl<T> RunWith<T> {
     }
 }
 
-input_macros::generate_input_struct!();
+input_macros::generate_output_type_input_struct!();
 
 pub fn read_args() -> RunWith<Input> {
     let cli = Cli::parse();