proptest_state_machine/

strategy.rs

//-
// Copyright 2023 The proptest developers
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Strategies used for abstract state machine testing.

use std::sync::atomic::{self, AtomicUsize};
use std::sync::Arc;

use proptest::bits::{BitSetLike, VarBitSet};
use proptest::collection::SizeRange;
use proptest::num::sample_uniform_incl;
use proptest::std_facade::fmt::{Debug, Formatter, Result};
use proptest::std_facade::Vec;
use proptest::strategy::BoxedStrategy;
use proptest::strategy::{NewTree, Strategy, ValueTree};
use proptest::test_runner::TestRunner;

/// This trait is used to model system under test as an abstract state machine.
///
/// The key to how this works is that the set of next valid transitions depends
/// on its current state (it's not the same as generating a random sequence of
/// transitions) and just like other prop strategies, the state machine strategy
/// attempts to shrink the transitions to find the minimal reproducible example
/// when it encounters a case that breaks any of the defined properties.
///
/// This is achieved with the [`ReferenceStateMachine::transitions`] that takes
/// the current state as an argument and can be used to decide which transitions
/// are valid from this state, together with the
/// [`ReferenceStateMachine::preconditions`], which are checked during generation
/// of transitions and during shrinking.
///
/// Hence, the `preconditions` only needs to contain checks for invariants that
/// depend on the current state and may be broken by shrinking and it doesn't
/// need to cover invariants that do not depend on the current state.
///
/// The reference state machine generation runs before the generated transitions
/// are attempted to be executed against the SUT (the concrete state machine)
/// as defined by [`crate::StateMachineTest`].
pub trait ReferenceStateMachine: 'static {
    /// The reference state machine's state type. This should contain the minimum
    /// required information needed to implement the state machine. It is used
    /// to drive the generations of transitions to decide which transitions are
    /// valid for the current state.
    type State: Clone + Debug;

    /// The reference state machine's transition type. This is typically an enum
    /// with its variants containing the parameters required to apply the
    /// transition, if any.
    type Transition: Clone + Debug;

    // TODO Instead of the boxed strategies, this could use
    // <https://github.com/rust-lang/rust/issues/63063> once stabilized:
    // type StateStrategy = impl Strategy<Value = Self::State>;
    // type TransitionStrategy = impl Strategy<Value = Self::Transition>;

    /// The initial state may be generated by any strategy. For a constant
    /// initial state, use [`proptest::strategy::Just`].
    fn init_state() -> BoxedStrategy<Self::State>;

    /// Generate the initial transitions.
    fn transitions(state: &Self::State) -> BoxedStrategy<Self::Transition>;

    /// Apply a transition in the reference state.
    fn apply(state: Self::State, transition: &Self::Transition) -> Self::State;

    /// Pre-conditions may be specified to control which transitions are valid
    /// from the current state. If not overridden, this allows any transition.
    /// The pre-conditions are checked in the generated transitions and during
    /// shrinking.
    ///
    /// The pre-conditions checking relies on proptest global rejection
    /// filtering, which comes with some [disadvantages](https://altsysrq.github.io/proptest-book/proptest/tutorial/filtering.html).
    /// This means that pre-conditions that are hard to satisfy might slow down
    /// the test or even fail by exceeding the maximum rejection count.
    fn preconditions(
        state: &Self::State,
        transition: &Self::Transition,
    ) -> bool {
        // This is to avoid `unused_variables` warning
        let _ = (state, transition);

        true
    }

    /// A sequential strategy runs the state machine transitions generated from
    /// the reference model sequentially in a test over a concrete state, which
    /// can be implemented with the help of
    /// [`crate::StateMachineTest`] trait.
    ///
    /// You typically never need to override this method.
    fn sequential_strategy(
        size: impl Into<SizeRange>,
    ) -> Sequential<
        Self::State,
        Self::Transition,
        BoxedStrategy<Self::State>,
        BoxedStrategy<Self::Transition>,
    > {
        Sequential::new(
            size.into(),
            Self::init_state,
            Self::preconditions,
            Self::transitions,
            Self::apply,
        )
    }
}

