Auto merge of #17203 - kilpkonn:collapse_terms, r=Veykril

Fix OOM caused by term search

The issue came from multi Cartesian product for exprs with many (25+) arguments, each having multiple options.
The solution is two fold:
### Avoid blowing up in Cartesian product
**Before the logic was:**
    1. Find expressions for each argument/param - there may be many
    2. Take the Cartesian product (which blows up in some cases)
    4. If there are more than 2 options throw them away by squashing them to `Many`
**Now the logic is:**
    1. Find expressions for each argument/param and squash them to `Many` if there are more than 2 as otherwise we are guaranteed to also have more than 2 after taking the product which means squashing them anyway.
    2. Take the Cartesian product on iterator
    3. Start consuming it one by one
    4. If there are more than 2 options throw them away by squashing them to `Many`  (same as before)

This is also why I had to update some tests as the expressions get squashed to many more eagerly.

### Use fuel to avoid long search times and high memory usage
Now all the tactics use `should_continue:  Fn() -> bool` to chech if they should keep iterating _(Similarly to chalk)_.
This reduces the search times by a magnitude, for example from ~139ms/hole to ~14ms/hole for `ripgrep` crate.
There are slightly less expressions found, but I think speed gain worth it for usability.
Also note that syntactic hits decreases more because of squashing so you simple need to run search multiple times to get full terms.
Also the worst case time (For example `nalgebra` crate cus it has tons of generics) has search times mostly under 200ms.

