myrmic_sdk/host_functions/db/
table.rs

1//! Typed access to a `tb` table, plus lazy iterators over its rows and keys.
2
3use alloc::collections::BTreeMap;
4use alloc::string::String;
5use alloc::vec::Vec;
6
7use core::borrow::Borrow;
8use core::marker::PhantomData;
9
10use crate::{Codec, Postcard, Result, Sri};
11
12use crate::db::{Cursor, Scope, TbOrderBy, tb_append, tb_count, tb_delete, tb_get, tb_list};
13
14/// Encode a value into the raw entity-id bytes used to key a table.
15pub trait AsEid {
16    /// The raw entity-id bytes for this value.
17    fn as_eid(&self) -> Vec<u8>;
18}
19
20/// A key type that can be reconstructed from the raw entity-id bytes of a
21/// stored row, as well as encoded into them.
22pub trait TableKey: AsEid + Sized {
23    /// Reconstruct a key from its raw entity-id bytes, as stored in the table.
24    fn from_eid(eid: &[u8]) -> Self;
25}
26
27impl AsEid for str {
28    fn as_eid(&self) -> Vec<u8> {
29        self.as_bytes().to_vec()
30    }
31}
32
33impl AsEid for String {
34    fn as_eid(&self) -> Vec<u8> {
35        self.as_bytes().to_vec()
36    }
37}
38
39impl TableKey for String {
40    fn from_eid(eid: &[u8]) -> Self {
41        String::from_utf8_lossy(eid).into_owned()
42    }
43}
44
45impl AsEid for [u8] {
46    fn as_eid(&self) -> Vec<u8> {
47        self.to_vec()
48    }
49}
50
51impl AsEid for Vec<u8> {
52    fn as_eid(&self) -> Vec<u8> {
53        self.clone()
54    }
55}
56
57impl TableKey for Vec<u8> {
58    fn from_eid(eid: &[u8]) -> Self {
59        eid.to_vec()
60    }
61}
62
63impl AsEid for Sri {
64    fn as_eid(&self) -> Vec<u8> {
65        self.to_bytes().to_vec()
66    }
67}
68
69impl TableKey for Sri {
70    fn from_eid(eid: &[u8]) -> Self {
71        let mut bytes = [0u8; 16];
72        let n = eid.len().min(16);
73        bytes[..n].copy_from_slice(&eid[..n]);
74        Sri::from_bytes(bytes)
75    }
76}
77
78/// A typed handle to a `tb` table, storing serde values `V` keyed by `K`.
79///
80/// Construct one with [`Table::new`] (or [`Table::new_in`] for a non-default
81/// [`Scope`]) and operate on it via [`get`](Self::get),
82/// [`insert`](Self::insert), [`iter`](Self::iter), etc.
83pub struct Table<V, K = String, C = Postcard> {
84    scope: Scope,
85    name: &'static str,
86    _marker: PhantomData<(V, K, C)>,
87}
88
89impl<V, K, C> Table<V, K, C> {
90    /// A handle to the table `name` in the default private [`Scope`].
91    pub const fn new(name: &'static str) -> Self {
92        Self::new_in(name, Scope::private())
93    }
94
95    /// A handle to the table `name` in an explicit `scope`.
96    pub const fn new_in(name: &'static str, scope: Scope) -> Self {
97        Self {
98            scope,
99            name,
100            _marker: PhantomData,
101        }
102    }
103
104    fn scope(&self) -> Scope {
105        self.scope.clone()
106    }
107}
108
109impl<V, K, C> Table<V, K, C>
110where
111    V: serde::Serialize + serde::de::DeserializeOwned,
112    C: Codec,
113{
114    /// Fetches the value stored under `key`, or `None` if there is no such row.
115    pub fn get<Q>(&self, key: &Q) -> Result<Option<V>>
116    where
117        K: Borrow<Q>,
118        Q: AsEid + ?Sized,
119    {
120        let mut req = [0u8; 256];
121        let mut resp = alloc::vec![0u8; 4096];
122        tb_get(
123            self.scope(),
124            String::from(self.name),
125            key.as_eid(),
126            &mut req,
127            &mut resp,
128        )?
129        .map(|b| C::decode(&b))
130        .transpose()
131    }
132
133    /// Inserts `val` under a host-generated entity id.
134    pub fn insert(&self, val: &V) -> Result<()> {
135        self.insert_impl(None, C::encode(val)?)
136    }
137
138    /// Inserts `val` under `key`, replacing any existing value.
139    pub fn insert_with<Q>(&self, key: &Q, val: &V) -> Result<()>
140    where
141        K: Borrow<Q>,
142        Q: AsEid + ?Sized,
143    {
144        self.insert_impl(Some(key.as_eid()), C::encode(val)?)
145    }
146
147    /// Neither entry point reports the row id back, so both append: the host
148    /// batches the write into the handler's transaction instead of spending a
149    /// round trip to tell us an id we would drop.
150    fn insert_impl(&self, key: Option<Vec<u8>>, val: Vec<u8>) -> Result<()> {
151        let mut req = alloc::vec![0u8; val.len() + 1024];
152        tb_append(self.scope(), String::from(self.name), key, val, &mut req)
153            .map_err(|_| "Table::insert")?;
154        Ok(())
155    }
156
157    /// Deletes the row keyed `key`, if any.
158    pub fn delete<Q>(&self, key: &Q) -> Result<()>
159    where
160        K: Borrow<Q>,
161        Q: AsEid + ?Sized,
162    {
163        let mut buf = [0u8; 256];
164        tb_delete(
165            self.scope(),
166            String::from(self.name),
167            key.as_eid(),
168            &mut buf,
169        )
170        .map_err(|_| "Table::delete")?;
171        Ok(())
172    }
173
174    /// The number of rows in the table.
175    pub fn count(&self) -> Result<usize> {
176        let mut req = [0u8; 256];
177        let mut resp = [0u8; 64];
178        tb_count(self.scope(), String::from(self.name), &mut req, &mut resp)
179            .map_err(|_| "Table::count")
180    }
181
182    /// All values, collected in ascending key order.
183    pub fn list(&self) -> Result<Vec<V>> {
184        let mut out = Vec::new();
185        self.for_each(|v| out.push(v))?;
186        Ok(out)
187    }
188
189    /// Like [`Self::list`], but returns values in descending key order.
190    pub fn list_rev(&self) -> Result<Vec<V>> {
191        let mut out = Vec::new();
192        self.for_each_rev(|v| out.push(v))?;
193        Ok(out)
194    }
195
196    /// Visit every value, in ascending key order.
197    pub fn for_each<F: FnMut(V)>(&self, mut f: F) -> Result<()> {
198        let scope = self.scope();
199        for_each_ordered::<_, V, C>(&scope, self.name, TbOrderBy::KeyAsc, |_eid, val: V| f(val))
200    }
201
202    /// Visit every value, in descending key order.
203    pub fn for_each_rev<F: FnMut(V)>(&self, mut f: F) -> Result<()> {
204        let scope = self.scope();
205        for_each_ordered::<_, V, C>(&scope, self.name, TbOrderBy::KeyDesc, |_eid, val: V| f(val))
206    }
207
208    /// Internal iterating through undecoded `(eid, value)` byte rows.
209    fn iter_raw(&self, order: TbOrderBy) -> IterRaw {
210        IterRaw::new(self.scope(), String::from(self.name), order)
211    }
212
213    /// Lazily stream `(key, value)` pairs in ascending key order.
214    pub fn iter(&self) -> Iter<V, K, C> {
215        Iter {
216            raw: self.iter_raw(TbOrderBy::KeyAsc),
217            _marker: PhantomData,
218        }
219    }
220
221    /// Like [`Self::iter`], but streams `(key, value)` pairs in descending key order.
222    pub fn iter_rev(&self) -> Iter<V, K, C> {
223        Iter {
224            raw: self.iter_raw(TbOrderBy::KeyDesc),
225            _marker: PhantomData,
226        }
227    }
228
229    /// Lazily stream the table's keys in ascending order.
230    pub fn keys(&self) -> Keys<K> {
231        Keys {
232            raw: self.iter_raw(TbOrderBy::KeyAsc),
233            _marker: PhantomData,
234        }
235    }
236
237    /// Like [`Self::keys`], but streams the table's keys in descending order.
238    pub fn keys_rev(&self) -> Keys<K> {
239        Keys {
240            raw: self.iter_raw(TbOrderBy::KeyDesc),
241            _marker: PhantomData,
242        }
243    }
244
245    /// All keys in the table, collected.
246    pub fn ids(&self) -> Result<Vec<K>>
247    where
248        K: TableKey,
249    {
250        self.keys().collect()
251    }
252
253    /// Materialise the whole table into a `key -> value` map by collecting
254    /// [`Self::iter`].
255    /// Prefer the lazy [`Self::iter`]/[`Self::for_each`] for large tables.
256    pub fn to_map(&self) -> Result<BTreeMap<K, V>>
257    where
258        K: TableKey + Ord,
259    {
260        self.iter().collect()
261    }
262}
263
264/// Internal iterator that doesn't decode anything,
265/// _just_ deals with getting rows lazily out of the table.
266struct IterRaw {
267    scope: Scope,
268    table: String,
269    cursor: Option<Cursor>,
270    order: TbOrderBy,
271    done: bool,
272}
273
274impl IterRaw {
275    fn new(scope: Scope, table: String, order: TbOrderBy) -> Self {
276        Self {
277            scope,
278            table,
279            cursor: None,
280            order,
281            done: false,
282        }
283    }
284}
285
286impl Iterator for IterRaw {
287    type Item = Result<(Vec<u8>, Vec<u8>)>;
288
289    fn next(&mut self) -> Option<Self::Item> {
290        if self.done {
291            return None;
292        }
293        let mut req = [0u8; 256];
294        let mut resp = alloc::vec![0u8; 4096];
295        let row = tb_list(
296            self.scope.clone(),
297            self.table.clone(),
298            self.cursor.take(),
299            Some(1),
300            Some(self.order),
301            &mut req,
302            &mut resp,
303        );
304
305        let rows = match row {
306            Ok(rows) => rows,
307            Err(_) => {
308                self.done = true;
309                return Some(Err("tb_list"));
310            }
311        };
312        let Some((eid, val)) = rows.into_iter().next() else {
313            self.done = true;
314            return None;
315        };
316        self.cursor = Some(Cursor::After(eid.clone()));
317        Some(Ok((eid, val)))
318    }
319}
320
321/// Lazy iterator over a [`Table`] yielding decoded `(key, value)` pairs
322pub struct Iter<V, K, C = Postcard> {
323    raw: IterRaw,
324    _marker: PhantomData<(V, K, C)>,
325}
326
327impl<V, K, C> Iterator for Iter<V, K, C>
328where
329    V: serde::de::DeserializeOwned,
330    K: TableKey,
331    C: Codec,
332{
333    type Item = Result<(K, V)>;
334
335    fn next(&mut self) -> Option<Self::Item> {
336        let (eid, val) = match self.raw.next()? {
337            Ok(pair) => pair,
338            Err(e) => return Some(Err(e)),
339        };
340        Some(C::decode::<V>(&val).map(|v| (K::from_eid(&eid), v)))
341    }
342}
343
344/// Lazy iterator over a [`Table`]'s keys
345pub struct Keys<K = String> {
346    raw: IterRaw,
347    _marker: PhantomData<K>,
348}
349
350impl<K> Iterator for Keys<K>
351where
352    K: TableKey,
353{
354    type Item = Result<K>;
355
356    fn next(&mut self) -> Option<Self::Item> {
357        match self.raw.next()? {
358            Ok((eid, _val)) => Some(Ok(K::from_eid(&eid))),
359            Err(e) => Some(Err(e)),
360        }
361    }
362}
363
364/// Visits every row of `table` in `scope` in ascending key order, passing the
365/// raw entity id and the [`Postcard`]-decoded value; rows that fail to decode
366/// are skipped.
367pub fn for_each<F, T>(scope: &Scope, table: &str, f: F) -> Result<()>
368where
369    F: for<'a> FnMut(&'a [u8], T),
370    T: serde::de::DeserializeOwned,
371{
372    for_each_ordered::<_, T, Postcard>(scope, table, TbOrderBy::KeyAsc, f)
373}
374
375fn for_each_ordered<F, T, C>(scope: &Scope, table: &str, order: TbOrderBy, mut f: F) -> Result<()>
376where
377    F: for<'a> FnMut(&'a [u8], T),
378    T: serde::de::DeserializeOwned,
379    C: Codec,
380{
381    for_each_raw_ordered(scope, table, order, |eid, val| {
382        if let Ok(s) = C::decode::<T>(&val) {
383            f(eid, s);
384        }
385    })
386}
387
388/// Visits every `(eid, value)` row of `table` in `scope` in ascending key
389/// order, without decoding the values.
390pub fn for_each_raw<F>(scope: &Scope, table: &str, f: F) -> Result<()>
391where
392    F: for<'a> FnMut(&'a [u8], Vec<u8>),
393{
394    for_each_raw_ordered(scope, table, TbOrderBy::KeyAsc, f)
395}
396
397fn for_each_raw_ordered<F>(scope: &Scope, table: &str, order: TbOrderBy, mut f: F) -> Result<()>
398where
399    F: for<'a> FnMut(&'a [u8], Vec<u8>),
400{
401    for row in IterRaw::new(scope.clone(), String::from(table), order) {
402        let (eid, val) = row?;
403        f(eid.as_slice(), val);
404    }
405    Ok(())
406}
myrmic_sdk/host_functions/db/table.rs

myrmic_sdk/host_functions/db/
table.rs