/// In a sequential state machine strategy, we first generate an acceptable
/// sequence of transitions. That is a sequence that satisfies the given
/// pre-conditions. The acceptability of each transition in the sequence depends
/// on the current state of the state machine, which is updated by the
/// transitions with the `next` function.
///
/// The shrinking strategy is to iteratively apply `Shrink::InitialState`,
/// `Shrink::DeleteTransition` and `Shrink::Transition`.
///
/// 1. We start by trying to delete transitions from the back of the list that
///    were never seen by the test, if any. Note that because proptest expects
///    deterministic results in for reproducible issues, unlike the following
///    steps this step will not be undone on `complicate`. If there were any
///    unseen transitions, then the next step will start at trying to delete
///    the transition before the last one seen as we know that the last
///    transition cannot be deleted as it's the one that has failed.
/// 2. Then, we keep trying to delete transitions from the back of the list, until
///    we can do so no further (reached the beginning of the list)..
/// 3. Then, we again iteratively attempt to shrink the individual transitions,
///    but this time starting from the front of the list - i.e. from the first
///    transition to be applied.
/// 4. Finally, we try to shrink the initial state until it's not possible to
///    shrink it any further.
///
/// For `complicate`, we attempt to undo the last shrink operation, if there was
/// any.
pub struct Sequential<State, Transition, StateStrategy, TransitionStrategy> {
    size: SizeRange,
    init_state: Arc<dyn Fn() -> StateStrategy + Send + Sync>,
    preconditions: Arc<dyn Fn(&State, &Transition) -> bool + Send + Sync>,
    transitions: Arc<dyn Fn(&State) -> TransitionStrategy + Send + Sync>,
    next: Arc<dyn Fn(State, &Transition) -> State + Send + Sync>,
}

impl<State, Transition, StateStrategy, TransitionStrategy>
    Sequential<State, Transition, StateStrategy, TransitionStrategy>
where
    State: 'static,
    Transition: 'static,
    StateStrategy: 'static,
    TransitionStrategy: 'static,
{
    pub fn new(
        size: SizeRange,
        init_state: impl Fn() -> StateStrategy + 'static + Send + Sync,
        preconditions: impl Fn(&State, &Transition) -> bool + 'static + Send + Sync,
        transitions: impl Fn(&State) -> TransitionStrategy + 'static + Send + Sync,
        next: impl Fn(State, &Transition) -> State + 'static + Send + Sync,
    ) -> Self {
        Self {
            size,
            init_state: Arc::new(init_state),
            preconditions: Arc::new(preconditions),
            transitions: Arc::new(transitions),
            next: Arc::new(next),
        }
    }
}

impl<State, Transition, StateStrategy, TransitionStrategy> Debug
    for Sequential<State, Transition, StateStrategy, TransitionStrategy>
{
    fn fmt(&self, f: &mut Formatter) -> Result {
        f.debug_struct("Sequential")
            .field("size", &self.size)
            .finish()
    }
}

impl<
        State: Clone + Debug,
        Transition: Clone + Debug,
        StateStrategy: Strategy<Value = State>,
        TransitionStrategy: Strategy<Value = Transition>,
    > Strategy
    for Sequential<State, Transition, StateStrategy, TransitionStrategy>
{
    type Tree = SequentialValueTree<
        State,
        Transition,
        StateStrategy::Tree,
        TransitionStrategy::Tree,
    >;
    type Value = (State, Vec<Transition>, Option<Arc<AtomicUsize>>);

    fn new_tree(&self, runner: &mut TestRunner) -> NewTree<Self> {
        // Generate the initial state value tree
        let initial_state = (self.init_state)().new_tree(runner)?;
        let last_valid_initial_state = initial_state.current();

        let (min_size, end) = self.size.start_end_incl();
        // Sample the maximum number of the transitions from the size range
        let max_size = sample_uniform_incl(runner, min_size, end);
        let mut transitions = Vec::with_capacity(max_size);
        let mut acceptable_transitions = Vec::with_capacity(max_size);
        let included_transitions = VarBitSet::saturated(max_size);
        let shrinkable_transitions = VarBitSet::saturated(max_size);

        // Sample the transitions until we reach the `max_size`
        let mut state = initial_state.current();
        while transitions.len() < max_size {
            // Apply the current state to find the current transition
            let transition_tree =
                (self.transitions)(&state).new_tree(runner)?;
            let transition = transition_tree.current();

            // If the pre-conditions are satisfied, use the transition
            if (self.preconditions)(&state, &transition) {
                transitions.push(transition_tree);
                state = (self.next)(state, &transition);
                acceptable_transitions
                    .push((TransitionState::Accepted, transition));
            } else {
                runner.reject_local("Pre-conditions were not satisfied")?;
            }
        }

        // The maximum index into the vectors and bit sets
        let max_ix = max_size - 1;

        Ok(SequentialValueTree {
            initial_state,
            is_initial_state_shrinkable: true,
            last_valid_initial_state,
            preconditions: self.preconditions.clone(),
            next: self.next.clone(),
            transitions,
            acceptable_transitions,
            included_transitions,
            shrinkable_transitions,
            max_ix,
            // On a failure, we start by shrinking transitions from the back
            // which is less likely to invalidate pre-conditions
            shrink: Shrink::DeleteTransition(max_ix),
            last_shrink: None,
            seen_transitions_counter: Some(Default::default()),
        })
    }
}

