ids 0.1.4

//! A hash map.

use core::{
    fmt::Debug,
    hash::{BuildHasher, Hash, Hasher},
    mem::replace,
    ops::{Deref, DerefMut},
    ptr::NonNull,
};

use alloc::{boxed::Box, sync::Arc, vec, vec::Vec};

use crate::cowarc::Cowarc;

/// The branch factor used for [`DefaultHashMap`].
pub const DEFAULT_BRANCH_FACTOR: usize = 8;

/// A [`HashMap`] with branch factor [`DEFAULT_BRANCH_FACTOR`].
pub type DefaultHashMap<K, V, H> = HashMap<K, V, H, DEFAULT_BRANCH_FACTOR>;

/// A hash map from keys in `K` to values in `V`, with a hasher `H` and a branch
/// factor `F`.
///
/// The hash map is implemented by a tree into which we index by successive
/// parts of the hash. If a node at depth `d` is an internal node, then a key
/// with hash `h` belongs under index `floor(h / F ** d) % d` in its array.
#[derive(Clone)]
pub struct HashMap<K: Eq + Hash, V, H: BuildHasher, const F: usize> {
    root: Pointer<K, V, F>,
    build_hasher: H,
}

impl<K: Eq + Hash, V, H: BuildHasher, const F: usize> HashMap<K, V, H, F> {
    /// Creates a new hash map.
    pub const fn new(build_hasher: H) -> Self {
        Self {
            root: None,
            build_hasher,
        }
    }

    /// Counts the number of entries in the hash map.
    pub fn len(&self) -> usize {
        fn count_rec<K, V, const F: usize>(pointer: &Pointer<K, V, F>) -> usize {
            match pointer.as_ref().map(Arc::as_ref) {
                Some(Node::SubTrie { entries }) => entries.iter().map(count_rec).sum(),
                Some(Node::Leaf { .. }) => 1,
                None => 0,
            }
        }

        count_rec(&self.root)
    }

    /// Iterates across all key-value pairs in the hash map.
    ///
    /// The iteration order is by hash value of keys.
    pub fn iter(&self) -> Iter<K, V, F> {
        Iter {
            stack: if let Some(root) = self.root.as_ref() {
                vec![(&**root, 0)]
            } else {
                Vec::new()
            },
        }
    }
}

impl<K: Eq + Hash + Clone, V: Clone, H: BuildHasher, const F: usize> HashMap<K, V, H, F> {
    /// Returns an [`Entry`] corresponding to `key`, which either indicates that
    /// there is no value associated to that key and allows for one to be
    /// relatively quickly inserted, or provides access to the value already
    /// present.
    ///
    /// This internally uses [`Cowarc`], which leads to on average `O(log(N))`
    /// allocations for keeping track of nodes that may need to be duplicated.
    /// It may be further optimized in the future.
    pub fn entry(&mut self, key: K) -> Entry<K, V, F> {
        let mut hasher = self.build_hasher.build_hasher();
        key.hash(&mut hasher);
        let mut hash = hasher.finish();
        let mut pointer = EntryPointer::Root(&mut self.root);

        while let Some(node) = pointer.get() {
            match node {
                Node::SubTrie { .. } => {
                    pointer = EntryPointer::Child {
                        node: pointer.into_cowarc(),
                        index: (hash % F as u64) as usize,
                    };
                    hash /= F as u64;
                }
                Node::Leaf {
                    partial_hash,
                    key: other_key,
                    ..
                } => {
                    if *other_key == key {
                        return Entry::Occupied(OccupiedEntry {
                            node: pointer.into_cowarc(),
                        });
                    } else if *partial_hash == hash {
                        panic!("hash collision")
                    } else {
                        break;
                    }
                }
            }
        }

        Entry::Vacant(VacantEntry {
            key,
            pointer,
            partial_hash: hash,
        })
    }

    /// Inserts a new value into the hash map, returning a mutable reference to
    /// it as well as the previous value associated with the given `key`, if
    /// there was one.
    ///
    /// This is currently not very optimized and just wraps a call to
    /// [`Self::entry`].
    pub fn insert(&mut self, key: K, value: V) -> (&mut V, Option<V>) {
        match self.entry(key) {
            Entry::Vacant(vacant_entry) => (vacant_entry.insert(value), None),
            Entry::Occupied(occupied_entry) => {
                let entry = OccupiedEntry::into_mut(occupied_entry);
                let old_value = replace(entry, value);
                (entry, Some(old_value))
            }
        }
    }

    /// Gets the value associated with a `key`, if it exists.
    pub fn get(&mut self, key: K) -> Option<&V> {
        match self.entry(key) {
            Entry::Vacant(_) => None,
            Entry::Occupied(entry) => Some(OccupiedEntry::into_ref(entry)),
        }
    }

