Skip to content

Commit

Permalink
Intern projections in mir place
Browse files Browse the repository at this point in the history
  • Loading branch information
@HKalbasi
HKalbasi committed Sep 8, 2023
1 parent e4c4693 commit 9708a29
Show file tree
Hide file tree
Showing 7 changed files with 253 additions and 157 deletions.
137 changes: 103 additions & 34 deletions crates/hir-ty/src/mir.rs
View file
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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 }
}
}

Expand Down Expand Up @@ -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,
Expand All @@ -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(_) => (),
}
Expand All @@ -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 { .. }
Expand All @@ -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 => (),
Expand All @@ -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();
Expand Down
12 changes: 7 additions & 5 deletions crates/hir-ty/src/mir/borrowck.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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>`
Expand Down Expand Up @@ -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;
}
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
31 changes: 18 additions & 13 deletions crates/hir-ty/src/mir/eval.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -485,17 +485,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;
}
Expand Down Expand Up @@ -656,7 +657,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 {
Expand Down Expand Up @@ -837,7 +838,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(_)
Expand Down Expand Up @@ -889,7 +892,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;
Expand Down Expand Up @@ -970,7 +975,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)?,
Expand Down Expand Up @@ -1646,7 +1651,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) => {
Expand Down Expand Up @@ -2468,7 +2473,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 {
Expand Down

0 comments on commit 9708a29

Please sign in to comment.