From 32bf7af83e3d6f201ebea87006f9c33bbc211f7b Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Fri, 4 Mar 2022 19:49:08 +0100
Subject: [PATCH 1/2] Support locals with multiple declaration sites

---
 crates/hir/src/lib.rs                         |  25 ++-
 crates/hir_def/src/body.rs                    |  18 +-
 crates/hir_def/src/body/lower.rs              |  84 ++++++--
 crates/ide/src/highlight_related.rs           | 201 +++++++++++++-----
 crates/ide/src/rename.rs                      |  49 +++++
 .../src/handlers/inline_local_variable.rs     |   1 +
 crates/ide_db/src/rename.rs                   |  14 +-
 crates/ide_db/src/search.rs                   |  35 ++-
 8 files changed, 346 insertions(+), 81 deletions(-)

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 45a4c6a50e..60e3548d49 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2036,6 +2036,11 @@ impl GenericDef {
     }
 }
 
+/// A single local definition.
+///
+/// If the definition of this is part of a "MultiLocal", that is a local that has multiple declarations due to or-patterns
+/// then this only references a single one of those.
+/// To retrieve the other locals you should use [`Local::associated_locals`]
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub struct Local {
     pub(crate) parent: DefWithBodyId,
@@ -2107,12 +2112,28 @@ impl Local {
         Type::new(db, krate, def, ty)
     }
 
+    pub fn associated_locals(self, db: &dyn HirDatabase) -> Box<[Local]> {
+        let body = db.body(self.parent);
+        body.ident_patterns_for(&self.pat_id)
+            .iter()
+            .map(|&pat_id| Local { parent: self.parent, pat_id })
+            .collect()
+    }
+
+    /// If this local is part of a multi-local, retrieve the representative local.
+    /// That is the local that references are being resolved to.
+    pub fn representative(self, db: &dyn HirDatabase) -> Local {
+        let body = db.body(self.parent);
+        Local { pat_id: body.pattern_representative(self.pat_id), ..self }
+    }
+
     pub fn source(self, db: &dyn HirDatabase) -> InFile<Either<ast::IdentPat, ast::SelfParam>> {
         let (_body, source_map) = db.body_with_source_map(self.parent);
         let src = source_map.pat_syntax(self.pat_id).unwrap(); // Hmm...
         let root = src.file_syntax(db.upcast());
-        src.map(|ast| {
-            ast.map_left(|it| it.cast().unwrap().to_node(&root)).map_right(|it| it.to_node(&root))
+        src.map(|ast| match ast {
+            Either::Left(it) => Either::Left(it.cast().unwrap().to_node(&root)),
+            Either::Right(it) => Either::Right(it.to_node(&root)),
         })
     }
 }
diff --git a/crates/hir_def/src/body.rs b/crates/hir_def/src/body.rs
index a2f64eda06..1f5e45eb2d 100644
--- a/crates/hir_def/src/body.rs
+++ b/crates/hir_def/src/body.rs
@@ -237,10 +237,11 @@ pub struct Mark {
 }
 
 /// The body of an item (function, const etc.).
-#[derive(Debug, Eq, PartialEq)]
+#[derive(Debug, PartialEq, Eq)]
 pub struct Body {
     pub exprs: Arena<Expr>,
     pub pats: Arena<Pat>,
+    pub or_pats: FxHashMap<PatId, Arc<[PatId]>>,
     pub labels: Arena<Label>,
     /// The patterns for the function's parameters. While the parameter types are
     /// part of the function signature, the patterns are not (they don't change
@@ -355,6 +356,18 @@ impl Body {
             .map(move |block| (*block, db.block_def_map(*block).expect("block ID without DefMap")))
     }
 
+    pub fn pattern_representative(&self, pat: PatId) -> PatId {
+        self.or_pats.get(&pat).and_then(|pats| pats.first().copied()).unwrap_or(pat)
+    }
+
+    /// Retrieves all ident patterns this pattern shares the ident with.
+    pub fn ident_patterns_for<'slf>(&'slf self, pat: &'slf PatId) -> &'slf [PatId] {
+        match self.or_pats.get(pat) {
+            Some(pats) => &**pats,
+            None => std::slice::from_ref(pat),
+        }
+    }
+
     fn new(
         db: &dyn DefDatabase,
         expander: Expander,
@@ -365,8 +378,9 @@ impl Body {
     }
 
     fn shrink_to_fit(&mut self) {
-        let Self { _c: _, body_expr: _, block_scopes, exprs, labels, params, pats } = self;
+        let Self { _c: _, body_expr: _, block_scopes, or_pats, exprs, labels, params, pats } = self;
         block_scopes.shrink_to_fit();
+        or_pats.shrink_to_fit();
         exprs.shrink_to_fit();
         labels.shrink_to_fit();
         params.shrink_to_fit();
diff --git a/crates/hir_def/src/body/lower.rs b/crates/hir_def/src/body/lower.rs
index 06ad7ce4cd..873ae285ca 100644
--- a/crates/hir_def/src/body/lower.rs
+++ b/crates/hir_def/src/body/lower.rs
@@ -12,6 +12,7 @@ use hir_expand::{
 };
 use la_arena::Arena;
 use profile::Count;
+use rustc_hash::FxHashMap;
 use syntax::{
     ast::{
         self, ArrayExprKind, AstChildren, HasArgList, HasLoopBody, HasName, LiteralKind,
@@ -92,6 +93,7 @@ pub(super) fn lower(
             body_expr: dummy_expr_id(),
             block_scopes: Vec::new(),
             _c: Count::new(),
+            or_pats: Default::default(),
         },
         expander,
         statements_in_scope: Vec::new(),
@@ -704,13 +706,40 @@ impl ExprCollector<'_> {
     }
 
     fn collect_pat(&mut self, pat: ast::Pat) -> PatId {
+        let mut name_to_pat = FxHashMap::default();
+        let pat_id = self.collect_pat_(&mut name_to_pat, false, pat);
+        for pats in name_to_pat.into_values() {
+            let pats = Arc::<[_]>::from(pats);
+            for &pat in &*pats {
+                self.body.or_pats.insert(pat, pats.clone());
+            }
+        }
+        pat_id
+    }
+
+    fn collect_pat_opt(&mut self, pat: Option<ast::Pat>) -> PatId {
+        match pat {
+            Some(pat) => self.collect_pat(pat),
+            None => self.missing_pat(),
+        }
+    }
+
+    fn collect_pat_(
+        &mut self,
+        name_to_pat: &mut FxHashMap<Name, Vec<PatId>>,
+        in_or_pat: bool,
+        pat: ast::Pat,
+    ) -> PatId {
         let pattern = match &pat {
             ast::Pat::IdentPat(bp) => {
                 let name = bp.name().map(|nr| nr.as_name()).unwrap_or_else(Name::missing);
+
+                let key = in_or_pat.then(|| name.clone());
                 let annotation =
                     BindingAnnotation::new(bp.mut_token().is_some(), bp.ref_token().is_some());
-                let subpat = bp.pat().map(|subpat| self.collect_pat(subpat));
-                if annotation == BindingAnnotation::Unannotated && subpat.is_none() {
+                let subpat =
+                    bp.pat().map(|subpat| self.collect_pat_(name_to_pat, in_or_pat, subpat));
+                let pattern = if annotation == BindingAnnotation::Unannotated && subpat.is_none() {
                     // This could also be a single-segment path pattern. To
                     // decide that, we need to try resolving the name.
                     let (resolved, _) = self.expander.def_map.resolve_path(
@@ -740,12 +769,19 @@ impl ExprCollector<'_> {
                     }
                 } else {
                     Pat::Bind { name, mode: annotation, subpat }
+                };
+
+                let ptr = AstPtr::new(&pat);
+                let pat = self.alloc_pat(pattern, Either::Left(ptr));
+                if let Some(key) = key {
+                    name_to_pat.entry(key).or_default().push(pat);
                 }
+                return pat;
             }
             ast::Pat::TupleStructPat(p) => {
                 let path =
                     p.path().and_then(|path| self.expander.parse_path(self.db, path)).map(Box::new);
-                let (args, ellipsis) = self.collect_tuple_pat(p.fields());
+                let (args, ellipsis) = self.collect_tuple_pat(name_to_pat, in_or_pat, p.fields());
                 Pat::TupleStruct { path, args, ellipsis }
             }
             ast::Pat::RefPat(p) => {
@@ -759,12 +795,12 @@ impl ExprCollector<'_> {
                 path.map(Pat::Path).unwrap_or(Pat::Missing)
             }
             ast::Pat::OrPat(p) => {
-                let pats = p.pats().map(|p| self.collect_pat(p)).collect();
+                let pats = p.pats().map(|p| self.collect_pat_(name_to_pat, true, p)).collect();
                 Pat::Or(pats)
             }
-            ast::Pat::ParenPat(p) => return self.collect_pat_opt(p.pat()),
+            ast::Pat::ParenPat(p) => return self.collect_pat_opt_(name_to_pat, in_or_pat, p.pat()),
             ast::Pat::TuplePat(p) => {
-                let (args, ellipsis) = self.collect_tuple_pat(p.fields());
+                let (args, ellipsis) = self.collect_tuple_pat(name_to_pat, in_or_pat, p.fields());
                 Pat::Tuple { args, ellipsis }
             }
             ast::Pat::WildcardPat(_) => Pat::Wild,
@@ -777,7 +813,7 @@ impl ExprCollector<'_> {
                     .fields()
                     .filter_map(|f| {
                         let ast_pat = f.pat()?;
-                        let pat = self.collect_pat(ast_pat);
+                        let pat = self.collect_pat_(name_to_pat, in_or_pat, ast_pat);
                         let name = f.field_name()?.as_name();
                         Some(RecordFieldPat { name, pat })
                     })
@@ -796,9 +832,15 @@ impl ExprCollector<'_> {
 
                 // FIXME properly handle `RestPat`
                 Pat::Slice {
-                    prefix: prefix.into_iter().map(|p| self.collect_pat(p)).collect(),
-                    slice: slice.map(|p| self.collect_pat(p)),
-                    suffix: suffix.into_iter().map(|p| self.collect_pat(p)).collect(),
+                    prefix: prefix
+                        .into_iter()
+                        .map(|p| self.collect_pat_(name_to_pat, in_or_pat, p))
+                        .collect(),
+                    slice: slice.map(|p| self.collect_pat_(name_to_pat, in_or_pat, p)),
+                    suffix: suffix
+                        .into_iter()
+                        .map(|p| self.collect_pat_(name_to_pat, in_or_pat, p))
+                        .collect(),
                 }
             }
             ast::Pat::LiteralPat(lit) => {
@@ -821,7 +863,7 @@ impl ExprCollector<'_> {
                 Pat::Missing
             }
             ast::Pat::BoxPat(boxpat) => {
-                let inner = self.collect_pat_opt(boxpat.pat());
+                let inner = self.collect_pat_opt_(name_to_pat, in_or_pat, boxpat.pat());
                 Pat::Box { inner }
             }
             ast::Pat::ConstBlockPat(const_block_pat) => {
@@ -837,7 +879,7 @@ impl ExprCollector<'_> {
                     let macro_ptr = AstPtr::new(&call);
                     let mut pat = None;
                     self.collect_macro_call(call, macro_ptr, true, |this, expanded_pat| {
-                        pat = Some(this.collect_pat_opt(expanded_pat));
+                        pat = Some(this.collect_pat_opt_(name_to_pat, in_or_pat, expanded_pat));
                     });
 
                     match pat {
@@ -854,21 +896,31 @@ impl ExprCollector<'_> {
         self.alloc_pat(pattern, Either::Left(ptr))
     }
 
-    fn collect_pat_opt(&mut self, pat: Option<ast::Pat>) -> PatId {
+    fn collect_pat_opt_(
+        &mut self,
+        name_to_pat: &mut FxHashMap<Name, Vec<PatId>>,
+        in_or_pat: bool,
+        pat: Option<ast::Pat>,
+    ) -> PatId {
         match pat {
-            Some(pat) => self.collect_pat(pat),
+            Some(pat) => self.collect_pat_(name_to_pat, in_or_pat, pat),
             None => self.missing_pat(),
         }
     }
 
-    fn collect_tuple_pat(&mut self, args: AstChildren<ast::Pat>) -> (Box<[PatId]>, Option<usize>) {
+    fn collect_tuple_pat(
+        &mut self,
+        name_to_pat: &mut FxHashMap<Name, Vec<PatId>>,
+        in_or_pat: bool,
+        args: AstChildren<ast::Pat>,
+    ) -> (Box<[PatId]>, Option<usize>) {
         // Find the location of the `..`, if there is one. Note that we do not
         // consider the possibility of there being multiple `..` here.
         let ellipsis = args.clone().position(|p| matches!(p, ast::Pat::RestPat(_)));
         // We want to skip the `..` pattern here, since we account for it above.
         let args = args
             .filter(|p| !matches!(p, ast::Pat::RestPat(_)))
-            .map(|p| self.collect_pat(p))
+            .map(|p| self.collect_pat_(name_to_pat, in_or_pat, p))
             .collect();
 
         (args, ellipsis)
diff --git a/crates/ide/src/highlight_related.rs b/crates/ide/src/highlight_related.rs
index f886ff7837..3690fe11d0 100644
--- a/crates/ide/src/highlight_related.rs
+++ b/crates/ide/src/highlight_related.rs
@@ -98,24 +98,37 @@ fn highlight_references(
             range,
             category: access,
         });
+    let mut res = FxHashSet::default();
 
-    let declarations = defs.iter().flat_map(|def| {
-        match def {
-            &Definition::Module(module) => {
+    let mut def_to_hl_range = |def| {
+        let hl_range = match def {
+            Definition::Module(module) => {
                 Some(NavigationTarget::from_module_to_decl(sema.db, module))
             }
             def => def.try_to_nav(sema.db),
         }
         .filter(|decl| decl.file_id == file_id)
-        .and_then(|decl| {
-            let range = decl.focus_range?;
+        .and_then(|decl| decl.focus_range)
+        .map(|range| {
             let category =
                 references::decl_mutability(&def, node, range).then(|| ReferenceCategory::Write);
-            Some(HighlightedRange { range, category })
-        })
-    });
+            HighlightedRange { range, category }
+        });
+        if let Some(hl_range) = hl_range {
+            res.insert(hl_range);
+        }
+    };
+    for &def in &defs {
+        match def {
+            Definition::Local(local) => local
+                .associated_locals(sema.db)
+                .iter()
+                .for_each(|&local| def_to_hl_range(Definition::Local(local))),
+            def => def_to_hl_range(def),
+        }
+    }
 
-    let res: FxHashSet<_> = declarations.chain(usages).collect();
+    res.extend(usages);
     if res.is_empty() {
         None
     } else {
@@ -332,6 +345,7 @@ mod tests {
 
     use super::*;
 
+    #[track_caller]
     fn check(ra_fixture: &str) {
         let config = HighlightRelatedConfig {
             break_points: true,
@@ -343,6 +357,7 @@ mod tests {
         check_with_config(ra_fixture, config);
     }
 
+    #[track_caller]
     fn check_with_config(ra_fixture: &str, config: HighlightRelatedConfig) {
         let (analysis, pos, annotations) = fixture::annotations(ra_fixture);
 
@@ -1053,13 +1068,15 @@ fn function(field: u32) {
             yield_points: true,
         };
 
-        let ra_fixture = r#"
+        check_with_config(
+            r#"
 fn foo() {
-    let x$0 = 5;
+    let x = 5;
     let y = x * 2;
-}"#;
-
-        check_with_config(ra_fixture, config);
+}
+"#,
+            config,
+        );
     }
 
     #[test]
@@ -1071,31 +1088,35 @@ fn foo() {
             yield_points: true,
         };
 
-        let ra_fixture = r#"
+        check_with_config(
+            r#"
 fn foo() {
-    let x$0 = 5;
+    let x = 5;
     let y = x * 2;
 
     loop {
         break;
     }
-}"#;
+}
+"#,
+            config.clone(),
+        );
 
-        check_with_config(ra_fixture, config.clone());
-
-        let ra_fixture = r#"
+        check_with_config(
+            r#"
 fn foo() {
     let x = 5;
     let y = x * 2;
 
-    loop$0 {
+    loop {
 //  ^^^^
         break;
 //      ^^^^^
     }
-}"#;
-
-        check_with_config(ra_fixture, config);
+}
+"#,
+            config,
+        );
     }
 
     #[test]
@@ -1107,27 +1128,31 @@ fn foo() {
             yield_points: true,
         };
 
-        let ra_fixture = r#"
+        check_with_config(
+            r#"
 async fn foo() {
-    let x$0 = 5;
+    let x = 5;
     let y = x * 2;
 
     0.await;
-}"#;
+}
+"#,
+            config.clone(),
+        );
 
-        check_with_config(ra_fixture, config.clone());
-
-        let ra_fixture = r#"
+        check_with_config(
+            r#"
     async fn foo() {
 //  ^^^^^
         let x = 5;
         let y = x * 2;
 
-        0.await$0;
+        0.await;
 //        ^^^^^
-}"#;
-
-        check_with_config(ra_fixture, config);
+}
+"#,
+            config,
+        );
     }
 
     #[test]
@@ -1139,9 +1164,10 @@ async fn foo() {
             yield_points: true,
         };
 
-        let ra_fixture = r#"
+        check_with_config(
+            r#"
 fn foo() -> i32 {
-    let x$0 = 5;
+    let x = 5;
     let y = x * 2;
 
     if true {
@@ -1149,12 +1175,14 @@ fn foo() -> i32 {
     }
 
     0?
-}"#;
+}
+"#,
+            config.clone(),
+        );
 
-        check_with_config(ra_fixture, config.clone());
-
-        let ra_fixture = r#"
-fn foo() ->$0 i32 {
+        check_with_config(
+            r#"
+fn foo() -> i32 {
     let x = 5;
     let y = x * 2;
 
@@ -1165,9 +1193,9 @@ fn foo() ->$0 i32 {
 
     0?
 //   ^
-"#;
-
-        check_with_config(ra_fixture, config);
+"#,
+            config,
+        );
     }
 
     #[test]
@@ -1179,14 +1207,16 @@ fn foo() ->$0 i32 {
             yield_points: true,
         };
 
-        let ra_fixture = r#"
+        check_with_config(
+            r#"
 fn foo() {
     loop {
-        break$0;
+        break;
     }
-}"#;
-
-        check_with_config(ra_fixture, config);
+}
+"#,
+            config,
+        );
     }
 
     #[test]
@@ -1198,12 +1228,14 @@ fn foo() {
             yield_points: false,
         };
 
-        let ra_fixture = r#"
-async$0 fn foo() {
+        check_with_config(
+            r#"
+async fn foo() {
     0.await;
-}"#;
-
-        check_with_config(ra_fixture, config);
+}
+"#,
+            config,
+        );
     }
 
     #[test]
@@ -1215,15 +1247,68 @@ async$0 fn foo() {
             yield_points: true,
         };
 
-        let ra_fixture = r#"
-fn foo() ->$0 i32 {
+        check_with_config(
+            r#"
+fn foo() -> i32 {
     if true {
         return -1;
     }
 
     42
-}"#;
+}"#,
+            config,
+        );
+    }
 
-        check_with_config(ra_fixture, config);
+    #[test]
+    fn test_hl_multi_local() {
+        check(
+            r#"
+fn foo((
+    foo$0
+  //^^^
+    | foo
+    //^^^
+    | foo
+    //^^^
+): ()) {
+    foo;
+  //^^^read
+    let foo;
+}
+"#,
+        );
+        check(
+            r#"
+fn foo((
+    foo
+  //^^^
+    | foo$0
+    //^^^
+    | foo
+    //^^^
+): ()) {
+    foo;
+  //^^^read
+    let foo;
+}
+"#,
+        );
+        check(
+            r#"
+fn foo((
+    foo
+  //^^^
+    | foo
+    //^^^
+    | foo
+    //^^^
+): ()) {
+    foo$0;
+  //^^^read
+    let foo;
+}
+"#,
+        );
     }
 }
diff --git a/crates/ide/src/rename.rs b/crates/ide/src/rename.rs
index daba317e42..83bc299adc 100644
--- a/crates/ide/src/rename.rs
+++ b/crates/ide/src/rename.rs
@@ -2084,4 +2084,53 @@ fn foo() {
 "#,
         )
     }
+
+    #[test]
+    fn rename_multi_local() {
+        check(
+            "bar",
+            r#"
+fn foo((foo$0 | foo | foo): ()) {
+    foo;
+    let foo;
+}
+"#,
+            r#"
+fn foo((bar | bar | bar): ()) {
+    bar;
+    let foo;
+}
+"#,
+        );
+        check(
+            "bar",
+            r#"
+fn foo((foo | foo$0 | foo): ()) {
+    foo;
+    let foo;
+}
+"#,
+            r#"
+fn foo((bar | bar | bar): ()) {
+    bar;
+    let foo;
+}
+"#,
+        );
+        check(
+            "bar",
+            r#"
+fn foo((foo | foo | foo): ()) {
+    foo$0;
+    let foo;
+}
+"#,
+            r#"
+fn foo((bar | bar | bar): ()) {
+    bar;
+    let foo;
+}
+"#,
+        );
+    }
 }
diff --git a/crates/ide_assists/src/handlers/inline_local_variable.rs b/crates/ide_assists/src/handlers/inline_local_variable.rs
index 1fce654df2..0b4f46caf4 100644
--- a/crates/ide_assists/src/handlers/inline_local_variable.rs
+++ b/crates/ide_assists/src/handlers/inline_local_variable.rs
@@ -206,6 +206,7 @@ fn inline_usage(
         return None;
     }
 
+    // FIXME: Handle multiple local definitions
     let bind_pat = match local.source(sema.db).value {
         Either::Left(ident) => ident,
         _ => return None,
diff --git a/crates/ide_db/src/rename.rs b/crates/ide_db/src/rename.rs
index 970ca2b6d7..a44fe04e74 100644
--- a/crates/ide_db/src/rename.rs
+++ b/crates/ide_db/src/rename.rs
@@ -293,8 +293,18 @@ fn rename_reference(
         (file_id, source_edit_from_references(references, def, new_name))
     }));
 
-    let (file_id, edit) = source_edit_from_def(sema, def, new_name)?;
-    source_change.insert_source_edit(file_id, edit);
+    let mut insert_def_edit = |def| {
+        let (file_id, edit) = source_edit_from_def(sema, def, new_name)?;
+        source_change.insert_source_edit(file_id, edit);
+        Ok(())
+    };
+    match def {
+        Definition::Local(l) => l
+            .associated_locals(sema.db)
+            .iter()
+            .try_for_each(|&local| insert_def_edit(Definition::Local(local))),
+        def => insert_def_edit(def),
+    }?;
     Ok(source_change)
 }
 
diff --git a/crates/ide_db/src/search.rs b/crates/ide_db/src/search.rs
index 7ea2bc63f2..3ff48520f4 100644
--- a/crates/ide_db/src/search.rs
+++ b/crates/ide_db/src/search.rs
@@ -310,6 +310,10 @@ impl Definition {
 
     pub fn usages<'a>(self, sema: &'a Semantics<RootDatabase>) -> FindUsages<'a> {
         FindUsages {
+            local_repr: match self {
+                Definition::Local(local) => Some(local.representative(sema.db)),
+                _ => None,
+            },
             def: self,
             sema,
             scope: None,
@@ -325,6 +329,7 @@ pub struct FindUsages<'a> {
     sema: &'a Semantics<'a, RootDatabase>,
     scope: Option<SearchScope>,
     include_self_kw_refs: Option<hir::Type>,
+    local_repr: Option<hir::Local>,
     search_self_mod: bool,
 }
 
@@ -593,6 +598,19 @@ impl<'a> FindUsages<'a> {
         sink: &mut dyn FnMut(FileId, FileReference) -> bool,
     ) -> bool {
         match NameRefClass::classify(self.sema, name_ref) {
+            Some(NameRefClass::Definition(def @ Definition::Local(local)))
+                if matches!(
+                    self.local_repr, Some(repr) if repr == local.representative(self.sema.db)
+                ) =>
+            {
+                let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
+                let reference = FileReference {
+                    range,
+                    name: ast::NameLike::NameRef(name_ref.clone()),
+                    category: ReferenceCategory::new(&def, name_ref),
+                };
+                sink(file_id, reference)
+            }
             Some(NameRefClass::Definition(def)) if def == self.def => {
                 let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
                 let reference = FileReference {
@@ -622,7 +640,7 @@ impl<'a> FindUsages<'a> {
                     Definition::Field(_) if field == self.def => {
                         ReferenceCategory::new(&field, name_ref)
                     }
-                    Definition::Local(l) if local == l => {
+                    Definition::Local(_) if matches!(self.local_repr, Some(repr) if repr == local.representative(self.sema.db)) => {
                         ReferenceCategory::new(&Definition::Local(local), name_ref)
                     }
                     _ => return false,
@@ -667,6 +685,21 @@ impl<'a> FindUsages<'a> {
                 };
                 sink(file_id, reference)
             }
+            Some(NameClass::Definition(def @ Definition::Local(local))) if def != self.def => {
+                if matches!(
+                    self.local_repr,
+                    Some(repr) if local.representative(self.sema.db) == repr
+                ) {
+                    let FileRange { file_id, range } = self.sema.original_range(name.syntax());
+                    let reference = FileReference {
+                        range,
+                        name: ast::NameLike::Name(name.clone()),
+                        category: None,
+                    };
+                    return sink(file_id, reference);
+                }
+                false
+            }
             // Resolve trait impl function definitions to the trait definition's version if self.def is the trait definition's
             Some(NameClass::Definition(def)) if def != self.def => {
                 /* poor man's try block */

From 4a866fc672927b9d3dd5a7211bfe45e9eebebdc2 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Fri, 4 Mar 2022 20:09:32 +0100
Subject: [PATCH 2/2] Simplify

---
 crates/hir_def/src/body.rs          |  4 +-
 crates/hir_def/src/body/lower.rs    | 74 +++++++++++------------------
 crates/ide/src/highlight_related.rs | 20 ++++----
 3 files changed, 40 insertions(+), 58 deletions(-)

diff --git a/crates/hir_def/src/body.rs b/crates/hir_def/src/body.rs
index 1f5e45eb2d..6d3c2c2c46 100644
--- a/crates/hir_def/src/body.rs
+++ b/crates/hir_def/src/body.rs
@@ -237,7 +237,7 @@ pub struct Mark {
 }
 
 /// The body of an item (function, const etc.).
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, Eq, PartialEq)]
 pub struct Body {
     pub exprs: Arena<Expr>,
     pub pats: Arena<Pat>,
@@ -353,7 +353,7 @@ impl Body {
     ) -> impl Iterator<Item = (BlockId, Arc<DefMap>)> + '_ {
         self.block_scopes
             .iter()
-            .map(move |block| (*block, db.block_def_map(*block).expect("block ID without DefMap")))
+            .map(move |&block| (block, db.block_def_map(block).expect("block ID without DefMap")))
     }
 
     pub fn pattern_representative(&self, pat: PatId) -> PatId {
diff --git a/crates/hir_def/src/body/lower.rs b/crates/hir_def/src/body/lower.rs
index 873ae285ca..46b2ba8a25 100644
--- a/crates/hir_def/src/body/lower.rs
+++ b/crates/hir_def/src/body/lower.rs
@@ -97,6 +97,8 @@ pub(super) fn lower(
         },
         expander,
         statements_in_scope: Vec::new(),
+        name_to_pat_grouping: Default::default(),
+        is_lowering_inside_or_pat: false,
     }
     .collect(params, body)
 }
@@ -107,6 +109,9 @@ struct ExprCollector<'a> {
     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,
 }
 
 impl ExprCollector<'_> {
@@ -706,14 +711,12 @@ impl ExprCollector<'_> {
     }
 
     fn collect_pat(&mut self, pat: ast::Pat) -> PatId {
-        let mut name_to_pat = FxHashMap::default();
-        let pat_id = self.collect_pat_(&mut name_to_pat, false, pat);
-        for pats in name_to_pat.into_values() {
+        let pat_id = self.collect_pat_(pat);
+        for (_, pats) in self.name_to_pat_grouping.drain() {
             let pats = Arc::<[_]>::from(pats);
-            for &pat in &*pats {
-                self.body.or_pats.insert(pat, pats.clone());
-            }
+            self.body.or_pats.extend(pats.iter().map(|&pat| (pat, pats.clone())));
         }
+        self.is_lowering_inside_or_pat = false;
         pat_id
     }
 
@@ -724,21 +727,15 @@ impl ExprCollector<'_> {
         }
     }
 
-    fn collect_pat_(
-        &mut self,
-        name_to_pat: &mut FxHashMap<Name, Vec<PatId>>,
-        in_or_pat: bool,
-        pat: ast::Pat,
-    ) -> PatId {
+    fn collect_pat_(&mut self, pat: ast::Pat) -> PatId {
         let pattern = match &pat {
             ast::Pat::IdentPat(bp) => {
                 let name = bp.name().map(|nr| nr.as_name()).unwrap_or_else(Name::missing);
 
-                let key = in_or_pat.then(|| name.clone());
+                let key = self.is_lowering_inside_or_pat.then(|| name.clone());
                 let annotation =
                     BindingAnnotation::new(bp.mut_token().is_some(), bp.ref_token().is_some());
-                let subpat =
-                    bp.pat().map(|subpat| self.collect_pat_(name_to_pat, in_or_pat, subpat));
+                let subpat = bp.pat().map(|subpat| self.collect_pat_(subpat));
                 let pattern = if annotation == BindingAnnotation::Unannotated && subpat.is_none() {
                     // This could also be a single-segment path pattern. To
                     // decide that, we need to try resolving the name.
@@ -774,14 +771,14 @@ impl ExprCollector<'_> {
                 let ptr = AstPtr::new(&pat);
                 let pat = self.alloc_pat(pattern, Either::Left(ptr));
                 if let Some(key) = key {
-                    name_to_pat.entry(key).or_default().push(pat);
+                    self.name_to_pat_grouping.entry(key).or_default().push(pat);
                 }
                 return pat;
             }
             ast::Pat::TupleStructPat(p) => {
                 let path =
                     p.path().and_then(|path| self.expander.parse_path(self.db, path)).map(Box::new);
-                let (args, ellipsis) = self.collect_tuple_pat(name_to_pat, in_or_pat, p.fields());
+                let (args, ellipsis) = self.collect_tuple_pat(p.fields());
                 Pat::TupleStruct { path, args, ellipsis }
             }
             ast::Pat::RefPat(p) => {
@@ -795,12 +792,13 @@ impl ExprCollector<'_> {
                 path.map(Pat::Path).unwrap_or(Pat::Missing)
             }
             ast::Pat::OrPat(p) => {
-                let pats = p.pats().map(|p| self.collect_pat_(name_to_pat, true, p)).collect();
+                self.is_lowering_inside_or_pat = true;
+                let pats = p.pats().map(|p| self.collect_pat_(p)).collect();
                 Pat::Or(pats)
             }
-            ast::Pat::ParenPat(p) => return self.collect_pat_opt_(name_to_pat, in_or_pat, p.pat()),
+            ast::Pat::ParenPat(p) => return self.collect_pat_opt_(p.pat()),
             ast::Pat::TuplePat(p) => {
-                let (args, ellipsis) = self.collect_tuple_pat(name_to_pat, in_or_pat, p.fields());
+                let (args, ellipsis) = self.collect_tuple_pat(p.fields());
                 Pat::Tuple { args, ellipsis }
             }
             ast::Pat::WildcardPat(_) => Pat::Wild,
@@ -813,7 +811,7 @@ impl ExprCollector<'_> {
                     .fields()
                     .filter_map(|f| {
                         let ast_pat = f.pat()?;
-                        let pat = self.collect_pat_(name_to_pat, in_or_pat, ast_pat);
+                        let pat = self.collect_pat_(ast_pat);
                         let name = f.field_name()?.as_name();
                         Some(RecordFieldPat { name, pat })
                     })
@@ -832,15 +830,9 @@ impl ExprCollector<'_> {
 
                 // FIXME properly handle `RestPat`
                 Pat::Slice {
-                    prefix: prefix
-                        .into_iter()
-                        .map(|p| self.collect_pat_(name_to_pat, in_or_pat, p))
-                        .collect(),
-                    slice: slice.map(|p| self.collect_pat_(name_to_pat, in_or_pat, p)),
-                    suffix: suffix
-                        .into_iter()
-                        .map(|p| self.collect_pat_(name_to_pat, in_or_pat, p))
-                        .collect(),
+                    prefix: prefix.into_iter().map(|p| self.collect_pat_(p)).collect(),
+                    slice: slice.map(|p| self.collect_pat_(p)),
+                    suffix: suffix.into_iter().map(|p| self.collect_pat_(p)).collect(),
                 }
             }
             ast::Pat::LiteralPat(lit) => {
@@ -863,7 +855,7 @@ impl ExprCollector<'_> {
                 Pat::Missing
             }
             ast::Pat::BoxPat(boxpat) => {
-                let inner = self.collect_pat_opt_(name_to_pat, in_or_pat, boxpat.pat());
+                let inner = self.collect_pat_opt_(boxpat.pat());
                 Pat::Box { inner }
             }
             ast::Pat::ConstBlockPat(const_block_pat) => {
@@ -879,7 +871,7 @@ impl ExprCollector<'_> {
                     let macro_ptr = AstPtr::new(&call);
                     let mut pat = None;
                     self.collect_macro_call(call, macro_ptr, true, |this, expanded_pat| {
-                        pat = Some(this.collect_pat_opt_(name_to_pat, in_or_pat, expanded_pat));
+                        pat = Some(this.collect_pat_opt_(expanded_pat));
                     });
 
                     match pat {
@@ -896,31 +888,21 @@ impl ExprCollector<'_> {
         self.alloc_pat(pattern, Either::Left(ptr))
     }
 
-    fn collect_pat_opt_(
-        &mut self,
-        name_to_pat: &mut FxHashMap<Name, Vec<PatId>>,
-        in_or_pat: bool,
-        pat: Option<ast::Pat>,
-    ) -> PatId {
+    fn collect_pat_opt_(&mut self, pat: Option<ast::Pat>) -> PatId {
         match pat {
-            Some(pat) => self.collect_pat_(name_to_pat, in_or_pat, pat),
+            Some(pat) => self.collect_pat_(pat),
             None => self.missing_pat(),
         }
     }
 
-    fn collect_tuple_pat(
-        &mut self,
-        name_to_pat: &mut FxHashMap<Name, Vec<PatId>>,
-        in_or_pat: bool,
-        args: AstChildren<ast::Pat>,
-    ) -> (Box<[PatId]>, Option<usize>) {
+    fn collect_tuple_pat(&mut self, args: AstChildren<ast::Pat>) -> (Box<[PatId]>, Option<usize>) {
         // Find the location of the `..`, if there is one. Note that we do not
         // consider the possibility of there being multiple `..` here.
         let ellipsis = args.clone().position(|p| matches!(p, ast::Pat::RestPat(_)));
         // We want to skip the `..` pattern here, since we account for it above.
         let args = args
             .filter(|p| !matches!(p, ast::Pat::RestPat(_)))
-            .map(|p| self.collect_pat_(name_to_pat, in_or_pat, p))
+            .map(|p| self.collect_pat_(p))
             .collect();
 
         (args, ellipsis)
diff --git a/crates/ide/src/highlight_related.rs b/crates/ide/src/highlight_related.rs
index 3690fe11d0..5724090e63 100644
--- a/crates/ide/src/highlight_related.rs
+++ b/crates/ide/src/highlight_related.rs
@@ -1071,7 +1071,7 @@ fn function(field: u32) {
         check_with_config(
             r#"
 fn foo() {
-    let x = 5;
+    let x$0 = 5;
     let y = x * 2;
 }
 "#,
@@ -1091,7 +1091,7 @@ fn foo() {
         check_with_config(
             r#"
 fn foo() {
-    let x = 5;
+    let x$0 = 5;
     let y = x * 2;
 
     loop {
@@ -1108,7 +1108,7 @@ fn foo() {
     let x = 5;
     let y = x * 2;
 
-    loop {
+    loop$0 {
 //  ^^^^
         break;
 //      ^^^^^
@@ -1131,7 +1131,7 @@ fn foo() {
         check_with_config(
             r#"
 async fn foo() {
-    let x = 5;
+    let x$0 = 5;
     let y = x * 2;
 
     0.await;
@@ -1147,7 +1147,7 @@ async fn foo() {
         let x = 5;
         let y = x * 2;
 
-        0.await;
+        0.await$0;
 //        ^^^^^
 }
 "#,
@@ -1167,7 +1167,7 @@ async fn foo() {
         check_with_config(
             r#"
 fn foo() -> i32 {
-    let x = 5;
+    let x$0 = 5;
     let y = x * 2;
 
     if true {
@@ -1182,7 +1182,7 @@ fn foo() -> i32 {
 
         check_with_config(
             r#"
-fn foo() -> i32 {
+fn foo() ->$0 i32 {
     let x = 5;
     let y = x * 2;
 
@@ -1211,7 +1211,7 @@ fn foo() -> i32 {
             r#"
 fn foo() {
     loop {
-        break;
+        break$0;
     }
 }
 "#,
@@ -1230,7 +1230,7 @@ fn foo() {
 
         check_with_config(
             r#"
-async fn foo() {
+async$0 fn foo() {
     0.await;
 }
 "#,
@@ -1249,7 +1249,7 @@ async fn foo() {
 
         check_with_config(
             r#"
-fn foo() -> i32 {
+fn foo() ->$0 i32 {
     if true {
         return -1;
     }