    /// Removes and returns the value associated with a `key`, if it exists.
    ///
    /// This is currently optimized for the case where the key exists; it may
    /// unnecessarily duplicate internal nodes if it does not. This is only a
    /// performance limitation and does not affect correctness.
    pub fn remove(&mut self, key: &K) -> Option<V> {
        fn remove_rec<K: Eq + Clone, V: Clone, const F: usize>(
            pointer: &mut Pointer<K, V, F>,
            key: &K,
            partial_hash: u64,
        ) -> Option<V> {
            if let Some(node) = pointer {
                match Arc::make_mut(node) {
                    Node::SubTrie { entries } => {
                        let result = remove_rec(
                            &mut entries[(partial_hash % F as u64) as usize],
                            key,
                            partial_hash / F as u64,
                        );
                        if result.is_some() && entries.iter().all(|ptr| ptr.is_none()) {
                            *pointer = None;
                        }
                        result
                    }
                    Node::Leaf { key: leaf_key, .. } => {
                        if leaf_key == key {
                            let Node::Leaf { value,.. } = Arc::unwrap_or_clone(pointer.take().unwrap()) else {
                                unreachable!()
                            };
                            Some(value)
                        } else {
                            None
                        }
                    }
                }
            } else {
                None
            }
        }

        let mut hasher = self.build_hasher.build_hasher();
        key.hash(&mut hasher);
        let hash = hasher.finish();

        remove_rec(&mut self.root, key, hash)
    }
}

impl<K: Eq + Hash + Debug, V: Debug, H: BuildHasher, const F: usize> Debug for HashMap<K, V, H, F> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

type Pointer<K, V, const F: usize> = Option<Arc<Node<K, V, F>>>;

/// An entry in a [`HashMap`] associated with a specific key.
pub enum Entry<'a, K, V, const F: usize> {
    /// The key is not present in the hash map.
    Vacant(VacantEntry<'a, K, V, F>),
    /// The key is present in the hash map.
    Occupied(OccupiedEntry<'a, K, V, F>),
}

/// An indication that a given key is not present in a hash map, with the
/// ability to insert it.
pub struct VacantEntry<'a, K, V, const F: usize> {
    key: K,
    pointer: EntryPointer<'a, K, V, F>,
    partial_hash: u64,
}

impl<'a, K, V, const F: usize> VacantEntry<'a, K, V, F> {
    /// The key associated with the entry.
    pub fn key(&self) -> &K {
        &self.key
    }

    /// Drops the entry and its mutable borrow of the hash map, returning the
    /// key.
    pub fn into_key(self) -> K {
        self.key
    }

    /// Inserts a value into the hash map at the key associated with this entry,
    /// returning a mutable reference to the value.
    ///
    /// This is just shorthand for [`Self::insert_entry`] followed by
    /// [`OccupiedEntry::into_mut`].
    pub fn insert(self, value: V) -> &'a mut V
    where
        K: Clone,
        V: Clone,
    {
        OccupiedEntry::into_mut(self.insert_entry(value))
    }

    /// Inserts a value into the hash map at the key associated with this entry,
    /// returning an [`OccupiedEntry`] referring to the value.
    pub fn insert_entry(mut self, value: V) -> OccupiedEntry<'a, K, V, F>
    where
        K: Clone,
        V: Clone,
    {
        loop {
            // This is a mutable reference to the Pointer to the node we are focusing on.
            let pointer = match &mut self.pointer {
                EntryPointer::Root(pointer) => pointer,
                EntryPointer::Child { node, index } => {
                    let Node::SubTrie { entries } = &mut **node else {
                        unreachable!()
                    };
                    &mut entries[*index]
                }
            };

            if let Some(pointer) = pointer {
                let node = Arc::make_mut(pointer);
                let mut leaf = replace(
                    node,
                    Node::SubTrie {
                        entries: [const { None }; F],
                    },
                );
                let Node::SubTrie { entries } = node else { unreachable!() };

                let index = match &mut leaf {
                    Node::SubTrie { .. } => unreachable!(),
                    Node::Leaf { partial_hash, .. } => {
                        let index = (*partial_hash % F as u64) as usize;
                        *partial_hash /= F as u64;
                        index
                    }
                };
                entries[index] = Some(Arc::new(leaf));

                self.pointer = EntryPointer::Child {
                    node: self.pointer.into_cowarc(),
                    index: (self.partial_hash % F as u64) as usize,
                };
                self.partial_hash /= F as u64;
            } else {
                *pointer = Some(Arc::new(Node::Leaf {
                    partial_hash: self.partial_hash,
                    key: self.key,
                    value,
                }));
                return OccupiedEntry {
                    node: self.pointer.into_cowarc(),
                };
            }
        }
    }
}

/// An indication that a given key is present in a hash map, providing access to
/// read or (copy-on-write) modify it.
pub struct OccupiedEntry<'a, K, V, const F: usize> {
    node: Cowarc<'a, Node<K, V, F>>,
}

