Skip to main content

kaspa_utils/
sim.rs

1//! Module with structs for supporting discrete event simulation in virtual time.
2//! Inspired by python's simpy library.
3//!
4//! Users should define the message type `T` required for the simulation, derive `Process<T>` with
5//! various simulation actor logic and plug the processes into a `Simulation<T>` instance.
6
7use std::boxed::Box;
8use std::collections::{BinaryHeap, HashMap, HashSet};
9
10/// Internal structure representing a scheduled simulator event
11struct Event<T> {
12    timestamp: u64,
13    dest: u64,
14    msg: Option<T>,
15}
16
17impl<T> Event<T> {
18    pub fn new(timestamp: u64, dest: u64, msg: Option<T>) -> Self {
19        Self { timestamp, dest, msg }
20    }
21}
22
23impl<T> PartialEq for Event<T> {
24    fn eq(&self, other: &Self) -> bool {
25        self.timestamp == other.timestamp && self.dest == other.dest
26    }
27}
28
29impl<T> Eq for Event<T> {}
30
31impl<T> PartialOrd for Event<T> {
32    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
33        Some(self.cmp(other))
34    }
35}
36
37impl<T> Ord for Event<T> {
38    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
39        // Reversing so that min timestamp is scheduled first
40        other.timestamp.cmp(&self.timestamp).then_with(|| other.dest.cmp(&self.dest))
41    }
42}
43
44/// Process resumption trigger
45pub enum Resumption<T> {
46    Initial,
47    Scheduled,
48    Message(T),
49}
50
51/// Process suspension reason
52pub enum Suspension {
53    Timeout(u64),
54    Idle,
55    Halt, // Halt the simulation
56}
57
58/// A simulation process
59pub trait Process<T> {
60    fn resume(&mut self, resumption: Resumption<T>, env: &mut Environment<T>) -> Suspension;
61}
62
63pub type BoxedProcess<T> = Box<dyn Process<T>>;
64
65/// The simulation environment
66#[derive(Default)]
67pub struct Environment<T> {
68    now: u64,
69    broadcast_delay: u64,
70    event_queue: BinaryHeap<Event<T>>,
71    process_ids: HashSet<u64>,
72}
73
74impl<T: Clone> Environment<T> {
75    pub fn new(delay: u64) -> Self {
76        Self::with_start_time(delay, 0)
77    }
78
79    pub fn with_start_time(delay: u64, start_time: u64) -> Self {
80        Self { now: start_time, broadcast_delay: delay, event_queue: BinaryHeap::new(), process_ids: HashSet::new() }
81    }
82
83    pub fn now(&self) -> u64 {
84        self.now
85    }
86
87    pub fn send(&mut self, delay: u64, dest: u64, msg: T) {
88        self.event_queue.push(Event::new(self.now + delay, dest, Some(msg)))
89    }
90
91    pub fn timeout(&mut self, timeout: u64, dest: u64) {
92        self.event_queue.push(Event::new(self.now + timeout, dest, None))
93    }
94
95    pub fn broadcast(&mut self, _sender: u64, msg: T) {
96        for &id in self.process_ids.iter() {
97            self.event_queue.push(Event::new(self.now + self.broadcast_delay, id, Some(msg.clone())));
98        }
99    }
100
101    fn next_event(&mut self) -> Event<T> {
102        let event = self.event_queue.pop().unwrap();
103        self.now = event.timestamp;
104        event
105    }
106}
107
108/// The simulation manager
109#[derive(Default)]
110pub struct Simulation<T> {
111    env: Environment<T>,
112    processes: HashMap<u64, BoxedProcess<T>>,
113}
114
115impl<T: Clone> Simulation<T> {
116    pub fn new(delay: u64) -> Self {
117        Self { env: Environment::new(delay), processes: HashMap::new() }
118    }
119
120    pub fn with_start_time(delay: u64, start_time: u64) -> Self {
121        Self { env: Environment::with_start_time(delay, start_time), processes: HashMap::new() }
122    }
123
124    pub fn register(&mut self, id: u64, process: BoxedProcess<T>) {
125        self.processes.insert(id, process);
126        self.env.process_ids.insert(id);
127    }
128
129    pub fn step(&mut self) -> bool {
130        let event = self.env.next_event();
131        let process = self.processes.get_mut(&event.dest).unwrap();
132        let op = if let Some(msg) = event.msg { Resumption::Message(msg) } else { Resumption::Scheduled };
133        match process.resume(op, &mut self.env) {
134            Suspension::Timeout(timeout) => {
135                self.env.timeout(timeout, event.dest);
136                true
137            }
138            Suspension::Idle => true,
139            Suspension::Halt => false,
140        }
141    }
142
143    pub fn run(&mut self, until: u64) {
144        for (&id, process) in self.processes.iter_mut() {
145            match process.resume(Resumption::Initial, &mut self.env) {
146                Suspension::Timeout(timeout) => self.env.timeout(timeout, id),
147                Suspension::Idle => {}
148                Suspension::Halt => panic!("not expecting halt on startup"),
149            }
150        }
151
152        while self.step() {
153            if self.env.now() > until {
154                break;
155            }
156        }
157        self.processes.clear();
158    }
159}