Restartable trait (#3004)

* separate into Restartable

* typo

* fixes, blanket impl

* aa

* lol

* aa

* fixer

* lol

fuzzers/forkserver/libafl-fuzz/src/stages/mutational_stage.rs

@@ -1,21 +1,21 @@
 use std::{borrow::Cow, marker::PhantomData};
 
 use libafl::{
-    stages::{MutationalStage, Stage},
+    stages::{MutationalStage, Restartable, Stage},
     Error,
 };
 use libafl_bolts::Named;
 
 #[derive(Debug)]
-pub enum SupportedMutationalStages<SM, P> {
+pub enum SupportedMutationalStages<P, SM> {
     StdMutational(SM, PhantomData<P>),
     PowerMutational(P, PhantomData<SM>),
 }
 
-impl<S, SM, P> MutationalStage<S> for SupportedMutationalStages<SM, P>
+impl<P, S, SM> MutationalStage<S> for SupportedMutationalStages<P, SM>
 where
-    SM: MutationalStage<S>,
     P: MutationalStage<S, Mutator = SM::Mutator>,
+    SM: MutationalStage<S>,
 {
     type Mutator = SM::Mutator;
     /// The mutator, added to this stage
@@ -44,10 +44,10 @@ where
     }
 }
 
-impl<SM, P> Named for SupportedMutationalStages<SM, P>
+impl<P, SM> Named for SupportedMutationalStages<P, SM>
 where
-    SM: Named,
     P: Named,
+    SM: Named,
 {
     fn name(&self) -> &Cow<'static, str> {
         match self {
@@ -57,10 +57,10 @@ where
     }
 }
 
-impl<E, EM, S, SM, P, Z> Stage<E, EM, S, Z> for SupportedMutationalStages<SM, P>
+impl<E, EM, P, S, SM, Z> Stage<E, EM, S, Z> for SupportedMutationalStages<P, SM>
 where
-    SM: Stage<E, EM, S, Z>,
     P: Stage<E, EM, S, Z>,
+    SM: Stage<E, EM, S, Z>,
 {
     #[inline]
     fn perform(
@@ -75,7 +75,13 @@ where
             Self::PowerMutational(p, _) => p.perform(fuzzer, executor, state, manager),
         }
     }
+}
 
+impl<P, S, SM> Restartable<S> for SupportedMutationalStages<P, SM>
+where
+    P: Restartable<S>,
+    SM: Restartable<S>,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         match self {
             Self::StdMutational(m, _) => m.should_restart(state),

libafl/src/stages/afl_stats.rs

@@ -30,7 +30,7 @@ use crate::{
     mutators::Tokens,
     observers::MapObserver,
     schedulers::{minimizer::IsFavoredMetadata, HasQueueCycles},
-    stages::{calibrate::UnstableEntriesMetadata, Stage},
+    stages::{calibrate::UnstableEntriesMetadata, Restartable, Stage},
     state::{HasCorpus, HasExecutions, HasImported, HasStartTime, Stoppable},
     std::string::ToString,
     Error, HasMetadata, HasNamedMetadata, HasScheduler,
@@ -432,7 +432,9 @@ where
 
         Ok(())
     }
+}
 
+impl<C, E, EM, I, O, S, Z> Restartable<S> for AflStatsStage<C, E, EM, I, O, S, Z> {
     fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
         Ok(true)
     }

libafl/src/stages/calibrate.rs

@@ -22,7 +22,7 @@ use crate::{
     monitors::stats::{AggregatorOps, UserStats, UserStatsValue},
     observers::{MapObserver, ObserversTuple},
     schedulers::powersched::SchedulerMetadata,
-    stages::{RetryCountRestartHelper, Stage},
+    stages::{Restartable, RetryCountRestartHelper, Stage},
     state::{HasCorpus, HasCurrentTestcase, HasExecutions},
     Error, HasMetadata, HasNamedMetadata,
 };
@@ -364,7 +364,12 @@ where
 
         Ok(())
     }
+}
 
+impl<C, E, I, O, OT, S> Restartable<S> for CalibrationStage<C, E, I, O, OT, S>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         // Calibration stage disallow restarts
         // If a testcase that causes crash/timeout in the queue, we need to remove it from the queue immediately.

libafl/src/stages/colorization.rs

