mun_hir/semantics/
source_to_def.rs

1use crate::{
2    code_model::src::HasSource,
3    ids::{DefWithBodyId, FunctionId, ItemDefinitionId, Lookup, StructId, TypeAliasId},
4    item_scope::ItemScope,
5    DefDatabase, FileId, HirDatabase, InFile, ModuleId,
6};
7use mun_syntax::{ast, match_ast, AstNode, SyntaxNode};
8use rustc_hash::FxHashMap;
9
10pub(super) type SourceToDefCache = FxHashMap<SourceToDefContainer, SourceToDefMap>;
11
12/// An object that can be used to efficiently find definitions of source objects. It is used to
13/// find HIR elements for corresponding AST elements.
14pub(super) struct SourceToDefContext<'a, 'db> {
15    pub(super) db: &'db dyn HirDatabase,
16    pub(super) cache: &'a mut SourceToDefCache,
17}
18
19impl SourceToDefContext<'_, '_> {
20    /// Find the container for the given syntax tree node.
21    pub(super) fn find_container(
22        &mut self,
23        src: InFile<&SyntaxNode>,
24    ) -> Option<SourceToDefContainer> {
25        for container in std::iter::successors(Some(src.cloned()), move |node| {
26            node.value.parent().map(|parent| node.with_value(parent))
27        })
28        .skip(1)
29        {
30            let res: SourceToDefContainer = match_ast! {
31                match (container.value) {
32                    ast::FunctionDef(it) => {
33                        let def = self.fn_to_def(container.with_value(it))?;
34                        DefWithBodyId::from(def).into()
35                    },
36                    _ => continue,
37                }
38            };
39            return Some(res);
40        }
41
42        let def = self.file_to_def(src.file_id)?;
43        Some(def.into())
44    }
45
46    /// Find the `FunctionId` associated with the specified syntax tree node.
47    fn fn_to_def(&mut self, src: InFile<ast::FunctionDef>) -> Option<FunctionId> {
48        let container = self.find_container(src.as_ref().map(|it| it.syntax()))?;
49        let db = self.db;
50        let def_map = &*self
51            .cache
52            .entry(container)
53            .or_insert_with(|| container.source_to_def_map(db));
54        def_map.functions.get(&src).copied()
55    }
56
57    /// Finds the `ModuleId` associated with the specified `file`
58    fn file_to_def(&self, file_id: FileId) -> Option<ModuleId> {
59        let source_root_id = self.db.file_source_root(file_id);
60        let packages = self.db.packages();
61        let result = packages
62            .iter()
63            .filter(|package_id| packages[*package_id].source_root == source_root_id)
64            .find_map(|package_id| {
65                let module_tree = self.db.module_tree(package_id);
66                let module_id = module_tree.module_for_file(file_id)?;
67                Some(ModuleId {
68                    package: package_id,
69                    local_id: module_id,
70                })
71            });
72        result
73    }
74}
75
76/// A container that holds other items.
77#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
78pub(crate) enum SourceToDefContainer {
79    DefWithBodyId(DefWithBodyId),
80    ModuleId(ModuleId),
81}
82
83impl From<DefWithBodyId> for SourceToDefContainer {
84    fn from(id: DefWithBodyId) -> Self {
85        SourceToDefContainer::DefWithBodyId(id)
86    }
87}
88
89impl From<ModuleId> for SourceToDefContainer {
90    fn from(id: ModuleId) -> Self {
91        SourceToDefContainer::ModuleId(id)
92    }
93}
94
95impl SourceToDefContainer {
96    fn source_to_def_map(self, db: &dyn HirDatabase) -> SourceToDefMap {
97        match self {
98            SourceToDefContainer::DefWithBodyId(id) => id.source_to_def_map(db),
99            SourceToDefContainer::ModuleId(id) => id.source_to_def_map(db),
100        }
101    }
102}
103
104/// A trait to construct a `SourceToDefMap` from a definition like a module.
105trait SourceToDef {
106    /// Returns all definitions in `self`.
107    fn source_to_def_map(&self, db: &dyn HirDatabase) -> SourceToDefMap;
108}
109
110impl SourceToDef for DefWithBodyId {
111    fn source_to_def_map(&self, _db: &dyn HirDatabase) -> SourceToDefMap {
112        // TODO: bodies dont yet contain items themselves
113        SourceToDefMap::default()
114    }
115}
116
117impl SourceToDef for ModuleId {
118    fn source_to_def_map(&self, db: &dyn HirDatabase) -> SourceToDefMap {
119        let package_defs = db.package_defs(self.package);
120        let module_scope = &package_defs[self.local_id];
121        module_scope.source_to_def_map(db)
122    }
123}
124
125impl SourceToDef for ItemScope {
126    fn source_to_def_map(&self, db: &dyn HirDatabase) -> SourceToDefMap {
127        let mut result = SourceToDefMap::default();
128        self.declarations()
129            .for_each(|item| add_module_def(db.upcast(), &mut result, item));
130        return result;
131
132        fn add_module_def(db: &dyn DefDatabase, map: &mut SourceToDefMap, item: ItemDefinitionId) {
133            match item {
134                ItemDefinitionId::FunctionId(id) => {
135                    let src = id.lookup(db).source(db);
136                    map.functions.insert(src, id);
137                }
138                ItemDefinitionId::StructId(id) => {
139                    let src = id.lookup(db).source(db);
140                    map.structs.insert(src, id);
141                }
142                ItemDefinitionId::TypeAliasId(id) => {
143                    let src = id.lookup(db).source(db);
144                    map.type_aliases.insert(src, id);
145                }
146                _ => {}
147            }
148        }
149    }
150}
151
152/// Holds conversion from source location to definitions in the HIR.
153#[derive(Default)]
154pub(crate) struct SourceToDefMap {
155    functions: FxHashMap<InFile<ast::FunctionDef>, FunctionId>,
156    structs: FxHashMap<InFile<ast::StructDef>, StructId>,
157    type_aliases: FxHashMap<InFile<ast::TypeAliasDef>, TypeAliasId>,
158}