1use convert_case::{Case, Casing};
7use proc_macro2::TokenStream;
8use quote::quote;
9
10use crate::codegen::Runtime;
11use crate::utils::{extract_qualified_type_name, extract_simple_type_name};
12
13#[derive(Debug, Clone, Copy)]
15pub enum RenderContext {
16 Constructor,
18 Extractor,
21 QueryExtractor,
24 Parameter,
26 ReturnType,
28 FieldType,
30 BuilderMethod,
32 PythonParameter,
34 NapiParameter,
36}
37
38pub trait TypeRenderer {
44 fn render(&self, ty: &UnifiedType) -> String;
45}
46
47pub struct RustRenderer(pub RenderContext);
49
50impl TypeRenderer for RustRenderer {
51 fn render(&self, ty: &UnifiedType) -> String {
52 unified_to_rust(ty, self.0)
53 }
54}
55
56pub struct PythonRenderer;
58
59impl TypeRenderer for PythonRenderer {
60 fn render(&self, ty: &UnifiedType) -> String {
61 unified_to_python_type(ty)
62 }
63}
64
65pub struct NapiRenderer;
67
68impl TypeRenderer for NapiRenderer {
69 fn render(&self, ty: &UnifiedType) -> String {
70 unified_to_napi(ty)
71 }
72}
73
74pub struct TypeScriptRenderer;
76
77impl TypeRenderer for TypeScriptRenderer {
78 fn render(&self, ty: &UnifiedType) -> String {
79 unified_to_typescript(ty)
80 }
81}
82
83pub fn unified_to_rust(unified_type: &UnifiedType, context: RenderContext) -> String {
92 let base_type_str = match &unified_type.base_type {
93 BaseType::String => {
94 if matches!(context, RenderContext::Constructor) && !unified_type.is_optional {
95 "impl Into<String>".to_string()
96 } else {
97 "String".to_string()
98 }
99 }
100 BaseType::Int32 => "i32".to_string(),
101 BaseType::Int64 => "i64".to_string(),
102 BaseType::Bool => "bool".to_string(),
103 BaseType::Float64 => "f64".to_string(),
104 BaseType::Float32 => "f32".to_string(),
105 BaseType::Bytes => "Vec<u8>".to_string(),
106 BaseType::Unit => "()".to_string(),
107 BaseType::Message(name) => extract_simple_type_name(name),
108 BaseType::Enum(name) => {
109 if matches!(
110 context,
111 RenderContext::Extractor | RenderContext::NapiParameter
112 ) {
113 "i32".to_string()
114 } else {
115 convert_protobuf_enum_to_rust_type(&format!("TYPE_ENUM:{name}"))
116 }
117 }
118 BaseType::OneOf(name) => extract_simple_type_name(name),
119 BaseType::Map(key_type, value_type) => {
120 if matches!(context, RenderContext::Constructor) && !unified_type.is_optional {
121 "impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>".to_string()
122 } else {
123 let key_str = unified_to_rust(key_type, context);
124 let value_str = unified_to_rust(value_type, context);
125 format!("HashMap<{key_str}, {value_str}>")
126 }
127 }
128 };
129
130 let mut result = base_type_str;
131
132 if unified_type.is_repeated && !matches!(context, RenderContext::BuilderMethod) {
134 result = format!("Vec<{result}>");
135 }
136
137 if should_wrap_in_option(context, unified_type) {
138 result = format!("Option<{result}>");
139 }
140
141 result
142}
143
144fn should_wrap_in_option(ctx: RenderContext, ty: &UnifiedType) -> bool {
156 let is_optional_non_builder = ty.is_optional && !matches!(ctx, RenderContext::BuilderMethod);
157 let is_ffi = matches!(
158 ctx,
159 RenderContext::PythonParameter | RenderContext::NapiParameter
160 );
161 let is_ffi_collection =
166 is_ffi && (matches!(ty.base_type, BaseType::Map(_, _)) || ty.is_repeated);
167 is_optional_non_builder || is_ffi_collection
168}
169
170pub fn field_assignment(
180 unified_type: &UnifiedType,
181 field_ident: &proc_macro2::Ident,
182 ctx: &RenderContext,
183 runtime: Runtime,
184) -> TokenStream {
185 if matches!(ctx, RenderContext::BuilderMethod) {
186 return flexible_optional_field_assignment(unified_type, field_ident, runtime);
187 }
188 match &unified_type.base_type {
189 BaseType::String if !unified_type.is_optional => quote! { #field_ident.into() },
190 BaseType::Enum(_) => match runtime {
191 Runtime::Prost if unified_type.is_repeated => {
192 quote! { #field_ident.into_iter().map(|v| v as i32).collect() }
193 }
194 Runtime::Prost => quote! { #field_ident as i32 },
195 Runtime::Buffa if unified_type.is_repeated => {
196 quote! { #field_ident.into_iter().map(buffa::EnumValue::Known).collect() }
197 }
198 Runtime::Buffa => quote! { buffa::EnumValue::Known(#field_ident) },
199 },
200 BaseType::Map(_, _) => quote! {
201 #field_ident.into_iter().map(|(k, v)| (k.into(), v.into())).collect()
202 },
203 BaseType::OneOf(_) if !unified_type.is_repeated => {
208 quote! { Some(#field_ident) }
209 }
210 BaseType::Message(_) if !unified_type.is_repeated => match runtime {
211 Runtime::Prost => quote! { Some(#field_ident) },
213 Runtime::Buffa => quote! { buffa::MessageField::some(#field_ident) },
215 },
216 _ => quote! { #field_ident },
217 }
218}
219
220pub fn unified_to_python_type(unified_type: &UnifiedType) -> String {
222 let base_type_str = match &unified_type.base_type {
223 BaseType::String => "str".to_string(),
224 BaseType::Int32 | BaseType::Int64 => "int".to_string(),
225 BaseType::Bool => "bool".to_string(),
226 BaseType::Float64 | BaseType::Float32 => "float".to_string(),
227 BaseType::Bytes => "bytes".to_string(),
228 BaseType::Unit => "None".to_string(),
229 BaseType::Message(name) => extract_qualified_type_name(name),
233 BaseType::Enum(name) => extract_qualified_type_name(name),
234 BaseType::OneOf(name) => extract_qualified_type_name(name),
235 BaseType::Map(key_type, value_type) => {
236 let key_str = unified_to_python_type(key_type);
237 let value_str = unified_to_python_type(value_type);
238 format!("Dict[{key_str}, {value_str}]")
239 }
240 };
241
242 let mut result = base_type_str;
243
244 if unified_type.is_repeated {
245 result = format!("List[{result}]");
246 }
247
248 if unified_type.is_optional {
249 result = format!("Optional[{result}]");
250 }
251
252 result
253}
254
255pub fn unified_to_napi(unified_type: &UnifiedType) -> String {
260 let base_type_str = match &unified_type.base_type {
261 BaseType::String => "String".to_string(),
262 BaseType::Int32 => "i32".to_string(),
263 BaseType::Int64 => "i64".to_string(),
264 BaseType::Bool => "bool".to_string(),
265 BaseType::Float64 => "f64".to_string(),
266 BaseType::Float32 => "f32".to_string(),
267 BaseType::Bytes => "Vec<u8>".to_string(),
268 BaseType::Unit => "()".to_string(),
269 BaseType::Message(name) => extract_simple_type_name(name),
270 BaseType::Enum(name) => convert_protobuf_enum_to_rust_type(&format!("TYPE_ENUM:{name}")),
271 BaseType::OneOf(name) => extract_simple_type_name(name),
272 BaseType::Map(key_type, value_type) => {
273 let key_str = unified_to_napi(key_type);
274 let value_str = unified_to_napi(value_type);
275 format!("HashMap<{key_str}, {value_str}>")
276 }
277 };
278
279 let mut result = base_type_str;
280
281 if unified_type.is_repeated {
282 result = format!("Vec<{result}>");
283 }
284
285 if unified_type.is_optional
286 || matches!(unified_type.base_type, BaseType::Map(_, _))
287 || unified_type.is_repeated
288 {
289 result = format!("Option<{result}>");
290 }
291
292 result
293}
294
295pub fn unified_to_typescript(unified_type: &UnifiedType) -> String {
299 let base_type_str = match &unified_type.base_type {
300 BaseType::String => "string".to_string(),
301 BaseType::Int32 | BaseType::Int64 => "number".to_string(),
302 BaseType::Bool => "boolean".to_string(),
303 BaseType::Float64 | BaseType::Float32 => "number".to_string(),
304 BaseType::Bytes => "Uint8Array".to_string(),
305 BaseType::Unit => "void".to_string(),
306 BaseType::Message(name) => extract_simple_type_name(name),
307 BaseType::Enum(_) => "number".to_string(),
308 BaseType::OneOf(name) => extract_simple_type_name(name),
309 BaseType::Map(key_type, value_type) => {
310 let key_str = unified_to_typescript(key_type);
311 let value_str = unified_to_typescript(value_type);
312 format!("Record<{key_str}, {value_str}>")
313 }
314 };
315
316 let mut result = base_type_str;
317
318 if unified_type.is_repeated {
319 result = format!("{result}[]");
320 }
321
322 if unified_type.is_optional {
323 result = format!("{result} | undefined");
324 }
325
326 result
327}
328
329fn flexible_optional_field_assignment(
330 unified_type: &UnifiedType,
331 field_ident: &proc_macro2::Ident,
332 runtime: Runtime,
333) -> TokenStream {
334 if unified_type.is_optional {
335 match (&unified_type.base_type, runtime) {
336 (BaseType::Enum(_), Runtime::Prost) => quote! { #field_ident.into().map(|e| e as i32) },
337 (BaseType::Enum(_), Runtime::Buffa) => {
338 quote! { #field_ident.into().map(buffa::EnumValue::Known) }
339 }
340 (BaseType::Message(_) | BaseType::OneOf(_), Runtime::Buffa) => quote! {
345 {
346 let #field_ident: ::core::option::Option<_> = #field_ident.into();
347 buffa::MessageField::from(#field_ident)
348 }
349 },
350 _ => quote! { #field_ident.into() },
351 }
352 } else {
353 match (&unified_type.base_type, runtime) {
354 (BaseType::String, _) => quote! { #field_ident.into() },
355 (
356 BaseType::Int32
357 | BaseType::Int64
358 | BaseType::Bool
359 | BaseType::Float64
360 | BaseType::Float32,
361 _,
362 ) => {
363 quote! { #field_ident.into() }
364 }
365 (BaseType::Enum(_), Runtime::Prost) => quote! { #field_ident as i32 },
366 (BaseType::Enum(_), Runtime::Buffa) => quote! { buffa::EnumValue::Known(#field_ident) },
367 _ => quote! { #field_ident },
370 }
371 }
372}
373
374pub(crate) fn convert_protobuf_enum_to_rust_type(proto_type: &str) -> String {
389 if let Some(enum_name) = proto_type.strip_prefix("TYPE_ENUM:") {
390 let enum_name = enum_name.trim_start_matches('.');
392
393 if let Some(last_dot) = enum_name.rfind('.') {
396 let parent_part = &enum_name[..last_dot];
397 let enum_simple_name = &enum_name[last_dot + 1..];
398
399 let parent_parts: Vec<&str> = parent_part.split('.').collect();
402 if let Some(last_part) = parent_parts.last() {
403 if last_part.chars().next().is_some_and(|c| c.is_uppercase())
406 && *last_part != "V1"
407 && !last_part
408 .chars()
409 .all(|c| c.is_lowercase() || c.is_numeric())
410 {
411 let snake_case_module = last_part.to_case(Case::Snake);
413
414 let enum_rust_name = if enum_simple_name.contains('_')
416 && enum_simple_name
417 .chars()
418 .all(|c| c.is_uppercase() || c == '_')
419 {
420 enum_simple_name.to_case(Case::Pascal)
421 } else {
422 enum_simple_name.to_string()
423 };
424
425 format!("{snake_case_module}::{enum_rust_name}")
426 } else {
427 convert_enum_name_to_rust(enum_simple_name)
429 }
430 } else {
431 convert_enum_name_to_rust(enum_simple_name)
433 }
434 } else {
435 convert_enum_name_to_rust(enum_name)
437 }
438 } else {
439 "i32".to_string() }
441}
442
443fn convert_enum_name_to_rust(enum_name: &str) -> String {
445 if enum_name.contains('_') && enum_name.chars().all(|c| c.is_uppercase() || c == '_') {
447 enum_name.to_case(Case::Pascal)
448 } else {
449 enum_name.to_string()
450 }
451}
452
453#[derive(Debug, Clone)]
461pub struct UnifiedType {
462 pub base_type: BaseType,
464 pub is_optional: bool,
466 pub is_repeated: bool,
468}
469
470#[derive(Debug, Clone)]
476pub enum BaseType {
477 String,
479 Int32,
481 Int64,
483 Bool,
485 Float64,
487 Float32,
489 Bytes,
491 Message(String),
493 Enum(String),
495 OneOf(String),
497 Map(Box<UnifiedType>, Box<UnifiedType>),
499 Unit,
501}
502
503impl UnifiedType {
504 pub fn type_ident(&self) -> syn::Ident {
506 let name = match &self.base_type {
507 BaseType::Message(n) | BaseType::Enum(n) | BaseType::OneOf(n) => {
508 n.split('.').next_back().unwrap_or(n)
509 }
510 BaseType::String => "String",
511 BaseType::Int32 => "i32",
512 BaseType::Int64 => "i64",
513 BaseType::Bool => "bool",
514 BaseType::Float64 => "f64",
515 BaseType::Float32 => "f32",
516 BaseType::Bytes => "Bytes",
517 BaseType::Unit => "()",
518 BaseType::Map(_, _) => "HashMap",
519 };
520 quote::format_ident!("{}", name)
521 }
522
523 pub fn string() -> Self {
525 Self {
526 base_type: BaseType::String,
527 is_optional: false,
528 is_repeated: false,
529 }
530 }
531
532 pub fn optional(mut self) -> Self {
534 self.is_optional = true;
535 self
536 }
537
538 pub fn repeated(mut self) -> Self {
540 self.is_repeated = true;
541 self
542 }
543
544 pub fn map(key: UnifiedType, value: UnifiedType) -> Self {
546 Self {
547 base_type: BaseType::Map(Box::new(key), Box::new(value)),
548 is_optional: false,
549 is_repeated: false,
550 }
551 }
552}