// MIT/Apache2 License

//! A stack-based, fixed-size map that puts nodes on an array rather than the heap. See documentation
//! for `TinyMap` for more information.

#![allow(incomplete_features)]
#![feature(const_generics)]
#![feature(const_in_array_repeat_expressions)]

#![warn(clippy::pedantic)]
#![allow(clippy::redundant_pattern_matching)] // I'm trying to reduce the amount of LLVM IR output

#![no_std]

use core::{cmp::Ordering, fmt, iter, mem};
use tinyvec::{ArrayVec, ArrayVecIterator};

// A node in the binary tree making up the map.
struct Node<K, V> {
    kv: (K, V),
    children: [Option<usize>; 2],
}

impl<K: Clone, V: Clone> Clone for Node<K, V> {
    #[inline]
    fn clone(&self) -> Self {
        Node {
            kv: self.kv.clone(),
            children: self.children,
        }
    }
}

impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for Node<K, V> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.kv, f)
    }
}

/// A binary tree that uses a stack-based array as storage for nodes.
///
/// Rather than using the heap to store nodes, this structure uses an `ArrayVec`. Note that this
/// is nowhere near as efficient as the standard library `HashMap` implementation. The purpose of
/// this structure is to offer an interface similar to `HashMap` to use in its place in `no_std` environments.
///
/// # Example
///
/// ```rust
/// use tinymap::TinyMap;
///
/// // Today, I'm making software for the lamp store. They want to match up a LampID to a
/// // Lamp. This software also needs to run on embedded hardware, where a Rust allocator
/// // hasn't been ported yet.
///
/// /// A representation of a lamp.
/// struct Lamp { wattage: u32, color: &'static str }
///
/// // The LampID is just a string that the manager uses to look up a lamp
/// // that the store has in stock.
/// let mut lamps: TinyMap<&'static str, Lamp, 3> = TinyMap::new();
///
/// // Let's take some of the lamps we have and put them in our map.
/// lamps.insert("32ZD", Lamp { wattage: 120, color: "red" });
/// lamps.insert("11HR", Lamp { wattage: 60, color: "gray" });
/// lamps.insert("2460", Lamp { wattage: 90, color: "blue" });
///
/// // The customer wants lamps #32ZD and #11HR?
/// assert_eq!(lamps.get(&"32ZD").unwrap().wattage, 120);
/// assert_eq!(lamps.get(&"11HR").unwrap().color, "gray");
///
/// // Let's add one more lamp! Nothing could go wrong!
/// let oops = lamps.try_insert("199K", Lamp { wattage: 80, color: "green" });
/// assert!(oops.is_err());
/// assert!(lamps.get(&"199k").is_none());
/// ```
pub struct TinyMap<K: PartialOrd, V, const N: usize> {
    arena: ArrayVec<[Option<Node<K, V>>; N]>,
    root: Option<usize>,
}

macro_rules! unwrap_tpn {
    ($self: expr, $e: expr) => {{
        match $self.try_push_node($e) {
            Ok(u) => Some(u),
            Err(n) => {
                let Node { kv, .. } = n;
                return Err(kv);
            }
        }
    }};
}

enum ChildType {
    Left,
    Right,
}

impl<K: PartialOrd, V, const N: usize> TinyMap<K, V, N> {
    /// Create a new [`TinyMap`].
    ///
    /// # Example
    ///
    /// ```rust
    /// # use tinymap::TinyMap;
    /// #[derive(PartialOrd, Ord, PartialEq, Eq)]
    /// struct Foo;
    /// struct Bar;
    ///
    /// let my_map: TinyMap<Foo, Bar, 1> = TinyMap::new();
    /// ```
    #[must_use]
    #[inline]
    pub fn new() -> Self {
        Self {
            arena: ArrayVec::from_array_len([None; N], 0),
            root: None,
        }
    }

