use panicgraph::{Category, Termination, util::Map};
use rustc_hir::{def::DefKind, def_id::DefId};
use rustc_middle::{
middle::codegen_fn_attrs::CodegenFnAttrFlags,
ty::{self, TyCtxt},
};
#[derive(Debug, Clone, Copy)]
pub struct Sink {
first: (Category, Termination),
second: Option<(Category, Termination)>,
}
impl Sink {
pub fn raises(self) -> impl Iterator<Item = (Category, Termination)> {
std::iter::once(self.first).chain(self.second)
}
const fn is_only(self, category: Category) -> bool {
self.second.is_none()
&& matches!(self.first, (c, _) if c as u8 == category as u8)
}
}
const fn unwind(category: Category) -> Sink {
Sink {
first: (category, Termination::Unwind),
second: None,
}
}
const fn abort(category: Category) -> Sink {
Sink {
first: (category, Termination::Abort),
second: None,
}
}
const fn funnel(first: Sink, second: Sink) -> Sink {
Sink {
first: first.first,
second: Some(second.first),
}
}
const EXACT: &[(Sink, &[(&str, &str)])] = &[
(
unwind(Category::Explicit),
&[
("core", "panicking::panic"),
("core", "panicking::panic_fmt"),
("core", "panicking::panic_str"),
("core", "panicking::panic_explicit"),
("core", "panicking::panic_display"),
("core", "panicking::assert_failed_inner"),
],
),
(
abort(Category::Explicit),
&[
("core", "panicking::panic_nounwind"),
("core", "panicking::panic_nounwind_fmt"),
("core", "panicking::panic_cannot_unwind"),
],
),
(
unwind(Category::Index),
&[
("core", "panicking::panic_bounds_check"),
("core", "slice::index::slice_index_fail"),
("core", "slice::index::slice_start_index_len_fail"),
("core", "slice::index::slice_end_index_len_fail"),
("core", "slice::index::slice_index_order_fail"),
],
),
(
unwind(Category::Unwrap),
&[
("core", "option::unwrap_failed"),
("core", "option::expect_failed"),
("core", "result::unwrap_failed"),
],
),
(
unwind(Category::StrBoundary),
&[("core", "str::slice_error_fail")],
),
(
unwind(Category::Borrow),
&[
("core", "cell::panic_already_borrowed"),
("core", "cell::panic_already_mutably_borrowed"),
],
),
(
unwind(Category::CapacityOverflow),
&[("alloc", "raw_vec::capacity_overflow")],
),
(
funnel(
unwind(Category::CapacityOverflow),
abort(Category::AllocFailure),
),
&[
("alloc", "raw_vec::handle_error"),
("alloc", "raw_vec::handle_reserve"),
],
),
(
abort(Category::AllocFailure),
&[
("alloc", "alloc::handle_alloc_error"),
("std", "alloc::handle_alloc_error"),
],
),
(
unwind(Category::Explicit),
&[
("std", "panic::resume_unwind"),
("std", "panicking::resume_unwind"),
("std", "panicking::rust_panic_without_hook"),
("std", "panicking::begin_panic"),
("std", "rt::begin_panic"),
],
),
];
const DISCARDED: &[(Category, &str, &str)] = &[
(Category::Poison, "std", "sync::poison::PoisonError"),
(Category::Fmt, "core", "fmt::Error"),
];
#[derive(Default)]
pub struct SinkTable {
cache: Map<DefId, Option<Sink>>,
}
impl SinkTable {
pub fn get(&mut self, tcx: TyCtxt<'_>, did: DefId) -> Option<Sink> {
*self
.cache
.entry(did)
.or_insert_with(|| Self::classify(tcx, did))
}
fn classify(tcx: TyCtxt<'_>, did: DefId) -> Option<Sink> {
let krate = tcx.crate_name(did.krate);
let krate = krate.as_str();
let path = Self::def_path(tcx, did);
for (sink, entries) in EXACT {
if entries.iter().any(|(k, p)| *k == krate && *p == path) {
return Some(*sink);
}
}
if krate == "core" && path.starts_with("panicking::") {
return Some(unwind(Category::Explicit));
}
if let Some(sink) =
Self::by_leaf_name(path.rsplit("::").next().unwrap_or(&path))
{
return Some(sink);
}
Self::opaque_divergence(tcx, did, krate)
}
fn opaque_divergence(
tcx: TyCtxt<'_>,
did: DefId,
krate: &str,
) -> Option<Sink> {
if !matches!(krate, "core" | "alloc") {
return None;
}
if !matches!(tcx.def_kind(did), DefKind::Fn | DefKind::AssocFn) {
return None;
}
if tcx.is_mir_available(did) {
return None;
}
if !tcx
.fn_sig(did)
.skip_binder()
.skip_binder()
.output()
.is_never()
{
return None;
}
let aborts = tcx
.codegen_fn_attrs(did)
.flags
.contains(CodegenFnAttrFlags::NEVER_UNWIND);
Some(if aborts {
abort(Category::Explicit)
} else {
unwind(Category::Explicit)
})
}
fn by_leaf_name(leaf: &str) -> Option<Sink> {
match leaf {
"capacity_overflow" => Some(unwind(Category::CapacityOverflow)),
"precondition_check" => Some(abort(Category::UbCheck)),
"handle_alloc_error" | "alloc_err" | "oom" => {
Some(abort(Category::AllocFailure))
}
"panic_arc_overflow" | "panic_rc_overflow" => {
Some(unwind(Category::RefCountOverflow))
}
_ => None,
}
}
pub fn refine_unwrap<'tcx>(
tcx: TyCtxt<'tcx>,
args: ty::GenericArgsRef<'tcx>,
sink: Sink,
) -> Sink {
if !sink.is_only(Category::Unwrap) {
return sink;
}
for arg in args {
let Some(ty) = arg.as_type() else { continue };
for step in ty.walk() {
let Some(inner) = step.as_type() else {
continue;
};
let ty::Adt(def, _) = inner.kind() else {
continue;
};
let did = def.did();
let krate = tcx.crate_name(did.krate);
for (category, wanted_krate, wanted_path) in DISCARDED {
if krate.as_str() == *wanted_krate
&& Self::def_path(tcx, did) == *wanted_path
{
return unwind(*category);
}
}
}
}
sink
}
pub fn def_path(tcx: TyCtxt<'_>, did: DefId) -> String {
let mut out = String::new();
for seg in &tcx.def_path(did).data {
let Some(name) = seg.data.get_opt_name() else {
continue;
};
if !out.is_empty() {
out.push_str("::");
}
out.push_str(name.as_str());
}
out
}
}