Benchmarks on `ripgrep` crate
Before:
```
Tail Expr syntactic hits: 291/1692 (17%)
Tail Exprs found: 1253/1692 (74%)
Term search avg time: 139ms
````
After:
```
Tail Expr syntactic hits: 239/1692 (14%)
Tail Exprs found: 1226/1692 (72%)
Term search avg time: 14ms
```

Author: bors

src/tools/rust-analyzer/crates/hir/src/term_search.rs

@@ -127,6 +127,13 @@ impl LookupTable {
             self.types_wishlist.insert(ty.clone());
         }
 
+        // Collapse suggestions if there are many
+        if let Some(res) = &res {
+            if res.len() > self.many_threshold {
+                return Some(vec![Expr::Many(ty.clone())]);
+            }
+        }
+
         res
     }
 
@@ -158,6 +165,13 @@ impl LookupTable {
             self.types_wishlist.insert(ty.clone());
         }
 
+        // Collapse suggestions if there are many
+        if let Some(res) = &res {
+            if res.len() > self.many_threshold {
+                return Some(vec![Expr::Many(ty.clone())]);
+            }
+        }
+
         res
     }
 
@@ -255,13 +269,13 @@ pub struct TermSearchConfig {
     pub enable_borrowcheck: bool,
     /// Indicate when to squash multiple trees to `Many` as there are too many to keep track
     pub many_alternatives_threshold: usize,
-    /// Depth of the search eg. number of cycles to run
-    pub depth: usize,
+    /// Fuel for term search in "units of work"
+    pub fuel: u64,
 }
 
 impl Default for TermSearchConfig {
     fn default() -> Self {
-        Self { enable_borrowcheck: true, many_alternatives_threshold: 1, depth: 6 }
+        Self { enable_borrowcheck: true, many_alternatives_threshold: 1, fuel: 400 }
     }
 }
 
@@ -280,8 +294,7 @@ impl Default for TermSearchConfig {
 ///    transformation tactics. For example functions take as from set of types (arguments) to some
 ///    type (return type). Other transformations include methods on type, type constructors and
 ///    projections to struct fields (field access).
-/// 3. Once we manage to find path to type we are interested in we continue for single round to see
-///    if we can find more paths that take us to the `goal` type.
+/// 3. If we run out of fuel (term search takes too long) we stop iterating.
 /// 4. Return all the paths (type trees) that take us to the `goal` type.
 ///
 /// Note that there are usually more ways we can get to the `goal` type but some are discarded to
@@ -297,21 +310,31 @@ pub fn term_search<DB: HirDatabase>(ctx: &TermSearchCtx<'_, DB>) -> Vec<Expr> {
     });
 
     let mut lookup = LookupTable::new(ctx.config.many_alternatives_threshold, ctx.goal.clone());
+    let fuel = std::cell::Cell::new(ctx.config.fuel);
+
+    let should_continue = &|| {
+        let remaining = fuel.get();
+        fuel.set(remaining.saturating_sub(1));
+        if remaining == 0 {
+            tracing::debug!("fuel exhausted");
+        }
+        remaining > 0
+    };
 
     // Try trivial tactic first, also populates lookup table
     let mut solutions: Vec<Expr> = tactics::trivial(ctx, &defs, &mut lookup).collect();
     // Use well known types tactic before iterations as it does not depend on other tactics
     solutions.extend(tactics::famous_types(ctx, &defs, &mut lookup));
 
-    for _ in 0..ctx.config.depth {
+    while should_continue() {
         lookup.new_round();
 
-        solutions.extend(tactics::type_constructor(ctx, &defs, &mut lookup));
-        solutions.extend(tactics::free_function(ctx, &defs, &mut lookup));
-        solutions.extend(tactics::impl_method(ctx, &defs, &mut lookup));
-        solutions.extend(tactics::struct_projection(ctx, &defs, &mut lookup));
-        solutions.extend(tactics::impl_static_method(ctx, &defs, &mut lookup));
-        solutions.extend(tactics::make_tuple(ctx, &defs, &mut lookup));
+        solutions.extend(tactics::type_constructor(ctx, &defs, &mut lookup, should_continue));
+        solutions.extend(tactics::free_function(ctx, &defs, &mut lookup, should_continue));
+        solutions.extend(tactics::impl_method(ctx, &defs, &mut lookup, should_continue));
+        solutions.extend(tactics::struct_projection(ctx, &defs, &mut lookup, should_continue));
+        solutions.extend(tactics::impl_static_method(ctx, &defs, &mut lookup, should_continue));
+        solutions.extend(tactics::make_tuple(ctx, &defs, &mut lookup, should_continue));
 
         // Discard not interesting `ScopeDef`s for speedup
         for def in lookup.exhausted_scopedefs() {

src/tools/rust-analyzer/crates/hir/src/term_search/expr.rs

@@ -211,13 +211,13 @@ impl Expr {
                 }
             }
             Expr::Method { func, target, params, .. } => {
-                if target.contains_many_in_illegal_pos() {
+                if self.contains_many_in_illegal_pos(db) {
                     return Ok(many_formatter(&target.ty(db)));
                 }
 
                 let func_name = func.name(db).display(db.upcast()).to_string();
                 let self_param = func.self_param(db).unwrap();
-                let target = target.gen_source_code(
+                let target_str = target.gen_source_code(
                     sema_scope,
                     many_formatter,
                     prefer_no_std,
@@ -236,17 +236,20 @@ impl Expr {
                     Some(trait_) => {
                         let trait_name = mod_item_path_str(sema_scope, &ModuleDef::Trait(trait_))?;
                         let target = match self_param.access(db) {
-                            crate::Access::Shared => format!("&{target}"),
-                            crate::Access::Exclusive => format!("&mut {target}"),
-                            crate::Access::Owned => target,
+                            crate::Access::Shared if !target.is_many() => format!("&{target_str}"),
+                            crate::Access::Exclusive if !target.is_many() => {
+                                format!("&mut {target_str}")
+                            }
+                            crate::Access::Owned => target_str,
+                            _ => many_formatter(&target.ty(db)),
                         };
                         let res = match args.is_empty() {
                             true => format!("{trait_name}::{func_name}({target})",),
                             false => format!("{trait_name}::{func_name}({target}, {args})",),
                         };
                         Ok(res)
                     }
-                    None => Ok(format!("{target}.{func_name}({args})")),
+                    None => Ok(format!("{target_str}.{func_name}({args})")),
                 }
             }
             Expr::Variant { variant, generics, params } => {
@@ -381,7 +384,7 @@ impl Expr {
                 Ok(res)
             }
             Expr::Field { expr, field } => {
-                if expr.contains_many_in_illegal_pos() {
+                if expr.contains_many_in_illegal_pos(db) {
                     return Ok(many_formatter(&expr.ty(db)));
                 }
 
@@ -395,7 +398,7 @@ impl Expr {
                 Ok(format!("{strukt}.{field}"))
             }
             Expr::Reference(expr) => {
-                if expr.contains_many_in_illegal_pos() {
+                if expr.contains_many_in_illegal_pos(db) {
                     return Ok(many_formatter(&expr.ty(db)));
                 }
 
@@ -466,10 +469,15 @@ impl Expr {
     /// macro!().bar()
     /// &macro!()
     /// ```