    /// Retrieves the current number of elements within the binary tree.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use tinymap::TinyMap;
    /// #[derive(PartialOrd, Ord, PartialEq, Eq)]
    /// struct MyIndex(u32);
    /// struct MyData { foo: &'static str }
    ///
    /// let mut my_map: TinyMap<MyIndex, MyData, 3> = TinyMap::new();
    /// my_map.insert(MyIndex(12), MyData { foo: "Leroy" });
    /// my_map.insert(MyIndex(13), MyData { foo: "Barbara" });
    /// my_map.insert(MyIndex(12), MyData { foo: "Joanne" });
    /// assert_eq!(my_map.len(), 2);
    /// ```
    pub fn len(&self) -> usize {
        self.arena
            .iter()
            .filter(|n| match n {
                Some(_) => true,
                None => false,
            })
            .count()
    }

    /// Check to see if there are any elements in this map.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use tinymap::TinyMap;
    /// let my_love_life: TinyMap<&'static str, u32, 10> = TinyMap::new();
    /// assert!(my_love_life.is_empty());
    /// ```
    pub fn is_empty(&self) -> bool {
        match self.root {
            Some(_) => false,
            None => true,
        }
    }

    #[must_use]
    #[inline]
    fn node_at(&self, index: usize) -> Option<&Node<K, V>> {
        self.arena[index].as_ref()
    }

    #[must_use]
    #[inline]
    fn node_at_mut(&mut self, index: usize) -> Option<&mut Node<K, V>> {
        self.arena[index].as_mut()
    }

    #[must_use]
    #[inline]
    fn root(&self) -> Option<usize> {
        self.root
    }

    #[must_use]
    #[inline]
    fn node_by_key(&self, key: &K) -> Option<&Node<K, V>> {
        let mut current = self.root();
        loop {
            match current {
                None => return None,
                Some(val) => {
                    let node = self.node_at(val).expect("Invalid node tree");
                    current = match node.kv.0.partial_cmp(key) {
                        None => return None,
                        Some(Ordering::Equal) => return Some(node),
                        Some(Ordering::Less) => node.children[0],
                        Some(Ordering::Greater) => node.children[1],
                    };
                }
            }
        }
    }

    #[must_use]
    #[inline]
    fn node_by_key_mut(&mut self, key: &K) -> Option<&mut Node<K, V>> {
        let mut current = self.root();
        loop {
            match current {
                None => return None,
                Some(val) => {
                    let node = self.node_at(val).expect("Invalid node tree");
                    current = match node.kv.0.partial_cmp(key) {
                        None => return None,
                        Some(Ordering::Equal) => {
                            return Some(self.node_at_mut(val).expect("Invalid node tree"));
                        }
                        Some(Ordering::Less) => node.children[0],
                        Some(Ordering::Greater) => node.children[1],
                    };
                }
            }
        }
    }

    /// Get a reference to an item stored inside of this map.
    #[must_use]
    #[inline]
    pub fn get(&self, key: &K) -> Option<&V> {
        match self.node_by_key(key) {
            Some(node) => Some(&node.kv.1),
            None => None,
        }
    }

