Add helper function for Capture Esclations and expressions

Co-authored-by: Dhruv Jauhar <dhruvjhr@gmail.com>

3 files changed, 110 insertions, 67 deletions

diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
index 5f2d3b7818e..52b49184cf1 100644
--- a/compiler/rustc_middle/src/ty/mod.rs
+++ b/compiler/rustc_middle/src/ty/mod.rs
@@ -765,7 +765,23 @@ pub struct UpvarBorrow<'tcx> {
 
 #[derive(PartialEq, Clone, Debug, Copy, TyEncodable, TyDecodable, HashStable)]
 pub struct CaptureInfo<'tcx> {
-    /// Expr Id pointing to use that resulting in selecting the current capture kind
+    /// Expr Id pointing to use that resulted in selecting the current capture kind
+    ///
+    /// If the user doesn't enable feature `capture_disjoint_fields` (RFC 2229) then, it is
+    /// possible that we don't see the use of a particular place resulting in expr_id being
+    /// None. In such case we fallback on uvpars_mentioned for span.
+    ///
+    /// Eg:
+    /// ```rust
+    /// let x = ...;
+    ///
+    /// let c = || {
+    ///     let _ = x
+    /// }
+    /// ```
+    ///
+    /// In this example, if `capture_disjoint_fields` is **not** set, then x will be captured,
+    /// but we won't see it being used during capture analysis, since it's essentially a discard.
     pub expr_id: Option<hir::HirId>,
 
     /// Capture mode that was selected
diff --git a/compiler/rustc_typeck/src/check/upvar.rs b/compiler/rustc_typeck/src/check/upvar.rs
index 9e4e6565361..63657971485 100644
--- a/compiler/rustc_typeck/src/check/upvar.rs
+++ b/compiler/rustc_typeck/src/check/upvar.rs
@@ -284,30 +284,17 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             let var_hir_id = upvar_id.var_path.hir_id;
             closure_captures.insert(var_hir_id, upvar_id);
 
-            let mut new_capture_kind = capture_info.capture_kind;
-            if let Some(existing_capture_kind) =
+            let new_capture_kind = if let Some(capture_kind) =
                 self.typeck_results.borrow_mut().upvar_capture_map.get(&upvar_id)
             {
-                // FIXME(@azhng): refactor this later
-                new_capture_kind = match (existing_capture_kind, new_capture_kind) {
-                    (ty::UpvarCapture::ByValue(Some(_)), _) => *existing_capture_kind,
-                    (_, ty::UpvarCapture::ByValue(Some(_))) => new_capture_kind,
-                    (ty::UpvarCapture::ByValue(_), _) | (_, ty::UpvarCapture::ByValue(_)) => {
-                        ty::UpvarCapture::ByValue(None)
-                    }
-                    (ty::UpvarCapture::ByRef(existing_ref), ty::UpvarCapture::ByRef(new_ref)) => {
-                        match (existing_ref.kind, new_ref.kind) {
-                            // Take RHS:
-                            (ty::ImmBorrow, ty::UniqueImmBorrow | ty::MutBorrow)
-                            | (ty::UniqueImmBorrow, ty::MutBorrow) => new_capture_kind,
-                            // Take LHS:
-                            (ty::ImmBorrow, ty::ImmBorrow)
-                            | (ty::UniqueImmBorrow, ty::ImmBorrow | ty::UniqueImmBorrow)
-                            | (ty::MutBorrow, _) => *existing_capture_kind,
-                        }
-                    }
-                };
-            }
+                // upvar_capture_map only stores the UpvarCapture (CaptureKind),
+                // so we create a fake capture info with no expression.
+                let fake_capture_info =
+                    ty::CaptureInfo { expr_id: None, capture_kind: capture_kind.clone() };
+                self.determine_capture_info(fake_capture_info, capture_info.clone()).capture_kind
+            } else {
+                capture_info.capture_kind
+            };
             self.typeck_results.borrow_mut().upvar_capture_map.insert(upvar_id, new_capture_kind);
         }
 
