From 8423b5e923769e40bb86e129ebdbf85efb51ff26 Mon Sep 17 00:00:00 2001
From: qthree <qthree3@gmail.com>
Date: Thu, 7 Aug 2025 00:06:05 +0700
Subject: [PATCH] Add option to prettify generated Rust code

---
 engine/src/lib.rs               |  4 ++--
 engine/src/output_generators.rs | 24 +++++++++++++++++++++---
 parser/src/config.rs            |  1 +
 parser/src/directives.rs        |  7 +++++++
 src/lib.rs                      |  6 ++++++
 5 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/engine/src/lib.rs b/engine/src/lib.rs
index 26a126de8..7d923e961 100644
--- a/engine/src/lib.rs
+++ b/engine/src/lib.rs
@@ -419,8 +419,8 @@ impl IncludeCppEngine {
             config: &self.config,
             rs: match &self.state {
                 State::NotGenerated => panic!("Generate first"),
-                State::Generated(gen_results) => gen_results.item_mod.to_token_stream(),
-                State::ParseOnly => TokenStream2::new(),
+                State::Generated(gen_results) => Some(&gen_results.item_mod),
+                State::ParseOnly => None,
             },
         }
     }
diff --git a/engine/src/output_generators.rs b/engine/src/output_generators.rs
index 09c143ab8..4165eb4bd 100644
--- a/engine/src/output_generators.rs
+++ b/engine/src/output_generators.rs
@@ -8,13 +8,26 @@
 
 use autocxx_parser::{IncludeCppConfig, MultiBindings};
 use proc_macro2::TokenStream;
+use quote::ToTokens;
+use syn::ItemMod;
 
 /// Opaque structure representing the Rust which needs to be generated
 /// for a given `include_cpp!` macro. You will want to pass this into
 /// either [`generate_rs_single`] or [`generate_rs_archive`].
 pub struct RsOutput<'a> {
     pub(crate) config: &'a IncludeCppConfig,
-    pub(crate) rs: TokenStream,
+    pub(crate) rs: Option<&'a ItemMod>,
+}
+
+impl RsOutput<'_> {
+    fn to_token_stream(&self) -> TokenStream {
+        self.rs.map(ToTokens::to_token_stream).unwrap_or_default()
+    }
+    fn pretty_print(&self) -> String {
+        self.rs
+            .map(crate::rust_pretty_printer::pretty_print)
+            .unwrap_or_default()
+    }
 }
 
 /// Creates an on-disk archive (actually a JSON file) of the Rust side of the bindings
@@ -24,7 +37,7 @@ pub struct RsOutput<'a> {
 pub fn generate_rs_archive<'a>(rs_outputs: impl Iterator<Item = RsOutput<'a>>) -> String {
     let mut multi_bindings = MultiBindings::default();
     for rs in rs_outputs {
-        multi_bindings.insert(rs.config, rs.rs);
+        multi_bindings.insert(rs.config, rs.to_token_stream());
     }
     serde_json::to_string(&multi_bindings).expect("Unable to encode JSON archive")
 }
@@ -39,8 +52,13 @@ pub struct RsInclude {
 /// to a file which can simply be `include!`ed by `autocxx_macro` when you give
 /// it the `AUTOCXX_RS_FILE` environment variable.
 pub fn generate_rs_single(rs_output: RsOutput) -> RsInclude {
+    let code = if rs_output.config.prettify {
+        rs_output.pretty_print()
+    } else {
+        rs_output.to_token_stream().to_string()
+    };
     RsInclude {
-        code: rs_output.rs.to_string(),
+        code,
         filename: rs_output.config.get_rs_filename(),
     }
 }
diff --git a/parser/src/config.rs b/parser/src/config.rs
index 925039c16..7d88ab6ef 100644
--- a/parser/src/config.rs
+++ b/parser/src/config.rs
@@ -215,6 +215,7 @@ pub struct IncludeCppConfig {
     pub unsafe_policy: UnsafePolicy,
     pub parse_only: bool,
     pub exclude_impls: bool,
+    pub prettify: bool,
     pub(crate) pod_requests: Vec<String>,
     pub allowlist: Allowlist,
     pub(crate) blocklist: Vec<String>,
diff --git a/parser/src/directives.rs b/parser/src/directives.rs
index a76cb0656..5211a44e4 100644
--- a/parser/src/directives.rs
+++ b/parser/src/directives.rs
@@ -100,6 +100,13 @@ pub(crate) fn get_directives() -> &'static DirectivesMap {
                 |config| &config.exclude_utilities,
             )),
         );
+        need_exclamation.insert(
+            "pretty".into(),
+            Box::new(BoolFlag(
+                |config| &mut config.prettify,
+                |config| &config.prettify,
+            )),
+        );
         need_exclamation.insert("name".into(), Box::new(ModName));
         need_exclamation.insert("concrete".into(), Box::new(Concrete));
         need_exclamation.insert("rust_type".into(), Box::new(RustType { output: false }));
diff --git a/src/lib.rs b/src/lib.rs
index 6f42fe52f..bffe57f23 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -183,6 +183,12 @@ macro_rules! exclude_utilities {
     ($($tt:tt)*) => { $crate::usage!{$($tt)*} };
 }
 
+/// Prettify generated code.
+#[macro_export]
+macro_rules! pretty {
+    ($($tt:tt)*) => { $crate::usage!{$($tt)*} };
+}
+
 /// Entirely block some type from appearing in the generated
 /// code. This can be useful if there is a type which is not
 /// understood by bindgen or autocxx, and incorrect code is