selene-core 0.9.0-rc.1

Backend for selene-server
Documentation
use std::{collections::HashMap, path::PathBuf, sync::Arc};

use lunar_lib::{
    database::{
        CompareAndSwapTransaction, DatabaseEntry, Db, DbIdIterExt, Entry, TransactionError,
        caching::{Cacheable, IdCacheIterExt},
    },
    define_db,
    id::{ID_SIZE, Id},
    iterator_ext::IteratorExtensions,
    paths::data_dir,
};

mod tx_extensions;
pub(crate) use tx_extensions::*;

#[cfg(debug_assertions)]
pub mod validator;

define_db!(Selene {
    fn path() -> Option<PathBuf> {
        Some(data_dir().join("selene_data"))
    }
});

pub trait SeleneEntryExt<T: DatabaseEntry> {
    fn to_db_entry(self) -> Entry<T>;
}

impl<T: DatabaseEntry> SeleneEntryExt<T> for crate::library::Entry<T> {
    fn to_db_entry(self) -> Entry<T> {
        self.into()
    }
}

pub trait Searchable: DatabaseEntry {
    const SEARCH_INDEX: &'static str;

    fn search_name(&self) -> Option<&str>;

    fn get_search_key(&self) -> Option<[u8; 8]> {
        let search_name = self.search_name().map(str::to_ascii_lowercase)?;
        let bytes = search_name.as_bytes();
        let len = bytes.len().min(8);
        let mut key = [0u8; 8];
        key[..len].copy_from_slice(&bytes[..len]);
        Some(key)
    }

    fn update_search_index(
        entry: &Entry<Self>,
        cas_tx: &mut CompareAndSwapTransaction<Self::DbInner>,
    ) -> Result<(), TransactionError> {
        let Some(key) = entry.get_search_key() else {
            return Ok(());
        };

        let request = cas_tx.get_or_new_index(Self::SEARCH_INDEX);
        request.fetch_and_update(&key, |old| {
            let existing = old.as_deref().unwrap_or_default();

            if existing.as_chunks::<{ ID_SIZE }>().0.contains(&*entry.id()) {
                return old;
            }

            let mut new = Vec::with_capacity(existing.len() + ID_SIZE);
            new.extend_from_slice(existing);
            new.extend_from_slice(&*entry.id());

            Some(new.into())
        })?;

        Ok(())
    }

    fn remove_search_index(
        entry: &Entry<Self>,
        cas_tx: &mut CompareAndSwapTransaction<Self::DbInner>,
    ) -> Result<(), TransactionError> {
        let Some(key) = entry.get_search_key() else {
            return Ok(());
        };

        let request = cas_tx.get_or_new_index(Self::SEARCH_INDEX);
        request.fetch_and_update(&key, |old| {
            let new = old.as_deref().map(|old| {
                let chunks = old.as_chunks::<{ ID_SIZE }>().0;

                if let Some(idx) = chunks.iter().position(|c| *c == *entry.id()) {
                    let mut v = Vec::with_capacity((chunks.len() - 1) * { ID_SIZE });
                    v.extend_from_slice(&old[..idx * { ID_SIZE }]);
                    v.extend_from_slice(&old[(idx + 1) * { ID_SIZE }..]);
                    v
                } else {
                    old.to_vec()
                }
            });

            if let Some(new) = new
                && !new.is_empty()
            {
                Some(new.into())
            } else {
                None
            }
        })?;

        Ok(())
    }

    fn search(
        db: &Db<Self::DbInner>,
        query: &str,
        limit: usize,
        offset: usize,
    ) -> Result<Vec<Entry<Self>>, TransactionError> {
        let query = query.to_ascii_lowercase();
        let len = query.len().min(8);
        let key = &query.as_bytes()[..len];

        let items = db
            .index(Self::SEARCH_INDEX)
            .scan_prefix(key)
            .try_map(|i| -> Result<_, TransactionError> {
                let (_, v) = i.map_err(TransactionError::from)?;
                let ids = v
                    .as_chunks::<{ ID_SIZE }>()
                    .0
                    .iter()
                    .copied()
                    .map(Id::<Self>::from)
                    .db_get_iter(db)
                    .filter_map(|item| match item {
                        Ok(item) => item
                            .search_name()
                            .expect("Items without names are not indexed")
                            .to_ascii_lowercase()
                            .starts_with(&query)
                            .then_some(Ok(item)),
                        Err(err) => Some(Err(err)),
                    })
                    .try_to_vec()?;

                Ok(ids)
            })?
            .flatten()
            .skip(offset)
            .take(limit)
            .to_vec();

        Ok(items)
    }