@@ -353,6 +340,60 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     fn should_log_capture_analysis(&self, closure_def_id: DefId) -> bool {
         self.tcx.has_attr(closure_def_id, sym::rustc_capture_analysis)
     }
+
+    /// Helper function to determine if we need to escalate CaptureKind from
+    /// CaptureInfo A to B and returns the escalated CaptureInfo.
+    /// (Note: CaptureInfo contains CaptureKind and an expression that led to capture it in that way)
+    ///
+    /// If both `CaptureKind`s are considered equivalent, then the CaptureInfo is selected based
+    /// on the `CaptureInfo` containing an associated expression id.
+    ///
+    /// If both the CaptureKind and Expression are considered to be equivalent,
+    /// then `CaptureInfo` A is preferred.
+    fn determine_capture_info(
+        &self,
+        capture_info_a: ty::CaptureInfo<'tcx>,
+        capture_info_b: ty::CaptureInfo<'tcx>,
+    ) -> ty::CaptureInfo<'tcx> {
+        // If the capture kind is equivalent then, we don't need to escalate and can compare the
+        // expressions.
+        let eq_capture_kind = match (capture_info_a.capture_kind, capture_info_b.capture_kind) {
+            (ty::UpvarCapture::ByValue(_), ty::UpvarCapture::ByValue(_)) => true,
+            (ty::UpvarCapture::ByRef(ref_a), ty::UpvarCapture::ByRef(ref_b)) => {
+                ref_a.kind == ref_b.kind
+            }
+            _ => false,
+        };
+
+        if eq_capture_kind {
+            match (capture_info_a.expr_id, capture_info_b.expr_id) {
+                (Some(_), _) | (None, None) => capture_info_a,
+                (None, Some(_)) => capture_info_b,
+            }
+        } else {
+            match (capture_info_a.capture_kind, capture_info_b.capture_kind) {
+                (ty::UpvarCapture::ByValue(_), _) => capture_info_a,
+                (_, ty::UpvarCapture::ByValue(_)) => capture_info_b,
+                (ty::UpvarCapture::ByRef(ref_a), ty::UpvarCapture::ByRef(ref_b)) => {
+                    match (ref_a.kind, ref_b.kind) {
+                        // Take LHS:
+                        (ty::UniqueImmBorrow | ty::MutBorrow, ty::ImmBorrow)
+                        | (ty::MutBorrow, ty::UniqueImmBorrow) => capture_info_a,
+
+                        // Take RHS:
+                        (ty::ImmBorrow, ty::UniqueImmBorrow | ty::MutBorrow)
+                        | (ty::UniqueImmBorrow, ty::MutBorrow) => capture_info_b,
+
+                        (ty::ImmBorrow, ty::ImmBorrow)
+                        | (ty::UniqueImmBorrow, ty::UniqueImmBorrow)
+                        | (ty::MutBorrow, ty::MutBorrow) => {
+                            bug!("Expected unequal capture kinds");
+                        }
+                    }
+                }
+            }
+        }
+    }
 }
 
 struct InferBorrowKind<'a, 'tcx> {
@@ -426,16 +467,10 @@ impl<'a, 'tcx> InferBorrowKind<'a, 'tcx> {
             capture_kind: ty::UpvarCapture::ByValue(Some(usage_span)),
         };
 
-        let curr_info = self.capture_information.get(&place_with_id.place);
-        let updated_info = match curr_info {
-            Some(info) => match info.capture_kind {
-                ty::UpvarCapture::ByRef(_) | ty::UpvarCapture::ByValue(None) => capture_info,
-                _ => *info,
-            },
-            None => capture_info,
-        };
+        let curr_info = self.capture_information[&place_with_id.place];
+        let updated_info = self.fcx.determine_capture_info(curr_info, capture_info);
 
-        self.capture_information.insert(place_with_id.place.clone(), updated_info);
+        self.capture_information[&place_with_id.place] = updated_info;
     }
 
     /// Indicates that `place_with_id` is being directly mutated (e.g., assigned