-    fn contains_many_in_illegal_pos(&self) -> bool {
+    fn contains_many_in_illegal_pos(&self, db: &dyn HirDatabase) -> bool {
         match self {
-            Expr::Method { target, .. } => target.contains_many_in_illegal_pos(),
-            Expr::Field { expr, .. } => expr.contains_many_in_illegal_pos(),
+            Expr::Method { target, func, .. } => {
+                match func.as_assoc_item(db).and_then(|it| it.container_or_implemented_trait(db)) {
+                    Some(_) => false,
+                    None => target.is_many(),
+                }
+            }
+            Expr::Field { expr, .. } => expr.contains_many_in_illegal_pos(db),
             Expr::Reference(target) => target.is_many(),
             Expr::Many(_) => true,
             _ => false,

src/tools/rust-analyzer/crates/hir/src/term_search/tactics.rs

@@ -4,6 +4,7 @@
 //! * `ctx` - Context for the term search
 //! * `defs` - Set of items in scope at term search target location
 //! * `lookup` - Lookup table for types
+//! * `should_continue` - Function that indicates when to stop iterating
 //! And they return iterator that yields type trees that unify with the `goal` type.
 
 use std::iter;
@@ -97,16 +98,19 @@ pub(super) fn trivial<'a, DB: HirDatabase>(
 /// * `ctx` - Context for the term search
 /// * `defs` - Set of items in scope at term search target location
 /// * `lookup` - Lookup table for types
+/// * `should_continue` - Function that indicates when to stop iterating
 pub(super) fn type_constructor<'a, DB: HirDatabase>(
     ctx: &'a TermSearchCtx<'a, DB>,
     defs: &'a FxHashSet<ScopeDef>,
     lookup: &'a mut LookupTable,
+    should_continue: &'a dyn std::ops::Fn() -> bool,
 ) -> impl Iterator<Item = Expr> + 'a {
     let db = ctx.sema.db;
     let module = ctx.scope.module();
     fn variant_helper(
         db: &dyn HirDatabase,
         lookup: &mut LookupTable,
+        should_continue: &dyn std::ops::Fn() -> bool,
         parent_enum: Enum,
         variant: Variant,
         config: &TermSearchConfig,
@@ -152,6 +156,7 @@ pub(super) fn type_constructor<'a, DB: HirDatabase>(
             .chain((non_default_type_params_len == 0).then_some(Vec::new()));
 
         generic_params
+            .filter(|_| should_continue())
             .filter_map(move |generics| {
                 // Insert default type params
                 let mut g = generics.into_iter();
@@ -194,8 +199,14 @@ pub(super) fn type_constructor<'a, DB: HirDatabase>(
     defs.iter()
         .filter_map(move |def| match def {
             ScopeDef::ModuleDef(ModuleDef::Variant(it)) => {
-                let variant_exprs =
-                    variant_helper(db, lookup, it.parent_enum(db), *it, &ctx.config);
+                let variant_exprs = variant_helper(
+                    db,
+                    lookup,
+                    should_continue,
+                    it.parent_enum(db),
+                    *it,
+                    &ctx.config,
+                );
                 if variant_exprs.is_empty() {
                     return None;
                 }
@@ -213,7 +224,9 @@ pub(super) fn type_constructor<'a, DB: HirDatabase>(
                 let exprs: Vec<(Type, Vec<Expr>)> = enum_
                     .variants(db)
                     .into_iter()
-                    .flat_map(|it| variant_helper(db, lookup, *enum_, it, &ctx.config))
+                    .flat_map(|it| {
+                        variant_helper(db, lookup, should_continue, *enum_, it, &ctx.config)
+                    })
                     .collect();
 
                 if exprs.is_empty() {
@@ -271,6 +284,7 @@ pub(super) fn type_constructor<'a, DB: HirDatabase>(
                     .chain((non_default_type_params_len == 0).then_some(Vec::new()));
 
                 let exprs = generic_params
+                    .filter(|_| should_continue())
                     .filter_map(|generics| {
                         // Insert default type params
                         let mut g = generics.into_iter();
@@ -345,10 +359,12 @@ pub(super) fn type_constructor<'a, DB: HirDatabase>(
 /// * `ctx` - Context for the term search
 /// * `defs` - Set of items in scope at term search target location
 /// * `lookup` - Lookup table for types
+/// * `should_continue` - Function that indicates when to stop iterating
 pub(super) fn free_function<'a, DB: HirDatabase>(
     ctx: &'a TermSearchCtx<'a, DB>,
     defs: &'a FxHashSet<ScopeDef>,
     lookup: &'a mut LookupTable,
+    should_continue: &'a dyn std::ops::Fn() -> bool,
 ) -> impl Iterator<Item = Expr> + 'a {
     let db = ctx.sema.db;
     let module = ctx.scope.module();
@@ -390,6 +406,7 @@ pub(super) fn free_function<'a, DB: HirDatabase>(
                     .permutations(non_default_type_params_len);
 
                 let exprs: Vec<_> = generic_params
+                    .filter(|_| should_continue())
                     .filter_map(|generics| {
                         // Insert default type params
                         let mut g = generics.into_iter();
@@ -474,10 +491,12 @@ pub(super) fn free_function<'a, DB: HirDatabase>(
 /// * `ctx` - Context for the term search
 /// * `defs` - Set of items in scope at term search target location
 /// * `lookup` - Lookup table for types
+/// * `should_continue` - Function that indicates when to stop iterating
 pub(super) fn impl_method<'a, DB: HirDatabase>(
     ctx: &'a TermSearchCtx<'a, DB>,
     _defs: &'a FxHashSet<ScopeDef>,
     lookup: &'a mut LookupTable,
+    should_continue: &'a dyn std::ops::Fn() -> bool,
 ) -> impl Iterator<Item = Expr> + 'a {
     let db = ctx.sema.db;
     let module = ctx.scope.module();
@@ -554,6 +573,7 @@ pub(super) fn impl_method<'a, DB: HirDatabase>(
                 .permutations(non_default_fn_type_params_len);
 
             let exprs: Vec<_> = generic_params
+                .filter(|_| should_continue())
                 .filter_map(|generics| {
                     // Insert default type params
                     let mut g = generics.into_iter();
@@ -645,17 +665,20 @@ pub(super) fn impl_method<'a, DB: HirDatabase>(
 /// * `ctx` - Context for the term search
 /// * `defs` - Set of items in scope at term search target location
 /// * `lookup` - Lookup table for types
+/// * `should_continue` - Function that indicates when to stop iterating
 pub(super) fn struct_projection<'a, DB: HirDatabase>(
     ctx: &'a TermSearchCtx<'a, DB>,
     _defs: &'a FxHashSet<ScopeDef>,
     lookup: &'a mut LookupTable,
+    should_continue: &'a dyn std::ops::Fn() -> bool,
 ) -> impl Iterator<Item = Expr> + 'a {
     let db = ctx.sema.db;
     let module = ctx.scope.module();
     lookup
         .new_types(NewTypesKey::StructProjection)
         .into_iter()
         .map(|ty| (ty.clone(), lookup.find(db, &ty).expect("Expr not in lookup")))
+        .filter(|_| should_continue())
         .flat_map(move |(ty, targets)| {
             ty.fields(db).into_iter().filter_map(move |(field, filed_ty)| {
                 if !field.is_visible_from(db, module) {
@@ -716,10 +739,12 @@ pub(super) fn famous_types<'a, DB: HirDatabase>(
 /// * `ctx` - Context for the term search
 /// * `defs` - Set of items in scope at term search target location
 /// * `lookup` - Lookup table for types
+/// * `should_continue` - Function that indicates when to stop iterating
 pub(super) fn impl_static_method<'a, DB: HirDatabase>(
     ctx: &'a TermSearchCtx<'a, DB>,
     _defs: &'a FxHashSet<ScopeDef>,
     lookup: &'a mut LookupTable,
+    should_continue: &'a dyn std::ops::Fn() -> bool,
 ) -> impl Iterator<Item = Expr> + 'a {
     let db = ctx.sema.db;
     let module = ctx.scope.module();
@@ -728,6 +753,7 @@ pub(super) fn impl_static_method<'a, DB: HirDatabase>(
         .clone()
         .into_iter()
         .chain(iter::once(ctx.goal.clone()))
+        .filter(|_| should_continue())
         .flat_map(|ty| {
             Impl::all_for_type(db, ty.clone()).into_iter().map(move |imp| (ty.clone(), imp))
         })
@@ -801,6 +827,7 @@ pub(super) fn impl_static_method<'a, DB: HirDatabase>(
                 .permutations(non_default_fn_type_params_len);
 
             let exprs: Vec<_> = generic_params
+                .filter(|_| should_continue())
                 .filter_map(|generics| {
                     // Insert default type params
                     let mut g = generics.into_iter();
@@ -884,10 +911,12 @@ pub(super) fn impl_static_method<'a, DB: HirDatabase>(
 /// * `ctx` - Context for the term search
 /// * `defs` - Set of items in scope at term search target location
 /// * `lookup` - Lookup table for types
+/// * `should_continue` - Function that indicates when to stop iterating
 pub(super) fn make_tuple<'a, DB: HirDatabase>(
     ctx: &'a TermSearchCtx<'a, DB>,
     _defs: &'a FxHashSet<ScopeDef>,
     lookup: &'a mut LookupTable,
+    should_continue: &'a dyn std::ops::Fn() -> bool,
 ) -> impl Iterator<Item = Expr> + 'a {
     let db = ctx.sema.db;
     let module = ctx.scope.module();
@@ -896,6 +925,7 @@ pub(super) fn make_tuple<'a, DB: HirDatabase>(
         .types_wishlist()
         .clone()
         .into_iter()
+        .filter(|_| should_continue())
         .filter(|ty| ty.is_tuple())
         .filter_map(move |ty| {
             // Double check to not contain unknown
@@ -915,6 +945,7 @@ pub(super) fn make_tuple<'a, DB: HirDatabase>(
             let exprs: Vec<Expr> = param_exprs
                 .into_iter()
                 .multi_cartesian_product()
+                .filter(|_| should_continue())
                 .map(|params| {
                     let tys: Vec<Type> = params.iter().map(|it| it.ty(db)).collect();
                     let tuple_ty = Type::new_tuple(module.krate().into(), &tys);

src/tools/rust-analyzer/crates/ide-assists/src/assist_config.rs

@@ -16,4 +16,5 @@ pub struct AssistConfig {
     pub prefer_no_std: bool,
     pub prefer_prelude: bool,
     pub assist_emit_must_use: bool,
+    pub term_search_fuel: u64,
 }

src/tools/rust-analyzer/crates/ide-assists/src/handlers/term_search.rs

@@ -1,5 +1,5 @@
 //! Term search assist
-use hir::term_search::TermSearchCtx;
+use hir::term_search::{TermSearchConfig, TermSearchCtx};
 use ide_db::{
     assists::{AssistId, AssistKind, GroupLabel},
     famous_defs::FamousDefs,
@@ -34,7 +34,7 @@ pub(crate) fn term_search(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<
         sema: &ctx.sema,
         scope: &scope,
         goal: target_ty,
-        config: Default::default(),
+        config: TermSearchConfig { fuel: ctx.config.term_search_fuel, ..Default::default() },
     };
     let paths = hir::term_search::term_search(&term_search_ctx);