    fn search_cache(
        db: &Db<Self::DbInner>,
        query: &str,
        limit: usize,
        offset: usize,
    ) -> Result<Vec<Arc<Entry<Self>>>, TransactionError>
    where
        Self: Cacheable,
    {
        let query = query.to_ascii_lowercase();
        let len = query.len().min(8);
        let key = &query.as_bytes()[..len];

        let items = db
            .index(Self::SEARCH_INDEX)
            .scan_prefix(key)
            .try_map(|i| -> Result<_, TransactionError> {
                let (_, v) = i.map_err(TransactionError::from)?;
                let ids = v
                    .as_chunks::<{ ID_SIZE }>()
                    .0
                    .iter()
                    .copied()
                    .map(Id::<Self>::from)
                    .cache_get_iter(db)
                    .filter_map(|item| match item {
                        Ok(item) => item
                            .search_name()
                            .expect("Items without names are not indexed")
                            .to_ascii_lowercase()
                            .starts_with(&query)
                            .then_some(Ok(item)),
                        Err(err) => Some(Err(err)),
                    })
                    .try_to_vec()?;

                Ok(ids)
            })?
            .flatten()
            .skip(offset)
            .take(limit)
            .to_vec();

        Ok(items)
    }
}

pub trait GenreIndexable: DatabaseEntry {
    const GENRE_INDEX: &'static str;

    fn get_with_genre(
        db: &Db<Self::DbInner>,
        genre: impl AsRef<str>,
        limit: usize,
        offset: usize,
    ) -> Result<Vec<Entry<Self>>, TransactionError> {
        let value = db.index(Self::GENRE_INDEX).get(genre.as_ref().as_bytes())?;
        if value.is_none() {
            return Ok(Vec::new());
        }

        if let Some(value) = value {
            let entries = value
                .as_chunks::<{ ID_SIZE }>()
                .0
                .iter()
                .skip(offset)
                .take(limit)
                .copied()
                .map(Id::<Self>::from)
                .db_get(db)?;
            Ok(entries)
        } else {
            Ok(Vec::new())
        }
    }

    fn get_all_genres(db: &Db<Self::DbInner>) -> Result<HashMap<String, usize>, TransactionError> {
        let hash_map = db
            .index(Self::GENRE_INDEX)
            .iter()
            .map(|r| -> Result<(String, usize), TransactionError> {
                let (genre, keys) = r?;

                let genre =
                    String::from_utf8(genre.to_vec()).expect("Stored genre key is UTF-8 bytes");
                let entries = keys.len() / { ID_SIZE };

                Ok((genre, entries))
            })
            .collect::<Result<HashMap<_, _>, _>>()?;

        Ok(hash_map)
    }

    fn add_to_genres(
        entry: &Entry<Self>,
        genres: &[String],
        cas_tx: &mut CompareAndSwapTransaction<Self::DbInner>,
    ) -> Result<(), TransactionError> {
        let request = cas_tx.get_or_new_index(Self::GENRE_INDEX);
        for genre in genres {
            request.fetch_and_update(genre.as_bytes(), |old| {
                let existing = old.as_deref().unwrap_or_default();

                if existing.as_chunks::<{ ID_SIZE }>().0.contains(&*entry.id()) {
                    return old;
                }

                let mut new = Vec::with_capacity(existing.len() + ID_SIZE);
                new.extend_from_slice(existing);
                new.extend_from_slice(&*entry.id());

                Some(new.into())
            })?;
        }

        Ok(())
    }

    fn remove_from_genres(
        entry: &Entry<Self>,
        genres: &[String],
        cas_tx: &mut CompareAndSwapTransaction<Self::DbInner>,
    ) -> Result<(), TransactionError> {
        let request = cas_tx.get_or_new_index(Self::GENRE_INDEX);
        for genre in genres {
            request.fetch_and_update(genre.as_bytes(), |old| {
                let new = old.as_deref().map(|old| {
                    let chunks = old.as_chunks::<{ ID_SIZE }>().0;

                    if let Some(idx) = chunks.iter().position(|c| *c == *entry.id()) {
                        let mut v = Vec::with_capacity((chunks.len() - 1) * { ID_SIZE });
                        v.extend_from_slice(&old[..idx * { ID_SIZE }]);
                        v.extend_from_slice(&old[(idx + 1) * { ID_SIZE }..]);
                        v
                    } else {
                        old.to_vec()
                    }
                });

                if let Some(new) = new
                    && !new.is_empty()
                {
                    Some(new.into())
                } else {
                    None
                }
            })?;
        }

        Ok(())
    }
}