hir/
term_search.rs

1//! Term search
2
3use hir_def::type_ref::Mutability;
4use hir_ty::db::HirDatabase;
5use itertools::Itertools;
6use rustc_hash::{FxHashMap, FxHashSet};
7
8use crate::{ModuleDef, ScopeDef, Semantics, SemanticsScope, Type};
9
10mod expr;
11pub use expr::Expr;
12
13mod tactics;
14
15/// Key for lookup table to query new types reached.
16#[derive(Debug, Hash, PartialEq, Eq)]
17enum NewTypesKey {
18    ImplMethod,
19    StructProjection,
20}
21
22/// Helper enum to squash big number of alternative trees into `Many` variant as there is too many
23/// to take into account.
24#[derive(Debug)]
25enum AlternativeExprs<'db> {
26    /// There are few trees, so we keep track of them all
27    Few(FxHashSet<Expr<'db>>),
28    /// There are too many trees to keep track of
29    Many,
30}
31
32impl<'db> AlternativeExprs<'db> {
33    /// Construct alternative trees
34    ///
35    /// # Arguments
36    /// `threshold` - threshold value for many trees (more than that is many)
37    /// `exprs` - expressions iterator
38    fn new(threshold: usize, exprs: impl Iterator<Item = Expr<'db>>) -> AlternativeExprs<'db> {
39        let mut it = AlternativeExprs::Few(Default::default());
40        it.extend_with_threshold(threshold, exprs);
41        it
42    }
43
44    /// Get type trees stored in alternative trees (or `Expr::Many` in case of many)
45    ///
46    /// # Arguments
47    /// `ty` - Type of expressions queried (this is used to give type to `Expr::Many`)
48    fn exprs(&self, ty: &Type<'db>) -> Vec<Expr<'db>> {
49        match self {
50            AlternativeExprs::Few(exprs) => exprs.iter().cloned().collect(),
51            AlternativeExprs::Many => vec![Expr::Many(ty.clone())],
52        }
53    }
54
55    /// Extend alternative expressions
56    ///
57    /// # Arguments
58    /// `threshold` - threshold value for many trees (more than that is many)
59    /// `exprs` - expressions iterator
60    fn extend_with_threshold(&mut self, threshold: usize, exprs: impl Iterator<Item = Expr<'db>>) {
61        match self {
62            AlternativeExprs::Few(tts) => {
63                for it in exprs {
64                    if tts.len() > threshold {
65                        *self = AlternativeExprs::Many;
66                        break;
67                    }
68
69                    tts.insert(it);
70                }
71            }
72            AlternativeExprs::Many => (),
73        }
74    }
75
76    fn is_many(&self) -> bool {
77        matches!(self, AlternativeExprs::Many)
78    }
79}
80
81/// # Lookup table for term search
82///
83/// Lookup table keeps all the state during term search.
84/// This means it knows what types and how are reachable.
85///
86/// The secondary functionality for lookup table is to keep track of new types reached since last
87/// iteration as well as keeping track of which `ScopeDef` items have been used.
88/// Both of them are to speed up the term search by leaving out types / ScopeDefs that likely do
89/// not produce any new results.
90#[derive(Default, Debug)]
91struct LookupTable<'db> {
92    /// All the `Expr`s in "value" produce the type of "key"
93    data: FxHashMap<Type<'db>, AlternativeExprs<'db>>,
94    /// New types reached since last query by the `NewTypesKey`
95    new_types: FxHashMap<NewTypesKey, Vec<Type<'db>>>,
96    /// Types queried but not present
97    types_wishlist: FxHashSet<Type<'db>>,
98    /// Threshold to squash trees to `Many`
99    many_threshold: usize,
100}
101
102impl<'db> LookupTable<'db> {
103    /// Initialize lookup table
104    fn new(many_threshold: usize, goal: Type<'db>) -> Self {
105        let mut res = Self { many_threshold, ..Default::default() };
106        res.new_types.insert(NewTypesKey::ImplMethod, Vec::new());
107        res.new_types.insert(NewTypesKey::StructProjection, Vec::new());
108        res.types_wishlist.insert(goal);
109        res
110    }
111
112    /// Find all `Expr`s that unify with the `ty`
113    fn find(&mut self, db: &'db dyn HirDatabase, ty: &Type<'db>) -> Option<Vec<Expr<'db>>> {
114        let res = self
115            .data
116            .iter()
117            .find(|(t, _)| t.could_unify_with_deeply(db, ty))
118            .map(|(t, tts)| tts.exprs(t));
119
120        if res.is_none() {
121            self.types_wishlist.insert(ty.clone());
122        }
123
124        // Collapse suggestions if there are many
125        if let Some(res) = &res
126            && res.len() > self.many_threshold
127        {
128            return Some(vec![Expr::Many(ty.clone())]);
129        }
130
131        res
132    }
133
134    /// Same as find but automatically creates shared reference of types in the lookup
135    ///
136    /// For example if we have type `i32` in data and we query for `&i32` it map all the type
137    /// trees we have for `i32` with `Expr::Reference` and returns them.
138    fn find_autoref(&mut self, db: &'db dyn HirDatabase, ty: &Type<'db>) -> Option<Vec<Expr<'db>>> {
139        let res = self
140            .data
141            .iter()
142            .find(|(t, _)| t.could_unify_with_deeply(db, ty))
143            .map(|(t, it)| it.exprs(t))
144            .or_else(|| {
145                self.data
146                    .iter()
147                    .find(|(t, _)| {
148                        t.add_reference(Mutability::Shared).could_unify_with_deeply(db, ty)
149                    })
150                    .map(|(t, it)| {
151                        it.exprs(t)
152                            .into_iter()
153                            .map(|expr| Expr::Reference(Box::new(expr)))
154                            .collect()
155                    })
156            });
157
158        if res.is_none() {
159            self.types_wishlist.insert(ty.clone());
160        }
161
162        // Collapse suggestions if there are many
163        if let Some(res) = &res
164            && res.len() > self.many_threshold
165        {
166            return Some(vec![Expr::Many(ty.clone())]);
167        }
168
169        res
170    }
171
172    /// Insert new type trees for type
173    ///
174    /// Note that the types have to be the same, unification is not enough as unification is not
175    /// transitive. For example Vec<i32> and FxHashSet<i32> both unify with Iterator<Item = i32>,
176    /// but they clearly do not unify themselves.
177    fn insert(&mut self, ty: Type<'db>, exprs: impl Iterator<Item = Expr<'db>>) {
178        match self.data.get_mut(&ty) {
179            Some(it) => {
180                it.extend_with_threshold(self.many_threshold, exprs);
181                if it.is_many() {
182                    self.types_wishlist.remove(&ty);
183                }
184            }
185            None => {
186                self.data.insert(ty.clone(), AlternativeExprs::new(self.many_threshold, exprs));
187                for it in self.new_types.values_mut() {
188                    it.push(ty.clone());
189                }
190            }
191        }
192    }
193
194    /// Iterate all the reachable types
195    fn iter_types(&self) -> impl Iterator<Item = Type<'db>> + '_ {
196        self.data.keys().cloned()
197    }
198
199    /// Query new types reached since last query by key
200    ///
201    /// Create new key if you wish to query it to avoid conflicting with existing queries.
202    fn new_types(&mut self, key: NewTypesKey) -> Vec<Type<'db>> {
203        match self.new_types.get_mut(&key) {
204            Some(it) => std::mem::take(it),
205            None => Vec::new(),
206        }
207    }
208
209    /// Types queried but not found
210    fn types_wishlist(&mut self) -> &FxHashSet<Type<'db>> {
211        &self.types_wishlist
212    }
213}
214
215/// Context for the `term_search` function
216#[derive(Debug)]
217pub struct TermSearchCtx<'db, DB: HirDatabase> {
218    /// Semantics for the program
219    pub sema: &'db Semantics<'db, DB>,
220    /// Semantic scope, captures context for the term search
221    pub scope: &'db SemanticsScope<'db>,
222    /// Target / expected output type
223    pub goal: Type<'db>,
224    /// Configuration for term search
225    pub config: TermSearchConfig,
226}
227
228/// Configuration options for the term search
229#[derive(Debug, Clone, Copy)]
230pub struct TermSearchConfig {
231    /// Enable borrow checking, this guarantees the outputs of the `term_search` to borrow-check
232    pub enable_borrowcheck: bool,
233    /// Indicate when to squash multiple trees to `Many` as there are too many to keep track
234    pub many_alternatives_threshold: usize,
235    /// Fuel for term search in "units of work"
236    pub fuel: u64,
237}
238
239impl Default for TermSearchConfig {
240    fn default() -> Self {
241        Self { enable_borrowcheck: true, many_alternatives_threshold: 1, fuel: 1200 }
242    }
243}
244
245/// # Term search
246///
247/// Search for terms (expressions) that unify with the `goal` type.
248///
249/// # Arguments
250/// * `ctx` - Context for term search
251///
252/// Internally this function uses Breadth First Search to find path to `goal` type.
253/// The general idea is following:
254/// 1. Populate lookup (frontier for BFS) from values (local variables, statics, constants, etc)
255///    as well as from well knows values (such as `true/false` and `()`)
256/// 2. Iteratively expand the frontier (or contents of the lookup) by trying different type
257///    transformation tactics. For example functions take as from set of types (arguments) to some
258///    type (return type). Other transformations include methods on type, type constructors and
259///    projections to struct fields (field access).
260/// 3. If we run out of fuel (term search takes too long) we stop iterating.
261/// 4. Return all the paths (type trees) that take us to the `goal` type.
262///
263/// Note that there are usually more ways we can get to the `goal` type but some are discarded to
264/// reduce the memory consumption. It is also unlikely anyone is willing ti browse through
265/// thousands of possible responses so we currently take first 10 from every tactic.
266pub fn term_search<'db, DB: HirDatabase>(ctx: &'db TermSearchCtx<'db, DB>) -> Vec<Expr<'db>> {
267    let module = ctx.scope.module();
268    let mut defs = FxHashSet::default();
269    defs.insert(ScopeDef::ModuleDef(ModuleDef::Module(module)));
270
271    ctx.scope.process_all_names(&mut |_, def| {
272        defs.insert(def);
273    });
274
275    let mut lookup = LookupTable::new(ctx.config.many_alternatives_threshold, ctx.goal.clone());
276    let fuel = std::cell::Cell::new(ctx.config.fuel);
277
278    let should_continue = &|| {
279        let remaining = fuel.get();
280        fuel.set(remaining.saturating_sub(1));
281        if remaining == 0 {
282            tracing::debug!("fuel exhausted");
283        }
284        remaining > 0
285    };
286
287    // Try trivial tactic first, also populates lookup table
288    let mut solutions: Vec<Expr<'db>> = tactics::trivial(ctx, &defs, &mut lookup).collect();
289    // Use well known types tactic before iterations as it does not depend on other tactics
290    solutions.extend(tactics::famous_types(ctx, &defs, &mut lookup));
291    solutions.extend(tactics::assoc_const(ctx, &defs, &mut lookup));
292
293    while should_continue() {
294        solutions.extend(tactics::data_constructor(ctx, &defs, &mut lookup, should_continue));
295        solutions.extend(tactics::free_function(ctx, &defs, &mut lookup, should_continue));
296        solutions.extend(tactics::impl_method(ctx, &defs, &mut lookup, should_continue));
297        solutions.extend(tactics::struct_projection(ctx, &defs, &mut lookup, should_continue));
298        solutions.extend(tactics::impl_static_method(ctx, &defs, &mut lookup, should_continue));
299        solutions.extend(tactics::make_tuple(ctx, &defs, &mut lookup, should_continue));
300    }
301
302    solutions.into_iter().filter(|it| !it.is_many()).unique().collect()
303}