use crate::analysis::path::{CallBinding, CallConstArg, PathTree};
use crate::compat::{FxHashMap, FxHashSet};
use crate::helpers::mir_utils;
use crate::limit::{HINTED_CALLEE_BLOCK_LIMIT, LOCAL_CALLEE_BLOCK_LIMIT};
#[cfg(rapx_has_attr_ir)]
use rustc_attr_ir::InlineAttr;
#[cfg(not(rapx_has_attr_ir))]
use rustc_hir::attrs::InlineAttr;
use rustc_hir::def_id::DefId;
use rustc_middle::{
mir::{BasicBlock, Body, Operand, Terminator, TerminatorKind},
ty::{EarlyBinder, Instance, TyCtxt, TypingEnv},
};
pub(crate) struct CallSiteIndex<'tcx> {
tcx: TyCtxt<'tcx>,
blocks: Vec<(DefId, usize)>,
bodies: FxHashMap<usize, Body<'tcx>>,
bindings: FxHashMap<usize, CallBinding>,
pub(crate) entries: FxHashMap<usize, Vec<usize>>,
}
impl<'tcx> CallSiteIndex<'tcx> {
pub(crate) fn new(tcx: TyCtxt<'tcx>, root: DefId, body: Option<Body<'tcx>>) -> Self {
let block_count = body
.as_ref()
.map(|body| body.basic_blocks.len())
.unwrap_or_else(|| tcx.optimized_mir(root).basic_blocks.len());
let blocks = (0..block_count)
.map(|index| (root, index))
.collect();
let mut bodies = FxHashMap::default();
if let Some(body) = body {
bodies.insert(0, body);
}
Self {
tcx,
blocks,
bodies,
bindings: FxHashMap::default(),
entries: FxHashMap::default(),
}
}
pub(crate) fn attach_metadata(self, tree: &mut PathTree) {
let mut parents = FxHashMap::default();
for (&call_block, entries) in &self.entries {
let parent = self.blocks[call_block].0;
for &entry in entries {
parents.insert(entry, parent);
}
}
let expanded_calls = self.entries.keys().copied().collect();
tree.set_block_metadata(self.blocks, self.bindings, parents, expanded_calls);
}
pub(crate) fn blocks(&self) -> &[(DefId, usize)] {
&self.blocks
}
pub(crate) fn tcx(&self) -> TyCtxt<'tcx> {
self.tcx
}
pub(crate) fn body(&self, block: usize) -> &Body<'tcx> {
let (def_id, local) = self.blocks[block];
self.bodies
.get(&(block - local))
.unwrap_or_else(|| self.tcx.optimized_mir(def_id))
}
pub(crate) fn has_body_override(&self, entry: usize) -> bool {
self.bodies.contains_key(&entry)
}
fn terminator(&self, block: usize) -> Option<&Terminator<'tcx>> {
let (_, local) = self.blocks[block];
self.body(block)
.basic_blocks
.get(BasicBlock::from(local))
.and_then(|data| data.terminator.as_ref())
}
pub(crate) fn normal_successor(&self, block: usize) -> Option<usize> {
let (_, local) = self.blocks[block];
let TerminatorKind::Call {
target: Some(target),
..
} = &self.terminator(block)?.kind
else {
return None;
};
Some(block - local + target.as_usize())
}
pub(crate) fn collect(&mut self) {
let mut expanded = FxHashSet::default();
expanded.insert(self.blocks[0].0);
let mut next_block = 0;
loop {
let block_count = self.blocks.len();
let mut pending = Vec::new();
for block in next_block..block_count {
let Some(terminator) = self.terminator(block) else {
continue;
};
let TerminatorKind::Call {
func,
target: Some(_),
..
} = &terminator.kind
else {
continue;
};
let Some(base_callee) = mir_utils::dep_callee_def_id(func) else {
continue;
};
let caller = self.blocks[block].0;
let cross_crate = base_callee.as_local().is_none();
let is_trait_method = self
.tcx
.opt_associated_item(base_callee)
.and_then(|item| item.trait_container(self.tcx))
.is_some();
let modeled = crate::verify::call_summary::effects::is_symbolically_modeled(Some(
base_callee,
));
let resolve_to_impl = !cross_crate || !is_trait_method || !modeled;
if caller.as_local().is_none() && !(is_trait_method && resolve_to_impl) {
continue;
}
let callee = if resolve_to_impl {
mir_utils::dep_callee_resolved_def_id(self.tcx, caller, func)
.unwrap_or(base_callee)
} else {
base_callee
};
if is_trait_method {
let is_virtual = mir_utils::resolve_callee_instance(self.tcx, caller, func)
.is_some_and(|instance| mir_utils::instance_def_id(self.tcx, instance).1);
if is_virtual {
if !modeled {
for candidate in mir_utils::trait_impl_candidates(self.tcx, base_callee)
.into_iter()
.take(crate::limit::trait_impl_limit())
{
if self.tcx.is_mir_available(candidate)
&& !expanded.contains(&candidate)
{
pending.push((block, candidate, None));
}
}
}
continue;
}
}
let instance =
if cross_crate && crate::def_ids::other_api::is_slice_get(Some(base_callee)) {
mir_utils::resolve_callee_instance(self.tcx, caller, func)
} else {
None
};
let cross_crate_trait_default =
cross_crate && is_trait_method && self.tcx.is_mir_available(base_callee);
if crate::def_ids::intrinsic_api::is_intrinsic(self.tcx, callee)
|| !self.tcx.is_mir_available(callee)
|| expanded.contains(&callee)
|| crate::verify::call_summary::effects::is_symbolically_modeled(Some(callee))
|| crate::def_ids::other_api::is_ub_check(self.tcx, callee)
|| cross_crate_trait_default
{
continue;
}
let small_local = !cross_crate && {
let limit = if matches!(
self.tcx.codegen_fn_attrs(callee).inline,
InlineAttr::Hint | InlineAttr::Always
) {
HINTED_CALLEE_BLOCK_LIMIT
} else {
LOCAL_CALLEE_BLOCK_LIMIT
};
self.tcx.optimized_mir(callee).basic_blocks.len() <= limit
};
if cross_crate || small_local {
pending.push((block, callee, instance));
}
}
next_block = block_count;
if pending.is_empty() {
break;
}
for (_, callee, _) in &pending {
expanded.insert(*callee);
}
for (block, callee, instance) in pending {
self.insert(block, callee, instance);
}
}
for entries in self.entries.values_mut() {
entries.sort_unstable();
}
}
fn insert(&mut self, call_block: usize, callee: DefId, instance: Option<Instance<'tcx>>) {
let Some(terminator) = self.terminator(call_block) else {
return;
};
let TerminatorKind::Call {
args, destination, ..
} = &terminator.kind
else {
return;
};
let arg_locals = args
.iter()
.filter_map(|arg| arg.node.place().map(|place| place.local.as_usize()))
.collect();
let arg_consts = args
.iter()
.enumerate()
.filter_map(|(index, arg)| {
if let Some(value) = mir_utils::operand_const_u64(&arg.node) {
return Some((index + 1, CallConstArg::Literal(value)));
}
let Operand::Constant(constant) = &arg.node else {
return None;
};
let rustc_middle::mir::Const::Ty(_, value) = constant.const_ else {
return None;
};
matches!(value.kind(), rustc_middle::ty::ConstKind::Param(_)).then(|| {
(
index + 1,
CallConstArg::Symbol(format!("{:?}", constant.const_)),
)
})
})
.collect();
let arg_field_offsets = args
.iter()
.enumerate()
.filter_map(|(index, arg)| {
let Operand::Constant(constant) = &arg.node else {
return None;
};
mir_utils::offset_of_container(self.tcx, &constant.const_).map(|_| index + 1)
})
.collect();
let binding = CallBinding {
arg_locals,
arg_consts,
arg_field_offsets,
dest_local: destination.local.as_usize(),
};
let parent = self.blocks[call_block].0;
let instantiated_body = instance.map(|instance| {
let mir = self.tcx.optimized_mir(instance.def_id());
instance.instantiate_mir_and_normalize_erasing_regions(
self.tcx,
TypingEnv::post_analysis(self.tcx, parent),
EarlyBinder::bind(self.tcx, mir.clone()),
)
});
let block_count = instantiated_body
.as_ref()
.map(|body| body.basic_blocks.len())
.unwrap_or_else(|| self.tcx.optimized_mir(callee).basic_blocks.len());
let entry = self.blocks.len();
self.blocks
.extend((0..block_count).map(|local| (callee, local)));
if let Some(body) = instantiated_body {
self.bodies.insert(entry, body);
}
self.bindings.insert(entry, binding);
self.entries.entry(call_block).or_default().push(entry);
}
}