Merge pull request #18254 from ChayimFriedman2/fix-mut

fix: Nail destructuring assignment once and for all

Author: Veykril

src/tools/rust-analyzer/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,7 +23,8 @@ use crate::{
     db::DefDatabase,
     expander::Expander,
     hir::{
-        dummy_expr_id, Binding, BindingId, Expr, ExprId, Label, LabelId, Pat, PatId, RecordFieldPat,
+        dummy_expr_id, Array, AsmOperand, Binding, BindingId, Expr, ExprId, ExprOrPatId, Label,
+        LabelId, Pat, PatId, RecordFieldPat, Statement,
     },
     item_tree::AttrOwner,
     nameres::DefMap,
@@ -67,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
@@ -83,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>,
@@ -286,7 +293,8 @@ impl Body {
             | Pat::Path(..)
             | Pat::ConstBlock(..)
             | Pat::Wild
-            | Pat::Missing => {}
+            | Pat::Missing
+            | Pat::Expr(_) => {}
             &Pat::Bind { subpat, .. } => {
                 if let Some(subpat) = subpat {
                     f(subpat);
@@ -322,6 +330,143 @@ impl Body {
             None => true,
         }
     }
+
+    pub fn walk_child_exprs(&self, expr_id: ExprId, mut f: impl FnMut(ExprId)) {
+        let expr = &self[expr_id];
+        match expr {
+            Expr::Continue { .. }
+            | Expr::Const(_)
+            | Expr::Missing
+            | Expr::Path(_)
+            | Expr::OffsetOf(_)
+            | Expr::Literal(_)
+            | Expr::Underscore => {}
+            Expr::InlineAsm(it) => it.operands.iter().for_each(|(_, op)| match op {
+                AsmOperand::In { expr, .. }
+                | AsmOperand::Out { expr: Some(expr), .. }
+                | AsmOperand::InOut { expr, .. } => f(*expr),
+                AsmOperand::SplitInOut { in_expr, out_expr, .. } => {
+                    f(*in_expr);
+                    if let Some(out_expr) = out_expr {
+                        f(*out_expr);
+                    }
+                }
+                AsmOperand::Out { expr: None, .. }
+                | AsmOperand::Const(_)
+                | AsmOperand::Label(_)
+                | AsmOperand::Sym(_) => (),
+            }),
+            Expr::If { condition, then_branch, else_branch } => {
+                f(*condition);
+                f(*then_branch);
+                if let &Some(else_branch) = else_branch {
+                    f(else_branch);
+                }
+            }
+            Expr::Let { expr, .. } => {
+                f(*expr);
+            }
+            Expr::Block { statements, tail, .. }
+            | Expr::Unsafe { statements, tail, .. }
+            | Expr::Async { statements, tail, .. } => {
+                for stmt in statements.iter() {
+                    match stmt {
+                        Statement::Let { initializer, else_branch, pat, .. } => {
+                            if let &Some(expr) = initializer {
+                                f(expr);
+                            }
+                            if let &Some(expr) = else_branch {
+                                f(expr);
+                            }
+                            self.walk_exprs_in_pat(*pat, &mut f);
+                        }
+                        Statement::Expr { expr: expression, .. } => f(*expression),
+                        Statement::Item => (),
+                    }
+                }
+                if let &Some(expr) = tail {
+                    f(expr);
+                }
+            }
+            Expr::Loop { body, .. } => f(*body),
+            Expr::Call { callee, args, .. } => {
+                f(*callee);
+                args.iter().copied().for_each(f);
+            }
+            Expr::MethodCall { receiver, args, .. } => {
+                f(*receiver);
+                args.iter().copied().for_each(f);
+            }
+            Expr::Match { expr, arms } => {
+                f(*expr);
+                arms.iter().map(|arm| arm.expr).for_each(f);
+            }
+            Expr::Break { expr, .. }
+            | Expr::Return { expr }
+            | Expr::Yield { expr }
+            | Expr::Yeet { expr } => {
+                if let &Some(expr) = expr {
+                    f(expr);
+                }
+            }
+            Expr::Become { expr } => f(*expr),
+            Expr::RecordLit { fields, spread, .. } => {
+                for field in fields.iter() {
+                    f(field.expr);
+                }
+                if let &Some(expr) = spread {
+                    f(expr);
+                }
+            }
+            Expr::Closure { body, .. } => {
+                f(*body);
+            }
+            Expr::BinaryOp { lhs, rhs, .. } => {
+                f(*lhs);
+                f(*rhs);
+            }
+            Expr::Range { lhs, rhs, .. } => {
+                if let &Some(lhs) = rhs {
+                    f(lhs);
+                }
+                if let &Some(rhs) = lhs {
+                    f(rhs);
+                }
+            }
+            Expr::Index { base, index, .. } => {
+                f(*base);
+                f(*index);
+            }
+            Expr::Field { expr, .. }
+            | Expr::Await { expr }
+            | Expr::Cast { expr, .. }
+            | Expr::Ref { expr, .. }
+            | Expr::UnaryOp { expr, .. }
+            | Expr::Box { expr } => {
+                f(*expr);
+            }
+            Expr::Tuple { exprs, .. } => exprs.iter().copied().for_each(f),
+            Expr::Array(a) => match a {
+                Array::ElementList { elements, .. } => elements.iter().copied().for_each(f),
+                Array::Repeat { initializer, repeat } => {
+                    f(*initializer);
+                    f(*repeat)
+                }
+            },
+            &Expr::Assignment { target, value } => {
+                self.walk_exprs_in_pat(target, &mut f);
+                f(value);
+            }
+        }
+    }
+
+    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);
+            }
+        });
+    }
 }
 
 impl Default for Body {
@@ -375,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()
     }
@@ -395,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)
     }
 
@@ -428,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()
     }
@@ -444,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.