feat: support opt snake case style (#223)

pilota-build/Cargo.toml

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

pilota-build/src/lib.rs

@@ -85,6 +85,7 @@ pub struct Builder<MkB, P> {
     change_case: bool,
     keep_unknown_fields: Vec<std::path::PathBuf>,
     dedups: Vec<FastStr>,
+    nonstandard_snake_case: bool,
 }
 
 impl Builder<MkThriftBackend, ThriftParser> {
@@ -103,6 +104,7 @@ impl Builder<MkThriftBackend, ThriftParser> {
             change_case: true,
             keep_unknown_fields: Vec::default(),
             dedups: Vec::default(),
+            nonstandard_snake_case: false,
         }
     }
 }
@@ -123,6 +125,7 @@ impl Builder<MkProtobufBackend, ProtobufParser> {
             change_case: true,
             keep_unknown_fields: Vec::default(),
             dedups: Vec::default(),
+            nonstandard_snake_case: false,
         }
     }
 }
@@ -135,6 +138,12 @@ where
         self.parser.include_dirs(include_dirs);
         self
     }
+
+    pub fn nonstandard_snake_case(mut self, flag: bool) -> Self {
+        self.parser.nonstandard_snake_case(flag);
+        self.nonstandard_snake_case = flag;
+        self
+    }
 }
 
 impl<MkB, P> Builder<MkB, P> {
@@ -149,6 +158,7 @@ impl<MkB, P> Builder<MkB, P> {
             change_case: self.change_case,
             keep_unknown_fields: self.keep_unknown_fields,
             dedups: self.dedups,
+            nonstandard_snake_case: self.nonstandard_snake_case,
         }
     }
 
@@ -240,6 +250,7 @@ where
         self.compile_with_config(services, out)
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub fn build_cx(
         services: Vec<IdlService>,
         out: Option<Output>,
@@ -250,6 +261,7 @@ where
         change_case: bool,
         keep_unknown_fields: Vec<PathBuf>,
         dedups: Vec<FastStr>,
+        nonstandard_snake_case: bool,
     ) -> Context {
         let mut db = RootDatabase::default();
         parser.inputs(services.iter().map(|s| &s.path));
@@ -318,7 +330,13 @@ where
 
         cx.keep(keep_unknown_fields);
 
-        cx.build(Arc::from(services), source_type, change_case, dedups)
+        cx.build(
+            Arc::from(services),
+            source_type,
+            change_case,
+            dedups,
+            nonstandard_snake_case,
+        )
     }
 
     pub fn compile_with_config(self, services: Vec<IdlService>, out: Output) {
@@ -334,6 +352,7 @@ where
             self.change_case,
             self.keep_unknown_fields,
             self.dedups,
+            self.nonstandard_snake_case,
         );
 
         cx.exec_plugin(BoxedPlugin);
@@ -414,6 +433,7 @@ where
             self.change_case,
             self.keep_unknown_fields,
             self.dedups,
