1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
|
use proc_macro::TokenStream;
#[path = "../../../guests_macro/src/parse_fn.rs"]
mod parse_fn;
use crate::parse_fn::FunctionDefinition;
/// Parses the `guests/type.txt` type note, created from the guest
/// Returns a tuple of the arguments group and the return type
fn new_fd() -> FunctionDefinition {
FunctionDefinition::new(&include_str!("../../../guests/type.txt").parse().unwrap())
}
static DERIVES: &str = "#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]";
/// Creates an Output type def and three Input structures from the guest
/// type.txt file.
///
/// # Usage
///
/// Inside zkvms_host_io:
///
/// ```rust
/// input_macros::generate_output_type_input_struct!();
/// ```
///
/// # Example output
///
/// ```rust
/// pub type Output = (... ...);
///
/// pub type Return = ...;
///
/// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
/// pub struct PublicInput {
/// pub ...: ...,
/// pub ...: ...,
/// ...
/// }
///
/// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
/// pub struct PrivateInput {
/// pub ...: ...,
/// pub ...: ...,
/// ...
/// }
///
/// #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
/// pub struct Input {
/// pub ...: ...,
/// pub ...: ...,
/// ...
/// }
///
/// // Converts Input to a tuple
/// impl From<Input> for (...) {
/// fn from(input: Input) -> (...) {
/// (
/// input....,
/// input....,
/// ...
/// )
/// }
/// }
/// ```
#[proc_macro]
pub fn generate_output_type_input_struct(_: TokenStream) -> TokenStream {
let fd = new_fd();
let sep = if fd.types().is_empty() { "" } else { ", " };
let output_type = format!(
"pub type Output = ({} {} {});",
fd.grouped_public_types(),
sep,
fd.return_type
)
.to_string();
let return_type = format!("pub type Return = {};", fd.return_type).to_string();
let public_attrs = fd
.public_arguments()
.iter()
.map(|x| format!("pub {x},"))
.collect::<String>();
let public_input_type =
format!("{} pub struct PublicInput {{ {} }}", DERIVES, public_attrs).to_string();
let private_attrs = fd
.private_arguments()
.iter()
.map(|x| format!("pub {x},"))
.collect::<String>();
let private_input_type = format!(
"{} pub struct PrivateInput {{ {} }}",
DERIVES, private_attrs
)
.to_string();
let attrs = fd
.arguments()
.iter()
.map(|x| format!("pub {x},"))
.collect::<String>();
let convertion = fd
.patterns()
.clone()
.iter()
.map(|x| format!("input.{x},"))
.collect::<String>();
let types = fd.grouped_types();
let struct_def = &format!(
"
{DERIVES} pub struct Input {{
{attrs}
}}
impl From<Input> for ({types}) {{
fn from(input: Input) -> ({types}) {{
({convertion})
}}
}}
"
)
.to_string();
(output_type + &return_type + &public_input_type + &private_input_type + &struct_def)
.parse::<TokenStream>()
.unwrap()
}
/// Repeats the given item as many times as fields there are, while replacing
/// all `.yield` occurences with the fields value (field name).
fn foreach_field(item: TokenStream, fields: Vec<TokenStream>) -> TokenStream {
let expr = format!("{}", item);
let mut out = String::new();
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()
}
/// Repeats the given code as many times as fields there are in the Input
/// struct, while replacing all `.yield` occurences with the concrete
/// field name.
#[proc_macro]
pub fn foreach_input_field(item: TokenStream) -> TokenStream {
foreach_field(item, new_fd().patterns().clone())
}
/// Repeats the given code as many times as fields there are in the
/// PublicInput struct, while replacing all `.yield` occurences with the
/// concrete field name.
#[proc_macro]
pub fn foreach_public_input_field(item: TokenStream) -> TokenStream {
foreach_field(item, new_fd().public_patterns().clone())
}
/// Repeats the given code as many times as fields there are in the
/// PrivateInput struct, while replacing all `.yield` occurences with the
/// concrete field name.
#[proc_macro]
pub fn foreach_private_input_field(item: TokenStream) -> TokenStream {
foreach_field(item, new_fd().private_patterns().clone())
}
/// Assuming the `run_info` variable is present, it creates a block with all
/// needed code to properly benchmark the input code, according to all command
/// parameters.
#[proc_macro]
pub fn benchmarkable(item: TokenStream) -> TokenStream {
format!(
r#"
{{
use std::time::Instant;
let mut starts = Vec::new();
let mut ends = Vec::new();
for i in 1..=run_info.runs {{
if run_info.benchmarking {{
starts.push(Instant::now());
}}
{item}
if run_info.benchmarking {{
ends.push(Instant::now());
}}
}}
if run_info.benchmarking {{
zkvms_host_io::emit_benchmark_results(run_info, starts, ends);
}}
}}
"#
)
.parse()
.unwrap()
}
|
