use std::collections::VecDeque;
use rustc_middle::{
mir::{self, BasicBlock, BinOp, UnwindAction},
ty::{self, Ty, TyCtxt},
};
use crate::value::{
self, Against, Bounds, Fact, Known, LenRel, Ranks, STOPS, Taught,
Thresholds, Value, truncate,
};
const PRECISE: u32 = 2;
pub const STEPS: usize = 2 * 2 * (PRECISE as usize + STOPS + 2) + NAMING_PLANES;
const NAMING_PLANES: usize = 6 * 2 + 2 * (PRECISE as usize + 3);
pub type State<'tcx> = Vec<Fact<'tcx>>;
#[derive(Debug, Clone, Copy)]
pub struct Subject<'tcx> {
pub read: mir::Local,
pub ty: Ty<'tcx>,
pub width: u32,
pub compared: [Option<Compared<'tcx>>; READINGS],
}
pub const READINGS: usize = 8;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Compared<'tcx> {
pub op: BinOp,
pub local: mir::Local,
pub against: Against<'tcx>,
pub source: Option<mir::Local>,
pub arm: Option<bool>,
}
pub struct Work<'tcx> {
entry: Vec<Option<State<'tcx>>>,
queued: Vec<bool>,
changes: Vec<u32>,
queue: VecDeque<BasicBlock>,
stops: Thresholds,
}
impl<'tcx> Work<'tcx> {
pub fn new(blocks: usize, stops: Thresholds) -> Self {
Self {
entry: vec![None; blocks],
queued: vec![false; blocks],
changes: vec![0; blocks],
queue: VecDeque::new(),
stops,
}
}
pub fn merge(&mut self, bb: BasicBlock, incoming: State<'tcx>) {
let widen = self
.changes
.get(bb.as_usize())
.is_some_and(|seen| *seen >= PRECISE);
let Some(slot) = self.entry.get_mut(bb.as_usize()) else {
return;
};
match slot {
None => *slot = Some(incoming),
Some(existing) => {
let mut changed = false;
for (held, arriving) in existing.iter_mut().zip(&incoming) {
let mut next = held.joined(*arriving);
if next == *held {
continue;
}
if widen {
next = next.widened(*held, &self.stops);
}
*held = next;
changed = true;
}
if !changed {
return;
}
}
}
if let Some(seen) = self.changes.get_mut(bb.as_usize()) {
*seen = seen.saturating_add(1);
}
if let Some(queued) = self.queued.get_mut(bb.as_usize())
&& !*queued
{
*queued = true;
self.queue.push_back(bb);
}
}
pub fn is_drained(&self) -> bool {
self.queue.is_empty()
}
pub fn pop(&mut self) -> Option<(BasicBlock, State<'tcx>)> {
let bb = self.queue.pop_front()?;
if let Some(queued) = self.queued.get_mut(bb.as_usize()) {
*queued = false;
}
let state = self.entry.get(bb.as_usize())?.clone()?;
Some((bb, state))
}
}
pub fn escaping<'tcx>(
tcx: TyCtxt<'tcx>,
env: ty::TypingEnv<'tcx>,
mir: &mir::Body<'tcx>,
) -> Vec<bool> {
let mut escaped = vec![false; mir.local_decls.len()];
for block in mir.basic_blocks.iter() {
for stmt in &block.statements {
let mir::StatementKind::Assign(pair) = &stmt.kind else {
continue;
};
let (place, shared) = match &pair.1 {
mir::Rvalue::Ref(_, kind, place) => {
(place, !matches!(kind, mir::BorrowKind::Mut { .. }))
}
mir::Rvalue::Reborrow(_, mutability, place) => {
(place, *mutability == mir::Mutability::Not)
}
mir::Rvalue::RawPtr(_, place) => (place, false),
_ => continue,
};
if place.projection.first() == Some(&mir::ProjectionElem::Deref) {
continue;
}
if shared && place.ty(&mir.local_decls, tcx).ty.is_freeze(tcx, env)
{
continue;
}
if let Some(slot) = escaped.get_mut(place.local.as_usize()) {
*slot = true;
}
}
}
escaped
}
pub fn put<'tcx>(state: &mut State<'tcx>, slot: mir::Local, fact: Fact<'tcx>) {
if let Some(cell) = state.get_mut(slot.as_usize()) {
*cell = fact;
}
}
pub fn root_of(state: &State<'_>, local: mir::Local) -> mir::Local {
state
.get(local.as_usize())
.and_then(|fact| fact.same)
.unwrap_or(local)
}
pub fn refined<'tcx>(
state: &State<'tcx>,
subject: Option<&Subject<'tcx>>,
value: Option<u128>,
matched: bool,
) -> State<'tcx> {
let mut next = state.clone();
let (Some(subject), Some(value)) = (subject, value) else {
return next;
};
let read = Known {
bits: truncate(value, subject.width),
ty: subject.ty,
width: subject.width,
};
learn(
&mut next,
subject.read,
Taught::Value(settle(read, matched)),
);
for compared in subject.compared.into_iter().flatten() {
let truth = if matched { value == 1 } else { value == 0 };
if compared.arm.is_some_and(|only| only != truth) {
continue;
}
if let Some(fact) = value::fact_of(compared.op, compared.against, truth)
{
learn(&mut next, compared.local, fact);
}
}
next
}
pub fn settle(known: Known<'_>, holds: bool) -> Value<'_> {
if holds {
Value::Exact(known)
} else {
Value::other_than(known)
}
}
pub fn learn<'tcx>(
state: &mut State<'tcx>,
local: mir::Local,
taught: Taught<'tcx>,
) {
let behind = match state.get(local.as_usize()).and_then(|slot| slot.value) {
Some(Value::Length(of)) => Some(of),
_ => None,
};
if let (Taught::Value(value), Some(of)) = (taught, behind) {
stretch(state, of, value);
return;
}
if let Taught::Alike(other) = taught {
let Some(mine) = behind else {
return;
};
if mine == other {
return;
}
if let Some(slot) = state.get_mut(mine.as_usize()) {
slot.paired = Some(other);
}
if let Some(slot) = state.get_mut(other.as_usize()) {
slot.paired = Some(mine);
}
return;
}
let over = state.get(local.as_usize()).and_then(|slot| slot.over);
let Some(slot) = state.get_mut(local.as_usize()) else {
return;
};
match taught {
Taught::Value(value) => {
slot.value =
Some(slot.value.map_or(value, |held| held.refined(value)));
}
Taught::Order(rel, of) => {
let trivial = matches!(
(slot.value, slot.order.first()),
(Some(Value::Length(mine)), Some((held, named)))
if mine == named && held == LenRel::AT_MOST
);
if trivial {
slot.order = Ranks::of(rel, of);
} else {
slot.order.add(rel, of);
}
if let Some((base, step)) = over
&& let Ok(step) = u64::try_from(step)
&& let Some(under) = state.get_mut(base.as_usize())
{
under.order.add(rel.lowered(step), of);
}
}
Taught::Apart(of) => {
if slot.order.against(of) == Some(LenRel::AT_MOST) {
slot.order.add(LenRel::BELOW, of);
}
}
Taught::Alike(_) => {}
}
}
pub fn stretch<'tcx>(
state: &mut State<'tcx>,
of: mir::Local,
taught: Value<'tcx>,
) {
let Some(slot) = state.get_mut(of.as_usize()) else {
return;
};
let whole = taught
.anchor()
.and_then(|end| Bounds::new(end.type_min(), end.type_max()));
let Some(held) = slot.extent.or(whole) else {
return;
};
let narrowed = Value::Within(held).refined(taught).bounds();
slot.extent = narrowed;
for slot in state.iter_mut() {
if slot.value == Some(Value::Length(of)) {
slot.extent = narrowed;
}
}
}
pub fn writes(stmt: &mir::Statement<'_>, local: mir::Local) -> bool {
match &stmt.kind {
mir::StatementKind::Assign(pair) => pair.0.local == local,
mir::StatementKind::SetDiscriminant { place, .. } => {
place.local == local
}
mir::StatementKind::StorageLive(other)
| mir::StatementKind::StorageDead(other) => *other == local,
mir::StatementKind::FakeRead(..)
| mir::StatementKind::PlaceMention(..)
| mir::StatementKind::AscribeUserType(..)
| mir::StatementKind::Coverage(..)
| mir::StatementKind::ConstEvalCounter
| mir::StatementKind::Nop
| mir::StatementKind::BackwardIncompatibleDropHint { .. } => false,
mir::StatementKind::Intrinsic(intrinsic) => {
!matches!(&**intrinsic, mir::NonDivergingIntrinsic::Assume(..))
}
}
}
pub fn forget(state: &mut State<'_>, local: mir::Local) {
for slot in state.iter_mut() {
if slot.value.is_some_and(|value| value.leans_on(local)) {
slot.value = None;
}
slot.order.forget(local);
if slot.same == Some(local) {
slot.same = None;
}
if slot.paired == Some(local) {
slot.paired = None;
}
if slot.spans == Some(local) {
slot.spans = None;
}
if slot.over.is_some_and(|(of, _)| of == local) {
slot.over = None;
}
}
if let Some(slot) = state.get_mut(local.as_usize()) {
*slot = Fact::default();
}
}
pub fn retire(state: &mut State<'_>, local: mir::Local) {
let held = state.get(local.as_usize()).copied().unwrap_or_default();
forget(state, local);
if let Some(slot) = state.get_mut(local.as_usize()) {
slot.same = held.same;
slot.extent = held.extent;
slot.paired = held.paired;
slot.spans = held.spans;
if matches!(held.value, Some(Value::Length(_))) {
slot.value = held.value;
}
}
}
pub fn unwind_to<'tcx>(
unwind: UnwindAction,
state: State<'tcx>,
work: &mut Work<'tcx>,
) {
if let UnwindAction::Cleanup(target) = unwind {
work.merge(target, state);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Step {
Deref,
Field(u32),
Variant(u32),
At(mir::Local),
}
const REACH: usize = 3;
const TRACKED: usize = 32;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Path {
pub base: mir::Local,
steps: [Option<Step>; REACH],
}
impl Path {
const fn step_of(element: mir::PlaceElem<'_>) -> Option<Step> {
Some(match element {
mir::ProjectionElem::Deref => Step::Deref,
mir::ProjectionElem::Field(field, _) => Step::Field(field.as_u32()),
mir::ProjectionElem::Downcast(_, variant) => {
Step::Variant(variant.as_u32())
}
mir::ProjectionElem::Index(local) => Step::At(local),
_ => return None,
})
}
fn of(place: &mir::Place<'_>) -> Option<Self> {
if place.projection.is_empty() || place.projection.len() > REACH {
return None;
}
let mut steps = [None; REACH];
for (slot, element) in steps.iter_mut().zip(place.projection) {
*slot = Some(Self::step_of(element)?);
}
Some(Self {
base: place.local,
steps,
})
}
pub fn under(
place: &mir::Place<'_>,
variant: Option<u32>,
field: u32,
) -> Option<Self> {
let mut steps = [None; REACH];
let mut at = 0usize;
for element in place.projection {
*steps.get_mut(at)? = Some(Self::step_of(element)?);
at = at.checked_add(1)?;
}
if let Some(variant) = variant {
*steps.get_mut(at)? = Some(Step::Variant(variant));
at = at.checked_add(1)?;
}
*steps.get_mut(at)? = Some(Step::Field(field));
Some(Self {
base: place.local,
steps,
})
}
pub const fn rebased(self, base: mir::Local) -> Self {
Self { base, ..self }
}
pub fn behind_pointer(self) -> bool {
self.steps.iter().flatten().any(|step| *step == Step::Deref)
}
pub fn indexed_by(self, local: mir::Local) -> bool {
self.steps
.iter()
.flatten()
.any(|step| *step == Step::At(local))
}
pub fn portable(self) -> bool {
!self
.steps
.iter()
.flatten()
.any(|step| matches!(step, Step::At(_)))
}
}
pub struct Places {
paths: Vec<Path>,
first: usize,
}
impl Places {
pub fn of<'tcx>(tcx: TyCtxt<'tcx>, mir: &mir::Body<'tcx>) -> Self {
let mut collect = Collect {
tcx,
found: Vec::new(),
};
mir::visit::Visitor::visit_body(&mut collect, mir);
Self {
paths: collect.found,
first: mir.local_decls.len(),
}
}
pub const fn len(&self) -> usize {
self.paths.len()
}
pub fn slot(&self, place: &mir::Place<'_>) -> Option<mir::Local> {
self.at(Path::of(place)?)
}
pub fn at(&self, path: Path) -> Option<mir::Local> {
let at = self.paths.iter().position(|held| *held == path)?;
Some(mir::Local::from_usize(self.first.saturating_add(at)))
}
pub fn path(&self, slot: mir::Local) -> Option<Path> {
self.paths
.get(slot.as_usize().checked_sub(self.first)?)
.copied()
}
pub fn each(&self) -> impl Iterator<Item = (mir::Local, Path)> + '_ {
self.paths.iter().enumerate().map(|(at, path)| {
(mir::Local::from_usize(self.first.saturating_add(at)), *path)
})
}
}
struct Collect<'tcx> {
tcx: TyCtxt<'tcx>,
found: Vec<Path>,
}
impl Collect<'_> {
fn add(&mut self, path: Path) {
if self.found.len() >= TRACKED || self.found.contains(&path) {
return;
}
self.found.push(path);
}
}
impl<'tcx> mir::visit::Visitor<'tcx> for Collect<'tcx> {
fn visit_place(
&mut self,
place: &mir::Place<'tcx>,
_: mir::visit::PlaceContext,
_: mir::Location,
) {
if let Some(path) = Path::of(place) {
self.add(path);
}
}
fn visit_assign(
&mut self,
place: &mir::Place<'tcx>,
rvalue: &mir::Rvalue<'tcx>,
location: mir::Location,
) {
self.super_assign(place, rvalue, location);
let mir::Rvalue::Aggregate(kind, fields) = rvalue else {
return;
};
let variant = match &**kind {
mir::AggregateKind::Tuple => None,
mir::AggregateKind::Adt(did, variant, ..) => {
let def = self.tcx.adt_def(*did);
if def.is_union() {
return;
}
def.is_enum().then(|| variant.as_u32())
}
_ => return,
};
for index in fields.indices() {
if let Some(path) = Path::under(place, variant, index.as_u32()) {
self.add(path);
}
}
}
}
pub fn sweep_base(state: &mut State<'_>, places: &Places, base: mir::Local) {
for (slot, path) in places.each() {
if path.base == base {
forget(state, slot);
}
}
}
pub fn sweep_indexed(
state: &mut State<'_>,
places: &Places,
local: mir::Local,
) {
for (slot, path) in places.each() {
if path.indexed_by(local) {
forget(state, slot);
}
}
}