Store patterns desugared from destructuring assignments in source map

And few more fixups.

I was worried this will lead to more memory usage since `ExprOrPatId` is double the size of `ExprId`, but this does not regress `analysis-stats .`. If this turns out to be a problem, we can easily use the high bit to encode this information.

Author: ChayimFriedman2

crates/hir-def/src/body.rs

@@ -10,6 +10,7 @@ use std::ops::{Deref, Index};
 
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
+use either::Either;
 use hir_expand::{name::Name, ExpandError, InFile};
 use la_arena::{Arena, ArenaMap, Idx, RawIdx};
 use rustc_hash::FxHashMap;
@@ -22,8 +23,8 @@ use crate::{
     db::DefDatabase,
     expander::Expander,
     hir::{
-        dummy_expr_id, Array, AsmOperand, Binding, BindingId, Expr, ExprId, Label, LabelId, Pat,
-        PatId, RecordFieldPat, Statement,
+        dummy_expr_id, Array, AsmOperand, Binding, BindingId, Expr, ExprId, ExprOrPatId, Label,
+        LabelId, Pat, PatId, RecordFieldPat, Statement,
     },
     item_tree::AttrOwner,
     nameres::DefMap,
@@ -68,9 +69,12 @@ pub type LabelSource = InFile<LabelPtr>;
 pub type FieldPtr = AstPtr<ast::RecordExprField>;
 pub type FieldSource = InFile<FieldPtr>;
 
-pub type PatFieldPtr = AstPtr<ast::RecordPatField>;
+pub type PatFieldPtr = AstPtr<Either<ast::RecordExprField, ast::RecordPatField>>;
 pub type PatFieldSource = InFile<PatFieldPtr>;
 
+pub type ExprOrPatPtr = AstPtr<Either<ast::Expr, ast::Pat>>;
+pub type ExprOrPatSource = InFile<ExprOrPatPtr>;
+
 /// An item body together with the mapping from syntax nodes to HIR expression
 /// IDs. This is needed to go from e.g. a position in a file to the HIR
 /// expression containing it; but for type inference etc., we want to operate on
@@ -84,11 +88,13 @@ pub type PatFieldSource = InFile<PatFieldPtr>;
 /// this properly for macros.
 #[derive(Default, Debug, Eq, PartialEq)]
 pub struct BodySourceMap {
-    expr_map: FxHashMap<ExprSource, ExprId>,
+    // AST expressions can create patterns in destructuring assignments. Therefore, `ExprSource` can also map
+    // to `PatId`, and `PatId` can also map to `ExprSource` (the other way around is unaffected).
+    expr_map: FxHashMap<ExprSource, ExprOrPatId>,
     expr_map_back: ArenaMap<ExprId, ExprSource>,
 
     pat_map: FxHashMap<PatSource, PatId>,
-    pat_map_back: ArenaMap<PatId, PatSource>,
+    pat_map_back: ArenaMap<PatId, ExprOrPatSource>,
 
     label_map: FxHashMap<LabelSource, LabelId>,
     label_map_back: ArenaMap<LabelId, LabelSource>,
@@ -372,7 +378,7 @@ impl Body {
                             if let &Some(expr) = else_branch {
                                 f(expr);
                             }
-                            walk_exprs_in_pat(self, *pat, &mut f);
+                            self.walk_exprs_in_pat(*pat, &mut f);
                         }
                         Statement::Expr { expr: expression, .. } => f(*expression),
                         Statement::Item => (),
@@ -448,18 +454,18 @@ impl Body {
                 }
             },
             &Expr::Assignment { target, value } => {
-                walk_exprs_in_pat(self, target, &mut f);
+                self.walk_exprs_in_pat(target, &mut f);
                 f(value);
             }
         }
+    }
 