@@ -22,7 +22,7 @@ use crate::{
     mutators::mutations::buffer_copy,
     nonzero,
     observers::ObserversTuple,
-    stages::{RetryCountRestartHelper, Stage},
+    stages::{Restartable, RetryCountRestartHelper, Stage},
     state::{HasCorpus, HasCurrentTestcase, HasRand},
     Error, HasMetadata, HasNamedMetadata,
 };
@@ -98,7 +98,12 @@ where
 
         Ok(())
     }
+}
 
+impl<C, E, EM, I, O, S, Z> Restartable<S> for ColorizationStage<C, E, EM, I, O, S, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         // This is a deterministic stage
         // Once it failed, then don't retry,

libafl/src/stages/concolic.rs

@@ -18,7 +18,7 @@ use crate::{
     corpus::HasCurrentCorpusId,
     executors::{Executor, HasObservers},
     observers::{concolic::ConcolicObserver, ObserversTuple},
-    stages::{RetryCountRestartHelper, Stage, TracingStage},
+    stages::{Restartable, RetryCountRestartHelper, Stage, TracingStage},
     state::{HasCorpus, HasCurrentTestcase, HasExecutions, MaybeHasClientPerfMonitor},
     Error, HasMetadata, HasNamedMetadata,
 };
@@ -76,7 +76,12 @@ where
         }
         Ok(())
     }
+}
 
+impl<EM, I, TE, S, Z> Restartable<S> for ConcolicTracingStage<'_, EM, I, TE, S, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         // This is a deterministic stage
         // Once it failed, then don't retry,
@@ -418,7 +423,13 @@ where
         }
         Ok(())
     }
+}
 
+#[cfg(feature = "concolic_mutation")]
+impl<I, S, Z> Restartable<S> for SimpleConcolicMutationalStage<I, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
     #[inline]
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         // This is a deterministic stage

libafl/src/stages/dump.rs

@@ -15,7 +15,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     corpus::{Corpus, CorpusId, Testcase},
     inputs::Input,
-    stages::Stage,
+    stages::{Restartable, Stage},
     state::{HasCorpus, HasRand, HasSolutions},
     Error, HasMetadata,
 };
@@ -60,7 +60,11 @@ where
     ) -> Result<(), Error> {
         self.dump_state_to_disk(state)
     }
+}
 
+impl<CB1, EM, I, S, Z> Restartable<S>
+    for DumpToDiskStage<CB1, fn(&Testcase<I>, &CorpusId) -> String, EM, I, S, Z>
+{
     #[inline]
     fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
         // Not executing the target, so restart safety is not needed

libafl/src/stages/generalization.rs

@@ -23,7 +23,7 @@ use crate::{
     mark_feature_time,
     observers::{CanTrack, MapObserver, ObserversTuple},
     require_novelties_tracking,
-    stages::{RetryCountRestartHelper, Stage},
+    stages::{Restartable, RetryCountRestartHelper, Stage},
     start_timer,
     state::{HasCorpus, HasExecutions, MaybeHasClientPerfMonitor},
     Error, HasMetadata, HasNamedMetadata,
@@ -62,6 +62,23 @@ impl<C, EM, I, O, OT, S, Z> Named for GeneralizationStage<C, EM, I, O, OT, S, Z>
     }
 }
 
+impl<C, EM, I, O, OT, S, Z> Restartable<S> for GeneralizationStage<C, EM, I, O, OT, S, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
+    #[inline]
+    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
+        // TODO: We need to be able to resume better if something crashes or times out
+        RetryCountRestartHelper::should_restart::<S>(state, &self.name, 3)
+    }
+
+    #[inline]
+    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
+        // TODO: We need to be able to resume better if something crashes or times out
+        RetryCountRestartHelper::clear_progress::<S>(state, &self.name)
+    }
+}
+
 impl<C, E, EM, O, S, Z> Stage<E, EM, S, Z>
     for GeneralizationStage<C, EM, BytesInput, O, E::Observers, S, Z>
 where
