From 088d8b12755897c56d25726a5c55c0592b3de978 Mon Sep 17 00:00:00 2001
From: jokemanfire <hu.dingyang@zte.com.cn>
Date: Fri, 27 Sep 2024 22:09:52 +0800
Subject: [PATCH] Add gen mod for more convenient to use

while use customize ,you can use gen_mod to generate mod.rs.

Signed-off-by: jokemanfire <hu.dingyang@zte.com.cn>
---
 .gitignore              |  1 -
 compiler/src/codegen.rs | 31 ++++++++++++++++++++++++++++++-
 compiler/src/lib.rs     |  2 ++
 example/build.rs        |  7 +++++--
 4 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index 14dabaed..ff70e924 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,5 +6,4 @@ Cargo.lock
 .idea
 *.o
 example/protocols/**/*.rs
-!example/protocols/**/mod.rs
 src/ttrpc.rs
diff --git a/compiler/src/codegen.rs b/compiler/src/codegen.rs
index 382f5ae7..5c87ed83 100644
--- a/compiler/src/codegen.rs
+++ b/compiler/src/codegen.rs
@@ -36,7 +36,11 @@
 
 #![allow(dead_code)]
 
-use std::collections::HashMap;
+use std::{
+    collections::{HashMap, HashSet},
+    fs,
+    io::BufRead,
+};
 
 use crate::Customize;
 use protobuf::{
@@ -724,6 +728,31 @@ pub fn gen_and_write(
 ) -> io::Result<()> {
     let results = gen(file_descriptors, files_to_generate, customize);
 
+    if customize.gen_mod {
+        let file_path = out_dir.join("mod.rs");
+        let mut set = HashSet::new();
+        //if mod file exists
+        if let Ok(file) = File::open(&file_path) {
+            let reader = io::BufReader::new(file);
+            reader.lines().for_each(|line| {
+                let _ = line.map(|r| set.insert(r));
+            });
+        }
+        let mut file_write = fs::OpenOptions::new()
+            .create(true)
+            .write(true)
+            .truncate(true)
+            .open(&file_path)?;
+        for r in &results {
+            let prefix_name: Vec<&str> = r.name.split('.').collect();
+            set.insert(format!("pub mod {};", prefix_name[0]));
+        }
+        for item in &set {
+            writeln!(file_write, "{}", item)?;
+        }
+        file_write.flush()?;
+    }
+
     for r in &results {
         let mut file_path = out_dir.to_owned();
         file_path.push(&r.name);
diff --git a/compiler/src/lib.rs b/compiler/src/lib.rs
index 25a84e0c..444fc612 100644
--- a/compiler/src/lib.rs
+++ b/compiler/src/lib.rs
@@ -36,4 +36,6 @@ pub struct Customize {
     pub async_client: bool,
     /// Indicates whether to generate async code for server.
     pub async_server: bool,
+    /// Gen mod rs in mod.rs
+    pub gen_mod: bool,
 }
diff --git a/example/build.rs b/example/build.rs
index a90274b4..5026b4b1 100644
--- a/example/build.rs
+++ b/example/build.rs
@@ -18,8 +18,9 @@ fn main() {
         "protocols/protos/google/protobuf/empty.proto",
         "protocols/protos/oci.proto",
     ];
-
-    let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(false);
+    ///This could be add while the new ttrpc compiler support it.
+    ///If need use now change  ttrpc_codegen's cargo.toml ttrpc-compiler = "0.6.1" to ttrpc-compiler = "../compiler"
+    let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(true);
 
     Codegen::new()
         .out_dir("protocols/sync")
@@ -27,6 +28,7 @@ fn main() {
         .include("protocols/protos")
         .rust_protobuf()
         .customize(Customize {
+            gen_mod:true,//This could be add while the new ttrpc compiler support it.
             ..Default::default()
         })
         .rust_protobuf_customize(protobuf_customized.clone())
@@ -42,6 +44,7 @@ fn main() {
         .include("protocols/protos")
         .rust_protobuf()
         .customize(Customize {
+            gen_mod:true,  //This could be add while the new ttrpc compiler support it.
             async_all: true,
             ..Default::default()
         })