From 5c299a7590767c52df82535c4cd1712119506ba1 Mon Sep 17 00:00:00 2001
From: sgdxbc <sung@comp.nus.edu.sg>
Date: Thu, 22 Aug 2024 19:34:42 +0800
Subject: [PATCH] Checkpoint simulated states

---
 src/model/search/state.rs |   8 +--
 src/model/simulate.rs     | 119 +++++++++++++++++++++++++++++++++++++-
 src/pbft/tests.rs         |   2 +
 3 files changed, 122 insertions(+), 7 deletions(-)

diff --git a/src/model/search/state.rs b/src/model/search/state.rs
index 1c78d2a5..e35f9df8 100644
--- a/src/model/search/state.rs
+++ b/src/model/search/state.rs
@@ -11,16 +11,18 @@ use crate::{
 #[derive_where(Default)]
 pub struct Schedule<M> {
     envelops: Vec<TimerEnvelop<M>>,
-    count: u32,
+    count: TimerId,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 struct TimerEnvelop<M> {
-    id: u32,
+    id: TimerId,
     period: Duration,
     event: M,
 }
 
+pub type TimerId = u32;
+
 impl<M> Schedule<M> {
     pub fn new() -> Self {
         Self::default()
@@ -54,8 +56,6 @@ impl<M: Into<N>, N> ScheduleEvent<M> for Schedule<N> {
     }
 }
 
-pub type TimerId = u32;
-
 impl<M> Schedule<M> {
     fn remove(&mut self, id: u32) -> anyhow::Result<TimerEnvelop<M>> {
         let Some(pos) = self.envelops.iter().position(|envelop| envelop.id == id) else {
diff --git a/src/model/simulate.rs b/src/model/simulate.rs
index 1aef8c59..3ee78caa 100644
--- a/src/model/simulate.rs
+++ b/src/model/simulate.rs
@@ -1,17 +1,130 @@
+use std::{
+    collections::{BTreeSet, HashMap},
+    time::Duration,
+};
+
 use derive_more::{Display, Error};
+use derive_where::derive_where;
+
+use crate::{
+    event::{ActiveTimer, ScheduleEvent, SendEvent},
+    net::events::Cast,
+};
 
 #[derive(Debug, Display, Error)]
 pub struct ProgressExhausted;
 
 pub trait State {
-    fn step(&mut self) -> anyhow::Result<()>;
+    type Event;
 
-    fn progress(&mut self) -> anyhow::Result<()> {
+    fn step(&mut self, temporal: &mut Temporal<Self::Event>) -> anyhow::Result<()>;
+
+    fn progress(&mut self, temporal: &mut Temporal<Self::Event>) -> anyhow::Result<()> {
         loop {
-            match self.step() {
+            match self.step(temporal) {
                 Err(err) if err.is::<ProgressExhausted>() => return Ok(()),
                 result => result?,
             }
         }
     }
 }
+
+pub type TimerId = u32;
+
+#[derive_where(Default)]
+pub struct Temporal<M> {
+    now: Duration,
+    count: TimerId,
+    timeline: BTreeSet<(Duration, TimerId)>,
+    timers: HashMap<TimerId, TimerEnvelop<M>>,
+}
+
+struct TimerEnvelop<M> {
+    event: M,
+    period: Duration,
+    at: Duration,
+}
+
+impl<M> ScheduleEvent<M> for Temporal<M> {
+    fn set(&mut self, period: Duration, event: M) -> anyhow::Result<ActiveTimer>
+    where
+        M: Send + Clone + 'static,
+    {
+        self.count += 1;
+        let id = self.count;
+        let at = self.now + period;
+        let envelop = TimerEnvelop { event, period, at };
+        let replaced = self.timers.insert(id, envelop);
+        assert!(replaced.is_none());
+        let inserted = self.timeline.insert((at, id));
+        assert!(inserted);
+        Ok(ActiveTimer(id))
+    }
+
+    fn set_internal(
+        &mut self,
+        _: Duration,
+        _: impl FnMut() -> M + Send + 'static,
+    ) -> anyhow::Result<ActiveTimer> {
+        anyhow::bail!("unimplemented")
+    }
+
+    fn unset(&mut self, ActiveTimer(id): ActiveTimer) -> anyhow::Result<()> {
+        let Some(envelop) = self.timers.remove(&id) else {
+            anyhow::bail!("missing timer envelop")
+        };
+        let removed = self.timeline.remove(&(envelop.at, id));
+        assert!(removed);
+        Ok(())
+    }
+}
+
+impl<M> Temporal<M> {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn pop(&mut self) -> anyhow::Result<M>
+    where
+        M: Clone,
+    {
+        let Some((at, id)) = self.timeline.pop_first() else {
+            anyhow::bail!(ProgressExhausted)
+        };
+        self.now = at;
+        let Some(envelop) = self.timers.get_mut(&id) else {
+            unreachable!()
+        };
+        envelop.at = self.now + envelop.period;
+        let event = envelop.event.clone();
+        let inserted = self.timeline.insert((envelop.at, id));
+        assert!(inserted);
+        Ok(event)
+    }
+}
+
+#[derive(Debug)]
+#[derive_where(Default)]
+pub struct NetworkState<A, M> {
+    messages: Vec<(A, M)>,
+}
+
+impl<A, M> NetworkState<A, M> {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    #[cfg(test)]
+    pub fn choose(&mut self, u: &mut arbtest::arbitrary::Unstructured) -> anyhow::Result<(A, M)> {
+        anyhow::ensure!(!self.messages.is_empty(), ProgressExhausted);
+        let i = u.choose_index(self.messages.len()).unwrap();
+        Ok(self.messages.swap_remove(i))
+    }
+}
+
+impl<A, M: Into<N>, N> SendEvent<Cast<A, M>> for NetworkState<A, N> {
+    fn send(&mut self, Cast(remote, message): Cast<A, M>) -> anyhow::Result<()> {
+        self.messages.push((remote, message.into()));
+        Ok(())
+    }
+}
diff --git a/src/pbft/tests.rs b/src/pbft/tests.rs
index f7fc9c30..e6b3a750 100644
--- a/src/pbft/tests.rs
+++ b/src/pbft/tests.rs
@@ -380,3 +380,5 @@ mod search {
         }
     }
 }
+
+mod simulate {}
\ No newline at end of file