//! Simple sorted set.
//!
//! ## Rust rewrite notes
//!
//! This was originally a full-blown red-black tree, but it was really only used
//! to keep a sorted list. It has been rewritten as a sorted `Vec`. We
//! technically no longer have O(log(n)) insertion time. Given set sizes and
//! actual execution time, that shouldn't matter.
//!
//! Prefer a proper Rust collection over this for new code.

use std::{
    ffi::{c_int, c_void},
    ptr::NonNull,
};

#[derive(Default)]
pub struct Bag {
    inner: Vec<(c_int, NonNull<c_void>)>,
}

/// Sentinel iterator value. In C, just check for negative values.
pub const NOT_FOUND: c_int = -1;

impl Bag {
    /// Returns `Ok(index)` if an item keyed by `key` is in `self`, else
    /// `Err(index)` if an item keyed by `key` could be inserted at `index`.
    fn search(&self, key: c_int) -> Result<usize, usize> {
        self.inner.binary_search_by_key(&key, |(key, _)| *key)
    }

    /// Sets a value associated with `key` to `value`. Clobbers any extant value.
    fn set(&mut self, key: c_int, value: NonNull<c_void>) {
        match self.search(key) {
            Ok(i) => self.inner[i] = (key, value),
            Err(i) => self.inner.insert(i, (key, value)),
        }
    }
}

pub mod ffi {
    use super::{Bag, NOT_FOUND};

    use std::{
        ffi::{c_int, c_void},
        ptr::{self, NonNull},
    };

    /// Basic constructor. Free with [`WMFreeBag`].
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMCreateTreeBag() -> *mut Bag {
        Box::leak(Box::new(Bag::default()))
    }

    /// Retrieves the value associated with `key`, or null.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMGetFromBag(bag: *mut Bag, key: c_int) -> *mut c_void {
        if bag.is_null() {
            return ptr::null_mut();
        }
        let bag = unsafe { &mut *bag };
        bag
            .search(key)
            .ok()
            .and_then(|index| bag.inner.get(index).copied())
            .map(|(_, value)| value.as_ptr())
            .unwrap_or(ptr::null_mut())
    }

    /// Sets the value associated with `key` to `item`. If `item` is null,
    /// removes any extant value.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMSetInBag(bag: *mut Bag, key: c_int, item: *mut c_void) {
        if bag.is_null() {
            return;
        }
        let bag = unsafe { &mut *bag };

        if let Some(item) = NonNull::new(item) {
            bag.set(key, item);
        } else {
            if let Some(i) = bag.search(key).ok() {
                bag.inner.remove(i);
            }
        }
    }

    /// Clears the contents of `bag`.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMEmptyBag(bag: *mut Bag) {
        if bag.is_null() {
            return;
        }
        unsafe {
            (*bag).inner.clear();
        }
    }

    /// Deletes `bag`.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMFreeBag(bag: *mut Bag) {
        if bag.is_null() {
            return;
        }
        unsafe {
            let _ = Box::from_raw(bag);
        }
    }

    /// Initializes `ptr` to the first element of `bag` and returns that
    /// element. Sets `ptr` to [`NOT_FOUND`] and returns null if `bag` is empty.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMBagFirst(bag: *mut Bag, ptr: *mut c_int) -> *mut c_void {
        if ptr.is_null() {
            return ptr::null_mut();
        }
        let ptr = unsafe { &mut *ptr };

        if bag.is_null() {
            *ptr = NOT_FOUND;
            return ptr::null_mut();
        }
        let bag = unsafe { &mut *bag };

        if let Some((_, value)) = bag.inner.first() {
            *ptr = 0;
            value.as_ptr()
        } else {
            *ptr = NOT_FOUND;
            ptr::null_mut()
        }
    }

    /// Initializes `ptr` to the last element of `bag` and returns that
    /// element. Sets `ptr` to [`NOT_FOUND`] and returns null if `bag` is empty.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMBagLast(bag: *mut Bag, ptr: *mut c_int) -> *mut c_void {
        if ptr.is_null() {
            return ptr::null_mut();
        }
        let ptr = unsafe { &mut *ptr };

        if bag.is_null() {
            *ptr = NOT_FOUND;
            return ptr::null_mut();
        }
        let bag = unsafe { &mut *bag };

        if let Some((_, value)) = bag.inner.last() {
            *ptr = (bag.inner.len() - 1) as c_int;
            value.as_ptr()
        } else {
            *ptr = NOT_FOUND;
            ptr::null_mut()
        }
    }

    /// Advances `ptr` to the next element of `bag` and returns it. Sets `ptr`
    /// to [`NOT_FOUND`] and returns null if `ptr` is already at the end of
    /// `bag`.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMBagNext(bag: *mut Bag, ptr: *mut c_int) -> *mut c_void {
        if ptr.is_null() {
            return ptr::null_mut();
        }
        let ptr = unsafe { &mut *ptr };

        if bag.is_null() {
            *ptr = NOT_FOUND;
            return ptr::null_mut();
        }
        let bag = unsafe { &mut *bag };

        if *ptr < 0 {
            return ptr::null_mut();
        }

        *ptr += 1;
        if let Some((_, value)) = bag.inner.get(*ptr as usize) {
            value.as_ptr()
        } else {
            *ptr = NOT_FOUND;
            ptr::null_mut()
        }
    }

    /// Decrements `ptr` to the previous element of `bag` and returns it. Sets `ptr`
    /// to [`NOT_FOUND`] and returns null if `ptr` is already at the beginning of
    /// `bag`.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMBagPrevious(bag: *mut Bag, ptr: *mut c_int) -> *mut c_void {
        if ptr.is_null() {
            return ptr::null_mut();
        }
        let ptr = unsafe { &mut *ptr };

        if bag.is_null() {
            *ptr = NOT_FOUND;
            return ptr::null_mut();
        }
        let bag = unsafe { &mut *bag };

        if *ptr <= 0 {
            *ptr = NOT_FOUND;
            return ptr::null_mut();
        }
        *ptr -= 1;
        if let Some((_, value)) = bag.inner.get(*ptr as usize) {
            value.as_ptr()
        } else {
            *ptr = NOT_FOUND;
            ptr::null_mut()
        }
    }

    /// Sets `ptr` to the element of `bag` with `key` and returns the associated
    /// value. Sets `ptr` to [`NOT_FOUND`] and returns null if there is no such
    /// element.
    #[unsafe(no_mangle)]
    pub unsafe extern "C" fn WMBagIteratorAtIndex(
        bag: *mut Bag,
        key: c_int,
        ptr: *mut c_int,
    ) -> *mut c_void {
        if ptr.is_null() {
            return ptr::null_mut();
        }
        let ptr = unsafe { &mut *ptr };

        if bag.is_null() {
            *ptr = NOT_FOUND;
            return ptr::null_mut();
        }
        let bag = unsafe { &mut *bag };

        if key < 0 {
            *ptr = NOT_FOUND;
            return ptr::null_mut();
        }

        if let Some(index) = bag.search(key).ok() {
            *ptr = index as c_int;
            bag.inner[index].1.as_ptr()
        } else {
            *ptr = NOT_FOUND;
            ptr::null_mut()
        }
    }
}

