Generate mod.rs file for each mode when splitting

cloudwego · Oct 13, 2024 · 1f74c61 · 1f74c61
1 parent 63df82e
commit 1f74c61
Show file tree

Hide file tree

Showing 24 changed files with 132 additions and 88 deletions.
diff --git a/pilota-build/src/codegen/mod.rs b/pilota-build/src/codegen/mod.rs
@@ -6,6 +6,7 @@ use std::{
 };
 
 use ahash::AHashMap;
+use dashmap::mapref::one::RefMut;
 use dashmap::DashMap;
 use faststr::FastStr;
 use itertools::Itertools;
@@ -478,44 +479,11 @@ where
 
             let _enter = span.enter();
             let mut dup = AHashMap::default();
-            for def_id in def_ids.iter() {
-                if this.split {
-                    let mut item_stream = String::new();
-                    let node = this.db.node(def_id.def_id).unwrap();
-                    let name_prefix = match node.kind {
-                        NodeKind::Item(ref item) => match item.as_ref() {
-                            Item::Message(_) => "message",
-                            Item::Enum(_) => "enum",
-                            Item::Service(_) => "service",
-                            Item::NewType(_) => "new_type",
-                            Item::Const(_) => "const",
-                            Item::Mod(_) => "mod",
-                        },
-                        NodeKind::Variant(_) => "variant",
-                        NodeKind::Field(_) => "field",
-                        NodeKind::Method(_) => "method",
-                        NodeKind::Arg(_) => "arg",
-                    };
-
-                    let base_mod_name = p.iter().map(|s| s.to_string()).join("/");
-                    let mod_dir = base_dir.join(base_mod_name.clone());
-
-                    let file_name = format!("{}_{}.rs", name_prefix, node.name());
-                    this.write_item(&mut item_stream, *def_id, &mut dup);
 
-                    let full_path = mod_dir.join(file_name.clone());
-                    std::fs::create_dir_all(mod_dir).unwrap();
-
-                    let mut file =
-                        std::io::BufWriter::new(std::fs::File::create(full_path.clone()).unwrap());
-                    file.write_all(item_stream.as_bytes()).unwrap();
-                    file.flush().unwrap();
-                    fmt_file(full_path);
-
-                    stream.push_str(
-                        format!("include!(\"{}/{}\");\n", base_mod_name, file_name).as_str(),
-                    );
-                } else {
+            if this.split {
+                Self::write_split_mod(this, base_dir, p, def_ids, &mut stream, &mut dup);
+            } else {
+                for def_id in def_ids.iter() {
                     this.write_item(&mut stream, *def_id, &mut dup)
                 }
             }
@@ -557,6 +525,62 @@ where
         write_stream(&mut pkgs, stream, &pkg_node);
     }
 
+    fn write_split_mod(
+        this: &mut Codegen<B>,
+        base_dir: &Path,
+        p: &Arc<[FastStr]>,
+        def_ids: &Vec<CodegenItem>,
+        stream: &mut RefMut<Arc<[FastStr]>, String>,
+        mut dup: &mut AHashMap<FastStr, Vec<DefId>>,
+    ) {
+        let base_mod_name = p.iter().map(|s| s.to_string()).join("/");
+        let mod_file_name = format!("{}/mod.rs", base_mod_name);
+        let mut mod_stream = String::new();
+
+        for def_id in def_ids.iter() {
+            let mut item_stream = String::new();
+            let node = this.db.node(def_id.def_id).unwrap();
+            let name_prefix = match node.kind {
+                NodeKind::Item(ref item) => match item.as_ref() {
+                    Item::Message(_) => "message",
+                    Item::Enum(_) => "enum",
+                    Item::Service(_) => "service",
+                    Item::NewType(_) => "new_type",
+                    Item::Const(_) => "const",
+                    Item::Mod(_) => "mod",
+                },
+                NodeKind::Variant(_) => "variant",
+                NodeKind::Field(_) => "field",
+                NodeKind::Method(_) => "method",
+                NodeKind::Arg(_) => "arg",
+            };
+
+            let mod_dir = base_dir.join(base_mod_name.clone());
+
+            let file_name = format!("{}_{}.rs", name_prefix, node.name());
+            this.write_item(&mut item_stream, *def_id, &mut dup);
+
+            let full_path = mod_dir.join(file_name.clone());
+            std::fs::create_dir_all(mod_dir).unwrap();
+
+            let mut file =
+                std::io::BufWriter::new(std::fs::File::create(full_path.clone()).unwrap());
+            file.write_all(item_stream.as_bytes()).unwrap();
+            file.flush().unwrap();
+            fmt_file(full_path);
+
+            mod_stream.push_str(format!("include!(\"{}\");\n", file_name).as_str());
+        }
+
+        let mod_path = base_dir.join(&mod_file_name);
+        let mut mod_file = std::io::BufWriter::new(std::fs::File::create(&mod_path).unwrap());
+        mod_file.write_all(mod_stream.as_bytes()).unwrap();
+        mod_file.flush().unwrap();
+        fmt_file(&mod_path);
+
+        stream.push_str(format!("include!(\"{}\");\n", mod_file_name).as_str());
+    }
+
     pub fn write_file(self, ns_name: Symbol, file_name: impl AsRef<Path>) {
         let base_dir = file_name.as_ref().parent().unwrap();
         let mut stream = String::default();

diff --git a/pilota-build/src/test/mod.rs b/pilota-build/src/test/mod.rs
@@ -1,8 +1,8 @@
 #![cfg(test)]
 
 use std::fs::File;
-use std::{fs, path::Path};
 use std::process::Command;
+use std::{fs, path::Path};
 use tempfile::tempdir;
 
 use crate::{plugin::SerdePlugin, IdlService};
@@ -83,19 +83,22 @@ fn check_cargo_build(target: impl AsRef<Path>) {
     // `cargo build` produces the `target` directory and the `Cargo.lock` file
     // To not pollute the test data, copy the directory to a temporary one
     copy_dir_recursively(&target, &tmp_target).unwrap();
-    
+
     println!("Running cargo build in {}", tmp_target.display());
 
     let result = Command::new(std::env::var("CARGO").unwrap_or_else(|_| "cargo".to_owned()))
         .current_dir(&tmp_target)
         .arg("build")
         .output();
-    
+
     match result {
         Ok(status) => {
             if !status.status.success() {
                 eprintln!("{}", String::from_utf8_lossy(&status.stderr));
-                panic!("cargo build returned non-zero exit code for {}. See above for more details", tmp_target.display());
+                panic!(
+                    "cargo build returned non-zero exit code for {}. See above for more details",
+                    tmp_target.display()
+                );
             }
         }
         Err(_) => {
@@ -152,7 +155,7 @@ fn test_with_builder_workspace<F: FnOnce(&Path, &Path)>(
         File::create(cargo_toml_path).unwrap();
 
         f(source.as_ref(), target.as_ref());
-        
+
         check_cargo_build(target)
     } else {
         let dir = tempdir().unwrap();
@@ -174,7 +177,7 @@ fn test_with_builder_workspace<F: FnOnce(&Path, &Path)>(
 
         f(source.as_ref(), &path);
         diff_dir(&target, &base_dir_tmp);
-        
+
         check_cargo_build(target)
     }
 }

diff --git a/pilota-build/test_data/thrift_with_split/wrapper_arc.rs b/pilota-build/test_data/thrift_with_split/wrapper_arc.rs
@@ -2,12 +2,6 @@ pub mod wrapper_arc {
     #![allow(warnings, clippy::all)]
 
     pub mod wrapper_arc {
-        include!("wrapper_arc/message_A.rs");
-        include!("wrapper_arc/service_TestService.rs");
-        include!("wrapper_arc/enum_TestServiceTestResultRecv.rs");
-        include!("wrapper_arc/message_TestServiceTestArgsRecv.rs");
-        include!("wrapper_arc/enum_TestServiceTestResultSend.rs");
-        include!("wrapper_arc/message_TEST.rs");
-        include!("wrapper_arc/message_TestServiceTestArgsSend.rs");
+        include!("wrapper_arc/mod.rs");
     }
 }
diff --git a/pilota-build/test_data/thrift_with_split/wrapper_arc/mod.rs b/pilota-build/test_data/thrift_with_split/wrapper_arc/mod.rs
@@ -0,0 +1,7 @@
+include!("message_A.rs");
+include!("service_TestService.rs");
+include!("enum_TestServiceTestResultRecv.rs");
+include!("message_TestServiceTestArgsRecv.rs");
+include!("enum_TestServiceTestResultSend.rs");
+include!("message_TEST.rs");
+include!("message_TestServiceTestArgsSend.rs");
diff --git a/...a-build/test_data/thrift_workspace_with_split/output/article/src/article/image/cdn/mod.rs b/...a-build/test_data/thrift_workspace_with_split/output/article/src/article/image/cdn/mod.rs
@@ -0,0 +1 @@
+include!("message_CDN.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/article/src/article/image/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/article/src/article/image/mod.rs
@@ -0,0 +1 @@
+include!("message_Image.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/article/src/article/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/article/src/article/mod.rs
@@ -0,0 +1,9 @@
+include!("enum_Status.rs");
+include!("message_ArticleServiceGetArticleArgsRecv.rs");
+include!("enum_ArticleServiceGetArticleResultSend.rs");
+include!("message_GetArticleResponse.rs");
+include!("message_ArticleServiceGetArticleArgsSend.rs");
+include!("message_GetArticleRequest.rs");
+include!("service_ArticleService.rs");
+include!("enum_ArticleServiceGetArticleResultRecv.rs");
+include!("message_Article.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/article/src/author/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/article/src/author/mod.rs
@@ -0,0 +1 @@
+include!("message_Author.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/article/src/common/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/article/src/common/mod.rs
@@ -0,0 +1 @@
+include!("message_CommonData.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/article/src/gen.rs b/pilota-build/test_data/thrift_workspace_with_split/output/article/src/gen.rs
@@ -2,31 +2,23 @@ pub mod gen {
     #![allow(warnings, clippy::all)]
 
     pub mod article {
-        include!("article/enum_Status.rs");
-        include!("article/message_ArticleServiceGetArticleArgsRecv.rs");
-        include!("article/enum_ArticleServiceGetArticleResultSend.rs");
-        include!("article/message_GetArticleResponse.rs");
-        include!("article/message_ArticleServiceGetArticleArgsSend.rs");
-        include!("article/message_GetArticleRequest.rs");
-        include!("article/service_ArticleService.rs");
-        include!("article/enum_ArticleServiceGetArticleResultRecv.rs");
-        include!("article/message_Article.rs");
+        include!("article/mod.rs");
 
         pub mod image {
-            include!("article/image/message_Image.rs");
+            include!("article/image/mod.rs");
 
             pub mod cdn {
-                include!("article/image/cdn/message_CDN.rs");
+                include!("article/image/cdn/mod.rs");
             }
         }
     }
 
     pub mod author {
-        include!("author/message_Author.rs");
+        include!("author/mod.rs");
     }
 
     pub mod common {
-        include!("common/message_CommonData.rs");
+        include!("common/mod.rs");
     }
     pub use article::*;
 }
diff --git a/...ta-build/test_data/thrift_workspace_with_split/output/author/src/article/image/cdn/mod.rs b/...ta-build/test_data/thrift_workspace_with_split/output/author/src/article/image/cdn/mod.rs
@@ -0,0 +1 @@
+include!("message_CDN.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/author/src/article/image/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/author/src/article/image/mod.rs
@@ -0,0 +1 @@
+include!("message_Image.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/author/src/author/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/author/src/author/mod.rs
@@ -0,0 +1,8 @@
+include!("service_AuthorService.rs");
+include!("enum_AuthorServiceGetAuthorResultRecv.rs");
+include!("message_AuthorServiceGetAuthorArgsRecv.rs");
+include!("enum_AuthorServiceGetAuthorResultSend.rs");
+include!("message_GetAuthorResponse.rs");
+include!("message_AuthorServiceGetAuthorArgsSend.rs");
+include!("message_GetAuthorRequest.rs");
+include!("message_Author.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/author/src/common/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/author/src/common/mod.rs
@@ -0,0 +1 @@
+include!("message_CommonData.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/author/src/gen.rs b/pilota-build/test_data/thrift_workspace_with_split/output/author/src/gen.rs
@@ -4,27 +4,20 @@ pub mod gen {
     pub mod article {
 
         pub mod image {
-            include!("article/image/message_Image.rs");
+            include!("article/image/mod.rs");
 
             pub mod cdn {
-                include!("article/image/cdn/message_CDN.rs");
+                include!("article/image/cdn/mod.rs");
             }
         }
     }
 
     pub mod author {
-        include!("author/service_AuthorService.rs");
-        include!("author/enum_AuthorServiceGetAuthorResultRecv.rs");
-        include!("author/message_AuthorServiceGetAuthorArgsRecv.rs");
-        include!("author/enum_AuthorServiceGetAuthorResultSend.rs");
-        include!("author/message_GetAuthorResponse.rs");
-        include!("author/message_AuthorServiceGetAuthorArgsSend.rs");
-        include!("author/message_GetAuthorRequest.rs");
-        include!("author/message_Author.rs");
+        include!("author/mod.rs");
     }
 
     pub mod common {
-        include!("common/message_CommonData.rs");
+        include!("common/mod.rs");
     }
     pub use author::*;
 }
diff --git a/...ta-build/test_data/thrift_workspace_with_split/output/common/src/article/image/cdn/mod.rs b/...ta-build/test_data/thrift_workspace_with_split/output/common/src/article/image/cdn/mod.rs
@@ -0,0 +1 @@
+include!("message_CDN.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/common/src/article/image/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/common/src/article/image/mod.rs
@@ -0,0 +1 @@
+include!("message_Image.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/common/src/author/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/common/src/author/mod.rs
@@ -0,0 +1 @@
+include!("message_Author.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/common/src/common/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/common/src/common/mod.rs
@@ -0,0 +1 @@
+include!("message_CommonData.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/common/src/gen.rs b/pilota-build/test_data/thrift_workspace_with_split/output/common/src/gen.rs
@@ -4,19 +4,19 @@ pub mod gen {
     pub mod article {
 
         pub mod image {
-            include!("article/image/message_Image.rs");
+            include!("article/image/mod.rs");
 
             pub mod cdn {
-                include!("article/image/cdn/message_CDN.rs");
+                include!("article/image/cdn/mod.rs");
             }
         }
     }
 
     pub mod author {
-        include!("author/message_Author.rs");
+        include!("author/mod.rs");
     }
 
     pub mod common {
-        include!("common/message_CommonData.rs");
+        include!("common/mod.rs");
     }
 }
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/image/src/article/image/cdn/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/image/src/article/image/cdn/mod.rs
@@ -0,0 +1 @@
+include!("message_CDN.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/image/src/article/image/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/image/src/article/image/mod.rs
@@ -0,0 +1,8 @@
+include!("message_ImageServiceGetImageArgsRecv.rs");
+include!("enum_ImageServiceGetImageResultSend.rs");
+include!("message_GetImageResponse.rs");
+include!("message_ImageServiceGetImageArgsSend.rs");
+include!("message_GetImageRequest.rs");
+include!("service_ImageService.rs");
+include!("enum_ImageServiceGetImageResultRecv.rs");
+include!("message_Image.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/image/src/common/mod.rs b/pilota-build/test_data/thrift_workspace_with_split/output/image/src/common/mod.rs
@@ -0,0 +1 @@
+include!("message_CommonData.rs");
diff --git a/pilota-build/test_data/thrift_workspace_with_split/output/image/src/gen.rs b/pilota-build/test_data/thrift_workspace_with_split/output/image/src/gen.rs
@@ -4,23 +4,16 @@ pub mod gen {
     pub mod article {
 
         pub mod image {
-            include!("article/image/message_ImageServiceGetImageArgsRecv.rs");
-            include!("article/image/enum_ImageServiceGetImageResultSend.rs");
-            include!("article/image/message_GetImageResponse.rs");
-            include!("article/image/message_ImageServiceGetImageArgsSend.rs");
-            include!("article/image/message_GetImageRequest.rs");
-            include!("article/image/service_ImageService.rs");
-            include!("article/image/enum_ImageServiceGetImageResultRecv.rs");
-            include!("article/image/message_Image.rs");
+            include!("article/image/mod.rs");
 
             pub mod cdn {
-                include!("article/image/cdn/message_CDN.rs");
+                include!("article/image/cdn/mod.rs");
             }
         }
     }
 
     pub mod common {
-        include!("common/message_CommonData.rs");
+        include!("common/mod.rs");
     }
     pub use article::image::*;
 }