From 3cc97f9bca61db3f745bc4236a77109d05963b42 Mon Sep 17 00:00:00 2001
From: Tuetuopay <tuetuopay@me.com>
Date: Fri, 25 Jun 2021 08:48:01 +0000
Subject: [PATCH] build: Add mod_rs option to generate a mod.rs file

This mod.rs will `include!` all generated .rs files with proper module
hierarchy.
---
 prost-build/src/lib.rs           | 24 +++++++++-
 prost-build/src/lib_generator.rs | 79 ++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 prost-build/src/lib_generator.rs

diff --git a/prost-build/src/lib.rs b/prost-build/src/lib.rs
index f93beb9e3..a146e0488 100644
--- a/prost-build/src/lib.rs
+++ b/prost-build/src/lib.rs
@@ -113,6 +113,7 @@ mod ast;
 mod code_generator;
 mod extern_paths;
 mod ident;
+mod lib_generator;
 mod message_graph;
 mod path;
 
@@ -135,6 +136,7 @@ pub use crate::ast::{Comments, Method, Service};
 use crate::code_generator::CodeGenerator;
 use crate::extern_paths::ExternPaths;
 use crate::ident::to_snake;
+use crate::lib_generator::{LibGenerator, Mod};
 use crate::message_graph::MessageGraph;
 use crate::path::PathMap;
 
@@ -225,7 +227,8 @@ const COMPILE_WELL_KNOWN_TYPES: &str = "compile_well_known_types";
 const DISABLE_COMMENTS: &str = "disable_comments";
 const EXTERN_PATH: &str = "extern_path";
 const RETAIN_ENUM_PREFIX: &str = "retain_enum_prefix";
-pub const PROTOC_OPTS: [&str; 8] = [
+const MOD_RS: &str = "mod_rs";
+pub const PROTOC_OPTS: [&str; 9] = [
     BTREE_MAP,
     BYTES,
     FIELD_ATTR,
@@ -234,6 +237,7 @@ pub const PROTOC_OPTS: [&str; 8] = [
     DISABLE_COMMENTS,
     EXTERN_PATH,
     RETAIN_ENUM_PREFIX,
+    MOD_RS,
 ];
 
 /// Configuration options for Protobuf code generation.
@@ -252,6 +256,7 @@ pub struct Config {
     extern_paths: Vec<(String, String)>,
     protoc_args: Vec<OsString>,
     disable_comments: PathMap<()>,
+    mod_rs: bool,
 }
 
 impl Config {
@@ -288,6 +293,7 @@ impl Config {
                 [DISABLE_COMMENTS, v] => self.disable_comments.insert(v.to_string(), ()),
                 [EXTERN_PATH, k, v] => self.extern_paths.push((k.to_string(), v.to_string())),
                 [RETAIN_ENUM_PREFIX] => self.strip_enum_prefix = false,
+                [MOD_RS] => self.mod_rs = true,
                 _ if log_unknown => eprintln!("prost: Unknown option `{}`", opt.join("=")),
                 _ => (),
             }
@@ -735,6 +741,12 @@ impl Config {
         self
     }
 
+    /// Generate a `mod.rs` file suitable for a full module.
+    pub fn mod_rs(&mut self) -> &mut Self {
+        self.mod_rs = true;
+        self
+    }
+
     /// Compile `.proto` files into Rust files during a Cargo build with additional code generator
     /// configuration options.
     ///
@@ -900,6 +912,15 @@ impl Config {
             }
         }
 
+        if self.mod_rs {
+            let mut mods = Mod::default();
+            for (module, _) in &modules {
+                mods.push(module);
+            }
+            let mut buf = modules.entry(vec!["mod".to_owned()]).or_default();
+            LibGenerator::generate_librs(self, &mods, &mut buf);
+        }
+
         Ok(modules)
     }
 
@@ -927,6 +948,7 @@ impl default::Default for Config {
             extern_paths: Vec::new(),
             protoc_args: Vec::new(),
             disable_comments: PathMap::default(),
+            mod_rs: false,
         }
     }
 }
diff --git a/prost-build/src/lib_generator.rs b/prost-build/src/lib_generator.rs
new file mode 100644
index 000000000..8fb7658ed
--- /dev/null
+++ b/prost-build/src/lib_generator.rs
@@ -0,0 +1,79 @@
+use std::collections::BTreeMap;
+
+use crate::{Config, Module};
+
+#[derive(Debug)]
+pub struct Mod {
+    submodules: BTreeMap<String, Mod>,
+    contents: Vec<Module>,
+}
+
+impl Mod {
+    pub fn push(&mut self, module: &Module) {
+        match module.as_slice() {
+            [] => (),
+            [name, left @ ..] => self.add(module.to_owned(), name, left),
+        }
+    }
+
+    fn add(&mut self, module: Module, name: &str, left: &[String]) {
+        let sub = self.submodules.entry(name.to_owned()).or_default();
+        match left {
+            [] => sub.contents.push(module),
+            [name, left @ ..] => sub.add(module, name, left),
+        }
+    }
+}
+
+impl Default for Mod {
+    fn default() -> Self {
+        Self {
+            submodules: BTreeMap::new(),
+            contents: Vec::new(),
+        }
+    }
+}
+
+pub struct LibGenerator<'a> {
+    config: &'a mut Config,
+    depth: u8,
+    buf: &'a mut String,
+}
+
+impl<'a> LibGenerator<'a> {
+    pub fn generate_librs(config: &'a mut Config, mods: &Mod, buf: &'a mut String) {
+        let mut generator = LibGenerator {
+            config,
+            depth: 0,
+            buf,
+        };
+        generator.push_mod(mods);
+    }
+
+    fn push_mod(&mut self, mods: &Mod) {
+        for (name, mods) in &mods.submodules {
+            self.push_indent();
+            self.buf.push_str("pub mod ");
+            self.buf.push_str(name);
+            self.buf.push_str(" {\n");
+            self.depth += 1;
+            self.push_mod(&mods);
+            self.depth -= 1;
+            self.push_indent();
+            self.buf.push_str("}\n");
+        }
+
+        for package in mods.contents.iter().map(|content| content.join(".")) {
+            self.push_indent();
+            self.buf.push_str("include!(\"");
+            self.buf.push_str(&package);
+            self.buf.push_str(".rs\");\n");
+        }
+    }
+
+    fn push_indent(&mut self) {
+        for _ in 0..self.depth {
+            self.buf.push_str("    ");
+        }
+    }
+}