Fix closure kind inference

Author: ShoyuVanilla

crates/hir-ty/src/infer/closure.rs

@@ -922,7 +922,7 @@ impl InferenceContext<'_> {
         }
     }
 
-    fn closure_kind(&self) -> FnTrait {
+    fn closure_kind_from_capture(&self) -> FnTrait {
         let mut r = FnTrait::Fn;
         for it in &self.current_captures {
             r = cmp::min(
@@ -939,7 +939,7 @@ impl InferenceContext<'_> {
         r
     }
 
-    fn analyze_closure(&mut self, closure: ClosureId) -> FnTrait {
+    fn analyze_closure(&mut self, closure: ClosureId, predicate: Option<FnTrait>) -> FnTrait {
         let (_, root) = self.db.lookup_intern_closure(closure.into());
         self.current_closure = Some(closure);
         let Expr::Closure { body, capture_by, .. } = &self.body[root] else {
@@ -957,7 +957,13 @@ impl InferenceContext<'_> {
         }
         self.restrict_precision_for_unsafe();
         // closure_kind should be done before adjust_for_move_closure
-        let closure_kind = self.closure_kind();
+        let closure_kind = {
+            let from_capture = self.closure_kind_from_capture();
+            // if predicate.unwrap_or(FnTrait::Fn) < from_capture {
+            //     // TODO: Diagnostic
+            // }
+            predicate.unwrap_or(from_capture)
+        };
         match capture_by {
             CaptureBy::Value => self.adjust_for_move_closure(),
             CaptureBy::Ref => (),
@@ -973,7 +979,9 @@ impl InferenceContext<'_> {
         let deferred_closures = self.sort_closures();
         for (closure, exprs) in deferred_closures.into_iter().rev() {
             self.current_captures = vec![];
-            let kind = self.analyze_closure(closure);
+
+            let predicate = self.table.get_closure_fn_trait_predicate(closure);
+            let kind = self.analyze_closure(closure, predicate);
 
             for (derefed_callee, callee_ty, params, expr) in exprs {
                 if let &Expr::Call { callee, .. } = &self.body[expr] {

crates/hir-ty/src/infer/unify.rs

@@ -1,10 +1,10 @@
 //! Unification and canonicalization logic.
 
-use std::{fmt, iter, mem};
+use std::{cmp, fmt, iter, mem};
 
 use chalk_ir::{
     cast::Cast, fold::TypeFoldable, interner::HasInterner, zip::Zip, CanonicalVarKind, FloatTy,
-    IntTy, TyVariableKind, UniverseIndex,
+    IntTy, TyVariableKind, UniverseIndex, WhereClause,
 };
 use chalk_solve::infer::ParameterEnaVariableExt;
 use either::Either;
@@ -14,11 +14,12 @@ use triomphe::Arc;
 
 use super::{InferOk, InferResult, InferenceContext, TypeError};
 use crate::{
-    consteval::unknown_const, db::HirDatabase, fold_tys_and_consts, static_lifetime,
-    to_chalk_trait_id, traits::FnTrait, AliasEq, AliasTy, BoundVar, Canonical, Const, ConstValue,
-    DebruijnIndex, GenericArg, GenericArgData, Goal, Guidance, InEnvironment, InferenceVar,
-    Interner, Lifetime, ParamKind, ProjectionTy, ProjectionTyExt, Scalar, Solution, Substitution,
-    TraitEnvironment, Ty, TyBuilder, TyExt, TyKind, VariableKind,
+    chalk_db::TraitId, consteval::unknown_const, db::HirDatabase, fold_tys_and_consts,
+    static_lifetime, to_chalk_trait_id, traits::FnTrait, AliasEq, AliasTy, BoundVar, Canonical,
+    ClosureId, Const, ConstValue, DebruijnIndex, DomainGoal, GenericArg, GenericArgData, Goal,
+    GoalData, Guidance, InEnvironment, InferenceVar, Interner, Lifetime, ParamKind, ProjectionTy,
+    ProjectionTyExt, Scalar, Solution, Substitution, TraitEnvironment, Ty, TyBuilder, TyExt,
+    TyKind, VariableKind,
 };
 
 impl InferenceContext<'_> {
@@ -146,6 +147,8 @@ pub(crate) struct InferenceTable<'a> {
     /// Double buffer used in [`Self::resolve_obligations_as_possible`] to cut down on
     /// temporary allocations.
     resolve_obligations_buffer: Vec<Canonicalized<InEnvironment<Goal>>>,
+    fn_trait_predicates: Vec<(Ty, FnTrait)>,
+    cached_fn_trait_ids: Option<CachedFnTraitIds>,
 }
 
 pub(crate) struct InferenceTableSnapshot {
@@ -154,15 +157,36 @@ pub(crate) struct InferenceTableSnapshot {
     type_variable_table_snapshot: Vec<TypeVariableFlags>,
 }
 
+#[derive(Clone)]
+struct CachedFnTraitIds {
+    fn_trait: TraitId,
+    fn_mut_trait: TraitId,
+    fn_once_trait: TraitId,
+}
+
+impl CachedFnTraitIds {
+    fn new(db: &dyn HirDatabase, trait_env: &Arc<TraitEnvironment>) -> Option<Self> {
+        let fn_trait = FnTrait::Fn.get_id(db, trait_env.krate.clone()).map(to_chalk_trait_id)?;
+        let fn_mut_trait =
+            FnTrait::FnMut.get_id(db, trait_env.krate.clone()).map(to_chalk_trait_id)?;
+        let fn_once_trait =
+            FnTrait::FnOnce.get_id(db, trait_env.krate.clone()).map(to_chalk_trait_id)?;
+        Some(Self { fn_trait, fn_mut_trait, fn_once_trait })
+    }
+}
+
 impl<'a> InferenceTable<'a> {
     pub(crate) fn new(db: &'a dyn HirDatabase, trait_env: Arc<TraitEnvironment>) -> Self {
+        let cached_fn_trait_ids = CachedFnTraitIds::new(db, &trait_env);
         InferenceTable {
             db,
             trait_env,
             var_unification_table: ChalkInferenceTable::new(),
             type_variable_table: Vec::new(),
             pending_obligations: Vec::new(),
             resolve_obligations_buffer: Vec::new(),
+            fn_trait_predicates: Vec::new(),
+            cached_fn_trait_ids,
         }
     }
 
@@ -498,6 +522,22 @@ impl<'a> InferenceTable<'a> {
     }
 
     fn register_obligation_in_env(&mut self, goal: InEnvironment<Goal>) {
+        if let Some(fn_trait_ids) = &self.cached_fn_trait_ids {
+            if let GoalData::DomainGoal(DomainGoal::Holds(WhereClause::Implemented(trait_ref))) =
+                goal.goal.data(Interner)
+            {
+                if let Some(ty) = trait_ref.substitution.type_parameters(Interner).next() {
+                    if trait_ref.trait_id == fn_trait_ids.fn_trait {
+                        self.fn_trait_predicates.push((ty, FnTrait::Fn));
+                    } else if trait_ref.trait_id == fn_trait_ids.fn_mut_trait {
+                        self.fn_trait_predicates.push((ty, FnTrait::FnMut));
+                    } else if trait_ref.trait_id == fn_trait_ids.fn_once_trait {
+                        self.fn_trait_predicates.push((ty, FnTrait::FnOnce));
+                    }
+                }
+            }
+        }
+
         let canonicalized = self.canonicalize(goal);
         if !self.try_resolve_obligation(&canonicalized) {
             self.pending_obligations.push(canonicalized);
@@ -791,6 +831,23 @@ impl<'a> InferenceTable<'a> {
             _ => c,
         }
     }
+
+    pub(super) fn get_closure_fn_trait_predicate(
+        &mut self,
+        closure_id: ClosureId,
+    ) -> Option<FnTrait> {
+        let predicates = mem::take(&mut self.fn_trait_predicates);
+        let res = predicates.iter().filter_map(|(ty, fn_trait)| {
+            if matches!(self.resolve_completely(ty.clone()).kind(Interner), TyKind::Closure(c, ..) if *c == closure_id) {
+                Some(*fn_trait)
+            } else {
+                None
+            }
+        }).fold(None, |acc, x| Some(cmp::max(acc.unwrap_or(FnTrait::FnOnce), x)));
+        self.fn_trait_predicates = predicates;
+
+        return res;
+    }
 }
 
 impl fmt::Debug for InferenceTable<'_> {

crates/hir-ty/src/mir/eval/tests.rs

@@ -636,6 +636,28 @@ fn main() {
     );
 }
 
+#[test]
+fn check_check_check() {
+    check_pass(
+        r#"
+//- minicore: fn
+struct Foo<F: FnOnce(i32)>(F);
+fn should_not_reach() {
+    _ // FIXME: replace this function with panic when that works
+}
+
+fn main() {
+    let mut a = 0;
+    let foo = Foo(move |val| {
+        a = val;
+    });
+    foo.0(10);
+    foo.0(20);
+}
+        "#,
+    )
+}
+
 #[test]
 fn closure_capture_array_const_generic() {
     check_pass(