/// A shrinking operation
#[derive(Clone, Copy, Debug)]
enum Shrink {
    /// Shrink the initial state
    InitialState,
    /// Delete a transition at given index
    DeleteTransition(usize),
    /// Shrink a transition at given index
    Transition(usize),
}
use Shrink::*;

/// The state of a transition in the model
#[derive(Clone, Copy, Debug)]
enum TransitionState {
    /// The transition that is equal to the result of `ValueTree::current()`
    /// and satisfies the pre-conditions
    Accepted,
    /// The transition has been simplified, but rejected by pre-conditions
    SimplifyRejected,
    /// The transition has been complicated, but rejected by pre-conditions
    ComplicateRejected,
}
use TransitionState::*;

/// The generated value tree for a sequential state machine.
pub struct SequentialValueTree<
    State,
    Transition,
    StateValueTree,
    TransitionValueTree,
> {
    /// The initial state value tree
    initial_state: StateValueTree,
    /// Can the `initial_state` be shrunk any further?
    is_initial_state_shrinkable: bool,
    /// The last initial state that has been accepted by the pre-conditions.
    /// We have to store this every time before attempt to shrink to be able
    /// to back to it in case the shrinking is rejected.
    last_valid_initial_state: State,
    /// The pre-conditions predicate
    preconditions: Arc<dyn Fn(&State, &Transition) -> bool>,
    /// The function from current state and a transition to an updated state
    next: Arc<dyn Fn(State, &Transition) -> State>,
    /// The list of transitions' value trees
    transitions: Vec<TransitionValueTree>,
    /// The sequence of included transitions with their shrinking state
    acceptable_transitions: Vec<(TransitionState, Transition)>,
    /// The bit-set of transitions that have not been deleted by shrinking
    included_transitions: VarBitSet,
    /// The bit-set of transitions that can be shrunk further
    shrinkable_transitions: VarBitSet,
    /// The maximum index in the `transitions` vector (its size - 1)
    max_ix: usize,
    /// The next shrink operation to apply
    shrink: Shrink,
    /// The last applied shrink operation, if any
    last_shrink: Option<Shrink>,
    /// The number of transitions that were seen by the test runner.
    /// On a test run this is shared with `StateMachineTest::test_sequential`
    /// which increments the inner counter value on every transition. If the
    /// test fails, the counter is used to remove any unseen transitions before
    /// shrinking and this field is set to `None` as it's no longer needed for
    /// shrinking.
    seen_transitions_counter: Option<Arc<AtomicUsize>>,
}

