Skip to main content

sqlite_containers/
sqlite_set.rs

1// SQLiteSet
2// This file is part of the 'SQLite-based containers for Rust' project (sqlite-containers)
3// SPDX-License-Identifier: Unlicense
4
5use crate::common::{Error, SizeT, check_constraint_violation};
6use crate::rusqlite::{Connection, Transaction};
7
8// ---------------------------------------------------------------------------
9// Statements
10// ---------------------------------------------------------------------------
11
12const SQL_CREATE_TBL: &str = "CREATE TABLE data (key TEXT PRIMARY KEY NOT NULL) WITHOUT ROWID;";
13const SQL_CREATE_NOC: &str = "CREATE TABLE data (key TEXT PRIMARY KEY NOT NULL COLLATE NOCASE) WITHOUT ROWID;";
14const SQL_COUNT_KEYS: &str = "SELECT COUNT(*) FROM data;";
15const SQL_INSERT_KEY: &str = "INSERT INTO data (key) VALUES (?1);";
16const SQL_EXISTS_KEY: &str = "SELECT 1 FROM data WHERE key = ? LIMIT 1;";
17const SQL_EXISTS_NOC: &str = "SELECT 1 FROM data WHERE key COLLATE NOCASE = ? LIMIT 1;";
18const SQL_QUERY_KEYS: &str = "SELECT key FROM data;";
19const SQL_DELETE_KEY: &str = "DELETE FROM data WHERE key = ?;";
20const SQL_DELETE_NOC: &str = "DELETE FROM data WHERE key COLLATE NOCASE = ?;";
21const SQL_DELETE_ALL: &str = "DELETE FROM data;";
22
23// ---------------------------------------------------------------------------
24// SQLiteSet
25// ---------------------------------------------------------------------------
26
27/// A [hash set](https://doc.rust-lang.org/std/collections/struct.HashSet.html) with [string](https://doc.rust-lang.org/beta/std/string/struct.String.html) keys, backed by an SQLite in-memory database.
28///
29/// By default, `SQLiteSet` treats its keys as case-sensitive, but a case-insensitive variant is available. Even when using the case-sensitive set variant, for some operations a dedicated "case-insensitive" version is provided.
30///
31/// <div class="warning">
32///
33/// **Important:** For bulk inserts or updates, it is **strongly recommended** to use an explicit [transaction](Self::transaction). Without one, SQLite executes each insert or update in its own transaction, which can significantly degrade performance.
34///
35/// </div>
36pub struct SQLiteSet {
37    connection: Connection,
38}
39
40impl SQLiteSet {
41    /// Creates a new, empty SQLite-backed hash set with case-sensitive keys.
42    #[inline]
43    pub fn new() -> Result<Self, Error> {
44        Ok(Self { connection: Self::initialize_connection(false)? })
45    }
46
47    /// Creates a new, empty SQLite-backed hash set with case-insensitive keys.
48    #[inline]
49    pub fn with_nocase() -> Result<Self, Error> {
50        Ok(Self { connection: Self::initialize_connection(true)? })
51    }
52
53    #[inline]
54    fn initialize_connection(no_case: bool) -> Result<Connection, Error> {
55        let connection = Connection::open_in_memory()?;
56        connection.pragma_update(None, "journal_mode", "OFF")?;
57        connection.pragma_update(None, "synchronous", "OFF")?;
58        connection.pragma_update(None, "temp_store", "MEMORY")?;
59        if !no_case {
60            connection.execute(SQL_CREATE_TBL, [])?;
61        } else {
62            connection.execute(SQL_CREATE_NOC, [])?;
63        }
64        Ok(connection)
65    }
66
67    /// Starts a new SQLite transaction for this set.
68    ///
69    /// Please note that using an explicit SQLite transaction allows for much more efficient bulk inserts &#x1F680;
70    ///
71    /// Returns the new [`SQLiteSetTransaction`] instance.
72    pub fn transaction(&mut self) -> Result<SQLiteSetTransaction<'_>, Error> {
73        SQLiteSetTransaction::from(&mut self.connection)
74    }
75
76    /// Inserts the given key into the set.
77    ///
78    /// Returns `true`, if the set did not already contain the key; otherwise returns `false`.
79    #[inline]
80    pub fn insert(&mut self, key: &str) -> Result<bool, Error> {
81        let mut insert = self.connection.prepare_cached(SQL_INSERT_KEY)?;
82        match insert.execute([key]) {
83            Ok(_) => Ok(true),
84            Err(error) => check_constraint_violation(error),
85        }
86    }
87
88    /// Checks whether the set contains the specified key.
89    ///
90    /// For case-sensitive sets, the check is case-sensitive; for case-insensitive sets, the check is case-insensitive.
91    ///
92    /// Returns `true`, if the set contains the key; otherwise returns `false`.
93    #[inline]
94    pub fn contains(&self, key: &str) -> Result<bool, Error> {
95        let mut contains = self.connection.prepare_cached(SQL_EXISTS_KEY)?;
96        Ok(contains.exists([key])?)
97    }
98
99    /// This is the "case-insensitive" version of the [`contains()`](Self::contains) function.
100    ///
101    /// The check is *always* performed case-insensitive.
102    ///
103    /// Returns `true`, if the set contains the key; otherwise returns `false`.
104    #[inline]
105    pub fn contains_nocase(&self, key: &str) -> Result<bool, Error> {
106        let mut contains = self.connection.prepare_cached(SQL_EXISTS_NOC)?;
107        Ok(contains.exists([key])?)
108    }
109
110    /// Removes the specified key from the set, if present.
111    ///
112    /// For case-sensitive sets, the key is treated as case-sensitive; for case-insensitive sets, it is treated as case-insensitive.
113    ///
114    /// Returns `true`, if the set contained the key; otherwise returns `false`.
115    #[inline]
116    pub fn remove(&mut self, key: &str) -> Result<bool, Error> {
117        let mut contains = self.connection.prepare_cached(SQL_DELETE_KEY)?;
118        Ok(contains.execute([key])? != 0)
119    }
120
121    /// This is the "case-insensitive" version of the [`remove()`](Self::remove) function.
122    ///
123    /// The key is *always* treated as case-insensitive.
124    ///
125    /// Returns `true`, if the set contained the key; otherwise returns `false`.
126    #[inline]
127    pub fn remove_nocase(&mut self, key: &str) -> Result<bool, Error> {
128        let mut contains = self.connection.prepare_cached(SQL_DELETE_NOC)?;
129        Ok(contains.execute([key])? != 0)
130    }
131
132    /// Invokes the given `callback` function for each key that is currently contained in the set.
133    ///
134    /// This function does **not** guarantee a specific iteration order.
135    #[inline]
136    pub fn for_each<F>(&self, mut callback: F) -> Result<(), Error>
137    where
138        F: FnMut(&str),
139    {
140        let mut iter = self.connection.prepare_cached(SQL_QUERY_KEYS)?;
141        let mut result = iter.query([])?;
142        while let Some(current_item) = result.next()? {
143            let key: String = current_item.get(0)?;
144            callback(&key);
145        }
146        Ok(())
147    }
148
149    /// Searches the set for the first key that satisfies the given `predicate`.
150    ///
151    /// Returns the first key that satisfies the given predicate, or `None` if **no** key satisfies the predicate or the set is empty.
152    ///
153    /// This function does **not** guarantee a specific iteration order.
154    ///
155    /// Also, the predicate is **not** always tested on *all* keys, because the function returns at the first match.
156    #[inline]
157    pub fn find<P>(&self, predicate: P) -> Result<Option<String>, Error>
158    where
159        P: Fn(&str) -> bool,
160    {
161        let mut iter = self.connection.prepare_cached(SQL_QUERY_KEYS)?;
162        let mut result = iter.query([])?;
163        while let Some(current_item) = result.next()? {
164            let key: String = current_item.get(0)?;
165            if predicate(&key) {
166                return Ok(Some(key));
167            }
168        }
169        Ok(None)
170    }
171
172    /// Returns the number of unique keys in the set.
173    #[inline]
174    pub fn len(&self) -> Result<SizeT, Error> {
175        let mut query_count = self.connection.prepare_cached(SQL_COUNT_KEYS)?;
176        let count: i64 = query_count.query_one([], |row| row.get(0))?;
177        Ok(count.try_into().unwrap_or_default())
178    }
179
180    /// Returns `true` if the set contains **no** keys; otherwise returns `false`.
181    #[inline]
182    pub fn is_empty(&self) -> Result<bool, Error> {
183        Ok(self.len()? == 0)
184    }
185
186    /// Removes *all* keys from the set.
187    #[inline]
188    pub fn clear(&mut self) -> Result<(), Error> {
189        let mut clear = self.connection.prepare_cached(SQL_DELETE_ALL)?;
190        clear.execute([])?;
191        Ok(())
192    }
193}
194
195impl Default for SQLiteSet {
196    /// Returns a new, empty set, as created by the [`SQLiteSet::new()`] function.
197    ///
198    /// # Panics
199    ///
200    /// Panics if a new `SQLiteSet` instance could **not** be created, e.g., because of an SQLite error.
201    #[inline]
202    fn default() -> Self {
203        Self::new().expect("Failed to create SQLiteSet instance!")
204    }
205}
206
207// ---------------------------------------------------------------------------
208// SQLiteSet Transaction
209// ---------------------------------------------------------------------------
210
211/// Represents an active SQLite transaction for a [`SQLiteSet`].
212///
213/// Most functions provided by this struct mirror the corresponding functions of the `SQLiteSet` struct.
214///
215/// The transaction is committed when the `SQLiteSetTransaction` is dropped.
216pub struct SQLiteSetTransaction<'a> {
217    transaction: Transaction<'a>,
218}
219
220impl<'a> SQLiteSetTransaction<'a> {
221    #[inline]
222    fn from(connection: &'a mut Connection) -> Result<Self, Error> {
223        let mut transaction = connection.transaction()?;
224        transaction.set_drop_behavior(rusqlite::DropBehavior::Commit);
225        Ok(Self { transaction })
226    }
227
228    /// Drops the `SQLiteSetTransaction`, thereby committing the SQLite transaction.
229    #[inline]
230    pub fn commit(self) {}
231
232    /// This function is equivalent to [`SQLiteSet::insert()`].
233    #[inline]
234    pub fn insert(&mut self, key: &str) -> Result<bool, Error> {
235        let mut insert = self.transaction.prepare_cached(SQL_INSERT_KEY)?;
236        match insert.execute([key]) {
237            Ok(_) => Ok(true),
238            Err(error) => check_constraint_violation(error),
239        }
240    }
241
242    /// This function is equivalent to [`SQLiteSet::contains()`].
243    #[inline]
244    pub fn contains(&self, key: &str) -> Result<bool, Error> {
245        let mut contains = self.transaction.prepare_cached(SQL_EXISTS_KEY)?;
246        Ok(contains.exists([key])?)
247    }
248
249    /// This function is equivalent to [`SQLiteSet::contains_nocase()`].
250    #[inline]
251    pub fn contains_nocase(&self, key: &str) -> Result<bool, Error> {
252        let mut contains = self.transaction.prepare_cached(SQL_EXISTS_NOC)?;
253        Ok(contains.exists([key])?)
254    }
255
256    /// This function is equivalent to [`SQLiteSet::remove()`].
257    #[inline]
258    pub fn remove(&mut self, key: &str) -> Result<bool, Error> {
259        let mut contains = self.transaction.prepare_cached(SQL_DELETE_KEY)?;
260        Ok(contains.execute([key])? != 0)
261    }
262
263    /// This function is equivalent to [`SQLiteSet::remove_nocase()`].
264    #[inline]
265    pub fn remove_nocase(&mut self, key: &str) -> Result<bool, Error> {
266        let mut contains = self.transaction.prepare_cached(SQL_DELETE_NOC)?;
267        Ok(contains.execute([key])? != 0)
268    }
269
270    /// This function is equivalent to [`SQLiteSet::for_each()`].
271    #[inline]
272    pub fn for_each<F>(&self, mut callback: F) -> Result<(), Error>
273    where
274        F: FnMut(&str),
275    {
276        let mut iter = self.transaction.prepare_cached(SQL_QUERY_KEYS)?;
277        let mut result = iter.query([])?;
278        while let Some(current_item) = result.next()? {
279            let key: String = current_item.get(0)?;
280            callback(&key);
281        }
282        Ok(())
283    }
284
285    /// This function is equivalent to [`SQLiteSet::find()`].
286    #[inline]
287    pub fn find<P>(&self, predicate: P) -> Result<Option<String>, Error>
288    where
289        P: Fn(&str) -> bool,
290    {
291        let mut iter = self.transaction.prepare_cached(SQL_QUERY_KEYS)?;
292        let mut result = iter.query([])?;
293        while let Some(current_item) = result.next()? {
294            let key: String = current_item.get(0)?;
295            if predicate(&key) {
296                return Ok(Some(key));
297            }
298        }
299        Ok(None)
300    }
301
302    /// This function is equivalent to [`SQLiteSet::len()`].
303    #[inline]
304    pub fn len(&self) -> Result<SizeT, Error> {
305        let mut query_count = self.transaction.prepare_cached(SQL_COUNT_KEYS)?;
306        let count: i64 = query_count.query_one([], |row| row.get(0))?;
307        Ok(count.try_into().unwrap_or_default())
308    }
309
310    /// This function is equivalent to [`SQLiteSet::is_empty()`].
311    #[inline]
312    pub fn is_empty(&self) -> Result<bool, Error> {
313        Ok(self.len()? == 0)
314    }
315
316    /// This function is equivalent to [`SQLiteSet::clear()`].
317    #[inline]
318    pub fn clear(&mut self) -> Result<(), Error> {
319        let mut clear = self.transaction.prepare_cached(SQL_DELETE_ALL)?;
320        clear.execute([])?;
321        Ok(())
322    }
323}