fix nested modules generation

Signed-off-by: Adphi <[email protected]>
neoeinstein · Oct 14, 2024 · a07e335 · a07e335
1 parent ceb53be
commit a07e335
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 13 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/protoc-gen-prost-validate/Cargo.toml b/protoc-gen-prost-validate/Cargo.toml
@@ -13,10 +13,13 @@ prost-build.workspace = true
 prost-types.workspace = true
 prost.workspace = true
 protoc-gen-prost = { version = "0.4.0", path = "../protoc-gen-prost" }
-prost-validate-build = { version = "0.2.3" }
-prost-validate-derive-core = { version = "0.2.3" }
-prost-validate-types = { version = "0.2.3" }
+prost-validate-build = { version = "0.2.4" }
+prost-validate-derive-core = { version = "0.2.4" }
+prost-validate-types = { version = "0.2.4" }
 syn = "2.0.75"
 proc-macro2 = "1.0.86"
 prost-reflect = "0.14.2"
+quote = "1"
 prettyplease = "0.2.9"
+
+[dev-dependencies]
diff --git a/protoc-gen-prost-validate/src/generator.rs b/protoc-gen-prost-validate/src/generator.rs
@@ -6,6 +6,7 @@ use prost_types::compiler::code_generator_response::File;
 use prost_validate_build::Builder;
 use protoc_gen_prost::{Error, Generator, ModuleRequest, ModuleRequestSet, Result};
 use std::collections::HashMap;
+use syn::__private::quote::quote;
 use syn::__private::ToTokens;
 
 pub struct ProstValidateGenerator {
@@ -86,7 +87,8 @@ impl ProstValidateGenerator {
             .items
             .iter()
             .filter_map(filter_item)
-            .map(|stream| prost_validate_derive_core::derive(stream))
+            .flatten()
+            .map(|item| prost_validate_derive_core::derive_with_module(item.stream, item.module))
             .for_each(|stream| stream.to_tokens(&mut file_stream));
 
         if file_stream.is_empty() {
@@ -122,22 +124,54 @@ impl ProstValidateGenerator {
     }
 }
 
-fn filter_item(item: &syn::Item) -> Option<TokenStream> {
+struct Item {
+    module: Option<TokenStream>,
+    stream: TokenStream,
+}
+
+fn filter_item(item: &syn::Item) -> Option<Vec<Item>> {
+    do_filter_item(item, None)
+}
+
+fn do_filter_item(item: &syn::Item, module: Option<TokenStream>) -> Option<Vec<Item>> {
     match item {
         syn::Item::Struct(s) => {
             if has_validator_derive(&s.attrs) {
-                Some(item.to_token_stream())
+                Some(vec![Item {
+                    stream: item.to_token_stream(),
+                    module,
+                }])
             } else {
                 None
             }
         }
         syn::Item::Enum(e) => {
             if has_validator_derive(&e.attrs) {
-                Some(item.to_token_stream())
+                Some(vec![Item {
+                    stream: item.to_token_stream(),
+                    module,
+                }])
             } else {
                 None
             }
         }
+        syn::Item::Mod(m) => {
+            // if the parent module is not None, we need to prefix the module path to the item identity
+            let module = match module {
+                Some(m2) => {
+                    let ident = m.ident.to_token_stream();
+                    Some(quote! { #m2::#ident })
+                }
+                None => Some(m.ident.to_token_stream()),
+            };
+            let mut items = Vec::new();
+            for item in m.content.as_ref()?.1.iter() {
+                if let Some(item) = do_filter_item(item, module.clone()) {
+                    items.extend(item);
+                }
+            }
+            Some(items)
+        }
         _ => None,
     }
 }