    /// Get a mutable reference to an item stored inside of this map.
    #[must_use]
    #[inline]
    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        match self.node_by_key_mut(key) {
            Some(node) => Some(&mut node.kv.1),
            None => None,
        }
    }

    /// Tell whether or not this map contains a certain key.
    #[must_use]
    #[inline]
    pub fn contains_key(&self, key: &K) -> bool {
        match self.node_by_key(key) {
            Some(_) => true,
            None => false,
        }
    }

    // push a node into the arena and get its index back
    #[inline]
    fn try_push_node(&mut self, node: Node<K, V>) -> Result<usize, Node<K, V>> {
        match self.arena.try_push(Some(node)) {
            Some(mut node) => {
                // look for a node to replace
                let pos = self.arena.iter().position(Option::is_none);
                match pos {
                    Some(pos) => {
                        mem::swap(&mut self.arena[pos], &mut node);
                        Ok(pos)
                    }
                    None => Err(node.unwrap()),
                }
            }
            None => Ok(self.arena.len() - 1),
        }
    }

    #[inline]
    fn insert_from_root(
        &mut self,
        mut node: Node<K, V>,
        mut current: usize,
    ) -> Result<Option<V>, (K, V)>
    where
        K: Ord,
    {
        let mut next;
        loop {
            let cmp_node = self.node_at(current).unwrap();
            let ct = match cmp_node.kv.0.cmp(&node.kv.0) {
                Ordering::Less => {
                    next = cmp_node.children[0];
                    ChildType::Left
                }
                Ordering::Greater => {
                    next = cmp_node.children[1];
                    ChildType::Right
                }
                Ordering::Equal => {
                    mem::swap(&mut self.node_at_mut(current).unwrap().kv.1, &mut node.kv.1);
                    return Ok(Some(node.kv.1));
                }
            };

            match next {
                None => {
                    let index = unwrap_tpn!(self, node);
                    let cmp_node_mut = self.node_at_mut(current).unwrap();
                    let slot = match ct {
                        ChildType::Left => &mut cmp_node_mut.children[0],
                        ChildType::Right => &mut cmp_node_mut.children[1],
                    };
                    *slot = index;

                    return Ok(None);
                }
                Some(next) => {
                    current = next;
                }
            }
        }
    }

    /// Insert a node into this binary tree. This function will return the value previously
    /// in the key's slot, if applicable.
    ///
    /// # Errors
    ///
    /// If the backing array is full, this function will return back the key-value pair passed
    /// in.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use tinymap::TinyMap;
    /// /// The representation of a mattress.
    /// #[derive(PartialOrd, Ord, PartialEq, Eq)]
    /// struct Mattress { width: u32, height: u32, brand: &'static str }
    ///
    /// let mut mattress_ratings: TinyMap<Mattress, u64, 12> = TinyMap::new();
    /// let mut limit = 0;
    /// loop {
    ///     // The mattress testing robot tests more mattresses than we could ever imagine.
    ///     // But how long until it tests too many?
    ///     let mattress = Mattress {
    ///         width: limit + 1,
    ///         height: (limit + 1) * 2,
    ///         brand: "Smith Mattresses",
    ///     };
    ///     let res = mattress_ratings.try_insert(mattress, 27);
    ///     match res {
    ///         Ok(_) => { limit += 1; },
    ///         Err((reject, _)) => {
    ///             assert_eq!(reject.width, 13);
    ///             break;
    ///         }
    ///     }
    /// }
    ///
    /// assert_eq!(limit, 12);
    /// ```
    #[inline]
    pub fn try_insert(&mut self, key: K, value: V) -> Result<Option<V>, (K, V)>
    where
        K: Ord,
    {
        let node = Node {
            kv: (key, value),
            children: [None; 2],
        };

        match self.root {
            None => {
                // if the root doesn't exist, this is an empty map
                self.root = unwrap_tpn!(self, node);
                Ok(None)
            }
            Some(root) => self.insert_from_root(node, root),
        }
    }

    /// Insert a node into this binary tree.
    ///
    /// # Panics
    ///
    /// Unlike `try_insert`, this function will panic if the backing array is full.
    ///
    /// # Example
    ///
    /// ```rust, should_panic
    /// # use tinymap::TinyMap;
    /// /// Representation of a viking.
    /// struct Viking { kill_count: u32, body_count: u32, name: &'static str }
    /// let mut famous_vikings: TinyMap<u32, Viking, 25> = TinyMap::new();
    ///
    /// for i in 0..100 {
    ///    // D'oh! I knew I shouldn't have run my school assignment on my microwave!
    ///    famous_vikings.insert(i, Viking { kill_count: i, body_count: i * 2, name: "John" });
    /// }
    /// ```
    #[inline]
    pub fn insert(&mut self, key: K, value: V) -> Option<V>
    where
        K: Ord,
    {
        match self.try_insert(key, value) {
            Err(_) => panic!("Unable to push node into binary tree"),
            Ok(res) => res,
        }
    }

    /// Remove a node entry from the binary tree. This function returns the key-value pair
    /// that was removed.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use tinymap::TinyMap;
    /// fn help_humans() { println!("Let's solve world hunger!"); }
    /// fn kill_humans() { panic!("Kill! Kill! Kill!"); }
    ///
    /// // My robot needs to know what to do!
    /// let mut robot_behaviors: TinyMap<&'static str, &dyn Fn(), 2> = TinyMap::new();
    /// robot_behaviors.insert("help", &help_humans);
    /// robot_behaviors.insert("kill", &kill_humans);
    ///
    /// // ...wait. It probably shouldn't do that.
    /// let removed = robot_behaviors.remove_entry(&"kill").unwrap();
    ///
    /// assert_eq!(removed.0, "kill");
    /// ```
    #[inline]
    pub fn remove_entry(&mut self, key: &K) -> Option<(K, V)>
    where
        K: Ord,
    {
        const ERR_MSG: &str = "Invalid binary tree path";

        enum ParentChildRelation {
            ChildIsRoot,
            Left,
            Right,
        }

        let mut parent_index: Option<usize> = None;
        let mut last_relationship = ParentChildRelation::ChildIsRoot;
        let mut current: Option<usize> = self.root;
        loop {
            let c = match current {
                Some(c) => c,
                None => return None,
            };
            let cmp_node = self.node_at(c).expect(ERR_MSG);
            match cmp_node.kv.0.partial_cmp(key) {
                None => return None,
                Some(Ordering::Equal) => {
                    // create a node to replace the node in the tree
                    let mut replacement_node = match cmp_node.children {
                        [None, None] => None,
                        [Some(child), None] | [None, Some(child)] => Some(child),
                        [Some(child1), Some(child2)] => {
                            // take the node out of child2
                            let mut reloc_node = None;
                            mem::swap(&mut reloc_node, &mut self.arena[child2]);

                            // insert child2 under child1
                            if let Err(_) = self.insert_from_root(reloc_node.unwrap(), child1) {
                                panic!("This should not happen!");
                            }

                            // use child1 as the replacement node
                            Some(child1)
                        }
                    };

                    let slot = match (last_relationship, parent_index) {
                        (ParentChildRelation::ChildIsRoot, _) => &mut self.root,
                        (ParentChildRelation::Left, Some(p)) => {
                            &mut self.node_at_mut(p).expect(ERR_MSG).children[0]
                        }
                        (ParentChildRelation::Right, Some(p)) => {
                            &mut self.node_at_mut(p).expect(ERR_MSG).children[1]
                        }
                        _ => unreachable!(),
                    };

                    mem::swap(slot, &mut replacement_node);
                    let n = self.arena[replacement_node.unwrap()].take();
                    return Some(n.unwrap().kv);
                }
                Some(Ordering::Less) => {
                    parent_index = current;
                    current = cmp_node.children[0];
                    last_relationship = ParentChildRelation::Left;
                }
                Some(Ordering::Greater) => {
                    parent_index = current;
                    current = cmp_node.children[1];
                    last_relationship = ParentChildRelation::Right;
                }
            }
        }
    }

    /// Remove a value from the binary tree. This is similar to `remove_entry`, but it does not
    /// keep the value.
    #[inline]
    pub fn remove(&mut self, key: &K) -> Option<V>
    where
        K: Ord,
    {
        match self.remove_entry(key) {
            Some((_k, v)) => Some(v),
            None => None,
        }
    }

    /// Clear out the binary tree.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use tinymap::TinyMap;
    /// struct Album { name: &'static str }
    ///
    /// let mut albums_by_artist: TinyMap<&'static str, Album, 5> = TinyMap::new();
    ///
    /// albums_by_artist.insert("Owl City", Album { name: "Ocean Lights" });
    /// albums_by_artist.insert("Billy Joel", Album { name: "Glass Houses" });
    /// albums_by_artist.insert("Frank Sinatra", Album { name: "My Way" });
    ///
    /// // You know what? I've decided I'm not really that into music anymore.
    /// albums_by_artist.clear();
    ///
    /// assert_eq!(albums_by_artist.len(), 0);
    /// ```
    #[inline]
    pub fn clear(&mut self) {
        self.arena.clear();
        self.root = None;
    }

    /// Iterate over the elements of this binary tree in arbitrary order.
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = (&K, &V)> {
        self.arena.iter().filter_map(|node| match node.as_ref() {
            None => None,
            Some(Node {
                kv: (ref k, ref v), ..
            }) => Some((k, v)),
        })
    }

    /// Iterate over the elements of this binary tree in arbitrary order, mutably.
    #[inline]
    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&K, &mut V)> {
        self.arena
            .iter_mut()
            .filter_map(|node| match node.as_mut() {
                None => None,
                Some(Node {
                    kv: (ref k, ref mut v),
                    ..
                }) => Some((k, v)),
            })
    }

    /// Iterate over the keys of this binary tree in arbitrary order.
    #[inline]
    pub fn keys(&self) -> impl Iterator<Item = &K> {
        self.arena.iter().filter_map(|node| match node.as_ref() {
            None => None,
            Some(Node { kv: (ref k, _), .. }) => Some(k),
        })
    }

    /// Iterate over the values of this binary tree in arbitrary order.
    #[inline]
    pub fn values(&self) -> impl Iterator<Item = &V> {
        self.arena.iter().filter_map(|node| match node.as_ref() {
            None => None,
            Some(Node { kv: (_, ref v), .. }) => Some(v),
        })
    }

    /// Iterate over the values of this binary tree in arbitrary order, mutably.
    #[inline]
    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut V> {
        self.arena
            .iter_mut()
            .filter_map(|node| match node.as_mut() {
                None => None,
                Some(Node {
                    kv: (_, ref mut v), ..
                }) => Some(v),
            })
    }
}

