From a673cfb1c33c35d4e151329f882a00b12a03da14 Mon Sep 17 00:00:00 2001 From: Kamen Mladenov Date: Fri, 31 Jan 2025 12:58:02 +0200 Subject: feat: Add output type to zkvms_host_io macro --- zkvms_host_io/input_macros/src/lib.rs | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) (limited to 'zkvms_host_io/input_macros/src/lib.rs') 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::().unwrap() +fn get_types() -> (TokenStream, TokenStream) { + let types: Vec<&str> = include_str!("../../../guests/type.txt") + .split('\n') + .collect(); + (types[0].parse::().unwrap(), types[1].parse::().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::().unwrap() + (output_type + &struct_def).parse::().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(); -- cgit v1.2.3