impl<
        State: Clone + Debug,
        Transition: Clone + Debug,
        StateValueTree: ValueTree<Value = State>,
        TransitionValueTree: ValueTree<Value = Transition>,
    >
    SequentialValueTree<State, Transition, StateValueTree, TransitionValueTree>
{
    /// Try to apply the next `self.shrink`. Returns `true` if a shrink has been
    /// applied.
    fn try_simplify(&mut self) -> bool {
        if let Some(seen_transitions_counter) =
            self.seen_transitions_counter.as_ref()
        {
            let seen_count =
                seen_transitions_counter.load(atomic::Ordering::SeqCst);

            let included_count = self.included_transitions.count();

            if seen_count < included_count {
                // the test runner did not see all the transitions so we can
                // delete the transitions that were not seen because they were
                // not executed

                let mut kept_count = 0;
                for ix in 0..self.transitions.len() {
                    if self.included_transitions.test(ix) {
                        // transition at ix was part of test

                        if kept_count < seen_count {
                            // transition at xi was seen by the test or we are
                            // still below minimum size for the test
                            kept_count += 1;
                        } else {
                            // transition at ix was never seen
                            self.included_transitions.clear(ix);
                            self.shrinkable_transitions.clear(ix);
                        }
                    }
                }
                // Set the next shrink based on how many transitions were seen:
                // - If 0 seen: go directly to shrinking the initial state.
                // - If 1 seen: can't delete any more, so shrink individual transitions.
                // - If >1 seen: delete the transition before the last seen transition.
                //   (subtract 2 from `kept_count` because the last seen transition
                //   caused the failure).
                if kept_count == 0 {
                    self.shrink = InitialState;
                } else if kept_count == 1 {
                    self.shrink = Transition(0);
                } else {
                    self.shrink = DeleteTransition(
                        kept_count.checked_sub(2).unwrap_or_default(),
                    );
                }
            }

            // Remove the seen transitions counter for shrinking runs
            self.seen_transitions_counter = None;
        }

        if let DeleteTransition(ix) = self.shrink {
            // Delete the index from the included transitions
            self.included_transitions.clear(ix);

            self.last_shrink = Some(self.shrink);
            self.shrink = if ix == 0 {
                // Reached the beginning of the list, move on to shrinking
                Transition(0)
            } else {
                // Try to delete the previous transition next
                DeleteTransition(ix - 1)
            };
            // If this delete is not acceptable, undo it and try again
            if !self
                .check_acceptable(None, self.last_valid_initial_state.clone())
            {
                self.included_transitions.set(ix);
                self.last_shrink = None;
                return self.try_simplify();
            }
            // If the delete was accepted, remove this index from shrinkable
            // transitions
            self.shrinkable_transitions.clear(ix);
            return true;
        }

        while let Transition(ix) = self.shrink {
            if self.shrinkable_transitions.count() == 0 {
                // Move on to shrinking the initial state
                self.shrink = Shrink::InitialState;
                break;
            }

            if !self.included_transitions.test(ix) {
                // No use shrinking something we're not including
                self.shrink = self.next_shrink_transition(ix);
                continue;
            }

            if let Some((SimplifyRejected, _trans)) =
                self.acceptable_transitions.get(ix)
            {
                // This transition is already simplified and rejected
                self.shrink = self.next_shrink_transition(ix);
            } else if self.transitions[ix].simplify() {
                self.last_shrink = Some(self.shrink);
                if self.check_acceptable(
                    Some(ix),
                    self.last_valid_initial_state.clone(),
                ) {
                    self.acceptable_transitions[ix] =
                        (Accepted, self.transitions[ix].current());
                    return true;
                } else {
                    let (state, _trans) =
                        self.acceptable_transitions.get_mut(ix).unwrap();
                    *state = SimplifyRejected;
                    self.shrinkable_transitions.clear(ix);
                    self.shrink = self.next_shrink_transition(ix);
                    return self.simplify();
                }
            } else {
                self.shrinkable_transitions.clear(ix);
                self.shrink = self.next_shrink_transition(ix);
            }
        }

        if let InitialState = self.shrink {
            if self.initial_state.simplify() {
                if self.check_acceptable(None, self.initial_state.current()) {
                    self.last_valid_initial_state =
                        self.initial_state.current();
                    self.last_shrink = Some(self.shrink);
                    return true;
                } else {
                    // If the shrink is not acceptable, clear it out
                    self.last_shrink = None;

                    // `initial_state` is "dirty" here but we won't ever use it again because it is unshrinkable from here.
                }
            }
            self.is_initial_state_shrinkable = false;
            // Nothing left to do
            return false;
        }

        // This statement should never be reached
        panic!("Unexpected shrink state");
    }

    /// Find if there's any acceptable included transition that is not current,
    /// starting from the given index. Expects that all the included transitions
    /// are currently being rejected (when `can_simplify` returns `false`).
    fn try_to_find_acceptable_transition(&mut self, ix: usize) -> bool {
        let mut ix_to_check = ix;
        loop {
            if self.included_transitions.test(ix_to_check)
                && self.check_acceptable(
                    Some(ix_to_check),
                    self.last_valid_initial_state.clone(),
                )
            {
                self.acceptable_transitions[ix_to_check] =
                    (Accepted, self.transitions[ix_to_check].current());
                return true;
            }
            // Move on to the next transition
            if ix_to_check == self.max_ix {
                ix_to_check = 0;
            } else {
                ix_to_check += 1;
            }
            // We're back to where we started, there nothing left to do
            if ix_to_check == ix {
                return false;
            }
        }
    }

    /// Check if the sequence of included transitions is acceptable by the
    /// pre-conditions. When `ix` is not `None`, the transition at the given
    /// index is taken from its current value.
    fn check_acceptable(&self, ix: Option<usize>, mut state: State) -> bool {
        let transitions = self.get_included_acceptable_transitions(ix);
        for transition in transitions.iter() {
            let is_acceptable = (self.preconditions)(&state, transition);
            if is_acceptable {
                state = (self.next)(state, transition);
            } else {
                return false;
            }
        }
        true
    }

    /// The currently included and acceptable transitions. When `ix` is not
    /// `None`, the transition at this index is taken from its current value
    /// which may not be acceptable by the pre-conditions, instead of its
    /// acceptable value.
    fn get_included_acceptable_transitions(
        &self,
        ix: Option<usize>,
    ) -> Vec<Transition> {
        self.acceptable_transitions
            .iter()
            .enumerate()
            // Filter out deleted transitions
            .filter(|&(this_ix, _)| self.included_transitions.test(this_ix))
            // Map the indices to the values
            .map(|(this_ix, (_, transition))| match ix {
                Some(ix) if this_ix == ix => self.transitions[ix].current(),
                _ => transition.clone(),
            })
            .collect()
    }

    /// Find if the initial state is still shrinkable or if any of the
    /// simplifications and complications of the included transitions have not
    /// yet been rejected.
    fn can_simplify(&self) -> bool {
        self.is_initial_state_shrinkable ||
             // If there are some transitions whose shrinking has not yet been
             // rejected, we can try to shrink them further
             !self
                .acceptable_transitions
                .iter()
                .enumerate()
                // Filter out deleted transitions
                .filter(|&(ix, _)| self.included_transitions.test(ix))
                .all(|(_, (state, _transition))| {
                    matches!(state, SimplifyRejected | ComplicateRejected)
                })
    }

    /// Find the next shrink transition. Loops back to the front of the list
    /// when the end is reached, because sometimes a transition might become
    /// acceptable only after a transition that comes before it in the sequence
    /// gets shrunk.
    fn next_shrink_transition(&self, current_ix: usize) -> Shrink {
        if current_ix == self.max_ix {
            // Either loop back to the start of the list...
            Transition(0)
        } else {
            // ...or move on to the next transition
            Transition(current_ix + 1)
        }
    }
}

