use rustc_hash::FxHashMap;
use rustc_hir::def_id::DefId;
use rustc_middle::mir::{Local, Operand, Rvalue, StatementKind};
use rustc_middle::ty::{Ty, TyCtxt, TyKind};
pub(crate) type TagId = usize;
#[derive(Clone, Debug)]
pub(crate) struct AliasNode {
pub parent: Option<TagId>,
pub fields: Vec<usize>,
pub local: Local,
}
#[derive(Clone, Debug)]
pub(crate) struct AliasTree {
pub nodes: Vec<AliasNode>,
tag_of_local: FxHashMap<Local, TagId>,
}
impl AliasTree {
pub(crate) fn build<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Self {
let body = tcx.optimized_mir(def_id);
let mut tree = AliasTree {
nodes: Vec::new(),
tag_of_local: FxHashMap::default(),
};
for local_index in 1..=body.arg_count {
let local = Local::from_usize(local_index);
let ty = body.local_decls[local].ty;
if classify(ty).is_some() {
tree.add(None, Vec::new(), local);
}
}
for block in body.basic_blocks.iter() {
for statement in &block.statements {
let StatementKind::Assign(assign) = &statement.kind else {
continue;
};
let (target, rvalue) = assign.as_ref();
match rvalue {
Rvalue::Use(Operand::Copy(place), ..)
| Rvalue::Use(Operand::Move(place), ..)
| Rvalue::CopyForDeref(place)
if field_projection(place).is_empty() =>
{
if let Some(tag) = tree.tag_of_local.get(&place.local).copied() {
tree.tag_of_local.insert(target.local, tag);
}
}
_ => {
if let Some(place) = crate::helpers::mir_utils::rvalue_source_place(rvalue)
&& let Some(parent) = tree.tag_of_local.get(&place.local).copied()
{
let ty = body.local_decls[target.local].ty;
if classify(ty).is_some() {
tree.add(Some(parent), field_projection(place), target.local);
}
}
}
}
}
if let rustc_middle::mir::TerminatorKind::Call {
args, destination, ..
} = &block.terminator().kind
{
let dest_local = destination.local;
let dest_ty = body.local_decls[dest_local].ty;
if let Some(is_owned) = classify(dest_ty) {
if is_owned {
tree.add(None, Vec::new(), dest_local);
} else if let Some(first_arg) = args.first()
&& let Some(place) = first_arg.node.place()
&& let Some(parent) = tree.tag_of_local.get(&place.local).copied()
{
tree.add(Some(parent), field_projection(&place), dest_local);
}
}
}
}
tree
}
fn add(&mut self, parent: Option<TagId>, fields: Vec<usize>, local: Local) -> TagId {
let tag = self.nodes.len();
self.nodes.push(AliasNode {
parent,
fields,
local,
});
self.tag_of_local.insert(local, tag);
tag
}
pub(crate) fn tag_of(&self, local: Local) -> Option<TagId> {
self.tag_of_local.get(&local).copied()
}
fn resolve_to_root(&self, tag: TagId) -> (Local, Vec<usize>) {
let mut cur = tag;
let mut fields: Vec<usize> = Vec::new();
let mut guard = 0;
loop {
let node = &self.nodes[cur];
let mut combined = node.fields.clone();
combined.extend(fields.iter().copied());
fields = combined;
match node.parent {
Some(parent) => cur = parent,
None => return (node.local, fields),
}
guard += 1;
if guard > self.nodes.len() {
return (self.nodes[cur].local, fields);
}
}
}
pub(crate) fn resolve_local_to_root(&self, local: Local) -> (usize, Vec<usize>) {
match self.tag_of(local) {
Some(tag) => {
let (root, fields) = self.resolve_to_root(tag);
(root.as_usize(), fields)
}
None => (local.as_usize(), Vec::new()),
}
}
}
fn field_projection(place: &rustc_middle::mir::Place<'_>) -> Vec<usize> {
place
.projection
.iter()
.filter_map(|p| match p {
rustc_middle::mir::ProjectionElem::Field(idx, _) => Some(idx.as_usize()),
_ => None,
})
.collect()
}
fn classify(ty: Ty<'_>) -> Option<bool> {
match ty.kind() {
TyKind::Ref(_, _, _) | TyKind::RawPtr(_, _) => Some(false),
TyKind::Adt(_, _) => Some(true),
_ => None,
}
}