-        fn walk_exprs_in_pat(this: &Body, pat_id: PatId, f: &mut impl FnMut(ExprId)) {
-            this.walk_pats(pat_id, &mut |pat| {
-                if let Pat::Expr(expr) | Pat::ConstBlock(expr) = this[pat] {
-                    f(expr);
-                }
-            });
-        }
+    pub fn walk_exprs_in_pat(&self, pat_id: PatId, f: &mut impl FnMut(ExprId)) {
+        self.walk_pats(pat_id, &mut |pat| {
+            if let Pat::Expr(expr) | Pat::ConstBlock(expr) = self[pat] {
+                f(expr);
+            }
+        });
     }
 }
 
@@ -514,11 +520,18 @@ impl Index<BindingId> for Body {
 // FIXME: Change `node_` prefix to something more reasonable.
 // Perhaps `expr_syntax` and `expr_id`?
 impl BodySourceMap {
+    pub fn expr_or_pat_syntax(&self, id: ExprOrPatId) -> Result<ExprOrPatSource, SyntheticSyntax> {
+        match id {
+            ExprOrPatId::ExprId(id) => self.expr_syntax(id).map(|it| it.map(AstPtr::wrap_left)),
+            ExprOrPatId::PatId(id) => self.pat_syntax(id),
+        }
+    }
+
     pub fn expr_syntax(&self, expr: ExprId) -> Result<ExprSource, SyntheticSyntax> {
         self.expr_map_back.get(expr).cloned().ok_or(SyntheticSyntax)
     }
 
-    pub fn node_expr(&self, node: InFile<&ast::Expr>) -> Option<ExprId> {
+    pub fn node_expr(&self, node: InFile<&ast::Expr>) -> Option<ExprOrPatId> {
         let src = node.map(AstPtr::new);
         self.expr_map.get(&src).cloned()
     }
@@ -534,7 +547,7 @@ impl BodySourceMap {
         self.expansions.iter().map(|(&a, &b)| (a, b))
     }
 
-    pub fn pat_syntax(&self, pat: PatId) -> Result<PatSource, SyntheticSyntax> {
+    pub fn pat_syntax(&self, pat: PatId) -> Result<ExprOrPatSource, SyntheticSyntax> {
         self.pat_map_back.get(pat).cloned().ok_or(SyntheticSyntax)
     }
 
@@ -567,7 +580,7 @@ impl BodySourceMap {
         self.pat_field_map_back[&pat]
     }
 
-    pub fn macro_expansion_expr(&self, node: InFile<&ast::MacroExpr>) -> Option<ExprId> {
+    pub fn macro_expansion_expr(&self, node: InFile<&ast::MacroExpr>) -> Option<ExprOrPatId> {
         let src = node.map(AstPtr::new).map(AstPtr::upcast::<ast::MacroExpr>).map(AstPtr::upcast);
         self.expr_map.get(&src).copied()
     }
@@ -583,16 +596,20 @@ impl BodySourceMap {
         node: InFile<&ast::FormatArgsExpr>,
     ) -> Option<&[(syntax::TextRange, Name)]> {
         let src = node.map(AstPtr::new).map(AstPtr::upcast::<ast::Expr>);
-        self.template_map.as_ref()?.0.get(self.expr_map.get(&src)?).map(std::ops::Deref::deref)
+        self.template_map
+            .as_ref()?
+            .0
+            .get(&self.expr_map.get(&src)?.as_expr()?)
+            .map(std::ops::Deref::deref)
     }
 
     pub fn asm_template_args(
         &self,
         node: InFile<&ast::AsmExpr>,
     ) -> Option<(ExprId, &[Vec<(syntax::TextRange, usize)>])> {
         let src = node.map(AstPtr::new).map(AstPtr::upcast::<ast::Expr>);
-        let expr = self.expr_map.get(&src)?;
-        Some(*expr).zip(self.template_map.as_ref()?.1.get(expr).map(std::ops::Deref::deref))
+        let expr = self.expr_map.get(&src)?.as_expr()?;
+        Some(expr).zip(self.template_map.as_ref()?.1.get(&expr).map(std::ops::Deref::deref))
     }
 
     /// Get a reference to the body source map's diagnostics.

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

@@ -424,7 +424,7 @@ impl ExprCollector<'_> {
                 let inner = self.collect_expr_opt(e.expr());
                 // make the paren expr point to the inner expression as well for IDE resolution
                 let src = self.expander.in_file(syntax_ptr);
-                self.source_map.expr_map.insert(src, inner);
+                self.source_map.expr_map.insert(src, inner.into());
                 inner
             }
             ast::Expr::ReturnExpr(e) => {
@@ -660,7 +660,7 @@ impl ExprCollector<'_> {
                         // Make the macro-call point to its expanded expression so we can query
                         // semantics on syntax pointers to the macro
                         let src = self.expander.in_file(syntax_ptr);
-                        self.source_map.expr_map.insert(src, id);
+                        self.source_map.expr_map.insert(src, id.into());
                         id
                     }
                     None => self.alloc_expr(Expr::Missing, syntax_ptr),
@@ -686,9 +686,12 @@ impl ExprCollector<'_> {
 
     fn collect_expr_as_pat(&mut self, expr: ast::Expr) -> PatId {
         self.maybe_collect_expr_as_pat(&expr).unwrap_or_else(|| {
-            let syntax_ptr = AstPtr::new(&expr);
+            let src = self.expander.in_file(AstPtr::new(&expr).wrap_left());
             let expr = self.collect_expr(expr);
-            self.alloc_pat_from_expr(Pat::Expr(expr), syntax_ptr)
+            // Do not use `alloc_pat_from_expr()` here, it will override the entry in `expr_map`.
+            let id = self.body.pats.alloc(Pat::Expr(expr));
+            self.source_map.pat_map_back.insert(id, src);
+            id
         })
     }
 
@@ -743,16 +746,12 @@ impl ExprCollector<'_> {
             ast::Expr::MacroExpr(e) => {
                 let e = e.macro_call()?;
                 let macro_ptr = AstPtr::new(&e);
+                let src = self.expander.in_file(AstPtr::new(expr));
                 let id = self.collect_macro_call(e, macro_ptr, true, |this, expansion| {
-                    expansion.map(|it| this.collect_expr_as_pat(it))
+                    this.collect_expr_as_pat_opt(expansion)
                 });
-                match id {
-                    Some(id) => {
-                        // FIXME: Insert pat into source map.
-                        id
-                    }
-                    None => self.alloc_pat_from_expr(Pat::Missing, syntax_ptr),
-                }
+                self.source_map.expr_map.insert(src, id.into());
+                id
             }
             ast::Expr::RecordExpr(e) => {
                 let path =
@@ -767,9 +766,8 @@ impl ExprCollector<'_> {
                         let field_expr = f.expr()?;
                         let pat = self.collect_expr_as_pat(field_expr);
                         let name = f.field_name()?.as_name();
-                        // FIXME: Enable this.
-                        // let src = self.expander.in_file(AstPtr::new(&f));
-                        // self.source_map.pat_field_map_back.insert(pat, src);
+                        let src = self.expander.in_file(AstPtr::new(&f).wrap_left());
+                        self.source_map.pat_field_map_back.insert(pat, src);
                         Some(RecordFieldPat { name, pat })
                     })
                     .collect();
@@ -813,7 +811,10 @@ impl ExprCollector<'_> {
                                 None => Either::Left(this.missing_pat()),
                             },
                         );
-                        // FIXME: Insert pat into source map.
+                        if let Either::Left(pat) = pat {
+                            let src = this.expander.in_file(AstPtr::new(&expr).wrap_left());
+                            this.source_map.pat_map_back.insert(pat, src);
+                        }
                         pat
                     }
                     None => {
@@ -1236,7 +1237,7 @@ impl ExprCollector<'_> {
             // Make the macro-call point to its expanded expression so we can query
             // semantics on syntax pointers to the macro
             let src = self.expander.in_file(syntax_ptr);
-            self.source_map.expr_map.insert(src, tail);
+            self.source_map.expr_map.insert(src, tail.into());
         })
     }
 
