use crate::{
CharType, Error, Size, SizeExt, Type,
common::{SQL_TYPE_TEXT, SQL_TYPES, SqlType, check_constraint_violation},
rusqlite::{Connection, OptionalExtension, ToSql, Transaction, params},
};
use cfg_if::cfg_if;
use std::{borrow::Borrow, collections::BTreeMap, marker::PhantomData, sync::LazyLock};
static SQL_CREATE_TBL: LazyLock<BTreeMap<(SqlType, SqlType), String>> = LazyLock::new(|| {
let mut map = BTreeMap::<(SqlType, SqlType), String>::new();
for key_type in SQL_TYPES {
for value_type in SQL_TYPES {
cfg_if! {
if #[cfg(feature = "with_rowid")] {
map.insert((key_type, value_type), format!("CREATE TABLE data (key {key_type} PRIMARY KEY NOT NULL, value {value_type} NOT NULL);"));
} else {
map.insert((key_type, value_type), format!("CREATE TABLE data (key {key_type} PRIMARY KEY NOT NULL, value {value_type} NOT NULL) WITHOUT ROWID;"));
}
}
}
}
map
});
static SQL_CREATE_NOC: LazyLock<BTreeMap<SqlType, String>> = LazyLock::new(|| {
let mut map = BTreeMap::<SqlType, String>::new();
for value_type in SQL_TYPES {
cfg_if! {
if #[cfg(feature = "with_rowid")] {
map.insert(value_type, format!("CREATE TABLE data (key TEXT PRIMARY KEY NOT NULL COLLATE NOCASE, value {value_type} NOT NULL);"));
} else {
map.insert(value_type, format!("CREATE TABLE data (key TEXT PRIMARY KEY NOT NULL COLLATE NOCASE, value {value_type} NOT NULL) WITHOUT ROWID;"));
}
}
}
map
});
const SQL_COUNT_KEYS: &str = "SELECT COUNT(*) FROM data;";
const SQL_INSERT_KEY: &str = "INSERT INTO data (key, value) VALUES (?1, ?2);";
const SQL_UPSERT_KEY: &str = "INSERT INTO data (key, value) VALUES (?1, ?2) ON CONFLICT(key) DO UPDATE SET value = excluded.value;";
const SQL_EXISTS_KEY: &str = "SELECT 1 FROM data WHERE key = ? LIMIT 1;";
const SQL_EXISTS_NOC: &str = "SELECT 1 FROM data WHERE key COLLATE NOCASE = ? LIMIT 1;";
const SQL_LOOKUP_KEY: &str = "SELECT value FROM data WHERE key = ? LIMIT 1;";
const SQL_LOOKUP_NOC: &str = "SELECT value FROM data WHERE key COLLATE NOCASE = ? LIMIT 1;";
const SQL_QUERY_KEYS: &str = "SELECT key FROM data;";
const SQL_WHERE_KEYS: &str = "SELECT key FROM data WHERE key LIKE ?1 ESCAPE '\\';";
const SQL_QUERY_PAIR: &str = "SELECT key, value FROM data;";
const SQL_WHERE_PAIR: &str = "SELECT key, value FROM data WHERE key LIKE ?1 ESCAPE '\\';";
const SQL_DELETE_KEY: &str = "DELETE FROM data WHERE key = ?;";
const SQL_DELETE_NOC: &str = "DELETE FROM data WHERE key COLLATE NOCASE = ?;";
const SQL_DELETE_ALL: &str = "DELETE FROM data;";
pub struct SQLiteMap<K, V>
where
K: Type,
V: Type,
{
_types: (PhantomData<K>, PhantomData<V>),
connection: Connection,
}
impl<K, V> SQLiteMap<K, V>
where
K: Type,
V: Type,
{
#[inline]
pub fn new() -> Result<Self, Error> {
Ok(Self { _types: (PhantomData, PhantomData), connection: Self::initialize_connection::<false>()? })
}
#[inline]
pub fn with_nocase() -> Result<Self, Error>
where
K: CharType,
{
Ok(Self { _types: (PhantomData, PhantomData), connection: Self::initialize_connection::<true>()? })
}
#[inline]
fn initialize_connection<const NO_CASE: bool>() -> Result<Connection, Error> {
let connection = Connection::open_in_memory()?;
connection.pragma_update(None, "journal_mode", "OFF")?;
connection.pragma_update(None, "synchronous", "OFF")?;
connection.pragma_update(None, "temp_store", "MEMORY")?;
if !NO_CASE {
connection.execute(SQL_CREATE_TBL.get(&(K::SQL_TYPE, V::SQL_TYPE)).unwrap(), [])?;
} else {
assert_eq!(K::SQL_TYPE, SQL_TYPE_TEXT);
connection.execute(SQL_CREATE_NOC.get(&V::SQL_TYPE).unwrap(), [])?;
}
Ok(connection)
}
#[inline]
pub fn transaction(&mut self) -> Result<SQLiteMapTransaction<'_, K, V>, Error> {
SQLiteMapTransaction::from(&mut self.connection)
}
#[inline]
pub fn insert<Q, R>(&mut self, key: &Q, value: &R) -> Result<bool, Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
V: Borrow<R>,
R: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_INSERT_KEY)?;
match statement.execute(params![key, value]) {
Ok(_) => Ok(true),
Err(error) => check_constraint_violation(error),
}
}
#[inline]
pub fn update<Q, R>(&mut self, key: &Q, value: &R) -> Result<(), Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
V: Borrow<R>,
R: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_UPSERT_KEY)?;
statement.execute(params![key, value])?;
Ok(())
}
#[inline]
pub fn contains_key<Q>(&self, key: &Q) -> Result<bool, Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_EXISTS_KEY)?;
Ok(statement.exists([key])?)
}
#[inline]
pub fn contains_key_nocase<Q>(&self, key: &Q) -> Result<bool, Error>
where
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_EXISTS_NOC)?;
Ok(statement.exists([key])?)
}
#[inline]
pub fn get<Q>(&self, key: &Q) -> Result<Option<V>, Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_LOOKUP_KEY)?;
Ok(statement.query_one([key], |row| row.get(0)).optional()?)
}
#[inline]
pub fn get_nocase<Q>(&self, key: &Q) -> Result<Option<V>, Error>
where
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_LOOKUP_NOC)?;
Ok(statement.query_one([key], |row| row.get(0)).optional()?)
}
#[inline]
pub fn remove<Q>(&mut self, key: &Q) -> Result<bool, Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_DELETE_KEY)?;
Ok(statement.execute([key])? != 0)
}
#[inline]
pub fn remove_nocase<Q>(&mut self, key: &Q) -> Result<bool, Error>
where
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_DELETE_NOC)?;
Ok(statement.execute([key])? != 0)
}
#[inline]
pub fn for_each_key<F>(&self, mut callback: F) -> Result<(), Error>
where
F: FnMut(K),
{
let mut statement = self.connection.prepare_cached(SQL_QUERY_KEYS)?;
for key in statement.query_map([], |row| row.get(0))? {
callback(key?);
}
Ok(())
}
#[inline]
pub fn for_each_entry<F>(&self, mut callback: F) -> Result<(), Error>
where
F: FnMut(K, V),
{
let mut statement = self.connection.prepare_cached(SQL_QUERY_PAIR)?;
let mut result = statement.query([])?;
while let Some(key_value) = result.next()? {
let (key, value): (K, V) = (key_value.get(0)?, key_value.get(1)?);
callback(key, value);
}
Ok(())
}
#[inline]
pub fn find_key<P>(&self, predicate: P) -> Result<Option<K>, Error>
where
P: Fn(&K) -> bool,
{
let mut statement = self.connection.prepare_cached(SQL_QUERY_KEYS)?;
for key in statement.query_map([], |row| row.get(0))? {
let key: K = key?;
if predicate(key.borrow()) {
return Ok(Some(key));
}
}
Ok(None)
}
#[inline]
pub fn find_entry<P>(&self, predicate: P) -> Result<Option<(K, V)>, Error>
where
P: Fn(&K, &V) -> bool,
{
let mut statement = self.connection.prepare_cached(SQL_QUERY_PAIR)?;
let mut result = statement.query([])?;
while let Some(key_value) = result.next()? {
let (key, value): (K, V) = (key_value.get(0)?, key_value.get(1)?);
if predicate(&key, &value) {
return Ok(Some((key, value)));
}
}
Ok(None)
}
#[inline]
pub fn filter_keys<F, Q>(&self, pattern: &Q, mut callback: F) -> Result<(), Error>
where
F: FnMut(K),
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_WHERE_KEYS)?;
for key in statement.query_map([pattern], |row| row.get(0))? {
callback(key?);
}
Ok(())
}
#[inline]
pub fn filter_entries<F, Q>(&self, pattern: &Q, mut callback: F) -> Result<(), Error>
where
F: FnMut(K, V),
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.connection.prepare_cached(SQL_WHERE_PAIR)?;
let mut result = statement.query([pattern])?;
while let Some(key_value) = result.next()? {
let (key, value): (K, V) = (key_value.get(0)?, key_value.get(1)?);
callback(key, value);
}
Ok(())
}
#[inline]
pub fn len(&self) -> Result<Size, Error> {
let mut statement = self.connection.prepare_cached(SQL_COUNT_KEYS)?;
let count: i64 = statement.query_one([], |row| row.get(0))?;
Ok(Size::from_i64(count))
}
#[inline]
pub fn is_empty(&self) -> Result<bool, Error> {
Ok(self.len()? == 0)
}
#[inline]
pub fn clear(&mut self) -> Result<(), Error> {
let mut statement = self.connection.prepare_cached(SQL_DELETE_ALL)?;
statement.execute([])?;
Ok(())
}
}
pub struct SQLiteMapTransaction<'a, K, V>
where
K: Type,
V: Type,
{
_types: (PhantomData<K>, PhantomData<V>),
transaction: Transaction<'a>,
}
impl<'a, K, V> SQLiteMapTransaction<'a, K, V>
where
K: Type,
V: Type,
{
#[inline]
fn from(connection: &'a mut Connection) -> Result<Self, Error> {
let mut transaction = connection.transaction()?;
transaction.set_drop_behavior(rusqlite::DropBehavior::Commit);
Ok(Self { _types: (PhantomData, PhantomData), transaction })
}
#[inline]
pub fn commit(self) {}
#[inline]
pub fn insert<Q, R>(&mut self, key: &Q, value: &R) -> Result<bool, Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
V: Borrow<R>,
R: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_INSERT_KEY)?;
match statement.execute(params![key, value]) {
Ok(_) => Ok(true),
Err(error) => check_constraint_violation(error),
}
}
#[inline]
pub fn update<Q, R>(&mut self, key: &Q, value: &R) -> Result<(), Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
V: Borrow<R>,
R: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_UPSERT_KEY)?;
statement.execute(params![key, value])?;
Ok(())
}
#[inline]
pub fn contains_key<Q>(&self, key: &Q) -> Result<bool, Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_EXISTS_KEY)?;
Ok(statement.exists([key])?)
}
#[inline]
pub fn contains_key_nocase<Q>(&self, key: &Q) -> Result<bool, Error>
where
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_EXISTS_NOC)?;
Ok(statement.exists([key])?)
}
#[inline]
pub fn get<Q>(&self, key: &Q) -> Result<Option<V>, Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_LOOKUP_KEY)?;
Ok(statement.query_one([key], |row| row.get(0)).optional()?)
}
#[inline]
pub fn get_nocase<Q>(&self, key: &Q) -> Result<Option<V>, Error>
where
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_LOOKUP_NOC)?;
Ok(statement.query_one([key], |row| row.get(0)).optional()?)
}
#[inline]
pub fn remove<Q>(&mut self, key: &Q) -> Result<bool, Error>
where
K: Borrow<Q>,
Q: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_DELETE_KEY)?;
Ok(statement.execute([key])? != 0)
}
#[inline]
pub fn remove_nocase<Q>(&mut self, key: &Q) -> Result<bool, Error>
where
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_DELETE_NOC)?;
Ok(statement.execute([key])? != 0)
}
#[inline]
pub fn for_each_key<F>(&self, mut callback: F) -> Result<(), Error>
where
F: FnMut(K),
{
let mut statement = self.transaction.prepare_cached(SQL_QUERY_KEYS)?;
for key in statement.query_map([], |row| row.get(0))? {
callback(key?);
}
Ok(())
}
#[inline]
pub fn for_each_entry<F>(&self, mut callback: F) -> Result<(), Error>
where
F: FnMut(K, V),
{
let mut statement = self.transaction.prepare_cached(SQL_QUERY_PAIR)?;
let mut result = statement.query([])?;
while let Some(key_value) = result.next()? {
let (key, value): (K, V) = (key_value.get(0)?, key_value.get(1)?);
callback(key, value);
}
Ok(())
}
#[inline]
pub fn find_key<P>(&self, predicate: P) -> Result<Option<K>, Error>
where
P: Fn(&K) -> bool,
{
let mut statement = self.transaction.prepare_cached(SQL_QUERY_KEYS)?;
for key in statement.query_map([], |row| row.get(0))? {
let key: K = key?;
if predicate(key.borrow()) {
return Ok(Some(key));
}
}
Ok(None)
}
#[inline]
pub fn find_entry<P>(&self, predicate: P) -> Result<Option<(K, V)>, Error>
where
P: Fn(&K, &V) -> bool,
{
let mut statement = self.transaction.prepare_cached(SQL_QUERY_PAIR)?;
let mut result = statement.query([])?;
while let Some(key_value) = result.next()? {
let (key, value): (K, V) = (key_value.get(0)?, key_value.get(1)?);
if predicate(&key, &value) {
return Ok(Some((key, value)));
}
}
Ok(None)
}
#[inline]
pub fn filter_keys<F, Q>(&self, pattern: &Q, mut callback: F) -> Result<(), Error>
where
F: FnMut(K),
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_WHERE_KEYS)?;
for key in statement.query_map([pattern], |row| row.get(0))? {
callback(key?);
}
Ok(())
}
#[inline]
pub fn filter_entries<F, Q>(&self, pattern: &Q, mut callback: F) -> Result<(), Error>
where
F: FnMut(K, V),
K: Borrow<Q> + CharType,
Q: ?Sized + ToSql,
{
let mut statement = self.transaction.prepare_cached(SQL_WHERE_PAIR)?;
let mut result = statement.query([pattern])?;
while let Some(key_value) = result.next()? {
let (key, value): (K, V) = (key_value.get(0)?, key_value.get(1)?);
callback(key, value);
}
Ok(())
}
#[inline]
pub fn len(&self) -> Result<Size, Error> {
let mut statement = self.transaction.prepare_cached(SQL_COUNT_KEYS)?;
let count: i64 = statement.query_one([], |row| row.get(0))?;
Ok(Size::from_i64(count))
}
#[inline]
pub fn is_empty(&self) -> Result<bool, Error> {
Ok(self.len()? == 0)
}
#[inline]
pub fn clear(&mut self) -> Result<(), Error> {
let mut statement = self.transaction.prepare_cached(SQL_DELETE_ALL)?;
statement.execute([])?;
Ok(())
}
}