hir_def/

upvars.rs

//! A simple query to collect tall locals (upvars) a closure use.

use base_db::SourceDatabase;
use hir_expand::mod_path::PathKind;
use rustc_hash::FxHashSet;

use crate::{
    DefWithBodyId, ExpressionStoreOwnerId, GenericDefId, ThinFxHashMap, VariantId,
    expr_store::{ExpressionStore, StoreVisitor, StoreVisitorExt, path::Path},
    hir::{BindingId, Expr, ExprId, PatId},
    resolver::{HasResolver, Resolver, ValueNs},
    type_ref::TypeRefId,
};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
// Kept sorted.
pub struct Upvars(Box<[BindingId]>);

impl Upvars {
    fn new(upvars: &FxHashSet<BindingId>) -> Upvars {
        let mut upvars = upvars.iter().copied().collect::<Box<[_]>>();
        upvars.sort_unstable();
        Upvars(upvars)
    }

    #[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 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.
pub fn upvars_mentioned(
    db: &dyn SourceDatabase,
    owner: ExpressionStoreOwnerId,
) -> &ThinFxHashMap<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]
    pub fn signature_upvars_mentioned(
        db: &dyn SourceDatabase,
        owner: GenericDefId,
    ) -> ThinFxHashMap<ExprId, Upvars> {
        upvars_mentioned_impl(db, owner.into())
    }

    #[salsa::tracked]
    pub fn body_upvars_mentioned(
        db: &dyn SourceDatabase,
        owner: DefWithBodyId,
    ) -> ThinFxHashMap<ExprId, Upvars> {
        upvars_mentioned_impl(db, owner.into())
    }

    #[salsa::tracked]
    pub fn variant_fields_upvars_mentioned(
        db: &dyn SourceDatabase,
        owner: VariantId,
    ) -> ThinFxHashMap<ExprId, Upvars> {
        upvars_mentioned_impl(db, owner.into())
    }
}

pub fn upvars_mentioned_impl(
    db: &dyn SourceDatabase,
    owner: ExpressionStoreOwnerId,
) -> ThinFxHashMap<ExprId, Upvars> {
    let store = ExpressionStore::of(db, owner);
    if store.expr_roots().next().is_none() {
        return ThinFxHashMap::new();
    }
    let resolver = owner.resolver(db);
    let mut visitor = UpvarsMentionedVisitor {
        db,
        resolver,
        owner,
        store,
        closures_map: ThinFxHashMap::default(),
        current_closure: None,
    };
    visitor.on_exprs(store.expr_roots());
    let mut result = visitor.closures_map;
    result.shrink_to_fit();
    result
}

struct UpvarsMentionedVisitor<'db> {
    db: &'db dyn SourceDatabase,
    resolver: Resolver<'db>,
    owner: ExpressionStoreOwnerId,
    store: &'db ExpressionStore,
    closures_map: ThinFxHashMap<ExprId, Upvars>,
    current_closure: Option<(ExprId, FxHashSet<BindingId>)>,
}

impl UpvarsMentionedVisitor<'_> {
    fn resolve_maybe_upvar(&mut self, expr: ExprId, path: &Path) {
        let Some((current_closure, upvars)) = &mut self.current_closure else { return };

        if let Path::BarePath(mod_path) = path
        && matches!(mod_path.kind, PathKind::Plain | PathKind::SELF)
        // `self` is length zero.
        && mod_path.segments().len() <= 1
        {
            // Could be a variable.
            let guard = self.resolver.update_to_inner_scope(self.db, self.owner, expr);
            let resolution = self.resolver.resolve_path_in_value_ns_fully(
                self.db,
                path,
                self.store.expr_or_pat_path_hygiene(expr.into()),
            );
            if let Some(ValueNs::LocalBinding(local)) = resolution
                && self.store.binding_owner(local) != Some(*current_closure)
            {
                upvars.insert(local);
            }
            self.resolver.reset_to_guard(guard);
        }
    }
}

