From 05c17d3159f579b3abf14d995caba37ac28ad028 Mon Sep 17 00:00:00 2001
From: Alwin Berger <alwin.berger@tu-dortmund.de>
Date: Mon, 12 Aug 2024 12:07:31 +0200
Subject: [PATCH] report state space exporation

---
 fuzzers/FRET/src/systemstate/mod.rs    |   1 +
 fuzzers/FRET/src/systemstate/report.rs | 133 +++++++++++++++++++++++++
 2 files changed, 134 insertions(+)
 create mode 100644 fuzzers/FRET/src/systemstate/report.rs

diff --git a/fuzzers/FRET/src/systemstate/mod.rs b/fuzzers/FRET/src/systemstate/mod.rs
index aac3894421..5928da1d76 100644
--- a/fuzzers/FRET/src/systemstate/mod.rs
+++ b/fuzzers/FRET/src/systemstate/mod.rs
@@ -18,6 +18,7 @@ pub mod feedbacks;
 pub mod schedulers;
 pub mod stg;
 pub mod mutational;
+pub mod report;
 
 // Constants
 const NUM_PRIOS: usize = 5;
diff --git a/fuzzers/FRET/src/systemstate/report.rs b/fuzzers/FRET/src/systemstate/report.rs
new file mode 100644
index 0000000000..aec71d004b
--- /dev/null
+++ b/fuzzers/FRET/src/systemstate/report.rs
@@ -0,0 +1,133 @@
+//! Stage to compute/report AFL stats
+
+use core::{marker::PhantomData, time::Duration};
+
+use libafl_bolts::current_time;
+
+use libafl::{
+    corpus::{Corpus, HasCurrentCorpusId}, events::EventFirer, prelude::minimizer::TopRatedsMetadata, schedulers::minimizer::IsFavoredMetadata, stages::Stage, state::{HasCorpus, HasImported, UsesState}, Error, HasMetadata
+};
+use libafl::{
+    events::Event,
+    monitors::{AggregatorOps, UserStats, UserStatsValue},
+};
+use std::borrow::Cow;
+use serde_json::json;
+
+/// The [`AflStatsStage`] is a simple stage that computes and reports some stats.
+#[derive(Debug, Clone)]
+pub struct SchedulerStatsStage<E, EM, Z> {
+    last_report_time: Duration,
+    // the interval that we report all stats
+    stats_report_interval: Duration,
+
+    phantom: PhantomData<(E, EM, Z)>,
+}
+
+impl<E, EM, Z> UsesState for SchedulerStatsStage<E, EM, Z>
+where
+    E: UsesState,
+{
+    type State = E::State;
+}
+
+impl<E, EM, Z> Stage<E, EM, Z> for SchedulerStatsStage<E, EM, Z>
+where
+    E: UsesState,
+    EM: EventFirer<State = Self::State>,
+    Z: UsesState<State = Self::State>,
+    Self::State: HasImported + HasCorpus + HasMetadata,
+{
+    fn perform(
+        &mut self,
+        _fuzzer: &mut Z,
+        _executor: &mut E,
+        state: &mut <Self as UsesState>::State,
+        _manager: &mut EM,
+    ) -> Result<(), Error> {
+        let Some(corpus_idx) = state.current_corpus_id()? else {
+            return Err(Error::illegal_state(
+                "state is not currently processing a corpus index",
+            ));
+        };
+
+
+        let corpus_size = state.corpus().count();
+
+        let cur = current_time();
+
+        if cur.checked_sub(self.last_report_time).unwrap_or_default() > self.stats_report_interval {
+            if let Some(meta) = state.metadata_map().get::<TopRatedsMetadata>() {
+                let kc = meta.map.keys().count();
+                let mut v : Vec<_> = meta.map.values().collect();
+                v.sort_unstable();
+                v.dedup();
+                let vc = v.len();
+                #[cfg(feature = "std")]
+                {
+                    let json = json!({
+                        "relevant":vc,
+                        "objects":kc,
+                    });
+                    _manager.fire(
+                        state,
+                        Event::UpdateUserStats {
+                            name: Cow::from("Minimizer"),
+                            value: UserStats::new(
+                                UserStatsValue::String(Cow::from(json.to_string())),
+                                AggregatorOps::None,
+                            ),
+                            phantom: PhantomData,
+                        },
+                    )?;
+                }
+                #[cfg(not(feature = "std"))]
+                log::info!(
+                    "pending: {}, pend_favored: {}, own_finds: {}, imported: {}",
+                    pending_size,
+                    pend_favored_size,
+                    self.own_finds_size,
+                    self.imported_size
+                );
+                self.last_report_time = cur;
+            }
+        }
+
+        Ok(())
+    }
+
+    #[inline]
+    fn restart_progress_should_run(&mut self, _state: &mut <Self as UsesState>::State) -> Result<bool, Error> {
+        // Not running the target so we wont't crash/timeout and, hence, don't need to restore anything
+        Ok(true)
+    }
+
+    #[inline]
+    fn clear_restart_progress(&mut self, _state: &mut <Self as UsesState>::State) -> Result<(), Error> {
+        // Not running the target so we wont't crash/timeout and, hence, don't need to restore anything
+        Ok(())
+    }
+}
+
+impl<E, EM, Z> SchedulerStatsStage<E, EM, Z> {
+    /// create a new instance of the [`AflStatsStage`]
+    #[must_use]
+    pub fn new(interval: Duration) -> Self {
+        Self {
+            stats_report_interval: interval,
+            ..Default::default()
+        }
+    }
+}
+
+impl<E, EM, Z> Default for SchedulerStatsStage<E, EM, Z> {
+    /// the default instance of the [`AflStatsStage`]
+    #[must_use]
+    fn default() -> Self {
+        Self {
+            last_report_time: current_time(),
+            stats_report_interval: Duration::from_secs(3),
+            phantom: PhantomData,
+        }
+    }
+}