use super::*;
use core::mem;
#[derive(Clone, Debug, Default)]
pub struct BTreeMap<KM, VM> {
key_mutator: KM,
value_mutator: VM,
}
pub fn btree_map<KM, VM>(key_mutator: KM, value_mutator: VM) -> BTreeMap<KM, VM> {
BTreeMap {
key_mutator,
value_mutator,
}
}
impl<KM, VM, K, V> Mutate<alloc::collections::BTreeMap<K, V>> for BTreeMap<KM, VM>
where
KM: Generate<K>,
VM: Generate<V> + Mutate<V>,
K: Ord,
{
#[inline]
fn mutation_count(
&self,
value: &alloc::collections::BTreeMap<K, V>,
shrink: bool,
) -> core::option::Option<u32> {
let mut count = 0u32;
count += !shrink as u32;
count += !value.is_empty() as u32;
for (k, v) in value.iter() {
count += self.value_mutator.mutation_count(v, shrink)?;
count += self.key_mutator.mutation_count(k, shrink)?;
}
Some(count)
}
#[inline]
fn mutate(
&mut self,
c: &mut Candidates,
value: &mut alloc::collections::BTreeMap<K, V>,
) -> Result<()> {
if !c.shrink() {
c.mutation(|ctx| {
let k = self.key_mutator.generate(ctx)?;
let v = self.value_mutator.generate(ctx)?;
value.insert(k, v);
Ok(())
})?;
}
if !value.is_empty() {
c.mutation(|ctx| {
if ctx.rng().gen_bool() {
value.pop_first();
} else {
value.pop_last();
}
Ok(())
})?;
}
for (_k, v) in value.iter_mut() {
self.value_mutator.mutate(c, v)?;
}
if !value.is_empty() {
let mut early_exit = None;
for (mut key, val) in mem::take(value) {
if early_exit.is_none() {
match self.key_mutator.mutate(c, &mut key) {
Ok(()) => {}
Err(e) if e.is_early_exit() => {
early_exit = Some(e);
}
Err(e) => return Err(e),
}
}
value.insert(key, val);
}
if let Some(e) = early_exit {
return Err(e);
}
}
Ok(())
}
}
impl<KM, VM, K, V> Generate<alloc::collections::BTreeMap<K, V>> for BTreeMap<KM, VM>
where
KM: Generate<K>,
VM: Generate<V> + Mutate<V>,
K: Ord,
{
#[inline]
fn generate(&mut self, ctx: &mut Context) -> Result<alloc::collections::BTreeMap<K, V>> {
self.generate_via_mutate(ctx, 1)
}
}
impl<K, V> DefaultMutate for alloc::collections::BTreeMap<K, V>
where
K: DefaultMutate + Ord,
K::DefaultMutate: Generate<K>,
V: DefaultMutate,
V::DefaultMutate: Generate<V>,
{
type DefaultMutate = BTreeMap<K::DefaultMutate, V::DefaultMutate>;
}