use alloc::{borrow::Cow, vec::Vec};
use core::{
fmt::Debug,
hash::{Hash, Hasher},
ops::{Deref, DerefMut},
};
use libafl_bolts::{AsSlice, AsSliceMut, HasLen, Named, Truncate, ownedref::OwnedMutSlice};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use crate::{
Error,
executors::ExitKind,
observers::{DifferentialObserver, Observer},
};
pub mod const_map;
pub use const_map::*;
pub mod variable_map;
pub use variable_map::*;
pub mod hitcount_map;
pub use hitcount_map::*;
pub mod multi_map;
pub use multi_map::*;
pub mod owned_map;
pub use owned_map::*;
pub trait CanTrack {
type WithIndexTracking: CanTrack;
type WithNoveltiesTracking: CanTrack;
const INDICES: bool;
const NOVELTIES: bool;
fn track_indices(self) -> Self::WithIndexTracking;
fn track_novelties(self) -> Self::WithNoveltiesTracking;
}
#[derive(Debug, Copy, Clone, Deserialize, Serialize)]
pub struct ExplicitTracking<T, const ITH: bool, const NTH: bool>(T);
impl<T, const ITH: bool, const NTH: bool> CanTrack for ExplicitTracking<T, ITH, NTH> {
type WithIndexTracking = ExplicitTracking<T, true, NTH>;
type WithNoveltiesTracking = ExplicitTracking<T, ITH, true>;
const INDICES: bool = ITH;
const NOVELTIES: bool = NTH;
fn track_indices(self) -> Self::WithIndexTracking {
ExplicitTracking::<T, true, NTH>(self.0)
}
fn track_novelties(self) -> Self::WithNoveltiesTracking {
ExplicitTracking::<T, ITH, true>(self.0)
}
}
impl<T, const ITH: bool, const NTH: bool> Named for ExplicitTracking<T, ITH, NTH>
where
T: Named,
{
fn name(&self) -> &Cow<'static, str> {
self.0.name()
}
}
impl<S, I, T, const ITH: bool, const NTH: bool> Observer<I, S> for ExplicitTracking<T, ITH, NTH>
where
T: Observer<I, S>,
{
fn flush(&mut self) -> Result<(), Error> {
self.0.flush()
}
fn pre_exec(&mut self, state: &mut S, input: &I) -> Result<(), Error> {
self.0.pre_exec(state, input)
}
fn post_exec(&mut self, state: &mut S, input: &I, exit_kind: &ExitKind) -> Result<(), Error> {
self.0.post_exec(state, input, exit_kind)
}
fn pre_exec_child(&mut self, state: &mut S, input: &I) -> Result<(), Error> {
self.0.pre_exec_child(state, input)
}
fn post_exec_child(
&mut self,
state: &mut S,
input: &I,
exit_kind: &ExitKind,
) -> Result<(), Error> {
self.0.post_exec_child(state, input, exit_kind)
}
}
impl<T, OTA, OTB, I, S, const ITH: bool, const NTH: bool> DifferentialObserver<OTA, OTB, I, S>
for ExplicitTracking<T, ITH, NTH>
where
T: DifferentialObserver<OTA, OTB, I, S>,
{
fn pre_observe_first(&mut self, observers: &mut OTA) -> Result<(), Error> {
self.as_mut().pre_observe_first(observers)
}
fn post_observe_first(&mut self, observers: &mut OTA) -> Result<(), Error> {
self.as_mut().post_observe_first(observers)
}
fn pre_observe_second(&mut self, observers: &mut OTB) -> Result<(), Error> {
self.as_mut().pre_observe_second(observers)
}
fn post_observe_second(&mut self, observers: &mut OTB) -> Result<(), Error> {
self.as_mut().post_observe_second(observers)
}
}
impl<T, const ITH: bool, const NTH: bool> AsRef<T> for ExplicitTracking<T, ITH, NTH> {
fn as_ref(&self) -> &T {
&self.0
}
}
impl<T, const ITH: bool, const NTH: bool> AsMut<T> for ExplicitTracking<T, ITH, NTH> {
fn as_mut(&mut self) -> &mut T {
&mut self.0
}
}
pub mod macros {
pub use const_format::{concatcp, str_repeat};
pub use const_panic::{FmtArg, concat_panic};
#[macro_export]
macro_rules! require_index_tracking {
($name: literal, $obs: ident) => {
struct TrackingEnabledCheck<O: $crate::observers::CanTrack> {
phantom: ::core::marker::PhantomData<O>,
}
impl<O: $crate::observers::CanTrack> TrackingEnabledCheck<O> {
#[rustfmt::skip]
const MESSAGE: &'static str = {
const LINE_OFFSET: usize = line!().ilog10() as usize + 2;
const SPACING: &str = $crate::observers::map::macros::str_repeat!(" ", LINE_OFFSET);
$crate::observers::map::macros::concatcp!(
"\n",
SPACING, "|\n",
SPACING, "= note: index tracking is required by ", $name, "\n",
SPACING, "= note: see the documentation of CanTrack for details\n",
SPACING, "|\n",
SPACING, "= hint: call `.track_indices()` on the map observer passed to ", $name, " at the point where it is defined\n",
SPACING, "|\n",
SPACING, "| ",
)
};
const TRACKING_ENABLED: bool = {
if !O::INDICES {
panic!("{}", Self::MESSAGE)
} else {
true
}
};
#[inline(always)]
fn check_enabled() {
if !Self::TRACKING_ENABLED {
unreachable!("{}", Self::MESSAGE);
}
}
}
TrackingEnabledCheck::<$obs>::check_enabled(); };
}
#[macro_export]
macro_rules! require_novelties_tracking {
($name: literal, $obs: ident) => {
struct TrackingEnabledCheck<O: $crate::observers::CanTrack> {
phantom: ::core::marker::PhantomData<O>,
}
impl<O: $crate::observers::CanTrack> TrackingEnabledCheck<O> {
#[rustfmt::skip]
const MESSAGE: &'static str = {
const LINE_OFFSET: usize = line!().ilog10() as usize + 2;
const SPACING: &str =
$crate::observers::map::macros::str_repeat!(" ", LINE_OFFSET);
$crate::observers::map::macros::concatcp!(
"\n",
SPACING, "|\n",
SPACING, "= note: novelty tracking is required by ", $name, "\n",
SPACING, "= note: see the documentation of CanTrack for details\n",
SPACING, "|\n",
SPACING, "= hint: call `.track_novelties()` on the map observer passed to ", $name, " at the point where it is defined\n",
SPACING, "|\n",
SPACING, "| ",
)
};
const TRACKING_ENABLED: bool = {
if !O::NOVELTIES {
panic!("{}", Self::MESSAGE)
} else {
true
}
};
#[inline(always)]
fn check_enabled() {
if !Self::TRACKING_ENABLED {
unreachable!("{}", Self::MESSAGE);
}
}
}
TrackingEnabledCheck::<$obs>::check_enabled(); };
}
}
pub trait MapObserver:
HasLen + Named + Serialize + DeserializeOwned + AsRef<Self> + AsMut<Self>
{
type Entry: PartialEq + Copy;
fn get(&self, idx: usize) -> Self::Entry;
fn set(&mut self, idx: usize, val: Self::Entry);
fn usable_count(&self) -> usize;
fn count_bytes(&self) -> u64;
fn initial(&self) -> Self::Entry;
fn reset_map(&mut self) -> Result<(), Error>;
fn to_vec(&self) -> Vec<Self::Entry>;
fn how_many_set(&self, indexes: &[usize]) -> usize;
}
pub trait VarLenMapObserver: MapObserver {
fn map_slice(&self) -> &[Self::Entry];
fn map_slice_mut(&mut self) -> &mut [Self::Entry];
fn size(&self) -> &usize;
fn size_mut(&mut self) -> &mut usize;
}
pub trait ConstLenMapObserver<const N: usize>: MapObserver {
const LENGTH: usize = N;
fn map_slice(&self) -> &[Self::Entry; N];
fn map_slice_mut(&mut self) -> &mut [Self::Entry; N];
}
impl<M> CanTrack for M
where
M: MapObserver,
{
type WithIndexTracking = ExplicitTracking<Self, true, false>;
type WithNoveltiesTracking = ExplicitTracking<Self, false, true>;
const INDICES: bool = false;
const NOVELTIES: bool = false;
fn track_indices(self) -> Self::WithIndexTracking {
ExplicitTracking::<Self, true, false>(self)
}
fn track_novelties(self) -> Self::WithNoveltiesTracking {
ExplicitTracking::<Self, false, true>(self)
}
}
#[derive(Clone, Serialize, Deserialize, Debug)]
#[expect(clippy::unsafe_derive_deserialize)]
pub struct StdMapObserver<'a, T, const DIFFERENTIAL: bool> {
map: OwnedMutSlice<'a, T>,
initial: T,
name: Cow<'static, str>,
}
impl<I, S, T> Observer<I, S> for StdMapObserver<'_, T, false>
where
Self: MapObserver,
{
#[inline]
fn pre_exec(&mut self, _state: &mut S, _input: &I) -> Result<(), Error> {
self.reset_map()
}
}
impl<I, S, T> Observer<I, S> for StdMapObserver<'_, T, true> {}
impl<T, const DIFFERENTIAL: bool> Named for StdMapObserver<'_, T, DIFFERENTIAL> {
#[inline]
fn name(&self) -> &Cow<'static, str> {
&self.name
}
}
impl<T, const DIFFERENTIAL: bool> HasLen for StdMapObserver<'_, T, DIFFERENTIAL> {
#[inline]
fn len(&self) -> usize {
self.map.as_slice().len()
}
}
impl<T, const DIFFERENTIAL: bool> Hash for StdMapObserver<'_, T, DIFFERENTIAL>
where
T: Hash,
{
#[inline]
fn hash<H: Hasher>(&self, hasher: &mut H) {
self.as_slice().hash(hasher);
}
}
impl<T, const DIFFERENTIAL: bool> AsRef<Self> for StdMapObserver<'_, T, DIFFERENTIAL> {
fn as_ref(&self) -> &Self {
self
}
}
impl<T, const DIFFERENTIAL: bool> AsMut<Self> for StdMapObserver<'_, T, DIFFERENTIAL> {
fn as_mut(&mut self) -> &mut Self {
self
}
}
impl<T, const DIFFERENTIAL: bool> MapObserver for StdMapObserver<'_, T, DIFFERENTIAL>
where
T: PartialEq + Copy + Hash + Serialize + DeserializeOwned + Debug,
{
type Entry = T;
#[inline]
fn get(&self, pos: usize) -> T {
self.as_slice()[pos]
}
fn set(&mut self, pos: usize, val: T) {
self.map.as_slice_mut()[pos] = val;
}
fn count_bytes(&self) -> u64 {
let initial = self.initial();
let cnt = self.usable_count();
let map = self.as_slice();
let mut res = 0;
for x in &map[0..cnt] {
if *x != initial {
res += 1;
}
}
res
}
#[inline]
fn usable_count(&self) -> usize {
self.as_slice().len()
}
#[inline]
fn initial(&self) -> T {
self.initial
}
fn to_vec(&self) -> Vec<T> {
self.as_slice().to_vec()
}
#[inline]
fn reset_map(&mut self) -> Result<(), Error> {
let initial = self.initial();
let cnt = self.usable_count();
let map = self.as_slice_mut();
for x in &mut map[0..cnt] {
*x = initial;
}
Ok(())
}
fn how_many_set(&self, indexes: &[usize]) -> usize {
let initial = self.initial();
let cnt = self.usable_count();
let map = self.as_slice();
let mut res = 0;
for i in indexes {
if *i < cnt && map[*i] != initial {
res += 1;
}
}
res
}
}
impl<T, const DIFFERENTIAL: bool> Truncate for StdMapObserver<'_, T, DIFFERENTIAL> {
fn truncate(&mut self, new_len: usize) {
self.map.truncate(new_len);
}
}
impl<T, const DIFFERENTIAL: bool> Deref for StdMapObserver<'_, T, DIFFERENTIAL> {
type Target = [T];
fn deref(&self) -> &[T] {
&self.map
}
}
impl<T, const DIFFERENTIAL: bool> DerefMut for StdMapObserver<'_, T, DIFFERENTIAL> {
fn deref_mut(&mut self) -> &mut [T] {
&mut self.map
}
}
impl<'a, T, const DIFFERENTIAL: bool> StdMapObserver<'a, T, DIFFERENTIAL>
where
T: Default,
{
#[must_use]
unsafe fn maybe_differential<S>(name: S, map: &'a mut [T]) -> Self
where
S: Into<Cow<'static, str>>,
{
unsafe {
let len = map.len();
let ptr = map.as_mut_ptr();
Self::maybe_differential_from_mut_ptr(name, ptr, len)
}
}
#[must_use]
fn maybe_differential_from_mut_slice<S>(name: S, map: OwnedMutSlice<'a, T>) -> Self
where
S: Into<Cow<'static, str>>,
{
StdMapObserver {
name: name.into(),
map,
initial: T::default(),
}
}
#[must_use]
fn maybe_differential_owned<S>(name: S, map: Vec<T>) -> Self
where
S: Into<Cow<'static, str>>,
{
Self {
map: OwnedMutSlice::from(map),
name: name.into(),
initial: T::default(),
}
}
#[must_use]
fn maybe_differential_from_ownedref<S>(name: S, map: OwnedMutSlice<'a, T>) -> Self
where
S: Into<Cow<'static, str>>,
{
Self {
map,
name: name.into(),
initial: T::default(),
}
}
unsafe fn maybe_differential_from_mut_ptr<S>(name: S, map_ptr: *mut T, len: usize) -> Self
where
S: Into<Cow<'static, str>>,
{
unsafe {
Self::maybe_differential_from_mut_slice(
name,
OwnedMutSlice::from_raw_parts_mut(map_ptr, len),
)
}
}
pub fn initial_mut(&mut self) -> &mut T {
&mut self.initial
}
pub fn map(&self) -> &OwnedMutSlice<'a, T> {
&self.map
}
pub fn map_mut(&mut self) -> &mut OwnedMutSlice<'a, T> {
&mut self.map
}
}
impl<'a, T> StdMapObserver<'a, T, false>
where
T: Default,
{
#[must_use]
pub unsafe fn new<S>(name: S, map: &'a mut [T]) -> Self
where
S: Into<Cow<'static, str>>,
{
unsafe { Self::maybe_differential(name, map) }
}
pub fn from_mut_slice<S>(name: S, map: OwnedMutSlice<'a, T>) -> Self
where
S: Into<Cow<'static, str>>,
{
Self::maybe_differential_from_mut_slice(name, map)
}
#[must_use]
pub fn owned<S>(name: S, map: Vec<T>) -> Self
where
S: Into<Cow<'static, str>>,
{
Self::maybe_differential_owned(name, map)
}
#[must_use]
pub fn from_ownedref<S>(name: S, map: OwnedMutSlice<'a, T>) -> Self
where
S: Into<Cow<'static, str>>,
{
Self::maybe_differential_from_ownedref(name, map)
}
pub unsafe fn from_mut_ptr<S>(name: S, map_ptr: *mut T, len: usize) -> Self
where
S: Into<Cow<'static, str>>,
{
unsafe { Self::maybe_differential_from_mut_ptr(name, map_ptr, len) }
}
}
impl<'a, T> StdMapObserver<'a, T, true>
where
T: Default,
{
#[must_use]
pub unsafe fn differential<S>(name: S, map: &'a mut [T]) -> Self
where
S: Into<Cow<'static, str>>,
{
unsafe { Self::maybe_differential(name, map) }
}
#[must_use]
pub fn differential_owned<S>(name: S, map: Vec<T>) -> Self
where
S: Into<Cow<'static, str>>,
{
Self::maybe_differential_owned(name, map)
}
#[must_use]
pub fn differential_from_ownedref<S>(name: S, map: OwnedMutSlice<'a, T>) -> Self
where
S: Into<Cow<'static, str>>,
{
Self::maybe_differential_from_ownedref(name, map)
}
pub unsafe fn differential_from_mut_ptr<S>(name: S, map_ptr: *mut T, len: usize) -> Self
where
S: Into<Cow<'static, str>>,
{
unsafe { Self::maybe_differential_from_mut_ptr(name, map_ptr, len) }
}
}
impl<OTA, OTB, I, S, T> DifferentialObserver<OTA, OTB, I, S> for StdMapObserver<'_, T, true> {}