Skip to main content

spongefish_derive/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2
3use proc_macro::TokenStream;
4use proc_macro2::TokenStream as TokenStream2;
5use quote::quote;
6use syn::{parse_macro_input, parse_quote, Data, DeriveInput, Member, Result, Type};
7
8/// One field of the deriving struct, with the member it is reached by.
9///
10/// [`Member`] is what lets named and tuple structs share every code path in
11/// this crate. `self.#member` reads `self.field` and `self.0` alike, and braced
12/// construction — `Self { #member: .. }` — builds a tuple struct as readily as
13/// a named one. The field-less case falls out for free: `Self {}` constructs a
14/// unit struct.
15struct StructField<'a> {
16    member: Member,
17    ty: &'a Type,
18    /// Set by `#[spongefish(skip)]`: the field is left out of every map, and
19    /// filled with `Default::default()` on the way back.
20    skip: bool,
21}
22
23/// The deriving struct's fields, in declaration order.
24///
25/// Errors if the input is not a struct, or if a field carries a malformed
26/// `#[spongefish(..)]`.
27fn struct_fields<'a>(input: &'a DeriveInput, derive: &str) -> Result<Vec<StructField<'a>>> {
28    let Data::Struct(data) = &input.data else {
29        return Err(syn::Error::new_spanned(
30            &input.ident,
31            format!("{derive} can only be derived for structs"),
32        ));
33    };
34
35    data.fields
36        .iter()
37        .enumerate()
38        .map(|(index, field)| {
39            let member = field
40                .ident
41                .clone()
42                .map_or_else(|| Member::Unnamed(syn::Index::from(index)), Member::Named);
43            Ok(StructField {
44                member,
45                ty: &field.ty,
46                skip: has_skip_attribute(&field.attrs)?,
47            })
48        })
49        .collect()
50}
51
52/// The types that must be bounded for an impl to hold: every field that
53/// the generated maps actually touch, and no others.
54///
55/// A skipped field is untouched by the generated code, so bounding it would
56/// reject types the derive can perfectly well handle.
57fn bounded_types<'a>(fields: &'a [StructField<'a>]) -> Vec<&'a Type> {
58    fields
59        .iter()
60        .filter(|field| !field.skip)
61        .map(|field| field.ty)
62        .collect()
63}
64
65/// Wraps `items` in an impl of `trait_path` for the deriving type, adding a
66/// `#ty: #trait_path` predicate for each of `bounded`.
67fn impl_block(
68    input: &DeriveInput,
69    trait_path: &TokenStream2,
70    bounded: &[&Type],
71    items: &TokenStream2,
72) -> TokenStream2 {
73    let name = &input.ident;
74    let mut generics = input.generics.clone();
75
76    if !bounded.is_empty() {
77        let where_clause = generics.make_where_clause();
78        for ty in bounded {
79            where_clause.predicates.push(parse_quote!(#ty: #trait_path));
80        }
81    }
82
83    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
84    quote! {
85        impl #impl_generics #trait_path for #name #ty_generics #where_clause {
86            #items
87        }
88    }
89}
90
91/// The compile-time width reserved for a field inside the derived `Repr`.
92///
93/// `FromUniform` exposes no associated constant for the length of its `Repr`, and
94/// `AsMut::as_mut(..).len()` — the width the sponge actually fills — is not a
95/// const expression, so the buffer size must be spelled with `size_of`. The
96/// generated `from_uniform` checks at run time that the two agree (see
97/// [`from_uniform_field_expr`]).
98fn field_repr_size(field_type: &Type) -> TokenStream2 {
99    quote! {
100        ::core::mem::size_of::<<#field_type as ::spongefish::FromUniform>::Repr>()
101    }
102}
103
104/// Samples one field out of `bytes`, advancing the shared `offset` cursor.
105///
106/// The width comes from `AsMut::<[u8]>::as_mut(..).len()`.
107///
108/// # Panics
109///
110/// The bounds check will panic if `Repr` is a slice whose length disagrees with its `size_of`.
111fn from_uniform_field_expr(field_type: &Type) -> TokenStream2 {
112    quote! {
113        {
114            let mut field_buf = <#field_type as ::spongefish::FromUniform>::Repr::default();
115            let field_size = ::core::convert::AsMut::<[u8]>::as_mut(&mut field_buf).len();
116            let start = offset;
117            let end = start + field_size;
118            assert!(
119                end <= bytes.len(),
120                "`FromUniform` derive: field representation is wider than the derived buffer; \
121                 `Repr` must satisfy `size_of::<Repr>() == Repr::default().as_mut().len()`"
122            );
123            ::core::convert::AsMut::<[u8]>::as_mut(&mut field_buf)
124                .copy_from_slice(&bytes[start..end]);
125            offset = end;
126            <#field_type as ::spongefish::FromUniform>::from_uniform(field_buf)
127        }
128    }
129}
130
131fn generate_encoding_impl(input: &DeriveInput) -> Result<TokenStream2> {
132    let fields = struct_fields(input, "Encoding")?;
133    let bounded = bounded_types(&fields);
134
135    let field_encodings = fields.iter().filter(|field| !field.skip).map(|field| {
136        let member = &field.member;
137        quote! {
138            output.extend_from_slice(self.#member.encode().as_ref());
139        }
140    });
141
142    let trait_path = quote!(::spongefish::Encoding);
143    Ok(impl_block(
144        input,
145        &trait_path,
146        &bounded,
147        &quote! {
148            fn encode(&self) -> impl AsRef<[u8]> {
149                // Sized up front from the struct's own width. That is only a
150                // hint — a field whose encoding is wider or narrower than
151                // its in-memory size just makes the vector grow or over-
152                // reserve — but for the fixed-width codecs that carry
153                // prover messages it is exact, which turns several
154                // reallocations per message into one allocation.
155                let mut output = ::spongefish::__private::Vec::with_capacity(
156                    ::core::mem::size_of::<Self>(),
157                );
158                #(#field_encodings)*
159                output
160            }
161        },
162    ))
163}
164
165fn generate_from_uniform_impl(input: &DeriveInput) -> Result<TokenStream2> {
166    let fields = struct_fields(input, "FromUniform")?;
167    let bounded = bounded_types(&fields);
168
169    let field_inits = fields.iter().map(|field| {
170        let member = &field.member;
171        if field.skip {
172            return quote!(#member: Default::default(),);
173        }
174        let from_uniform_field = from_uniform_field_expr(field.ty);
175        quote!(#member: #from_uniform_field,)
176    });
177
178    let size_components = bounded.iter().copied().map(field_repr_size);
179    let size_calc = if bounded.is_empty() {
180        quote!(0usize)
181    } else {
182        quote!(#(#size_components)+*)
183    };
184
185    // Fields are read in declaration order from a single `offset` cursor, so
186    // the offsets and the buffer size can never drift apart silently: the final
187    // check pins the total to the buffer length.
188    let body = if bounded.is_empty() {
189        quote! {
190            let _ = buf;
191            Self { #(#field_inits)* }
192        }
193    } else {
194        quote! {
195            // Keep the preimage in its wiping owner instead of copying it
196            // into an unprotected array on the stack.
197            let bytes = buf.as_ref();
198            let mut offset = 0usize;
199            let value = Self { #(#field_inits)* };
200            assert_eq!(
201                offset,
202                bytes.len(),
203                "`FromUniform` derive: field representations do not cover the derived buffer; \
204                 every `Repr` must satisfy `size_of::<Repr>() == Repr::default().as_mut().len()`"
205            );
206            value
207        }
208    };
209
210    let trait_path = quote!(::spongefish::FromUniform);
211    Ok(impl_block(
212        input,
213        &trait_path,
214        &bounded,
215        &quote! {
216            type Repr = ::spongefish::ByteArray<{ #size_calc }>;
217
218            fn from_uniform(buf: Self::Repr) -> Self {
219                #body
220            }
221        },
222    ))
223}
224
225fn generate_from_narg_impl(input: &DeriveInput) -> Result<TokenStream2> {
226    let fields = struct_fields(input, "FromNarg")?;
227    let bounded = bounded_types(&fields);
228
229    let field_inits = fields.iter().map(|field| {
230        let member = &field.member;
231        if field.skip {
232            return quote!(#member: Default::default(),);
233        }
234        let field_type = field.ty;
235        quote! {
236            #member: ::spongefish::NargReader::read::<#field_type>(reader)?,
237        }
238    });
239
240    let trait_path = quote!(::spongefish::FromNarg);
241    Ok(impl_block(
242        input,
243        &trait_path,
244        &bounded,
245        &quote! {
246            fn from_narg(
247                reader: &mut ::spongefish::NargReader<'_>,
248            ) -> ::core::result::Result<Self, ::spongefish::VerificationError> {
249                // A unit struct, or one whose every field is skipped,
250                // reads nothing at all.
251                let _ = &reader;
252                Ok(Self { #(#field_inits)* })
253            }
254        },
255    ))
256}
257
258fn generate_unit_impl(input: &DeriveInput) -> Result<TokenStream2> {
259    let fields = struct_fields(input, "Unit")?;
260    let bounded = bounded_types(&fields);
261
262    let zero_fields = fields.iter().map(|field| {
263        let member = &field.member;
264        if field.skip {
265            return quote!(#member: ::core::default::Default::default(),);
266        }
267        let field_type = field.ty;
268        quote!(#member: <#field_type as ::spongefish::Unit>::ZERO,)
269    });
270
271    let trait_path = quote!(::spongefish::Unit);
272    Ok(impl_block(
273        input,
274        &trait_path,
275        &bounded,
276        &quote! {
277            const ZERO: Self = Self { #(#zero_fields)* };
278        },
279    ))
280}
281
282/// Turns a generated impl into tokens, reporting failure as a `compile_error!`
283/// on the offending span rather than as a proc-macro panic.
284fn expand(generated: Result<TokenStream2>) -> TokenStream {
285    TokenStream::from(generated.unwrap_or_else(syn::Error::into_compile_error))
286}
287
288/// Derive [`Encoding`](https://docs.rs/spongefish/latest/spongefish/trait.Encoding.html) for structs.
289///
290/// Fields marked with `#[spongefish(skip)]` are omitted from the encoding.
291/// Any other `#[spongefish(..)]` form is a compile error.
292///
293/// # Security
294///
295/// Skip only values that are recomputable, or genuinely irrelevant to the statement being proven.
296///
297/// A skipped field is not bound by the Fiat-Shamir transformation, and therefore is never part of
298/// the inputs to the random oracle, nor of the NARG string.
299///
300/// ```
301/// use spongefish::Encoding;
302/// # use spongefish_derive::Encoding;
303///
304/// #[derive(Encoding)]
305/// struct Rgb {
306///     r: u8,
307///     g: u8,
308///     b: u8,
309/// }
310///
311/// let colors = Rgb { r: 1, g: 2, b: 3 };
312/// let data = colors.encode();
313/// assert_eq!(data.as_ref(), [1, 2, 3]);
314/// ```
315#[proc_macro_derive(Encoding, attributes(spongefish))]
316pub fn derive_encoding(input: TokenStream) -> TokenStream {
317    let input = parse_macro_input!(input as DeriveInput);
318    expand(generate_encoding_impl(&input))
319}
320
321/// Derive macro for the [`FromUniform`](https://docs.rs/spongefish/latest/spongefish/trait.FromUniform.html) trait.
322///
323/// Generates an implementation that samples struct fields sequentially from a fixed-size buffer.
324/// Fields can be skipped using `#[spongefish(skip)]`.
325#[proc_macro_derive(FromUniform, attributes(spongefish))]
326pub fn derive_from_uniform(input: TokenStream) -> TokenStream {
327    let input = parse_macro_input!(input as DeriveInput);
328    expand(generate_from_uniform_impl(&input))
329}
330
331/// Derive macro for the [`FromNarg`](https://docs.rs/spongefish/latest/spongefish/trait.FromNarg.html) trait.
332///
333/// Generates an implementation that parses struct fields sequentially from the NARG string.
334/// Fields can be skipped using `#[spongefish(skip)]`.
335/// Parsing returns `VerificationError` on failure; use `NargReader::read` to
336/// reject poisoned readers and automatically record nested field errors.
337#[proc_macro_derive(FromNarg, attributes(spongefish))]
338pub fn derive_from_narg(input: TokenStream) -> TokenStream {
339    let input = parse_macro_input!(input as DeriveInput);
340    expand(generate_from_narg_impl(&input))
341}
342
343/// Derive macro that generates [`Encoding`](https://docs.rs/spongefish/latest/spongefish/trait.Encoding.html),
344/// [`FromUniform`](https://docs.rs/spongefish/latest/spongefish/trait.FromUniform.html), and
345/// [`FromNarg`](https://docs.rs/spongefish/latest/spongefish/trait.FromNarg.html) in one go.
346#[proc_macro_derive(Codec, attributes(spongefish))]
347pub fn derive_codec(input: TokenStream) -> TokenStream {
348    let input = parse_macro_input!(input as DeriveInput);
349    expand((|| {
350        let encoding = generate_encoding_impl(&input)?;
351        let from_uniform = generate_from_uniform_impl(&input)?;
352        let from_narg = generate_from_narg_impl(&input)?;
353        Ok(quote! {
354            #encoding
355            #from_uniform
356            #from_narg
357        })
358    })())
359}
360
361/// Derive [`Unit`](https://docs.rs/spongefish/latest/spongefish/trait.Unit.html) for structs.
362///
363/// ```
364/// use spongefish::Unit;
365/// # use spongefish_derive::Unit;
366///
367/// #[derive(Clone, Unit)]
368/// struct Rgb {
369///     r: u8,
370///     g: u8,
371///     b: u8,
372/// }
373///
374/// assert_eq!((Rgb::ZERO.r, Rgb::ZERO.g, Rgb::ZERO.b), (0, 0, 0));
375/// ```
376#[proc_macro_derive(Unit, attributes(spongefish))]
377pub fn derive_unit(input: TokenStream) -> TokenStream {
378    let input = parse_macro_input!(input as DeriveInput);
379    expand(generate_unit_impl(&input))
380}
381
382/// Whether a field carries `#[spongefish(skip)]`.
383///
384/// # Errors
385///
386/// Any `#[spongefish(..)]` that is not exactly `skip` is a compile error in the
387/// deriving crate, reported on the attribute itself.
388///
389/// # Security
390///
391/// Dropping a field from the encoding map is security-relevant, so it must be
392/// spelled out rather than inferred from an attribute that failed to parse:
393/// `#[spongefish()]` parses zero nested metas, and a misspelling like
394/// `#[spongefish(skpi)]` names no option at all. Neither may be read as "skip",
395/// and neither may be silently ignored. See [`spongefish::Encoding`].
396fn has_skip_attribute(attrs: &[syn::Attribute]) -> Result<bool> {
397    for attr in attrs {
398        if !attr.path().is_ident("spongefish") {
399            continue;
400        }
401
402        let mut skip = false;
403        attr.parse_nested_meta(|meta| {
404            if meta.path.is_ident("skip") {
405                skip = true;
406                Ok(())
407            } else {
408                Err(meta.error("unknown `spongefish` option; expected `skip`"))
409            }
410        })?;
411
412        if !skip {
413            return Err(syn::Error::new_spanned(
414                attr,
415                "empty `#[spongefish(..)]`: write `#[spongefish(skip)]` or remove the attribute",
416            ));
417        }
418        return Ok(true);
419    }
420
421    Ok(false)
422}