@@ -76,18 +93,6 @@ where
         + HasCurrentCorpusId
         + MaybeHasClientPerfMonitor,
 {
-    #[inline]
-    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
-        // TODO: We need to be able to resume better if something crashes or times out
-        RetryCountRestartHelper::should_restart::<S>(state, &self.name, 3)
-    }
-
-    #[inline]
-    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
-        // TODO: We need to be able to resume better if something crashes or times out
-        RetryCountRestartHelper::clear_progress::<S>(state, &self.name)
-    }
-
     #[inline]
     #[expect(clippy::too_many_lines)]
     fn perform(

libafl/src/stages/generation.rs

@@ -6,7 +6,12 @@
 
 use core::marker::PhantomData;
 
-use crate::{generators::Generator, stages::Stage, state::HasRand, Error, Evaluator};
+use crate::{
+    generators::Generator,
+    stages::{Restartable, Stage},
+    state::HasRand,
+    Error, Evaluator,
+};
 
 /// A [`Stage`] that generates a single input via a [`Generator`] and evaluates
 /// it using the fuzzer, possibly adding it to the corpus.
@@ -40,7 +45,9 @@ where
         fuzzer.evaluate_filtered(state, executor, manager, &input)?;
         Ok(())
     }
+}
 
+impl<G, I, S, Z> Restartable<S> for GenStage<G, I, S, Z> {
     fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
         // It's a random generation stage
         // so you can restart for whatever times you want

libafl/src/stages/logics.rs

@@ -3,7 +3,7 @@
 use core::marker::PhantomData;
 
 use crate::{
-    stages::{Stage, StageId, StagesTuple},
+    stages::{Restartable, Stage, StageId, StagesTuple},
     state::HasNestedStage,
     Error,
 };
@@ -59,7 +59,12 @@ where
 
         Ok(())
     }
+}
 
+impl<CB, E, EM, ST, S, Z> Restartable<S> for WhileStage<CB, E, EM, ST, S, Z>
+where
+    S: HasNestedStage,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         NestedStageRetryCountRestartHelper::should_restart(state, self)
     }
@@ -112,7 +117,12 @@ where
         }
         Ok(())
     }
+}
 
+impl<CB, E, EM, ST, S, Z> Restartable<S> for IfStage<CB, E, EM, ST, S, Z>
+where
+    S: HasNestedStage,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         NestedStageRetryCountRestartHelper::should_restart(state, self)
     }
@@ -188,7 +198,12 @@ where
 
         Ok(())
     }
+}
 
+impl<CB, E, EM, ST1, ST2, S, Z> Restartable<S> for IfElseStage<CB, E, EM, ST1, ST2, S, Z>
+where
+    S: HasNestedStage,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         NestedStageRetryCountRestartHelper::should_restart(state, self)
     }
@@ -238,7 +253,12 @@ where
             Ok(())
         }
     }
+}
 
+impl<E, EM, ST, S, Z> Restartable<S> for OptionalStage<E, EM, ST, S, Z>
+where
+    S: HasNestedStage,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         NestedStageRetryCountRestartHelper::should_restart(state, self)
     }
@@ -295,8 +315,8 @@ mod test {
         events::NopEventManager,
         executors::test::NopExecutor,
         stages::{
-            ClosureStage, CorpusId, HasCurrentCorpusId, IfElseStage, IfStage, Stage, StagesTuple,
+            ClosureStage, CorpusId, HasCurrentCorpusId, IfElseStage, IfStage, Restartable, Stage,
-            WhileStage,
+            StagesTuple, WhileStage,
         },
         state::{HasCurrentStageId, StdState},
         HasMetadata, NopFuzzer,
@@ -362,7 +382,12 @@ mod test {
             );
             Ok(())
         }
+    }
 
+    impl<S> Restartable<S> for ResumeSucceededStage<S>
+    where
+        S: HasMetadata,
+    {
         fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
             TestProgress::should_restart(state, self)
         }
@@ -395,7 +420,12 @@ mod test {
             }
             Ok(())
         }
+    }
 
+    impl<S> Restartable<S> for ResumeFailedStage<S>
+    where
+        S: HasMetadata,
+    {
         fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
             TestProgress::should_restart(state, self)
         }
@@ -546,7 +576,9 @@ mod test {
         ) -> Result<(), Error> {
             panic!("Test failed; panic stage should never be executed.");
         }
+    }
 
+    impl<S> Restartable<S> for PanicStage<S> {
         fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
             Ok(true)
         }
@@ -555,7 +587,6 @@ mod test {
             Ok(())
         }
     }
