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::capacity_allocator::CapacityAllocator;
34use crate::{connect_tx, port_rx, take_option};
35
36mod byte_store;
37mod object_store;
38
39pub use byte_store::ByteStore;
40pub use object_store::ObjectStore;
41
42type ObjectToCapacity<T> = fn(&T) -> usize;
43
44#[derive(Clone)]
45struct State<T>
46where
47    T: SimObject,
48{
49    entity: Rc<Entity>,
50    capacity: CapacityAllocator,
51    data: Rc<RefCell<VecDeque<T>>>,
52    error_on_overflow: Rc<RefCell<bool>>,
53    object_to_capacity: ObjectToCapacity<T>,
54}
55
56/// A component that can support a configurable number of capacity units.
57///
58/// Use [ObjectStore] to build an object-counted store and [ByteStore] to build
59/// a byte-counted store.
60#[derive(EntityGet, EntityDisplay)]
61pub struct Store<T>
62where
63    T: SimObject,
64{
65    entity: Rc<Entity>,
66    spawner: Spawner,
67    capacity: CapacityAllocator,
68    data: Rc<RefCell<VecDeque<T>>>,
69    error_on_overflow: Rc<RefCell<bool>>,
70    object_to_capacity: ObjectToCapacity<T>,
71    tx: RefCell<Option<OutPort<T>>>,
72    rx: RefCell<Option<InPort<T>>>,
73}
74
75impl<T> Store<T>
76where
77    T: SimObject,
78{
79    fn new(
80        engine: &Engine,
81        clock: &Clock,
82        entity: &Rc<Entity>,
83        aka: Option<&Aka>,
84        capacity: usize,
85        capacity_unit: &str,
86        object_to_capacity: ObjectToCapacity<T>,
87    ) -> Result<Self, SimError> {
88        if capacity == 0 {
89            return sim_error!("Unsupported Store with capacity of 0");
90        }
91        let capacity = CapacityAllocator::for_entity(entity, capacity, capacity_unit)?;
92        Ok(Self {
93            entity: entity.clone(),
94            spawner: engine.spawner(),
95            capacity,
96            data: Rc::new(RefCell::new(VecDeque::new())),
97            error_on_overflow: Rc::new(RefCell::new(false)),
98            object_to_capacity,
99            tx: RefCell::new(Some(OutPort::new_with_renames(entity, "tx", aka))),
100            rx: RefCell::new(Some(InPort::new_with_renames(
101                engine, clock, entity, "rx", aka,
102            ))),
103        })
104    }
105
106    pub fn connect_port_tx(&self, port_state: PortStateResult<T>) -> SimResult {
107        connect_tx!(self.tx, connect ; port_state)
108    }
109
110    pub fn port_rx(&self) -> PortStateResult<T> {
111        port_rx!(self.rx, state)
112    }
113
114    #[must_use]
115    pub fn capacity_used(&self) -> usize {
116        self.capacity.used()
117    }
118
119    pub fn set_error_on_overflow(&self) {
120        *self.error_on_overflow.borrow_mut() = true;
121    }
122
123    #[must_use]
124    pub fn get_level_change_event(&self) -> Repeated<usize> {
125        self.capacity.level_change_event()
126    }
127
128    fn state(&self) -> State<T> {
129        State {
130            entity: self.entity.clone(),
131            capacity: self.capacity.clone(),
132            data: self.data.clone(),
133            error_on_overflow: self.error_on_overflow.clone(),
134            object_to_capacity: self.object_to_capacity,
135        }
136    }
137}
138
139#[async_trait(?Send)]
140impl<T> Runnable for Store<T>
141where
142    T: SimObject,
143{
144    async fn run(&self) -> SimResult {
145        let rx = take_option!(self.rx);
146        let state = self.state();
147        self.spawner.spawn(async move { state.run_rx(rx).await });
148
149        let tx = take_option!(self.tx);
150        let state = self.state();
151        self.spawner.spawn(async move { state.run_tx(tx).await });
152        Ok(())
153    }
154}
155
156impl<T> State<T>
157where
158    T: SimObject,
159{
160    fn push_value(&self, value: T) -> SimResult {
161        let units = (self.object_to_capacity)(&value);
162        self.capacity.allocate(units)?;
163        self.entity.track_enter(value.id());
164        self.data.borrow_mut().push_back(value);
165        Ok(())
166    }
167
168    fn pop_value(&self) -> Result<T, SimError> {
169        let value = self.data.borrow_mut().pop_front().unwrap();
170        self.capacity.release((self.object_to_capacity)(&value));
171        self.entity.track_exit(value.id());
172        Ok(value)
173    }
174
175    async fn run_rx(&self, mut rx: InPort<T>) -> SimResult {
176        loop {
177            let value = rx.start_get()?.await;
178            let units = (self.object_to_capacity)(&value);
179            if !*self.error_on_overflow.borrow() {
180                self.capacity.wait_for_capacity(units).await?;
181            }
182            self.push_value(value)?;
183            rx.finish_get();
184        }
185    }
186
187    async fn run_tx(&self, mut tx: OutPort<T>) -> SimResult {
188        let level_change = self.capacity.level_change_event();
189        loop {
190            let level = self.data.borrow().len();
191            if level > 0 {
192                tx.try_put()?.await;
193                let value = self.pop_value()?;
194                tx.put(value)?.await;
195            } else {
196                level_change.listen().await;
197            }
198        }
199    }
200}
gwr_components/store/mod.rs

gwr_components/store/
mod.rs
