zksync_cs_derive 0.32.11

Derive macro for boojum crate
Documentation 
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

use proc_macro2::{Span, TokenStream};
use proc_macro_error::abort_call_site;
use quote::quote;
use syn::{parse_macro_input, punctuated::Punctuated, token::Comma, DeriveInput, GenericParam, Generics, Type, WhereClause};

use crate::utils::*;

const BOUND_ATTR_NAME: &'static str = "WitnessHookBound";

pub(crate) fn derive_witness_hook(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let derived_input = parse_macro_input!(input as DeriveInput);
    let DeriveInput { ident, generics, data, attrs, .. } = derived_input.clone();

    let mut struct_initializations = TokenStream::new();
    let mut field_selections = TokenStream::new();
    let mut field_callings = TokenStream::new();

    let extra_bound = if let Some(bound) = fetch_attr_from_list(BOUND_ATTR_NAME, &attrs) {
        let bound = syn::parse_str::<WhereClause>(&bound).expect("must parse bound as WhereClause");

        Some(bound)
    } else {
        None
    };

    let bound = merge_where_clauses(generics.where_clause.clone(), extra_bound);

    match data {
        syn::Data::Struct(ref struct_data) => match struct_data.fields {
            syn::Fields::Named(ref named_fields) => {
                for field in named_fields.named.iter() {
                    let field_ident = field.ident.clone().expect("should have a field elem ident");
                    match field.ty {
                        Type::Array(ref _array_ty) => {
                            let field_select = quote! {
                                let #field_ident = WitnessHookable::<F>::witness_hook(&self.#field_ident, cs);
                            };
                            field_selections.extend(field_select);
                            let field_calling = quote! {
                                let #field_ident = (#field_ident)()?;
                            };
                            field_callings.extend(field_calling);
                        }
                        Type::Path(_) => {
                            let field_select = quote! {
                                let #field_ident = WitnessHookable::<F>::witness_hook(&self.#field_ident, cs);
                            };
                            field_selections.extend(field_select);
                            let field_calling = quote! {
                                let #field_ident = (#field_ident)()?;
                            };
                            field_callings.extend(field_calling);
                        }
                        _ => abort_call_site!("only array and path types are allowed"),
                    };

                    let init_field = quote! {
                        #field_ident,
                    };

                    struct_initializations.extend(init_field);
                }
            }
            _ => abort_call_site!("only named fields are allowed!"),
        },
        _ => abort_call_site!("only struct types are allowed!"),
    }

    let comma = Comma(Span::call_site());

    let field_generic_param = syn::parse_str::<GenericParam>(&"F: SmallField").unwrap();
    let has_engine_param = has_proper_small_field_parameter(&generics.params, &field_generic_param);
    if has_engine_param == false {
        panic!("Expected to have `F: SmallField` somewhere in bounds");
    }

    // add CS to func generic params
    let mut function_generic_params = Punctuated::new();
    let cs_generic_param = syn::parse_str::<GenericParam>(&"CS: ConstraintSystem<F>").unwrap();
    function_generic_params.push(cs_generic_param.clone());
    function_generic_params.push_punct(comma.clone());

    let type_params_of_allocated_struct = get_type_params_from_generics(&generics, &comma);

    let function_generics = Generics {
        lt_token: Some(syn::token::Lt(Span::call_site())),
        params: function_generic_params,
        gt_token: Some(syn::token::Gt(Span::call_site())),
        where_clause: None,
    };

    let expanded = quote! {
        impl #generics WitnessHookable<F> for #ident<#type_params_of_allocated_struct> #bound {
            fn witness_hook #function_generics(
                &self,
                cs: &CS,
            ) -> Box<dyn FnOnce() -> Option<Self::Witness> + 'static> {
                #field_selections;

                Box::new(
                    move || {
                        #field_callings;

                        Some(
                            Self::Witness {
                                #struct_initializations
                            }
                        )
                    }
                )
            }
        }
    };

    proc_macro::TokenStream::from(expanded)
}