spokes 0.3.0

use std::collections::HashMap;
use std::hash::Hash;

use crate::ArcInfo;

use super::ArcStorage;

/// Arc storage with hash maps.
#[derive(Clone, Debug)]
pub struct HashMapStorage<I, A> {
    forward: HashMap<I, HashMap<I, usize>>,
    reverse: HashMap<I, HashMap<I, usize>>,
    arc_attributes: HashMap<usize, ArcInfo<I, A>>,
    next_arc_id: usize,
}

impl<I, A> HashMapStorage<I, A> {
    /// Create a new `HashMapStorage`.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

impl<I, A> PartialEq for HashMapStorage<I, A>
where
    I: PartialEq + Hash + Eq + Copy,
    A: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.arc_iter()
            .all(|arc| other.arc(arc.tail, arc.head) == Some(&arc.attributes))
    }
}

impl<I, A> Default for HashMapStorage<I, A> {
    fn default() -> Self {
        Self {
            forward: HashMap::default(),
            reverse: HashMap::default(),
            arc_attributes: HashMap::default(),
            next_arc_id: Default::default(),
        }
    }
}

impl<I, A> ArcStorage<I, A> for HashMapStorage<I, A>
where
    I: Hash + Eq + Copy,
{
    fn with_capacity(n_nodes: usize, m_arcs: usize) -> Self {
        Self {
            forward: HashMap::with_capacity(n_nodes),
            reverse: HashMap::with_capacity(n_nodes),
            arc_attributes: HashMap::with_capacity(m_arcs),
            next_arc_id: 0,
        }
    }

    fn m_arcs(&self) -> usize {
        self.arc_attributes.len()
    }

    fn arc(&self, tail: I, head: I) -> Option<&A> {
        self.forward
            .get(&tail)
            .and_then(|x| x.get(&head))
            .and_then(|id| self.arc_attributes.get(id).map(|a| &a.attributes))
    }

    fn arc_mut(&mut self, tail: I, head: I) -> Option<&mut A> {
        self.forward
            .get(&tail)
            .and_then(|x| x.get(&head))
            .and_then(|id| self.arc_attributes.get_mut(id).map(|a| &mut a.attributes))
    }

    fn contains_arc(&self, tail: I, head: I) -> bool {
        self.forward
            .get(&tail)
            .and_then(|x| x.get(&head))
            .is_some_and(|id| self.arc_attributes.contains_key(id))
    }

    fn add_arc<T: Into<crate::ArcInfo<I, A>>>(&mut self, arc: T) {
        let arc = arc.into();
        let id = self.next_arc_id;
        self.next_arc_id += 1;

        self.forward
            .entry(arc.tail)
            .or_default()
            .insert(arc.head, id);

        self.reverse
            .entry(arc.head)
            .or_default()
            .insert(arc.tail, id);

        self.arc_attributes.insert(id, arc);
    }

    fn remove_arc(&mut self, tail: &I, head: &I) -> Option<A> {
        let forward = self.forward.get_mut(tail)?;
        let reverse = self.reverse.get_mut(head)?;

        let forward_id = forward.get(head).copied();
        let reverse_id = reverse.get(tail).copied();

        match (forward_id, reverse_id) {
            (None, None) => None,
            (Some(x), Some(y)) if x == y => {
                forward.remove(head);
                reverse.remove(tail);
                self.arc_attributes.remove(&x).map(|a| a.attributes)
            }
            (Some(_) | None, Some(_)) | (Some(_), None) => unreachable!(),
        }
    }

    fn arc_iter<'a>(&'a self) -> impl Iterator<Item = &'a ArcInfo<I, A>> + 'a
    where
        A: 'a,
        I: 'a,
    {
        self.forward
            .iter()
            .flat_map(|(_tail, f)| f.values().map(|id| &self.arc_attributes[id]))
    }

    fn arc_iter_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut ArcInfo<I, A>> + 'a
    where
        A: 'a,
        I: 'a,
    {
        self.arc_attributes.values_mut()
    }

    fn forward_arcs<'a>(&'a self, node: &'a I) -> impl Iterator<Item = &'a ArcInfo<I, A>> + 'a
    where
        A: 'a,
        I: 'a,
    {
        self.forward
            .get(node)
            .map_or_else(std::collections::hash_map::Iter::default, |f| f.iter())
            .map(|(_, id)| self.arc_attributes.get(id).unwrap())
    }

    fn reverse_arcs<'a>(&'a self, node: &'a I) -> impl Iterator<Item = &'a ArcInfo<I, A>> + 'a
    where
        A: 'a,
        I: 'a,
    {
        self.reverse
            .get(node)
            .map_or_else(std::collections::hash_map::Iter::default, |f| f.iter())
            .map(|(_, id)| self.arc_attributes.get(id).unwrap())
    }

    fn has_forward_arcs(&self, node: &I) -> bool {
        self.forward.get(node).is_some_and(|a| !a.is_empty())
    }

    fn has_reverse_arcs(&self, node: &I) -> bool {
        self.reverse.get(node).is_some_and(|a| !a.is_empty())
    }

    fn remove_arcs_with_node(&mut self, id: &I) {
        if let Some(reverse) = self.forward.remove(id) {
            for (rid, arc_id) in reverse {
                if let Some(reverse_tails) = self.reverse.get_mut(&rid) {
                    reverse_tails.remove(id);
                }
                self.arc_attributes.remove(&arc_id);
            }
        }
    }
}

impl<I, A> FromIterator<ArcInfo<I, A>> for HashMapStorage<I, A>
where
    I: Hash + Eq + Copy,
{
    fn from_iter<T: IntoIterator<Item = ArcInfo<I, A>>>(iter: T) -> Self {
        let mut arcs = Self::new();
        for arc in iter {
            arcs.add_arc(arc);
        }

        arcs
    }
}

impl<I, A> IntoIterator for HashMapStorage<I, A>
where
    I: Hash + Eq,
{
    type Item = ArcInfo<I, A>;

    type IntoIter = std::collections::hash_map::IntoValues<usize, ArcInfo<I, A>>;

    fn into_iter(self) -> Self::IntoIter {
        self.arc_attributes.into_values()
    }
}