-
     #[test]
     fn check_resumability_if_else_if() {
         let once = RefCell::new(true);

libafl/src/stages/mod.rs

@@ -87,20 +87,11 @@ pub mod verify_timeouts;
 /// A stage is one step in the fuzzing process.
 /// Multiple stages will be scheduled one by one for each input.
 pub trait Stage<E, EM, S, Z> {
-    /// This method will be called before every call to [`Stage::perform`].
-    /// Initialize the restart tracking for this stage, _if it is not yet initialized_.
-    /// On restart, this will be called again.
-    /// As long as [`Stage::clear_progress`], all subsequent calls happen on restart.
-    /// Returns `true`, if the stage's [`Stage::perform`] method should run, else `false`.
-    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error>;
-
-    /// Clear the current status tracking of the associated stage
-    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error>;
-
     /// Run the stage.
     ///
-    /// Before a call to perform, [`Stage::should_restart`] will be (must be!) called.
+    /// If you want this stage to restart, then
-    /// After returning (so non-target crash or timeout in a restarting case), [`Stage::clear_progress`] gets called.
+    /// Before a call to perform, [`Restartable::should_restart`] will be (must be!) called.
+    /// After returning (so non-target crash or timeout in a restarting case), [`Restartable::clear_progress`] gets called.
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -110,6 +101,19 @@ pub trait Stage<E, EM, S, Z> {
     ) -> Result<(), Error>;
 }
 
+/// Restartable trait takes care of stage restart.
+pub trait Restartable<S> {
+    /// This method will be called before every call to [`Stage::perform`].
+    /// Initialize the restart tracking for this stage, _if it is not yet initialized_.
+    /// On restart, this will be called again.
+    /// As long as [`Restartable::clear_progress`], all subsequent calls happen on restart.
+    /// Returns `true`, if the stage's [`Stage::perform`] method should run, else `false`.
+    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error>;
+
+    /// Clear the current status tracking of the associated stage
+    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error>;
+}
+
 /// A tuple holding all `Stages` used for fuzzing.
 pub trait StagesTuple<E, EM, S, Z> {
     /// Performs all `Stages` in this tuple.
@@ -145,7 +149,7 @@ where
 
 impl<Head, Tail, E, EM, S, Z> StagesTuple<E, EM, S, Z> for (Head, Tail)
 where
-    Head: Stage<E, EM, S, Z>,
+    Head: Stage<E, EM, S, Z> + Restartable<S>,
     Tail: StagesTuple<E, EM, S, Z> + HasConstLen,
     S: HasCurrentStageId + Stoppable,
     EM: SendExiting,
@@ -169,10 +173,7 @@ where
 
                 let stage = &mut self.0;
 
-                if stage.should_restart(state)? {
+                stage.perform_restartable(fuzzer, executor, state, manager)?;
-                    stage.perform(fuzzer, executor, state, manager)?;
-                }
-                stage.clear_progress(state)?;
 
                 state.clear_stage_id()?;
             }
@@ -185,10 +186,7 @@ where
 
                 let stage = &mut self.0;
 
-                if stage.should_restart(state)? {
+                stage.perform_restartable(fuzzer, executor, state, manager)?;
-                    stage.perform(fuzzer, executor, state, manager)?;
-                }
-                stage.clear_progress(state)?;
 
                 state.clear_stage_id()?;
             }
@@ -240,7 +238,36 @@ impl<E, EM, S, Z> IntoVec<Box<dyn Stage<E, EM, S, Z>>> for Vec<Box<dyn Stage<E,
     }
 }
 
-impl<E, EM, S, Z> StagesTuple<E, EM, S, Z> for Vec<Box<dyn Stage<E, EM, S, Z>>>
+trait RestartableStage<E, EM, S, Z>: Stage<E, EM, S, Z> + Restartable<S> {
+    fn perform_restartable(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut S,
+        manager: &mut EM,
+    ) -> Result<(), Error>;
+}
+
+impl<E, EM, S, ST, Z> RestartableStage<E, EM, S, Z> for ST
+where
+    ST: Stage<E, EM, S, Z> + Restartable<S>,
+{
+    /// Run the stage, calling [`Stage::should_restart`] and [`Stage::clear_progress`] appropriately
+    fn perform_restartable(
+        &mut self,
+        fuzzer: &mut Z,
+        executor: &mut E,
+        state: &mut S,
+        manager: &mut EM,
+    ) -> Result<(), Error> {
+        if self.should_restart(state)? {
+            self.perform(fuzzer, executor, state, manager)?;
+        }
+        self.clear_progress(state)
+    }
+}
+
+impl<E, EM, S, Z> StagesTuple<E, EM, S, Z> for Vec<Box<dyn RestartableStage<E, EM, S, Z>>>
 where
     EM: SendExiting,
     S: HasCurrentStageId + Stoppable,
@@ -261,10 +288,7 @@ where
                 manager.on_shutdown()?;
                 return Err(Error::shutting_down());
             }
