ab_contracts_macros_impl/contract/
init.rs

1use crate::contract::method::{MethodDetails, MethodType};
2use crate::contract::{ContractDetails, Method, MethodOutput};
3use proc_macro2::Ident;
4use quote::format_ident;
5use std::collections::HashMap;
6use syn::spanned::Spanned;
7use syn::{Attribute, Error, FnArg, Meta, Signature, Type};
8
9pub(super) fn process_init_fn(
10    self_type: Type,
11    _trait_name: Option<&Ident>,
12    fn_sig: &mut Signature,
13    fn_attrs: &[Attribute],
14    contract_details: &mut ContractDetails,
15) -> Result<MethodOutput, Error> {
16    let mut methods_details = MethodDetails::new(MethodType::Init, self_type);
17
18    for input in &mut fn_sig.inputs {
19        let input_span = input.span();
20        // TODO: Moving this outside of the loop causes confusing lifetime issues
21        // TODO: Simplify `-> Result<(), Error>` to `-> _` once https://github.com/rust-lang/rust/issues/135864 is resolved
22        let supported_attrs = HashMap::<_, fn(_, _, _) -> Result<(), Error>>::from_iter([
23            (format_ident!("env"), MethodDetails::process_env_arg_rw as _),
24            (format_ident!("tmp"), MethodDetails::process_tmp_arg as _),
25            (
26                format_ident!("slot"),
27                MethodDetails::process_slot_arg_rw as _,
28            ),
29            (
30                format_ident!("input"),
31                MethodDetails::process_input_arg as _,
32            ),
33            (
34                format_ident!("output"),
35                MethodDetails::process_output_arg as _,
36            ),
37        ]);
38
39        match input {
40            FnArg::Receiver(_receiver) => {
41                return Err(Error::new(
42                    fn_sig.span(),
43                    "`#[init]` must return `Self`, not take it as an argument",
44                ));
45            }
46            FnArg::Typed(pat_type) => {
47                let mut attrs = pat_type.attrs.extract_if(.., |attr| match &attr.meta {
48                    Meta::Path(path) => {
49                        path.leading_colon.is_none()
50                            && path.segments.len() == 1
51                            && supported_attrs.contains_key(&path.segments[0].ident)
52                    }
53                    Meta::List(_meta_list) => false,
54                    Meta::NameValue(_meta_name_value) => false,
55                });
56
57                let Some(attr) = attrs.next() else {
58                    return Err(Error::new(
59                        input_span,
60                        "Each `#[init]` argument must be annotated with exactly one of: `#[env]` \
61                        or `#[input]`, in that order",
62                    ));
63                };
64
65                if let Some(next_attr) = attrs.take(1).next() {
66                    return Err(Error::new(
67                        next_attr.span(),
68                        "Each `#[init]` argument must be annotated with exactly one of: `#[env]` \
69                        or `#[input]`, in that order",
70                    ));
71                }
72
73                let processor = supported_attrs
74                    .get(&attr.path().segments[0].ident)
75                    .expect("Matched above to be one of the supported attributes; qed");
76
77                processor(&mut methods_details, input_span, &*pat_type)?;
78            }
79        }
80    }
81
82    methods_details.process_return(&fn_sig.output)?;
83
84    let guest_ffi = methods_details.generate_guest_ffi(fn_sig, None)?;
85    let trait_ext_components =
86        methods_details.generate_trait_ext_components(fn_sig, fn_attrs, None)?;
87
88    contract_details.methods.push(Method {
89        original_ident: fn_sig.ident.clone(),
90        methods_details,
91    });
92
93    Ok(MethodOutput {
94        guest_ffi,
95        trait_ext_components,
96    })
97}