We can traverse bindings before `lower_match_tree` now

2 files changed, 98 insertions, 50 deletions

diff --git a/compiler/rustc_mir_build/src/build/matches/mod.rs b/compiler/rustc_mir_build/src/build/matches/mod.rs
index 68244136d1a..f182e1c57f4 100644
--- a/compiler/rustc_mir_build/src/build/matches/mod.rs
+++ b/compiler/rustc_mir_build/src/build/matches/mod.rs
@@ -21,6 +21,7 @@ use rustc_span::symbol::Symbol;
 use rustc_span::{BytePos, Pos, Span};
 use rustc_target::abi::VariantIdx;
 use tracing::{debug, instrument};
+use util::visit_bindings;
 
 // helper functions, broken out by category:
 mod simplify;
@@ -725,55 +726,49 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         set_match_place: bool,
     ) -> BlockAnd<()> {
         let mut candidate = Candidate::new(initializer.clone(), irrefutable_pat, false, self);
-        let fake_borrow_temps = self.lower_match_tree(
-            block,
-            irrefutable_pat.span,
-            &initializer,
-            irrefutable_pat.span,
-            false,
-            &mut [&mut candidate],
-        );
 
         // For matches and function arguments, the place that is being matched
         // can be set when creating the variables. But the place for
         // let PATTERN = ... might not even exist until we do the assignment.
         // so we set it here instead.
         if set_match_place {
-            let mut next = Some(&candidate);
-            while let Some(candidate_ref) = next.take() {
-                for binding in &candidate_ref.extra_data.bindings {
+            // `try_to_place` may fail if it is unable to resolve the given `PlaceBuilder` inside a
+            // closure. In this case, we don't want to include a scrutinee place.
+            // `scrutinee_place_builder` will fail for destructured assignments. This is because a
+            // closure only captures the precise places that it will read and as a result a closure
+            // may not capture the entire tuple/struct and rather have individual places that will
+            // be read in the final MIR.
+            // Example:
+            // ```
+            // let foo = (0, 1);
+            // let c = || {
+            //    let (v1, v2) = foo;
+            // };
+            // ```
+            if let Some(place) = initializer.try_to_place(self) {
+                visit_bindings(&[&mut candidate], |binding: &Binding<'_>| {
                     let local = self.var_local_id(binding.var_id, OutsideGuard);
-                    // `try_to_place` may fail if it is unable to resolve the given
-                    // `PlaceBuilder` inside a closure. In this case, we don't want to include
-                    // a scrutinee place. `scrutinee_place_builder` will fail for destructured
-                    // assignments. This is because a closure only captures the precise places
-                    // that it will read and as a result a closure may not capture the entire
-                    // tuple/struct and rather have individual places that will be read in the
-                    // final MIR.
-                    // Example:
-                    // ```
-                    // let foo = (0, 1);
-                    // let c = || {
-                    //    let (v1, v2) = foo;
-                    // };
-                    // ```
-                    if let Some(place) = initializer.try_to_place(self) {
-                        let LocalInfo::User(BindingForm::Var(VarBindingForm {
-                            opt_match_place: Some((ref mut match_place, _)),
-                            ..
-                        })) = **self.local_decls[local].local_info.as_mut().assert_crate_local()
-                        else {
-                            bug!("Let binding to non-user variable.")
-                        };
+                    if let LocalInfo::User(BindingForm::Var(VarBindingForm {
+                        opt_match_place: Some((ref mut match_place, _)),
+                        ..
+                    })) = **self.local_decls[local].local_info.as_mut().assert_crate_local()
+                    {
                         *match_place = Some(place);
-                    }
-                }
-                // All of the subcandidates should bind the same locals, so we
-                // only visit the first one.
-                next = candidate_ref.subcandidates.get(0)
+                    } else {
+                        bug!("Let binding to non-user variable.")
+                    };
+                });
             }
         }
 
+        let fake_borrow_temps = self.lower_match_tree(
+            block,
+            irrefutable_pat.span,
+            &initializer,
+            irrefutable_pat.span,
+            false,
+            &mut [&mut candidate],
+        );
         self.bind_pattern(
             self.source_info(irrefutable_pat.span),
             candidate,
diff --git a/compiler/rustc_mir_build/src/build/matches/util.rs b/compiler/rustc_mir_build/src/build/matches/util.rs
index 50f4ca2d819..e0ba736f198 100644
--- a/compiler/rustc_mir_build/src/build/matches/util.rs
+++ b/compiler/rustc_mir_build/src/build/matches/util.rs
@@ -1,3 +1,5 @@
+use std::marker::PhantomData;
+
 use crate::build::expr::as_place::{PlaceBase, PlaceBuilder};
 use crate::build::matches::{Binding, Candidate, FlatPat, MatchPair, TestCase};
 use crate::build::Builder;
@@ -267,18 +269,6 @@ impl<'pat, 'tcx> MatchPair<'pat, 'tcx> {
     }
 }
 
-pub(super) struct FakeBorrowCollector<'a, 'b, 'tcx> {
-    cx: &'a mut Builder<'b, 'tcx>,
-    /// Base of the scrutinee place. Used to distinguish bindings inside the scrutinee place from
-    /// bindings inside deref patterns.
-    scrutinee_base: PlaceBase,
-    /// Store for each place the kind of borrow to take. In case of conflicts, we take the strongest
-    /// borrow (i.e. Deep > Shallow).
-    /// Invariant: for any place in `fake_borrows`, all the prefixes of this place that are
-    /// dereferences are also borrowed with the same of stronger borrow kind.
-    fake_borrows: FxIndexMap<Place<'tcx>, FakeBorrowKind>,
-}
-
 /// Determine the set of places that have to be stable across match guards.
 ///
 /// Returns a list of places that need a fake borrow along with a local to store it.