+            self.nonstandard_snake_case,
         );
 
         std::thread::scope(|_scope| {

pilota-build/src/middle/context.rs

@@ -69,6 +69,7 @@ pub struct Context {
     pub entry_map: HashMap<DefLocation, Vec<(DefId, DefLocation)>>,
     pub plugin_gen: DashMap<DefLocation, String>,
     pub(crate) dedups: Vec<FastStr>,
+    pub(crate) nonstandard_snake_case: bool,
 }
 
 impl Clone for Context {
@@ -87,6 +88,7 @@ impl Clone for Context {
             entry_map: self.entry_map.clone(),
             plugin_gen: self.plugin_gen.clone(),
             dedups: self.dedups.clone(),
+            nonstandard_snake_case: self.nonstandard_snake_case,
         }
     }
 }
@@ -428,6 +430,7 @@ impl ContextBuilder {
         source_type: SourceType,
         change_case: bool,
         dedups: Vec<FastStr>,
+        nonstandard_snake_case: bool,
     ) -> Context {
         Context {
             adjusts: Default::default(),
@@ -446,6 +449,7 @@ impl ContextBuilder {
             entry_map: self.entry_map,
             plugin_gen: Default::default(),
             dedups,
+            nonstandard_snake_case,
         }
     }
 }
@@ -840,8 +844,8 @@ impl Context {
                 crate::rir::Item::Enum(e) => (&**e.name).enum_ident(),
                 crate::rir::Item::Service(s) => (&**s.name).trait_ident(),
                 crate::rir::Item::NewType(t) => (&**t.name).newtype_ident(),
-                crate::rir::Item::Const(c) => (&**c.name).const_ident(),
-                crate::rir::Item::Mod(m) => (&**m.name).mod_ident(),
+                crate::rir::Item::Const(c) => (&**c.name).const_ident(self.nonstandard_snake_case),
+                crate::rir::Item::Mod(m) => (&**m.name).mod_ident(self.nonstandard_snake_case),
             },
             NodeKind::Variant(v) => {
                 let parent = self.node(def_id).unwrap().parent.unwrap();
@@ -854,14 +858,14 @@ impl Context {
                     .unwrap_or(EnumMode::Enum)
                     == EnumMode::NewType
                 {
-                    (&**v.name).shouty_snake_case()
+                    (&**v.name).shouty_snake_case(self.nonstandard_snake_case)
                 } else {
                     (&**v.name).variant_ident()
                 }
             }
-            NodeKind::Field(f) => (&**f.name).field_ident(),
-            NodeKind::Method(m) => (&**m.name).fn_ident(),
-            NodeKind::Arg(a) => (&**a.name).field_ident(),
+            NodeKind::Field(f) => (&**f.name).field_ident(self.nonstandard_snake_case),
+            NodeKind::Method(m) => (&**m.name).fn_ident(self.nonstandard_snake_case),
+            NodeKind::Arg(a) => (&**a.name).field_ident(self.nonstandard_snake_case),
         }
         .into()
     }

pilota-build/src/middle/resolver.rs

@@ -51,7 +51,11 @@ impl PathResolver for DefaultPathResolver {
                 let file = cx.file(node.file_id).unwrap();
                 let package = &file.package;
                 if package.len() != 1 || !package.first().unwrap().0.is_empty() {
-                    segs.extend(package.iter().map(|s| (&*s.0).mod_ident().into()))
+                    segs.extend(
+                        package
+                            .iter()
+                            .map(|s| (&*s.0).mod_ident(cx.nonstandard_snake_case).into()),
+                    )
                 }
             }
 

pilota-build/src/parser/mod.rs

@@ -28,5 +28,7 @@ pub trait Parser {
 
     fn include_dirs(&mut self, dirs: Vec<PathBuf>);
 
+    fn nonstandard_snake_case(&mut self, _nonstandard: bool) {}
+
     fn parse(self) -> ParseResult;
 }

pilota-build/src/parser/protobuf/mod.rs

@@ -13,18 +13,20 @@ use super::Parser;
 use crate::{
     index::Idx,
     ir::{self, FieldKind, Item, Path, TyKind},
-    symbol::{EnumRepr, FileId, Ident, IdentName},
+    symbol::{EnumRepr, FileId, Ident},
     tags::{
         protobuf::{ClientStreaming, OneOf, ProstType, Repeated, ServerStreaming},
         PilotaName, Tags,
     },
+    IdentName,
 };
 
 #[derive(Default)]
 pub struct ProtobufParser {
     inner: protobuf_parse::Parser,
     include_dirs: Vec<PathBuf>,
     input_files: FxHashSet<PathBuf>,
+    nonstandard_snake_case: bool,
 }
 
 #[derive(PartialEq, Eq)]
@@ -38,6 +40,7 @@ struct Lower {
     files: FxHashMap<String, FileId>,
     cur_package: Option<String>,
     cur_syntax: Syntax,
+    nonstandard_snake_case: bool,
 }
 
 impl Default for Lower {
@@ -47,6 +50,7 @@ impl Default for Lower {
             files: Default::default(),
             cur_package: None,
             cur_syntax: Syntax::Proto3,
+            nonstandard_snake_case: false,
         }
     }
 }
