Use Places to express closure/generator Captures

Co-authored-by: Archer Zhang <[email protected]>

Author: arora-aman

compiler/rustc_middle/src/ty/context.rs

@@ -415,6 +415,10 @@ pub struct TypeckResults<'tcx> {
     /// entire variable.
     pub closure_captures: ty::UpvarListMap,
 
+    /// Given the closure ID this map provides the list of
+    /// `Place`s and how/why are they captured by the closure.
+    pub closure_capture_information: ty::CaptureInformationMap<'tcx>,
+
     /// Stores the type, expression, span and optional scope span of all types
     /// that are live across the yield of this generator (if a generator).
     pub generator_interior_types: Vec<GeneratorInteriorTypeCause<'tcx>>,
@@ -442,6 +446,7 @@ impl<'tcx> TypeckResults<'tcx> {
             tainted_by_errors: None,
             concrete_opaque_types: Default::default(),
             closure_captures: Default::default(),
+            closure_capture_information: Default::default(),
             generator_interior_types: Default::default(),
         }
     }
@@ -676,6 +681,7 @@ impl<'a, 'tcx> HashStable<StableHashingContext<'a>> for TypeckResults<'tcx> {
             tainted_by_errors,
             ref concrete_opaque_types,
             ref closure_captures,
+            ref closure_capture_information,
             ref generator_interior_types,
         } = *self;
 
@@ -709,6 +715,7 @@ impl<'a, 'tcx> HashStable<StableHashingContext<'a>> for TypeckResults<'tcx> {
             tainted_by_errors.hash_stable(hcx, hasher);
             concrete_opaque_types.hash_stable(hcx, hasher);
             closure_captures.hash_stable(hcx, hasher);
+            closure_capture_information.hash_stable(hcx, hasher);
             generator_interior_types.hash_stable(hcx, hasher);
         })
     }

compiler/rustc_middle/src/ty/mod.rs

@@ -6,6 +6,7 @@ pub use self::IntVarValue::*;
 pub use self::Variance::*;
 
 use crate::hir::exports::ExportMap;
+use crate::hir::place::Place as HirPlace;
 use crate::ich::StableHashingContext;
 use crate::middle::cstore::CrateStoreDyn;
 use crate::middle::resolve_lifetime::ObjectLifetimeDefault;
@@ -674,6 +675,12 @@ pub struct UpvarId {
     pub closure_expr_id: LocalDefId,
 }
 
+impl UpvarId {
+    pub fn new(var_hir_id: hir::HirId, closure_def_id: LocalDefId) -> UpvarId {
+        UpvarId { var_path: UpvarPath { hir_id: var_hir_id }, closure_expr_id: closure_def_id }
+    }
+}
+
 #[derive(Clone, PartialEq, Debug, TyEncodable, TyDecodable, Copy, HashStable)]
 pub enum BorrowKind {
     /// Data must be immutable and is aliasable.
@@ -756,9 +763,40 @@ pub struct UpvarBorrow<'tcx> {
     pub region: ty::Region<'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
+    pub expr_id: Option<hir::HirId>,
+
+    /// Capture mode that was selected
+    pub capture_kind: UpvarCapture<'tcx>,
+}
+
 pub type UpvarListMap = FxHashMap<DefId, FxIndexMap<hir::HirId, UpvarId>>;
 pub type UpvarCaptureMap<'tcx> = FxHashMap<UpvarId, UpvarCapture<'tcx>>;
 
+/// Consider closure where s.str1 is captured via an ImmutableBorrow and s.str2 via a MutableBorrow
+///
+/// ```rust
+/// // Assume that thte HirId for the variable definition is `V1`
+/// let mut s = SomeStruct { str1: format!("s1"), str2: format!("s2") }
+///
+/// let fix_s = |new_s2| {
+///     // Assume that the HirId for the expression `s.str1` is `E1`
+///     println!("Updating SomeStruct with str1=", s.str1);
+///     // Assume that the HirId for the expression `*s.str2` is `E2`
+///     s.str2 = new_s2;
+/// }
+/// ```
+///
+/// For closure `fix_s`, (at a high level) the IndexMap will contain:
+///
+/// Place { V1, [ProjectionKind::Field(Index=0, Variant=0)] } : CaptureKind { E1, ImmutableBorrow }
+/// Place { V1, [ProjectionKind::Field(Index=1, Variant=0)] } : CaptureKind { E2, MutableBorrow }
+///
+pub type CaptureInformationMap<'tcx> =
+    FxHashMap<DefId, FxIndexMap<HirPlace<'tcx>, CaptureInfo<'tcx>>>;
+
 #[derive(Clone, Copy, PartialEq, Eq)]
 pub enum IntVarValue {
     IntType(ast::IntTy),

compiler/rustc_mir_build/src/thir/cx/expr.rs

@@ -387,8 +387,9 @@ fn make_mirror_unadjusted<'a, 'tcx>(
                 }
             };
             let upvars = cx
-                .tcx
-                .upvars_mentioned(def_id)
+                .typeck_results()
+                .closure_captures
+                .get(&def_id)
                 .iter()
                 .flat_map(|upvars| upvars.iter())
                 .zip(substs.upvar_tys())