Skip to main content

elevator_core/dispatch/
scratch.rs

1//! Per-elevator scratch helper for custom dispatch strategies.
2//!
3//! Custom strategies that carry per-car state (idle counters, last-served
4//! ticks, learned weights, ...) typically reach for a
5//! `HashMap<EntityId, T>`. Each one then has to remember to drop the
6//! entry from `notify_removed`, or per-car state leaks every time an
7//! elevator is removed or reassigned.
8//!
9//! [`PrepareScratch`] is a typed wrapper around that pattern with
10//! batteries — `entry`, `get`, `get_mut`, `insert`, `remove`, `clear` —
11//! and a clear name so a strategy carrying multiple buckets reads as
12//! "the per-car scratch for X" rather than "another HashMap".
13
14use std::collections::HashMap;
15
16use crate::entity::EntityId;
17
18/// Per-elevator scratch storage, keyed by `EntityId`.
19///
20/// Custom strategies use `PrepareScratch<T>` to hold per-car state that
21/// is computed in `prepare_car` and read in `rank`. Drop entries from
22/// `notify_removed` so an elevator leaving the group doesn't leak.
23///
24/// # Example
25///
26/// ```
27/// use elevator_core::dispatch::{
28///     DispatchStrategy, PrepareScratch, RankContext, ElevatorGroup, DispatchManifest,
29/// };
30/// use elevator_core::entity::EntityId;
31/// use elevator_core::world::World;
32///
33/// #[derive(Default)]
34/// struct CarStats { idle_for: f64 }
35///
36/// #[derive(Default)]
37/// struct IdleAware {
38///     stats: PrepareScratch<CarStats>,
39/// }
40///
41/// impl DispatchStrategy for IdleAware {
42///     fn prepare_car(
43///         &mut self,
44///         car: EntityId,
45///         _car_position: f64,
46///         _group: &ElevatorGroup,
47///         _manifest: &DispatchManifest,
48///         _world: &World,
49///     ) {
50///         self.stats.entry(car).idle_for += 1.0;
51///     }
52///
53///     fn rank(&self, ctx: &RankContext<'_>) -> Option<f64> {
54///         let idle = self.stats.get(ctx.car).map_or(0.0, |s| s.idle_for);
55///         Some((ctx.car_position() - ctx.stop_position()).abs() - 0.01 * idle)
56///     }
57///
58///     fn notify_removed(&mut self, eid: EntityId) {
59///         self.stats.remove(eid);
60///     }
61/// }
62/// ```
63#[derive(Debug, Clone, Default)]
64pub struct PrepareScratch<T: Default> {
65    /// The underlying per-entity bucket.
66    inner: HashMap<EntityId, T>,
67}
68
69impl<T: Default> PrepareScratch<T> {
70    /// Empty scratch.
71    #[must_use]
72    pub fn new() -> Self {
73        Self {
74            inner: HashMap::new(),
75        }
76    }
77
78    /// Read-only access to `eid`'s scratch slot.
79    #[must_use]
80    pub fn get(&self, eid: EntityId) -> Option<&T> {
81        self.inner.get(&eid)
82    }
83
84    /// Mutable access to `eid`'s scratch slot, if it exists.
85    pub fn get_mut(&mut self, eid: EntityId) -> Option<&mut T> {
86        self.inner.get_mut(&eid)
87    }
88
89    /// Mutable access to `eid`'s scratch slot, inserting `T::default()`
90    /// if absent. Mirrors `HashMap::entry(...).or_default()` but takes
91    /// the typed `EntityId` directly and returns `&mut T` so callers can
92    /// chain field updates (`scratch.entry(car).field = …`).
93    ///
94    /// Intended for use inside `prepare_car`, where the framework
95    /// guarantees the elevator exists. Calling this from other
96    /// `&mut self` hooks (e.g. `pre_dispatch`) on an `eid` that has
97    /// not yet been through `prepare_car` will silently insert a
98    /// `T::default()` entry; prefer [`get`](Self::get) or
99    /// [`get_mut`](Self::get_mut) there.
100    pub fn entry(&mut self, eid: EntityId) -> &mut T {
101        self.inner.entry(eid).or_default()
102    }
103
104    /// Iterate `(EntityId, &T)` for every populated slot.
105    ///
106    /// Useful from `pre_dispatch` for fleet-wide scans (aging, decay,
107    /// tick-bound counters) without keeping a side-channel
108    /// `Vec<EntityId>` to track the populated set.
109    pub fn iter(&self) -> impl Iterator<Item = (EntityId, &T)> + '_ {
110        self.inner.iter().map(|(id, t)| (*id, t))
111    }
112
113    /// Iterate `(EntityId, &mut T)` for every populated slot.
114    pub fn iter_mut(&mut self) -> impl Iterator<Item = (EntityId, &mut T)> + '_ {
115        self.inner.iter_mut().map(|(id, t)| (*id, t))
116    }
117
118    /// Iterate `&T` over every populated slot.
119    pub fn values(&self) -> impl Iterator<Item = &T> + '_ {
120        self.inner.values()
121    }
122
123    /// Iterate `&mut T` over every populated slot.
124    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut T> + '_ {
125        self.inner.values_mut()
126    }
127
128    /// Replace the scratch value for `eid`, returning the previous value
129    /// if any.
130    pub fn insert(&mut self, eid: EntityId, value: T) -> Option<T> {
131        self.inner.insert(eid, value)
132    }
133
134    /// Drop the scratch entry for `eid`. Call from `notify_removed` so
135    /// the scratch doesn't outlive the elevator.
136    pub fn remove(&mut self, eid: EntityId) -> Option<T> {
137        self.inner.remove(&eid)
138    }
139
140    /// Drop every scratch entry.
141    pub fn clear(&mut self) {
142        self.inner.clear();
143    }
144
145    /// Number of scratch entries.
146    #[must_use]
147    pub fn len(&self) -> usize {
148        self.inner.len()
149    }
150
151    /// Whether the scratch is empty.
152    #[must_use]
153    pub fn is_empty(&self) -> bool {
154        self.inner.is_empty()
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use slotmap::KeyData;
162
163    fn fake_id(idx: u64) -> EntityId {
164        EntityId::from(KeyData::from_ffi(idx))
165    }
166
167    #[test]
168    fn entry_inserts_default_then_returns_mut() {
169        #[derive(Default, Debug, PartialEq)]
170        struct CarStats {
171            counter: u32,
172        }
173        let mut scratch: PrepareScratch<CarStats> = PrepareScratch::new();
174        let id = fake_id(0x4242_0000_0000_0001);
175        scratch.entry(id).counter = 7;
176        assert_eq!(scratch.get(id).map(|s| s.counter), Some(7));
177        scratch.entry(id).counter += 5;
178        assert_eq!(scratch.get(id).map(|s| s.counter), Some(12));
179    }
180
181    #[test]
182    fn remove_returns_old_value_and_drops_entry() {
183        let mut scratch: PrepareScratch<u32> = PrepareScratch::new();
184        let id = fake_id(0x1234_0000_0000_0001);
185        scratch.insert(id, 99);
186        assert_eq!(scratch.remove(id), Some(99));
187        assert!(scratch.get(id).is_none());
188        assert!(scratch.is_empty());
189    }
190}