Skip to main content

codegen/
rust_source.rs

1//! Typed Rust token rendering for generated source files.
2
3use proc_macro2::{Span, TokenStream};
4use syn::{Ident, LitFloat, LitInt, LitStr, Path};
5
6/// A structurally attached comment rendered before a generated Rust file.
7#[derive(Clone, Debug)]
8pub(crate) enum FileComment {
9    InnerDoc(Box<str>),
10    Line(Box<str>),
11}
12
13impl FileComment {
14    pub(crate) fn inner_doc(text: impl Into<Box<str>>) -> Result<Self, RustSourceError> {
15        Self::new(text, Self::InnerDoc)
16    }
17
18    pub(crate) fn line(text: impl Into<Box<str>>) -> Result<Self, RustSourceError> {
19        Self::new(text, Self::Line)
20    }
21
22    fn new(
23        text: impl Into<Box<str>>,
24        constructor: impl FnOnce(Box<str>) -> Self,
25    ) -> Result<Self, RustSourceError> {
26        let text = text.into();
27
28        if text.contains(['\r', '\n']) {
29            return Err(RustSourceError::invalid_comment());
30        }
31
32        Ok(constructor(text))
33    }
34}
35
36/// Failure to construct a complete formatted Rust source file.
37#[derive(Debug, thiserror::Error)]
38#[error(transparent)]
39pub(crate) struct RustSourceError(#[from] RustSourceErrorKind);
40
41#[derive(Debug, thiserror::Error)]
42enum RustSourceErrorKind {
43    #[error("generated Rust did not parse")]
44    Parse {
45        #[source]
46        source: syn::Error,
47    },
48    #[error("generated Rust comment contains a line break")]
49    InvalidComment,
50}
51
52impl RustSourceError {
53    fn invalid_comment() -> Self {
54        RustSourceErrorKind::InvalidComment.into()
55    }
56}
57
58/// Parses a complete token stream as a Rust file and formats it in-process.
59pub(crate) fn render_rust(
60    tokens: TokenStream,
61    preamble: &[FileComment],
62) -> Result<String, RustSourceError> {
63    let file =
64        syn::parse2::<syn::File>(tokens).map_err(|source| RustSourceErrorKind::Parse { source })?;
65    let formatted = prettyplease::unparse(&file);
66    let formatted_file =
67        syn::parse_file(&formatted).map_err(|source| RustSourceErrorKind::Parse { source })?;
68    let formatted = apply_statement_padding(&formatted, &formatted_file);
69    let preamble_len: usize = preamble
70        .iter()
71        .map(|comment| match comment {
72            FileComment::InnerDoc(text) => text.len() + 5,
73            FileComment::Line(text) => text.len() + 4,
74        })
75        .sum();
76    let mut source = String::with_capacity(preamble_len + formatted.len());
77
78    for comment in preamble {
79        match comment {
80            FileComment::InnerDoc(_) => source.push_str("//! "),
81            FileComment::Line(_) => source.push_str("// "),
82        }
83
84        let text = match comment {
85            FileComment::InnerDoc(text) | FileComment::Line(text) => text,
86        };
87
88        source.push_str(text);
89        source.push('\n');
90    }
91
92    source.push_str(&formatted);
93
94    Ok(source)
95}
96
97fn apply_statement_padding(source: &str, file: &syn::File) -> String {
98    use syn::visit::Visit as _;
99
100    #[derive(Default)]
101    struct PaddingVisitor {
102        lines: std::collections::BTreeSet<usize>,
103    }
104
105    impl<'ast> syn::visit::Visit<'ast> for PaddingVisitor {
106        fn visit_block(&mut self, i: &'ast syn::Block) {
107            use syn::spanned::Spanned as _;
108
109            for (index, statement) in i.stmts.iter().enumerate().skip(1) {
110                let previous = &i.stmts[index - 1];
111                let is_tail =
112                    index + 1 == i.stmts.len() && matches!(statement, syn::Stmt::Expr(_, None));
113                let is_return = statement_expression(statement)
114                    .is_some_and(|expression| matches!(expression, syn::Expr::Return(_)));
115                let is_multiline_control =
116                    statement_expression(statement).is_some_and(multiline_control_expression);
117                let previous_is_multiline_control =
118                    statement_expression(previous).is_some_and(multiline_control_expression);
119                let follows_let_run = matches!(previous, syn::Stmt::Local(_))
120                    && !matches!(statement, syn::Stmt::Local(_));
121
122                if is_tail
123                    || is_return
124                    || is_multiline_control
125                    || previous_is_multiline_control
126                    || follows_let_run
127                {
128                    self.lines.insert(statement.span().start().line);
129                }
130            }
131
132            syn::visit::visit_block(self, i);
133        }
134    }
135
136    fn statement_expression(statement: &syn::Stmt) -> Option<&syn::Expr> {
137        match statement {
138            syn::Stmt::Expr(expression, _) => Some(expression),
139            syn::Stmt::Local(_) | syn::Stmt::Item(_) | syn::Stmt::Macro(_) => None,
140        }
141    }
142
143    fn multiline_control_expression(expression: &syn::Expr) -> bool {
144        use syn::spanned::Spanned as _;
145
146        matches!(
147            expression,
148            syn::Expr::ForLoop(_)
149                | syn::Expr::If(_)
150                | syn::Expr::Loop(_)
151                | syn::Expr::Match(_)
152                | syn::Expr::While(_)
153        ) && expression.span().start().line != expression.span().end().line
154    }
155
156    let mut visitor = PaddingVisitor::default();
157
158    visitor.visit_file(file);
159
160    let mut padded = String::with_capacity(source.len() + visitor.lines.len());
161    let mut previous_line_was_blank = false;
162
163    for (index, line) in source.split_inclusive('\n').enumerate() {
164        let line_number = index + 1;
165
166        if visitor.lines.contains(&line_number) && !previous_line_was_blank {
167            padded.push('\n');
168        }
169
170        padded.push_str(line);
171        previous_line_was_blank = line.trim().is_empty();
172    }
173
174    padded
175}
176
177pub(crate) fn rust_ident(value: &str) -> Result<Ident, syn::Error> {
178    syn::parse_str(value)
179}
180
181pub(crate) fn rust_path(value: &str) -> Result<Path, syn::Error> {
182    syn::parse_str(value)
183}
184
185#[derive(Clone, Copy, Debug, Default)]
186pub(crate) struct ExpressionRequirements {
187    pub(crate) data: bool,
188    pub(crate) budget: bool,
189}
190
191pub(crate) fn expression_requirements(
192    tokens: TokenStream,
193) -> Result<ExpressionRequirements, syn::Error> {
194    use syn::visit::Visit as _;
195
196    #[derive(Default)]
197    struct RequirementVisitor {
198        requirements: ExpressionRequirements,
199    }
200
201    impl<'ast> syn::visit::Visit<'ast> for RequirementVisitor {
202        fn visit_expr_path(&mut self, i: &'ast syn::ExprPath) {
203            if i.qself.is_none() && i.path.segments.len() == 1 {
204                let ident = &i.path.segments[0].ident;
205
206                self.requirements.data |= ident == "data";
207                self.requirements.budget |= ident == "budget";
208            }
209
210            syn::visit::visit_expr_path(self, i);
211        }
212    }
213
214    let expression = syn::parse2::<syn::Expr>(tokens)?;
215    let mut visitor = RequirementVisitor::default();
216
217    visitor.visit_expr(&expression);
218
219    Ok(visitor.requirements)
220}
221
222pub(crate) fn result_expression(expression: TokenStream) -> Result<TokenStream, syn::Error> {
223    let expression = syn::parse2::<syn::Expr>(expression)?;
224
225    Ok(match expression {
226        syn::Expr::Try(try_expression) => {
227            let expression = try_expression.expr;
228
229            quote::quote!(#expression)
230        }
231        expression => quote::quote!(Ok(#expression)),
232    })
233}
234
235pub(crate) fn lit_str(value: &str) -> LitStr {
236    LitStr::new(value, Span::call_site())
237}
238
239pub(crate) fn lit_int(value: impl std::fmt::Display) -> LitInt {
240    LitInt::new(&group_integer(value), Span::call_site())
241}
242
243pub(crate) fn lit_f64(value: f64) -> LitFloat {
244    let source = crate::helpers::fmt_f64(value);
245
246    LitFloat::new(&group_float(&source), Span::call_site())
247}
248
249pub(crate) fn inspect_link(path: &str, id: impl std::fmt::Display) -> String {
250    const SCHEME: &str = "https";
251    const HOST: &str = "wowlab.gg";
252
253    format!("<{SCHEME}://{HOST}/inspect/{path}/{id}>")
254}
255
256pub(crate) fn documented_u32_constant(
257    name: &str,
258    id: u32,
259    url_path: &str,
260) -> anyhow::Result<TokenStream> {
261    let name = rust_ident(name)?;
262    let documentation = lit_str(&inspect_link(url_path, id));
263    let id = lit_int(id);
264
265    Ok(quote::quote! {
266        #[doc = #documentation]
267        pub(crate) const #name: u32 = #id;
268    })
269}
270
271fn group_integer(value: impl std::fmt::Display) -> String {
272    let digits = value.to_string();
273
274    if digits.len() < 6 {
275        return digits;
276    }
277
278    let first_group = digits.len() % 3;
279    let first_group = if first_group == 0 { 3 } else { first_group };
280    let mut formatted = String::with_capacity(digits.len() + digits.len() / 3);
281
282    formatted.push_str(&digits[..first_group]);
283
284    for chunk in digits.as_bytes()[first_group..].chunks(3) {
285        formatted.push('_');
286        formatted.push_str(std::str::from_utf8(chunk).expect("decimal digits are valid UTF-8"));
287    }
288
289    formatted
290}
291
292fn group_float(value: &str) -> String {
293    let (mantissa, exponent) = value
294        .find(['e', 'E'])
295        .map_or((value, ""), |index| value.split_at(index));
296    let Some((integer, fraction)) = mantissa.split_once('.') else {
297        return value.to_owned();
298    };
299    let integer = group_integer(integer);
300
301    if fraction.len() < 5 {
302        return format!("{integer}.{fraction}{exponent}");
303    }
304
305    let mut formatted = String::with_capacity(value.len() + fraction.len() / 3);
306
307    formatted.push_str(&integer);
308    formatted.push('.');
309
310    for (index, chunk) in fraction.as_bytes().chunks(3).enumerate() {
311        if index > 0 {
312            formatted.push('_');
313        }
314
315        formatted.push_str(std::str::from_utf8(chunk).expect("decimal digits are valid UTF-8"));
316    }
317
318    formatted.push_str(exponent);
319
320    formatted
321}
322
323#[cfg(test)]
324mod test_matchers {
325    use std::fmt::Debug;
326
327    use googletest::{
328        description::Description,
329        matcher::{Matcher, MatcherBase, MatcherResult},
330    };
331
332    #[derive(MatcherBase)]
333    pub(crate) struct SourceContains {
334        expected: String,
335    }
336
337    #[derive(MatcherBase)]
338    pub(crate) struct SourceStartsWith {
339        expected: String,
340    }
341
342    pub(crate) fn contains_source(expected: &str) -> SourceContains {
343        SourceContains {
344            expected: normalize_source(expected),
345        }
346    }
347
348    pub(crate) fn starts_with_source(expected: &str) -> SourceStartsWith {
349        SourceStartsWith {
350            expected: normalize_source(expected),
351        }
352    }
353
354    impl<Actual> Matcher<Actual> for SourceContains
355    where
356        Actual: AsRef<str> + Copy + Debug,
357    {
358        fn matches(&self, actual: Actual) -> MatcherResult {
359            normalize_source(actual.as_ref())
360                .contains(&self.expected)
361                .into()
362        }
363
364        fn describe(&self, sense: MatcherResult) -> Description {
365            match sense {
366                MatcherResult::Match => {
367                    format!("contains Rust source equivalent to {:?}", self.expected).into()
368                }
369                MatcherResult::NoMatch => format!(
370                    "does not contain Rust source equivalent to {:?}",
371                    self.expected
372                )
373                .into(),
374            }
375        }
376    }
377
378    impl<Actual> Matcher<Actual> for SourceStartsWith
379    where
380        Actual: AsRef<str> + Copy + Debug,
381    {
382        fn matches(&self, actual: Actual) -> MatcherResult {
383            normalize_source(actual.as_ref())
384                .starts_with(&self.expected)
385                .into()
386        }
387
388        fn describe(&self, sense: MatcherResult) -> Description {
389            match sense {
390                MatcherResult::Match => {
391                    format!("starts with Rust source equivalent to {:?}", self.expected).into()
392                }
393                MatcherResult::NoMatch => format!(
394                    "does not start with Rust source equivalent to {:?}",
395                    self.expected
396                )
397                .into(),
398            }
399        }
400    }
401
402    fn normalize_source(source: &str) -> String {
403        let characters = source.chars().collect::<Vec<_>>();
404        let mut compact = String::with_capacity(source.len());
405
406        for (index, character) in characters.iter().copied().enumerate() {
407            if character.is_whitespace() {
408                continue;
409            }
410
411            if character == '_'
412                && index > 0
413                && characters[index - 1].is_ascii_digit()
414                && characters.get(index + 1).is_some_and(char::is_ascii_digit)
415            {
416                continue;
417            }
418
419            compact.push(character);
420        }
421
422        for delimiter in [')', ']', '}'] {
423            let trailing_comma = format!(",{delimiter}");
424
425            while compact.contains(&trailing_comma) {
426                compact = compact.replace(&trailing_comma, &delimiter.to_string());
427            }
428        }
429
430        compact
431    }
432}
433
434#[cfg(test)]
435pub(crate) use test_matchers::{contains_source, starts_with_source};
436
437#[cfg(test)]
438#[path = "rust_source/tests.rs"]
439mod tests;