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