From 324346839e9e0b16cef9da5dcf99d92235dc0cc4 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                      | 10 ++++++---
 example/protocols/asynchronous/mod.rs | 15 -------------
 example/protocols/sync/mod.rs         | 13 -----------
 6 files changed, 39 insertions(+), 33 deletions(-)
 delete mode 100644 example/protocols/asynchronous/mod.rs
 delete mode 100644 example/protocols/sync/mod.rs

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 853fa757..a9cbaedb 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::{
@@ -723,6 +727,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..84454fa6 100644
--- a/example/build.rs
+++ b/example/build.rs
@@ -4,12 +4,15 @@
 //
 
 use std::{
-    fs::File,
+    fs::{self, File},
     io::{Read, Write},
 };
 use ttrpc_codegen::{Codegen, Customize, ProtobufCustomize};
 
 fn main() {
+    fs::create_dir_all("protocols/sync").unwrap();
+    fs::create_dir_all("protocols/asynchronous").unwrap();
+
     let mut protos = vec![
         "protocols/protos/github.com/gogo/protobuf/gogoproto/gogo.proto",
         "protocols/protos/github.com/kata-containers/agent/pkg/types/types.proto",
@@ -18,8 +21,7 @@ fn main() {
         "protocols/protos/google/protobuf/empty.proto",
         "protocols/protos/oci.proto",
     ];
-
-    let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(false);
+    let protobuf_customized = ProtobufCustomize::default().gen_mod_rs(true);
 
     Codegen::new()
         .out_dir("protocols/sync")
@@ -27,6 +29,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 +45,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()
         })
diff --git a/example/protocols/asynchronous/mod.rs b/example/protocols/asynchronous/mod.rs
deleted file mode 100644
index 964c24d4..00000000
--- a/example/protocols/asynchronous/mod.rs
+++ /dev/null
@@ -1,15 +0,0 @@
-// Copyright (c) 2020 Ant Financial
-//
-// SPDX-License-Identifier: Apache-2.0
-//
-
-pub mod agent;
-pub mod agent_ttrpc;
-pub mod empty;
-mod gogo;
-pub mod health;
-pub mod health_ttrpc;
-mod oci;
-pub mod streaming;
-pub mod streaming_ttrpc;
-pub mod types;
diff --git a/example/protocols/sync/mod.rs b/example/protocols/sync/mod.rs
deleted file mode 100644
index d598eb2d..00000000
--- a/example/protocols/sync/mod.rs
+++ /dev/null
@@ -1,13 +0,0 @@
-// Copyright (c) 2020 Ant Financial
-//
-// SPDX-License-Identifier: Apache-2.0
-//
-
-pub mod agent;
-pub mod agent_ttrpc;
-pub mod empty;
-mod gogo;
-pub mod health;
-pub mod health_ttrpc;
-mod oci;
-pub mod types;