-            if stage.should_restart(state)? {
+            stage.perform_restartable(fuzzer, executor, state, manager)
-                stage.perform(fuzzer, executor, state, manager)?;
-            }
-            stage.clear_progress(state)
         })
     }
 }
@@ -301,7 +325,12 @@ where
     ) -> Result<(), Error> {
         (self.closure)(fuzzer, executor, state, manager)
     }
+}
 
+impl<CB, E, EM, S, Z> Restartable<S> for ClosureStage<CB, E, EM, Z>
+where
+    S: HasNamedMetadata + HasCurrentCorpusId,
+{
     #[inline]
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         // There's no restart safety in the content of the closure.

libafl/src/stages/mutational.rs

@@ -18,7 +18,7 @@ use crate::{
     mark_feature_time,
     mutators::{MultiMutator, MutationResult, Mutator},
     nonzero,
-    stages::{RetryCountRestartHelper, Stage},
+    stages::{Restartable, RetryCountRestartHelper, Stage},
     start_timer,
     state::{HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, MaybeHasClientPerfMonitor},
     Error, HasMetadata, HasNamedMetadata,
@@ -170,7 +170,12 @@ where
 
         ret
     }
+}
 
+impl<E, EM, I1, I2, M, S, Z> Restartable<S> for StdMutationalStage<E, EM, I1, I2, M, S, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         RetryCountRestartHelper::should_restart(state, &self.name, 3)
     }
@@ -316,17 +321,6 @@ where
     S: HasRand + HasNamedMetadata + HasCurrentTestcase<I> + HasCurrentCorpusId,
     Z: Evaluator<E, EM, I, S>,
 {
-    #[inline]
-    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
-        // Make sure we don't get stuck crashing on a single testcase
-        RetryCountRestartHelper::should_restart(state, &self.name, 3)
-    }
-
-    #[inline]
-    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
-        RetryCountRestartHelper::clear_progress(state, &self.name)
-    }
-
     #[inline]
     fn perform(
         &mut self,
@@ -354,6 +348,22 @@ where
     }
 }
 
+impl<E, EM, I, M, S, Z> Restartable<S> for MultiMutationalStage<E, EM, I, M, S, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
+    #[inline]
+    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
+        // Make sure we don't get stuck crashing on a single testcase
+        RetryCountRestartHelper::should_restart(state, &self.name, 3)
+    }
+
+    #[inline]
+    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
+        RetryCountRestartHelper::clear_progress(state, &self.name)
+    }
+}
+
 impl<E, EM, I, M, S, Z> MultiMutationalStage<E, EM, I, M, S, Z> {
     /// Creates a new [`MultiMutationalStage`]
     pub fn new(mutator: M) -> Self {

libafl/src/stages/power.rs

@@ -19,7 +19,7 @@ use crate::{
     schedulers::{testcase_score::CorpusPowerTestcaseScore, TestcaseScore},
     stages::{
         mutational::{MutatedTransform, MutatedTransformPost},
-        MutationalStage, RetryCountRestartHelper, Stage,
+        MutationalStage, Restartable, RetryCountRestartHelper, Stage,
     },
     start_timer,
     state::{HasCurrentTestcase, HasExecutions, HasRand, MaybeHasClientPerfMonitor},
@@ -101,7 +101,12 @@ where
         let ret = self.perform_mutational(fuzzer, executor, state, manager);
         ret
     }
+}
 
+impl<E, F, EM, I, M, S, Z> Restartable<S> for PowerMutationalStage<E, F, EM, I, M, S, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         // Make sure we don't get stuck crashing on a single testcase
         RetryCountRestartHelper::should_restart(state, &self.name, 3)

libafl/src/stages/push/mod.rs

@@ -28,7 +28,7 @@ use crate::{
     executors::{Executor, ExitKind, HasObservers},
     observers::ObserversTuple,
     schedulers::Scheduler,
-    stages::{RetryCountRestartHelper, Stage},
+    stages::{Restartable, RetryCountRestartHelper, Stage},
     state::{HasCorpus, HasExecutions, HasLastReportTime, HasRand},
     Error, EvaluatorObservers, ExecutesInput, ExecutionProcessor, HasMetadata, HasScheduler,
 };
@@ -244,6 +244,22 @@ impl<CS, EM, I, OT, PS, Z> Named for PushStageAdapter<CS, EM, I, OT, PS, Z> {
     }
 }
 
