cuprate_database/backend/heed/
database.rs

1//! Implementation of `trait Database` for `heed`.
2
3//---------------------------------------------------------------------------------------------------- Import
4use std::cell::RefCell;
5
6use crate::{
7    backend::heed::types::HeedDb,
8    database::{DatabaseIter, DatabaseRo, DatabaseRw},
9    error::{DbResult, RuntimeError},
10    table::Table,
11};
12
13//---------------------------------------------------------------------------------------------------- Heed Database Wrappers
14// Q. Why does `HeedTableR{o,w}` exist?
15// A. These wrapper types combine `heed`'s database/table
16// types with its transaction types. It exists to match
17// `redb`, which has this behavior built-in.
18//
19// `redb` forces us to abstract read/write semantics
20// at the _opened table_ level, so, we must match that in `heed`,
21// which abstracts it at the transaction level.
22//
23// We must also maintain the ability for
24// write operations to also read, aka, `Rw`.
25
26/// An opened read-only database associated with a transaction.
27///
28/// Matches `redb::ReadOnlyTable`.
29pub(super) struct HeedTableRo<'tx, T: Table> {
30    /// An already opened database table.
31    pub(super) db: HeedDb<T::Key, T::Value>,
32    /// The associated read-only transaction that opened this table.
33    pub(super) tx_ro: &'tx heed::RoTxn<'tx>,
34}
35
36/// An opened read/write database associated with a transaction.
37///
38/// Matches `redb::Table` (read & write).
39pub(super) struct HeedTableRw<'env, 'tx, T: Table> {
40    /// An already opened database table.
41    pub(super) db: HeedDb<T::Key, T::Value>,
42    /// The associated read/write transaction that opened this table.
43    pub(super) tx_rw: &'tx RefCell<heed::RwTxn<'env>>,
44}
45
46//---------------------------------------------------------------------------------------------------- Shared functions
47// FIXME: we cannot just deref `HeedTableRw -> HeedTableRo` and
48// call the functions since the database is held by value, so
49// just use these generic functions that both can call instead.
50
51/// Shared [`DatabaseRo::get()`].
52#[inline]
53fn get<T: Table>(
54    db: &HeedDb<T::Key, T::Value>,
55    tx_ro: &heed::RoTxn<'_>,
56    key: &T::Key,
57) -> DbResult<T::Value> {
58    db.get(tx_ro, key)?.ok_or(RuntimeError::KeyNotFound)
59}
60
61/// Shared [`DatabaseRo::len()`].
62#[inline]
63fn len<T: Table>(db: &HeedDb<T::Key, T::Value>, tx_ro: &heed::RoTxn<'_>) -> DbResult<u64> {
64    Ok(db.len(tx_ro)?)
65}
66
67/// Shared [`DatabaseRo::first()`].
68#[inline]
69fn first<T: Table>(
70    db: &HeedDb<T::Key, T::Value>,
71    tx_ro: &heed::RoTxn<'_>,
72) -> DbResult<(T::Key, T::Value)> {
73    db.first(tx_ro)?.ok_or(RuntimeError::KeyNotFound)
74}
75
76/// Shared [`DatabaseRo::last()`].
77#[inline]
78fn last<T: Table>(
79    db: &HeedDb<T::Key, T::Value>,
80    tx_ro: &heed::RoTxn<'_>,
81) -> DbResult<(T::Key, T::Value)> {
82    db.last(tx_ro)?.ok_or(RuntimeError::KeyNotFound)
83}
84
85/// Shared [`DatabaseRo::is_empty()`].
86#[inline]
87fn is_empty<T: Table>(db: &HeedDb<T::Key, T::Value>, tx_ro: &heed::RoTxn<'_>) -> DbResult<bool> {
88    Ok(db.is_empty(tx_ro)?)
89}
90
91//---------------------------------------------------------------------------------------------------- DatabaseIter Impl
92impl<T: Table> DatabaseIter<T> for HeedTableRo<'_, T> {
93    /*
94    #[inline]
95    fn get_range<'a, Range>(
96        &'a self,
97        range: Range,
98    ) -> DbResult<impl Iterator<Item = DbResult<T::Value>> + 'a>
99    where
100        Range: RangeBounds<T::Key> + 'a,
101    {
102        Ok(self.db.range(self.tx_ro, &range)?.map(|res| Ok(res?.1)))
103    }
104
105     */
106
107    #[inline]
108    fn iter(&self) -> DbResult<impl Iterator<Item = DbResult<(T::Key, T::Value)>> + '_> {
109        Ok(self.db.iter(self.tx_ro)?.map(|res| Ok(res?)))
110    }
111
112    #[inline]
113    fn keys(&self) -> DbResult<impl Iterator<Item = DbResult<T::Key>> + '_> {
114        Ok(self.db.iter(self.tx_ro)?.map(|res| Ok(res?.0)))
115    }
116
117    #[inline]
118    fn values(&self) -> DbResult<impl Iterator<Item = DbResult<T::Value>> + '_> {
119        Ok(self.db.iter(self.tx_ro)?.map(|res| Ok(res?.1)))
120    }
121}
122
123//---------------------------------------------------------------------------------------------------- DatabaseRo Impl
124// SAFETY: `HeedTableRo: !Send` as it holds a reference to `heed::RoTxn: Send + !Sync`.
125unsafe impl<T: Table> DatabaseRo<T> for HeedTableRo<'_, T> {
126    #[inline]
127    fn get(&self, key: &T::Key) -> DbResult<T::Value> {
128        get::<T>(&self.db, self.tx_ro, key)
129    }
130
131    #[inline]
132    fn len(&self) -> DbResult<u64> {
133        len::<T>(&self.db, self.tx_ro)
134    }
135
136    #[inline]
137    fn first(&self) -> DbResult<(T::Key, T::Value)> {
138        first::<T>(&self.db, self.tx_ro)
139    }
140
141    #[inline]
142    fn last(&self) -> DbResult<(T::Key, T::Value)> {
143        last::<T>(&self.db, self.tx_ro)
144    }
145
146    #[inline]
147    fn is_empty(&self) -> DbResult<bool> {
148        is_empty::<T>(&self.db, self.tx_ro)
149    }
150}
151
152//---------------------------------------------------------------------------------------------------- DatabaseRw Impl
153// SAFETY: The `Send` bound only applies to `HeedTableRo`.
154// `HeedTableRw`'s write transaction is `!Send`.
155unsafe impl<T: Table> DatabaseRo<T> for HeedTableRw<'_, '_, T> {
156    #[inline]
157    fn get(&self, key: &T::Key) -> DbResult<T::Value> {
158        get::<T>(&self.db, &self.tx_rw.borrow(), key)
159    }
160
161    #[inline]
162    fn len(&self) -> DbResult<u64> {
163        len::<T>(&self.db, &self.tx_rw.borrow())
164    }
165
166    #[inline]
167    fn first(&self) -> DbResult<(T::Key, T::Value)> {
168        first::<T>(&self.db, &self.tx_rw.borrow())
169    }
170
171    #[inline]
172    fn last(&self) -> DbResult<(T::Key, T::Value)> {
173        last::<T>(&self.db, &self.tx_rw.borrow())
174    }
175
176    #[inline]
177    fn is_empty(&self) -> DbResult<bool> {
178        is_empty::<T>(&self.db, &self.tx_rw.borrow())
179    }
180}
181
182impl<T: Table> DatabaseRw<T> for HeedTableRw<'_, '_, T> {
183    #[inline]
184    fn put(&mut self, key: &T::Key, value: &T::Value) -> DbResult<()> {
185        Ok(self.db.put(&mut self.tx_rw.borrow_mut(), key, value)?)
186    }
187
188    #[inline]
189    fn delete(&mut self, key: &T::Key) -> DbResult<()> {
190        self.db.delete(&mut self.tx_rw.borrow_mut(), key)?;
191        Ok(())
192    }
193
194    #[inline]
195    fn take(&mut self, key: &T::Key) -> DbResult<T::Value> {
196        // LMDB/heed does not return the value on deletion.
197        // So, fetch it first - then delete.
198        let value = get::<T>(&self.db, &self.tx_rw.borrow(), key)?;
199        match self.db.delete(&mut self.tx_rw.borrow_mut(), key) {
200            Ok(true) => Ok(value),
201            Err(e) => Err(e.into()),
202            // We just `get()`'ed the value - it is
203            // incorrect for it to suddenly not exist.
204            Ok(false) => unreachable!(),
205        }
206    }
207
208    #[inline]
209    fn pop_first(&mut self) -> DbResult<(T::Key, T::Value)> {
210        let tx_rw = &mut self.tx_rw.borrow_mut();
211
212        // Get the value first...
213        let Some((key, value)) = self.db.first(tx_rw)? else {
214            return Err(RuntimeError::KeyNotFound);
215        };
216
217        // ...then remove it.
218        match self.db.delete(tx_rw, &key) {
219            Ok(true) => Ok((key, value)),
220            Err(e) => Err(e.into()),
221            // We just `get()`'ed the value - it is
222            // incorrect for it to suddenly not exist.
223            Ok(false) => unreachable!(),
224        }
225    }
226
227    #[inline]
228    fn pop_last(&mut self) -> DbResult<(T::Key, T::Value)> {
229        let tx_rw = &mut self.tx_rw.borrow_mut();
230
231        // Get the value first...
232        let Some((key, value)) = self.db.last(tx_rw)? else {
233            return Err(RuntimeError::KeyNotFound);
234        };
235
236        // ...then remove it.
237        match self.db.delete(tx_rw, &key) {
238            Ok(true) => Ok((key, value)),
239            Err(e) => Err(e.into()),
240            // We just `get()`'ed the value - it is
241            // incorrect for it to suddenly not exist.
242            Ok(false) => unreachable!(),
243        }
244    }
245}
246
247//---------------------------------------------------------------------------------------------------- Tests
248#[cfg(test)]
249mod test {
250    // use super::*;
251}
cuprate_database/backend/heed/database.rs

cuprate_database/backend/heed/
database.rs
