feat: allow duplicate name by dedup_list (#215)

pilota-build/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pilota-build"
-version = "0.9.9"
+version = "0.9.10"
 edition = "2021"
 description = "Compile thrift and protobuf idl into rust code at compile-time."
 documentation = "https://docs.rs/pilota-build"

pilota-build/src/codegen/mod.rs

@@ -7,6 +7,7 @@ use std::{
 
 use dashmap::DashMap;
 use faststr::FastStr;
+use fxhash::FxHashMap;
 use itertools::Itertools;
 use normpath::PathExt;
 use pkg_tree::PkgNode;
@@ -17,6 +18,7 @@ use traits::CodegenBackend;
 use self::workspace::Workspace;
 use crate::{
     db::RirDatabase,
+    dedup::def_id_equal,
     fmt::fmt_file,
     middle::{
         self,
@@ -130,42 +132,51 @@ where
         self.backend.codegen_struct_impl(def_id, stream, s);
     }
 
-    pub fn write_item(&self, stream: &mut String, item: CodegenItem) {
+    pub fn write_item(
+        &self,
+        stream: &mut String,
+        item: CodegenItem,
+        dup: &mut FxHashMap<FastStr, Vec<DefId>>,
+    ) {
         CUR_ITEM.set(&item.def_id, || match item.kind {
             CodegenKind::Direct => {
-                let def_id = item.def_id;
-                let item = self.item(def_id).unwrap();
-                tracing::trace!("write item {}", item.symbol_name());
-                self.with_adjust(def_id, |adjust| {
-                    let attrs = adjust.iter().flat_map(|a| a.attrs()).join("\n");
-
-                    let impls = adjust
-                        .iter()
-                        .flat_map(|a| &a.nested_items)
-                        .sorted()
-                        .join("\n");
-                    stream.push_str(&impls);
-                    stream.push_str(&attrs);
-                });
-
-                match &*item {
-                    middle::rir::Item::Message(s) => self.write_struct(def_id, stream, s),
-                    middle::rir::Item::Enum(e) => self.write_enum(def_id, stream, e),
-                    middle::rir::Item::Service(s) => self.write_service(def_id, stream, s),
-                    middle::rir::Item::NewType(t) => self.write_new_type(def_id, stream, t),
-                    middle::rir::Item::Const(c) => self.write_const(def_id, stream, c),
-                    middle::rir::Item::Mod(m) => {
-                        let mut inner = Default::default();
-                        m.items
+                if !self.duplicate(dup, item.def_id) {
+                    let def_id = item.def_id;
+                    let item = self.item(def_id).unwrap();
+                    tracing::trace!("write item {}", item.symbol_name());
+                    self.with_adjust(def_id, |adjust| {
+                        let attrs = adjust.iter().flat_map(|a| a.attrs()).join("\n");
+
+                        let impls = adjust
                             .iter()
-                            .for_each(|def_id| self.write_item(&mut inner, (*def_id).into()));
-
-                        let name = self.rust_name(def_id);
-                        stream.push_str(&format! {
-                            r#"pub mod {name} {{
+                            .flat_map(|a| &a.nested_items)
+                            .sorted()
+                            .join("\n");
+                        stream.push_str(&impls);
+                        stream.push_str(&attrs);
+                    });
+
+                    match &*item {
+                        middle::rir::Item::Message(s) => {
+                            self.write_struct(def_id, stream, s);
+                        }
+                        middle::rir::Item::Enum(e) => self.write_enum(def_id, stream, e),
+                        middle::rir::Item::Service(s) => self.write_service(def_id, stream, s),
+                        middle::rir::Item::NewType(t) => self.write_new_type(def_id, stream, t),
+                        middle::rir::Item::Const(c) => self.write_const(def_id, stream, c),
+                        middle::rir::Item::Mod(m) => {
+                            let mut inner = Default::default();
+                            m.items.iter().for_each(|def_id| {
+                                self.write_item(&mut inner, (*def_id).into(), dup)
+                            });
+
+                            let name = self.rust_name(def_id);
+                            stream.push_str(&format! {
+                                r#"pub mod {name} {{
                             {inner}
                         }}"#
-                        })
+                            })
+                        }
                     }
                 };
             }
@@ -180,6 +191,21 @@ where
         })
     }
 
+    fn duplicate(&self, dup: &mut FxHashMap<FastStr, Vec<DefId>>, def_id: DefId) -> bool {
+        let name = self.rust_name(def_id);
+        if !self.dedups.contains(&name.0) {
+            return false;
+        }
+        let dup = dup.entry(name.0).or_default();
+        for id in dup.iter() {
+            if def_id_equal(&self.nodes(), *id, def_id) {
+                return true;
+            }
+        }
+        dup.push(def_id);
+        false
+    }
+
     pub fn write_enum_as_new_type(
         &self,
         def_id: DefId,
@@ -457,8 +483,9 @@ where
             let span = tracing::span!(tracing::Level::TRACE, "write_mod", path = ?p);
 
             let _enter = span.enter();
+            let mut dup = FxHashMap::default();
             for def_id in def_ids.iter() {
-                this.write_item(&mut stream, *def_id)
+                this.write_item(&mut stream, *def_id, &mut dup)
             }
         });
 

pilota-build/src/dedup.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use fxhash::FxHashMap;
 
 use crate::{
@@ -6,9 +8,9 @@ use crate::{
     DefId,
 };
 
-type Nodes = FxHashMap<DefId, Node>;
+type Nodes = Arc<FxHashMap<DefId, Node>>;
 
-fn def_id_equal(nodes: &Nodes, def_id1: DefId, def_id2: DefId) -> bool {
+pub fn def_id_equal(nodes: &Nodes, def_id1: DefId, def_id2: DefId) -> bool {
     let node1 = nodes.get(&def_id1).unwrap();
     let node2 = nodes.get(&def_id2).unwrap();
     node_equal(nodes, node1, node2)
@@ -28,9 +30,7 @@ fn node_equal(nodes: &Nodes, n1: &Node, n2: &Node) -> bool {
         (crate::rir::NodeKind::Method(m1), crate::rir::NodeKind::Method(m2)) => {
             method_equal(nodes, m1, m2)
         }
-        (crate::rir::NodeKind::Arg(a1), crate::rir::NodeKind::Arg(a2)) => {
-            arg_equal(nodes, a1, a2)
-        }
+        (crate::rir::NodeKind::Arg(a1), crate::rir::NodeKind::Arg(a2)) => arg_equal(nodes, a1, a2),
         _ => false,
     }
 }

pilota-build/src/lib.rs

@@ -17,11 +17,12 @@ pub mod parser;
 mod resolve;
 mod symbol;
 
+use faststr::FastStr;
 pub use symbol::Symbol;
 pub mod tags;
 use std::{path::PathBuf, sync::Arc};
 
-// mod dedup;
+mod dedup;
 pub mod plugin;
 
 pub use codegen::{
@@ -82,6 +83,7 @@ pub struct Builder<MkB, P> {
     touches: Vec<(std::path::PathBuf, Vec<String>)>,
     change_case: bool,
     keep_unknown_fields: Vec<std::path::PathBuf>,
+    dedups: Vec<FastStr>,
 }
 
 impl Builder<MkThriftBackend, ThriftParser> {
@@ -99,6 +101,7 @@ impl Builder<MkThriftBackend, ThriftParser> {
             ignore_unused: true,
             change_case: true,
             keep_unknown_fields: Vec::default(),
+            dedups: Vec::default(),
         }
     }
 }
@@ -118,6 +121,7 @@ impl Builder<MkProtobufBackend, ProtobufParser> {
             ignore_unused: true,
             change_case: true,
             keep_unknown_fields: Vec::default(),
+            dedups: Vec::default(),
         }
     }
 }
@@ -143,6 +147,7 @@ impl<MkB, P> Builder<MkB, P> {
             touches: self.touches,
             change_case: self.change_case,
             keep_unknown_fields: self.keep_unknown_fields,
+            dedups: self.dedups,
         }
     }
 
@@ -185,6 +190,11 @@ impl<MkB, P> Builder<MkB, P> {
         self.keep_unknown_fields.extend(item);
         self
     }
+
+    pub fn dedup(mut self, item: impl IntoIterator<Item = FastStr>) -> Self {
+        self.dedups.extend(item);
+        self
+    }
 }
 
 pub enum Output {
@@ -238,6 +248,7 @@ where
         source_type: SourceType,
         change_case: bool,
         keep_unknown_fields: Vec<PathBuf>,
+        dedups: Vec<FastStr>,
     ) -> Context {
         let mut db = RootDatabase::default();
         parser.inputs(services.iter().map(|s| &s.path));
@@ -304,7 +315,7 @@ where
 
         cx.keep(keep_unknown_fields);
 
-        cx.build(Arc::from(services), source_type, change_case)
+        cx.build(Arc::from(services), source_type, change_case, dedups)
     }
 
     pub fn compile_with_config(self, services: Vec<IdlService>, out: Output) {
@@ -319,6 +330,7 @@ where
             self.source_type,
             self.change_case,
             self.keep_unknown_fields,
+            self.dedups,
         );
 
         cx.exec_plugin(BoxedPlugin);
@@ -398,6 +410,7 @@ where
             self.source_type,
             self.change_case,
             self.keep_unknown_fields,
+            self.dedups,
         );
 
         std::thread::scope(|_scope| {

pilota-build/src/middle/context.rs

@@ -68,6 +68,7 @@ pub struct Context {
     pub location_map: FxHashMap<DefId, DefLocation>,
     pub entry_map: HashMap<DefLocation, Vec<(DefId, DefLocation)>>,
     pub plugin_gen: DashMap<DefLocation, String>,
+    pub(crate) dedups: Vec<FastStr>,
 }
 
 impl Clone for Context {
@@ -85,6 +86,7 @@ impl Clone for Context {
             location_map: self.location_map.clone(),
             entry_map: self.entry_map.clone(),
             plugin_gen: self.plugin_gen.clone(),
+            dedups: self.dedups.clone(),
         }
     }
 }
@@ -411,6 +413,7 @@ impl ContextBuilder {
         services: Arc<[crate::IdlService]>,
         source_type: SourceType,
         change_case: bool,
+        dedups: Vec<FastStr>,
     ) -> Context {
         Context {
             adjusts: Default::default(),
@@ -428,6 +431,7 @@ impl ContextBuilder {
             location_map: self.location_map,
             entry_map: self.entry_map,
             plugin_gen: Default::default(),
+            dedups,
         }
     }
 }