feat: serde plugin support attributes

pilota-build/src/lib.rs

@@ -26,14 +26,14 @@ mod test;
 
 use codegen::protobuf::ProtobufBackend;
 pub use codegen::{thrift::ThriftBackend, traits::CodegenBackend, Codegen};
-use db::RootDatabase;
+use db::{RirDatabase, RootDatabase};
 use fmt::fmt_file;
+pub use middle::{context::Context, rir, ty};
 use middle::{
     context::{tls::CONTEXT, CollectMode},
     rir::NodeKind,
     type_graph::TypeGraph,
 };
-pub use middle::{rir, ty};
 use parser::{protobuf::ProtobufParser, thrift::ThriftParser, ParseResult, Parser};
 use plugin::{
     AutoDerivePlugin, BoxedPlugin, EnumNumPlugin, ImplDefaultPlugin, PredicateResult,
@@ -44,9 +44,7 @@ use resolve::{ResolveResult, Resolver};
 use salsa::{Durability, ParallelDatabase};
 pub use symbol::{DefId, IdentName};
 use syn::parse_quote;
-
-use crate::db::RirDatabase;
-pub use crate::middle::context::Context;
+pub use tags::TagId;
 
 pub trait MakeBackend: Sized {
     type Target: CodegenBackend;

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

@@ -452,7 +452,7 @@ impl ThriftLower {
         }
 
         annotations.iter().for_each(
-            |annotation| with_tags!(annotation -> crate::tags::PilotaName | crate::tags::RustType | crate::tags::RustWrapperArc),
+            |annotation| with_tags!(annotation -> crate::tags::PilotaName | crate::tags::RustType | crate::tags::RustWrapperArc | crate::tags::SerdeAttribute),
         );
 
         tags

pilota-build/src/plugin/mod.rs

@@ -7,7 +7,7 @@ use syn::{parse_quote, Attribute};
 
 use crate::{
     db::RirDatabase,
-    rir::{Field, Item},
+    rir::{EnumVariant, Field, Item},
     symbol::{DefId, EnumRepr},
     ty::{self, Ty, Visitor},
     Context, IdentName,
@@ -23,7 +23,11 @@ pub trait Plugin {
     }
 
     fn on_field(&mut self, cx: &mut Context, def_id: DefId, f: Arc<Field>) {
-        walk_filed(self, cx, def_id, f)
+        walk_field(self, cx, def_id, f)
+    }
+
+    fn on_variant(&mut self, cx: &mut Context, def_id: DefId, variant: Arc<EnumVariant>) {
+        walk_variant(self, cx, def_id, variant)
     }
 
     fn on_emit(&mut self, _cx: &mut Context) {}
@@ -56,6 +60,10 @@ impl Plugin for BoxClonePlugin {
         self.0.on_field(cx, def_id, f)
     }
 
+    fn on_variant(&mut self, cx: &mut Context, def_id: DefId, variant: Arc<EnumVariant>) {
+        self.0.on_variant(cx, def_id, variant)
+    }
+
     fn on_emit(&mut self, cx: &mut Context) {
         self.0.on_emit(cx)
     }
@@ -82,6 +90,10 @@ where
         (*self).on_field(cx, def_id, f)
     }
 
+    fn on_variant(&mut self, cx: &mut Context, def_id: DefId, variant: Arc<EnumVariant>) {
+        (*self).on_variant(cx, def_id, variant)
+    }
+
     fn on_emit(&mut self, cx: &mut Context) {
         (*self).on_emit(cx)
     }
@@ -94,18 +106,30 @@ pub fn walk_item<P: Plugin + ?Sized>(p: &mut P, cx: &mut Context, _def_id: DefId
             .fields
             .iter()
             .for_each(|f| p.on_field(cx, f.did, f.clone())),
+        Item::Enum(e) => e
+            .variants
+            .iter()
+            .for_each(|v| p.on_variant(cx, v.did, v.clone())),
         _ => {}
     }
 }
 
-pub fn walk_filed<P: Plugin + ?Sized>(
+pub fn walk_field<P: Plugin + ?Sized>(
     _p: &mut P,
     _cx: &mut Context,
     _def_id: DefId,
     _field: Arc<Field>,
 ) {
 }
 
+pub fn walk_variant<P: Plugin + ?Sized>(
+    _p: &mut P,
+    _cx: &mut Context,
+    _def_id: DefId,
+    _variant: Arc<EnumVariant>,
+) {
+}
+
 pub struct BoxedPlugin;
 
 impl Plugin for BoxedPlugin {

pilota-build/src/plugin/serde.rs

@@ -1,5 +1,9 @@
+use std::str::FromStr;
+
 use syn::parse_quote;
 
+use crate::tags::SerdeAttribute;
+
 #[derive(Clone, Copy)]
 pub struct SerdePlugin;
 
@@ -10,14 +14,55 @@ impl crate::Plugin for SerdePlugin {
         def_id: crate::DefId,
         item: std::sync::Arc<crate::rir::Item>,
     ) {
+        let attribute = cx
+            .node_tags(def_id)
+            .and_then(|tags| tags.get::<SerdeAttribute>().cloned());
+
         match &*item {
             crate::rir::Item::Message(_)
             | crate::rir::Item::Enum(_)
             | crate::rir::Item::NewType(_) => cx.with_adjust(def_id, |adj| {
-                adj.add_attrs(&[parse_quote!(#[derive(::serde::Serialize, ::serde::Deserialize)])])
+                adj.add_attrs(&[parse_quote!(#[derive(::serde::Serialize, ::serde::Deserialize)])]);
+                if let Some(attribute) = attribute {
+                    let attr = attribute.0.to_string().replace('\\', "");
+                    let tokens = proc_macro2::TokenStream::from_str(&attr).unwrap();
+                    adj.add_attrs(&[parse_quote!(#tokens)]);
+                }
             }),
             _ => {}
         };
         crate::plugin::walk_item(self, cx, def_id, item)
     }
+
+    fn on_field(
+        &mut self,
+        cx: &mut crate::Context,
+        def_id: crate::DefId,
+        f: std::sync::Arc<crate::rir::Field>,
+    ) {
+        if let Some(attribute) = cx
+            .tags(f.tags_id)
+            .and_then(|tags| tags.get::<SerdeAttribute>().cloned())
+        {
+            let attr = attribute.0.replace('\\', "");
+            let tokens = proc_macro2::TokenStream::from_str(&attr).unwrap();
+            cx.with_adjust(def_id, |adj| adj.add_attrs(&[parse_quote!(#tokens)]))
+        }
+    }
+
+    fn on_variant(
+        &mut self,
+        cx: &mut crate::Context,
+        def_id: crate::DefId,
+        variant: std::sync::Arc<crate::rir::EnumVariant>,
+    ) {
+        if let Some(attribute) = cx
+            .node_tags(variant.did)
+            .and_then(|tags| tags.get::<SerdeAttribute>().cloned())
+        {
+            let attr = attribute.0.replace('\\', "");
+            let tokens = proc_macro2::TokenStream::from_str(&attr).unwrap();
+            cx.with_adjust(def_id, |adj| adj.add_attrs(&[parse_quote!(#tokens)]))
+        }
+    }
 }

pilota-build/src/tags.rs

@@ -167,6 +167,21 @@ impl Annotation for RustWrapperArc {
     const KEY: &'static str = "pilota.rust_wrapper_arc";
 }
 
+#[derive(Clone)]
+pub struct SerdeAttribute(pub FastStr);
+
+impl FromStr for SerdeAttribute {
+    type Err = std::convert::Infallible;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(Self(FastStr::new(s)))
+    }
+}
+
+impl Annotation for SerdeAttribute {
+    const KEY: &'static str = "pilota.serde_attribute";
+}
+
 pub mod protobuf {
 
     #[derive(Copy, Clone, PartialEq, Eq)]

pilota-build/src/test/mod.rs

@@ -4,6 +4,8 @@ use std::path::Path;
 
 use tempfile::tempdir;
 
+use crate::plugin::SerdePlugin;
+
 fn diff_file(old: impl AsRef<Path>, new: impl AsRef<Path>) {
     let old_content = unsafe { String::from_utf8_unchecked(std::fs::read(old).unwrap()) };
 
@@ -55,6 +57,24 @@ fn test_thrift(source: impl AsRef<Path>, target: impl AsRef<Path>) {
     });
 }
 
+fn test_plugin_thrift(source: impl AsRef<Path>, target: impl AsRef<Path>) {
+    test_with_builder(source, target, |source, target| {
+        crate::Builder::thrift()
+            .ignore_unused(false)
+            .plugin(SerdePlugin)
+            .compile(&[source], target)
+    });
+}
+
+fn test_plugin_proto(source: impl AsRef<Path>, target: impl AsRef<Path>) {
+    test_with_builder(source, target, |source, target| {
+        crate::Builder::protobuf()
+            .ignore_unused(false)
+            .plugin(SerdePlugin)
+            .compile(&[source], target)
+    });
+}
+
 #[test]
 fn test_thrift_gen() {
     let test_data_dir = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
@@ -97,6 +117,31 @@ fn test_protobuf_gen() {
     });
 }
 
+#[test]
+fn test_plugin_gen() {
+    let test_data_dir = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("test_data")
+        .join("plugin");
+
+    test_data_dir.read_dir().unwrap().for_each(|f| {
+        let f = f.unwrap();
+
+        let path = f.path();
+
+        if let Some(ext) = path.extension() {
+            if ext == "thrift" {
+                let mut rs_path = path.clone();
+                rs_path.set_extension("rs");
+                test_plugin_thrift(path, rs_path);
+            } else if ext == "proto" {
+                let mut rs_path = path.clone();
+                rs_path.set_extension("rs");
+                test_plugin_proto(path, rs_path);
+            }
+        }
+    });
+}
+
 #[test]
 fn test_touch() {
     let file_path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))