@@ -342,6 +332,18 @@ pub(super) fn collect_fake_borrows<'tcx>(
         .collect()
 }
 
+pub(super) struct FakeBorrowCollector<'a, 'b, 'tcx> {
+    cx: &'a mut Builder<'b, 'tcx>,
+    /// Base of the scrutinee place. Used to distinguish bindings inside the scrutinee place from
+    /// bindings inside deref patterns.
+    scrutinee_base: PlaceBase,
+    /// Store for each place the kind of borrow to take. In case of conflicts, we take the strongest
+    /// borrow (i.e. Deep > Shallow).
+    /// Invariant: for any place in `fake_borrows`, all the prefixes of this place that are
+    /// dereferences are also borrowed with the same of stronger borrow kind.
+    fake_borrows: FxIndexMap<Place<'tcx>, FakeBorrowKind>,
+}
+
 impl<'a, 'b, 'tcx> FakeBorrowCollector<'a, 'b, 'tcx> {
     // Fake borrow this place and its dereference prefixes.
     fn fake_borrow(&mut self, place: Place<'tcx>, kind: FakeBorrowKind) {
@@ -455,6 +457,57 @@ impl<'a, 'b, 'tcx> FakeBorrowCollector<'a, 'b, 'tcx> {
     }
 }
 
+/// Visit all the bindings of these candidates. Because or-alternatives bind the same variables, we
+/// only explore the first one of each or-pattern.
+pub(super) fn visit_bindings<'tcx>(
+    candidates: &[&mut Candidate<'_, 'tcx>],
+    f: impl FnMut(&Binding<'tcx>),
+) {
+    let mut visitor = BindingsVisitor { f, phantom: PhantomData };
+    for candidate in candidates.iter() {
+        visitor.visit_candidate(candidate);
+    }
+}
+
+pub(super) struct BindingsVisitor<'tcx, F> {
+    f: F,
+    phantom: PhantomData<&'tcx ()>,
+}
+
+impl<'tcx, F> BindingsVisitor<'tcx, F>
+where
+    F: FnMut(&Binding<'tcx>),
+{
+    fn visit_candidate(&mut self, candidate: &Candidate<'_, 'tcx>) {
+        for binding in &candidate.extra_data.bindings {
+            (self.f)(binding)
+        }
+        for match_pair in &candidate.match_pairs {
+            self.visit_match_pair(match_pair);
+        }
+    }
+
+    fn visit_flat_pat(&mut self, flat_pat: &FlatPat<'_, 'tcx>) {
+        for binding in &flat_pat.extra_data.bindings {
+            (self.f)(binding)
+        }
+        for match_pair in &flat_pat.match_pairs {
+            self.visit_match_pair(match_pair);
+        }
+    }
+
+    fn visit_match_pair(&mut self, match_pair: &MatchPair<'_, 'tcx>) {
+        if let TestCase::Or { pats, .. } = &match_pair.test_case {
+            // All the or-alternatives should bind the same locals, so we only visit the first one.
+            self.visit_flat_pat(&pats[0])
+        } else {
+            for subpair in &match_pair.subpairs {
+                self.visit_match_pair(subpair);
+            }
+        }
+    }
+}
+
 #[must_use]
 pub(crate) fn ref_pat_borrow_kind(ref_mutability: Mutability) -> BorrowKind {
     match ref_mutability {