ab_contracts_macros_impl/contract/
update.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, parse_quote};
8
9pub(super) fn process_update_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::Update, 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                if receiver.reference.is_none() {
42                    return Err(Error::new(
43                        fn_sig.span(),
44                        "`#[update]` can't consume `Self`, use `&self` or `&mut self` instead",
45                    ));
46                }
47
48                methods_details.process_state_arg_rw(input_span, &receiver.ty)?;
49            }
50            FnArg::Typed(pat_type) => {
51                let mut attrs = pat_type.attrs.extract_if(.., |attr| match &attr.meta {
52                    Meta::Path(path) => {
53                        path.leading_colon.is_none()
54                            && path.segments.len() == 1
55                            && supported_attrs.contains_key(&path.segments[0].ident)
56                    }
57                    Meta::List(_meta_list) => false,
58                    Meta::NameValue(_meta_name_value) => false,
59                });
60
61                let Some(attr) = attrs.next() else {
62                    return Err(Error::new(
63                        input_span,
64                        "Each `#[update]` argument (except `&self`/`&mut self`) must be \
65                        annotated with exactly one of: `#[env]`, `#[slot]`, `#[input]` or \
66                        `#[output]`, in that order",
67                    ));
68                };
69
70                if let Some(next_attr) = attrs.take(1).next() {
71                    return Err(Error::new(
72                        next_attr.span(),
73                        "Each `#[update]` argument (except `&self`/`&mut self`) must be \
74                        annotated with exactly one of: `#[env]`, `#[slot]`, `#[input]` or \
75                        `#[output]`, in that order",
76                    ));
77                }
78
79                let processor = supported_attrs
80                    .get(&attr.path().segments[0].ident)
81                    .expect("Matched above to be one of the supported attributes; qed");
82
83                processor(&mut methods_details, input_span, &*pat_type)?;
84            }
85        }
86    }
87
88    methods_details.process_return(&fn_sig.output)?;
89
90    let guest_ffi = methods_details.generate_guest_ffi(fn_sig, trait_name)?;
91    let trait_ext_components =
92        methods_details.generate_trait_ext_components(fn_sig, fn_attrs, trait_name)?;
93
94    contract_details.methods.push(Method {
95        original_ident: fn_sig.ident.clone(),
96        methods_details,
97    });
98
99    Ok(MethodOutput {
100        guest_ffi,
101        trait_ext_components,
102    })
103}
104
105pub(super) fn process_update_fn_definition(
106    trait_name: &Ident,
107    fn_sig: &mut Signature,
108    fn_attrs: &[Attribute],
109    contract_details: &mut ContractDetails,
110) -> Result<MethodOutput, Error> {
111    let mut methods_details =
112        MethodDetails::new(MethodType::Update, Type::Path(parse_quote! { #trait_name }));
113
114    for input in &mut fn_sig.inputs {
115        let input_span = input.span();
116        // TODO: Moving this outside of the loop causes confusing lifetime issues
117        // TODO: Simplify `-> Result<(), Error>` to `-> _` once https://github.com/rust-lang/rust/issues/135864 is resolved
118        let supported_attrs = HashMap::<_, fn(_, _, _) -> Result<(), Error>>::from_iter([
119            (format_ident!("env"), MethodDetails::process_env_arg_rw as _),
120            (
121                format_ident!("input"),
122                MethodDetails::process_input_arg as _,
123            ),
124            (
125                format_ident!("output"),
126                MethodDetails::process_output_arg as _,
127            ),
128        ]);
129
130        match input {
131            FnArg::Receiver(_receiver) => {
132                return Err(Error::new(
133                    fn_sig.span(),
134                    "`#[update]` in trait definition can't access state through `Self` directly, \
135                    call implementation method through `Env` instead if necessary",
136                ));
137            }
138            FnArg::Typed(pat_type) => {
139                let mut attrs = pat_type.attrs.extract_if(.., |attr| match &attr.meta {
140                    Meta::Path(path) => {
141                        path.leading_colon.is_none()
142                            && path.segments.len() == 1
143                            && supported_attrs.contains_key(&path.segments[0].ident)
144                    }
145                    Meta::List(_meta_list) => false,
146                    Meta::NameValue(_meta_name_value) => false,
147                });
148
149                let Some(attr) = attrs.next() else {
150                    return Err(Error::new(
151                        input_span,
152                        "Each `#[update]` argument must be annotated with exactly one of: \
153                        `#[env]`, `#[input]` or `#[output]`, in that order",
154                    ));
155                };
156
157                if let Some(next_attr) = attrs.take(1).next() {
158                    return Err(Error::new(
159                        next_attr.span(),
160                        "Each `#[update]` argument must be annotated with exactly one of: \
161                        `#[env]`, `#[input]` or `#[output]`, in that order",
162                    ));
163                }
164
165                let processor = supported_attrs
166                    .get(&attr.path().segments[0].ident)
167                    .expect("Matched above to be one of the supported attributes; qed");
168
169                processor(&mut methods_details, input_span, &*pat_type)?;
170            }
171        }
172    }
173
174    methods_details.process_return(&fn_sig.output)?;
175
176    let guest_ffi = methods_details.generate_guest_trait_ffi(fn_sig, Some(trait_name))?;
177    let trait_ext_components =
178        methods_details.generate_trait_ext_components(fn_sig, fn_attrs, Some(trait_name))?;
179
180    contract_details.methods.push(Method {
181        original_ident: fn_sig.ident.clone(),
182        methods_details,
183    });
184
185    Ok(MethodOutput {
186        guest_ffi,
187        trait_ext_components,
188    })
189}