+impl<CS, EM, I, OT, PS, S, Z> Restartable<S> for PushStageAdapter<CS, EM, I, OT, PS, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
+    #[inline]
+    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
+        // TODO: Proper restart handling - call post_exec at the right time, etc...
+        RetryCountRestartHelper::no_retry(state, &self.name)
+    }
+
+    #[inline]
+    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
+        RetryCountRestartHelper::clear_progress(state, &self.name)
+    }
+}
+
 impl<CS, E, EM, I, OT, PS, S, Z> Stage<E, EM, S, Z> for PushStageAdapter<CS, EM, I, OT, PS, Z>
 where
     CS: Scheduler<I, S>,
@@ -263,17 +279,6 @@ where
         + EvaluatorObservers<E, EM, I, OT>
         + HasScheduler<I, S>,
 {
-    #[inline]
-    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
-        // TODO: Proper restart handling - call post_exec at the right time, etc...
-        RetryCountRestartHelper::no_retry(state, &self.name)
-    }
-
-    #[inline]
-    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
-        RetryCountRestartHelper::clear_progress(state, &self.name)
-    }
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,

libafl/src/stages/replay.rs

@@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     corpus::{Corpus, CorpusId},
-    stages::Stage,
+    stages::{Restartable, Stage},
     state::{HasCorpus, HasSolutions},
     Error, Evaluator,
 };
@@ -114,15 +114,6 @@ where
     Z: Evaluator<E, EM, I, S>,
     I: Clone,
 {
-    fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
-        Ok(true)
-    }
-
-    fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> {
-        self.restart_helper.clear();
-        Ok(())
-    }
-
     #[inline]
     fn perform(
         &mut self,
@@ -169,3 +160,14 @@ where
         Ok(())
     }
 }
+
+impl<I, S> Restartable<S> for ReplayStage<I> {
+    fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
+        Ok(true)
+    }
+
+    fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> {
+        self.restart_helper.clear();
+        Ok(())
+    }
+}

libafl/src/stages/sync.rs

@@ -21,7 +21,7 @@ use crate::{
     executors::{Executor, ExitKind, HasObservers},
     fuzzer::{Evaluator, EvaluatorObservers, ExecutionProcessor},
     inputs::{Input, InputConverter},
-    stages::{RetryCountRestartHelper, Stage},
+    stages::{Restartable, RetryCountRestartHelper, Stage},
     state::{
         HasCorpus, HasCurrentTestcase, HasExecutions, HasRand, HasSolutions,
         MaybeHasClientPerfMonitor, Stoppable,
@@ -85,18 +85,6 @@ where
         + HasCurrentCorpusId
         + MaybeHasClientPerfMonitor,
 {
-    #[inline]
-    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
-        // TODO: Needs proper crash handling for when an imported testcase crashes
-        // For now, Make sure we don't get stuck crashing on this testcase
-        RetryCountRestartHelper::no_retry(state, &self.name)
-    }
-
-    #[inline]
-    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
-        RetryCountRestartHelper::clear_progress(state, &self.name)
-    }
-
     #[inline]
     fn perform(
         &mut self,
@@ -158,6 +146,23 @@ where
     }
 }
 
+impl<CB, E, EM, I, S, Z> Restartable<S> for SyncFromDiskStage<CB, E, EM, I, S, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
+    #[inline]
+    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
+        // TODO: Needs proper crash handling for when an imported testcase crashes
+        // For now, Make sure we don't get stuck crashing on this testcase
+        RetryCountRestartHelper::no_retry(state, &self.name)
+    }
+
+    #[inline]
+    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
+        RetryCountRestartHelper::clear_progress(state, &self.name)
+    }
+}
+
 impl<CB, E, EM, I, S, Z> SyncFromDiskStage<CB, E, EM, I, S, Z> {
     /// Creates a new [`SyncFromDiskStage`]
     #[must_use]
@@ -249,18 +254,6 @@ where
     SP: ShMemProvider<ShMem = SHM>,
     Z: EvaluatorObservers<E, EM, I, S> + ExecutionProcessor<EM, I, E::Observers, S>,
 {
-    #[inline]
-    fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
-        // No restart handling needed - does not execute the target.
-        Ok(true)
-    }
-
-    #[inline]
-    fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> {
-        // Not needed - does not execute the target.
-        Ok(())
-    }
-
     #[inline]
     fn perform(
         &mut self,
@@ -320,6 +313,20 @@ where
     }
 }
 
