Use unit of work as fuel instead of time

kilpkonn · kilpkonn · commit f9f1d3140f5d · 2024-05-07T22:13:27.000+03:00
diff --git a/src/tools/rust-analyzer/crates/hir/src/term_search.rs b/src/tools/rust-analyzer/crates/hir/src/term_search.rs
@@ -1,7 +1,5 @@
 //! Term search
 
-use std::time::{Duration, Instant};
-
 use hir_def::type_ref::Mutability;
 use hir_ty::db::HirDatabase;
 use itertools::Itertools;
@@ -271,20 +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 i.e. number of cycles to run
-    pub depth: usize,
-    /// Time fuel for term search
-    pub fuel: Option<Duration>,
+    /// Fuel for term search
+    pub fuel: u64,
 }
 
 impl Default for TermSearchConfig {
     fn default() -> Self {
-        Self {
-            enable_borrowcheck: true,
-            many_alternatives_threshold: 1,
-            depth: 5,
-            fuel: Some(Duration::from_millis(100)),
-        }
+        Self { enable_borrowcheck: true, many_alternatives_threshold: 1, fuel: 400 }
     }
 }
 
@@ -319,40 +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 start = Instant::now();
+    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();
 
-        if ctx.config.fuel.is_some_and(|timeout| start.elapsed() > timeout) {
-            break;
-        }
-        solutions.extend(tactics::type_constructor(ctx, &defs, &mut lookup));
-        if ctx.config.fuel.is_some_and(|timeout| start.elapsed() > timeout) {
-            break;
-        }
-        solutions.extend(tactics::free_function(ctx, &defs, &mut lookup));
-        if ctx.config.fuel.is_some_and(|timeout| start.elapsed() > timeout) {
-            break;
-        }
-        solutions.extend(tactics::impl_method(ctx, &defs, &mut lookup));
-        if ctx.config.fuel.is_some_and(|timeout| start.elapsed() > timeout) {
-            break;
-        }
-        solutions.extend(tactics::struct_projection(ctx, &defs, &mut lookup));
-        if ctx.config.fuel.is_some_and(|timeout| start.elapsed() > timeout) {
-            break;
-        }
-        solutions.extend(tactics::impl_static_method(ctx, &defs, &mut lookup));
-        if ctx.config.fuel.is_some_and(|timeout| start.elapsed() > timeout) {
-            break;
-        }
-        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() {
diff --git a/src/tools/rust-analyzer/crates/hir/src/term_search/tactics.rs b/src/tools/rust-analyzer/crates/hir/src/term_search/tactics.rs
@@ -101,12 +101,14 @@ 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 +154,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 +197,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 +222,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 +282,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();
@@ -349,6 +361,7 @@ 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 +403,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();
@@ -478,6 +492,7 @@ 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 +569,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();
@@ -649,13 +665,15 @@ 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) {
@@ -720,6 +738,7 @@ 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 +747,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 +821,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();
@@ -888,6 +909,7 @@ 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 +918,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 +938,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);
diff --git a/src/tools/rust-analyzer/crates/ide-completion/src/completions/expr.rs b/src/tools/rust-analyzer/crates/ide-completion/src/completions/expr.rs
@@ -353,7 +353,7 @@ pub(crate) fn complete_expr(acc: &mut Completions, ctx: &CompletionContext<'_>)
             config: hir::term_search::TermSearchConfig {
                 enable_borrowcheck: false,
                 many_alternatives_threshold: 1,
-                depth: 3,
+                fuel: 200,
                 ..Default::default()
             },
         };
