Skip to main content

syn_mid/
func.rs

1// SPDX-License-Identifier: Apache-2.0 OR MIT
2
3// Based on https://github.com/dtolnay/syn/blob/2.0.37/src/item.rs.
4
5use alloc::{boxed::Box, vec::Vec};
6
7use proc_macro2::TokenStream;
8use syn::{
9    Abi, Attribute, Generics, Ident, Lifetime, ReturnType, Token, Type, Visibility,
10    punctuated::Punctuated, token,
11};
12
13use super::{Pat, PatType};
14
15ast_struct! {
16    /// A free-standing function: `fn process(n: usize) -> Result<()> { ...
17    /// }`.
18    pub struct ItemFn {
19        pub attrs: Vec<Attribute>,
20        pub vis: Visibility,
21        pub sig: Signature,
22        pub block: Box<Block>,
23    }
24}
25
26ast_struct! {
27    /// A braced block containing Rust statements.
28    pub struct Block {
29        pub brace_token: token::Brace,
30        /// Statements in a block
31        pub stmts: TokenStream,
32    }
33}
34
35ast_struct! {
36    /// A function signature in a trait or implementation: `unsafe fn
37    /// initialize(&self)`.
38    pub struct Signature {
39        pub constness: Option<Token![const]>,
40        pub asyncness: Option<Token![async]>,
41        pub unsafety: Option<Token![unsafe]>,
42        pub abi: Option<Abi>,
43        pub fn_token: Token![fn],
44        pub ident: Ident,
45        pub generics: Generics,
46        pub paren_token: token::Paren,
47        pub inputs: Punctuated<FnArg, Token![,]>,
48        pub variadic: Option<Variadic>,
49        pub output: ReturnType,
50    }
51}
52
53ast_enum_of_structs! {
54    /// An argument in a function signature: the `n: usize` in `fn f(n: usize)`.
55    pub enum FnArg {
56        /// The `self` argument of an associated method, whether taken by value
57        /// or by reference.
58        Receiver(Receiver),
59
60        /// A function argument accepted by pattern and type.
61        Typed(PatType),
62    }
63}
64
65ast_struct! {
66    /// The `self` argument of an associated method, whether taken by value
67    /// or by reference.
68    pub struct Receiver {
69        pub attrs: Vec<Attribute>,
70        pub reference: Option<(Token![&], Option<Lifetime>)>,
71        pub mutability: Option<Token![mut]>,
72        pub self_token: Token![self],
73        pub colon_token: Option<Token![:]>,
74        pub ty: Box<Type>,
75    }
76}
77
78ast_struct! {
79    /// The variadic argument of a foreign function.
80    pub struct Variadic {
81        pub attrs: Vec<Attribute>,
82        pub pat: Option<(Box<Pat>, Token![:])>,
83        pub dots: Token![...],
84        pub comma: Option<Token![,]>,
85    }
86}
87
88mod parsing {
89    use alloc::{boxed::Box, vec, vec::Vec};
90
91    use syn::{
92        Abi, Attribute, Error, Generics, Ident, Lifetime, Path, ReturnType, Token, Type, TypePath,
93        TypeReference, Visibility, braced, parenthesized,
94        parse::{Parse, ParseStream, Result, discouraged::Speculative as _},
95        punctuated::Punctuated,
96    };
97
98    use super::{Block, FnArg, ItemFn, Receiver, Signature, Variadic};
99    use crate::pat::{Pat, PatType, PatWild};
100
101    impl Parse for Block {
102        fn parse(input: ParseStream<'_>) -> Result<Self> {
103            let content;
104            Ok(Self { brace_token: braced!(content in input), stmts: content.parse()? })
105        }
106    }
107
108    impl Parse for Signature {
109        fn parse(input: ParseStream<'_>) -> Result<Self> {
110            let constness: Option<Token![const]> = input.parse()?;
111            let asyncness: Option<Token![async]> = input.parse()?;
112            let unsafety: Option<Token![unsafe]> = input.parse()?;
113            let abi: Option<Abi> = input.parse()?;
114            let fn_token: Token![fn] = input.parse()?;
115            let ident: Ident = input.parse()?;
116            let mut generics: Generics = input.parse()?;
117
118            let content;
119            let paren_token = parenthesized!(content in input);
120            let (inputs, variadic) = parse_fn_args(&content)?;
121
122            let output: ReturnType = input.parse()?;
123            generics.where_clause = input.parse()?;
124
125            Ok(Self {
126                constness,
127                asyncness,
128                unsafety,
129                abi,
130                fn_token,
131                ident,
132                generics,
133                paren_token,
134                inputs,
135                variadic,
136                output,
137            })
138        }
139    }
140
141    impl Parse for ItemFn {
142        fn parse(input: ParseStream<'_>) -> Result<Self> {
143            let attrs = input.call(Attribute::parse_outer)?;
144            let vis: Visibility = input.parse()?;
145            let sig: Signature = input.parse()?;
146            let block = input.parse()?;
147            Ok(Self { attrs, vis, sig, block: Box::new(block) })
148        }
149    }
150
151    impl Parse for FnArg {
152        fn parse(input: ParseStream<'_>) -> Result<Self> {
153            let allow_variadic = false;
154            let attrs = input.call(Attribute::parse_outer)?;
155            match parse_fn_arg_or_variadic(input, attrs, allow_variadic)? {
156                FnArgOrVariadic::FnArg(arg) => Ok(arg),
157                FnArgOrVariadic::Variadic(_) => unreachable!(),
158            }
159        }
160    }
161
162    enum FnArgOrVariadic {
163        FnArg(FnArg),
164        Variadic(Variadic),
165    }
166
167    fn parse_fn_arg_or_variadic(
168        input: ParseStream<'_>,
169        attrs: Vec<Attribute>,
170        allow_variadic: bool,
171    ) -> Result<FnArgOrVariadic> {
172        let ahead = input.fork();
173        if let Ok(mut receiver) = ahead.parse::<Receiver>() {
174            input.advance_to(&ahead);
175            receiver.attrs = attrs;
176            return Ok(FnArgOrVariadic::FnArg(FnArg::Receiver(receiver)));
177        }
178
179        // Hack to parse pre-2018 syntax in
180        // test/ui/rfc-2565-param-attrs/param-attrs-pretty.rs
181        // because the rest of the test case is valuable.
182        if input.peek(Ident) && input.peek2(Token![<]) {
183            let span = input.fork().parse::<Ident>()?.span();
184            return Ok(FnArgOrVariadic::FnArg(FnArg::Typed(PatType {
185                attrs,
186                pat: Box::new(Pat::Wild(PatWild {
187                    attrs: vec![],
188                    underscore_token: Token![_](span),
189                })),
190                colon_token: Token![:](span),
191                ty: input.parse()?,
192            })));
193        }
194
195        let pat = Box::new(Pat::parse_single(input)?);
196        let colon_token: Token![:] = input.parse()?;
197
198        if allow_variadic {
199            if let Some(dots) = input.parse::<Option<Token![...]>>()? {
200                return Ok(FnArgOrVariadic::Variadic(Variadic {
201                    attrs,
202                    pat: Some((pat, colon_token)),
203                    dots,
204                    comma: None,
205                }));
206            }
207        }
208
209        Ok(FnArgOrVariadic::FnArg(FnArg::Typed(PatType {
210            attrs,
211            pat,
212            colon_token,
213            ty: input.parse()?,
214        })))
215    }
216
217    impl Parse for Receiver {
218        fn parse(input: ParseStream<'_>) -> Result<Self> {
219            let reference = if input.peek(Token![&]) {
220                let ampersand: Token![&] = input.parse()?;
221                let lifetime: Option<Lifetime> = input.parse()?;
222                Some((ampersand, lifetime))
223            } else {
224                None
225            };
226            let mutability: Option<Token![mut]> = input.parse()?;
227            let self_token: Token![self] = input.parse()?;
228            let colon_token: Option<Token![:]> =
229                if reference.is_some() { None } else { input.parse()? };
230            let ty: Type = if colon_token.is_some() {
231                input.parse()?
232            } else {
233                let mut ty = Type::Path(TypePath {
234                    attrs: vec![],
235                    qself: None,
236                    path: Path::from(Ident::new("Self", self_token.span)),
237                });
238                if let Some((ampersand, lifetime)) = reference.as_ref() {
239                    ty = Type::Reference(TypeReference {
240                        attrs: vec![],
241                        and_token: Token![&](ampersand.span),
242                        lifetime: lifetime.clone(),
243                        mutability: mutability.as_ref().map(|m| Token![mut](m.span)),
244                        elem: Box::new(ty),
245                    });
246                }
247                ty
248            };
249            Ok(Self {
250                attrs: vec![],
251                reference,
252                mutability,
253                self_token,
254                colon_token,
255                ty: Box::new(ty),
256            })
257        }
258    }
259
260    fn parse_fn_args(
261        input: ParseStream<'_>,
262    ) -> Result<(Punctuated<FnArg, Token![,]>, Option<Variadic>)> {
263        let mut args = Punctuated::new();
264        let mut variadic = None;
265        let mut has_receiver = false;
266
267        while !input.is_empty() {
268            let attrs = input.call(Attribute::parse_outer)?;
269
270            if let Some(dots) = input.parse::<Option<Token![...]>>()? {
271                variadic = Some(Variadic {
272                    attrs,
273                    pat: None,
274                    dots,
275                    comma: if input.is_empty() { None } else { Some(input.parse()?) },
276                });
277                break;
278            }
279
280            let allow_variadic = true;
281            let arg = match parse_fn_arg_or_variadic(input, attrs, allow_variadic)? {
282                FnArgOrVariadic::FnArg(arg) => arg,
283                FnArgOrVariadic::Variadic(arg) => {
284                    variadic = Some(Variadic {
285                        comma: if input.is_empty() { None } else { Some(input.parse()?) },
286                        ..arg
287                    });
288                    break;
289                }
290            };
291
292            match &arg {
293                FnArg::Receiver(receiver) if has_receiver => {
294                    return Err(Error::new(
295                        receiver.self_token.span,
296                        "unexpected second method receiver",
297                    ));
298                }
299                FnArg::Receiver(receiver) if !args.is_empty() => {
300                    return Err(Error::new(receiver.self_token.span, "unexpected method receiver"));
301                }
302                FnArg::Receiver(_) => has_receiver = true,
303                FnArg::Typed(_) => {}
304            }
305            args.push_value(arg);
306
307            if input.is_empty() {
308                break;
309            }
310
311            let comma: Token![,] = input.parse()?;
312            args.push_punct(comma);
313        }
314
315        Ok((args, variadic))
316    }
317}
318
319mod printing {
320    use proc_macro2::TokenStream;
321    use quote::{ToTokens, TokenStreamExt as _};
322    use syn::{Token, Type};
323
324    use super::{Block, ItemFn, Receiver, Signature, Variadic};
325
326    impl ToTokens for ItemFn {
327        fn to_tokens(&self, tokens: &mut TokenStream) {
328            tokens.append_all(&self.attrs);
329            self.vis.to_tokens(tokens);
330            self.sig.to_tokens(tokens);
331            self.block.to_tokens(tokens);
332        }
333    }
334
335    impl ToTokens for Block {
336        fn to_tokens(&self, tokens: &mut TokenStream) {
337            self.brace_token.surround(tokens, |tokens| {
338                tokens.append_all(self.stmts.clone());
339            });
340        }
341    }
342
343    impl ToTokens for Signature {
344        fn to_tokens(&self, tokens: &mut TokenStream) {
345            self.constness.to_tokens(tokens);
346            self.asyncness.to_tokens(tokens);
347            self.unsafety.to_tokens(tokens);
348            self.abi.to_tokens(tokens);
349            self.fn_token.to_tokens(tokens);
350            self.ident.to_tokens(tokens);
351            self.generics.to_tokens(tokens);
352            self.paren_token.surround(tokens, |tokens| {
353                self.inputs.to_tokens(tokens);
354                if let Some(variadic) = &self.variadic {
355                    if !self.inputs.empty_or_trailing() {
356                        <Token![,]>::default().to_tokens(tokens);
357                    }
358                    variadic.to_tokens(tokens);
359                }
360            });
361            self.output.to_tokens(tokens);
362            self.generics.where_clause.to_tokens(tokens);
363        }
364    }
365
366    impl ToTokens for Receiver {
367        fn to_tokens(&self, tokens: &mut TokenStream) {
368            tokens.append_all(&self.attrs);
369            if let Some((ampersand, lifetime)) = &self.reference {
370                ampersand.to_tokens(tokens);
371                lifetime.to_tokens(tokens);
372            }
373            self.mutability.to_tokens(tokens);
374            self.self_token.to_tokens(tokens);
375            if let Some(colon_token) = &self.colon_token {
376                colon_token.to_tokens(tokens);
377                self.ty.to_tokens(tokens);
378            } else {
379                let consistent = match (&self.reference, &self.mutability, &*self.ty) {
380                    (Some(_), mutability, Type::Reference(ty)) => {
381                        mutability.is_some() == ty.mutability.is_some()
382                            && match &*ty.elem {
383                                Type::Path(ty) => ty.qself.is_none() && ty.path.is_ident("Self"),
384                                _ => false,
385                            }
386                    }
387                    (None, _, Type::Path(ty)) => ty.qself.is_none() && ty.path.is_ident("Self"),
388                    _ => false,
389                };
390                if !consistent {
391                    <Token![:]>::default().to_tokens(tokens);
392                    self.ty.to_tokens(tokens);
393                }
394            }
395        }
396    }
397
398    impl ToTokens for Variadic {
399        fn to_tokens(&self, tokens: &mut TokenStream) {
400            tokens.append_all(&self.attrs);
401            if let Some((pat, colon)) = &self.pat {
402                pat.to_tokens(tokens);
403                colon.to_tokens(tokens);
404            }
405            self.dots.to_tokens(tokens);
406            self.comma.to_tokens(tokens);
407        }
408    }
409}