Skip to main content

clifford_codegen/codegen/
format.rs

1//! Code formatting utilities.
2//!
3//! This module provides utilities for formatting generated Rust code
4//! using rustfmt.
5
6use proc_macro2::TokenStream;
7use std::io::Write;
8use std::process::{Command, Stdio};
9
10/// The Rust edition used for formatting generated code.
11const RUST_EDITION: &str = "2024";
12
13/// Formats a token stream using rustfmt.
14///
15/// Returns the formatted code as a string. If rustfmt is not available
16/// or fails, returns the unformatted token stream as a string.
17///
18/// # Example
19///
20/// ```
21/// use quote::quote;
22/// use clifford_codegen::codegen::format_tokens;
23///
24/// let tokens = quote! {
25///     pub struct Foo { x: i32 }
26/// };
27///
28/// let formatted = format_tokens(&tokens);
29/// // formatted will be properly indented and spaced
30/// ```
31pub fn format_tokens(tokens: &TokenStream) -> String {
32    let unformatted = tokens.to_string();
33    format_code(&unformatted).unwrap_or(unformatted)
34}
35
36/// Formats Rust code using rustfmt.
37///
38/// Returns `None` if rustfmt is not available or fails.
39fn format_code(code: &str) -> Option<String> {
40    let mut child = Command::new("rustfmt")
41        .arg(format!("--edition={}", RUST_EDITION))
42        .stdin(Stdio::piped())
43        .stdout(Stdio::piped())
44        .stderr(Stdio::null())
45        .spawn()
46        .ok()?;
47
48    // Write code to stdin
49    if let Some(mut stdin) = child.stdin.take() {
50        stdin.write_all(code.as_bytes()).ok()?;
51    }
52
53    // Read formatted output
54    let output = child.wait_with_output().ok()?;
55
56    if output.status.success() {
57        String::from_utf8(output.stdout).ok()
58    } else {
59        None
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use quote::quote;
67
68    #[test]
69    fn formats_simple_struct() {
70        let tokens = quote! {
71            pub struct Foo{x:i32,y:i32}
72        };
73
74        let formatted = format_tokens(&tokens);
75
76        // Check that it compiles (basic structure is preserved)
77        assert!(formatted.contains("struct Foo"));
78        assert!(formatted.contains("x"));
79        assert!(formatted.contains("y"));
80    }
81
82    #[test]
83    fn handles_complex_code() {
84        let tokens = quote! {
85            impl<T: Float> Vector<T> {
86                pub fn new(x: T, y: T) -> Self {
87                    Self { x, y }
88                }
89
90                pub fn x(&self) -> T {
91                    self.x
92                }
93            }
94        };
95
96        let formatted = format_tokens(&tokens);
97        assert!(formatted.contains("impl<T: Float> Vector<T>"));
98        assert!(formatted.contains("pub fn new"));
99    }
100}