Generate modules to make easy importing

src/generator.rs

@@ -3,6 +3,7 @@ use crate::proto_gen::google::protobuf::descriptor::{
 };
 use crate::{ProtoAccessor, Result, RotoError};
 use std::str;
+use std::collections::{HashMap, HashSet};
 
 pub fn to_pascal_case(s: &str) -> String {
     s.split('_')
@@ -75,14 +76,37 @@ fn map_type_to_rust_builder(field_type: i32) -> (String, String) {
     }
 }
 
-pub fn generate_rust_code(set: &FileDescriptorSet) -> String {
-    let mut output = String::new();
-    output.push_str("use crate::{ProtoAccessor, ProtoBuilder, Result, RotoError, read_varint, RepeatedFieldIterator};\n");
-    output.push_str("use std::str;\n\n");
+pub fn generate_rust_code(
+    set: &FileDescriptorSet,
+    files_to_generate: Option<&[String]>,
+    generate_mod_files: bool,
+) -> Vec<(String, String)> {
+    let mut generated_files = Vec::new();
 
     for file_res in set.file() {
         let (file_data, _) = file_res.expect("Failed to iterate file");
         let file_proto = FileDescriptorProto::new(file_data).expect("Failed to parse FileDescriptorProto");
+        let proto_name = file_proto.name().expect("File proto name missing");
+
+        if let Some(filter) = files_to_generate {
+            if !filter.contains(&proto_name.to_string()) {
+                continue;
+            }
+        }
+
+        let rust_file_name = format!("{}.rs", proto_name.replace(".proto", ""));
+
+        let mut output = String::new();
+        output.push_str("use crate::{ProtoAccessor, ProtoBuilder, Result, RotoError, read_varint, RepeatedFieldIterator};\n");
+        output.push_str("use std::str;\n\n");
+
+        for dep_res in file_proto.dependency() {
+            let (dep_data, _) = dep_res.expect("Failed to iterate dependency");
+            let dep_name = str::from_utf8(dep_data).expect("Dependency name invalid utf8");
+            let dep_mod_path = dep_name.replace(".proto", "").replace('/', "::");
+            output.push_str(&format!("use crate::{};\n", dep_mod_path));
+        }
+        output.push_str("\n");
 
         // Enums
         for enum_res in file_proto.enum_type() {
@@ -293,7 +317,54 @@ pub fn generate_rust_code(set: &FileDescriptorSet) -> String {
                 "    pub fn finish(self) -> Result<&'b mut [u8]> {{\n        self.builder.finish()\n    }}\n}}\n\n"
             ));
         }
+        generated_files.push((rust_file_name, output));
     }
 
-    output
+    if !generate_mod_files {
+        return generated_files;
+    }
+
+    let mut all_paths: Vec<String> = generated_files.iter().map(|(p, _)| p.clone()).collect();
+    all_paths.sort();
+
+    let mut mod_files: HashMap<String, HashSet<String>> = HashMap::new();
+    for path in &all_paths {
+        let parts: Vec<&str> = path.split('/').collect();
+        let mut current_dir = String::new();
+        for i in 0..parts.len() - 1 {
+            if !current_dir.is_empty() {
+                current_dir.push('/');
+            }
+            current_dir.push_str(parts[i]);
+            let mod_path = format!("{}/mod.rs", current_dir);
+            let sub_mod = parts[i + 1].replace(".rs", "");
+            mod_files.entry(mod_path).or_default().insert(sub_mod);
+        }
+    }
+
+    let mut root_mods = HashSet::new();
+    for path in &all_paths {
+        let parts: Vec<&str> = path.split('/').collect();
+        root_mods.insert(parts[0].replace(".rs", ""));
+    }
+
+    let mut root_mod_content = String::new();
+    let mut sorted_root_mods: Vec<_> = root_mods.into_iter().collect();
+    sorted_root_mods.sort();
+    for m in sorted_root_mods {
+        root_mod_content.push_str(&format!("pub mod {};\n", m));
+    }
+    generated_files.push(("mod.rs".to_string(), root_mod_content));
+
+    for (mod_path, sub_mods) in mod_files {
+        let mut content = String::new();
+        let mut sorted_subs: Vec<_> = sub_mods.into_iter().collect();
+        sorted_subs.sort();
+        for sub in sorted_subs {
+            content.push_str(&format!("pub mod {};\n", sub));
+        }
+        generated_files.push((mod_path, content));
+    }
+
+    generated_files
 }