#[cfg(test)]
mod test {
    use super::{ffi, Bag, NOT_FOUND};

    use std::{
        ffi::{c_int, c_void},
        ptr::NonNull,
    };

    fn void_of<T>(t: &mut T) -> NonNull<c_void> {
        NonNull::new(t as *mut _ as *mut c_void).unwrap()
    }

    #[test]
    fn insert_out_of_order() {
        let mut b = Bag::default();
        let mut x = 3u16;
        let mut y = 4u16;
        let mut z = 5u16;
        b.set(7, void_of(&mut z));
        b.set(3, void_of(&mut y));
        b.set(5, void_of(&mut x));

        assert_eq!(
            b.inner,
            &[
                (3, void_of(&mut y)),
                (5, void_of(&mut x)),
                (7, void_of(&mut z))
            ]
        );
    }

    #[test]
    fn iterate_forward() {
        let mut b = Bag::default();
        let mut x = 3u16;
        let mut y = 4u16;
        let mut z = 5u16;
        b.set(7, void_of(&mut z));
        b.set(3, void_of(&mut y));
        b.set(5, void_of(&mut x));

        let mut iterator: c_int = NOT_FOUND;
        unsafe {
            assert_eq!(
                ffi::WMBagFirst(&mut b, &mut iterator),
                void_of(&mut y).as_ptr()
            );
            assert_eq!(iterator, 0);
            assert_eq!(
                ffi::WMBagNext(&mut b, &mut iterator),
                void_of(&mut x).as_ptr()
            );
            assert_eq!(iterator, 1);
            assert_eq!(
                ffi::WMBagNext(&mut b, &mut iterator),
                void_of(&mut z).as_ptr()
            );
            assert_eq!(iterator, 2);
            assert!(ffi::WMBagNext(&mut b, &mut iterator).is_null());
            assert_eq!(iterator, NOT_FOUND);
        }
    }

    #[test]
    fn iterate_backward() {
        let mut b = Bag::default();
        let mut x = 3u16;
        let mut y = 4u16;
        let mut z = 5u16;
        b.set(7, void_of(&mut z));
        b.set(3, void_of(&mut y));
        b.set(5, void_of(&mut x));

        let mut iterator: c_int = NOT_FOUND;
        unsafe {
            assert_eq!(
                ffi::WMBagLast(&mut b, &mut iterator),
                void_of(&mut z).as_ptr()
            );
            assert_eq!(iterator, 2);
            assert_eq!(
                ffi::WMBagPrevious(&mut b, &mut iterator),
                void_of(&mut x).as_ptr()
            );
            assert_eq!(iterator, 1);
            assert_eq!(
                ffi::WMBagPrevious(&mut b, &mut iterator),
                void_of(&mut y).as_ptr()
            );
            assert_eq!(iterator, 0);
            assert!(ffi::WMBagPrevious(&mut b, &mut iterator).is_null());
            assert_eq!(iterator, NOT_FOUND);
        }
    }

    #[test]
    fn find_iterate() {
        let mut b = Bag::default();
        let mut x = 3u16;
        let mut y = 4u16;
        let mut z = 5u16;
        b.set(7, void_of(&mut z));
        b.set(3, void_of(&mut y));
        b.set(5, void_of(&mut x));

        let mut iterator: c_int = NOT_FOUND;
        unsafe {
            assert_eq!(
                ffi::WMBagIteratorAtIndex(&mut b, 5, &mut iterator),
                void_of(&mut x).as_ptr()
            );
            assert_eq!(iterator, 1);
            assert_eq!(
                ffi::WMBagNext(&mut b, &mut iterator),
                void_of(&mut z).as_ptr()
            );
            assert_eq!(iterator, 2);
        }
    }
}