@@ -1500,7 +1501,7 @@ impl ExprCollector<'_> {
                         let ast_pat = f.pat()?;
                         let pat = self.collect_pat(ast_pat, binding_list);
                         let name = f.field_name()?.as_name();
-                        let src = self.expander.in_file(AstPtr::new(&f));
+                        let src = self.expander.in_file(AstPtr::new(&f).wrap_right());
                         self.source_map.pat_field_map_back.insert(pat, src);
                         Some(RecordFieldPat { name, pat })
                     })
@@ -2187,7 +2188,7 @@ impl ExprCollector<'_> {
         let src = self.expander.in_file(ptr);
         let id = self.body.exprs.alloc(expr);
         self.source_map.expr_map_back.insert(id, src);
-        self.source_map.expr_map.insert(src, id);
+        self.source_map.expr_map.insert(src, id.into());
         id
     }
     // FIXME: desugared exprs don't have ptr, that's wrong and should be fixed.
@@ -2215,14 +2216,17 @@ impl ExprCollector<'_> {
         binding
     }
 
-    fn alloc_pat_from_expr(&mut self, pat: Pat, _ptr: ExprPtr) -> PatId {
-        // FIXME: Insert into source map.
-        self.body.pats.alloc(pat)
+    fn alloc_pat_from_expr(&mut self, pat: Pat, ptr: ExprPtr) -> PatId {
+        let src = self.expander.in_file(ptr);
+        let id = self.body.pats.alloc(pat);
+        self.source_map.expr_map.insert(src, id.into());
+        self.source_map.pat_map_back.insert(id, src.map(AstPtr::wrap_left));
+        id
     }
     fn alloc_pat(&mut self, pat: Pat, ptr: PatPtr) -> PatId {
         let src = self.expander.in_file(ptr);
         let id = self.body.pats.alloc(pat);
-        self.source_map.pat_map_back.insert(id, src);
+        self.source_map.pat_map_back.insert(id, src.map(AstPtr::wrap_right));
         self.source_map.pat_map.insert(src, id);
         id
     }

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

@@ -333,6 +333,8 @@ mod tests {
 
         let expr_id = source_map
             .node_expr(InFile { file_id: file_id.into(), value: &marker.into() })
+            .unwrap()
+            .as_expr()
             .unwrap();
         let scope = scopes.scope_for(expr_id);
 
@@ -488,8 +490,11 @@ fn foo() {
 
         let expr_scope = {
             let expr_ast = name_ref.syntax().ancestors().find_map(ast::Expr::cast).unwrap();
-            let expr_id =
-                source_map.node_expr(InFile { file_id: file_id.into(), value: &expr_ast }).unwrap();
+            let expr_id = source_map
+                .node_expr(InFile { file_id: file_id.into(), value: &expr_ast })
+                .unwrap()
+                .as_expr()
+                .unwrap();
             scopes.scope_for(expr_id).unwrap()
         };
 

crates/hir-def/src/hir.rs

@@ -48,6 +48,22 @@ pub enum ExprOrPatId {
     ExprId(ExprId),
     PatId(PatId),
 }
+
+impl ExprOrPatId {
+    pub fn as_expr(self) -> Option<ExprId> {
+        match self {
+            Self::ExprId(v) => Some(v),
+            _ => None,
+        }
+    }
+
+    pub fn as_pat(self) -> Option<PatId> {
+        match self {
+            Self::PatId(v) => Some(v),
+            _ => None,
+        }
+    }
+}
 stdx::impl_from!(ExprId, PatId for ExprOrPatId);
 
 #[derive(Debug, Clone, Eq, PartialEq)]

crates/hir-def/src/test_db.rs

@@ -198,7 +198,10 @@ impl TestDB {
             .filter_map(|node| {
                 let block = ast::BlockExpr::cast(node)?;
                 let expr = ast::Expr::from(block);
-                let expr_id = source_map.node_expr(InFile::new(position.file_id.into(), &expr))?;
+                let expr_id = source_map
+                    .node_expr(InFile::new(position.file_id.into(), &expr))?
+                    .as_expr()
+                    .unwrap();
                 let scope = scopes.scope_for(expr_id).unwrap();
                 Some(scope)
             });