@@ -532,42 +567,28 @@ impl<'a, 'tcx> InferBorrowKind<'a, 'tcx> {
         diag_expr_id: hir::HirId,
         kind: ty::BorrowKind,
     ) {
-        let capture_info = self
-            .capture_information
-            .get(&place_with_id.place)
-            .unwrap_or_else(|| bug!("Upar capture info missing"));
-        // We init capture_information for each element
+        let curr_capture_info = self.capture_information[&place_with_id.place];
 
         debug!(
-            "adjust_upvar_borrow_kind(place={:?}, , diag_expr_id={:?}, capture_info={:?}, kind={:?})",
-            place_with_id, diag_expr_id, capture_info, kind
+            "adjust_upvar_borrow_kind(place={:?}, diag_expr_id={:?}, capture_info={:?}, kind={:?})",
+            place_with_id, diag_expr_id, curr_capture_info, kind
         );
 
-        match capture_info.capture_kind {
-            ty::UpvarCapture::ByValue(_) => {
-                // Upvar is already by-value, the strongest criteria.
-            }
-            ty::UpvarCapture::ByRef(upvar_borrow) => {
-                match (upvar_borrow.kind, kind) {
-                    // Take RHS:
-                    (ty::ImmBorrow, ty::UniqueImmBorrow | ty::MutBorrow)
-                    | (ty::UniqueImmBorrow, ty::MutBorrow) => {
-                        if let Some(ty::CaptureInfo { expr_id, capture_kind }) =
-                            self.capture_information.get_mut(&place_with_id.place)
-                        {
-                            *expr_id = Some(diag_expr_id);
-                            if let ty::UpvarCapture::ByRef(borrow_kind) = capture_kind {
-                                borrow_kind.kind = kind;
-                            }
-                        }
-                    }
-                    // Take LHS:
-                    (ty::ImmBorrow, ty::ImmBorrow)
-                    | (ty::UniqueImmBorrow, ty::ImmBorrow | ty::UniqueImmBorrow)
-                    | (ty::MutBorrow, _) => {}
-                }
-            }
-        }
+        if let ty::UpvarCapture::ByValue(_) = curr_capture_info.capture_kind {
+            // It's already captured by value, we don't need to do anything here
+            return;
+        } else if let ty::UpvarCapture::ByRef(curr_upvar_borrow) = curr_capture_info.capture_kind {
+            // Use the same region as the current capture information
+            // Doesn't matter since only one of the UpvarBorrow will be used.
+            let new_upvar_borrow = ty::UpvarBorrow { kind, region: curr_upvar_borrow.region };
+
+            let capture_info = ty::CaptureInfo {
+                expr_id: Some(diag_expr_id),
+                capture_kind: ty::UpvarCapture::ByRef(new_upvar_borrow),
+            };
+            let updated_info = self.fcx.determine_capture_info(curr_capture_info, capture_info);
+            self.capture_information[&place_with_id.place] = updated_info;
+        };
     }
 
     fn adjust_closure_kind(
@@ -622,7 +643,6 @@ impl<'a, 'tcx> InferBorrowKind<'a, 'tcx> {
 
             debug!("Capturing new place {:?}", place_with_id);
 
-            let tcx = self.fcx.tcx;
             let capture_kind =
                 self.fcx.init_capture_kind(self.capture_clause, upvar_id, self.closure_span);
 
diff --git a/src/test/ui/closures/2229_closure_analysis/slice-pat.stdout b/src/test/ui/closures/2229_closure_analysis/slice-pat.stdout
index ba9b2bca196..3e352e2c525 100644
--- a/src/test/ui/closures/2229_closure_analysis/slice-pat.stdout
+++ b/src/test/ui/closures/2229_closure_analysis/slice-pat.stdout
@@ -12,9 +12,16 @@ For closure=DefId(0:5 ~ slice_pat[317d]::main::{closure#0}): capture information
             },
         ],
     }: CaptureInfo {
-        expr_id: None,
+        expr_id: Some(
+            HirId {
+                owner: DefId(0:3 ~ slice_pat[317d]::main),
+                local_id: 179,
+            },
+        ),
         capture_kind: ByValue(
-            None,
+            Some(
+                $DIR/slice-pat.rs:21:33: 21:36 (#0),
+            ),
         ),
     },
 }