shuttle 0.8.1

use crate::runtime::task::{Task, TaskId, DEFAULT_INLINE_TASKS};
use crate::scheduler::{Schedule, Scheduler};
use smallvec::SmallVec;

#[derive(Clone, Debug, PartialEq, Eq)]
enum ScheduleRecord {
    Task(Option<TaskId>, SmallVec<[TaskId; DEFAULT_INLINE_TASKS]>, bool),
    Random(u64),
}

/// An `UncontrolledNondeterminismCheckScheduler` checks whether a given program exhibits uncontrolled
/// nondeterminism by wrapping an inner `Scheduler`, and, for each schedule generated by that scheduler,
/// replaying the schedule once. When doing the replay, we check that the schedule is still
/// valid and that the set of runnable tasks is the same at each step.
/// Violations of these checks means that the program exhibits nondeterminism which is not
/// under Shuttle's control. Note that the opposite is not true — there are no guarantees that
/// the program under test does not have uncontrolled nondeterminism if it passes a run of
/// the `UncontrolledNondeterminismCheckScheduler`, even in the case where the wrapped `scheduler` is exhaustive.
#[derive(Debug)]
pub struct UncontrolledNondeterminismCheckScheduler<S: Scheduler> {
    scheduler: Box<S>,
    recording: bool,
    previous_schedule: Vec<ScheduleRecord>,
    current_step: usize,
}

impl<S: Scheduler> UncontrolledNondeterminismCheckScheduler<S> {
    /// Create a new `UncontrolledNondeterminismCheckScheduler` by wrapping the given `Scheduler` implementation.
    pub fn new(scheduler: S) -> Self {
        Self {
            scheduler: Box::new(scheduler),
            previous_schedule: Vec::new(),
            recording: false,
            current_step: 0,
        }
    }
}

impl<S: Scheduler> Scheduler for UncontrolledNondeterminismCheckScheduler<S> {
    fn new_execution(&mut self) -> Option<Schedule> {
        // Dummy schedule. We do this instead of doing `self.scheduler.new_execution`
        // as that would cause the schedule to be ran half as many times as intended.
        let mut out = Some(Schedule::new(0));

        if !self.recording {
            // Start a new recording
            if self.current_step != self.previous_schedule.len() {
                panic!("possible nondeterminism: current execution ended earlier than expected (expected length {} but ended after {})", self.previous_schedule.len(), self.current_step);
            }

            self.previous_schedule.clear();
            out = self.scheduler.new_execution();
        }

        self.recording = !self.recording;
        self.current_step = 0;

        out
    }

    fn next_task(
        &mut self,
        runnable_tasks: &[&Task],
        current_task: Option<TaskId>,
        is_yielding: bool,
    ) -> Option<TaskId> {
        if self.recording {
            let choice = self.scheduler.next_task(runnable_tasks, current_task, is_yielding);
            let runnable_ids = runnable_tasks
                .iter()
                .map(|t| t.id())
                .collect::<SmallVec<[TaskId; DEFAULT_INLINE_TASKS]>>();
            self.previous_schedule
                .push(ScheduleRecord::Task(choice, runnable_ids, is_yielding));

            choice
        } else {
            if self.current_step >= self.previous_schedule.len() {
                panic!(
                    "possible nondeterminism: current execution should have ended after {} steps, whereas current step count is {}",
                    self.previous_schedule.len(),
                    self.current_step
                );
            }

            match &self.previous_schedule[self.current_step] {
                ScheduleRecord::Task(maybe_id, runnables, was_yielding) => {
                    let runnable_ids = runnable_tasks
                        .iter()
                        .map(|t| t.id())
                        .collect::<SmallVec<[TaskId; DEFAULT_INLINE_TASKS]>>();
                    if *runnables.as_slice() != *runnable_ids {
                        panic!("possible nondeterminism: set of runnable tasks is different than expected.\nExpected:\n{runnables:?}\nbut got:\n{runnable_ids:?}");
                    }

                    if *was_yielding != is_yielding {
                        panic!("possible nondeterminism: `next_task` was called with `is_yielding` equal to {was_yielding} in the original execution, and {is_yielding} in the current execution");
                    }

                    self.current_step += 1;

                    *maybe_id
                }
                ScheduleRecord::Random(_) => {
                    panic!("possible nondeterminism: next step was context switch, but recording expected random number generation")
                }
            }
        }
    }

    fn next_u64(&mut self) -> u64 {
        if self.recording {
            let next = self.scheduler.next_u64();
            self.previous_schedule.push(ScheduleRecord::Random(next));

            next
        } else {
            if self.current_step >= self.previous_schedule.len() {
                panic!(
                    "possible nondeterminism: current execution should have ended after {} steps, whereas current step count is {}",
                    self.previous_schedule.len(),
                    self.current_step
                );
            }

            match self.previous_schedule[self.current_step] {
                ScheduleRecord::Task(..) => panic!("possible nondeterminism: next step was random number generation, but recording expected context switch"),
                ScheduleRecord::Random(num) => {
                    self.current_step += 1;
                    num
                }
            }
        }
    }
}