pilota_build/
db.rs

1pub mod cached_queries;
2mod salsa_ids;
3
4use std::{fmt::Debug, path::PathBuf, sync::Arc};
5
6pub use cached_queries::CachedQueries;
7use faststr::FastStr;
8use rustc_hash::{FxHashMap, FxHashSet};
9pub use salsa_ids::{IntoSalsa, SalsaDefId, SalsaFileId, SalsaTyKind};
10
11use crate::{
12    middle::{
13        context::{CrateId, DefLocation},
14        ext::pb::{Extendee, ExtendeeIndex},
15        ty::{CodegenTy, TyKind},
16    },
17    rir::{self, File, Item, Node},
18    symbol::{DefId, FileId},
19    tags::{TagId, Tags},
20};
21
22pub type ItemPath = Arc<Vec<FastStr>>;
23pub type TypeGraph = crate::middle::type_graph::TypeGraph;
24pub type WorkspaceGraph = crate::middle::workspace_graph::WorkspaceGraph;
25
26fn empty_type_graph() -> TypeGraph {
27    TypeGraph::from_items(std::iter::empty::<(DefId, Arc<Item>)>())
28}
29
30fn empty_workspace_graph() -> WorkspaceGraph {
31    WorkspaceGraph::from_items(std::iter::empty::<(DefId, Arc<Item>)>())
32}
33
34// 数据库定义 - 使用新的 #[salsa::db] 宏
35#[salsa::db]
36#[derive(Clone)]
37pub struct RootDatabase {
38    storage: salsa::Storage<Self>,
39    // 直接在数据库中存储数据
40    input_files: Arc<Vec<FileId>>,
41    nodes: Arc<FxHashMap<DefId, rir::Node>>,
42    files: Arc<FxHashMap<FileId, Arc<rir::File>>>,
43    file_ids_map: Arc<FxHashMap<Arc<PathBuf>, FileId>>,
44    file_paths: Arc<FxHashMap<FileId, Arc<PathBuf>>>,
45    file_names: Arc<FxHashMap<FileId, FastStr>>,
46    type_graph: Arc<TypeGraph>,
47    args: Arc<FxHashSet<DefId>>,
48    tags_map: Arc<FxHashMap<TagId, Arc<Tags>>>,
49    workspace_graph: Arc<WorkspaceGraph>,
50    pb_ext_indexes: Arc<FxHashMap<ExtendeeIndex, Arc<Extendee>>>,
51    pb_exts_used: Arc<FxHashSet<ExtendeeIndex>>,
52}
53
54impl Default for RootDatabase {
55    fn default() -> Self {
56        RootDatabase {
57            storage: salsa::Storage::new(None),
58            nodes: Arc::new(FxHashMap::default()),
59            files: Arc::new(FxHashMap::default()),
60            file_ids_map: Arc::new(FxHashMap::default()),
61            file_names: Arc::new(FxHashMap::default()),
62            type_graph: Arc::new(empty_type_graph()),
63            tags_map: Arc::new(FxHashMap::default()),
64            input_files: Arc::new(Vec::new()),
65            args: Arc::new(FxHashSet::default()),
66            workspace_graph: Arc::new(empty_workspace_graph()),
67            file_paths: Arc::new(FxHashMap::default()),
68            pb_ext_indexes: Arc::new(FxHashMap::default()),
69            pb_exts_used: Arc::new(FxHashSet::default()),
70        }
71    }
72}
73
74impl RootDatabase {
75    pub fn with_nodes(mut self, nodes: FxHashMap<DefId, rir::Node>) -> Self {
76        self.nodes = Arc::new(nodes);
77        self
78    }
79
80    pub fn with_workspace_graph(mut self, g: WorkspaceGraph) -> Self {
81        self.workspace_graph = Arc::new(g);
82        self
83    }
84
85    pub fn with_input_files(mut self, input_files: Vec<FileId>) -> Self {
86        self.input_files = Arc::new(input_files);
87        self
88    }
89
90    pub fn with_files(mut self, files: impl Iterator<Item = (FileId, Arc<File>)>) -> Self {
91        self.files = Arc::new(files.collect());
92        self
93    }
94
95    pub fn with_file_ids_map(mut self, file_ids_map: FxHashMap<Arc<PathBuf>, FileId>) -> Self {
96        self.file_ids_map = Arc::new(file_ids_map);
97        self
98    }
99
100    pub fn with_file_paths(mut self, file_paths: FxHashMap<FileId, Arc<PathBuf>>) -> Self {
101        self.file_paths = Arc::new(file_paths);
102        self
103    }
104
105    pub fn with_file_names(mut self, file_names: FxHashMap<FileId, FastStr>) -> Self {
106        self.file_names = Arc::new(file_names);
107        self
108    }
109
110    pub fn with_tags(
111        mut self,
112        tags_map: FxHashMap<TagId, Arc<Tags>>,
113        type_graph: TypeGraph,
114    ) -> Self {
115        self.tags_map = Arc::new(tags_map);
116        self.type_graph = Arc::new(type_graph);
117        self
118    }
119
120    pub fn with_args(mut self, args: FxHashSet<DefId>) -> Self {
121        self.args = Arc::new(args);
122        self
123    }
124
125    pub fn with_pb_ext_indexes(
126        mut self,
127        pb_ext_indexes: FxHashMap<ExtendeeIndex, Arc<Extendee>>,
128    ) -> Self {
129        self.pb_ext_indexes = Arc::new(pb_ext_indexes);
130        self
131    }
132
133    pub fn with_pb_exts_used(mut self, pb_exts_used: FxHashSet<ExtendeeIndex>) -> Self {
134        self.pb_exts_used = Arc::new(pb_exts_used);
135        self
136    }
137
138    pub fn collect_def_ids(
139        &self,
140        input: &[DefId],
141        locations: Option<&FxHashMap<DefId, DefLocation>>,
142    ) -> FxHashMap<DefId, DefLocation> {
143        use crate::middle::ty::Visitor;
144        struct PathCollector<'a> {
145            map: &'a mut FxHashMap<DefId, DefLocation>,
146            visiting: &'a mut FxHashSet<DefId>,
147            db: &'a RootDatabase,
148            locations: Option<&'a FxHashMap<DefId, DefLocation>>,
149        }
150
151        impl crate::ty::Visitor for PathCollector<'_> {
152            fn visit_path(&mut self, path: &crate::rir::Path) {
153                collect(self.db, path.did, self.map, self.visiting, self.locations)
154            }
155        }
156
157        fn collect(
158            db: &RootDatabase,
159            def_id: DefId,
160            map: &mut FxHashMap<DefId, DefLocation>,
161            visiting: &mut FxHashSet<DefId>,
162            locations: Option<&FxHashMap<DefId, DefLocation>>,
163        ) {
164            if map.contains_key(&def_id) {
165                return;
166            }
167            if let Some(locations) = locations {
168                map.insert(def_id, locations[&def_id].clone());
169            } else if !matches!(&*db.item(def_id).unwrap(), rir::Item::Mod(_)) {
170                let file_id = db.node(def_id).unwrap().file_id;
171
172                if db.input_files().contains(&file_id) {
173                    let type_graph = db.workspace_graph();
174                    let node = type_graph.node_map[&def_id];
175                    for from in type_graph
176                        .graph
177                        .neighbors_directed(node, petgraph::Direction::Incoming)
178                    {
179                        let from_def_id = type_graph.id_map[&from];
180                        let from_file_id = db.node(from_def_id).unwrap().file_id;
181
182                        if from_file_id != file_id {
183                            map.insert(def_id, DefLocation::Dynamic);
184                            break;
185                        } else {
186                            if !map.contains_key(&from_def_id) && !visiting.contains(&from_def_id) {
187                                visiting.insert(from_def_id);
188                                collect(db, from_def_id, map, visiting, locations);
189                                visiting.remove(&from_def_id);
190                            }
191                            if map
192                                .get(&from_def_id)
193                                .map(|v| match v {
194                                    DefLocation::Fixed(_, _) => false,
195                                    DefLocation::Dynamic => true,
196                                })
197                                .unwrap_or(true)
198                            {
199                                map.insert(def_id, DefLocation::Dynamic);
200                                break;
201                            }
202                        }
203                    }
204                    map.entry(def_id).or_insert_with(|| {
205                        let file = db.file(file_id).unwrap();
206                        DefLocation::Fixed(CrateId { main_file: file_id }, file.package.clone())
207                    });
208                } else {
209                    map.insert(def_id, DefLocation::Dynamic);
210                }
211            }
212
213            let node = db.node(def_id).unwrap();
214            tracing::trace!("collecting {:?}", node.expect_item().symbol_name());
215
216            node.related_nodes
217                .iter()
218                .for_each(|def_id| collect(db, *def_id, map, visiting, locations));
219
220            let item = node.expect_item();
221
222            match item {
223                rir::Item::Message(m) => m.fields.iter().for_each(|f| {
224                    PathCollector {
225                        db,
226                        map,
227                        visiting,
228                        locations,
229                    }
230                    .visit(&f.ty)
231                }),
232                rir::Item::Enum(e) => e.variants.iter().flat_map(|v| &v.fields).for_each(|ty| {
233                    PathCollector {
234                        db,
235                        map,
236                        visiting,
237                        locations,
238                    }
239                    .visit(ty)
240                }),
241                rir::Item::Service(s) => {
242                    s.extend
243                        .iter()
244                        .for_each(|p| collect(db, p.did, map, visiting, locations));
245                    s.methods
246                        .iter()
247                        .flat_map(|m| m.args.iter().map(|f| &f.ty).chain(std::iter::once(&m.ret)))
248                        .for_each(|ty| {
249                            PathCollector {
250                                db,
251                                map,
252                                visiting,
253                                locations,
254                            }
255                            .visit(ty)
256                        });
257                }
258                rir::Item::NewType(n) => PathCollector {
259                    db,
260                    map,
261                    visiting,
262                    locations,
263                }
264                .visit(&n.ty),
265                rir::Item::Const(c) => {
266                    PathCollector {
267                        db,
268                        map,
269                        visiting,
270                        locations,
271                    }
272                    .visit(&c.ty);
273                }
274                rir::Item::Mod(m) => {
275                    m.items
276                        .iter()
277                        .for_each(|i| collect(db, *i, map, visiting, locations));
278                }
279            }
280        }
281        let mut map = FxHashMap::default();
282        let mut visiting = FxHashSet::default();
283
284        input.iter().for_each(|def_id| {
285            visiting.insert(*def_id);
286            collect(self, *def_id, &mut map, &mut visiting, locations);
287            visiting.remove(def_id);
288        });
289
290        map
291    }
292}
293
294// 实现 salsa::Database trait
295#[salsa::db]
296impl salsa::Database for RootDatabase {}
297
298// 定义 RirDatabase trait
299pub trait RirDatabase: salsa::Database {
300    // 访问数据库中的数据
301    fn nodes(&self) -> &Arc<FxHashMap<DefId, rir::Node>>;
302    fn files(&self) -> &Arc<FxHashMap<FileId, Arc<rir::File>>>;
303    fn file_ids_map(&self) -> &Arc<FxHashMap<Arc<PathBuf>, FileId>>;
304    fn file_paths(&self) -> &Arc<FxHashMap<FileId, Arc<PathBuf>>>;
305    fn type_graph(&self) -> &Arc<TypeGraph>;
306    fn tags_map(&self) -> &Arc<FxHashMap<TagId, Arc<Tags>>>;
307    fn input_files(&self) -> &Arc<Vec<FileId>>;
308    fn args(&self) -> &Arc<FxHashSet<DefId>>;
309    fn workspace_graph(&self) -> &Arc<WorkspaceGraph>;
310    fn pb_ext_indexes(&self) -> &Arc<FxHashMap<ExtendeeIndex, Arc<Extendee>>>;
311    fn pb_exts_used(&self) -> &Arc<FxHashSet<ExtendeeIndex>>;
312
313    // 查询方法
314    fn node(&self, def_id: DefId) -> Option<Node>;
315
316    fn file(&self, file_id: FileId) -> Option<Arc<File>>;
317
318    fn file_id(&self, path: PathBuf) -> Option<FileId> {
319        self.file_ids_map().get(&path).cloned()
320    }
321
322    fn file_name(&self, file_id: FileId) -> Option<FastStr>;
323
324    fn item(&self, def_id: DefId) -> Option<Arc<Item>>;
325
326    fn expect_item(&self, def_id: DefId) -> Arc<Item> {
327        self.item(def_id).unwrap()
328    }
329
330    fn codegen_item_ty(&self, ty: TyKind) -> CodegenTy;
331
332    fn codegen_const_ty(&self, ty: TyKind) -> CodegenTy;
333
334    fn codegen_ty(&self, def_id: DefId) -> CodegenTy;
335
336    fn service_methods(&self, def_id: DefId) -> Arc<[Arc<rir::Method>]>;
337
338    fn is_arg(&self, def_id: DefId) -> bool;
339
340    fn pb_ext(&self, index: &ExtendeeIndex) -> Option<Arc<Extendee>>;
341
342    fn pb_ext_used(&self, index: &ExtendeeIndex) -> bool;
343}
344
345// 为 RootDatabase 实现 RirDatabase trait
346impl RirDatabase for RootDatabase {
347    fn nodes(&self) -> &Arc<FxHashMap<DefId, rir::Node>> {
348        &self.nodes
349    }
350
351    fn files(&self) -> &Arc<FxHashMap<FileId, Arc<rir::File>>> {
352        &self.files
353    }
354
355    fn file_ids_map(&self) -> &Arc<FxHashMap<Arc<PathBuf>, FileId>> {
356        &self.file_ids_map
357    }
358
359    fn file_paths(&self) -> &Arc<FxHashMap<FileId, Arc<PathBuf>>> {
360        &self.file_paths
361    }
362
363    fn type_graph(&self) -> &Arc<TypeGraph> {
364        &self.type_graph
365    }
366
367    fn tags_map(&self) -> &Arc<FxHashMap<TagId, Arc<Tags>>> {
368        &self.tags_map
369    }
370
371    fn input_files(&self) -> &Arc<Vec<FileId>> {
372        &self.input_files
373    }
374
375    fn args(&self) -> &Arc<FxHashSet<DefId>> {
376        &self.args
377    }
378
379    fn workspace_graph(&self) -> &Arc<WorkspaceGraph> {
380        &self.workspace_graph
381    }
382
383    fn pb_ext_indexes(&self) -> &Arc<FxHashMap<ExtendeeIndex, Arc<Extendee>>> {
384        &self.pb_ext_indexes
385    }
386
387    fn pb_exts_used(&self) -> &Arc<FxHashSet<ExtendeeIndex>> {
388        &self.pb_exts_used
389    }
390
391    // 使用缓存实现查询方法
392    fn node(&self, def_id: DefId) -> Option<Node> {
393        use cached_queries::{CachedQueries, get_node};
394        let salsa_id = def_id.into_salsa(self as &dyn CachedQueries);
395        get_node(self as &dyn CachedQueries, salsa_id)
396    }
397
398    fn file(&self, file_id: FileId) -> Option<Arc<File>> {
399        use cached_queries::{CachedQueries, get_file};
400        let salsa_id = file_id.into_salsa(self as &dyn CachedQueries);
401        get_file(self as &dyn CachedQueries, salsa_id)
402    }
403
404    fn file_name(&self, file_id: FileId) -> Option<FastStr> {
405        self.file_names.get(&file_id).cloned()
406    }
407
408    fn item(&self, def_id: DefId) -> Option<Arc<Item>> {
409        use cached_queries::{CachedQueries, get_item};
410        let salsa_id = def_id.into_salsa(self as &dyn CachedQueries);
411        get_item(self as &dyn CachedQueries, salsa_id)
412    }
413
414    fn service_methods(&self, def_id: DefId) -> Arc<[Arc<rir::Method>]> {
415        use cached_queries::{CachedQueries, get_service_methods};
416        let salsa_id = def_id.into_salsa(self as &dyn CachedQueries);
417        get_service_methods(self as &dyn CachedQueries, salsa_id)
418    }
419
420    fn is_arg(&self, def_id: DefId) -> bool {
421        use cached_queries::{CachedQueries, is_arg_cached};
422        let salsa_id = def_id.into_salsa(self as &dyn CachedQueries);
423        is_arg_cached(self as &dyn CachedQueries, salsa_id)
424    }
425
426    fn pb_ext(&self, index: &ExtendeeIndex) -> Option<Arc<Extendee>> {
427        self.pb_ext_indexes().get(index).cloned()
428    }
429
430    fn pb_ext_used(&self, index: &ExtendeeIndex) -> bool {
431        self.pb_exts_used().contains(index)
432    }
433
434    fn codegen_item_ty(&self, ty: TyKind) -> CodegenTy {
435        use cached_queries::{CachedQueries, codegen_item_ty_cached};
436        let salsa_ty = ty.into_salsa(self as &dyn CachedQueries);
437        codegen_item_ty_cached(self as &dyn CachedQueries, salsa_ty)
438    }
439
440    fn codegen_const_ty(&self, ty: TyKind) -> CodegenTy {
441        use cached_queries::{CachedQueries, codegen_const_ty_cached};
442        let salsa_ty = ty.into_salsa(self as &dyn CachedQueries);
443        codegen_const_ty_cached(self as &dyn CachedQueries, salsa_ty)
444    }
445
446    fn codegen_ty(&self, def_id: DefId) -> CodegenTy {
447        use cached_queries::{CachedQueries, codegen_ty_cached};
448        let salsa_id = def_id.into_salsa(self as &dyn CachedQueries);
449        codegen_ty_cached(self as &dyn CachedQueries, salsa_id)
450    }
451}
452
453impl Debug for RootDatabase {
454    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
455        write!(f, "RootDatabase {{ .. }}")
456    }
457}
458
459#[cfg(test)]
460mod tests {
461    use std::sync::Arc;
462
463    use pilota::Bytes;
464    use rustc_hash::FxHashMap;
465
466    use super::*;
467    use crate::{
468        middle::{
469            context::{CrateId, DefLocation},
470            ext::{FileExts, ItemExts},
471            rir::{self, FieldKind},
472            ty::{Ty, TyKind},
473        },
474        symbol::{DefId, FileId, Ident, Symbol},
475        tags::TagId,
476    };
477
478    fn make_item_path(parts: &[&str]) -> rir::ItemPath {
479        let symbols: Vec<Symbol> = parts
480            .iter()
481            .map(|p| Symbol::from(FastStr::new((*p).to_string())))
482            .collect();
483        let boxed: Box<[Symbol]> = symbols.into_boxed_slice();
484        rir::ItemPath::from(boxed)
485    }
486
487    fn make_message_item(name: &str, fields: Vec<Arc<rir::Field>>) -> Arc<rir::Item> {
488        Arc::new(rir::Item::Message(rir::Message {
489            name: Ident::from(FastStr::new(name.to_string())),
490            fields,
491            is_wrapper: false,
492            item_exts: ItemExts::Thrift,
493            leading_comments: FastStr::new(""),
494            trailing_comments: FastStr::new(""),
495        }))
496    }
497
498    fn make_node(file_id: FileId, item: Arc<rir::Item>, related_nodes: Vec<DefId>) -> rir::Node {
499        rir::Node {
500            file_id,
501            kind: rir::NodeKind::Item(item),
502            parent: None,
503            tags: TagId::from_u32(0),
504            related_nodes,
505        }
506    }
507
508    fn make_file(file_id: FileId, package: rir::ItemPath, items: Vec<DefId>) -> Arc<rir::File> {
509        Arc::new(rir::File {
510            package,
511            items,
512            file_id,
513            uses: vec![],
514            descriptor: Bytes::new(),
515            extensions: FileExts::Thrift,
516            comments: FastStr::new(""),
517        })
518    }
519
520    #[test]
521    fn collect_def_ids_uses_provided_locations() {
522        let root = DefId::from_u32(1);
523        let child = DefId::from_u32(2);
524        let file_id = FileId::from_u32(10);
525
526        let root_item = make_message_item("Root", Vec::new());
527        let child_item = make_message_item("Child", Vec::new());
528
529        let mut nodes = FxHashMap::default();
530        nodes.insert(root, make_node(file_id, root_item.clone(), vec![child]));
531        nodes.insert(child, make_node(file_id, child_item.clone(), vec![]));
532
533        let package = make_item_path(&["pkg", "root"]);
534        let files = vec![(
535            file_id,
536            make_file(file_id, package.clone(), vec![root, child]),
537        )];
538
539        let workspace_graph = WorkspaceGraph::from_items(
540            vec![(root, root_item.clone()), (child, child_item.clone())].into_iter(),
541        );
542
543        let db = RootDatabase::default()
544            .with_nodes(nodes)
545            .with_files(files.into_iter())
546            .with_workspace_graph(workspace_graph)
547            .with_input_files(vec![file_id]);
548
549        let mut provided = FxHashMap::default();
550        provided.insert(
551            root,
552            DefLocation::Fixed(CrateId { main_file: file_id }, package.clone()),
553        );
554        provided.insert(child, DefLocation::Dynamic);
555
556        let result = db.collect_def_ids(&[root], Some(&provided));
557
558        assert_eq!(result.get(&root), provided.get(&root));
559        assert_eq!(result.get(&child), provided.get(&child));
560    }
561
562    #[test]
563    fn collect_def_ids_infers_locations_from_workspace() {
564        let file_main = FileId::from_u32(20);
565        let file_other = FileId::from_u32(30);
566
567        let def_main = DefId::from_u32(100);
568        let def_standalone = DefId::from_u32(200);
569        let def_referrer = DefId::from_u32(300);
570
571        let main_item = make_message_item("Main", Vec::new());
572        let standalone_item = make_message_item("Standalone", Vec::new());
573
574        let dep_ty = Ty {
575            kind: TyKind::Path(rir::Path {
576                kind: rir::DefKind::Type,
577                did: def_main,
578            }),
579            tags_id: TagId::from_u32(0),
580        };
581        let dep_field = Arc::new(rir::Field {
582            did: DefId::from_u32(400),
583            name: Ident::from(FastStr::new("dep".to_string())),
584            id: 1,
585            ty: dep_ty,
586            kind: FieldKind::Required,
587            tags_id: TagId::from_u32(0),
588            default: None,
589            item_exts: ItemExts::Thrift,
590            leading_comments: FastStr::new(""),
591            trailing_comments: FastStr::new(""),
592        });
593        let referrer_item = make_message_item("Ref", vec![dep_field]);
594
595        let mut nodes = FxHashMap::default();
596        nodes.insert(def_main, make_node(file_main, main_item.clone(), vec![]));
597        nodes.insert(
598            def_standalone,
599            make_node(file_main, standalone_item.clone(), vec![]),
600        );
601        nodes.insert(
602            def_referrer,
603            make_node(file_other, referrer_item.clone(), vec![]),
604        );
605
606        let main_package = make_item_path(&["pkg", "main"]);
607        let other_package = make_item_path(&["pkg", "other"]);
608        let files = vec![
609            (
610                file_main,
611                make_file(
612                    file_main,
613                    main_package.clone(),
614                    vec![def_main, def_standalone],
615                ),
616            ),
617            (
618                file_other,
619                make_file(file_other, other_package, vec![def_referrer]),
620            ),
621        ];
622
623        let workspace_graph = WorkspaceGraph::from_items(
624            vec![
625                (def_main, main_item.clone()),
626                (def_standalone, standalone_item.clone()),
627                (def_referrer, referrer_item.clone()),
628            ]
629            .into_iter(),
630        );
631
632        let db = RootDatabase::default()
633            .with_nodes(nodes)
634            .with_files(files.into_iter())
635            .with_workspace_graph(workspace_graph)
636            .with_input_files(vec![file_main]);
637
638        let result = db.collect_def_ids(&[def_main, def_standalone], None);
639
640        assert_eq!(result.get(&def_main), Some(&DefLocation::Dynamic));
641
642        let expected_fixed = DefLocation::Fixed(
643            CrateId {
644                main_file: file_main,
645            },
646            main_package,
647        );
648        assert_eq!(result.get(&def_standalone), Some(&expected_fixed));
649    }
650}
pilota_build/db.rs

pilota_build/
db.rs
