feat(zkvms_host_io/input_macros): Implement forach_input_field macro

1 files changed, 13 insertions, 0 deletions

diff --git a/zkvms_host_io/input_macros/src/lib.rs b/zkvms_host_io/input_macros/src/lib.rs
index 51fd613..2585653 100644
--- a/zkvms_host_io/input_macros/src/lib.rs
+++ b/zkvms_host_io/input_macros/src/lib.rs
@@ -34,3 +34,16 @@ pub fn generate_input_struct(_: TokenStream) -> TokenStream {
 
     struct_def.parse::<TokenStream>().unwrap()
 }
+
+#[proc_macro]
+pub fn foreach_input_field(item: TokenStream) -> TokenStream {
+    let arg_patterns = args_divide(&get_args()).0;
+
+    let expr = format!("{}", item);
+    let mut out = String::new();
+    for field in arg_patterns {
+        // Unquoted yield is a keyword, so it is not allowed as field name
+        out += &expr.replace(".yield", &format!(".{field}"));
+    }
+    out.parse::<TokenStream>().unwrap()
+}