use std::{
fmt::{self},
hash::Hash,
path::Path,
};
use flowistry_pdg::{CallString, RichLocation};
use internment::Intern;
use petgraph::{dot, graph::DiGraph};
use rustc_hash::FxHashSet;
use rustc_hir::def_id::DefId;
use rustc_index::IndexVec;
use rustc_middle::{
mir::{Body, HasLocalDecls, Local, LocalDecl, LocalDecls, Location, Place},
ty::{GenericArgsRef, Instance, TyCtxt},
};
use rustc_span::Span;
use rustc_utils::PlaceExt;
pub use flowistry_pdg::{SourceUse, TargetUse};
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct OneHopLocation {
pub location: RichLocation,
pub in_child: Option<(DefId, bool)>,
}
impl fmt::Display for OneHopLocation {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.location)?;
if let Some((_, is_start)) = self.in_child {
write!(f, " -> {}", if is_start { "start" } else { "end" })?;
}
Ok(())
}
}
impl From<RichLocation> for OneHopLocation {
fn from(location: RichLocation) -> Self {
OneHopLocation {
location,
in_child: None,
}
}
}
impl From<Location> for OneHopLocation {
fn from(at: Location) -> Self {
OneHopLocation {
location: RichLocation::Location(at),
in_child: None,
}
}
}
impl<T: Copy> From<&'_ T> for OneHopLocation
where
OneHopLocation: From<T>,
{
fn from(at: &'_ T) -> Self {
OneHopLocation::from(*at)
}
}
#[derive(Clone, Debug)]
pub struct DepNode<'tcx, Loc> {
pub place: Place<'tcx>,
pub at: Loc,
pub(crate) place_pretty: Option<Intern<String>>,
pub is_split: bool,
pub span: Span,
}
impl<Loc: PartialEq> PartialEq for DepNode<'_, Loc> {
fn eq(&self, other: &Self) -> bool {
let Self {
place,
at,
place_pretty: _,
span,
is_split,
} = self;
let eq = (place, at).eq(&(&other.place, &other.at));
if eq {
debug_assert_eq!(span, &other.span);
debug_assert_eq!(is_split, &other.is_split);
}
eq
}
}
impl<Loc: Eq> Eq for DepNode<'_, Loc> {}
impl<Loc: Hash> Hash for DepNode<'_, Loc> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
let Self {
place,
at,
place_pretty: _,
span: _,
is_split: _,
} = self;
(place, at).hash(state)
}
}
impl<'tcx> DepNode<'tcx, OneHopLocation> {
pub fn new(
place: Place<'tcx>,
at: OneHopLocation,
tcx: TyCtxt<'tcx>,
body: &Body<'tcx>,
context: DefId,
is_split: bool,
) -> Self {
let span = match at.location {
RichLocation::Location(loc) => {
let expanded_span = body
.stmt_at(loc)
.either(|s| s.source_info.span, |t| t.source_info.span);
tcx.sess.source_map().stmt_span(expanded_span, body.span)
}
RichLocation::Start | RichLocation::End => tcx.def_span(context),
};
DepNode {
place,
at,
place_pretty: place.to_string(tcx, body).map(Intern::new),
span,
is_split,
}
}
}
impl<'a, Loc> DepNode<'a, Loc> {
pub fn place_pretty(&self) -> Option<&str> {
self.place_pretty.map(|s| s.as_ref().as_str())
}
pub fn map_at<'b, Loc2, F: FnOnce(&'b Loc) -> Loc2>(&'b self, f: F) -> DepNode<'a, Loc2> {
DepNode {
place: self.place,
at: f(&self.at),
place_pretty: self.place_pretty,
span: self.span,
is_split: self.is_split,
}
}
}
impl<Loc: fmt::Display> fmt::Display for DepNode<'_, Loc> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.place_pretty() {
Some(s) => s.fmt(f)?,
None => write!(f, "{:?}", self.place)?,
};
write!(f, " @ {}", self.at)
}
}
#[derive(PartialEq, Eq, Hash, Clone, Copy, Debug)]
pub enum DepEdgeKind {
Control,
Data,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct DepEdge<Loc> {
pub kind: DepEdgeKind,
pub at: Loc,
pub source_use: SourceUse,
pub target_use: TargetUse,
}
impl<Loc> DepEdge<Loc> {
pub fn data(at: Loc, source_use: SourceUse, target_use: TargetUse) -> Self {
DepEdge {
kind: DepEdgeKind::Data,
at,
source_use,
target_use,
}
}
pub fn control(at: Loc, source_use: SourceUse, target_use: TargetUse) -> Self {
DepEdge {
kind: DepEdgeKind::Control,
at,
source_use,
target_use,
}
}
pub fn is_control(&self) -> bool {
self.kind == DepEdgeKind::Control
}
pub fn is_data(&self) -> bool {
self.kind == DepEdgeKind::Data
}
}
impl<Loc> DepEdge<Loc> {
pub fn map_at<'a, Loc2, F: FnOnce(&'a Loc) -> Loc2>(&'a self, f: F) -> DepEdge<Loc2> {
DepEdge {
kind: self.kind,
at: f(&self.at),
source_use: self.source_use,
target_use: self.target_use,
}
}
}
impl<Loc: fmt::Display> fmt::Display for DepEdge<Loc> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:?}\n@ {}", self.kind, self.at)
}
}
#[derive(Debug)]
pub struct DepGraph<'tcx> {
pub graph: DiGraph<DepNode<'tcx, CallString>, DepEdge<CallString>>,
}
impl Clone for DepGraph<'_> {
fn clone(&self) -> Self {
DepGraph {
graph: self.graph.clone(),
}
}
fn clone_from(&mut self, source: &Self) {
self.graph.clone_from(&source.graph);
}
}
impl<'tcx> DepGraph<'tcx> {
pub fn new(graph: DiGraph<DepNode<'tcx, CallString>, DepEdge<CallString>>) -> Self {
Self { graph }
}
}
impl DepGraph<'_> {
pub fn generate_graphviz(&self, path: impl AsRef<Path>) -> anyhow::Result<()> {
let graph_dot = format!(
"{}",
dot::Dot::with_attr_getters(
&self.graph,
&[],
&|_, _| "fontname=\"Courier New\"".to_string(),
&|_, (_, _)| "fontname=\"Courier New\",shape=box".to_string(),
)
);
rustc_utils::mir::body::run_dot(path.as_ref(), &graph_dot.into_bytes())
}
}
#[derive(Debug, Clone)]
pub struct PartialGraph<'tcx, K> {
pub(crate) nodes: FxHashSet<DepNode<'tcx, OneHopLocation>>,
pub(crate) edges: FxHashSet<(
DepNode<'tcx, OneHopLocation>,
DepNode<'tcx, OneHopLocation>,
DepEdge<OneHopLocation>,
)>,
pub(crate) generics: GenericArgsRef<'tcx>,
pub(crate) def_id: DefId,
arg_count: usize,
local_decls: IndexVec<Local, LocalDecl<'tcx>>,
pub(crate) k: K,
pub(crate) inlined_calls: Vec<(Location, Instance<'tcx>, K, Vec<GraphConnectionPoint<'tcx>>)>,
}
type GraphConnectionPoint<'tcx> = (DepNode<'tcx, OneHopLocation>, DepEdge<OneHopLocation>);
impl<'tcx, K> HasLocalDecls<'tcx> for PartialGraph<'tcx, K> {
fn local_decls(&self) -> &LocalDecls<'tcx> {
&self.local_decls
}
}
impl<'tcx, K> PartialGraph<'tcx, K> {
pub fn new(
generics: GenericArgsRef<'tcx>,
def_id: DefId,
arg_count: usize,
local_decls: &LocalDecls<'tcx>,
k: K,
) -> Self {
Self {
nodes: Default::default(),
edges: Default::default(),
generics,
def_id,
arg_count,
local_decls: local_decls.to_owned(),
inlined_calls: Default::default(),
k,
}
}
pub(crate) fn parentable_srcs<'a>(&'a self) -> FxHashSet<(DepNode<'tcx, bool>, Option<u8>)> {
self.edges
.iter()
.filter(|&(n, _, _)| n.at.location.is_start())
.map(|(n, _, _)| {
n.map_at(|_| {
assert!(n.at.in_child.is_none());
true
})
})
.filter_map(move |a| {
let as_arg = as_arg(&a, self.arg_count)?;
Some((a, as_arg))
})
.collect()
}
pub(crate) fn parentable_dsts<'a>(&'a self) -> FxHashSet<(DepNode<'tcx, bool>, Option<u8>)> {
self.edges
.iter()
.filter(|&(_, n, _)| n.at.location.is_end())
.map(|(_, n, _)| {
assert!(n.at.in_child.is_none());
n.map_at(|_| false)
})
.filter_map(move |a| {
let arg = as_arg(&a, self.arg_count)?;
Some((a, arg))
})
.collect()
}
}
fn as_arg<Loc>(node: &DepNode<'_, Loc>, arg_count: usize) -> Option<Option<u8>> {
let local = node.place.local.as_usize();
if node.place.local == rustc_middle::mir::RETURN_PLACE {
Some(None)
} else if local > 0 && (local - 1) < arg_count {
Some(Some(node.place.local.as_u32() as u8 - 1))
} else {
None
}
}