+impl<I, IC, ICB, S, SHM, SP> Restartable<S> for SyncFromBrokerStage<I, IC, ICB, S, SHM, SP> {
+    #[inline]
+    fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
+        // No restart handling needed - does not execute the target.
+        Ok(true)
+    }
+
+    #[inline]
+    fn clear_progress(&mut self, _state: &mut S) -> Result<(), Error> {
+        // Not needed - does not execute the target.
+        Ok(())
+    }
+}
+
 impl<I, IC, ICB, S, SHM, SP> SyncFromBrokerStage<I, IC, ICB, S, SHM, SP> {
     /// Creates a new [`SyncFromBrokerStage`]
     #[must_use]

libafl/src/stages/time_tracker.rs

@@ -3,7 +3,10 @@ use std::{marker::PhantomData, time::Duration};
 
 use libafl_bolts::{current_time, Error};
 
-use crate::{stages::Stage, HasMetadata};
+use crate::{
+    stages::{Restartable, Stage},
+    HasMetadata,
+};
 /// Track an inner Stage's execution time
 #[derive(Debug)]
 pub struct TimeTrackingStageWrapper<T, S, ST> {
@@ -48,7 +51,12 @@ where
         }
         Ok(())
     }
+}
 
+impl<T, S, ST> Restartable<S> for TimeTrackingStageWrapper<T, S, ST>
+where
+    ST: Restartable<S>,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         self.inner.should_restart(state)
     }

libafl/src/stages/tmin.rs

@@ -30,7 +30,7 @@ use crate::{
     schedulers::RemovableScheduler,
     stages::{
         mutational::{MutatedTransform, MutatedTransformPost},
-        ExecutionCountRestartHelper, Stage,
+        ExecutionCountRestartHelper, Restartable, Stage,
     },
     start_timer,
     state::{
@@ -82,14 +82,6 @@ where
     M: Mutator<I, S>,
     I: Input + Hash + HasLen,
 {
-    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
-        self.restart_helper.should_restart(state, &self.name)
-    }
-
-    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
-        self.restart_helper.clear_progress::<S>(state, &self.name)
-    }
-
     fn perform(
         &mut self,
         fuzzer: &mut Z,
@@ -106,6 +98,19 @@ where
     }
 }
 
+impl<E, EM, F, FF, I, M, S, Z> Restartable<S> for StdTMinMutationalStage<E, EM, F, FF, I, M, S, Z>
+where
+    S: HasNamedMetadata + HasExecutions,
+{
+    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
+        self.restart_helper.should_restart(state, &self.name)
+    }
+
+    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
+        self.restart_helper.clear_progress::<S>(state, &self.name)
+    }
+}
+
 impl<E, EM, F, FF, I, M, S, Z> FeedbackFactory<F, E::Observers>
     for StdTMinMutationalStage<E, EM, F, FF, I, M, S, Z>
 where

libafl/src/stages/tracing.rs

@@ -16,7 +16,7 @@ use crate::{
     inputs::Input,
     mark_feature_time,
     observers::ObserversTuple,
-    stages::{RetryCountRestartHelper, Stage},
+    stages::{Restartable, RetryCountRestartHelper, Stage},
     start_timer,
     state::{HasCorpus, HasCurrentTestcase, HasExecutions, MaybeHasClientPerfMonitor},
     Error, HasNamedMetadata,
@@ -93,7 +93,12 @@ where
     ) -> Result<(), Error> {
         self.trace(fuzzer, state, manager)
     }
+}
 
+impl<EM, I, TE, S, Z> Restartable<S> for TracingStage<EM, I, TE, S, Z>
+where
+    S: HasNamedMetadata + HasCurrentCorpusId,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         RetryCountRestartHelper::no_retry(state, &self.name)
     }
@@ -174,14 +179,6 @@ where
         + HasCurrentCorpusId
         + MaybeHasClientPerfMonitor,
 {
-    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
-        RetryCountRestartHelper::no_retry(state, &self.name)
-    }
-
-    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
-        RetryCountRestartHelper::clear_progress(state, &self.name)
-    }
-
     #[inline]
     fn perform(
         &mut self,
@@ -219,6 +216,19 @@ where
     }
 }
 
