Auto merge of #12668 - Veykril:mac-source-map, r=Veykril

fix: Simplify macro statement expansion handling

I only meant to fix https://github.com/rust-lang/rust-analyzer/issues/12644 but that somehow turned into a rewrite of the statement handling ... at least this fixes a few more issues in the IDE layer now

crates/hir-def/src/body.rs

@@ -19,7 +19,6 @@ use la_arena::{Arena, ArenaMap};
 use limit::Limit;
 use profile::Count;
 use rustc_hash::FxHashMap;
-use smallvec::SmallVec;
 use syntax::{ast, AstNode, AstPtr, SyntaxNodePtr};
 
 use crate::{
@@ -294,10 +293,6 @@ pub struct BodySourceMap {
     field_map: FxHashMap<InFile<AstPtr<ast::RecordExprField>>, ExprId>,
     field_map_back: FxHashMap<ExprId, InFile<AstPtr<ast::RecordExprField>>>,
 
-    /// Maps a macro call to its lowered expressions, a single one if it expands to an expression,
-    /// or multiple if it expands to MacroStmts.
-    macro_call_to_exprs: FxHashMap<InFile<AstPtr<ast::MacroCall>>, SmallVec<[ExprId; 1]>>,
-
     expansions: FxHashMap<InFile<AstPtr<ast::MacroCall>>, HirFileId>,
 
     /// Diagnostics accumulated during body lowering. These contain `AstPtr`s and so are stored in
@@ -466,9 +461,9 @@ impl BodySourceMap {
         self.field_map.get(&src).cloned()
     }
 
-    pub fn macro_expansion_expr(&self, node: InFile<&ast::MacroCall>) -> Option<&[ExprId]> {
+    pub fn macro_expansion_expr(&self, node: InFile<&ast::MacroExpr>) -> Option<ExprId> {
-        let src = node.map(AstPtr::new);
+        let src = node.map(AstPtr::new).map(AstPtr::upcast::<ast::MacroExpr>).map(AstPtr::upcast);
-        self.macro_call_to_exprs.get(&src).map(|it| &**it)
+        self.expr_map.get(&src).copied()
     }
 
     /// Get a reference to the body source map's diagnostics.

crates/hir-def/src/body/lower.rs

@@ -13,7 +13,6 @@ use hir_expand::{
 use la_arena::Arena;
 use profile::Count;
 use rustc_hash::FxHashMap;
-use smallvec::smallvec;
 use syntax::{
     ast::{
         self, ArrayExprKind, AstChildren, HasArgList, HasLoopBody, HasName, LiteralKind,
@@ -97,7 +96,6 @@ pub(super) fn lower(
             or_pats: Default::default(),
         },
         expander,
-        statements_in_scope: Vec::new(),
         name_to_pat_grouping: Default::default(),
         is_lowering_inside_or_pat: false,
     }
@@ -109,7 +107,6 @@ struct ExprCollector<'a> {
     expander: Expander,
     body: Body,
     source_map: BodySourceMap,
-    statements_in_scope: Vec<Statement>,
     // a poor-mans union-find?
     name_to_pat_grouping: FxHashMap<Name, Vec<PatId>>,
     is_lowering_inside_or_pat: bool,
@@ -514,27 +511,25 @@ impl ExprCollector<'_> {
             ast::Expr::MacroExpr(e) => {
                 let e = e.macro_call()?;
                 let macro_ptr = AstPtr::new(&e);
-                let id = self.collect_macro_call(e, macro_ptr.clone(), true, |this, expansion| {
+                let id = self.collect_macro_call(e, macro_ptr, true, |this, expansion| {
                     expansion.map(|it| this.collect_expr(it))
                 });
                 match id {
                     Some(id) => {
-                        self.source_map
+                        // Make the macro-call point to its expanded expression so we can query
-                            .macro_call_to_exprs
+                        // semantics on syntax pointers to the macro
-                            .insert(self.expander.to_source(macro_ptr), smallvec![id]);
+                        let src = self.expander.to_source(syntax_ptr);
+                        self.source_map.expr_map.insert(src, id);
                         id
                     }
-                    None => self.alloc_expr(Expr::Missing, syntax_ptr.clone()),
+                    None => self.alloc_expr(Expr::Missing, syntax_ptr),
                 }
             }
             ast::Expr::MacroStmts(e) => {
-                e.statements().for_each(|s| self.collect_stmt(s));
+                let statements = e.statements().filter_map(|s| self.collect_stmt(s)).collect();
-                let tail = e
+                let tail = e.expr().map(|e| self.collect_expr(e));
-                    .expr()
-                    .map(|e| self.collect_expr(e))
-                    .unwrap_or_else(|| self.alloc_expr(Expr::Missing, syntax_ptr.clone()));
 
-                self.alloc_expr(Expr::MacroStmts { tail }, syntax_ptr)
+                self.alloc_expr(Expr::MacroStmts { tail, statements }, syntax_ptr)
             }
             ast::Expr::UnderscoreExpr(_) => self.alloc_expr(Expr::Underscore, syntax_ptr),
         })
@@ -607,11 +602,11 @@ impl ExprCollector<'_> {
         }
     }
 
-    fn collect_stmt(&mut self, s: ast::Stmt) {
+    fn collect_stmt(&mut self, s: ast::Stmt) -> Option<Statement> {
         match s {
             ast::Stmt::LetStmt(stmt) => {
                 if self.check_cfg(&stmt).is_none() {
-                    return;
+                    return None;
                 }
                 let pat = self.collect_pat_opt(stmt.pat());
                 let type_ref =
@@ -621,70 +616,61 @@ impl ExprCollector<'_> {
                     .let_else()
                     .and_then(|let_else| let_else.block_expr())
                     .map(|block| self.collect_block(block));
-                self.statements_in_scope.push(Statement::Let {
+                Some(Statement::Let { pat, type_ref, initializer, else_branch })
-                    pat,
-                    type_ref,
-                    initializer,
-                    else_branch,
-                });
             }
             ast::Stmt::ExprStmt(stmt) => {
-                if let Some(expr) = stmt.expr() {
+                let expr = stmt.expr();
-                    if self.check_cfg(&expr).is_none() {
+                if let Some(expr) = &expr {
-                        return;
+                    if self.check_cfg(expr).is_none() {
+                        return None;
                     }
                 }
                 let has_semi = stmt.semicolon_token().is_some();
-                // Note that macro could be expended to multiple statements
+                // Note that macro could be expanded to multiple statements
-                if let Some(ast::Expr::MacroExpr(e)) = stmt.expr() {
+                if let Some(expr @ ast::Expr::MacroExpr(mac)) = &expr {
-                    let m = match e.macro_call() {
+                    let mac_call = mac.macro_call()?;
-                        Some(it) => it,
+                    let syntax_ptr = AstPtr::new(expr);
-                        None => return,
+                    let macro_ptr = AstPtr::new(&mac_call);
-                    };
+                    let stmt = self.collect_macro_call(
-                    let macro_ptr = AstPtr::new(&m);
+                        mac_call,
-                    let syntax_ptr = AstPtr::new(&stmt.expr().unwrap());
+                        macro_ptr,
-
+                        false,
-                    let prev_stmt = self.statements_in_scope.len();
+                        |this, expansion: Option<ast::MacroStmts>| match expansion {
-                    self.collect_macro_call(m, macro_ptr.clone(), false, |this, expansion| {
-                        match expansion {
                             Some(expansion) => {
-                                let statements: ast::MacroStmts = expansion;
+                                let statements = expansion
+                                    .statements()
+                                    .filter_map(|stmt| this.collect_stmt(stmt))
+                                    .collect();
+                                let tail = expansion.expr().map(|expr| this.collect_expr(expr));
 
-                                statements.statements().for_each(|stmt| this.collect_stmt(stmt));
+                                let mac_stmts = this.alloc_expr(
-                                if let Some(expr) = statements.expr() {
+                                    Expr::MacroStmts { tail, statements },
-                                    let expr = this.collect_expr(expr);
+                                    AstPtr::new(&ast::Expr::MacroStmts(expansion)),
-                                    this.statements_in_scope
+                                );
-                                        .push(Statement::Expr { expr, has_semi });
-                                }
-                            }
-                            None => {
-                                let expr = this.alloc_expr(Expr::Missing, syntax_ptr.clone());
-                                this.statements_in_scope.push(Statement::Expr { expr, has_semi });
-                            }
-                        }
-                    });
 
-                    let mut macro_exprs = smallvec![];
+                                Some(mac_stmts)
-                    for stmt in &self.statements_in_scope[prev_stmt..] {
-                        match *stmt {
-                            Statement::Let { initializer, else_branch, .. } => {
-                                macro_exprs.extend(initializer);
-                                macro_exprs.extend(else_branch);
                             }
-                            Statement::Expr { expr, .. } => macro_exprs.push(expr),
+                            None => None,
+                        },
+                    );
+
+                    let expr = match stmt {
+                        Some(expr) => {
+                            // Make the macro-call point to its expanded expression so we can query
+                            // semantics on syntax pointers to the macro
+                            let src = self.expander.to_source(syntax_ptr);
+                            self.source_map.expr_map.insert(src, expr);
+                            expr
                         }
-                    }
+                        None => self.alloc_expr(Expr::Missing, syntax_ptr),
-                    if !macro_exprs.is_empty() {
+                    };
-                        self.source_map
+                    Some(Statement::Expr { expr, has_semi })
-                            .macro_call_to_exprs
-                            .insert(self.expander.to_source(macro_ptr), macro_exprs);
-                    }
                 } else {
-                    let expr = self.collect_expr_opt(stmt.expr());
+                    let expr = self.collect_expr_opt(expr);
-                    self.statements_in_scope.push(Statement::Expr { expr, has_semi });
+                    Some(Statement::Expr { expr, has_semi })
                 }
             }
-            ast::Stmt::Item(_item) => {}
+            ast::Stmt::Item(_item) => None,
         }
     }
 
@@ -703,25 +689,27 @@ impl ExprCollector<'_> {
         };
         let prev_def_map = mem::replace(&mut self.expander.def_map, def_map);
         let prev_local_module = mem::replace(&mut self.expander.module, module);
-        let prev_statements = std::mem::take(&mut self.statements_in_scope);
 
-        block.statements().for_each(|s| self.collect_stmt(s));
+        let mut statements: Vec<_> =
-        block.tail_expr().and_then(|e| {
+            block.statements().filter_map(|s| self.collect_stmt(s)).collect();
-            let expr = self.maybe_collect_expr(e)?;
+        let tail = block.tail_expr().and_then(|e| self.maybe_collect_expr(e));
-            self.statements_in_scope.push(Statement::Expr { expr, has_semi: false });
+        let tail = tail.or_else(|| {
-            Some(())
+            let stmt = statements.pop()?;
+            if let Statement::Expr { expr, has_semi: false } = stmt {
+                return Some(expr);
+            }
+            statements.push(stmt);
+            None
         });
 
-        let mut tail = None;
-        if let Some(Statement::Expr { expr, has_semi: false }) = self.statements_in_scope.last() {
-            tail = Some(*expr);
-            self.statements_in_scope.pop();
-        }
-        let tail = tail;
-        let statements = std::mem::replace(&mut self.statements_in_scope, prev_statements).into();
         let syntax_node_ptr = AstPtr::new(&block.into());
         let expr_id = self.alloc_expr(
-            Expr::Block { id: block_id, statements, tail, label: None },
+            Expr::Block {
+                id: block_id,
+                statements: statements.into_boxed_slice(),
+                tail,
+                label: None,
+            },
             syntax_node_ptr,
         );
 
@@ -903,10 +891,12 @@ impl ExprCollector<'_> {
             ast::Pat::MacroPat(mac) => match mac.macro_call() {
                 Some(call) => {
                     let macro_ptr = AstPtr::new(&call);
+                    let src = self.expander.to_source(Either::Left(AstPtr::new(&pat)));
                     let pat =
                         self.collect_macro_call(call, macro_ptr, true, |this, expanded_pat| {
                             this.collect_pat_opt_(expanded_pat)
                         });
+                    self.source_map.pat_map.insert(src, pat);
                     return pat;
                 }
                 None => Pat::Missing,

crates/hir-def/src/body/scope.rs

@@ -145,27 +145,28 @@ fn compute_block_scopes(
     tail: Option<ExprId>,
     body: &Body,
     scopes: &mut ExprScopes,
-    mut scope: ScopeId,
+    scope: &mut ScopeId,
 ) {
     for stmt in statements {
         match stmt {
             Statement::Let { pat, initializer, else_branch, .. } => {
                 if let Some(expr) = initializer {
-                    compute_expr_scopes(*expr, body, scopes, &mut scope);
+                    compute_expr_scopes(*expr, body, scopes, scope);
                 }
                 if let Some(expr) = else_branch {
-                    compute_expr_scopes(*expr, body, scopes, &mut scope);
+                    compute_expr_scopes(*expr, body, scopes, scope);
                 }
-                scope = scopes.new_scope(scope);
+
-                scopes.add_bindings(body, scope, *pat);
+                *scope = scopes.new_scope(*scope);
+                scopes.add_bindings(body, *scope, *pat);
             }
             Statement::Expr { expr, .. } => {
-                compute_expr_scopes(*expr, body, scopes, &mut scope);
+                compute_expr_scopes(*expr, body, scopes, scope);
             }
         }
     }
     if let Some(expr) = tail {
-        compute_expr_scopes(expr, body, scopes, &mut scope);
+        compute_expr_scopes(expr, body, scopes, scope);
     }
 }
 
@@ -175,12 +176,15 @@ fn compute_expr_scopes(expr: ExprId, body: &Body, scopes: &mut ExprScopes, scope
 
     scopes.set_scope(expr, *scope);
     match &body[expr] {
+        Expr::MacroStmts { statements, tail } => {
+            compute_block_scopes(statements, *tail, body, scopes, scope);
+        }
         Expr::Block { statements, tail, id, label } => {
-            let scope = scopes.new_block_scope(*scope, *id, make_label(label));
+            let mut scope = scopes.new_block_scope(*scope, *id, make_label(label));
             // Overwrite the old scope for the block expr, so that every block scope can be found
             // via the block itself (important for blocks that only contain items, no expressions).
             scopes.set_scope(expr, scope);
-            compute_block_scopes(statements, *tail, body, scopes, scope);
+            compute_block_scopes(statements, *tail, body, scopes, &mut scope);
         }
         Expr::For { iterable, pat, body: body_expr, label } => {
             compute_expr_scopes(*iterable, body, scopes, scope);

crates/hir-def/src/expr.rs

@@ -199,7 +199,8 @@ pub enum Expr {
         body: ExprId,
     },
     MacroStmts {
-        tail: ExprId,
+        statements: Box<[Statement]>,
+        tail: Option<ExprId>,
     },
     Array(Array),
     Literal(Literal),
@@ -254,7 +255,7 @@ impl Expr {
             Expr::Let { expr, .. } => {
                 f(*expr);
             }
-            Expr::Block { statements, tail, .. } => {
+            Expr::MacroStmts { tail, statements } | Expr::Block { statements, tail, .. } => {
                 for stmt in statements.iter() {
                     match stmt {
                         Statement::Let { initializer, .. } => {
@@ -344,7 +345,6 @@ impl Expr {
                     f(*repeat)
                 }
             },
-            Expr::MacroStmts { tail } => f(*tail),
             Expr::Literal(_) => {}
             Expr::Underscore => {}
         }

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

@@ -774,7 +774,9 @@ impl<'a> InferenceContext<'a> {
                     None => self.table.new_float_var(),
                 },
             },
-            Expr::MacroStmts { tail } => self.infer_expr_inner(*tail, expected),
+            Expr::MacroStmts { tail, statements } => {
+                self.infer_block(tgt_expr, statements, *tail, expected)
+            }
             Expr::Underscore => {
                 // Underscore expressions may only appear in assignee expressions,
                 // which are handled by `infer_assignee_expr()`, so any underscore

crates/hir-ty/src/tests/macros.rs

@@ -1,6 +1,8 @@
 use expect_test::expect;
 use test_utils::{bench, bench_fixture, skip_slow_tests};
 
+use crate::tests::check_infer_with_mismatches;
+
 use super::{check_infer, check_types};
 
 #[test]
@@ -191,6 +193,8 @@ fn expr_macro_def_expanded_in_various_places() {
             !0..6 '1isize': isize
             !0..6 '1isize': isize
             !0..6 '1isize': isize
+            !0..6 '1isize': isize
+            !0..6 '1isize': isize
             39..442 '{     ...!(); }': ()
             73..94 'spam!(...am!())': {unknown}
             100..119 'for _ ...!() {}': ()
@@ -272,6 +276,8 @@ fn expr_macro_rules_expanded_in_various_places() {
             !0..6 '1isize': isize
             !0..6 '1isize': isize
             !0..6 '1isize': isize
+            !0..6 '1isize': isize
+            !0..6 '1isize': isize
             53..456 '{     ...!(); }': ()
             87..108 'spam!(...am!())': {unknown}
             114..133 'for _ ...!() {}': ()
@@ -306,7 +312,6 @@ fn expr_macro_expanded_in_stmts() {
         }
         "#,
         expect![[r#"
-            !0..8 'leta=();': ()
             !0..8 'leta=();': ()
             !3..4 'a': ()
             !5..7 '()': ()
@@ -335,11 +340,11 @@ fn recurisve_macro_expanded_in_stmts() {
         }
         "#,
         expect![[r#"
-            !0..7 'leta=3;': {unknown}
+            !0..7 'leta=3;': ()
-            !0..7 'leta=3;': {unknown}
+            !0..13 'ng!{[leta=3]}': ()
-            !0..13 'ng!{[leta=3]}': {unknown}
+            !0..13 'ng!{[leta=]3}': ()
-            !0..13 'ng!{[leta=]3}': {unknown}
+            !0..13 'ng!{[leta]=3}': ()
-            !0..13 'ng!{[leta]=3}': {unknown}
+            !0..13 'ng!{[let]a=3}': ()
             !3..4 'a': i32
             !5..6 '3': i32
             196..237 '{     ...= a; }': ()
@@ -364,8 +369,8 @@ fn recursive_inner_item_macro_rules() {
         "#,
         expect![[r#"
             !0..1 '1': i32
-            !0..26 'macro_...>{1};}': {unknown}
+            !0..7 'mac!($)': ()
-            !0..26 'macro_...>{1};}': {unknown}
+            !0..26 'macro_...>{1};}': ()
             107..143 '{     ...!(); }': ()
             129..130 'a': i32
         "#]],
@@ -1244,3 +1249,28 @@ fn infinitely_recursive_macro_type() {
         "#]],
     );
 }
+
+#[test]
+fn cfg_tails() {
+    check_infer_with_mismatches(
+        r#"
+//- /lib.rs crate:foo cfg:feature=foo
+struct S {}
+
+impl S {
+    fn new2(bar: u32) -> Self {
+        #[cfg(feature = "foo")]
+        { Self { } }
+        #[cfg(not(feature = "foo"))]
+        { Self { } }
+    }
+}
+"#,
+        expect![[r#"
+            34..37 'bar': u32
+            52..170 '{     ...     }': S
+            62..106 '#[cfg(... { } }': S
+            96..104 'Self { }': S
+        "#]],
+    );
+}

crates/hir-ty/src/tests/regression.rs

@@ -573,6 +573,7 @@ fn issue_6811() {
         }
         "#,
         expect![[r#"
+            !0..16 'let_a=...t_b=1;': ()
             !3..5 '_a': i32
             !6..7 '1': i32
             !11..13 '_b': i32

crates/hir/src/lib.rs

@@ -2223,6 +2223,7 @@ impl Local {
         let src = source_map.pat_syntax(self.pat_id).unwrap(); // Hmm...
         let root = src.file_syntax(db.upcast());
         src.map(|ast| match ast {
+            // Suspicious unwrap
             Either::Left(it) => Either::Left(it.cast().unwrap().to_node(&root)),
             Either::Right(it) => Either::Right(it.to_node(&root)),
         })

crates/hir/src/source_analyzer.rs

@@ -539,24 +539,26 @@ impl SourceAnalyzer {
                 _ => (),
             }
         }
+        let macro_expr = match macro_call
+            .map(|it| it.syntax().parent().and_then(ast::MacroExpr::cast))
+            .transpose()
+        {
+            Some(it) => it,
+            None => return false,
+        };
+
         if let (Some((def, body, sm)), Some(infer)) = (&self.def, &self.infer) {
-            if let Some(expr_ids) = sm.macro_expansion_expr(macro_call) {
+            if let Some(expanded_expr) = sm.macro_expansion_expr(macro_expr.as_ref()) {
                 let mut is_unsafe = false;
-                for &expr_id in expr_ids {
+                unsafe_expressions(
-                    unsafe_expressions(
+                    db,
-                        db,
+                    infer,
-                        infer,
+                    *def,
-                        *def,
+                    body,
-                        body,
+                    expanded_expr,
-                        expr_id,
+                    &mut |UnsafeExpr { inside_unsafe_block, .. }| is_unsafe |= !inside_unsafe_block,
-                        &mut |UnsafeExpr { inside_unsafe_block, .. }| {
+                );
-                            is_unsafe |= !inside_unsafe_block
+                return is_unsafe;
-                        },
-                    );
-                    if is_unsafe {
-                        return true;
-                    }
-                }
             }
         }
         false

crates/ide-completion/src/tests/special.rs

@@ -810,3 +810,33 @@ fn main() {
         "#]],
     )
 }
+
+#[test]
+fn regression_12644() {
+    check(
+        r#"
+macro_rules! __rust_force_expr {
+    ($e:expr) => {
+        $e
+    };
+}
+macro_rules! vec {
+    ($elem:expr) => {
+        __rust_force_expr!($elem)
+    };
+}
+
+struct Struct;
+impl Struct {
+    fn foo(self) {}
+}
+
+fn f() {
+    vec![Struct].$0;
+}
+"#,
+        expect![[r#"
+            me foo() fn(self)
+        "#]],
+    );
+}

crates/ide/src/highlight_related.rs

@@ -689,7 +689,7 @@ fn foo() ->$0 u32 {
     never();
  // ^^^^^^^
     never!();
- // FIXME sema doesn't give us types for macrocalls
+ // ^^^^^^^^
 
     Never.never();
  // ^^^^^^^^^^^^^

crates/ide/src/hover/tests.rs

@@ -4920,3 +4920,34 @@ impl T for () {
         "#]],
     );
 }
+
+#[test]
+fn hover_ranged_macro_call() {
+    check_hover_range(
+        r#"
+macro_rules! __rust_force_expr {
+    ($e:expr) => {
+        $e
+    };
+}
+macro_rules! vec {
+    ($elem:expr) => {
+        __rust_force_expr!($elem)
+    };
+}
+
+struct Struct;
+impl Struct {
+    fn foo(self) {}
+}
+
+fn f() {
+    $0vec![Struct]$0;
+}
+"#,
+        expect![[r#"
+            ```rust
+            Struct
+            ```"#]],
+    );
+}

crates/syntax/src/ptr.rs

@@ -67,6 +67,13 @@ impl<N: AstNode> AstPtr<N> {
         Some(AstPtr { raw: self.raw, _ty: PhantomData })
     }
 
+    pub fn upcast<M: AstNode>(self) -> AstPtr<M>
+    where
+        N: Into<M>,
+    {
+        AstPtr { raw: self.raw, _ty: PhantomData }
+    }
+
     /// Like `SyntaxNodePtr::cast` but the trait bounds work out.
     pub fn try_from_raw(raw: SyntaxNodePtr) -> Option<AstPtr<N>> {
         N::can_cast(raw.kind()).then(|| AstPtr { raw, _ty: PhantomData })