Support non-body expression store for upvars_mentioned()

Author: ChayimFriedman2

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

@@ -1718,9 +1718,6 @@ fn coerce<'db>(
 
 fn is_capturing_closure(db: &dyn HirDatabase, closure: InternedClosureId) -> bool {
     let InternedClosure(owner, expr) = closure.loc(db);
-    let Some(body_owner) = owner.as_def_with_body() else {
-        return false;
-    };
-    upvars_mentioned(db, body_owner)
+    upvars_mentioned(db, owner)
         .is_some_and(|upvars| upvars.get(&expr).is_some_and(|upvars| !upvars.is_empty()))
 }

crates/hir-ty/src/upvars.rs

@@ -1,8 +1,8 @@
 //! A simple query to collect tall locals (upvars) a closure use.
 
 use hir_def::{
-    DefWithBodyId,
-    expr_store::{Body, path::Path},
+    DefWithBodyId, ExpressionStoreOwnerId, GenericDefId, VariantId,
+    expr_store::{ExpressionStore, path::Path},
     hir::{BindingId, Expr, ExprId, ExprOrPatId, Pat},
     resolver::{HasResolver, Resolver, ValueNs},
 };
@@ -36,18 +36,89 @@ impl Upvars {
     pub fn is_empty(&self) -> bool {
         self.0.is_empty()
     }
+
+    #[inline]
+    pub fn as_ref(&self) -> UpvarsRef<'_> {
+        UpvarsRef(&self.0)
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
+// Kept sorted.
+pub struct UpvarsRef<'db>(&'db [BindingId]);
+
+impl UpvarsRef<'_> {
+    #[inline]
+    pub fn contains(self, local: BindingId) -> bool {
+        self.0.binary_search(&local).is_ok()
+    }
+
+    #[inline]
+    pub fn iter(self) -> impl ExactSizeIterator<Item = BindingId> {
+        self.0.iter().copied()
+    }
+
+    #[inline]
+    pub fn is_empty(self) -> bool {
+        self.0.is_empty()
+    }
+
+    #[inline]
+    pub const fn empty() -> Self {
+        UpvarsRef(&[])
+    }
 }
 
 /// Returns a map from `Expr::Closure` to its upvars.
-#[salsa::tracked(returns(as_deref))]
 pub fn upvars_mentioned(
     db: &dyn HirDatabase,
-    owner: DefWithBodyId,
+    owner: ExpressionStoreOwnerId,
+) -> Option<&FxHashMap<ExprId, Upvars>> {
+    return match owner {
+        ExpressionStoreOwnerId::Signature(owner) => signature_upvars_mentioned(db, owner),
+        ExpressionStoreOwnerId::Body(owner) => body_upvars_mentioned(db, owner),
+        ExpressionStoreOwnerId::VariantFields(owner) => variant_fields_upvars_mentioned(db, owner),
+    };
+
+    #[salsa::tracked(returns(as_deref))]
+    pub fn signature_upvars_mentioned(
+        db: &dyn HirDatabase,
+        owner: GenericDefId,
+    ) -> Option<Box<FxHashMap<ExprId, Upvars>>> {
+        upvars_mentioned_impl(db, owner.into())
+    }
+
+    #[salsa::tracked(returns(as_deref))]
+    pub fn body_upvars_mentioned(
+        db: &dyn HirDatabase,
+        owner: DefWithBodyId,
+    ) -> Option<Box<FxHashMap<ExprId, Upvars>>> {
+        upvars_mentioned_impl(db, owner.into())
+    }
+
+    #[salsa::tracked(returns(as_deref))]
+    pub fn variant_fields_upvars_mentioned(
+        db: &dyn HirDatabase,
+        owner: VariantId,
+    ) -> Option<Box<FxHashMap<ExprId, Upvars>>> {
+        upvars_mentioned_impl(db, owner.into())
+    }
+}
+
+pub fn upvars_mentioned_impl(
+    db: &dyn HirDatabase,
+    owner: ExpressionStoreOwnerId,
 ) -> Option<Box<FxHashMap<ExprId, Upvars>>> {
-    let body = Body::of(db, owner);
+    let store = ExpressionStore::of(db, owner);
+    if store.const_expr_origins().is_empty() {
+        // Save constructing a Resolver.
+        return None;
+    }
     let mut resolver = owner.resolver(db);
     let mut result = FxHashMap::default();
-    handle_expr_outside_closure(db, &mut resolver, owner, body, body.root_expr(), &mut result);
+    for root_expr in store.expr_roots() {
+        handle_expr_outside_closure(db, &mut resolver, owner, store, root_expr, &mut result);
+    }
     return if result.is_empty() {
         None
     } else {
@@ -58,8 +129,8 @@ pub fn upvars_mentioned(
     fn handle_expr_outside_closure<'db>(
         db: &'db dyn HirDatabase,
         resolver: &mut Resolver<'db>,
-        owner: DefWithBodyId,
-        body: &Body,
+        owner: ExpressionStoreOwnerId,
+        body: &ExpressionStore,
         expr: ExprId,
         closures_map: &mut FxHashMap<ExprId, Upvars>,
     ) {
@@ -89,8 +160,8 @@ pub fn upvars_mentioned(
     fn handle_expr_inside_closure<'db>(
         db: &'db dyn HirDatabase,
         resolver: &mut Resolver<'db>,
-        owner: DefWithBodyId,
-        body: &Body,
+        owner: ExpressionStoreOwnerId,
+        body: &ExpressionStore,
         current_closure: ExprId,
         expr: ExprId,
         upvars: &mut FxHashSet<BindingId>,
@@ -170,8 +241,8 @@ pub fn upvars_mentioned(
 fn resolve_maybe_upvar<'db>(
     db: &'db dyn HirDatabase,
     resolver: &mut Resolver<'db>,
-    owner: DefWithBodyId,
-    body: &Body,
+    owner: ExpressionStoreOwnerId,
+    body: &ExpressionStore,
     current_closure: ExprId,
     expr: ExprId,
     id: ExprOrPatId,
@@ -198,7 +269,7 @@ fn resolve_maybe_upvar<'db>(
 #[cfg(test)]
 mod tests {
     use expect_test::{Expect, expect};
-    use hir_def::{ModuleDefId, expr_store::Body, nameres::crate_def_map};
+    use hir_def::{DefWithBodyId, ModuleDefId, expr_store::Body, nameres::crate_def_map};
     use itertools::Itertools;
     use span::Edition;
     use test_fixture::WithFixture;
@@ -220,7 +291,7 @@ mod tests {
                 .exactly_one()
                 .unwrap_or_else(|_| panic!("expected one function"));
             let (body, source_map) = Body::with_source_map(&db, func.into());
-            let Some(upvars) = upvars_mentioned(&db, func.into()) else {
+            let Some(upvars) = upvars_mentioned(&db, DefWithBodyId::from(func).into()) else {
                 expectation.assert_eq("");
                 return;
             };