composable-indexes 0.9.0

use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec::Vec;

use proptest::prelude::Arbitrary;

use crate::core::{Collection, Index, Key};

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DBOperation<T> {
    InsertOrUpdate(Key, T),
    Delete(Key),
}

#[derive(Debug, Clone)]
pub struct TestOps<T: Clone> {
    pub operations: Vec<DBOperation<T>>,
}

impl<T: Clone + 'static> TestOps<T> {
    pub fn apply<Ix: Index<T>>(&self, db: &mut Collection<T, Ix>) {
        self.operations.iter().cloned().for_each(|op| match op {
            DBOperation::InsertOrUpdate(key, value) => {
                db.update_by_key(key, |_existing| value);
            }
            DBOperation::Delete(key) => {
                db.delete_by_key(key);
            }
        });
    }

    pub fn end_state(&self) -> BTreeMap<Key, T> {
        let mut ret = BTreeMap::new();

        self.operations.iter().for_each(|op| match op {
            DBOperation::InsertOrUpdate(key, value) => {
                ret.insert(*key, value.clone());
            }
            DBOperation::Delete(key) => {
                ret.remove(key);
            }
        });

        ret
    }
}

impl<T: Arbitrary + Clone + 'static> proptest::arbitrary::Arbitrary for TestOps<T> {
    type Strategy = proptest::strategy::BoxedStrategy<Self>;
    type Parameters = ();

    fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
        use proptest::prelude::*;
        prop::collection::vec((0u32..50, any::<Option<T>>()), 0..400)
            .prop_map(|ops| {
                // Ensure the first operation per key is Some
                let mut seen_keys: BTreeSet<u32> = BTreeSet::new();
                let mut operations = Vec::new();
                for (key, value) in ops {
                    if seen_keys.insert(key) {
                        if value.is_some() {
                            operations.push((key, value));
                        } else {
                            // Ignore if the first value is some
                        }
                    } else {
                        operations.push((key, value));
                    }
                }

                operations
            })
            .prop_map(|ops| {
                let ops = ops
                    .into_iter()
                    .map(|(k, v)| {
                        let k = Key::unsafe_from_u64(k as u64);
                        if let Some(v) = v {
                            DBOperation::<T>::InsertOrUpdate(k, v)
                        } else {
                            DBOperation::Delete(k)
                        }
                    })
                    .collect::<Vec<DBOperation<T>>>();

                TestOps { operations: ops }
            })
            .boxed()
    }
}