ringbuf-basedrop 0.1.1

use crate::RingBuffer;
use alloc::sync::Arc;
use alloc::{collections::BTreeSet, vec, vec::Vec};
use basedrop::Collector;
use core::cell::RefCell;

#[derive(Debug)]
struct Dropper {
    id: i32,
    set: Arc<RefCellWrapper<BTreeSet<i32>>>,
}

impl Dropper {
    fn new(set: Arc<RefCellWrapper<BTreeSet<i32>>>, id: i32) -> Self {
        if !set.0.borrow_mut().insert(id) {
            panic!("value {} already exists", id);
        }
        Self { set, id }
    }
}

impl Drop for Dropper {
    fn drop(&mut self) {
        if !self.set.0.borrow_mut().remove(&self.id) {
            panic!("value {} already removed", self.id);
        }
    }
}

#[derive(Debug)]
struct RefCellWrapper<T>(RefCell<T>);

unsafe impl<T> Send for RefCellWrapper<T> {}
unsafe impl<T> Sync for RefCellWrapper<T> {}

#[test]
fn single() {
    let mut collector = Collector::new();

    let set = Arc::new(RefCellWrapper(RefCell::new(BTreeSet::new())));

    let cap = 3;
    let buf = RingBuffer::new(cap);

    assert_eq!(set.0.borrow().len(), 0);

    {
        let (mut prod, mut cons) = buf.split(&collector.handle());

        prod.push(Dropper::new(Arc::clone(&set), 1)).unwrap();
        collector.collect();
        assert_eq!(set.0.borrow().len(), 1);
        prod.push(Dropper::new(Arc::clone(&set), 2)).unwrap();
        collector.collect();
        assert_eq!(set.0.borrow().len(), 2);
        prod.push(Dropper::new(Arc::clone(&set), 3)).unwrap();
        collector.collect();
        assert_eq!(set.0.borrow().len(), 3);

        cons.pop().unwrap();
        collector.collect();
        assert_eq!(set.0.borrow().len(), 2);
        cons.pop().unwrap();
        collector.collect();
        assert_eq!(set.0.borrow().len(), 1);

        prod.push(Dropper::new(Arc::clone(&set), 4)).unwrap();
        collector.collect();
        assert_eq!(set.0.borrow().len(), 2);
    }

    collector.collect();
    assert_eq!(set.0.borrow().len(), 0);
}

#[test]
fn multiple_each() {
    let mut collector = Collector::new();

    let set = Arc::new(RefCellWrapper(RefCell::new(BTreeSet::new())));

    let cap = 5;
    let buf = RingBuffer::new(cap);

    assert_eq!(set.0.borrow().len(), 0);

    {
        let (mut prod, mut cons) = buf.split(&collector.handle());
        let mut id = 0;
        let mut cnt = 0;

        assert_eq!(
            prod.push_each(|| {
                if cnt < 4 {
                    id += 1;
                    cnt += 1;
                    Some(Dropper::new(Arc::clone(&set), id))
                } else {
                    None
                }
            }),
            4
        );
        assert_eq!(cnt, 4);
        collector.collect();
        assert_eq!(cnt, set.0.borrow().len());

        assert_eq!(
            cons.pop_each(
                |_| {
                    cnt -= 1;
                    true
                },
                Some(2)
            ),
            2
        );
        assert_eq!(cnt, 2);
        collector.collect();
        assert_eq!(cnt, set.0.borrow().len());

        assert_eq!(
            prod.push_each(|| {
                id += 1;
                cnt += 1;
                Some(Dropper::new(Arc::clone(&set), id))
            }),
            3
        );
        assert_eq!(cnt, 5);
        collector.collect();
        assert_eq!(cnt, set.0.borrow().len());

        assert_eq!(
            cons.pop_each(
                |_| {
                    cnt -= 1;
                    true
                },
                None
            ),
            5
        );
        assert_eq!(cnt, 0);
        collector.collect();
        assert_eq!(cnt, set.0.borrow().len());

        assert_eq!(
            prod.push_each(|| {
                id += 1;
                cnt += 1;
                Some(Dropper::new(Arc::clone(&set), id))
            }),
            5
        );
        assert_eq!(cnt, 5);
        collector.collect();
        assert_eq!(cnt, set.0.borrow().len());
    }

    collector.collect();
    assert_eq!(set.0.borrow().len(), 0);
}

/// Test the `pop_each` with internal function that returns false
#[test]
fn pop_each_test1() {
    let collector = Collector::new();

    let cap = 10usize;
    let (mut producer, mut consumer) = RingBuffer::new(cap).split(&collector.handle());

    for i in 0..cap {
        producer.push((i, vec![0u8; 1000])).unwrap();
    }

    for _ in 0..cap {
        let removed = consumer.pop_each(|_val| -> bool { false }, None);
        assert_eq!(removed, 1);
    }

    assert_eq!(consumer.len(), 0);
}

/// Test the `pop_each` with capped pop
#[test]
fn pop_each_test2() {
    let collector = Collector::new();

    let cap = 10usize;
    let (mut producer, mut consumer) = RingBuffer::new(cap).split(&collector.handle());

    for i in 0..cap {
        producer.push((i, vec![0u8; 1000])).unwrap();
    }

    for _ in 0..cap {
        let removed = consumer.pop_each(|_val| -> bool { true }, Some(1));
        assert_eq!(removed, 1);
    }

    assert_eq!(consumer.len(), 0);
}

/// Test the `push_each` with internal function that adds only 1 element.
#[test]
fn push_each_test1() {
    let collector = Collector::new();

    let cap = 10usize;
    let (mut producer, mut consumer) = RingBuffer::new(cap).split(&collector.handle());

    for i in 0..cap {
        let mut count = 0;
        // Add 1 element at a time
        let added = producer.push_each(|| -> Option<(usize, Vec<u8>)> {
            if count == 0 {
                count += 1;
                Some((i, vec![0u8; 1000]))
            } else {
                None
            }
        });
        assert_eq!(added, 1);
    }

    for _ in 0..cap {
        consumer.pop().unwrap();
    }

    assert_eq!(consumer.len(), 0);
}

/// Test the `push_each` with split internal buffer
#[test]
fn push_each_test2() {
    let collector = Collector::new();

    let cap = 10usize;
    let cap_half = 5usize;
    let (mut producer, mut consumer) = RingBuffer::new(cap).split(&collector.handle());

    // Fill the entire buffer
    for i in 0..cap {
        producer.push((i, vec![0u8; 1000])).unwrap();
    }

    // Remove half elements
    for _ in 0..cap_half {
        consumer.pop().unwrap();
    }

    // Re add half elements one by one and check the return count.
    for i in 0..cap_half {
        let mut count = 0;
        // Add 1 element at a time
        let added = producer.push_each(|| -> Option<(usize, Vec<u8>)> {
            if count == 0 {
                count += 1;
                Some((i, vec![0u8; 1000]))
            } else {
                None
            }
        });
        assert_eq!(added, 1);
    }

    for _ in 0..cap {
        consumer.pop().unwrap();
    }

    assert_eq!(consumer.len(), 0);
}