impl<'a, K: Clone, V: Clone, const F: usize> OccupiedEntry<'a, K, V, F> {
    /// The key associated with the entry. This key is taken from the map
    /// itself; the copy of the key which was passed into [`HashMap::entry`] is
    /// already dropped.
    pub fn key(entry: &Self) -> &K {
        let Node::Leaf { key, .. } = &*entry.node else {
            unreachable!()
        };
        key
    }

    /// Converts the entry into a shared reference to the value.
    pub fn into_ref(entry: Self) -> &'a V {
        let Node::Leaf { value,..} = Cowarc::into_ref(entry.node) else {
            unreachable!();
        };
        value
    }

    /// Converts the entry into a mutable reference to the value.
    pub fn into_mut(entry: Self) -> &'a mut V {
        let Node::Leaf { value,..} = Cowarc::into_mut(entry.node) else {
            unreachable!();
        };
        value
    }
}

impl<'a, K, V, const F: usize> Deref for OccupiedEntry<'a, K, V, F> {
    type Target = V;

    fn deref(&self) -> &Self::Target {
        let Node::Leaf { value,..} = &*self.node else {
            unreachable!();
        };
        value
    }
}

impl<'a, K: Clone, V: Clone, const F: usize> DerefMut for OccupiedEntry<'a, K, V, F> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        let Node::Leaf { value,..} = &mut *self.node else {
            unreachable!();
        };
        value
    }
}

enum EntryPointer<'a, K, V, const F: usize> {
    Root(&'a mut Pointer<K, V, F>),
    Child {
        node: Cowarc<'a, Node<K, V, F>>,
        index: usize,
    },
}

impl<'a, K: Clone, V: Clone, const F: usize> EntryPointer<'a, K, V, F> {
    fn get(&self) -> Option<&Node<K, V, F>> {
        match self {
            EntryPointer::Root(node) => node.as_ref().map(|v| &**v),
            EntryPointer::Child { node, index } => {
                let Node::SubTrie { entries } = &**node else { unreachable!() };
                entries[*index].as_ref().map(|v| &**v)
            }
        }
    }

    fn into_cowarc(self) -> Cowarc<'a, Node<K, V, F>> {
        match self {
            EntryPointer::Root(node) => Cowarc::new_root(node.as_mut().unwrap()),
            EntryPointer::Child { mut node, index } => {
                let Node::SubTrie { entries } = &*node else { unreachable!() };
                let ptr = NonNull::from(&**entries[index].as_ref().unwrap());
                let map = Box::new(move || {
                    let Node::SubTrie { entries } = &mut *node else { unreachable!() };
                    Arc::make_mut(entries[index].as_mut().unwrap()).into()
                });
                unsafe { Cowarc::new_child(ptr, map) }
            }
        }
    }
}

#[derive(Clone)]
enum Node<K, V, const F: usize> {
    SubTrie { entries: [Pointer<K, V, F>; F] },
    Leaf { partial_hash: u64, key: K, value: V },
}

/// An [`Iterator`] over key-value pairs in a [`HashMap`].
pub struct Iter<'a, K, V, const F: usize> {
    stack: Vec<(&'a Node<K, V, F>, usize)>,
}

impl<'a, K, V, const F: usize> Clone for Iter<'a, K, V, F> {
    fn clone(&self) -> Self {
        Self {
            stack: self.stack.clone(),
        }
    }
}

impl<'a, K, V, const F: usize> Iterator for Iter<'a, K, V, F> {
    type Item = (&'a K, &'a V);

    fn next(&mut self) -> Option<Self::Item> {
        'outer: while let Some((node, index)) = self.stack.last_mut() {
            match &**node {
                Node::SubTrie { entries } => {
                    while *index < F {
                        let prev_index = *index;
                        *index += 1;
                        if let Some(node) = entries[prev_index].as_ref() {
                            self.stack.push((node, 0));
                            continue 'outer;
                        }
                    }
                    self.stack.pop();
                }
                Node::Leaf { key, value, .. } => {
                    self.stack.pop();
                    return Some((key, value));
                }
            }
        }
        None
    }
}

#[cfg(test)]
mod test {
    use std::hash::RandomState;

    use crate::hash_map::HashMap;

    #[test]
    fn simple_hash_map() {
        // Using the smallest possible size forces as much sharing as possible, to
        // stress-test the algorithms.
        let mut hm = HashMap::<_, _, _, 2>::new(RandomState::new());

        let mut hm2 = hm.clone();
        hm.insert("foo", "bar");

        assert_eq!(hm.get("foo"), Some(&"bar"));
        assert_eq!(hm2.get("foo"), None);

        hm.insert("key1", "value1");
        hm.insert("key2", "value2");
        hm.insert("key3", "value3");
        hm.insert("key4", "value4");

        assert_eq!(hm.len(), 5);

        println!("hm = {:#?}", hm);
        println!("hm2 = {:#?}", hm2);

        let hm3 = hm.clone();
        hm.remove(&"foo");

        println!("hm = {:#?}", hm);
        println!("hm2 = {:#?}", hm2);
        println!("hm3 = {:#?}", hm3);
    }
}