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::{
6    CharType, Error, Size, SizeExt, Type,
7    common::{SQL_TYPE_TEXT, SQL_TYPES, SqlType, check_constraint_violation},
8    rusqlite::{Connection, ToSql, Transaction},
9};
10use cfg_if::cfg_if;
11use std::{borrow::Borrow, collections::BTreeMap, marker::PhantomData, sync::LazyLock};
12
13// ---------------------------------------------------------------------------
14// Statements
15// ---------------------------------------------------------------------------
16
17static SQL_CREATE_TBL: LazyLock<BTreeMap<SqlType, String>> = LazyLock::new(|| {
18    let mut map = BTreeMap::<SqlType, String>::new();
19    for sql_type in SQL_TYPES {
20        cfg_if! {
21            if #[cfg(feature = "with_rowid")] {
22                map.insert(sql_type, format!("CREATE TABLE data (key {sql_type} PRIMARY KEY NOT NULL);"));
23            } else {
24                map.insert(sql_type, format!("CREATE TABLE data (key {sql_type} PRIMARY KEY NOT NULL) WITHOUT ROWID;"));
25            }
26        }
27    }
28    map
29});
30
31cfg_if! {
32    if #[cfg(feature = "with_rowid")] {
33        const SQL_CREATE_NOC: &str = "CREATE TABLE data (key TEXT PRIMARY KEY NOT NULL COLLATE NOCASE);";
34    } else {
35        const SQL_CREATE_NOC: &str = "CREATE TABLE data (key TEXT PRIMARY KEY NOT NULL COLLATE NOCASE) WITHOUT ROWID;";
36    }
37}
38
39const SQL_COUNT_KEYS: &str = "SELECT COUNT(*) FROM data;";
40const SQL_INSERT_KEY: &str = "INSERT INTO data (key) VALUES (?1);";
41const SQL_EXISTS_KEY: &str = "SELECT 1 FROM data WHERE key = ?1 LIMIT 1;";
42const SQL_EXISTS_NOC: &str = "SELECT 1 FROM data WHERE key COLLATE NOCASE = ? LIMIT 1;";
43const SQL_QUERY_KEYS: &str = "SELECT key FROM data;";
44const SQL_WHERE_KEYS: &str = "SELECT key FROM data WHERE key LIKE ?1 ESCAPE '\\';";
45const SQL_DELETE_KEY: &str = "DELETE FROM data WHERE key = ?1;";
46const SQL_DELETE_NOC: &str = "DELETE FROM data WHERE key COLLATE NOCASE = ?1;";
47const SQL_DELETE_ALL: &str = "DELETE FROM data;";
48
49// ---------------------------------------------------------------------------
50// SQLiteSet
51// ---------------------------------------------------------------------------
52
53/// A [set](https://en.wikipedia.org/wiki/Set_(abstract_data_type)) that stores keys of type [`T`](crate::Type), backed by an SQLite in-memory database.
54///
55/// It serves as a drop-in replacement for [`HashSet`](std::collections::HashSet).
56///
57/// By default, `SQLiteSet<String>` 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.
58///
59/// &nbsp;
60///
61/// <div class="warning">
62///
63/// 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.
64///
65/// </div>
66pub struct SQLiteSet<T>
67where
68    T: Type,
69{
70    _type: PhantomData<T>,
71    connection: Connection,
72}
73
74impl<T> SQLiteSet<T>
75where
76    T: Type,
77{
78    /// Creates a new, empty SQLite-backed hash set with case-sensitive keys.
79    #[inline]
80    pub fn new() -> Result<Self, Error> {
81        Ok(Self { _type: PhantomData, connection: Self::initialize_connection::<false>()? })
82    }
83
84    /// Creates a new, empty SQLite-backed hash set with case-insensitive keys.
85    ///
86    /// Note that only ASCII characters are case folded. SQLite does **not** attempt to do full UTF case folding!
87    #[inline]
88    pub fn with_nocase() -> Result<Self, Error>
89    where
90        T: CharType,
91    {
92        Ok(Self { _type: PhantomData, connection: Self::initialize_connection::<true>()? })
93    }
94
95    #[inline]
96    fn initialize_connection<const NO_CASE: bool>() -> Result<Connection, Error> {
97        let connection = Connection::open_in_memory()?;
98        connection.pragma_update(None, "journal_mode", "OFF")?;
99        connection.pragma_update(None, "synchronous", "OFF")?;
100        connection.pragma_update(None, "temp_store", "MEMORY")?;
101        if !NO_CASE {
102            connection.execute(SQL_CREATE_TBL.get(&T::SQL_TYPE).unwrap(), [])?;
103        } else {
104            assert_eq!(T::SQL_TYPE, SQL_TYPE_TEXT);
105            connection.execute(SQL_CREATE_NOC, [])?;
106        }
107        Ok(connection)
108    }
109
110    /// Starts a new SQLite transaction for this set.
111    ///
112    /// Please note that using an explicit SQLite transaction allows for much more efficient bulk inserts &#x1F680;
113    ///
114    /// Returns the new [`SQLiteSetTransaction`] instance.
115    pub fn transaction(&mut self) -> Result<SQLiteSetTransaction<'_, T>, Error> {
116        SQLiteSetTransaction::from(&mut self.connection)
117    }
118
119    /// Inserts the given key into the set.
120    ///
121    /// Returns `true`, if the set did not already contain the key; otherwise returns `false`.
122    #[inline]
123    pub fn insert<Q>(&mut self, key: &Q) -> Result<bool, Error>
124    where
125        T: Borrow<Q>,
126        Q: ?Sized + ToSql,
127    {
128        let mut statement = self.connection.prepare_cached(SQL_INSERT_KEY)?;
129        match statement.execute([key]) {
130            Ok(_) => Ok(true),
131            Err(error) => check_constraint_violation(error),
132        }
133    }
134
135    /// Checks whether the set contains the specified key.
136    ///
137    /// For case-sensitive sets, the check is case-sensitive; for case-insensitive sets, the check is case-insensitive.
138    ///
139    /// Returns `true`, if the set contains the key; otherwise returns `false`.
140    #[inline]
141    pub fn contains<Q>(&self, key: &Q) -> Result<bool, Error>
142    where
143        T: Borrow<Q>,
144        Q: ?Sized + ToSql,
145    {
146        let mut statement = self.connection.prepare_cached(SQL_EXISTS_KEY)?;
147        Ok(statement.exists([key])?)
148    }
149
150    /// This is the "case-insensitive" version of the [`contains()`](Self::contains) function.
151    ///
152    /// The check is *always* performed case-insensitive.
153    ///
154    /// Returns `true`, if the set contains the key; otherwise returns `false`.
155    #[inline]
156    pub fn contains_nocase<Q>(&self, key: &Q) -> Result<bool, Error>
157    where
158        T: Borrow<Q> + CharType,
159        Q: ?Sized + ToSql,
160    {
161        let mut statement = self.connection.prepare_cached(SQL_EXISTS_NOC)?;
162        Ok(statement.exists([key])?)
163    }
164
165    /// Removes the specified key from the set, if present.
166    ///
167    /// For case-sensitive sets, the key is treated as case-sensitive; for case-insensitive sets, it is treated as case-insensitive.
168    ///
169    /// Returns `true`, if the set contained the key; otherwise returns `false`.
170    #[inline]
171    pub fn remove<Q>(&mut self, key: &Q) -> Result<bool, Error>
172    where
173        T: Borrow<Q>,
174        Q: ?Sized + ToSql,
175    {
176        let mut statement = self.connection.prepare_cached(SQL_DELETE_KEY)?;
177        Ok(statement.execute([key])? != 0)
178    }
179
180    /// This is the "case-insensitive" version of the [`remove()`](Self::remove) function.
181    ///
182    /// The key is *always* treated as case-insensitive.
183    ///
184    /// Returns `true`, if the set contained the key; otherwise returns `false`.
185    #[inline]
186    pub fn remove_nocase<Q>(&mut self, key: &Q) -> Result<bool, Error>
187    where
188        T: Borrow<Q> + CharType,
189        Q: ?Sized + ToSql,
190    {
191        let mut statement = self.connection.prepare_cached(SQL_DELETE_NOC)?;
192        Ok(statement.execute([key])? != 0)
193    }
194
195    /// Invokes the given `callback` function for each key that is currently contained in the set.
196    ///
197    /// This function does **not** guarantee a specific iteration order.
198    #[inline]
199    pub fn for_each<F>(&self, mut callback: F) -> Result<(), Error>
200    where
201        F: FnMut(T),
202    {
203        let mut statement = self.connection.prepare_cached(SQL_QUERY_KEYS)?;
204        for key in statement.query_map([], |row| row.get(0))? {
205            callback(key?);
206        }
207        Ok(())
208    }
209
210    /// Searches the set for the first key that satisfies the given `predicate`.
211    ///
212    /// Returns the first key that satisfies the given predicate, or `None` if **no** key satisfies the predicate or the set is empty.
213    ///
214    /// This function does **not** guarantee a specific iteration order.
215    ///
216    /// Also, the predicate is **not** always tested on *all* keys, because the function returns at the first match.
217    #[inline]
218    pub fn find<P, Q>(&self, predicate: P) -> Result<Option<T>, Error>
219    where
220        P: Fn(&Q) -> bool,
221        T: Borrow<Q>,
222        Q: ?Sized,
223    {
224        let mut statement = self.connection.prepare_cached(SQL_QUERY_KEYS)?;
225        for key in statement.query_map([], |row| row.get(0))? {
226            let key: T = key?;
227            if predicate(key.borrow()) {
228                return Ok(Some(key));
229            }
230        }
231        Ok(None)
232    }
233
234    /// Invokes the given `callback` function for each key that matches the given pattern.
235    ///
236    /// In the pattern string, the special character `%` matches zero or more characters, the special character `_` matches exactly one character, and the special character `\` serves as an escape character. All other characters are literals.
237    ///
238    /// Please note that SQLite's pattern matching is always *case-insensitive*.
239    ///
240    /// For example, the pattern `A%` matches all keys that begin with the character "A", whereas the pattern `H_nt` matches all keys with a length of exactly four characters where the first character is "H" and the last two characters are "nt".
241    ///
242    /// This function does **not** guarantee a specific iteration order.
243    #[inline]
244    pub fn filter<F, Q>(&self, pattern: &Q, mut callback: F) -> Result<(), Error>
245    where
246        F: FnMut(T),
247        T: Borrow<Q> + CharType,
248        Q: ?Sized + ToSql,
249    {
250        let mut statement = self.connection.prepare_cached(SQL_WHERE_KEYS)?;
251        for key in statement.query_map([pattern], |row| row.get(0))? {
252            callback(key?);
253        }
254        Ok(())
255    }
256
257    /// Returns the number of unique keys in the set.
258    #[inline]
259    pub fn len(&self) -> Result<Size, Error> {
260        let mut statement = self.connection.prepare_cached(SQL_COUNT_KEYS)?;
261        let count: i64 = statement.query_one([], |row| row.get(0))?;
262        Ok(Size::from_i64(count))
263    }
264
265    /// Returns `true` if the set contains **no** keys; otherwise returns `false`.
266    #[inline]
267    pub fn is_empty(&self) -> Result<bool, Error> {
268        Ok(self.len()? == 0)
269    }
270
271    /// Removes *all* keys from the set.
272    #[inline]
273    pub fn clear(&mut self) -> Result<(), Error> {
274        let mut statement = self.connection.prepare_cached(SQL_DELETE_ALL)?;
275        statement.execute([])?;
276        Ok(())
277    }
278}
279
280// ---------------------------------------------------------------------------
281// SQLiteSet Transaction
282// ---------------------------------------------------------------------------
283
284/// Represents an active SQLite transaction for a [`SQLiteSet`].
285///
286/// Most functions provided by this struct mirror the corresponding functions of the `SQLiteSet` struct.
287///
288/// The transaction is committed when the `SQLiteSetTransaction` is dropped.
289pub struct SQLiteSetTransaction<'a, T>
290where
291    T: Type,
292{
293    _type: PhantomData<T>,
294    transaction: Transaction<'a>,
295}
296
297impl<'a, T> SQLiteSetTransaction<'a, T>
298where
299    T: Type,
300{
301    #[inline]
302    fn from(connection: &'a mut Connection) -> Result<Self, Error> {
303        let mut transaction = connection.transaction()?;
304        transaction.set_drop_behavior(rusqlite::DropBehavior::Commit);
305        Ok(Self { _type: PhantomData, transaction })
306    }
307
308    /// Drops the `SQLiteSetTransaction`, thereby committing the SQLite transaction.
309    #[inline]
310    pub fn commit(self) {}
311
312    /// This function is equivalent to [`SQLiteSet::insert()`].
313    #[inline]
314    pub fn insert<Q>(&mut self, key: &Q) -> Result<bool, Error>
315    where
316        T: Borrow<Q>,
317        Q: ?Sized + ToSql,
318    {
319        let mut statement = self.transaction.prepare_cached(SQL_INSERT_KEY)?;
320        match statement.execute([key]) {
321            Ok(_) => Ok(true),
322            Err(error) => check_constraint_violation(error),
323        }
324    }
325
326    /// This function is equivalent to [`SQLiteSet::contains()`].
327    #[inline]
328    pub fn contains<Q>(&self, key: &Q) -> Result<bool, Error>
329    where
330        T: Borrow<Q>,
331        Q: ?Sized + ToSql,
332    {
333        let mut statement = self.transaction.prepare_cached(SQL_EXISTS_KEY)?;
334        Ok(statement.exists([key])?)
335    }
336
337    /// This function is equivalent to [`SQLiteSet::contains_nocase()`].
338    #[inline]
339    pub fn contains_nocase<Q>(&self, key: &Q) -> Result<bool, Error>
340    where
341        T: Borrow<Q> + CharType,
342        Q: ?Sized + ToSql,
343    {
344        let mut statement = self.transaction.prepare_cached(SQL_EXISTS_NOC)?;
345        Ok(statement.exists([key])?)
346    }
347
348    /// This function is equivalent to [`SQLiteSet::remove()`].
349    #[inline]
350    pub fn remove<Q>(&mut self, key: &Q) -> Result<bool, Error>
351    where
352        T: Borrow<Q>,
353        Q: ?Sized + ToSql,
354    {
355        let mut statement = self.transaction.prepare_cached(SQL_DELETE_KEY)?;
356        Ok(statement.execute([key])? != 0)
357    }
358
359    /// This function is equivalent to [`SQLiteSet::remove_nocase()`].
360    #[inline]
361    pub fn remove_nocase<Q>(&mut self, key: &Q) -> Result<bool, Error>
362    where
363        T: Borrow<Q> + CharType,
364        Q: ?Sized + ToSql,
365    {
366        let mut statement = self.transaction.prepare_cached(SQL_DELETE_NOC)?;
367        Ok(statement.execute([key])? != 0)
368    }
369
370    /// This function is equivalent to [`SQLiteSet::for_each()`].
371    #[inline]
372    pub fn for_each<F>(&self, mut callback: F) -> Result<(), Error>
373    where
374        F: FnMut(T),
375    {
376        let mut statement = self.transaction.prepare_cached(SQL_QUERY_KEYS)?;
377        for key in statement.query_map([], |row| row.get(0))? {
378            callback(key?);
379        }
380        Ok(())
381    }
382
383    /// This function is equivalent to [`SQLiteSet::find()`].
384    #[inline]
385    pub fn find<P, Q>(&self, predicate: P) -> Result<Option<T>, Error>
386    where
387        P: Fn(&Q) -> bool,
388        T: Borrow<Q>,
389        Q: ?Sized,
390    {
391        let mut statement = self.transaction.prepare_cached(SQL_QUERY_KEYS)?;
392        for key in statement.query_map([], |row| row.get(0))? {
393            let key: T = key?;
394            if predicate(key.borrow()) {
395                return Ok(Some(key));
396            }
397        }
398        Ok(None)
399    }
400
401    /// This function is equivalent to [`SQLiteSet::filter()`].
402    #[inline]
403    pub fn filter<F, Q>(&self, pattern: &Q, mut callback: F) -> Result<(), Error>
404    where
405        F: FnMut(T),
406        T: Borrow<Q> + CharType,
407        Q: ?Sized + ToSql,
408    {
409        let mut statement = self.transaction.prepare_cached(SQL_WHERE_KEYS)?;
410        for key in statement.query_map([pattern], |row| row.get(0))? {
411            callback(key?);
412        }
413        Ok(())
414    }
415
416    /// This function is equivalent to [`SQLiteSet::len()`].
417    #[inline]
418    pub fn len(&self) -> Result<Size, Error> {
419        let mut statement = self.transaction.prepare_cached(SQL_COUNT_KEYS)?;
420        let count: i64 = statement.query_one([], |row| row.get(0))?;
421        Ok(Size::from_i64(count))
422    }
423
424    /// This function is equivalent to [`SQLiteSet::is_empty()`].
425    #[inline]
426    pub fn is_empty(&self) -> Result<bool, Error> {
427        Ok(self.len()? == 0)
428    }
429
430    /// This function is equivalent to [`SQLiteSet::clear()`].
431    #[inline]
432    pub fn clear(&mut self) -> Result<(), Error> {
433        let mut statement = self.transaction.prepare_cached(SQL_DELETE_ALL)?;
434        statement.execute([])?;
435        Ok(())
436    }
437}