use super::PathTree;
use crate::compat::{FxHashMap, FxHashSet};
use crate::graphs::{
cfg::{CfgBlock, ControlFlowGraph},
scc::{Scc, SccInfo},
};
use rustc_middle::{
mir::{
AggregateKind, BasicBlock, Local, Operand, Rvalue, StatementKind, SwitchTargets,
Terminator, TerminatorKind, UnwindAction,
},
ty::{TyCtxt, TyKind, TypingEnv},
};
use rustc_span::def_id::DefId;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
const WHOLE_CFG_PATH_LIMIT: usize = 4000;
const WHOLE_CFG_PATH_DEPTH_LIMIT: usize = 256;
const SCC_PATH_CACHE_LIMIT: usize = 2048;
const SCC_MAX_DEPTH: usize = 128;
const SCC_MAX_SEEN_PATHS: usize = 128;
const SCC_MAX_PATH_LEN: usize = 200;
fn check_postfix_segment(
path: &[usize],
enter: usize,
segment_counts: &mut FxHashMap<Vec<usize>, usize>,
max_repeats: usize,
) -> bool {
let segment = extract_segment(path, enter);
let count = segment_counts.entry(segment).or_insert(0);
*count += 1;
*count == 1 || *count - 1 <= max_repeats
}
fn extract_segment(path: &[usize], enter: usize) -> Vec<usize> {
let prev_pos = path[..path.len() - 1]
.iter()
.rposition(|&node| node == enter)
.unwrap_or(0);
path[prev_pos + 1..path.len() - 1].to_vec()
}
#[derive(Clone, Debug)]
pub struct SccPath {
pub blocks: Vec<usize>,
pub exit_successors: Vec<usize>,
}
#[derive(Clone, Debug, Default)]
pub struct BlockConstantInfo {
pub assigned_locals: FxHashSet<usize>,
pub constants: FxHashMap<usize, usize>,
pub constraint_copies: FxHashMap<usize, usize>,
}
#[derive(Clone, Debug, Default)]
pub struct DiscriminantInfo {
pub source_of: FxHashMap<usize, usize>,
pub variant_count_of: FxHashMap<usize, usize>,
}
#[derive(Clone)]
pub struct PathGraph<'tcx> {
pub cfg: ControlFlowGraph<'tcx>,
pub block_info: Vec<BlockConstantInfo>,
pub disc_info: DiscriminantInfo,
}
impl<'tcx> PathGraph<'tcx> {
pub fn new(tcx: TyCtxt<'tcx>, def_id: DefId) -> PathGraph<'tcx> {
let body = tcx.optimized_mir(def_id);
let basicblocks = &body.basic_blocks;
let mut cfg_blocks = Vec::<CfgBlock>::new();
let mut block_info = Vec::new();
let mut disc_info = DiscriminantInfo::default();
for i in 0..basicblocks.len() {
let bb = &basicblocks[BasicBlock::from(i)];
let mut cfg_block = CfgBlock::new(i, bb.is_cleanup);
let mut info = BlockConstantInfo::default();
for stmt in &bb.statements {
if let StatementKind::Assign(assign) = &stmt.kind {
let (place, rvalue) = &**assign;
let dest = place.local.as_usize();
info.assigned_locals.insert(dest);
match rvalue {
Rvalue::Use(Operand::Constant(c), ..) => {
let typing_env = TypingEnv::post_analysis(tcx, def_id);
let val = match c.const_.ty().kind() {
TyKind::Bool => c
.const_
.try_eval_bool(tcx, typing_env)
.map(|b| if b { 1 } else { 0 }),
TyKind::Int(_) | TyKind::Uint(_) => {
c.const_.try_eval_bits(tcx, typing_env).map(|v| v as usize)
}
_ => None,
};
if let Some(val) = val {
info.constants.insert(dest, val);
}
}
Rvalue::Use(Operand::Copy(src) | Operand::Move(src), ..) => {
info.constraint_copies.insert(dest, src.local.as_usize());
}
Rvalue::Discriminant(rv_place) => {
disc_info.source_of.insert(dest, rv_place.local.as_usize());
let src_local = rv_place.local.as_usize();
if !disc_info.variant_count_of.contains_key(&src_local) {
let src_ty = body.local_decls[rv_place.local].ty;
if let TyKind::Adt(adt_def, _) = src_ty.kind() {
let num = adt_def.variants().len();
if num > 0 {
disc_info.variant_count_of.insert(src_local, num);
}
}
}
}
Rvalue::Aggregate(kind, _) => {
if let AggregateKind::Adt(_, variant_idx, _, _, _) = kind.as_ref() {
let discr = variant_idx.as_usize();
info.constants.insert(dest, discr);
if !disc_info.variant_count_of.contains_key(&dest) {
let dest_ty = body.local_decls[place.local].ty;
if let TyKind::Adt(adt_def, _) = dest_ty.kind() {
let num = adt_def.variants().len();
if num > 0 {
disc_info.variant_count_of.insert(dest, num);
}
}
}
}
}
_ => {}
}
}
}
let Some(terminator) = &bb.terminator else {
continue;
};
match terminator.kind.clone() {
TerminatorKind::Goto { ref target } => {
cfg_block.add_next(target.as_usize());
}
TerminatorKind::SwitchInt {
discr: _,
ref targets,
} => {
for (_, ref target) in targets.iter() {
cfg_block.add_next(target.as_usize());
}
cfg_block.add_next(targets.otherwise().as_usize());
}
TerminatorKind::Drop {
place: _,
target,
unwind,
replace: _,
drop: _,
#[cfg(not(rapx_rustc_ge_198))]
async_fut: _,
} => {
cfg_block.add_next(target.as_usize());
if let UnwindAction::Cleanup(target) = unwind {
cfg_block.add_next(target.as_usize());
}
}
TerminatorKind::Call {
ref target,
ref unwind,
..
} => {
if let Some(tt) = target {
cfg_block.add_next(tt.as_usize());
}
if let UnwindAction::Cleanup(tt) = unwind {
cfg_block.add_next(tt.as_usize());
}
}
TerminatorKind::Assert {
cond: _,
expected: _,
msg: _,
ref target,
ref unwind,
} => {
cfg_block.add_next(target.as_usize());
if let UnwindAction::Cleanup(target) = unwind {
cfg_block.add_next(target.as_usize());
}
}
TerminatorKind::Yield {
value: _,
ref resume,
resume_arg: _,
ref drop,
} => {
cfg_block.add_next(resume.as_usize());
if let Some(target) = drop {
cfg_block.add_next(target.as_usize());
}
}
TerminatorKind::FalseEdge {
ref real_target,
imaginary_target: _,
} => {
cfg_block.add_next(real_target.as_usize());
}
TerminatorKind::FalseUnwind {
ref real_target,
unwind: _,
} => {
cfg_block.add_next(real_target.as_usize());
}
TerminatorKind::InlineAsm {
template: _,
operands: _,
options: _,
line_spans: _,
ref unwind,
targets,
asm_macro: _,
} => {
for target in targets {
cfg_block.add_next(target.as_usize());
}
if let UnwindAction::Cleanup(target) = unwind {
cfg_block.add_next(target.as_usize());
}
}
_ => {}
}
cfg_blocks.push(cfg_block);
block_info.push(info);
}
let cfg = ControlFlowGraph::new(def_id, tcx, cfg_blocks);
PathGraph {
cfg,
block_info,
disc_info,
}
}
pub fn find_scc(&mut self) {
self.cfg.find_scc();
self.populate_all_child_sccs();
}
pub fn def_id(&self) -> DefId {
self.cfg.def_id
}
pub fn tcx(&self) -> TyCtxt<'tcx> {
self.cfg.tcx
}
pub fn cfg_block(&self, index: usize) -> &CfgBlock {
self.cfg.block(index)
}
pub fn cfg_block_mut(&mut self, index: usize) -> &mut CfgBlock {
self.cfg.block_mut(index)
}
pub fn terminator(&self, index: usize) -> Option<&Terminator<'tcx>> {
self.cfg.terminator(index)
}
pub fn is_cleanup_block(&self, index: usize) -> bool {
self.cfg
.blocks
.get(index)
.map(|b| b.is_cleanup)
.unwrap_or(false)
}
pub fn check_transition(
&self,
cur: usize,
next: usize,
constraints: &mut FxHashMap<usize, usize>,
) -> bool {
if cur >= self.cfg.blocks.len() || next >= self.cfg.blocks.len() {
return false;
}
if let Some(info) = self.block_info.get(cur) {
for local in &info.assigned_locals {
if let Some(&src) = info.constraint_copies.get(local) {
if let Some(&src_val) = constraints.get(&src) {
constraints.insert(*local, src_val);
continue;
}
if let Some(&dst_val) = constraints.get(local) {
constraints.insert(src, dst_val);
constraints.insert(*local, dst_val);
continue;
}
}
if let Some(&val) = info.constants.get(local) {
constraints.insert(*local, val);
continue;
}
constraints.remove(local);
}
}
let successors = &self.cfg.block(cur).next;
if !successors.contains(&next) {
if !self.is_unwind_target(cur, next) {
return false;
}
}
if !self.check_switch_transition(cur, next, constraints) {
return false;
}
true
}
fn check_switch_transition(
&self,
cur: usize,
next: usize,
constraints: &mut FxHashMap<usize, usize>,
) -> bool {
let Some(terminator) = self.cfg.terminator(cur) else {
return true;
};
match &terminator.kind {
TerminatorKind::SwitchInt { discr, targets } => {
let discr_local = discr.place().map(|p| p.local.as_usize());
let constraint_local = discr_local
.and_then(|l| self.disc_info.source_of.get(&l).copied())
.or(discr_local);
let all_targets: FxHashSet<usize> = targets
.iter()
.map(|(_, bb)| bb.as_usize())
.chain(std::iter::once(targets.otherwise().as_usize()))
.collect();
if !all_targets.contains(&next) {
return false;
}
let const_val = match discr {
Operand::Constant(c) => c
.const_
.try_eval_target_usize(
self.cfg.tcx,
TypingEnv::post_analysis(self.cfg.tcx, self.cfg.def_id),
)
.map(|v| v as usize),
_ => None,
};
if let Some(val) = const_val {
let expected = resolve_switch_target(targets, val as u128);
if next != expected {
return false;
}
if let Some(local) = constraint_local {
constraints.insert(local, val);
}
return true;
}
if let Some(local) = constraint_local {
if let Some(&known_val) = constraints.get(&local) {
let expected = resolve_switch_target(targets, known_val as u128);
if next != expected {
return false;
}
return true;
}
}
if next == targets.otherwise().as_usize() {
if let Some(local) = constraint_local {
if let Some(&num_variants) = self.disc_info.variant_count_of.get(&local) {
let all_covered = (0..num_variants)
.all(|v| targets.iter().any(|(tv, _)| tv == v as u128));
if all_covered {
return false;
}
}
}
}
self.learn_constraint_with_backprop(
cur,
constraint_local,
&targets,
next,
constraints,
);
true
}
_ => true,
}
}
fn learn_constraint_with_backprop(
&self,
cur: usize,
constraint_local: Option<usize>,
targets: &SwitchTargets,
next: usize,
constraints: &mut FxHashMap<usize, usize>,
) {
let Some(local) = constraint_local else {
return;
};
let Some((val, _)) = targets.iter().find(|(_, bb)| bb.as_usize() == next) else {
if let Some(inferred) = self.infer_otherwise_value(targets, local) {
constraints.insert(local, inferred);
self.backprop_constraint(cur, local, inferred, constraints);
}
return;
};
let val = val as usize;
constraints.insert(local, val);
self.backprop_constraint(cur, local, val, constraints);
}
fn backprop_constraint(
&self,
cur: usize,
local: usize,
val: usize,
constraints: &mut FxHashMap<usize, usize>,
) {
let Some(info) = self.block_info.get(cur) else {
return;
};
let mut current = local;
while let Some(&src) = info.constraint_copies.get(¤t) {
if current == src {
break;
}
constraints.insert(src, val);
current = src;
}
}
fn infer_otherwise_value(&self, targets: &SwitchTargets, discr_local: usize) -> Option<usize> {
let body = self.cfg.tcx.optimized_mir(self.cfg.def_id);
let discr_ty = body.local_decls[Local::from_usize(discr_local)].ty;
let possible_values: Vec<usize> = match discr_ty.kind() {
TyKind::Bool => vec![0, 1],
TyKind::Adt(adt_def, _) if adt_def.is_enum() => (0..adt_def.variants().len()).collect(),
_ => return None,
};
let explicit_values: FxHashSet<usize> = targets.iter().map(|(v, _)| v as usize).collect();
let remaining: Vec<usize> = possible_values
.into_iter()
.filter(|v| !explicit_values.contains(v))
.collect();
if remaining.len() == 1 {
Some(remaining[0])
} else {
None
}
}
fn is_unwind_target(&self, cur: usize, next: usize) -> bool {
let Some(terminator) = self.cfg.terminator(cur) else {
return false;
};
let unwind = match &terminator.kind {
TerminatorKind::Call { unwind, .. }
| TerminatorKind::Drop { unwind, .. }
| TerminatorKind::Assert { unwind, .. } => unwind,
_ => return false,
};
if let UnwindAction::Cleanup(target) = unwind {
return target.as_usize() == next;
}
false
}
fn populate_child_sccs(&mut self, enter: usize) {
let nodes: Vec<usize> = self.cfg.block(enter).scc.nodes.iter().cloned().collect();
let mut child_enters = Vec::new();
let mut seen = FxHashSet::default();
for node in nodes {
if let Some(block) = self.cfg.blocks.get(node) {
let node_enter = block.scc.enter;
let non_trivial = !block.scc.nodes.is_empty();
if node_enter != enter && non_trivial && seen.insert(node_enter) {
child_enters.push(node_enter);
}
}
}
self.cfg.block_mut(enter).scc.child_sccs = child_enters;
let child_count = self.cfg.block(enter).scc.child_sccs.len();
for i in 0..child_count {
let child_enter = self.cfg.block(enter).scc.child_sccs[i];
self.populate_child_sccs(child_enter);
}
}
fn populate_all_child_sccs(&mut self) {
let mut visited = FxHashSet::default();
let block_count = self.cfg.blocks.len();
for i in 0..block_count {
let scc = &self.cfg.block(i).scc;
let enter = scc.enter;
if scc.nodes.is_empty() || !visited.insert(enter) {
continue;
}
self.populate_child_sccs(enter);
}
}
}
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Default)]
pub struct ConstraintHash(u64);
impl ConstraintHash {
fn from_path(path: &[usize], graph: &PathGraph<'_>) -> Self {
let mut hasher = DefaultHasher::new();
let mut constraints: FxHashMap<usize, usize> = FxHashMap::default();
for &block in path.iter() {
if let Some(info) = graph.block_info.get(block) {
for local in &info.assigned_locals {
if let Some(&src) = info.constraint_copies.get(local) {
if let Some(&src_val) = constraints.get(&src) {
constraints.insert(*local, src_val);
continue;
}
if let Some(&dst_val) = constraints.get(local) {
constraints.insert(src, dst_val);
constraints.insert(*local, dst_val);
continue;
}
}
if let Some(&val) = info.constants.get(local) {
constraints.insert(*local, val);
continue;
}
constraints.remove(local);
}
}
}
let mut entries: Vec<(usize, usize)> = constraints.into_iter().collect();
entries.sort();
entries.hash(&mut hasher);
ConstraintHash(hasher.finish())
}
}
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct SccKey {
pub entry: usize,
pub repeat: usize,
pub constraint: ConstraintHash,
}
pub struct PathEnumerator<'g, 'tcx> {
graph: &'g PathGraph<'tcx>,
scc_paths: FxHashMap<SccKey, Vec<SccPath>>,
visited_sccs: FxHashSet<SccKey>,
}
impl<'g, 'tcx> PathEnumerator<'g, 'tcx> {
pub fn new(graph: &'g PathGraph<'tcx>) -> Self {
PathEnumerator {
graph,
scc_paths: FxHashMap::default(),
visited_sccs: FxHashSet::default(),
}
}
pub fn enumerate_paths(&mut self) -> PathTree {
self.enumerate_paths_repeat(0)
}
pub fn enumerate_paths_repeat(&mut self, postfix_repeat: usize) -> PathTree {
let mut tree = PathTree::new();
if self.graph.cfg.blocks.is_empty() {
return tree;
}
self.collect_whole_cfg_paths(0, &mut vec![0], &mut tree, 0, postfix_repeat, &FxHashMap::default());
tree
}
pub fn find_scc_paths_repeat(
&mut self,
start: usize,
scc: &SccInfo,
postfix_repeat: usize,
) -> Vec<SccPath> {
let cache_key = SccKey {
entry: scc.enter,
repeat: postfix_repeat,
constraint: ConstraintHash::default(),
};
if let Some(cached) = self.scc_paths.get(&cache_key) {
return cached.clone();
}
let mut out = Vec::new();
let mut seen: FxHashSet<Vec<usize>> = FxHashSet::default();
let mut path = vec![start];
let mut segment_counts = FxHashMap::default();
self.dfs_scc_tree(
scc,
start,
&mut path,
&mut segment_counts,
postfix_repeat,
&mut out,
&mut seen,
0,
);
if self.scc_paths.len() >= SCC_PATH_CACHE_LIMIT {
self.scc_paths.clear();
}
self.scc_paths.insert(cache_key, out.clone());
out
}
#[allow(clippy::too_many_arguments)]
fn dfs_scc_tree(
&mut self,
scc: &SccInfo,
cur: usize,
path: &mut Vec<usize>,
segment_counts: &mut FxHashMap<Vec<usize>, usize>,
postfix_repeat: usize,
out: &mut Vec<SccPath>,
seen_paths: &mut FxHashSet<Vec<usize>>,
depth: usize,
) {
if depth > SCC_MAX_DEPTH {
return;
}
if out.len() >= SCC_MAX_SEEN_PATHS {
return;
}
if path.len() > SCC_MAX_PATH_LEN {
return;
}
if cur != scc.enter && !scc.nodes.contains(&cur) {
return;
}
if cur == scc.enter && path.len() > 1 {
if !check_postfix_segment(path, scc.enter, segment_counts, postfix_repeat) {
if (postfix_repeat > 0 || segment_counts.len() > 1)
&& scc.exits.iter().any(|e| e.exit == cur)
{
self.record_unique_path(path, scc, out, seen_paths);
}
return;
}
}
if scc.exits.iter().any(|e| e.exit == cur) {
self.record_unique_path(path, scc, out, seen_paths);
}
let is_child = scc.child_sccs.contains(&cur);
if is_child {
let ctx = self.constraint_context(path);
if !self.visited_sccs.insert(SccKey { entry: cur, repeat: 0, constraint: ctx }) {
return;
}
let child_scc = self.graph.cfg_block(cur).scc.clone();
let child_paths = self.find_scc_paths_repeat(cur, &child_scc, postfix_repeat);
for child_path in &child_paths {
let orig_len = path.len();
if child_path.blocks.len() > 1 {
path.extend(&child_path.blocks[1..]);
}
let mut branch_counts = segment_counts.clone();
for &next in &child_path.exit_successors {
path.push(next);
self.dfs_scc_tree(
scc,
next,
path,
&mut branch_counts,
postfix_repeat,
out,
seen_paths,
depth + 1,
);
path.pop();
}
path.truncate(orig_len);
}
return;
}
let successors: Vec<usize> = self.graph.cfg.block(cur).next.iter().copied().collect();
let saved_counts = segment_counts.clone();
for next in successors {
if next != scc.enter && !scc.nodes.contains(&next) {
self.record_unique_path(path, scc, out, seen_paths);
continue;
}
let mut branch_counts = saved_counts.clone();
path.push(next);
self.dfs_scc_tree(
scc,
next,
path,
&mut branch_counts,
postfix_repeat,
out,
seen_paths,
depth + 1,
);
path.pop();
}
}
fn constraint_context(&self, path: &[usize]) -> ConstraintHash {
ConstraintHash::from_path(path, self.graph)
}
fn collect_whole_cfg_paths(
&mut self,
current: usize,
path: &mut Vec<usize>,
tree: &mut PathTree,
depth: usize,
postfix_repeat: usize,
constraints: &FxHashMap<usize, usize>,
) {
if current >= self.graph.cfg.blocks.len() {
return;
}
if depth > WHOLE_CFG_PATH_DEPTH_LIMIT || tree.len() >= WHOLE_CFG_PATH_LIMIT {
return;
}
let scc_info = self.graph.cfg_block(current).scc.clone();
let is_scc = current == scc_info.enter && !scc_info.nodes.is_empty();
if is_scc {
let scc = self.sort_scc_tree(&scc_info);
let segments = self.find_scc_paths_repeat(current, &scc, postfix_repeat);
if segments.is_empty() {
tree.insert(path);
return;
}
for seg in segments {
if tree.len() >= WHOLE_CFG_PATH_LIMIT {
break;
}
let orig_len = path.len();
let mut seg_constraints = constraints.clone();
let mut reachable = true;
if seg.blocks.len() > 1 {
for i in 0..seg.blocks.len() - 1 {
if !self.graph.check_transition(seg.blocks[i], seg.blocks[i + 1], &mut seg_constraints) {
reachable = false;
break;
}
}
if reachable {
path.extend_from_slice(&seg.blocks[1..]);
}
}
if reachable {
if seg.exit_successors.is_empty() {
tree.insert(path);
} else {
for &next in &seg.exit_successors {
let mut next_constraints = seg_constraints.clone();
let last = *path.last().unwrap();
if self.graph.check_transition(last, next, &mut next_constraints) {
path.push(next);
self.collect_whole_cfg_paths(next, path, tree, depth + 1, postfix_repeat, &next_constraints);
path.pop();
}
}
}
}
path.truncate(orig_len);
}
return;
}
let successors: Vec<usize> = self.graph.cfg_block(current).next.iter().copied().collect();
if successors.is_empty() {
tree.insert(path);
return;
}
for next in successors {
let mut next_constraints = constraints.clone();
if self.graph.check_transition(current, next, &mut next_constraints) {
path.push(next);
self.collect_whole_cfg_paths(next, path, tree, depth + 1, postfix_repeat, &next_constraints);
path.pop();
}
}
}
fn sort_scc_tree(&self, scc: &SccInfo) -> SccInfo {
self.graph.cfg_block(scc.enter).scc.clone()
}
fn record_unique_path(
&self,
path: &[usize],
scc: &SccInfo,
out: &mut Vec<SccPath>,
seen_paths: &mut FxHashSet<Vec<usize>>,
) {
if !seen_paths.insert(path.to_vec()) {
return;
}
let exit_successors = self.compute_exit_successors(path, scc);
out.push(SccPath {
blocks: path.to_vec(),
exit_successors,
});
}
fn compute_exit_successors(&self, path: &[usize], scc: &SccInfo) -> Vec<usize> {
let Some(&last) = path.last() else {
return vec![];
};
scc.exits
.iter()
.filter(|e| e.exit == last)
.map(|e| e.to)
.filter(|&n| {
!scc.child_sccs
.contains(&self.graph.cfg.block(n).scc.enter())
})
.collect()
}
}
fn resolve_switch_target(targets: &SwitchTargets, val: u128) -> usize {
targets
.iter()
.find(|(v, _)| *v == val)
.map(|(_, bb)| bb.as_usize())
.unwrap_or_else(|| targets.otherwise().as_usize())
}