feat(zkvms_host_io): Implement foreach_private_input_field and foreach_public_input_field macros

2 files changed, 37 insertions, 7 deletions

diff --git a/zkvms_host_io/input_macros/src/lib.rs b/zkvms_host_io/input_macros/src/lib.rs
index 49cff80..868c755 100644
--- a/zkvms_host_io/input_macros/src/lib.rs
+++ b/zkvms_host_io/input_macros/src/lib.rs
@@ -65,16 +65,46 @@ pub fn generate_output_type_input_struct(_: TokenStream) -> TokenStream {
     (output_type + &public_input_type + &private_input_type + &struct_def).parse::<TokenStream>().unwrap()
 }
 
-#[proc_macro]
-pub fn foreach_input_field(item: TokenStream) -> TokenStream {
-    let (args, _) = get_types();
-    let arg_patterns = args_divide(&args).0;
-
+fn foreach_field(item: TokenStream, fields: Vec<TokenStream>) -> TokenStream {
     let expr = format!("{}", item);
     let mut out = String::new();
-    for field in arg_patterns {
+    for field in fields {
         // Unquoted yield is a keyword, so it is not allowed as field name
         out += &expr.replace(".yield", &format!(".{field}"));
     }
     out.parse::<TokenStream>().unwrap()
 }
+
+#[proc_macro]
+pub fn foreach_input_field(item: TokenStream) -> TokenStream {
+    let (args, _) = get_types();
+    let arg_patterns = args_divide(&args).0;
+
+    foreach_field(item, arg_patterns)
+}
+
+#[proc_macro]
+pub fn foreach_public_input_field(item: TokenStream) -> TokenStream {
+    let (args, _) = get_types();
+
+    let public_inputs = toml::from_str::<toml::Table>(
+            include_str!(concat!(env!("INPUTS_DIR"), "/default_public_input.toml"))
+        )
+        .unwrap();
+    let public_patterns = args_divide_public(&args, &public_inputs.keys().collect()).0.0;
+
+    foreach_field(item, public_patterns)
+}
+
+#[proc_macro]
+pub fn foreach_private_input_field(item: TokenStream) -> TokenStream {
+    let (args, _) = get_types();
+
+    let public_inputs = toml::from_str::<toml::Table>(
+            include_str!(concat!(env!("INPUTS_DIR"), "/default_public_input.toml"))
+        )
+        .unwrap();
+    let private_patterns = args_divide_public(&args, &public_inputs.keys().collect()).1.0;
+
+    foreach_field(item, private_patterns)
+}
diff --git a/zkvms_host_io/src/lib.rs b/zkvms_host_io/src/lib.rs
index 2624352..abf9edf 100644
--- a/zkvms_host_io/src/lib.rs
+++ b/zkvms_host_io/src/lib.rs
@@ -3,7 +3,7 @@ use num_traits::NumCast;
 use serde::{ Serialize, Deserialize };
 use env_file_reader::read_str;
 use std::{env, option::Option, fs::read_to_string, collections::HashMap};
-pub use input_macros::foreach_input_field;
+pub use input_macros::{ foreach_input_field, foreach_public_input_field, foreach_private_input_field };
 
 static DEFAULT_PUBLIC_INPUT: &str = include_str!(concat!(env!("INPUTS_DIR"), "/default_public_input.toml"));
 static DEFAULT_PRIVATE_INPUT: &str = include_str!(concat!(env!("INPUTS_DIR"), "/default_private_input.toml"));