diff --git a/libafl/src/stages/logics.rs b/libafl/src/stages/logics.rs
index 9de107bd2a..5f431fca38 100644
--- a/libafl/src/stages/logics.rs
+++ b/libafl/src/stages/logics.rs
@@ -3,7 +3,7 @@
use core::marker::PhantomData;
use crate::{
- stages::{HasCurrentStage, HasNestedStageStatus, Stage, StagesTuple},
+ stages::{HasCurrentStage, HasNestedStageStatus, Stage, StageId, StagesTuple},
state::UsesState,
Error,
};
@@ -72,7 +72,8 @@ where
state: &mut E::State,
manager: &mut EM,
) -> Result<(), Error> {
- while state.current_stage()?.is_some() || (self.closure)(fuzzer, executor, state, manager)?
+ while state.current_stage_idx()?.is_some()
+ || (self.closure)(fuzzer, executor, state, manager)?
{
self.stages.perform_all(fuzzer, executor, state, manager)?;
}
@@ -150,7 +151,8 @@ where
state: &mut E::State,
manager: &mut EM,
) -> Result<(), Error> {
- if state.current_stage()?.is_some() || (self.closure)(fuzzer, executor, state, manager)? {
+ if state.current_stage_idx()?.is_some() || (self.closure)(fuzzer, executor, state, manager)?
+ {
self.if_stages
.perform_all(fuzzer, executor, state, manager)?;
}
@@ -231,21 +233,21 @@ where
state: &mut E::State,
manager: &mut EM,
) -> Result<(), Error> {
- let current = state.current_stage()?;
+ let current = state.current_stage_idx()?;
let fresh = current.is_none();
let closure_return = fresh && (self.closure)(fuzzer, executor, state, manager)?;
- if current == Some(0) || closure_return {
+ if current == Some(StageId(0)) || closure_return {
if fresh {
- state.set_stage(0)?;
+ state.set_current_stage_idx(StageId(0))?;
}
state.enter_inner_stage()?;
self.if_stages
.perform_all(fuzzer, executor, state, manager)?;
} else {
if fresh {
- state.set_stage(1)?;
+ state.set_current_stage_idx(StageId(1))?;
}
state.enter_inner_stage()?;
self.else_stages
diff --git a/libafl/src/stages/mod.rs b/libafl/src/stages/mod.rs
index 1b6813a9d8..96638b3fb3 100644
--- a/libafl/src/stages/mod.rs
+++ b/libafl/src/stages/mod.rs
@@ -5,7 +5,7 @@ Other stages may enrich [`crate::corpus::Testcase`]s with metadata.
*/
use alloc::{borrow::Cow, boxed::Box, vec::Vec};
-use core::marker::PhantomData;
+use core::{fmt, marker::PhantomData};
pub use calibrate::CalibrationStage;
pub use colorization::*;
@@ -153,7 +153,7 @@ where
stage: &mut S,
_: &mut EM,
) -> Result<(), Error> {
- if stage.current_stage()?.is_some() {
+ if stage.current_stage_idx()?.is_some() {
Err(Error::illegal_state(
"Got to the end of the tuple without completing resume.",
))
@@ -179,11 +179,11 @@ where
state: &mut Head::State,
manager: &mut EM,
) -> Result<(), Error> {
- match state.current_stage()? {
- Some(idx) if idx < Self::LEN => {
+ match state.current_stage_idx()? {
+ Some(idx) if idx < StageId(Self::LEN) => {
// do nothing; we are resuming
}
- Some(idx) if idx == Self::LEN => {
+ Some(idx) if idx == StageId(Self::LEN) => {
// perform the stage, but don't set it
let stage = &mut self.0;
@@ -191,12 +191,12 @@ where
state.clear_stage()?;
}
- Some(idx) if idx > Self::LEN => {
+ Some(idx) if idx > StageId(Self::LEN) => {
unreachable!("We should clear the stage index before we get here...");
}
// this is None, but the match can't deduce that
_ => {
- state.set_stage(Self::LEN)?;
+ state.set_current_stage_idx(StageId(Self::LEN))?;
let stage = &mut self.0;
stage.perform_restartable(fuzzer, executor, state, manager)?;
@@ -540,16 +540,27 @@ impl RetryRestartHelper {
}
}
+/// The index of a stage
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
+#[repr(transparent)]
+pub struct StageId(pub(crate) usize);
+
+impl fmt::Display for StageId {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "{}", self.0)
+ }
+}
+
/// Trait for types which track the current stage
pub trait HasCurrentStage {
/// Set the current stage; we have started processing this stage
- fn set_stage(&mut self, idx: usize) -> Result<(), Error>;
+ fn set_current_stage_idx(&mut self, idx: StageId) -> Result<(), Error>;
/// Clear the current stage; we are done processing this stage
fn clear_stage(&mut self) -> Result<(), Error>;
/// Fetch the current stage -- typically used after a state recovery or transfer
- fn current_stage(&self) -> Result