/// A consuming iterator for instances of `TinyMap`.
#[repr(transparent)]
pub struct TinyMapIterator<K, V, const N: usize> {
    inner: ArrayVecIterator<[Option<Node<K, V>>; N]>,
}

impl<K, V, const N: usize> Iterator for TinyMapIterator<K, V, N> {
    type Item = (K, V);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.find_map(|node| match node {
            None => None,
            Some(node) => Some(node.kv),
        })
    }

    #[inline]
    #[must_use]
    fn size_hint(&self) -> (usize, Option<usize>) {
        (0, self.inner.size_hint().1)
    }
}

impl<K: PartialOrd, V, const N: usize> Default for TinyMap<K, V, N> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl<K: PartialOrd + Clone, V: Clone, const N: usize> Clone for TinyMap<K, V, N> {
    #[inline]
    fn clone(&self) -> Self {
        TinyMap {
            arena: self.arena.clone(),
            root: self.root,
        }
    }
}

impl<K: PartialOrd + fmt::Debug, V: fmt::Debug, const N: usize> fmt::Debug for TinyMap<K, V, N> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.arena, f)
    }
}

impl<K: PartialOrd, V, const N: usize> iter::IntoIterator for TinyMap<K, V, N> {
    type Item = (K, V);
    type IntoIter = TinyMapIterator<K, V, N>;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        TinyMapIterator {
            inner: self.arena.into_iter(),
        }
    }
}

impl<K: Ord, V, const N: usize> iter::Extend<(K, V)> for TinyMap<K, V, N> {
    #[inline]
    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
        iter.into_iter().for_each(|(k, v)| {
            self.insert(k, v);
        });
    }
}

impl<K: Ord, V, const N: usize> iter::FromIterator<(K, V)> for TinyMap<K, V, N> {
    #[inline]
    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
        let mut map = TinyMap::new();
        map.extend(iter);
        map
    }
}

#[test]
fn test_remove() {
    let mut map: TinyMap<u32, u32, 25> = TinyMap::new();
    for i in 0..25 {
        map.insert(i, i);
    }

    for j in 10..15 {
        map.remove(&j);
    }

    let _test = map.get(&16).unwrap();
}