Auto merge of rust-lang#15575 - HKalbasi:mir, r=HKalbasi

Intern projections in mir place

I hope this reduces mir memory usage.

Author: bors

crates/hir-ty/src/mir.rs

@@ -1,6 +1,6 @@
 //! MIR definitions and implementation
 
-use std::{fmt::Display, iter};
+use std::{collections::hash_map::Entry, fmt::Display, iter};
 
 use crate::{
     consteval::usize_const,
@@ -37,6 +37,7 @@ pub use monomorphization::{
     monomorphize_mir_body_bad, monomorphized_mir_body_for_closure_query,
     monomorphized_mir_body_query, monomorphized_mir_body_recover,
 };
+use rustc_hash::FxHashMap;
 use smallvec::{smallvec, SmallVec};
 use stdx::{impl_from, never};
 use triomphe::Arc;
@@ -223,35 +224,93 @@ impl<V, T> ProjectionElem<V, T> {
 
 type PlaceElem = ProjectionElem<LocalId, Ty>;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct ProjectionId(u32);
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ProjectionStore {
+    id_to_proj: FxHashMap<ProjectionId, Box<[PlaceElem]>>,
+    proj_to_id: FxHashMap<Box<[PlaceElem]>, ProjectionId>,
+}
+
+impl Default for ProjectionStore {
+    fn default() -> Self {
+        let mut this = Self { id_to_proj: Default::default(), proj_to_id: Default::default() };
+        // Ensure that [] will get the id 0 which is used in `ProjectionId::Empty`
+        this.intern(Box::new([]));
+        this
+    }
+}
+
+impl ProjectionStore {
+    fn shrink_to_fit(&mut self) {
+        self.id_to_proj.shrink_to_fit();
+        self.proj_to_id.shrink_to_fit();
+    }
+
+    fn intern_if_exist(&self, projection: &[PlaceElem]) -> Option<ProjectionId> {
+        self.proj_to_id.get(projection).copied()
+    }
+
+    fn intern(&mut self, projection: Box<[PlaceElem]>) -> ProjectionId {
+        let new_id = ProjectionId(self.proj_to_id.len() as u32);
+        match self.proj_to_id.entry(projection) {
+            Entry::Occupied(id) => *id.get(),
+            Entry::Vacant(e) => {
+                let key_clone = e.key().clone();
+                e.insert(new_id);
+                self.id_to_proj.insert(new_id, key_clone);
+                new_id
+            }
+        }
+    }
+}
+
+impl ProjectionId {
+    const EMPTY: ProjectionId = ProjectionId(0);
+
+    fn lookup(self, store: &ProjectionStore) -> &[PlaceElem] {
+        store.id_to_proj.get(&self).unwrap()
+    }
+
+    fn project(self, projection: PlaceElem, store: &mut ProjectionStore) -> ProjectionId {
+        let mut current = self.lookup(store).to_vec();
+        current.push(projection);
+        store.intern(current.into())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct Place {
     pub local: LocalId,
-    pub projection: Box<[PlaceElem]>,
+    pub projection: ProjectionId,
 }
 
 impl Place {
-    fn is_parent(&self, child: &Place) -> bool {
-        self.local == child.local && child.projection.starts_with(&self.projection)
+    fn is_parent(&self, child: &Place, store: &ProjectionStore) -> bool {
+        self.local == child.local
+            && child.projection.lookup(store).starts_with(&self.projection.lookup(store))
     }
 
     /// The place itself is not included
-    fn iterate_over_parents(&self) -> impl Iterator<Item = Place> + '_ {
-        (0..self.projection.len())
-            .map(|x| &self.projection[0..x])
-            .map(|x| Place { local: self.local, projection: x.to_vec().into() })
+    fn iterate_over_parents<'a>(
+        &'a self,
+        store: &'a ProjectionStore,
+    ) -> impl Iterator<Item = Place> + 'a {
+        let projection = self.projection.lookup(store);
+        (0..projection.len()).map(|x| &projection[0..x]).filter_map(move |x| {
+            Some(Place { local: self.local, projection: store.intern_if_exist(x)? })
+        })
     }
 
-    fn project(&self, projection: PlaceElem) -> Place {
-        Place {
-            local: self.local,
-            projection: self.projection.iter().cloned().chain([projection]).collect(),
-        }
+    fn project(&self, projection: PlaceElem, store: &mut ProjectionStore) -> Place {
+        Place { local: self.local, projection: self.projection.project(projection, store) }
     }
 }
 
 impl From<LocalId> for Place {
     fn from(local: LocalId) -> Self {
-        Self { local, projection: vec![].into() }
+        Self { local, projection: ProjectionId::EMPTY }
     }
 }
 
@@ -997,6 +1056,7 @@ pub struct BasicBlock {
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct MirBody {
+    pub projection_store: ProjectionStore,
     pub basic_blocks: Arena<BasicBlock>,
     pub locals: Arena<Local>,
     pub start_block: BasicBlockId,
@@ -1009,11 +1069,15 @@ pub struct MirBody {
 }
 
 impl MirBody {
-    fn walk_places(&mut self, mut f: impl FnMut(&mut Place)) {
-        fn for_operand(op: &mut Operand, f: &mut impl FnMut(&mut Place)) {
+    fn walk_places(&mut self, mut f: impl FnMut(&mut Place, &mut ProjectionStore)) {
+        fn for_operand(
+            op: &mut Operand,
+            f: &mut impl FnMut(&mut Place, &mut ProjectionStore),
+            store: &mut ProjectionStore,
+        ) {
             match op {
                 Operand::Copy(p) | Operand::Move(p) => {
-                    f(p);
+                    f(p, store);
                 }
                 Operand::Constant(_) | Operand::Static(_) => (),
             }
@@ -1022,38 +1086,40 @@ impl MirBody {
             for statement in &mut block.statements {
                 match &mut statement.kind {
                     StatementKind::Assign(p, r) => {
-                        f(p);
+                        f(p, &mut self.projection_store);
                         match r {
                             Rvalue::ShallowInitBoxWithAlloc(_) => (),
                             Rvalue::ShallowInitBox(o, _)
                             | Rvalue::UnaryOp(_, o)
                             | Rvalue::Cast(_, o, _)
                             | Rvalue::Repeat(o, _)
-                            | Rvalue::Use(o) => for_operand(o, &mut f),
+                            | Rvalue::Use(o) => for_operand(o, &mut f, &mut self.projection_store),
                             Rvalue::CopyForDeref(p)
                             | Rvalue::Discriminant(p)
                             | Rvalue::Len(p)
-                            | Rvalue::Ref(_, p) => f(p),
+                            | Rvalue::Ref(_, p) => f(p, &mut self.projection_store),
                             Rvalue::CheckedBinaryOp(_, o1, o2) => {
-                                for_operand(o1, &mut f);
-                                for_operand(o2, &mut f);
+                                for_operand(o1, &mut f, &mut self.projection_store);
+                                for_operand(o2, &mut f, &mut self.projection_store);
                             }
                             Rvalue::Aggregate(_, ops) => {
                                 for op in ops.iter_mut() {
-                                    for_operand(op, &mut f);
+                                    for_operand(op, &mut f, &mut self.projection_store);
                                 }
                             }
                         }
                     }
-                    StatementKind::Deinit(p) => f(p),
+                    StatementKind::Deinit(p) => f(p, &mut self.projection_store),
                     StatementKind::StorageLive(_)
                     | StatementKind::StorageDead(_)
                     | StatementKind::Nop => (),
                 }
             }
             match &mut block.terminator {
                 Some(x) => match &mut x.kind {
-                    TerminatorKind::SwitchInt { discr, .. } => for_operand(discr, &mut f),
+                    TerminatorKind::SwitchInt { discr, .. } => {
+                        for_operand(discr, &mut f, &mut self.projection_store)
+                    }
                     TerminatorKind::FalseEdge { .. }
                     | TerminatorKind::FalseUnwind { .. }
                     | TerminatorKind::Goto { .. }
@@ -1063,23 +1129,24 @@ impl MirBody {
                     | TerminatorKind::Return
                     | TerminatorKind::Unreachable => (),
                     TerminatorKind::Drop { place, .. } => {
-                        f(place);
+                        f(place, &mut self.projection_store);
                     }
                     TerminatorKind::DropAndReplace { place, value, .. } => {
-                        f(place);
-                        for_operand(value, &mut f);
+                        f(place, &mut self.projection_store);
+                        for_operand(value, &mut f, &mut self.projection_store);
                     }
                     TerminatorKind::Call { func, args, destination, .. } => {
-                        for_operand(func, &mut f);
-                        args.iter_mut().for_each(|x| for_operand(x, &mut f));
-                        f(destination);
+                        for_operand(func, &mut f, &mut self.projection_store);
+                        args.iter_mut()
+                            .for_each(|x| for_operand(x, &mut f, &mut self.projection_store));
+                        f(destination, &mut self.projection_store);
                     }
                     TerminatorKind::Assert { cond, .. } => {
-                        for_operand(cond, &mut f);
+                        for_operand(cond, &mut f, &mut self.projection_store);
                     }
                     TerminatorKind::Yield { value, resume_arg, .. } => {
-                        for_operand(value, &mut f);
-                        f(resume_arg);
+                        for_operand(value, &mut f, &mut self.projection_store);
+                        f(resume_arg, &mut self.projection_store);
                     }
                 },
                 None => (),
@@ -1096,7 +1163,9 @@ impl MirBody {
             binding_locals,
             param_locals,
             closures,
+            projection_store,
         } = self;
+        projection_store.shrink_to_fit();
         basic_blocks.shrink_to_fit();
         locals.shrink_to_fit();
         binding_locals.shrink_to_fit();

crates/hir-ty/src/mir/borrowck.rs

@@ -88,7 +88,7 @@ fn moved_out_of_ref(db: &dyn HirDatabase, body: &MirBody) -> Vec<MovedOutOfRef>
         Operand::Copy(p) | Operand::Move(p) => {
             let mut ty: Ty = body.locals[p.local].ty.clone();
             let mut is_dereference_of_ref = false;
-            for proj in &*p.projection {
+            for proj in p.projection.lookup(&body.projection_store) {
                 if *proj == ProjectionElem::Deref && ty.as_reference().is_some() {
                     is_dereference_of_ref = true;
                 }
@@ -195,7 +195,7 @@ enum ProjectionCase {
 fn place_case(db: &dyn HirDatabase, body: &MirBody, lvalue: &Place) -> ProjectionCase {
     let mut is_part_of = false;
     let mut ty = body.locals[lvalue.local].ty.clone();
-    for proj in lvalue.projection.iter() {
+    for proj in lvalue.projection.lookup(&body.projection_store).iter() {
         match proj {
             ProjectionElem::Deref if ty.as_adt().is_none() => return ProjectionCase::Indirect, // It's indirect in case of reference and raw
             ProjectionElem::Deref // It's direct in case of `Box<T>`
@@ -254,7 +254,7 @@ fn ever_initialized_map(
         for statement in &block.statements {
             match &statement.kind {
                 StatementKind::Assign(p, _) => {
-                    if p.projection.len() == 0 && p.local == l {
+                    if p.projection.lookup(&body.projection_store).len() == 0 && p.local == l {
                         is_ever_initialized = true;
                     }
                 }
@@ -289,7 +289,9 @@ fn ever_initialized_map(
             | TerminatorKind::Return
             | TerminatorKind::Unreachable => (),
             TerminatorKind::Call { target, cleanup, destination, .. } => {
-                if destination.projection.len() == 0 && destination.local == l {
+                if destination.projection.lookup(&body.projection_store).len() == 0
+                    && destination.local == l
+                {
                     is_ever_initialized = true;
                 }
                 target
@@ -389,7 +391,7 @@ fn mutability_of_locals(
             | TerminatorKind::Assert { .. }
             | TerminatorKind::Yield { .. } => (),
             TerminatorKind::Call { destination, .. } => {
-                if destination.projection.len() == 0 {
+                if destination.projection.lookup(&body.projection_store).len() == 0 {
                     if ever_init_map.get(destination.local).copied().unwrap_or_default() {
                         push_mut_span(destination.local, MirSpan::Unknown);
                     } else {

crates/hir-ty/src/mir/eval.rs

@@ -46,8 +46,8 @@ use crate::{
 
 use super::{
     return_slot, AggregateKind, BasicBlockId, BinOp, CastKind, LocalId, MirBody, MirLowerError,
-    MirSpan, Operand, Place, PlaceElem, ProjectionElem, Rvalue, StatementKind, TerminatorKind,
-    UnOp,
+    MirSpan, Operand, Place, PlaceElem, ProjectionElem, ProjectionStore, Rvalue, StatementKind,
+    TerminatorKind, UnOp,
 };
 
 mod shim;
@@ -483,17 +483,18 @@ struct DropFlags {
 }
 
 impl DropFlags {
-    fn add_place(&mut self, p: Place) {
-        if p.iterate_over_parents().any(|it| self.need_drop.contains(&it)) {
+    fn add_place(&mut self, p: Place, store: &ProjectionStore) {
+        if p.iterate_over_parents(store).any(|it| self.need_drop.contains(&it)) {
             return;
         }
-        self.need_drop.retain(|it| !p.is_parent(it));
+        self.need_drop.retain(|it| !p.is_parent(it, store));
         self.need_drop.insert(p);
     }
 
-    fn remove_place(&mut self, p: &Place) -> bool {
+    fn remove_place(&mut self, p: &Place, store: &ProjectionStore) -> bool {
         // FIXME: replace parents with parts
-        if let Some(parent) = p.iterate_over_parents().find(|it| self.need_drop.contains(&it)) {
+        if let Some(parent) = p.iterate_over_parents(store).find(|it| self.need_drop.contains(&it))
+        {
             self.need_drop.remove(&parent);
             return true;
         }
@@ -654,7 +655,7 @@ impl Evaluator<'_> {
         let mut addr = locals.ptr[p.local].addr;
         let mut ty: Ty = locals.body.locals[p.local].ty.clone();
         let mut metadata: Option<IntervalOrOwned> = None; // locals are always sized
-        for proj in &*p.projection {
+        for proj in p.projection.lookup(&locals.body.projection_store) {
             let prev_ty = ty.clone();
             ty = self.projected_ty(ty, proj.clone());
             match proj {
@@ -835,7 +836,9 @@ impl Evaluator<'_> {
                                 let addr = self.place_addr(l, &locals)?;
                                 let result = self.eval_rvalue(r, &mut locals)?.to_vec(&self)?;
                                 self.write_memory(addr, &result)?;
-                                locals.drop_flags.add_place(l.clone());
+                                locals
+                                    .drop_flags
+                                    .add_place(l.clone(), &locals.body.projection_store);
                             }
                             StatementKind::Deinit(_) => not_supported!("de-init statement"),
                             StatementKind::StorageLive(_)
@@ -887,7 +890,9 @@ impl Evaluator<'_> {
                                 )?,
                                 it => not_supported!("unknown function type {it:?}"),
                             };
-                            locals.drop_flags.add_place(destination.clone());
+                            locals
+                                .drop_flags
+                                .add_place(destination.clone(), &locals.body.projection_store);
                             if let Some(stack_frame) = stack_frame {
                                 self.code_stack.push(my_stack_frame);
                                 current_block_idx = stack_frame.locals.body.start_block;
@@ -968,7 +973,7 @@ impl Evaluator<'_> {
     ) -> Result<()> {
         let mut remain_args = body.param_locals.len();
         for ((l, interval), value) in locals.ptr.iter().skip(1).zip(args) {
-            locals.drop_flags.add_place(l.into());
+            locals.drop_flags.add_place(l.into(), &locals.body.projection_store);
             match value {
                 IntervalOrOwned::Owned(value) => interval.write_from_bytes(self, &value)?,
                 IntervalOrOwned::Borrowed(value) => interval.write_from_interval(self, value)?,
@@ -1644,7 +1649,7 @@ impl Evaluator<'_> {
     fn eval_operand(&mut self, it: &Operand, locals: &mut Locals) -> Result<Interval> {
         Ok(match it {
             Operand::Copy(p) | Operand::Move(p) => {
-                locals.drop_flags.remove_place(p);
+                locals.drop_flags.remove_place(p, &locals.body.projection_store);
                 self.eval_place(p, locals)?
             }
             Operand::Static(st) => {
@@ -2508,7 +2513,7 @@ impl Evaluator<'_> {
 
     fn drop_place(&mut self, place: &Place, locals: &mut Locals, span: MirSpan) -> Result<()> {
         let (addr, ty, metadata) = self.place_addr_and_ty_and_metadata(place, locals)?;
-        if !locals.drop_flags.remove_place(place) {
+        if !locals.drop_flags.remove_place(place, &locals.body.projection_store) {
             return Ok(());
         }
         let metadata = match metadata {