treehouse 0.1.0

use crate::{err, utils};

use std::collections::BTreeSet;

use sled::transaction::{TransactionalTree, UnabortableTransactionError, ConflictableTransactionError};

pub trait MaybeTransactional {
    type Error;

    fn insert<K, V>(
        &self,
        key: K,
        value: V,
    ) -> Result<Option<sled::IVec>, Self::Error>
    where
        sled::IVec: From<K> + From<V>,
        K: AsRef<[u8]>;

    fn remove<K>(
        &self,
        key: K,
    ) -> Result<Option<sled::IVec>, Self::Error>
    where
        sled::IVec: From<K>,
        K: AsRef<[u8]>;

    fn get<K: AsRef<[u8]>>(
        &self,
        key: K,
    ) -> Result<Option<sled::IVec>, Self::Error>;

    fn apply_batch(
        &self,
        batch: sled::Batch,
    ) -> Result<(), Self::Error>;
}

impl MaybeTransactional for TransactionalTree {
    type Error = ConflictableTransactionError<err::Error>;

    fn insert<K, V>(
        &self,
        key: K,
        value: V,
    ) -> Result<Option<sled::IVec>, Self::Error>
    where
        sled::IVec: From<K> + From<V>,
        K: AsRef<[u8]> {
        self.insert(key, value).map_err(|e| e.into() )
    }

    fn remove<K>(
        &self,
        key: K,
    ) -> Result<Option<sled::IVec>, Self::Error>
    where
        sled::IVec: From<K>,
        K: AsRef<[u8]> {
        self.remove(key).map_err(|e| e.into() )
    }

    fn get<K: AsRef<[u8]>>(
        &self,
        key: K,
    ) -> Result<Option<sled::IVec>, Self::Error> {
        self.get(key).map_err(|e| e.into() )
    }

    fn apply_batch(
        &self,
        batch: sled::Batch,
    ) -> Result<(), Self::Error> {
        self.apply_batch(&batch).map_err(|e| e.into() )
    }
}

impl MaybeTransactional for sled::Tree {
    type Error = err::Error;

    fn insert<K, V>(
        &self,
        key: K,
        value: V,
    ) -> Result<Option<sled::IVec>, Self::Error>
    where
        sled::IVec: From<K> + From<V>,
        K: AsRef<[u8]> {
        self.insert(key, value).map_err(|e| e.into() )
    }

    fn remove<K>(
        &self,
        key: K,
    ) -> Result<Option<sled::IVec>, Self::Error>
    where
        sled::IVec: From<K>,
        K: AsRef<[u8]> {
        self.remove(key).map_err(|e| e.into() )
    }

    fn get<K: AsRef<[u8]>>(
        &self,
        key: K,
    ) -> Result<Option<sled::IVec>, Self::Error> {
        self.get(key).map_err(|e| e.into() )
    }

    fn apply_batch(
        &self,
        batch: sled::Batch,
    ) -> Result<(), Self::Error> {
        self.apply_batch(batch).map_err(|e| e.into() )
    }
}

pub struct TransactionalMeta<'a>(pub &'a TransactionalTree);

impl<'a> TransactionalMeta<'a> {

    pub fn add(&self, terms: Option<BTreeSet<Vec<u8>>>, id: u64, id_bytes: &[u8]) -> Result<(), ConflictableTransactionError<err::Error>> {
        if let Some(prev_terms) = self.insert_terms(id_bytes, &terms)? {
            
            let b_tree_set_default;

            let filtered_prev_terms = match terms {
                Some(ref terms) => prev_terms.difference(terms),
                None => {
                    b_tree_set_default = BTreeSet::default();
                    prev_terms.difference(&b_tree_set_default)
                }
            };

            for prev_term in filtered_prev_terms {
                self.remove_term_id(prev_term.as_slice(), id)?;
            }

            if let Some(terms) = terms {
                for term in terms.difference(&prev_terms) {
                    self.insert_term_id(term.as_slice(), id)?;
                }
            }

        } else {
            if let Some(terms) = terms {
                for term in terms {
                    self.insert_term_id(term, id)?;
                }
            }
        }

        Ok(())
    }

    fn insert_terms(&self, id_bytes: &[u8], terms: &Option<BTreeSet<Vec<u8>>>) -> Result<Option<BTreeSet<Vec<u8>>>, ConflictableTransactionError<err::Error>> {
        let key = crate::query::TERMS_PREFIX.into_iter().chain(id_bytes).copied().collect::<Vec<_>>();

        let serialized_prev_terms_opt = match terms {
            Some(ref terms) => {
                let serialized_terms = utils::serialize(terms)?;
                self.0.insert(key, serialized_terms)?
            },
            None => self.0.remove(key)?
        };

        let out = if let Some(serialized_prev_terms) = serialized_prev_terms_opt {
            let prev_terms: BTreeSet<Vec<u8>> = utils::deserialize(&serialized_prev_terms)?;
            Some(prev_terms)
        } else {
            None
        };

        Ok(out)
    }

    fn insert_term_id<T: Into<Vec<u8>>>(&self, term: T, id: u64) -> Result<(), ConflictableTransactionError<err::Error>> {
        let term = term.into();
        let term_val = match self.0.get(&term)? {
            Some(serialized_prev_ids) => {
                let mut prev_ids: Vec<u64> = utils::deserialize(&serialized_prev_ids)?;
                prev_ids.push(id);
                prev_ids
            },
            None => vec![id],
        };

        self.0.insert(term, utils::serialize(&term_val)?)?;

        Ok(())
    }

    fn remove_term_id<T: Into<Vec<u8>>>(&self, term: T, id: u64) -> Result<(), ConflictableTransactionError<err::Error>> {
        let term = term.into();
        let term_val = match self.0.get(&term)? {
            Some(serialized_prev_ids) => {
                let prev_ids: Vec<u64> = utils::deserialize(&serialized_prev_ids)?;
                Some(prev_ids.into_iter().filter(|x| x == &id ).collect::<Vec<_>>())
            }
            None => None        
        };

        if let Some(term_val) = term_val {
            self.0.insert(term, utils::serialize(&term_val)?)?;
        }

        Ok(())
    }
}