1use std::{collections::HashMap, ops::Deref, path::PathBuf, sync::Arc};
2
3use ahash::{AHashMap, HashSet};
4use anyhow::Context as _;
5use dashmap::DashMap;
6use faststr::FastStr;
7use itertools::Itertools;
8use normpath::PathExt;
9use rustc_hash::{FxHashMap, FxHashSet};
10
11use self::tls::with_cur_item;
12use super::{
13 adjust::Adjust,
14 resolver::{DefaultPathResolver, PathResolver, WorkspacePathResolver},
15 rir::NodeKind,
16 root_selector::{RootSelection, SelectionKind},
17};
18use crate::{
19 Plugin,
20 db::{RirDatabase, RootDatabase},
21 rir::{self, Field, Item, ItemPath, Literal},
22 symbol::{DefId, FileId, IdentName, ModPath, SPECIAL_NAMINGS, Symbol},
23 tags::{TagId, Tags},
24 ty::{AdtDef, AdtKind, CodegenTy, Visitor},
25};
26
27#[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Clone)]
28pub struct CrateId {
29 pub(crate) main_file: FileId,
30}
31
32#[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Clone)]
33pub enum DefLocation {
34 Fixed(CrateId, ItemPath),
35 Dynamic,
36}
37
38#[deprecated(
39 since = "0.13.8",
40 note = "`CollectMode` has been superseded by `crate::middle::root_selector::RootSelection`. \
41 This alias is kept temporarily for backward compatibility and will be removed in a \
42 future release."
43)]
44pub enum CollectMode {
45 All,
46 OnlyUsed {
47 touches: Vec<(std::path::PathBuf, Vec<String>)>,
48 },
49}
50
51#[derive(Debug)]
52pub struct WorkspaceInfo {
53 pub dir: PathBuf,
54 pub(crate) location_map: FxHashMap<DefId, DefLocation>,
55}
56
57#[derive(Debug)]
58pub enum Mode {
59 Workspace(WorkspaceInfo),
60 SingleFile { file_path: std::path::PathBuf },
61}
62
63pub struct Context {
64 pub db: RootDatabase,
65 pub source: Source,
66 pub config: Config,
67 pub cache: Cache,
68}
69
70#[derive(Clone)]
71pub struct Source {
72 pub source_type: SourceType,
73 pub services: Arc<[crate::IdlService]>,
74 pub mode: Arc<Mode>,
75 pub path_resolver: Arc<dyn PathResolver>,
76 pub selection_kind: SelectionKind,
77}
78
79#[derive(Clone)]
80pub struct Config {
81 pub change_case: bool,
82 pub split: bool,
83 pub with_descriptor: bool,
84 pub with_field_mask: bool,
85 pub touch_all: bool,
86 pub common_crate_name: FastStr,
87 pub with_comments: bool,
88}
89
90#[derive(Clone)]
91pub struct Cache {
92 pub adjusts: Arc<DashMap<DefId, Adjust>>,
93 pub mod_idxes: AHashMap<ModPath, DefId>, pub codegen_items: Vec<DefId>,
95 pub mod_items: AHashMap<ModPath, Vec<DefId>>,
96 pub def_mod: HashMap<DefId, ModPath>,
97 pub mod_files: HashMap<ModPath, Vec<FileId>>,
98 pub keep_unknown_fields: Arc<FxHashSet<DefId>>,
99 pub location_map: Arc<FxHashMap<DefId, DefLocation>>,
100 pub entry_map: Arc<HashMap<DefLocation, Vec<(DefId, DefLocation)>>>,
101 pub plugin_gen: Arc<DashMap<DefLocation, String>>,
102 pub dedups: Vec<FastStr>,
103 pub names: FxHashMap<DefId, usize>,
104}
105
106impl Clone for Context {
107 fn clone(&self) -> Self {
108 Self {
109 db: self.db.clone(),
110 source: self.source.clone(),
111 config: self.config.clone(),
112 cache: self.cache.clone(),
113 }
114 }
115}
116
117pub(crate) struct ContextBuilder {
118 db: RootDatabase,
119 pub(crate) codegen_items: Vec<DefId>,
120 input_items: Vec<DefId>,
121 mode: Mode,
122 keep_unknown_fields: FxHashSet<DefId>,
123 pub location_map: FxHashMap<DefId, DefLocation>,
124 entry_map: HashMap<DefLocation, Vec<(DefId, DefLocation)>>,
125}
126
127impl ContextBuilder {
128 pub fn new(db: RootDatabase, mode: Mode, input_items: Vec<DefId>) -> Self {
129 ContextBuilder {
130 db,
131 mode,
132 input_items,
133 codegen_items: Default::default(),
134 keep_unknown_fields: Default::default(),
135 location_map: Default::default(),
136 entry_map: Default::default(),
137 }
138 }
139 pub(crate) fn collect(&mut self, root_selection: RootSelection) {
140 match root_selection {
141 RootSelection::All => {
142 let nodes = self.db.nodes();
143 self.codegen_items
144 .extend(nodes.iter().filter_map(|(k, v)| match &v.kind {
145 NodeKind::Item(i) => {
146 if !matches!(&**i, Item::Mod(_)) {
147 Some(k)
148 } else {
149 None
150 }
151 }
152 _ => None,
153 }));
154 }
155 RootSelection::Service(root_defs) | RootSelection::Explicit(root_defs) => {
156 self.input_items.extend(root_defs);
157
158 let def_ids = self.collect_items(&self.input_items);
159 self.codegen_items.extend(def_ids.iter());
160 }
161 }
162 if matches!(self.mode, Mode::Workspace(_)) {
163 let location_map = self.workspace_collect_def_ids(&self.codegen_items);
164 self.location_map = location_map.clone();
165 self.entry_map = location_map
166 .clone()
167 .into_iter()
168 .into_group_map_by(|item| item.1.clone());
169 if let Mode::Workspace(info) = &mut self.mode {
170 info.location_map = location_map
171 }
172 }
173 }
174
175 pub(crate) fn collect_items(&self, input: &[DefId]) -> FxHashSet<DefId> {
176 struct PathCollector<'a> {
177 set: &'a mut FxHashSet<DefId>,
178 cx: &'a ContextBuilder,
179 file_ids: &'a mut FxHashSet<FileId>,
180 }
181
182 impl super::ty::Visitor for PathCollector<'_> {
183 fn visit_path(&mut self, path: &crate::rir::Path) {
184 collect(self.cx, path.did, self.set, self.file_ids)
185 }
186 }
187
188 fn collect(
189 cx: &ContextBuilder,
190 def_id: DefId,
191 set: &mut FxHashSet<DefId>,
192 file_ids: &mut FxHashSet<FileId>,
193 ) {
194 if set.contains(&def_id) {
195 return;
196 }
197
198 let node = cx.db.node(def_id).unwrap();
199
200 if !file_ids.contains(&node.file_id) {
201 file_ids.insert(node.file_id);
202
203 let file = cx.db.file(node.file_id).unwrap();
204 if file.extensions.has_used_options() {
205 file.extensions
206 .unwrap_as_pb()
207 .used_options
208 .0
209 .iter()
210 .for_each(|option| {
211 let extendee = cx.db.pb_ext(option).unwrap();
212 PathCollector { cx, set, file_ids }
213 .visit(&extendee.extendee_ty.item_ty);
214 });
215 }
216 }
217
218 match node.kind {
219 NodeKind::Item(_) => {}
220 _ => return collect(cx, node.parent.unwrap(), set, file_ids),
221 }
222
223 if !matches!(&*cx.db.item(def_id).unwrap(), rir::Item::Mod(_)) {
224 set.insert(def_id);
225 }
226
227 let node = cx.db.node(def_id).unwrap();
228 tracing::trace!("collecting {:?}", node.expect_item().symbol_name());
229
230 node.related_nodes
231 .iter()
232 .for_each(|def_id| collect(cx, *def_id, set, file_ids));
233
234 let item = node.expect_item();
235
236 match item {
237 rir::Item::Message(m) => {
238 m.fields.iter().for_each(|f| {
240 PathCollector { cx, set, file_ids }.visit(&f.ty);
241 if let Some(Literal::Path(p)) = &f.default {
242 PathCollector { cx, set, file_ids }.visit_path(p);
243 }
244 if f.item_exts.has_used_options() {
246 f.item_exts
247 .unwrap_as_pb()
248 .used_options
249 .0
250 .iter()
251 .for_each(|index| {
252 let extendee = cx.db.pb_ext(index).unwrap_or_else(|| {
253 panic!("extension:{:?} not found", index)
254 });
255 PathCollector { cx, set, file_ids }
256 .visit(&extendee.extendee_ty.item_ty);
257 });
258 }
259 });
260 if m.item_exts.has_used_options() {
262 m.item_exts
263 .unwrap_as_pb()
264 .used_options
265 .0
266 .iter()
267 .for_each(|index| {
268 let extendee = cx.db.pb_ext(index).unwrap();
269 PathCollector { cx, set, file_ids }
270 .visit(&extendee.extendee_ty.item_ty);
271 });
272 }
273 }
274 rir::Item::Enum(e) => {
275 e.variants.iter().for_each(|v| {
276 for ty in &v.fields {
277 PathCollector { cx, set, file_ids }.visit(ty);
278 }
279 if v.item_exts.has_used_options() {
280 v.item_exts
281 .unwrap_as_pb()
282 .used_options
283 .0
284 .iter()
285 .for_each(|index| {
286 let extendee = cx.db.pb_ext(index).unwrap();
287 PathCollector { cx, set, file_ids }
288 .visit(&extendee.extendee_ty.item_ty);
289 });
290 }
291 });
292 if e.item_exts.has_used_options() {
293 e.item_exts
294 .unwrap_as_pb()
295 .used_options
296 .0
297 .iter()
298 .for_each(|index| {
299 let extendee = cx.db.pb_ext(index).unwrap();
300 PathCollector { cx, set, file_ids }
301 .visit(&extendee.extendee_ty.item_ty);
302 });
303 }
304 }
305 rir::Item::Service(s) => {
306 s.extend
307 .iter()
308 .for_each(|p| collect(cx, p.did, set, file_ids));
309 s.methods.iter().for_each(|m| {
310 m.args
312 .iter()
313 .for_each(|f| PathCollector { cx, set, file_ids }.visit(&f.ty));
314 PathCollector { cx, set, file_ids }.visit(&m.ret);
316 if let Some(exceptions) = &m.exceptions {
318 PathCollector { cx, set, file_ids }.visit_path(exceptions);
319 }
320 if m.item_exts.has_used_options() {
322 m.item_exts
323 .unwrap_as_pb()
324 .used_options
325 .0
326 .iter()
327 .for_each(|index| {
328 let extendee = cx.db.pb_ext(index).unwrap();
329 PathCollector { cx, set, file_ids }
330 .visit(&extendee.extendee_ty.item_ty);
331 });
332 }
333 });
334
335 if s.item_exts.has_used_options() {
337 s.item_exts
338 .unwrap_as_pb()
339 .used_options
340 .0
341 .iter()
342 .for_each(|index| {
343 let extendee = cx.db.pb_ext(index).unwrap();
344 PathCollector { cx, set, file_ids }
345 .visit(&extendee.extendee_ty.item_ty);
346 });
347 }
348 }
349 rir::Item::NewType(n) => PathCollector { cx, set, file_ids }.visit(&n.ty),
350 rir::Item::Const(c) => {
351 PathCollector { cx, set, file_ids }.visit(&c.ty);
352 }
353 rir::Item::Mod(m) => {
354 m.items.iter().for_each(|i| collect(cx, *i, set, file_ids));
355 }
356 }
357 }
358 let mut set = FxHashSet::default();
359
360 let mut file_ids = FxHashSet::default();
361
362 input.iter().for_each(|def_id| {
363 collect(self, *def_id, &mut set, &mut file_ids);
364 });
365
366 self.db.nodes().iter().for_each(|(def_id, node)| {
367 if let NodeKind::Item(item) = &node.kind {
368 if let rir::Item::Const(_) = &**item {
369 collect(self, *def_id, &mut set, &mut file_ids);
370 }
371 }
372 });
373
374 set
375 }
376
377 pub(crate) fn workspace_collect_def_ids(
378 &self,
379 input: &[DefId],
380 ) -> FxHashMap<DefId, DefLocation> {
381 self.db.collect_def_ids(input, None)
382 }
383
384 pub(crate) fn keep(&mut self, keep_unknown_fields: Vec<PathBuf>) {
385 let mut file_ids = FxHashSet::default();
386 keep_unknown_fields.into_iter().for_each(|p| {
387 let path = p.normalize().unwrap().into_path_buf();
388 let file_id = {
389 let file_ids_map = self.db.file_ids_map();
390 *file_ids_map.get(&path).unwrap()
391 };
392 keep_files(self, &file_id, &mut file_ids);
393
394 fn keep_files(
395 cx: &mut ContextBuilder,
396 file_id: &FileId,
397 file_ids: &mut FxHashSet<FileId>,
398 ) {
399 if !file_ids.insert(*file_id) {
400 return;
401 }
402 let (uses, items_to_keep) = {
403 let files = cx.db.files();
404 let file = files.get(file_id).unwrap();
405 let uses = file.uses.clone();
406 let items_to_keep = file
407 .items
408 .iter()
409 .filter(|&&def_id| match cx.db.node(def_id) {
410 Some(rir::Node {
411 kind: rir::NodeKind::Item(_),
412 tags,
413 ..
414 }) => !matches!(
415 cx.db.tags_map().get(&tags).and_then(|tags| {
416 tags.get::<crate::tags::KeepUnknownFields>()
417 }),
418 Some(crate::tags::KeepUnknownFields(false))
419 ),
420 _ => true,
421 })
422 .cloned()
423 .collect::<Vec<_>>();
424 (uses, items_to_keep)
425 };
426
427 for f in &uses {
428 keep_files(cx, f, file_ids);
429 }
430
431 cx.keep_unknown_fields.extend(items_to_keep);
432 }
433 });
434 }
435
436 #[allow(clippy::too_many_arguments)]
437 pub(crate) fn build(
438 self,
439 services: Arc<[crate::IdlService]>,
440 source_type: SourceType,
441 selection_kind: SelectionKind,
442 change_case: bool,
443 dedups: Vec<FastStr>,
444 special_namings: Vec<FastStr>,
445 common_crate_name: FastStr,
446 split: bool,
447 with_descriptor: bool,
448 with_field_mask: bool,
449 touch_all: bool,
450 with_comments: bool,
451 ) -> Context {
452 let mode = Arc::new(self.mode);
453 SPECIAL_NAMINGS.get_or_init(|| special_namings);
454 let mut cx = Context {
455 db: self.db.clone(),
456 source: Source {
457 source_type,
458 services,
459 mode: mode.clone(),
460 path_resolver: match &*mode {
461 Mode::Workspace(_) => Arc::new(WorkspacePathResolver),
462 Mode::SingleFile { .. } => Arc::new(DefaultPathResolver),
463 },
464 selection_kind,
465 },
466 config: Config {
467 change_case,
468 split,
469 with_descriptor,
470 with_field_mask,
471 touch_all,
472 common_crate_name,
473 with_comments,
474 },
475 cache: Cache {
476 adjusts: Default::default(),
477 codegen_items: self.codegen_items,
478 keep_unknown_fields: Arc::new(self.keep_unknown_fields),
479 location_map: Arc::new(self.location_map),
480 entry_map: Arc::new(self.entry_map),
481 plugin_gen: Default::default(),
482 dedups,
483 names: Default::default(),
484 mod_idxes: Default::default(),
485 mod_items: Default::default(),
486 mod_files: Default::default(),
487 def_mod: Default::default(),
488 },
489 };
490 let mut map: FxHashMap<(Vec<DefId>, String), Vec<DefId>> = FxHashMap::default();
491 let mut mod_idxes = AHashMap::default();
492 cx.nodes()
493 .iter()
494 .for_each(|(def_id, node)| match &node.kind {
495 NodeKind::Item(item) => {
496 if let crate::rir::Item::Mod(_) = &**item {
497 mod_idxes.insert(
498 ModPath::from(Arc::from_iter(
499 cx.mod_path(*def_id).iter().map(|s| s.0.clone()),
500 )),
501 *def_id,
502 );
503 return;
504 }
505 if let Mode::Workspace(_) = &*cx.source.mode {
506 if !cx.cache.location_map.contains_key(def_id) {
507 return;
508 }
509 }
510 let rust_name = cx.item_path(*def_id).join("::");
511 map.entry((vec![], rust_name)).or_default().push(*def_id);
512 }
513 _ => {
514 let mut item_def_ids = vec![];
515 let mut item_def_id = *def_id;
516 while !matches!(cx.node(item_def_id).unwrap().kind, NodeKind::Item(_)) {
517 item_def_id = cx.node(item_def_id).unwrap().parent.unwrap();
518 item_def_ids.push(item_def_id);
519 }
520 let rust_name = cx.rust_name(*def_id).to_string();
521 map.entry((item_def_ids, rust_name))
522 .or_default()
523 .push(*def_id);
524 }
525 });
526 cx.cache.names.extend(
527 map.into_iter()
528 .filter(|(_, v)| v.len() > 1)
529 .map(|(_, v)| v)
530 .flat_map(|v| v.into_iter().enumerate().map(|(i, def_id)| (def_id, i)))
531 .collect::<HashMap<DefId, usize>>(),
532 );
533 cx.cache.mod_idxes.extend(mod_idxes);
534
535 if matches!(&*cx.source.mode, Mode::SingleFile { .. }) {
536 let mut mod_files = HashMap::<ModPath, HashSet<FileId>>::default();
537 let mod_items = cx
538 .cache
539 .codegen_items
540 .clone()
541 .into_iter()
542 .into_group_map_by(|def_id| {
543 let file_id = cx.node(*def_id).unwrap().file_id;
544
545 let mod_path = cx.mod_index(*def_id);
546 let set = mod_files.entry(mod_path.clone()).or_default();
547 if !set.contains(&file_id) {
548 set.insert(file_id);
549 }
550 cx.cache.def_mod.insert(*def_id, mod_path.clone());
551 mod_path
552 });
553 cx.cache.mod_items.extend(AHashMap::from_iter(mod_items));
554 cx.cache.mod_files.extend(
555 mod_files
556 .into_iter()
557 .map(|(k, v)| (k, v.into_iter().collect())),
558 );
559 }
560 cx
561 }
562}
563
564impl Deref for Context {
565 type Target = RootDatabase;
566
567 fn deref(&self) -> &Self::Target {
568 &self.db
569 }
570}
571
572#[derive(Clone, Copy)]
573pub enum SourceType {
574 Thrift,
575 Protobuf,
576}
577
578impl Context {
579 pub fn config_data(&self) -> &Config {
580 &self.config
581 }
582
583 pub fn cache_data(&self) -> &Cache {
584 &self.cache
585 }
586
587 pub fn source_data(&self) -> &Source {
588 &self.source
589 }
590
591 pub fn with_adjust<T, F>(&self, def_id: DefId, f: F) -> T
592 where
593 F: FnOnce(Option<&Adjust>) -> T,
594 {
595 match self.cache.adjusts.get(&def_id) {
596 Some(adj) => f(Some(&*adj)),
597 None => f(None),
598 }
599 }
600
601 pub fn with_adjust_mut<T, F>(&self, def_id: DefId, f: F) -> T
602 where
603 F: FnOnce(&mut Adjust) -> T,
604 {
605 let adjust = &mut *self.cache.adjusts.entry(def_id).or_default();
606 f(adjust)
607 }
608
609 pub fn tags(&self, tags_id: TagId) -> Option<Arc<Tags>> {
610 self.db.tags_map().get(&tags_id).cloned()
611 }
612
613 pub fn node_tags(&self, def_id: DefId) -> Option<Arc<Tags>> {
614 let tags_id = self.node(def_id).unwrap().tags;
615 self.tags(tags_id)
616 }
617
618 pub fn contains_tag<T: 'static>(&self, tags_id: TagId) -> bool {
619 self.tags(tags_id)
620 .and_then(|tags| tags.contains::<T>().then_some(true))
621 .is_some()
622 }
623
624 pub fn node_contains_tag<T: 'static>(&self, def_id: DefId) -> bool {
625 self.node_tags(def_id)
626 .and_then(|tags| tags.contains::<T>().then_some(true))
627 .is_some()
628 }
629
630 pub fn symbol_name(&self, def_id: DefId) -> Symbol {
631 let item = self.item(def_id).unwrap();
632 item.symbol_name()
633 }
634
635 fn get_codegen_ty_for_path(&self, def_id: DefId) -> CodegenTy {
636 let node = self.node(def_id).unwrap();
637 match &node.kind {
638 NodeKind::Item(item) => match &**item {
639 Item::Const(c) => {
640 let ty = self.codegen_const_ty(c.ty.kind.clone());
641 if ty.should_lazy_static() {
642 CodegenTy::LazyStaticRef(Arc::new(ty))
643 } else {
644 ty
645 }
646 }
647 Item::Enum(_) => CodegenTy::Adt(AdtDef {
648 did: def_id,
649 kind: AdtKind::Enum,
650 }),
651 Item::NewType(t) => CodegenTy::Adt(AdtDef {
652 did: def_id,
653 kind: AdtKind::NewType(Arc::new(self.codegen_item_ty(t.ty.kind.clone()))),
654 }),
655 Item::Message(_) => CodegenTy::Adt(AdtDef {
656 did: def_id,
657 kind: AdtKind::Struct,
658 }),
659 _ => panic!("Unexpected item type for path: {:?}", item),
660 },
661 NodeKind::Variant(_v) => {
662 let parent_def_id = node.parent.unwrap();
664 CodegenTy::Adt(AdtDef {
665 did: parent_def_id,
666 kind: AdtKind::Enum,
667 })
668 }
669 NodeKind::Field(_) | NodeKind::Method(_) | NodeKind::Arg(_) => {
670 panic!("Unexpected node kind for path: {:?}", node.kind)
671 }
672 }
673 }
674
675 fn convert_element_expr(&self, expr: &str, from: &CodegenTy, to: &CodegenTy) -> String {
676 if from == to {
677 if self.is_copy_ty(from) {
678 expr.to_string()
679 } else {
680 format!("({expr}).clone()", expr = expr)
681 }
682 } else if let Some((converted, _)) =
683 self.convert_codegen_ty_expr(expr.to_string().into(), from, to, false)
684 {
685 converted.to_string()
686 } else {
687 expr.to_string()
688 }
689 }
690
691 fn convert_owned_element_expr(&self, expr: &str, from: &CodegenTy, to: &CodegenTy) -> String {
692 if from == to {
693 return expr.to_string();
694 }
695
696 match (from, to) {
697 (
698 CodegenTy::Adt(AdtDef {
699 kind: AdtKind::NewType(inner),
700 ..
701 }),
702 _,
703 ) => {
704 let inner_expr = format!("({expr}).0", expr = expr);
705 self.convert_owned_element_expr(&inner_expr, inner, to)
706 }
707 (
708 _,
709 CodegenTy::Adt(AdtDef {
710 kind: AdtKind::NewType(inner),
711 did,
712 }),
713 ) => {
714 let inner_expr = self.convert_owned_element_expr(expr, from, inner);
715 let ctor = self.cur_related_item_path(*did);
716 format!("{ctor}({inner_expr})")
717 }
718 _ => self
719 .convert_codegen_ty_expr(expr.to_string().into(), from, to, true)
720 .map(|(converted, _)| converted.to_string())
721 .unwrap_or_else(|| expr.to_string()),
722 }
723 }
724
725 fn is_copy_ty(&self, ty: &CodegenTy) -> bool {
726 matches!(
727 ty,
728 CodegenTy::Bool
729 | CodegenTy::I8
730 | CodegenTy::I16
731 | CodegenTy::I32
732 | CodegenTy::I64
733 | CodegenTy::UInt32
734 | CodegenTy::UInt64
735 | CodegenTy::U8
736 | CodegenTy::F32
737 | CodegenTy::F64
738 | CodegenTy::OrderedF64
739 )
740 }
741
742 pub fn default_val(&self, f: &Field) -> Option<(FastStr, bool )> {
743 f.default.as_ref().map(|d| {
744 let ty = self.codegen_item_ty(f.ty.kind.clone());
745 match self
746 .lit_as_rvalue(d, &ty)
747 .with_context(|| format!("calc the default value for field {}", f.name))
748 {
749 Ok(v) => v,
750 Err(err) => {
751 panic!("{err}")
752 }
753 }
754 })
755 }
756
757 fn map_literal_expr(
758 &self,
759 entries: &[(Literal, Literal)],
760 key_ty: &Arc<CodegenTy>,
761 value_ty: &Arc<CodegenTy>,
762 btree: bool,
763 ) -> anyhow::Result<FastStr> {
764 let key_ty = &**key_ty;
765 let value_ty = &**value_ty;
766 let len = entries.len();
767 let kvs = entries
768 .iter()
769 .map(|(k, v)| {
770 let (k_expr, _) = self.lit_into_ty(k, key_ty)?;
771 let (v_expr, _) = self.lit_into_ty(v, value_ty)?;
772 anyhow::Ok(format!("map.insert({k_expr}, {v_expr});"))
773 })
774 .try_collect::<_, Vec<_>, _>()?
775 .join("");
776 let new = if btree {
777 "::std::collections::BTreeMap::new()".to_string()
778 } else {
779 format!("::pilota::AHashMap::with_capacity({len})")
780 };
781 Ok(format! {r#"{{
782 let mut map = {new};
783 {kvs}
784 map
785 }}"#}
786 .into())
787 }
788
789 fn lit_as_rvalue(
790 &self,
791 lit: &Literal,
792 ty: &CodegenTy,
793 ) -> anyhow::Result<(FastStr, bool )> {
794 anyhow::Ok(match (lit, ty) {
795 (Literal::Map(m), CodegenTy::LazyStaticRef(map)) => match &**map {
796 CodegenTy::Map(k_ty, v_ty) => (self.map_literal_expr(m, k_ty, v_ty, false)?, false),
797 CodegenTy::BTreeMap(k_ty, v_ty) => {
798 (self.map_literal_expr(m, k_ty, v_ty, true)?, false)
799 }
800 _ => panic!("invalid map type {map:?}"),
801 },
802 (Literal::Map(m), CodegenTy::Map(k_ty, v_ty)) => {
803 (self.map_literal_expr(m, k_ty, v_ty, false)?, false)
804 }
805 (Literal::Map(m), CodegenTy::BTreeMap(k_ty, v_ty)) => {
806 (self.map_literal_expr(m, k_ty, v_ty, true)?, false)
807 }
808 (Literal::List(l), CodegenTy::LazyStaticRef(map)) => match &**map {
809 CodegenTy::Map(_, _) => {
810 assert!(l.is_empty());
811 ("::pilota::AHashMap::new()".into(), false)
812 }
813 CodegenTy::BTreeMap(_, _) => {
814 assert!(l.is_empty());
815 ("::std::collections::BTreeMap::new()".into(), false)
816 }
817 CodegenTy::Set(inner) => {
818 if l.is_empty() {
819 ("::pilota::AHashSet::new()".into(), false)
820 } else {
821 let stream = self.list_stream(l, inner)?;
822 (
823 format!("::pilota::AHashSet::from([{stream}])").into(),
824 false,
825 )
826 }
827 }
828 CodegenTy::BTreeSet(inner) => {
829 if l.is_empty() {
830 ("::std::collections::BTreeSet::new()".into(), false)
831 } else {
832 let stream = self.list_stream(l, inner)?;
833 (
834 format!("::std::collections::BTreeSet::from([{stream}])").into(),
835 false,
836 )
837 }
838 }
839 _ => panic!("invalid map type {map:?}"),
840 },
841 (Literal::List(l), CodegenTy::Map(_, _)) => {
842 assert!(l.is_empty());
843 ("::pilota::AHashMap::new()".into(), false)
844 }
845 (Literal::List(l), CodegenTy::BTreeMap(_, _)) => {
846 assert!(l.is_empty());
847 ("::std::collections::BTreeMap::new()".into(), false)
848 }
849 _ => self.lit_into_ty(lit, ty)?,
850 })
851 }
852
853 fn ident_into_ty(
854 &self,
855 did: DefId,
856 ident_ty: &CodegenTy,
857 target: &CodegenTy,
858 ) -> (FastStr, bool ) {
859 let stream = self.cur_related_item_path(did);
860 if ident_ty == target {
861 return (stream.clone(), true);
862 }
863
864 let ident_norm = self.normalize_codegen_ty(ident_ty);
865 let target_norm = self.normalize_codegen_ty(target);
866
867 if ident_norm == target_norm {
868 if let Some((converted, is_const)) =
869 self.convert_codegen_ty_expr(stream.clone(), ident_ty, target, true)
870 {
871 return (converted, is_const);
872 }
873 }
874
875 match (ident_ty, target) {
876 (CodegenTy::Str, CodegenTy::FastStr) => (
877 format!("::pilota::FastStr::from_static_str({stream})").into(),
878 true,
879 ),
880 (
881 CodegenTy::Adt(AdtDef {
882 did: _,
883 kind: AdtKind::Enum,
884 }),
885 CodegenTy::I64,
886 )
887 | (
888 CodegenTy::Adt(AdtDef {
889 did: _,
890 kind: AdtKind::Enum,
891 }),
892 CodegenTy::I32,
893 )
894 | (
895 CodegenTy::Adt(AdtDef {
896 did: _,
897 kind: AdtKind::Enum,
898 }),
899 CodegenTy::I16,
900 )
901 | (
902 CodegenTy::Adt(AdtDef {
903 did: _,
904 kind: AdtKind::Enum,
905 }),
906 CodegenTy::I8,
907 ) => {
908 let stream = self.cur_related_item_path(did);
909 let target = match target {
910 CodegenTy::I64 => "i64",
911 CodegenTy::I32 => "i32",
912 CodegenTy::I16 => "i16",
913 CodegenTy::I8 => "i8",
914 _ => unreachable!(),
915 };
916 (format!("({stream}.inner() as {target})").into(), true)
917 }
918 _ => panic!("invalid convert {ident_ty:?} to {target:?}"),
919 }
920 }
921
922 fn normalize_codegen_ty(&self, ty: &CodegenTy) -> CodegenTy {
924 match ty {
925 CodegenTy::Adt(AdtDef {
926 kind: AdtKind::NewType(inner),
927 ..
928 }) => self.normalize_codegen_ty(inner),
929 CodegenTy::LazyStaticRef(inner) | CodegenTy::StaticRef(inner) => {
930 self.normalize_codegen_ty(inner)
931 }
932 CodegenTy::Str => CodegenTy::FastStr,
933 CodegenTy::FastStr => CodegenTy::FastStr,
934 CodegenTy::Vec(inner) => CodegenTy::Vec(Arc::new(self.normalize_codegen_ty(inner))),
935 CodegenTy::Array(inner, size) => {
936 CodegenTy::Array(Arc::new(self.normalize_codegen_ty(inner)), *size)
937 }
938 CodegenTy::Set(inner) => CodegenTy::Set(Arc::new(self.normalize_codegen_ty(inner))),
939 CodegenTy::BTreeSet(inner) => {
940 CodegenTy::BTreeSet(Arc::new(self.normalize_codegen_ty(inner)))
941 }
942 CodegenTy::Map(k, v) => CodegenTy::Map(
943 Arc::new(self.normalize_codegen_ty(k)),
944 Arc::new(self.normalize_codegen_ty(v)),
945 ),
946 CodegenTy::BTreeMap(k, v) => CodegenTy::BTreeMap(
947 Arc::new(self.normalize_codegen_ty(k)),
948 Arc::new(self.normalize_codegen_ty(v)),
949 ),
950 CodegenTy::Arc(inner) => CodegenTy::Arc(Arc::new(self.normalize_codegen_ty(inner))),
951 _ => ty.clone(),
952 }
953 }
954
955 fn convert_codegen_ty_expr(
956 &self,
957 expr: FastStr,
958 from: &CodegenTy,
959 to: &CodegenTy,
960 owned: bool,
961 ) -> Option<(FastStr, bool)> {
962 if from == to {
963 if owned || self.is_copy_ty(from) {
964 return Some((expr, true));
965 }
966 let expr_str = expr.to_string();
967 return Some((format!("({expr_str}).clone()").into(), false));
968 }
969
970 match (from, to) {
971 (
972 CodegenTy::Adt(AdtDef {
973 kind: AdtKind::NewType(inner),
974 ..
975 }),
976 _,
977 ) => {
978 if owned {
979 let expr_str = expr.to_string();
980 let inner_expr: FastStr = format!("({expr_str}).0").into();
981 let (converted, is_const) =
982 self.convert_codegen_ty_expr(inner_expr, inner, to, true)?;
983 Some((converted, is_const))
984 } else {
985 let expr_str = expr.to_string();
986 let cloned: FastStr = format!("{expr_str}.clone()").into();
987 let (converted, _) = self.convert_codegen_ty_expr(cloned, inner, to, true)?;
988 Some((converted, false))
989 }
990 }
991 (
992 _,
993 CodegenTy::Adt(AdtDef {
994 kind: AdtKind::NewType(inner),
995 did,
996 }),
997 ) => {
998 let (inner_expr, _) = self.convert_codegen_ty_expr(expr, from, inner, owned)?;
999 let ident = self.cur_related_item_path(*did);
1000 Some((format!("{ident}({inner_expr})").into(), false))
1001 }
1002 (CodegenTy::Str, CodegenTy::FastStr) => {
1003 let expr_str = expr.to_string();
1004 Some((
1005 format!("::pilota::FastStr::from_static_str({expr_str})").into(),
1006 true,
1007 ))
1008 }
1009 (CodegenTy::LazyStaticRef(inner), _) => {
1010 let expr_str = expr.to_string();
1011 let needs_clone = !expr_str.trim_end().ends_with(".clone()");
1012 let owned_expr: FastStr = if needs_clone {
1013 format!("{expr_str}.clone()").into()
1014 } else {
1015 expr_str.into()
1016 };
1017 let (converted, _) = self.convert_codegen_ty_expr(owned_expr, inner, to, true)?;
1018 Some((converted, false))
1019 }
1020 (CodegenTy::StaticRef(inner), _) => {
1021 let expr_str = expr.to_string();
1022 let needs_clone = !expr_str.trim_end().ends_with(".clone()");
1023 let owned_expr: FastStr = if needs_clone {
1024 format!("{expr_str}.clone()").into()
1025 } else {
1026 expr_str.into()
1027 };
1028 let (converted, _) = self.convert_codegen_ty_expr(owned_expr, inner, to, true)?;
1029 Some((converted, false))
1030 }
1031 (CodegenTy::Vec(from_inner), CodegenTy::Vec(to_inner)) => {
1032 if from_inner == to_inner {
1033 return Some((expr, owned));
1034 }
1035 let expr_str = expr.to_string();
1036 if owned {
1037 let body = self.convert_owned_element_expr("el.clone()", from_inner, to_inner);
1038 let converted = format!(
1039 r#"{{
1040 {expr_str}
1041 .iter()
1042 .map(|el| {body})
1043 .collect::<::std::vec::Vec<_>>()
1044 }}"#
1045 )
1046 .into();
1047 Some((converted, false))
1048 } else {
1049 let body = self.convert_element_expr("el", from_inner, to_inner);
1050 let converted = format!(
1051 r#"{{
1052 {expr_str}
1053 .iter()
1054 .map(|el| {body})
1055 .collect::<::std::vec::Vec<_>>()
1056 }}"#
1057 )
1058 .into();
1059 Some((converted, false))
1060 }
1061 }
1062 (CodegenTy::Set(from_inner), CodegenTy::Set(to_inner)) => {
1063 if from_inner == to_inner {
1064 return Some((expr, owned));
1065 }
1066 let expr_str = expr.to_string();
1067 if owned {
1068 let body = self.convert_owned_element_expr("el.clone()", from_inner, to_inner);
1069 let converted = format!(
1070 r#"{{
1071 {expr_str}
1072 .iter()
1073 .map(|el| {body})
1074 .collect::<::pilota::AHashSet<_>>()
1075 }}"#
1076 )
1077 .into();
1078 Some((converted, false))
1079 } else {
1080 let body = self.convert_element_expr("el", from_inner, to_inner);
1081 let converted = format!(
1082 r#"{{
1083 {expr_str}
1084 .iter()
1085 .map(|el| {body})
1086 .collect::<::pilota::AHashSet<_>>()
1087 }}"#
1088 )
1089 .into();
1090 Some((converted, false))
1091 }
1092 }
1093 (CodegenTy::BTreeSet(from_inner), CodegenTy::BTreeSet(to_inner)) => {
1094 if from_inner == to_inner {
1095 return Some((expr, owned));
1096 }
1097 let expr_str = expr.to_string();
1098 if owned {
1099 let body = self.convert_owned_element_expr("el.clone()", from_inner, to_inner);
1100 let converted = format!(
1101 r#"{{
1102 {expr_str}.iter()
1103 .map(|el| {body})
1104 .collect::<::std::collections::BTreeSet<_>>()
1105 }}"#
1106 )
1107 .into();
1108 Some((converted, false))
1109 } else {
1110 let body = self.convert_element_expr("el", from_inner, to_inner);
1111 let converted = format!(
1112 r#"{{
1113 {expr_str}.iter()
1114 .map(|el| {body})
1115 .collect::<::std::collections::BTreeSet<_>>()
1116 }}"#
1117 )
1118 .into();
1119 Some((converted, false))
1120 }
1121 }
1122 (CodegenTy::Map(from_k, from_v), CodegenTy::Map(to_k, to_v)) => {
1123 if from_k == to_k && from_v == to_v {
1124 return Some((expr, owned));
1125 }
1126 let expr_str = expr.to_string();
1127 if owned {
1128 let k_body = self.convert_owned_element_expr("k.clone()", from_k, to_k);
1129 let v_body = self.convert_owned_element_expr("v.clone()", from_v, to_v);
1130 let converted = format!(
1131 r#"{{
1132 {expr_str}
1133 .iter()
1134 .map(|(k, v)| ({k_body}, {v_body}))
1135 .collect::<::pilota::AHashMap<_, _>>()
1136 }}"#
1137 )
1138 .into();
1139 Some((converted, false))
1140 } else {
1141 let k_body = self.convert_element_expr("k", from_k, to_k);
1142 let v_body = self.convert_element_expr("v", from_v, to_v);
1143 let converted = format!(
1144 r#"{{
1145 {expr_str}
1146 .iter()
1147 .map(|(k, v)| ({k_body}, {v_body}))
1148 .collect::<::pilota::AHashMap<_, _>>()
1149 }}"#
1150 )
1151 .into();
1152 Some((converted, false))
1153 }
1154 }
1155 (CodegenTy::BTreeMap(from_k, from_v), CodegenTy::BTreeMap(to_k, to_v)) => {
1156 if from_k == to_k && from_v == to_v {
1157 return Some((expr, owned));
1158 }
1159 let expr_str = expr.to_string();
1160 if owned {
1161 let k_body = self.convert_owned_element_expr("k.clone()", from_k, to_k);
1162 let v_body = self.convert_owned_element_expr("v.clone()", from_v, to_v);
1163 let converted = format!(
1164 r#"{{
1165 {expr_str}
1166 .iter()
1167 .map(|(k, v)| ({k_body}, {v_body}))
1168 .collect::<::std::collections::BTreeMap<_, _>>()
1169 }}"#
1170 )
1171 .into();
1172 Some((converted, false))
1173 } else {
1174 let k_body = self.convert_element_expr("k", from_k, to_k);
1175 let v_body = self.convert_element_expr("v", from_v, to_v);
1176 let converted = format!(
1177 r#"{{
1178 {expr_str}
1179 .iter()
1180 .map(|(k, v)| ({k_body}, {v_body}))
1181 .collect::<::std::collections::BTreeMap<_, _>>()
1182 }}"#
1183 )
1184 .into();
1185 Some((converted, false))
1186 }
1187 }
1188 _ => None,
1189 }
1190 }
1191
1192 fn lit_into_ty(
1193 &self,
1194 lit: &Literal,
1195 ty: &CodegenTy,
1196 ) -> anyhow::Result<(FastStr, bool )> {
1197 Ok(match (lit, ty) {
1198 (Literal::Path(p), ty) => {
1199 let ident_ty = self.get_codegen_ty_for_path(p.did);
1200
1201 if matches!(
1202 ty,
1203 CodegenTy::Map(_, _)
1204 | CodegenTy::BTreeMap(_, _)
1205 | CodegenTy::Set(_)
1206 | CodegenTy::BTreeSet(_)
1207 ) {
1208 let normalized_ident_ty = self.normalize_codegen_ty(&ident_ty);
1209 if matches!(
1210 normalized_ident_ty,
1211 CodegenTy::Map(_, _)
1212 | CodegenTy::BTreeMap(_, _)
1213 | CodegenTy::Set(_)
1214 | CodegenTy::BTreeSet(_)
1215 ) {
1216 let stream = self.cur_related_item_path(p.did);
1217 if let Some((converted, _)) =
1218 self.convert_codegen_ty_expr(stream.clone(), &ident_ty, ty, false)
1219 {
1220 return Ok((converted, false));
1221 }
1222 }
1223 }
1224
1225 self.ident_into_ty(p.did, &ident_ty, ty)
1226 }
1227 (Literal::String(s), CodegenTy::Str) => {
1228 let escaped = s.escape_debug().to_string();
1229 (format!("\"{escaped}\"").into(), true)
1230 }
1231 (Literal::String(s), CodegenTy::String) => {
1232 let escaped = s.escape_debug().to_string();
1233 (format!("\"{escaped}\".to_string()").into(), false)
1234 }
1235 (Literal::String(s), CodegenTy::FastStr) => {
1236 let escaped = s.escape_debug().to_string();
1237 (
1238 format!("::pilota::FastStr::from_static_str(\"{escaped}\")").into(),
1239 true,
1240 )
1241 }
1242 (Literal::Int(i), CodegenTy::I8) => (format! { "{i}i8" }.into(), true),
1243 (Literal::Int(i), CodegenTy::I16) => (format! { "{i}i16" }.into(), true),
1244 (Literal::Int(i), CodegenTy::I32) => (format! { "{i}i32" }.into(), true),
1245 (Literal::Int(i), CodegenTy::I64) => (format! { "{i}i64" }.into(), true),
1246 (Literal::Int(i), CodegenTy::F32) => {
1247 let f = (*i) as f32;
1248 (format!("{f}f32").into(), true)
1249 }
1250 (Literal::Int(i), CodegenTy::F64) => {
1251 let f = (*i) as f64;
1252 (format!("{f}f64").into(), true)
1253 }
1254 (
1255 Literal::Int(i),
1256 CodegenTy::Adt(AdtDef {
1257 did,
1258 kind: AdtKind::Enum,
1259 }),
1260 ) => {
1261 let item = self.item(*did).unwrap();
1262 let e = match &*item {
1263 Item::Enum(e) => e,
1264 _ => panic!("invalid enum"),
1265 };
1266
1267 (
1268 e.variants.iter().find(|v| v.discr == Some(*i)).map_or_else(
1269 || panic!("invalid enum value"),
1270 |v| self.cur_related_item_path(v.did),
1271 ),
1272 true,
1273 )
1274 }
1275 (Literal::Float(f), CodegenTy::F32) => {
1276 let f = f.parse::<f32>().unwrap();
1277 (format! { "{f}f32" }.into(), true)
1278 }
1279 (Literal::Float(f), CodegenTy::F64) => {
1280 let f = f.parse::<f64>().unwrap();
1281 (format! { "{f}f64" }.into(), true)
1282 }
1283 (Literal::Float(f), CodegenTy::OrderedF64) => {
1284 let f = f.parse::<f64>().unwrap();
1285 (format! { "::pilota::OrderedFloat({f}f64)" }.into(), true)
1286 }
1287 (
1288 l,
1289 CodegenTy::Adt(AdtDef {
1290 kind: AdtKind::NewType(inner_ty),
1291 did,
1292 }),
1293 ) => {
1294 let ident = self.cur_related_item_path(*did);
1295 let (stream, is_const) = self.lit_into_ty(l, inner_ty)?;
1296 (format! { "{ident}({stream})" }.into(), is_const)
1297 }
1298 (Literal::Map(_), CodegenTy::StaticRef(map)) => match &**map {
1299 CodegenTy::Map(_, _) | CodegenTy::BTreeMap(_, _) => {
1300 let lazy_map =
1301 self.def_lit("INNER_MAP", lit, &mut CodegenTy::LazyStaticRef(map.clone()))?;
1302 let stream = format! {
1303 r#"{{
1304 {lazy_map}
1305 &*INNER_MAP
1306 }}"#
1307 }
1308 .into();
1309 (stream, false)
1310 }
1311 _ => panic!("invalid map type {map:?}"),
1312 },
1313 (Literal::List(els), CodegenTy::Array(inner, _)) => {
1314 let stream = els
1315 .iter()
1316 .map(|el| self.lit_into_ty(el, inner))
1317 .try_collect::<_, Vec<_>, _>()?;
1318 let is_const = stream.iter().all(|(_, is_const)| *is_const);
1319 let stream = stream.into_iter().map(|(s, _)| s).join(",");
1320
1321 (format! {"[{stream}]" }.into(), is_const)
1322 }
1323 (Literal::List(els), CodegenTy::Vec(inner)) => {
1324 let stream = self.list_stream(els, inner)?;
1325 (format! { "::std::vec![{stream}]" }.into(), false)
1326 }
1327 (Literal::List(els), CodegenTy::Set(inner)) => {
1328 let stream = self.list_stream(els, inner)?;
1329 (
1330 format! { "::pilota::AHashSet::from([{stream}])" }.into(),
1331 false,
1332 )
1333 }
1334 (Literal::List(els), CodegenTy::BTreeSet(inner)) => {
1335 let stream = self.list_stream(els, inner)?;
1336 (
1337 format! { "::std::collections::BTreeSet::from([{stream}])" }.into(),
1338 false,
1339 )
1340 }
1341 (Literal::Bool(b), CodegenTy::Bool) => (format! { "{b}" }.into(), true),
1342 (Literal::Int(i), CodegenTy::Bool) => {
1343 let b = *i != 0;
1344 (format! { "{b}" }.into(), true)
1345 }
1346 (Literal::String(s), CodegenTy::Bytes) => {
1347 let escaped = s.escape_debug().to_string();
1348 (
1349 format!("::pilota::Bytes::from_static(\"{escaped}\".as_bytes())").into(),
1350 true,
1351 )
1352 }
1353 (
1354 Literal::Map(m),
1355 CodegenTy::Adt(AdtDef {
1356 did,
1357 kind: AdtKind::Struct,
1358 }),
1359 ) => {
1360 let def = self.item(*did).unwrap();
1361 let def = match &*def {
1362 Item::Message(m) => m,
1363 _ => panic!(),
1364 };
1365
1366 let fields: Vec<_> = def
1367 .fields
1368 .iter()
1369 .map(|f| {
1370 let v = m.iter().find_map(|(k, v)| {
1371 let k = match k {
1372 Literal::String(s) => s,
1373 _ => panic!(),
1374 };
1375 if **k == **f.name { Some(v) } else { None }
1376 });
1377
1378 let name = self.rust_name(f.did);
1379
1380 if let Some(v) = v {
1381 let (mut v, is_const) =
1382 self.lit_into_ty(v, &self.codegen_item_ty(f.ty.kind.clone()))?;
1383
1384 if f.is_optional() {
1385 v = format!("Some({v})").into()
1386 }
1387 anyhow::Ok((format!("{name}: {v}"), is_const))
1388 } else if f.is_optional() {
1389 anyhow::Ok((format!("{name}: None"), true))
1390 } else {
1391 anyhow::Ok((format!("{name}: Default::default()"), false))
1392 }
1393 })
1394 .try_collect()?;
1395 let is_const = fields.iter().all(|(_, is_const)| *is_const);
1396 let mut fields = fields.into_iter().map(|f| f.0).collect::<Vec<_>>();
1397 if !def.is_wrapper && self.config.with_field_mask {
1398 fields.push("_field_mask: ::std::option::Option::None".to_string());
1399 }
1400 let fields = fields.join(",");
1401
1402 let name = self.cur_related_item_path(*did);
1403
1404 (
1405 format! {
1406 r#"{name} {{
1407 {fields}
1408 }}"#
1409 }
1410 .into(),
1411 is_const,
1412 )
1413 }
1414 (Literal::Map(m), CodegenTy::Map(k_ty, v_ty)) => {
1415 (self.map_literal_expr(m, k_ty, v_ty, false)?, false)
1416 }
1417 (Literal::Map(m), CodegenTy::BTreeMap(k_ty, v_ty)) => {
1418 (self.map_literal_expr(m, k_ty, v_ty, true)?, false)
1419 }
1420 (Literal::List(l), CodegenTy::Map(_, _)) => {
1421 assert!(l.is_empty());
1422 ("::pilota::AHashMap::new()".into(), false)
1423 }
1424 (Literal::List(l), CodegenTy::BTreeMap(_, _)) => {
1425 assert!(l.is_empty());
1426 ("::std::collections::BTreeMap::new()".into(), false)
1427 }
1428 _ => {
1429 let (def_path, idl_file) = with_cur_item(|def_id| {
1430 let def_path = self.item_path(def_id).iter().join("::");
1431 let file_path = self
1432 .db
1433 .node(def_id)
1434 .and_then(|node| {
1435 self.db
1436 .file_paths()
1437 .get(&node.file_id)
1438 .map(|path| path.display().to_string())
1439 })
1440 .unwrap_or_else(|| "<unknown>".to_string());
1441 (def_path, file_path)
1442 });
1443
1444 let error_message = format!(
1445 "unexpected literal {lit:?} with ty {ty}, def_path: {def_path}, idl_file: \
1446 {idl_file}"
1447 );
1448
1449 panic!("{error_message}");
1450 }
1451 })
1452 }
1453
1454 #[inline]
1455 fn list_stream(&self, els: &[Literal], inner: &Arc<CodegenTy>) -> anyhow::Result<String> {
1456 Ok(els
1457 .iter()
1458 .map(|el| self.lit_into_ty(el, inner))
1459 .try_collect::<_, Vec<_>, _>()?
1460 .into_iter()
1461 .map(|(s, _)| s)
1462 .join(","))
1463 }
1464
1465 pub(crate) fn def_lit(
1466 &self,
1467 name: &str,
1468 lit: &Literal,
1469 ty: &mut CodegenTy,
1470 ) -> anyhow::Result<String> {
1471 let should_lazy_static = ty.should_lazy_static();
1472 if let (Literal::List(lit), CodegenTy::Array(_, size)) = (lit, &mut *ty) {
1473 *size = lit.len()
1474 }
1475 Ok(if should_lazy_static {
1476 let lit = self.lit_as_rvalue(lit, ty)?.0;
1477 format! {r#"
1478 pub static {name}: ::std::sync::LazyLock<{ty}> = ::std::sync::LazyLock::new(|| {{
1479 {lit}
1480 }});
1481 "#}
1482 } else {
1483 let (lit, is_const) = self.lit_into_ty(lit, ty)?;
1484 if is_const {
1485 format!(r#"pub const {name}: {ty} = {lit};"#)
1486 } else {
1487 format! {r#"
1488 pub static {name}: ::std::sync::LazyLock<{ty}> = ::std::sync::LazyLock::new(|| {{
1489 {lit}
1490 }});
1491 "#}
1492 }
1493 })
1494 }
1495
1496 pub fn rust_name(&self, def_id: DefId) -> Symbol {
1497 let node = self.node(def_id).unwrap();
1498
1499 if let Some(name) = self
1500 .tags(node.tags)
1501 .and_then(|tags| tags.get::<crate::tags::PilotaName>().cloned())
1502 {
1503 return name.0.into();
1504 }
1505
1506 if !self.config.change_case || self.cache.names.contains_key(&def_id) {
1507 return node.name();
1508 }
1509
1510 match self.node(def_id).unwrap().kind {
1511 NodeKind::Item(item) => match &*item {
1512 crate::rir::Item::Message(m) => (&**m.name).struct_ident(),
1513 crate::rir::Item::Enum(e) => (&**e.name).enum_ident(),
1514 crate::rir::Item::Service(s) => (&**s.name).trait_ident(),
1515 crate::rir::Item::NewType(t) => (&**t.name).newtype_ident(),
1516 crate::rir::Item::Const(c) => (&**c.name).const_ident(),
1517 crate::rir::Item::Mod(m) => (&**m.name).mod_ident(),
1518 },
1519 NodeKind::Variant(v) => {
1520 let parent = self.node(def_id).unwrap().parent.unwrap();
1521 let item = self.expect_item(parent);
1522 match &*item {
1523 rir::Item::Enum(e) => {
1524 if e.repr.is_some() {
1525 (&**v.name).const_ident()
1526 } else {
1527 (&**v.name).variant_ident()
1528 }
1529 }
1530 _ => unreachable!(),
1531 }
1532 }
1533 NodeKind::Field(f) => (&**f.name).field_ident(),
1534 NodeKind::Method(m) => (&**m.name).fn_ident(),
1535 NodeKind::Arg(a) => (&**a.name).field_ident(),
1536 }
1537 .into()
1538 }
1539
1540 pub fn mod_path(&self, def_id: DefId) -> Arc<[Symbol]> {
1541 self.source.path_resolver.mod_prefix(self, def_id)
1542 }
1543
1544 pub fn mod_index(&self, def_id: DefId) -> ModPath {
1545 let mod_path = self
1546 .mod_path(def_id)
1547 .iter()
1548 .map(|s| s.0.clone())
1549 .collect_vec();
1550
1551 match &*self.source.mode {
1552 Mode::Workspace(_) => ModPath::from(&mod_path[1..]),
1553 Mode::SingleFile { .. } => ModPath::from(mod_path),
1554 }
1555 }
1556
1557 pub fn item_path(&self, def_id: DefId) -> Arc<[Symbol]> {
1558 self.source.path_resolver.path_for_def_id(self, def_id)
1559 }
1560
1561 fn related_path(&self, p1: &[Symbol], p2: &[Symbol]) -> FastStr {
1562 self.source.path_resolver.related_path(p1, p2)
1563 }
1564
1565 pub fn cur_related_item_path(&self, did: DefId) -> FastStr {
1566 let a = with_cur_item(|def_id| def_id);
1567 self.related_item_path(a, did)
1568 }
1569
1570 pub fn related_item_path(&self, a: DefId, b: DefId) -> FastStr {
1571 let cur_item_path = self.item_path(a);
1572 let mut mod_segs = vec![];
1573
1574 cur_item_path[..cur_item_path.len() - 1]
1575 .iter()
1576 .for_each(|p| {
1577 mod_segs.push(p.clone());
1578 });
1579
1580 let other_item_path = self.item_path(b);
1581 self.related_path(&mod_segs, &other_item_path)
1582 }
1583
1584 #[allow(clippy::single_match)]
1585 pub fn exec_plugin<P: Plugin>(&self, mut p: P) {
1586 p.on_codegen_uint(self, &self.cache.codegen_items);
1587
1588 p.on_emit(self)
1589 }
1590
1591 pub(crate) fn workspace_info(&self) -> &WorkspaceInfo {
1592 let Mode::Workspace(info) = &*self.source.mode else {
1593 panic!(
1594 "can not access workspace info in mode `{:?}`",
1595 self.source.mode
1596 )
1597 };
1598 info
1599 }
1600
1601 pub fn config(&self, crate_id: &CrateId) -> &serde_yaml::Value {
1615 &self.find_service(crate_id.main_file).config
1616 }
1617
1618 pub(crate) fn crate_name(&self, location: &DefLocation) -> FastStr {
1619 match location {
1620 DefLocation::Fixed(crate_id, _) => {
1621 let main_file = crate_id.main_file;
1622 let service = self.find_service(main_file);
1623 self.config(crate_id)
1624 .get("crate_name")
1625 .and_then(|s| s.as_str().map(FastStr::new))
1626 .unwrap_or_else(|| {
1627 service
1628 .path
1629 .file_stem()
1630 .unwrap()
1631 .to_str()
1632 .unwrap()
1633 .replace('.', "_")
1634 .into()
1635 })
1636 }
1637 DefLocation::Dynamic => self.config.common_crate_name.clone(),
1638 }
1639 }
1640
1641 fn find_service(&self, file_id: FileId) -> &crate::IdlService {
1642 self.source
1643 .services
1644 .iter()
1645 .find(|s| {
1646 let path = s
1647 .path
1648 .normalize()
1649 .unwrap_or_else(|err| {
1650 panic!("normalize path {} failed: {:?}", s.path.display(), err)
1651 })
1652 .into_path_buf();
1653 self.file_id(path.clone()).unwrap_or_else(|| {
1654 panic!(
1655 "file_id not found for path {} in file_ids_map {:?}",
1656 path.display(),
1657 self.file_ids_map()
1658 )
1659 }) == file_id
1660 })
1661 .unwrap()
1662 }
1663}
1664
1665pub mod tls {
1666
1667 use scoped_tls::scoped_thread_local;
1668
1669 use super::Context;
1670 use crate::DefId;
1671
1672 scoped_thread_local!(pub static CONTEXT: Context);
1673 scoped_thread_local!(pub static CUR_ITEM: DefId);
1674
1675 pub fn with_cx<T, F>(f: F) -> T
1676 where
1677 F: FnOnce(&Context) -> T,
1678 {
1679 CONTEXT.with(|cx| f(cx))
1680 }
1681
1682 pub fn with_cur_item<T, F>(f: F) -> T
1683 where
1684 F: FnOnce(DefId) -> T,
1685 {
1686 CUR_ITEM.with(|def_id| f(*def_id))
1687 }
1688}
1689
1690#[cfg(test)]
1691mod tests {
1692 use std::{collections::HashMap, path::PathBuf, sync::Arc};
1693
1694 use anyhow::Result;
1695 use faststr::FastStr;
1696 use pilota::Bytes;
1697 use rustc_hash::{FxHashMap, FxHashSet};
1698
1699 use super::*;
1700 use crate::{
1701 middle::{
1702 ext::{FileExts, ItemExts},
1703 rir::{self, FieldKind, Message},
1704 ty::{CodegenTy, Ty, TyKind},
1705 },
1706 symbol::{Ident, Symbol},
1707 };
1708
1709 fn make_test_context() -> Context {
1710 let mode = Arc::new(Mode::SingleFile {
1711 file_path: PathBuf::from("dummy.rs"),
1712 });
1713 let services: Arc<[crate::IdlService]> =
1714 Arc::from(Vec::<crate::IdlService>::new().into_boxed_slice());
1715 Context {
1716 db: RootDatabase::default(),
1717 source: Source {
1718 source_type: SourceType::Thrift,
1719 services,
1720 mode: mode.clone(),
1721 path_resolver: Arc::new(DefaultPathResolver),
1722 selection_kind: SelectionKind::Service,
1723 },
1724 config: Config {
1725 change_case: false,
1726 split: false,
1727 with_descriptor: false,
1728 with_field_mask: false,
1729 touch_all: false,
1730 common_crate_name: "common".into(),
1731 with_comments: false,
1732 },
1733 cache: Cache {
1734 adjusts: Arc::new(DashMap::default()),
1735 mod_idxes: AHashMap::new(),
1736 codegen_items: Vec::new(),
1737 mod_items: AHashMap::new(),
1738 def_mod: HashMap::new(),
1739 mod_files: HashMap::new(),
1740 keep_unknown_fields: Arc::new(FxHashSet::default()),
1741 location_map: Arc::new(FxHashMap::default()),
1742 entry_map: Arc::new(HashMap::default()),
1743 plugin_gen: Arc::new(DashMap::default()),
1744 dedups: Vec::new(),
1745 names: FxHashMap::default(),
1746 },
1747 }
1748 }
1749
1750 #[test]
1751 fn collect_items_traverses_field_dependencies() {
1752 let file_id = FileId::from_u32(0);
1753 let root_id = DefId::from_u32(0);
1754 let dep_id = DefId::from_u32(1);
1755 let field_id = DefId::from_u32(2);
1756 let tag_id = TagId::from_u32(0);
1757
1758 let dep_message = Arc::new(rir::Item::Message(Message {
1759 name: Ident::from("Dep"),
1760 fields: Vec::new(),
1761 is_wrapper: false,
1762 item_exts: ItemExts::Thrift,
1763 leading_comments: FastStr::new(""),
1764 trailing_comments: FastStr::new(""),
1765 }));
1766
1767 let dep_node = rir::Node {
1768 file_id,
1769 kind: rir::NodeKind::Item(dep_message.clone()),
1770 parent: None,
1771 tags: tag_id,
1772 related_nodes: Vec::new(),
1773 };
1774
1775 let field_ty = Ty {
1776 kind: TyKind::Path(rir::Path {
1777 kind: rir::DefKind::Type,
1778 did: dep_id,
1779 }),
1780 tags_id: tag_id,
1781 };
1782
1783 let field = Arc::new(rir::Field {
1784 did: field_id,
1785 name: Ident::from("dep"),
1786 id: 1,
1787 ty: field_ty,
1788 kind: FieldKind::Required,
1789 tags_id: tag_id,
1790 default: None,
1791 item_exts: ItemExts::Thrift,
1792 leading_comments: FastStr::new(""),
1793 trailing_comments: FastStr::new(""),
1794 });
1795
1796 let root_message = Arc::new(rir::Item::Message(Message {
1797 name: Ident::from("Root"),
1798 fields: vec![field],
1799 is_wrapper: false,
1800 item_exts: ItemExts::Thrift,
1801 leading_comments: FastStr::new(""),
1802 trailing_comments: FastStr::new(""),
1803 }));
1804
1805 let root_node = rir::Node {
1806 file_id,
1807 kind: rir::NodeKind::Item(root_message.clone()),
1808 parent: None,
1809 tags: tag_id,
1810 related_nodes: Vec::new(),
1811 };
1812
1813 let mut nodes = FxHashMap::default();
1814 nodes.insert(root_id, root_node);
1815 nodes.insert(dep_id, dep_node);
1816
1817 let package = ItemPath::from(Arc::<[Symbol]>::from(
1818 Vec::<Symbol>::new().into_boxed_slice(),
1819 ));
1820 let file = rir::File {
1821 package,
1822 items: vec![root_id, dep_id],
1823 file_id,
1824 uses: Vec::new(),
1825 descriptor: Bytes::new(),
1826 extensions: FileExts::Thrift,
1827 comments: FastStr::new(""),
1828 };
1829
1830 let file_arc = Arc::new(file);
1831
1832 let mut file_ids_map = FxHashMap::default();
1833 let normalized = Arc::new(PathBuf::from("/tmp/test.thrift"));
1834 file_ids_map.insert(normalized.clone(), file_id);
1835
1836 let mut file_paths = FxHashMap::default();
1837 file_paths.insert(file_id, normalized);
1838
1839 let mut file_names = FxHashMap::default();
1840 file_names.insert(file_id, FastStr::from_static_str("test.thrift"));
1841
1842 let db = RootDatabase::default()
1843 .with_nodes(nodes)
1844 .with_files(vec![(file_id, file_arc)].into_iter())
1845 .with_file_ids_map(file_ids_map)
1846 .with_file_paths(file_paths)
1847 .with_file_names(file_names)
1848 .with_input_files(vec![file_id]);
1849
1850 let builder = ContextBuilder::new(
1851 db,
1852 Mode::SingleFile {
1853 file_path: PathBuf::from("/tmp/output.rs"),
1854 },
1855 vec![root_id],
1856 );
1857
1858 let collected = builder.collect_items(&[root_id]);
1859
1860 assert!(collected.contains(&root_id));
1861 assert!(collected.contains(&dep_id));
1862 assert_eq!(collected.len(), 2);
1863 }
1864
1865 #[test]
1866 fn lit_into_ty_handles_basic_literals() -> Result<()> {
1867 let cx = make_test_context();
1868
1869 let (expr, is_const) =
1870 cx.lit_into_ty(&Literal::String(Arc::from("hello")), &CodegenTy::FastStr)?;
1871 assert_eq!(&*expr, "::pilota::FastStr::from_static_str(\"hello\")");
1872 assert!(is_const);
1873
1874 let list_lit = Literal::List(vec![Literal::Int(1), Literal::Int(2)]);
1875 let vec_ty = CodegenTy::Vec(Arc::new(CodegenTy::I32));
1876 let (expr, is_const) = cx.lit_into_ty(&list_lit, &vec_ty)?;
1877 assert_eq!(&*expr, "::std::vec![1i32,2i32]");
1878 assert!(!is_const);
1879
1880 Ok(())
1881 }
1882
1883 #[test]
1884 fn convert_codegen_vec_owned_iter_clone() {
1885 let cx = make_test_context();
1886 let from_ty = CodegenTy::Vec(Arc::new(CodegenTy::Str));
1887 let to_ty = CodegenTy::Vec(Arc::new(CodegenTy::FastStr));
1888 let (expr, _) = cx
1889 .convert_codegen_ty_expr(FastStr::from_static_str("items"), &from_ty, &to_ty, true)
1890 .expect("vec conversion should succeed");
1891 let rendered = expr.to_string();
1892 assert!(rendered.contains(".iter()"));
1893 assert!(!rendered.contains("into_iter"));
1894 assert!(rendered.contains("::pilota::FastStr::from_static_str"));
1895 }
1896
1897 #[test]
1898 fn convert_codegen_set_owned_iter_clone() {
1899 let cx = make_test_context();
1900 let from_ty = CodegenTy::Set(Arc::new(CodegenTy::Str));
1901 let to_ty = CodegenTy::Set(Arc::new(CodegenTy::FastStr));
1902 let (expr, _) = cx
1903 .convert_codegen_ty_expr(
1904 FastStr::from_static_str("set_items"),
1905 &from_ty,
1906 &to_ty,
1907 true,
1908 )
1909 .expect("set conversion should succeed");
1910 let rendered = expr.to_string();
1911 assert!(rendered.contains(".iter()"));
1912 assert!(!rendered.contains("into_iter"));
1913 assert!(rendered.contains("::pilota::FastStr::from_static_str"));
1914 }
1915
1916 #[test]
1917 fn convert_codegen_map_owned_iter_clone() {
1918 let cx = make_test_context();
1919 let from_ty = CodegenTy::Map(Arc::new(CodegenTy::Str), Arc::new(CodegenTy::Str));
1920 let to_ty = CodegenTy::Map(Arc::new(CodegenTy::FastStr), Arc::new(CodegenTy::FastStr));
1921 let (expr, _) = cx
1922 .convert_codegen_ty_expr(
1923 FastStr::from_static_str("map_items"),
1924 &from_ty,
1925 &to_ty,
1926 true,
1927 )
1928 .expect("map conversion should succeed");
1929 let rendered = expr.to_string();
1930 assert!(rendered.contains(".iter()"));
1931 assert!(!rendered.contains("into_iter"));
1932 assert!(rendered.contains("::pilota::FastStr::from_static_str(k.clone())"));
1933 assert!(rendered.contains("::pilota::FastStr::from_static_str(v.clone())"));
1934 }
1935}