xref: /aosp_15_r20/external/pigweed/pw_format/rust/pw_format_example_macro.rs (revision 61c4878ac05f98d0ceed94b57d316916de578985)
1 // Copyright 2023 The Pigweed Authors
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
4 // use this file except in compliance with the License. You may obtain a copy of
5 // the License at
6 //
7 //     https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
11 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12 // License for the specific language governing permissions and limitations under
13 // the License.
14 
15 use proc_macro::TokenStream;
16 use pw_format::macros::{
17     generate, Arg, FormatAndArgsFlavor, FormatMacroGenerator, FormatParams,
18     PrintfFormatStringParser, Result,
19 };
20 use quote::quote;
21 use syn::{
22     parse::{Parse, ParseStream},
23     parse_macro_input, Expr, Token,
24 };
25 
26 type TokenStream2 = proc_macro2::TokenStream;
27 
28 // Declare a struct to hold our proc macro arguments.
29 #[derive(Debug)]
30 struct MacroArgs {
31     prefix: Expr,
32     // Select printf style format strings.
33     format_and_args: FormatAndArgsFlavor<PrintfFormatStringParser>,
34 }
35 
36 // Implement `Parse` for our argument struct.
37 impl Parse for MacroArgs {
parse(input: ParseStream) -> syn::parse::Result<Self>38     fn parse(input: ParseStream) -> syn::parse::Result<Self> {
39         // Our prefix argument comes first argument and ends with a `,`.
40         let prefix: Expr = input.parse()?;
41         input.parse::<Token![,]>()?;
42 
43         // Prase the remaining arguments as a format string and arguments.
44         let format_and_args: FormatAndArgsFlavor<_> = input.parse()?;
45 
46         Ok(MacroArgs {
47             prefix,
48             format_and_args,
49         })
50     }
51 }
52 
53 // Our generator struct needs to track the prefix as well as the code
54 // fragments we've generated.
55 struct Generator {
56     prefix: Expr,
57     code_fragments: Vec<TokenStream2>,
58 }
59 
60 impl Generator {
new(prefix: Expr) -> Self61     pub fn new(prefix: Expr) -> Self {
62         Self {
63             prefix,
64             code_fragments: Vec::new(),
65         }
66     }
67 }
68 
69 // This toy example implements the generator by calling `format!()` at runtime.
70 impl FormatMacroGenerator for Generator {
71     // Wrap all our fragments in boilerplate and return the code.
finalize(self) -> Result<TokenStream2>72     fn finalize(self) -> Result<TokenStream2> {
73         // Create locally scoped alias so we can refer to them in `quote!()`.
74         let prefix = self.prefix;
75         let code_fragments = self.code_fragments;
76 
77         Ok(quote! {
78           {
79             // Initialize the result string with our prefix.
80             let mut result = String::new();
81             result.push_str(#prefix);
82 
83             // Enumerate all our code fragments.
84             #(#code_fragments);*;
85 
86             // Return the resulting string
87             result
88           }
89         })
90     }
91 
92     // Handles a string embedded in the format string.  This is different from
93     // `string_conversion` which is used to handle a string referenced by the
94     // format string (i.e. "%s'").
string_fragment(&mut self, string: &str) -> Result<()>95     fn string_fragment(&mut self, string: &str) -> Result<()> {
96         self.code_fragments.push(quote! {
97             result.push_str(#string);
98         });
99         Ok(())
100     }
101 
102     // This example ignores display type and width.
integer_conversion( &mut self, _params: &FormatParams, _signed: bool, _type_width: u8, expression: Arg, ) -> Result<()>103     fn integer_conversion(
104         &mut self,
105         _params: &FormatParams,
106         _signed: bool,
107         _type_width: u8, // in bits
108         expression: Arg,
109     ) -> Result<()> {
110         self.code_fragments.push(quote! {
111             result.push_str(&format!("{}", #expression));
112         });
113         Ok(())
114     }
115 
string_conversion(&mut self, expression: Arg) -> Result<()>116     fn string_conversion(&mut self, expression: Arg) -> Result<()> {
117         self.code_fragments.push(quote! {
118             result.push_str(&format!("{}", #expression));
119         });
120         Ok(())
121     }
122 
char_conversion(&mut self, expression: Arg) -> Result<()>123     fn char_conversion(&mut self, expression: Arg) -> Result<()> {
124         self.code_fragments.push(quote! {
125             result.push_str(&format!("{}", #expression));
126         });
127         Ok(())
128     }
129 
untyped_conversion(&mut self, expression: Arg, _params: &FormatParams) -> Result<()>130     fn untyped_conversion(&mut self, expression: Arg, _params: &FormatParams) -> Result<()> {
131         self.code_fragments.push(quote! {
132             result.push_str(&format!("{}", #expression));
133         });
134         Ok(())
135     }
136 }
137 
138 // Lastly we declare our proc macro.
139 #[proc_macro]
example_macro(tokens: TokenStream) -> TokenStream140 pub fn example_macro(tokens: TokenStream) -> TokenStream {
141     // Parse our proc macro arguments.
142     let input = parse_macro_input!(tokens as MacroArgs);
143 
144     // Create our generator.
145     let generator = Generator::new(input.prefix);
146 
147     // Call into `generate()` to handle the code generation.
148     match generate(generator, input.format_and_args.into()) {
149         Ok(token_stream) => token_stream.into(),
150         Err(e) => e.to_compile_error().into(),
151     }
152 }
153