spokes 0.4.0-rc.3

//! Node Storge Trait and Implementations

use std::{
    collections::{BTreeMap, HashMap},
    hash::BuildHasher,
};

/// Node Storage
pub trait NodeStorage<I, A>: Default {
    /// Add a node with id and attribute, returning an old attribute if a node was replaced.
    fn add_node(&mut self, id: I, attribute: A) -> Option<A>;

    /// Remove a node by id and return the attribute if such a node existed.
    fn remove_node(&mut self, id: &I) -> Option<A>;

    /// Check if a node is contained withing this storage system
    fn contains_node(&self, id: &I) -> bool;

    /// Number of nodes in this store
    fn n_nodes(&self) -> usize;

    /// Get a node's attribute by reference.
    fn node(&self, id: &I) -> Option<&A>;

    /// Get a mutable reference to a node's attribute.
    fn node_mut(&mut self, id: &I) -> Option<&mut A>;

    /// Iterate through nodes with ids and attributes
    fn iter_nodes<'a>(&'a self) -> impl Iterator<Item = (&'a I, &'a A)>
    where
        I: 'a,
        A: 'a;

    /// Iterate through mutable nodes with ids and attributes
    fn iter_nodes_mut<'a>(&'a mut self) -> impl Iterator<Item = (&'a I, &'a mut A)>
    where
        I: 'a,
        A: 'a;
}

impl<I, A, S> NodeStorage<I, A> for HashMap<I, A, S>
where
    I: Eq + std::hash::Hash,
    S: Default + BuildHasher,
{
    fn add_node(&mut self, id: I, attribute: A) -> Option<A> {
        self.insert(id, attribute)
    }

    fn remove_node(&mut self, id: &I) -> Option<A> {
        self.remove(id)
    }

    fn contains_node(&self, id: &I) -> bool {
        self.contains_key(id)
    }

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

    fn node(&self, id: &I) -> Option<&A> {
        self.get(id)
    }

    fn node_mut(&mut self, id: &I) -> Option<&mut A> {
        self.get_mut(id)
    }

    fn iter_nodes<'a>(&'a self) -> impl Iterator<Item = (&'a I, &'a A)>
    where
        I: 'a,
        A: 'a,
    {
        self.iter()
    }

    fn iter_nodes_mut<'a>(&'a mut self) -> impl Iterator<Item = (&'a I, &'a mut A)>
    where
        I: 'a,
        A: 'a,
    {
        self.iter_mut()
    }
}

impl<I, V> NodeStorage<I, V> for BTreeMap<I, V>
where
    I: Ord,
{
    fn add_node(&mut self, id: I, attribute: V) -> Option<V> {
        self.insert(id, attribute)
    }

    fn remove_node(&mut self, id: &I) -> Option<V> {
        self.remove(id)
    }

    fn contains_node(&self, id: &I) -> bool {
        self.contains_key(id)
    }

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

    fn node(&self, id: &I) -> Option<&V> {
        self.get(id)
    }

    fn node_mut(&mut self, id: &I) -> Option<&mut V> {
        self.get_mut(id)
    }

    fn iter_nodes<'a>(&'a self) -> impl Iterator<Item = (&'a I, &'a V)>
    where
        I: 'a,
        V: 'a,
    {
        self.iter()
    }

    fn iter_nodes_mut<'a>(&'a mut self) -> impl Iterator<Item = (&'a I, &'a mut V)>
    where
        I: 'a,
        V: 'a,
    {
        self.iter_mut()
    }
}

/// A Dense node store, i.e. stores nodes in a vec indexed by vec index.
#[derive(Clone, Debug)]
pub struct DenseNodeStorage<A> {
    store: Vec<Option<(usize, A)>>,
    len: usize,
}

impl<A> Default for DenseNodeStorage<A> {
    fn default() -> Self {
        Self {
            store: vec![],
            len: 0,
        }
    }
}

impl<A> IntoIterator for DenseNodeStorage<A> {
    type Item = (usize, A);
    type IntoIter = std::iter::Flatten<std::vec::IntoIter<Option<(usize, A)>>>;

    fn into_iter(self) -> Self::IntoIter {
        self.store.into_iter().flatten()
    }
}

impl<A> NodeStorage<usize, A> for DenseNodeStorage<A> {
    fn add_node(&mut self, id: usize, attribute: A) -> Option<A> {
        // ensure the vec is of sufficient size
        for _ in self.store.len()..=id {
            self.store.push(None);
        }

        if self.store[id].is_none() {
            self.len += 1;
        }

        self.store[id].replace((id, attribute)).map(|(_i, a)| a)
    }

    fn remove_node(&mut self, id: &usize) -> Option<A> {
        if self.store.get(*id).is_some() {
            self.len -= 1;
        }

        self.store
            .get_mut(*id)
            .and_then(|v| v.take().map(|(_, a)| a))
    }

    fn contains_node(&self, id: &usize) -> bool {
        self.store
            .get(*id)
            .is_some_and(std::option::Option::is_some)
    }

    fn iter_nodes<'a>(&'a self) -> impl Iterator<Item = (&'a usize, &'a A)>
    where
        A: 'a,
    {
        self.store.iter().flatten().map(|(i, a)| (i, a))
    }

    fn iter_nodes_mut<'a>(&'a mut self) -> impl Iterator<Item = (&'a usize, &'a mut A)>
    where
        A: 'a,
    {
        self.store.iter_mut().flatten().map(|(i, a)| (&*i, a))
    }

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

    fn node(&self, id: &usize) -> Option<&A> {
        self.store
            .get(*id)
            .as_ref()
            .and_then(|opt| opt.as_ref().map(|(_, a)| a))
    }

    fn node_mut(&mut self, id: &usize) -> Option<&mut A> {
        self.store
            .get_mut(*id)
            .and_then(|opt| opt.as_mut().map(|(_, a)| a))
    }
}