+impl<E, EM, I, SOT, S, Z> Restartable<S> for ShadowTracingStage<E, EM, I, SOT, S, Z>
+where
+    S: HasNamedMetadata + HasCurrentCorpusId,
+{
+    fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
+        RetryCountRestartHelper::no_retry(state, &self.name)
+    }
+
+    fn clear_progress(&mut self, state: &mut S) -> Result<(), Error> {
+        RetryCountRestartHelper::clear_progress(state, &self.name)
+    }
+}
+
 impl<E, EM, I, SOT, S, Z> ShadowTracingStage<E, EM, I, SOT, S, Z>
 where
     E: Executor<EM, I, S, Z> + HasObservers,

libafl/src/stages/tuneable.rs

@@ -14,7 +14,7 @@ use crate::{
     nonzero,
     stages::{
         mutational::{MutatedTransform, MutatedTransformPost, DEFAULT_MUTATIONAL_MAX_ITERATIONS},
-        ExecutionCountRestartHelper, MutationalStage, Stage,
+        ExecutionCountRestartHelper, MutationalStage, Restartable, Stage,
     },
     start_timer,
     state::{HasCurrentTestcase, HasExecutions, HasRand, MaybeHasClientPerfMonitor},
@@ -217,7 +217,12 @@ where
 
         ret
     }
+}
 
+impl<E, EM, I, M, S, Z> Restartable<S> for TuneableMutationalStage<E, EM, I, M, S, Z>
+where
+    S: HasNamedMetadata + HasExecutions,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         self.restart_helper.should_restart(state, &self.name)
     }

libafl/src/stages/unicode.rs

@@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     inputs::{BytesInput, HasTargetBytes},
-    stages::Stage,
+    stages::{Restartable, Stage},
     state::{HasCorpus, HasCurrentTestcase},
     HasMetadata,
 };
@@ -112,6 +112,18 @@ impl<E, EM, S, Z> Stage<E, EM, S, Z> for UnicodeIdentificationStage<BytesInput,
 where
     S: HasCorpus<BytesInput> + HasCurrentTestcase<BytesInput>,
 {
+    fn perform(
+        &mut self,
+        _fuzzer: &mut Z,
+        _executor: &mut E,
+        state: &mut S,
+        _manager: &mut EM,
+    ) -> Result<(), Error> {
+        UnicodeIdentificationStage::identify_unicode_in_current_testcase(state)
+    }
+}
+
+impl<S> Restartable<S> for UnicodeIdentificationStage<BytesInput, S> {
     #[inline]
     fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
         // Stage does not run the target. No reset helper needed.
@@ -123,14 +135,4 @@ where
         // Stage does not run the target. No reset helper needed.
         Ok(())
     }
-
-    fn perform(
-        &mut self,
-        _fuzzer: &mut Z,
-        _executor: &mut E,
-        state: &mut S,
-        _manager: &mut EM,
-    ) -> Result<(), Error> {
-        UnicodeIdentificationStage::identify_unicode_in_current_testcase(state)
-    }
 }

libafl/src/stages/verify_timeouts.rs

@@ -13,7 +13,7 @@ use crate::{
     executors::{Executor, HasObservers, HasTimeout},
     inputs::BytesInput,
     observers::ObserversTuple,
-    stages::Stage,
+    stages::{Restartable, Stage},
     Evaluator, HasMetadata,
 };
 
@@ -111,6 +111,9 @@ where
         *res = TimeoutsToVerify::<I>::new();
         Ok(())
     }
+}
+
+impl<E, I, S> Restartable<S> for VerifyTimeoutsStage<E, I, S> {
     fn should_restart(&mut self, _state: &mut S) -> Result<bool, Error> {
         Ok(true)
     }

libafl_targets/src/cmps/stages/aflpptracing.rs

@@ -6,7 +6,7 @@ use libafl::{
     executors::{Executor, HasObservers},
     inputs::BytesInput,
     observers::ObserversTuple,
-    stages::{colorization::TaintMetadata, RetryCountRestartHelper, Stage},
+    stages::{colorization::TaintMetadata, Restartable, RetryCountRestartHelper, Stage},
     state::{HasCorpus, HasCurrentTestcase},
     Error, HasMetadata, HasNamedMetadata,
 };
@@ -111,7 +111,12 @@ where
 
         Ok(())
     }
+}
 
+impl<EM, TE, S, Z> Restartable<S> for AFLppCmplogTracingStage<'_, EM, TE, S, Z>
+where
+    S: HasMetadata + HasNamedMetadata + HasCurrentCorpusId,
+{
     fn should_restart(&mut self, state: &mut S) -> Result<bool, Error> {
         // Tracing stage is always deterministic
         // don't restart