gwr_components/store/
mod.rs

1// Copyright (c) 2026 Graphcore Ltd. All rights reserved.
2
3//! A data store.
4//!
5//! [ObjectStore] builds an object-counted [Store], while [ByteStore] builds a
6//! byte-counted [Store] using
7//! [`total_bytes`](gwr_engine::traits::TotalBytes::total_bytes). The returned
8//! [Store] is the registered component in both cases.
9//!
10//! # Ports
11//!
12//! This component has the following ports:
13//!   - The `rx` port [InPort] which is used to put data into the store.
14//!   - The `tx` port [OutPort] which is used to get data out of the store.
15
16use std::cell::RefCell;
17use std::collections::VecDeque;
18use std::rc::Rc;
19
20use async_trait::async_trait;
21use gwr_engine::engine::Engine;
22use gwr_engine::events::repeated::Repeated;
23use gwr_engine::executor::Spawner;
24use gwr_engine::port::{InPort, OutPort, PortStateResult};
25use gwr_engine::sim_error;
26use gwr_engine::time::clock::Clock;
27use gwr_engine::traits::{Event, Runnable, SimObject};
28use gwr_engine::types::{SimError, SimResult};
29use gwr_model_builder::{EntityDisplay, EntityGet};
30use gwr_track::entity::Entity;
31use gwr_track::tracker::aka::Aka;
32
33use crate::{connect_tx, port_rx, take_option};
34
35mod byte_store;
36mod object_store;
37
38pub use byte_store::ByteStore;
39pub use object_store::ObjectStore;
40
41type ObjectToCapacity<T> = fn(&T) -> usize;
42
43struct State<T>
44where
45    T: SimObject,
46{
47    entity: Rc<Entity>,
48    capacity: usize,
49    capacity_unit: RefCell<String>,
50    used: RefCell<usize>,
51    data: RefCell<VecDeque<T>>,
52    error_on_overflow: RefCell<bool>,
53    level_change: Repeated<usize>,
54    object_to_capacity: ObjectToCapacity<T>,
55}
56
57impl<T> State<T>
58where
59    T: SimObject,
60{
61    fn new(entity: &Rc<Entity>, capacity: usize, object_to_capacity: ObjectToCapacity<T>) -> Self {
62        Self {
63            entity: entity.clone(),
64            capacity,
65            capacity_unit: RefCell::new("objects".to_string()),
66            used: RefCell::new(0),
67            data: RefCell::new(VecDeque::new()),
68            error_on_overflow: RefCell::new(false),
69            level_change: Repeated::new(usize::default()),
70            object_to_capacity,
71        }
72    }
73
74    fn has_capacity_for(&self, units: usize) -> bool {
75        units <= self.capacity - *self.used.borrow()
76    }
77
78    fn check_units_can_fit(&self, units: usize) -> SimResult {
79        if units > self.capacity {
80            let capacity_unit = self.capacity_unit.borrow();
81            return sim_error!(
82                "Cannot store {units} {capacity_unit} in {:?} with capacity {}",
83                self.entity.full_name(),
84                self.capacity
85            );
86        }
87        Ok(())
88    }
89
90    fn push_value(&self, value: T) -> SimResult {
91        let units = (self.object_to_capacity)(&value);
92        self.entity.track_enter(value.id());
93        if *self.error_on_overflow.borrow() {
94            if !self.has_capacity_for(units) {
95                return sim_error!("Overflow in {:?}", self.entity.full_name());
96            }
97        } else {
98            assert!(self.has_capacity_for(units));
99        }
100
101        self.data.borrow_mut().push_back(value);
102        *self.used.borrow_mut() += units;
103        self.level_change.notify_result(*self.used.borrow());
104        Ok(())
105    }
106
107    fn pop_value(&self) -> Result<T, SimError> {
108        let value = self.data.borrow_mut().pop_front().unwrap();
109        *self.used.borrow_mut() -= (self.object_to_capacity)(&value);
110        self.level_change.notify_result(*self.used.borrow());
111        self.entity.track_exit(value.id());
112        Ok(value)
113    }
114}
115
116/// A component that can support a configurable number of capacity units.
117///
118/// Use [ObjectStore] to build an object-counted store and [ByteStore] to build
119/// a byte-counted store.
120#[derive(EntityGet, EntityDisplay)]
121pub struct Store<T>
122where
123    T: SimObject,
124{
125    entity: Rc<Entity>,
126    spawner: Spawner,
127    state: Rc<State<T>>,
128    tx: RefCell<Option<OutPort<T>>>,
129    rx: RefCell<Option<InPort<T>>>,
130}
131
132impl<T> Store<T>
133where
134    T: SimObject,
135{
136    fn new(
137        engine: &Engine,
138        clock: &Clock,
139        entity: &Rc<Entity>,
140        aka: Option<&Aka>,
141        capacity: usize,
142        object_to_capacity: ObjectToCapacity<T>,
143    ) -> Result<Self, SimError> {
144        if capacity == 0 {
145            return sim_error!("Unsupported Store with capacity of 0");
146        }
147        Ok(Self {
148            entity: entity.clone(),
149            spawner: engine.spawner(),
150            state: Rc::new(State::new(entity, capacity, object_to_capacity)),
151            tx: RefCell::new(Some(OutPort::new_with_renames(entity, "tx", aka))),
152            rx: RefCell::new(Some(InPort::new_with_renames(
153                engine, clock, entity, "rx", aka,
154            ))),
155        })
156    }
157
158    pub fn connect_port_tx(&self, port_state: PortStateResult<T>) -> SimResult {
159        connect_tx!(self.tx, connect ; port_state)
160    }
161
162    pub fn port_rx(&self) -> PortStateResult<T> {
163        port_rx!(self.rx, state)
164    }
165
166    #[must_use]
167    pub fn capacity_used(&self) -> usize {
168        *self.state.used.borrow()
169    }
170
171    pub fn set_error_on_overflow(&self) {
172        *self.state.error_on_overflow.borrow_mut() = true;
173    }
174
175    pub fn set_capacity_unit(&self, capacity_unit: impl Into<String>) {
176        *self.state.capacity_unit.borrow_mut() = capacity_unit.into();
177    }
178
179    #[must_use]
180    pub fn get_level_change_event(&self) -> Repeated<usize> {
181        self.state.level_change.clone()
182    }
183}
184
185#[async_trait(?Send)]
186impl<T> Runnable for Store<T>
187where
188    T: SimObject,
189{
190    async fn run(&self) -> SimResult {
191        let rx = take_option!(self.rx);
192        let state = self.state.clone();
193        self.spawner.spawn(async move { run_rx(rx, state).await });
194
195        let tx = take_option!(self.tx);
196        let state = self.state.clone();
197        self.spawner.spawn(async move { run_tx(tx, state).await });
198        Ok(())
199    }
200}
201
202async fn run_rx<T>(mut rx: InPort<T>, state: Rc<State<T>>) -> SimResult
203where
204    T: SimObject,
205{
206    let level_change = state.level_change.clone();
207    loop {
208        let value = rx.start_get()?.await;
209        let units = (state.object_to_capacity)(&value);
210        state.check_units_can_fit(units)?;
211        while !state.has_capacity_for(units) && !*state.error_on_overflow.borrow() {
212            level_change.listen().await;
213        }
214        state.push_value(value)?;
215        rx.finish_get();
216    }
217}
218
219async fn run_tx<T>(mut tx: OutPort<T>, state: Rc<State<T>>) -> SimResult
220where
221    T: SimObject,
222{
223    let level_change = state.level_change.clone();
224    loop {
225        let level = state.data.borrow().len();
226        if level > 0 {
227            tx.try_put()?.await;
228            let value = state.pop_value()?;
229            tx.put(value)?.await;
230        } else {
231            level_change.listen().await;
232        }
233    }
234}
gwr_components/store/mod.rs

gwr_components/store/
mod.rs