impl<
        State: Clone + Debug,
        Transition: Clone + Debug,
        StateValueTree: ValueTree<Value = State>,
        TransitionValueTree: ValueTree<Value = Transition>,
    > ValueTree
    for SequentialValueTree<
        State,
        Transition,
        StateValueTree,
        TransitionValueTree,
    >
{
    type Value = (State, Vec<Transition>, Option<Arc<AtomicUsize>>);

    fn current(&self) -> Self::Value {
        if let Some(seen_transitions_counter) = &self.seen_transitions_counter {
            if seen_transitions_counter.load(atomic::Ordering::SeqCst) > 0 {
                panic!("Unexpected non-zero `seen_transitions_counter`");
            }
        }

        (
            self.last_valid_initial_state.clone(),
            // The current included acceptable transitions
            self.get_included_acceptable_transitions(None),
            self.seen_transitions_counter.clone(),
        )
    }

    fn simplify(&mut self) -> bool {
        let was_simplified = if self.can_simplify() {
            self.try_simplify()
        } else if let Some(Transition(ix)) = self.last_shrink {
            self.try_to_find_acceptable_transition(ix)
        } else {
            false
        };

        // reset seen transactions counter for next run
        self.seen_transitions_counter = Default::default();

        was_simplified
    }

    fn complicate(&mut self) -> bool {
        // reset seen transactions counter for next run
        self.seen_transitions_counter = Default::default();

        match &self.last_shrink {
            None => false,
            Some(DeleteTransition(ix)) => {
                // Undo the last item we deleted. Can't complicate any further,
                // so unset prev_shrink.
                self.included_transitions.set(*ix);
                self.shrinkable_transitions.set(*ix);
                self.last_shrink = None;
                true
            }
            Some(Transition(ix)) => {
                let ix = *ix;
                if self.transitions[ix].complicate() {
                    if self.check_acceptable(
                        Some(ix),
                        self.last_valid_initial_state.clone(),
                    ) {
                        self.acceptable_transitions[ix] =
                            (Accepted, self.transitions[ix].current());
                        // Don't unset prev_shrink; we may be able to complicate
                        // it again
                        return true;
                    } else {
                        let (state, _trans) =
                            self.acceptable_transitions.get_mut(ix).unwrap();
                        *state = ComplicateRejected;
                    }
                }
                // Can't complicate the last element any further
                self.last_shrink = None;
                false
            }
            Some(InitialState) => {
                if self.initial_state.complicate()
                    && self.check_acceptable(None, self.initial_state.current())
                {
                    self.last_valid_initial_state =
                        self.initial_state.current();
                    // Don't unset prev_shrink; we may be able to complicate
                    // it again
                    return true;
                }
                // Can't complicate the initial state any further
                self.last_shrink = None;
                false
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    use proptest::collection::hash_set;
    use proptest::prelude::*;

    use heap_state_machine::*;
    use std::collections::HashSet;

    /// A number of simplifications that can be applied in the `ValueTree`
    /// produced by [`deterministic_sequential_value_tree`]. It depends on the
    /// [`TRANSITIONS`] given to its `sequential_strategy`.
    ///
    /// This constant can be determined from the test
    /// `number_of_sequential_value_tree_simplifications`.
    const SIMPLIFICATIONS: usize = 32;
    /// Number of transitions in the [`deterministic_sequential_value_tree`].
    const TRANSITIONS: usize = 32;

    #[test]
    fn number_of_sequential_value_tree_simplifications() {
        let mut value_tree = deterministic_sequential_value_tree();
        value_tree
            .seen_transitions_counter
            .as_mut()
            .unwrap()
            .store(TRANSITIONS, atomic::Ordering::SeqCst);

        let mut i = 0;
        loop {
            let simplified = value_tree.simplify();
            if simplified {
                i += 1;
            } else {
                break;
            }
        }
        assert_eq!(i, SIMPLIFICATIONS);
    }

    proptest! {
        /// Test the simplifications and complication of the
        /// `SequentialValueTree` produced by
        /// `deterministic_sequential_value_tree`.
        ///
        /// The indices of simplification on which we'll attempt to complicate
        /// after simplification are selected from the randomly generated
        /// `complicate_ixs`.
        ///
        /// Every simplification and complication must satisfy pre-conditions of
        /// the state-machine.
        #[test]
        fn test_state_machine_sequential_value_tree(
            complicate_ixs in hash_set(0..SIMPLIFICATIONS, 0..SIMPLIFICATIONS)
        ) {
            test_state_machine_sequential_value_tree_aux(complicate_ixs)
        }
    }

    fn test_state_machine_sequential_value_tree_aux(
        complicate_ixs: HashSet<usize>,
    ) {
        println!("Complicate indices: {complicate_ixs:?}");

        let mut value_tree = deterministic_sequential_value_tree();

        let check_preconditions = |value_tree: &TestValueTree| {
            let (mut state, transitions, _seen_counter) = value_tree.current();
            let len = transitions.len();
            println!("Transitions {}", len);
            for (ix, transition) in transitions.into_iter().enumerate() {
                println!("Transition {}/{len} {transition:?}", ix + 1);
                // Every transition must satisfy the pre-conditions
                assert!(
                    <HeapStateMachine as ReferenceStateMachine>::preconditions(
                        &state,
                        &transition
                    )
                );

                // Apply the transition to update the state for the next transition
                state = <HeapStateMachine as ReferenceStateMachine>::apply(
                    state,
                    &transition,
                );
            }
        };

        let mut ix = 0_usize;
        loop {
            let simplified = value_tree.simplify();

            check_preconditions(&value_tree);

            if !simplified {
                break;
            }
            ix += 1;

            if complicate_ixs.contains(&ix) {
                loop {
                    let complicated = value_tree.complicate();

                    check_preconditions(&value_tree);

                    if !complicated {
                        break;
                    }
                }
            }
        }
    }

    proptest! {
        /// Test the initial simplifications of the `SequentialValueTree` produced
        /// by `deterministic_sequential_value_tree`.
        ///
        /// We want to make sure that we initially remove the transitions that
        /// where not seen.
        #[test]
        fn test_value_tree_initial_simplification(
            len in 10usize..100,
        ) {
            test_value_tree_initial_simplification_aux(len)
        }
    }

    fn test_value_tree_initial_simplification_aux(len: usize) {
        let sequential =
            <HeapStateMachine as ReferenceStateMachine>::sequential_strategy(
                ..len,
            );

        let mut runner = TestRunner::deterministic();
        let mut value_tree = sequential.new_tree(&mut runner).unwrap();

        let (_, transitions, mut seen_counter) = value_tree.current();

        let num_seen = transitions.len() / 2;
        let seen_counter = seen_counter.as_mut().unwrap();
        seen_counter.store(num_seen, atomic::Ordering::SeqCst);

        let mut seen_before_complication =
            transitions.into_iter().take(num_seen).collect::<Vec<_>>();

        assert!(value_tree.simplify());

        let (_, transitions, _seen_counter) = value_tree.current();

        let seen_after_first_complication =
            transitions.into_iter().collect::<Vec<_>>();

        // After the unseen transitions are removed, the shrink behavior depends
        // on how many transitions were seen:
        // - If > 1 seen: delete the transition before the last seen one
        // - If = 1 seen: can't delete any more, may start individual transition shrinking
        if seen_before_complication.len() > 1 {
            let last = seen_before_complication.pop().unwrap();
            seen_before_complication.pop();
            seen_before_complication.push(last);
            assert_eq!(
                seen_before_complication, seen_after_first_complication,
                "only seen transitions should be present after first simplification"
            );
        } else {
            // When there's only 1 seen transition, we expect it to be preserved.
            assert!(
                !seen_after_first_complication.is_empty(),
                "When only 1 transition was seen, at least 1 should remain after simplification"
            );
            assert!(
                matches!(value_tree.shrink, Transition(0)),
                "When only 1 transition was seen, shrink should be set to Transition(0)"
            );
        }
    }

    #[test]
    fn test_call_to_current_with_non_zero_seen_counter() {
        let result = std::panic::catch_unwind(|| {
            let value_tree = deterministic_sequential_value_tree();

            let (_, _transitions1, mut seen_counter) = value_tree.current();
            {
                let seen_counter = seen_counter.as_mut().unwrap();
                seen_counter.store(1, atomic::Ordering::SeqCst);
            }
            drop(seen_counter);

            let _transitions2 = value_tree.current();
        })
        .expect_err("should panic");

        let s = "Unexpected non-zero `seen_transitions_counter`";
        assert_eq!(result.downcast_ref::<&str>(), Some(&s));
    }

    /// The following is a definition of an reference state machine used for the
    /// tests.
    mod heap_state_machine {
        use std::vec::Vec;

        use crate::{ReferenceStateMachine, SequentialValueTree};
        use proptest::prelude::*;
        use proptest::test_runner::TestRunner;

        use super::TRANSITIONS;

        pub struct HeapStateMachine;

        pub type TestValueTree = SequentialValueTree<
            TestState,
            TestTransition,
            <BoxedStrategy<TestState> as Strategy>::Tree,
            <BoxedStrategy<TestTransition> as Strategy>::Tree,
        >;

        pub type TestState = Vec<i32>;

        #[derive(Clone, Debug, PartialEq)]
        pub enum TestTransition {
            PopNonEmpty,
            PopEmpty,
            Push(i32),
        }

        pub fn deterministic_sequential_value_tree() -> TestValueTree {
            let sequential =
                <HeapStateMachine as ReferenceStateMachine>::sequential_strategy(
                    TRANSITIONS,
                );
            let mut runner = TestRunner::deterministic();
            sequential.new_tree(&mut runner).unwrap()
        }

        impl ReferenceStateMachine for HeapStateMachine {
            type State = TestState;
            type Transition = TestTransition;

            fn init_state() -> BoxedStrategy<Self::State> {
                Just(vec![]).boxed()
            }

            fn transitions(
                state: &Self::State,
            ) -> BoxedStrategy<Self::Transition> {
                if state.is_empty() {
                    prop_oneof![
                        1 => Just(TestTransition::PopEmpty),
                        2 => (any::<i32>()).prop_map(TestTransition::Push),
                    ]
                    .boxed()
                } else {
                    prop_oneof![
                        1 => Just(TestTransition::PopNonEmpty),
                        2 => (any::<i32>()).prop_map(TestTransition::Push),
                    ]
                    .boxed()
                }
            }

            fn apply(
                mut state: Self::State,
                transition: &Self::Transition,
            ) -> Self::State {
                match transition {
                    TestTransition::PopEmpty => {
                        state.pop();
                    }
                    TestTransition::PopNonEmpty => {
                        state.pop();
                    }
                    TestTransition::Push(value) => state.push(*value),
                }
                state
            }

            fn preconditions(
                state: &Self::State,
                transition: &Self::Transition,
            ) -> bool {
                match transition {
                    TestTransition::PopEmpty => state.is_empty(),
                    TestTransition::PopNonEmpty => !state.is_empty(),
                    TestTransition::Push(_) => true,
                }
            }
        }
    }

    /// A tests that verifies that the strategy finds a simplest failing case, and
    /// that this simplest failing case is ultimately reported by the test runner,
    /// as opposed to reporting input that actually passes the test.
    ///
    /// This module defines a state machine test that is designed to fail.
    /// The reference state machine consists of a lower bound the acceptable value
    /// of a transition. And the test fails if an unacceptably low transition
    /// value is observed, given the reference state's limit.
    ///
    /// This intentionally-failing state machine test is then run inside a proptest
    /// to verify that it reports a simplest failing input when it fails.
    mod find_simplest_failure {
        use proptest::prelude::*;
        use proptest::strategy::BoxedStrategy;
        use proptest::test_runner::TestRng;
        use proptest::{
            collection,
            strategy::Strategy,
            test_runner::{Config, TestError, TestRunner},
        };

        use crate::{ReferenceStateMachine, StateMachineTest};

        const MIN_TRANSITION: u32 = 10;
        const MAX_TRANSITION: u32 = 20;

        const MIN_LIMIT: u32 = 2;
        const MAX_LIMIT: u32 = 50;

        #[derive(Debug, Default, Clone)]
        struct FailIfLessThan(u32);
        impl ReferenceStateMachine for FailIfLessThan {
            type State = Self;
            type Transition = u32;

            fn init_state() -> BoxedStrategy<Self> {
                (MIN_LIMIT..MAX_LIMIT).prop_map(FailIfLessThan).boxed()
            }

            fn transitions(_: &Self::State) -> BoxedStrategy<u32> {
                (MIN_TRANSITION..MAX_TRANSITION).boxed()
            }

            fn apply(state: Self::State, _: &Self::Transition) -> Self::State {
                state
            }
        }

        /// Defines a test that is intended to fail, so that we can inspect the
        /// failing input.
        struct FailIfLessThanTest;
        impl StateMachineTest for FailIfLessThanTest {
            type SystemUnderTest = ();
            type Reference = FailIfLessThan;

            fn init_test(ref_state: &FailIfLessThan) {
                println!();
                println!("starting {ref_state:?}");
            }

            fn apply(
                (): Self::SystemUnderTest,
                ref_state: &FailIfLessThan,
                transition: u32,
            ) -> Self::SystemUnderTest {
                // Fail on any transition that is less than the ref state's limit.
                let FailIfLessThan(limit) = ref_state;
                println!("{transition} < {}?", limit);
                if transition < ref_state.0 {
                    panic!("{transition} < {}", limit);
                }
            }
        }

        proptest! {
            #[test]
            fn test_returns_simplest_failure(
                seed in collection::vec(any::<u8>(), 32).no_shrink()) {

                // We need to explicitly run create a runner so that we can
                // inspect the output, and determine if it does return an input that
                // should fail, and is minimal.
                let mut runner = TestRunner::new_with_rng(
                    Config::default(), TestRng::from_seed(Default::default(), &seed));
                let result = runner.run(
                    &FailIfLessThan::sequential_strategy(10..50_usize),
                    |(ref_state, transitions, seen_counter)| {
                        Ok(FailIfLessThanTest::test_sequential(
                            Default::default(),
                            ref_state,
                            transitions,
                            seen_counter,
                        ))
                    },
                );
                if let Err(TestError::Fail(
                    _,
                    (FailIfLessThan(limit), transitions, _seen_counter),
                )) = result
                {
                    assert_eq!(transitions.len(), 1, "The minimal failing case should be ");
                    assert_eq!(limit, MIN_TRANSITION + 1);
                    assert!(transitions.into_iter().next().unwrap() < limit);
                } else {
                    prop_assume!(false,
                        "If the state machine doesn't fail as intended, we need a case that fails.");
                }
            }
        }
    }

    #[test]
    fn test_zero_seen_transitions_optimization() {
        // Test that when 0 transitions are seen, we go directly to InitialState shrinking
        let mut value_tree = deterministic_sequential_value_tree();

        // Simulate that no transitions were seen (kept_count = 0)
        value_tree
            .seen_transitions_counter
            .as_mut()
            .unwrap()
            .store(0, atomic::Ordering::SeqCst);

        // Call simplify - this should trigger the optimization
        let simplified = value_tree.simplify();

        assert_eq!(value_tree.included_transitions.count(), 0,
            "All transitions should be removed when none were seen");
        assert!(matches!(value_tree.shrink, InitialState),
            "Shrink should be set to InitialState when kept_count == 0");

        // The HeapStateMachine uses Just(vec![]) for initial state, which is not shrinkable
        // So simplify() should return false, but the optimization still works correctly
        assert!(!simplified,
            "Simplification should return false since initial state (Just(vec![])) is not shrinkable");

        let (_, transitions, _) = value_tree.current();
        assert!(transitions.is_empty(),
            "No transitions should remain when none were seen");
    }
}