opendp_tooling/bootstrap/
partial.rs

1use syn::{
2    punctuated::Punctuated, GenericArgument, ItemFn, PathArguments, ReturnType, Signature, Type,
3};
4
5use crate::bootstrap::signature::syn_fnarg_to_syn_pattype;
6
7pub fn generate_partial(mut item_fn: ItemFn) -> Option<ItemFn> {
8    if !supports_partial(&item_fn.sig) {
9        return None;
10    }
11
12    // update function name
13    item_fn.sig.ident = syn::Ident::new(
14        &item_fn.sig.ident.to_string().replacen("make_", "then_", 1),
15        item_fn.sig.ident.span(),
16    );
17
18    // update function arguments
19    let mut inputs = Vec::from_iter(item_fn.sig.inputs.into_iter());
20    let input_domain_arg = inputs.remove(0);
21    let input_metric_arg = inputs.remove(0);
22    item_fn.sig.inputs = Punctuated::from_iter(inputs);
23
24    // update return type
25    let syn::ReturnType::Type(_, fallible_type) = &mut item_fn.sig.output else {
26        return None;
27    };
28    let syn::Type::Path(path) = fallible_type.as_mut() else {
29        return None;
30    };
31    let PathArguments::AngleBracketed(args) = &mut path.path.segments.last_mut()?.arguments else {
32        return None;
33    };
34    let GenericArgument::Type(operator_type) = &mut args.args.first_mut()? else {
35        return None;
36    };
37    let syn::Type::Path(path) = operator_type else {
38        return None;
39    };
40
41    let pathargs = path.path.segments.last()?.clone().arguments;
42
43    let mut operator_type = path.path.segments.last()?.clone();
44    operator_type.ident = syn::Ident::new(
45        format!("Partial{}", operator_type.ident).as_str(),
46        operator_type.ident.span(),
47    );
48    let ret_type: Type = syn::parse_quote!(crate::core::#operator_type);
49    let body_operator_ident = operator_type.ident.clone();
50    item_fn.sig.output =
51        syn::ReturnType::Type(syn::token::RArrow::default(), Box::new(ret_type.clone()));
52
53    let old_block = item_fn.block.clone();
54    // update function body
55    item_fn.block = syn::parse_quote! {{
56        crate::core::#body_operator_ident::#pathargs::new(move |#input_domain_arg, #input_metric_arg| #old_block)
57    }};
58
59    Some(item_fn)
60}
61
62pub fn supports_partial(sig: &Signature) -> bool {
63    if sig.inputs.len() < 2 {
64        return false;
65    }
66
67    if !sig.ident.to_string().starts_with("make_") {
68        return false;
69    }
70
71    let Some((input_domain_type, input_metric_type)) = extract_domain_metric_types(&sig.output)
72    else {
73        return false;
74    };
75
76    let mut inputs = Vec::from_iter(sig.inputs.iter().cloned());
77
78    let Ok(first_arg) = syn_fnarg_to_syn_pattype(inputs.remove(0)) else {
79        return false;
80    };
81    let Ok(second_arg) = syn_fnarg_to_syn_pattype(inputs.remove(0)) else {
82        return false;
83    };
84
85    first_arg.1 == input_domain_type && second_arg.1 == input_metric_type
86}
87
88fn extract_domain_metric_types(output: &ReturnType) -> Option<(Type, Type)> {
89    let syn::ReturnType::Type(_, output_type) = output.clone() else {
90        return None;
91    };
92    let data_type = extract_parameters(*output_type, "Fallible")?
93        .first()?
94        .clone();
95
96    let supporting_types = extract_parameters(data_type.clone(), "Transformation")
97        .or_else(|| extract_parameters(data_type.clone(), "Measurement"))?;
98
99    let [input_domain_type, _, input_metric_type, _] =
100        <[Type; 4]>::try_from(supporting_types).ok()?;
101
102    Some((input_domain_type, input_metric_type))
103}
104
105fn extract_parameters(ty: Type, name: &str) -> Option<Vec<Type>> {
106    let syn::Type::Path(path) = ty else {
107        return None;
108    };
109
110    let segment = path.path.segments.last()?;
111    if segment.ident != name {
112        return None;
113    }
114    let syn::PathArguments::AngleBracketed(args) = &segment.arguments else {
115        return None;
116    };
117    args.args
118        .iter()
119        .map(|arg| match arg {
120            syn::GenericArgument::Type(ty) => Some(ty.clone()),
121            _ => None,
122        })
123        .collect()
124}
opendp_tooling/bootstrap/partial.rs

opendp_tooling/bootstrap/
partial.rs