@@ -354,7 +358,7 @@ impl Lower {
         } else {
             let name = item.name();
             let mut tags = Tags::default();
-            tags.insert(PilotaName(name.0.mod_ident()));
+            tags.insert(PilotaName(name.0.mod_ident(self.nonstandard_snake_case)));
             vec![
                 item,
                 Item {
@@ -492,12 +496,19 @@ impl Parser for ProtobufParser {
         self.inner.includes(dirs);
     }
 
+    fn nonstandard_snake_case(&mut self, nonstandard: bool) {
+        self.nonstandard_snake_case = nonstandard;
+    }
+
     fn parse(self) -> super::ParseResult {
         let descriptors = self.inner.parse_and_typecheck().unwrap().file_descriptors;
 
         let mut input_file_ids = vec![];
 
-        let mut lower = Lower::default();
+        let mut lower = Lower {
+            nonstandard_snake_case: self.nonstandard_snake_case,
+            ..Default::default()
+        };
 
         let files = lower.lower(&descriptors);
 

pilota-build/src/symbol.rs

@@ -162,21 +162,21 @@ pub trait IdentName {
         self.upper_camel_ident()
     }
 
-    fn mod_ident(&self) -> FastStr {
-        self.snake_ident()
+    fn mod_ident(&self, nonstandard: bool) -> FastStr {
+        self.snake_ident(nonstandard)
     }
 
     fn variant_ident(&self) -> FastStr {
         self.upper_camel_ident()
     }
-    fn fn_ident(&self) -> FastStr {
-        self.snake_ident()
+    fn fn_ident(&self, nonstandard: bool) -> FastStr {
+        self.snake_ident(nonstandard)
     }
-    fn field_ident(&self) -> FastStr {
-        self.snake_ident()
+    fn field_ident(&self, nonstandard: bool) -> FastStr {
+        self.snake_ident(nonstandard)
     }
-    fn const_ident(&self) -> FastStr {
-        self.shouty_snake_case()
+    fn const_ident(&self, nonstandard: bool) -> FastStr {
+        self.shouty_snake_case(nonstandard)
     }
 
     fn trait_ident(&self) -> FastStr {
@@ -188,8 +188,8 @@ pub trait IdentName {
     }
 
     fn upper_camel_ident(&self) -> FastStr;
-    fn snake_ident(&self) -> FastStr;
-    fn shouty_snake_case(&self) -> FastStr;
+    fn snake_ident(&self, nonstandard: bool) -> FastStr;
+    fn shouty_snake_case(&self, nonstandard: bool) -> FastStr;
 }
 
 impl IdentName for &str {
@@ -198,12 +198,22 @@ impl IdentName for &str {
         s.into()
     }
 
-    fn snake_ident(&self) -> FastStr {
-        self.to_snake_case().into()
+    fn snake_ident(&self, nonstandard: bool) -> FastStr {
+        if nonstandard {
+            to_snake_case(self)
+        } else {
+            self.to_snake_case()
+        }
+        .into()
     }
 
-    fn shouty_snake_case(&self) -> FastStr {
-        self.to_shouty_snake_case().into()
+    fn shouty_snake_case(&self, nonstandard: bool) -> FastStr {
+        if nonstandard {
+            to_snake_case(self).to_uppercase()
+        } else {
+            self.to_shouty_snake_case()
+        }
+        .into()
     }
 }
 
@@ -212,11 +222,55 @@ impl IdentName for FastStr {
         (&**self).upper_camel_ident()
     }
 
-    fn snake_ident(&self) -> FastStr {
-        (&**self).snake_ident()
+    fn snake_ident(&self, nonstandard: bool) -> FastStr {
+        (&**self).snake_ident(nonstandard)
+    }
+
+    fn shouty_snake_case(&self, nonstandard: bool) -> FastStr {
+        (&**self).shouty_snake_case(nonstandard)
+    }
+}
+
+// Taken from rustc.
+fn to_snake_case(mut str: &str) -> String {
+    let mut words = vec![];
+    // Preserve leading underscores
+    str = str.trim_start_matches(|c: char| {
+        if c == '_' {
+            words.push(String::new());
+            true
+        } else {
+            false
+        }
+    });
+    for s in str.split('_') {
+        let mut last_upper = false;
+        let mut buf = String::new();
+        if s.is_empty() {
+            continue;
+        }
+        for ch in s.chars() {
+            if !buf.is_empty() && buf != "'" && ch.is_uppercase() && !last_upper {
+                words.push(buf);
+                buf = String::new();
+            }
+            last_upper = ch.is_uppercase();
+            buf.extend(ch.to_lowercase());
+        }
+        words.push(buf);
     }
+    words.join("_")
+}
+
+#[cfg(test)]
+mod tests {
+    use heck::ToSnakeCase;
+
+    use crate::symbol::to_snake_case;
 
-    fn shouty_snake_case(&self) -> FastStr {
-        (&**self).shouty_snake_case()
+    #[test]
+    fn snake_case() {
+        assert_eq!("IDs".to_snake_case(), "i_ds");
+        assert_eq!(to_snake_case("IDs"), "ids");
     }
 }