impl StoreVisitor for UpvarsMentionedVisitor<'_> {
    fn on_expr(&mut self, expr: ExprId) {
        match &self.store[expr] {
            Expr::Path(path) => {
                self.resolve_maybe_upvar(expr, path);
                self.on_path(path);
            }
            Expr::Closure {
                body: body_expr,
                args,
                arg_types,
                ret_type,
                capture_by: _,
                closure_kind: _,
            } => {
                self.on_pats(args);
                arg_types.iter().for_each(|arg_ty| self.on_type_opt(*arg_ty));
                self.on_type_opt(*ret_type);

                let mut old_current_closure =
                    self.current_closure.replace((expr, FxHashSet::default()));
                self.on_expr(*body_expr);
                let closure_upvars = self.current_closure.take().unwrap().1;
                if !closure_upvars.is_empty() {
                    self.closures_map.insert(expr, Upvars::new(&closure_upvars));
                    // All nested closure's upvars are also upvars of the parent closure.
                    if let Some((current_closure, upvars)) = &mut old_current_closure {
                        upvars.extend(closure_upvars.iter().copied().filter(|local| {
                            self.store.binding_owner(*local) != Some(*current_closure)
                        }));
                    }
                }
                self.current_closure = old_current_closure;
            }
            _ => self.store.visit_expr_children(expr, self),
        }
    }

    fn on_anon_const_expr(&mut self, expr: ExprId) {
        // Anon consts don't contribute upvars.
        let old_current_closure = self.current_closure.take();
        self.on_expr(expr);
        self.current_closure = old_current_closure;
    }

    fn on_pat(&mut self, pat: PatId) {
        self.store.visit_pat_children(pat, self);
    }

    fn on_type(&mut self, ty: TypeRefId) {
        // Anon consts don't contribute upvars.
        let old_current_closure = self.current_closure.take();
        self.store.visit_type_ref_children(ty, &mut *self);
        self.current_closure = old_current_closure;
    }
}

#[cfg(test)]
mod tests {
    use expect_test::{Expect, expect};
    use itertools::Itertools;
    use span::Edition;
    use test_fixture::WithFixture;

    use crate::{
        AssocItemId, DefWithBodyId, ModuleDefId, expr_store::Body, nameres::crate_def_map,
        test_db::TestDB, upvars::upvars_mentioned,
    };

    #[track_caller]
    fn check(#[rust_analyzer::rust_fixture] ra_fixture: &str, expectation: Expect) {
        let db = TestDB::with_files(ra_fixture);
        let def_map = crate_def_map(&db, db.test_crate());
        let func = def_map
            .modules()
            .flat_map(|(_, module)| module.scope.declarations())
            .filter_map(|decl| match decl {
                ModuleDefId::FunctionId(func) => Some(func),
                _ => None,
            })
            .chain(def_map.modules().flat_map(|(_, module)| {
                module.scope.impls().flat_map(|impl_| &*impl_.impl_items(&db).items).filter_map(
                    |&(_, item)| match item {
                        AssocItemId::FunctionId(it) => Some(it),
                        _ => None,
                    },
                )
            }))
            .exactly_one()
            .unwrap_or_else(|_| panic!("expected one function"));
        let (body, source_map) = Body::with_source_map(&db, func.into());
        let upvars = upvars_mentioned(&db, DefWithBodyId::from(func).into());
        let mut closures = Vec::new();
        for (&closure, upvars) in upvars {
            let closure_range = source_map.expr_syntax(closure).unwrap().value.text_range();
            let upvars = upvars
                .iter()
                .map(|local| body[local].name.display(&db, Edition::CURRENT))
                .join(", ");
            closures.push((closure_range, upvars));
        }
        closures.sort_unstable_by_key(|(range, _)| (range.start(), range.end()));
        let closures =
            closures.into_iter().map(|(range, upvars)| format!("{range:?}: {upvars}")).join("\n");
        expectation.assert_eq(&closures);
    }

    #[test]
    fn simple() {
        check(
            r#"
struct foo;
fn foo(param: i32) {
    let local = "boo";
    || { param; foo };
    || local;
    || { param; local; param; local; };
    || 0xDEAFBEAF;
}
        "#,
            expect![[r#"
                60..77: param
                83..91: local
                97..131: param, local"#]],
        );
    }

    #[test]
    fn nested() {
        check(
            r#"
fn foo() {
    let (a, b);
    || {
        || a;
        || b;
    };
}
        "#,
            expect![[r#"
                31..69: a, b
                44..48: a
                58..62: b"#]],
        );
    }

    #[test]
    fn closure_var() {
        check(
            r#"
fn foo() {
    let upvar = 1;
    |closure_param: i32| {
        let closure_local = closure_param;
        closure_local + upvar
    };
}
        "#,
            expect!["34..135: upvar"],
        );
    }

    #[test]
    fn closure_var_nested() {
        check(
            r#"
fn foo() {
    let a = 1;
    |b: i32| {
        || {
            let c = 123;
            a + b + c
        }
    };
}
        "#,
            expect![[r#"
                30..116: a
                49..110: a, b"#]],
        );
    }

    #[test]
    fn self_upvar() {
        check(
            r#"
struct Foo(i32);
impl Foo {
    fn foo(&self) {
        || self.0;
    }
}
        "#,
            expect!["56..65: self"],
        );
    }

    #[test]
    fn const_block() {
        check(
            r#"
fn main() {
    || {
        const {
            let v = ();
            v
        }
    };
}
        "#,
            expect![""],
        );
    }

    #[test]
    fn async_fn_local() {
        check(
            r#"
async fn foo() {
    let v = ();
    v
}
        "#,
            expect![[r#""#]],
        );
    }
}