feat: generate faststr for thrift string (#73)

cloudwego · Dec 7, 2022 · d2e3759 · d2e3759
1 parent f14ab07
commit d2e3759
Show file tree

Hide file tree

Showing 24 changed files with 268 additions and 259 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pilota-build/Cargo.toml b/pilota-build/Cargo.toml
@@ -20,7 +20,6 @@ maintenance = { status = "actively-developed" }
 pilota-thrift-parser = { path = "../pilota-thrift-parser", version = "0.4" }
 
 heck = "0.4"
-smol_str = "0.1"
 syn = "1"
 normpath = "0.3"
 fxhash = "0.2"
@@ -39,6 +38,7 @@ petgraph = "0.6"
 # We will switch to the official one when https://github.com/stepancheg/rust-protobuf/pull/646 is fixed.
 protobuf-parse = { package = "protobuf-parse2", version = "4.0.0-alpha.2" }
 protobuf = { package = "protobuf2", version = "4.0.0-alpha.2" }
+faststr = "0.1"
 
 [build-dependencies]
 itertools = "0.10"

diff --git a/pilota-build/src/codegen/mod.rs b/pilota-build/src/codegen/mod.rs
@@ -1,5 +1,6 @@
 use std::{collections::HashMap, ops::Deref, sync::Arc};
 
+use faststr::FastStr;
 use fxhash::FxHashMap;
 use heck::ToShoutySnakeCase;
 use pkg_tree::PkgNode;
@@ -30,7 +31,7 @@ pub struct Codegen<B> {
     backend: B,
     zero_copy: bool,
     cx: Arc<Context>,
-    pkgs: FxHashMap<Arc<[smol_str::SmolStr]>, TokenStream>,
+    pkgs: FxHashMap<Arc<[FastStr]>, TokenStream>,
 }
 
 impl<B> Deref for Codegen<B> {
@@ -339,9 +340,9 @@ where
                 let s = &**s;
                 quote! { #s.to_string() }
             }
-            (Literal::String(s), CodegenTy::SmolStr) => {
+            (Literal::String(s), CodegenTy::FastStr) => {
                 let s = &**s;
-                quote! { ::pilota::SmolStr::new(#s) }
+                quote! { ::pilota::FastStr::new(#s) }
             }
             (Literal::Int(i), CodegenTy::I16) => {
                 let i = *i as i16;
@@ -392,7 +393,7 @@ where
         }
     }
 
-    pub(crate) fn write_mods(&mut self, mods: HashMap<Arc<[smol_str::SmolStr]>, Vec<DefId>>) {
+    pub(crate) fn write_mods(&mut self, mods: HashMap<Arc<[FastStr]>, Vec<DefId>>) {
         mods.iter().for_each(|(p, def_ids)| {
             let stream: &mut TokenStream =
                 unsafe { std::mem::transmute(self.pkgs.entry(p.clone()).or_default()) };
@@ -405,7 +406,7 @@ where
 
     pub fn link(mut self, ns_name: &str) -> TokenStream {
         fn write_stream(
-            pkgs: &mut FxHashMap<Arc<[smol_str::SmolStr]>, TokenStream>,
+            pkgs: &mut FxHashMap<Arc<[FastStr]>, TokenStream>,
             stream: &mut TokenStream,
             nodes: &[PkgNode],
         ) {

diff --git a/pilota-build/src/codegen/pkg_tree.rs b/pilota-build/src/codegen/pkg_tree.rs
@@ -1,14 +1,15 @@
 use std::sync::Arc;
 
+use faststr::FastStr;
 use itertools::Itertools;
 
 #[derive(Debug)]
 pub struct PkgNode {
-    pub path: Arc<[smol_str::SmolStr]>,
+    pub path: Arc<[FastStr]>,
     pub children: Arc<[PkgNode]>,
 }
 
-fn from_pkgs(base_path: &[smol_str::SmolStr], pkgs: &[Arc<[smol_str::SmolStr]>]) -> Arc<[PkgNode]> {
+fn from_pkgs(base_path: &[FastStr], pkgs: &[Arc<[FastStr]>]) -> Arc<[PkgNode]> {
     let groups = pkgs.iter().into_group_map_by(|p| p.first().unwrap());
 
     Arc::from_iter(groups.into_iter().map(|(k, v)| {
@@ -33,11 +34,11 @@ fn from_pkgs(base_path: &[smol_str::SmolStr], pkgs: &[Arc<[smol_str::SmolStr]>])
 }
 
 impl PkgNode {
-    pub fn from_pkgs(pkgs: &[Arc<[smol_str::SmolStr]>]) -> Arc<[PkgNode]> {
+    pub fn from_pkgs(pkgs: &[Arc<[FastStr]>]) -> Arc<[PkgNode]> {
         from_pkgs(&[], pkgs)
     }
 
-    pub fn ident(&self) -> smol_str::SmolStr {
+    pub fn ident(&self) -> FastStr {
         self.path.last().unwrap().clone()
     }
 }
diff --git a/pilota-build/src/codegen/protobuf/mod.rs b/pilota-build/src/codegen/protobuf/mod.rs
@@ -51,7 +51,7 @@ impl ProstPlugin {
             };
         }
         match &ty.kind {
-            ty::String | ty::SmolStr => quote!(string),
+            ty::String | ty::FastStr => quote!(string),
             ty::Bool => quote!(bool),
             ty::BytesVec | ty::Bytes => quote!(bytes),
             ty::I32 => quote!(int32),

diff --git a/pilota-build/src/codegen/thrift/decode_helper.rs b/pilota-build/src/codegen/thrift/decode_helper.rs
@@ -41,7 +41,7 @@ impl DecodeHelper {
     protocol_method!(read_bytes_vec);
     protocol_method!(read_byte);
     protocol_method!(read_string);
-    protocol_method!(read_smolstr);
+    protocol_method!(read_faststr);
     protocol_method!(read_list_begin);
     protocol_method!(read_list_end);
     protocol_method!(read_set_begin);

diff --git a/pilota-build/src/codegen/thrift/ty.rs b/pilota-build/src/codegen/thrift/ty.rs
@@ -10,7 +10,7 @@ use crate::{
 impl ThriftBackend {
     pub(crate) fn ttype(&self, ty: &Ty) -> TokenStream {
         match &ty.kind {
-            ty::String | ty::SmolStr => quote! {::pilota::thrift::TType::Binary},
+            ty::String | ty::FastStr => quote! {::pilota::thrift::TType::Binary},
             ty::Void => quote! {::pilota::thrift::TType::Void},
             ty::U8 => quote! {::pilota::thrift::TType::I8},
             ty::Bool => quote! {::pilota::thrift::TType::Bool},
@@ -42,8 +42,8 @@ impl ThriftBackend {
 
     pub(crate) fn codegen_encode_ty(&self, ty: &Ty, ident: &Ident) -> TokenStream {
         match &ty.kind {
-            ty::String => quote! { protocol.write_string(#ident)?; },
-            ty::SmolStr => quote! { protocol.write_smolstr(#ident.clone())?; },
+            ty::String => quote! { protocol.write_string(&#ident)?; },
+            ty::FastStr => quote! { protocol.write_faststr(#ident.clone())?; },
             ty::Void => quote! {
                 protocol.write_struct_begin(&*::pilota::thrift::VOID_IDENT)?;
                 protocol.write_struct_end()?;
@@ -118,7 +118,7 @@ impl ThriftBackend {
     pub(crate) fn codegen_ty_size(&self, ty: &Ty, ident: &Ident) -> TokenStream {
         match &ty.kind {
             ty::String => quote! { protocol.write_string_len(&#ident) },
-            ty::SmolStr => quote! { protocol.write_smolstr_len(#ident) },
+            ty::FastStr => quote! { protocol.write_faststr_len(#ident) },
             ty::Void => {
                 quote! { protocol.write_struct_begin_len(&*::pilota::thrift::VOID_IDENT) +  protocol.write_struct_end_len() }
             }
@@ -206,7 +206,7 @@ impl ThriftBackend {
     pub(crate) fn codegen_decode_ty(&self, helper: &DecodeHelper, ty: &Ty) -> TokenStream {
         match &ty.kind {
             ty::String => helper.codegen_read_string(),
-            ty::SmolStr => helper.codegen_read_smolstr(),
+            ty::FastStr => helper.codegen_read_faststr(),
             ty::Void => {
                 let read_struct_begin = helper.codegen_read_struct_begin();
                 let read_struct_end = helper.codegen_read_struct_end();