spokes 0.3.0

use crate::{ArcInfo, arc_storage::ArcStorage};

/// Forward Star representation of arcs in a networks
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ForwardAndReverseStar<A>
where
    A: Default,
{
    /// Jump list for faster lookups of arcs corresponding to a node
    jump: Vec<usize>,
    /// Reverse jump list for faster lookups of arcs corresponding to a node
    rjump: Vec<usize>,
    /// Trace lookup, re-ordering for sort by head then tail.
    trace: Vec<usize>,
    /// Arc Data
    data: Vec<ArcInfo<usize, A>>,
    /// Number of Arcs
    m_arcs: usize,
}

impl<A> ArcStorage<usize, A> for ForwardAndReverseStar<A>
where
    A: Default,
{
    /// Create a new `FowardAndRevereStar` arc representation with pre-allocated capacities.
    fn with_capacity(n: usize, m: usize) -> Self {
        Self {
            jump: Vec::with_capacity(n),
            rjump: Vec::with_capacity(n),
            trace: Vec::with_capacity(m),
            data: Vec::with_capacity(m),
            m_arcs: 0,
        }
    }

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

    fn arc(&self, tail: usize, head: usize) -> Option<&A> {
        let start = self.jump.get(tail).copied().unwrap_or(self.m_arcs);
        let end = self.jump.get(tail + 1).copied().unwrap_or(self.m_arcs);

        self.data[start..end]
            .binary_search_by_key(&head, |a| a.head)
            .map(|x| x + start)
            .ok()
            .map(|idx| &self.data[idx].attributes)
    }

    fn arc_mut(&mut self, tail: usize, head: usize) -> Option<&mut A> {
        let start = self.jump.get(tail).copied().unwrap_or(self.m_arcs);
        let end = self.jump.get(tail + 1).copied().unwrap_or(self.m_arcs);

        self.data[start..end]
            .binary_search_by_key(&head, |a| a.head)
            .map(|x| x + start)
            .ok()
            .map(|idx| &mut self.data[idx].attributes)
    }

    fn contains_arc(&self, tail: usize, head: usize) -> bool {
        let start = self.jump.get(tail).copied().unwrap_or(self.m_arcs);
        let end = self.jump.get(tail + 1).copied().unwrap_or(self.m_arcs);

        self.data[start..end]
            .binary_search_by_key(&head, |a| a.head)
            .is_ok()
    }

    fn add_arc<T: Into<ArcInfo<usize, A>>>(&mut self, arc: T) {
        let arc = arc.into();

        // Extend the jump list if the tail node is not included in the list
        if self.jump.len() <= arc.tail {
            self.jump
                .extend(std::iter::repeat_n(self.m_arcs, arc.tail - self.jump.len() + 1));
        }

        // Extend the rjump list if the tail node is not included in the list
        if self.rjump.len() <= arc.head {
            self.rjump
                .extend(std::iter::repeat_n(self.m_arcs, arc.head - self.rjump.len() + 1));
        }

        // Get intervals intervals corresponding to tail node in data table
        let f_start = self.jump[arc.tail];
        let f_end = self.jump.get(arc.tail + 1).copied().unwrap_or(self.m_arcs);

        // Get intervals intervals corresponding to head node in trace table
        let r_start = self.rjump[arc.head];
        let r_end = self.rjump.get(arc.head + 1).copied().unwrap_or(self.m_arcs);

        // insert the arc into the data
        let f_insertion_point: usize = self.data[f_start..f_end]
            .binary_search_by_key(&arc.head, |a| a.head)
            .unwrap_or_else(|e| e)
            + f_start;

        // Find insertion point in trace list
        let r_insertion_point: usize = self.trace[r_start..r_end]
            .binary_search_by(|&x| self.data[x].tail.cmp(&arc.tail))
            .unwrap_or_else(|e| e)
            + r_start;

        // Increment jump list after segment end
        for i in (arc.tail + 1)..self.jump.len() {
            self.jump[i] += 1;
        }

        // Increment rjump list after segment end
        for i in (arc.head + 1)..self.rjump.len() {
            self.rjump[i] += 1;
        }

        // Insert new arc data
        self.data.insert(f_insertion_point, arc);

        // increment all values in trace above f_insertion_point
        self.trace.iter_mut().for_each(|x| {
            if *x >= f_insertion_point {
                *x += 1;
            }
        });

        // Insert mapping to new value
        self.trace.insert(r_insertion_point, f_insertion_point);
        self.m_arcs += 1;
    }

    fn remove_arc(&mut self, tail: &usize, head: &usize) -> Option<A> {
        // Find location in data
        let f_start = self.jump[*tail];
        let f_end = self.jump.get(tail + 1).copied().unwrap_or(self.m_arcs);

        let data_idx = self.data[dbg!(f_start)..dbg!(f_end)]
            .binary_search_by_key(&head, |a| &a.head)
            .map(|x| x + f_start);

        let r_start = self.rjump[*head];
        let r_end = self.rjump.get(head + 1).copied().unwrap_or(self.m_arcs);

        let trace_idx = self.trace[dbg!(r_start)..dbg!(r_end)]
            .binary_search_by(|&idx| self.data[idx].tail.cmp(tail))
            .map(|x| x + r_start);

        match (data_idx, trace_idx) {
            (Ok(data_idx), Ok(trace_idx)) => {
                // the entry exists, so it's removed from the data-structure

                // Update jump
                for i in (tail + 1)..self.jump.len() {
                    self.jump[i] = self.jump[i].saturating_sub(1);
                }

                // Update rjump
                for i in (head + 1)..self.rjump.len() {
                    self.rjump[i] = self.rjump[i].saturating_sub(1);
                }

                // Decrement values above trace_idx
                for i in 0..self.m_arcs {
                    if self.trace[i] > self.trace[trace_idx] {
                        self.trace[i] = self.trace[i].saturating_sub(1);
                    }
                }
                // Remove entry in the trace
                self.trace.remove(trace_idx);

                self.m_arcs -= 1;
                Some(self.data.remove(data_idx).attributes)
            }
            (Err(_), Err(_)) => {
                // The arc doesn't exist, so return None
                None
            }
            (Ok(_), Err(_)) | (Err(_), Ok(_)) => {
                // The internal consistency of the ForwardAndReverseStar representation is broken.
                unreachable!();
            }
        }
    }

    fn remove_arcs_with_node(&mut self, id: &usize) {
        let to_remove: Vec<_> = self
            .forward_arcs(id)
            .chain(self.reverse_arcs(id))
            .map(|arc| (arc.tail, arc.head))
            .collect();

        self.remove_arcs(to_remove.iter().map(|(t, h)| (t, h)));
    }

    fn has_forward_arcs(&self, node: &usize) -> bool {
        let start: usize = self.jump.get(*node).copied().unwrap_or(self.m_arcs);
        let end: usize = self.jump.get(*node + 1).copied().unwrap_or(self.m_arcs);
        start > end
    }

    fn has_reverse_arcs(&self, node: &usize) -> bool {
        let start: usize = self.rjump.get(*node).copied().unwrap_or(self.m_arcs);
        let end: usize = self.rjump.get(*node + 1).copied().unwrap_or(self.m_arcs);
        start > end
    }

    fn arc_iter<'a>(&'a self) -> impl Iterator<Item = &'a ArcInfo<usize, A>> + 'a
    where
        A: 'a,
    {
        self.data.iter()
    }

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

    fn forward_arcs<'a>(&'a self, node: &usize) -> impl Iterator<Item = &'a ArcInfo<usize, A>> + 'a
    where
        A: 'a,
    {
        let start: usize = self.jump.get(*node).copied().unwrap_or(self.m_arcs);
        let end: usize = self.jump.get(*node + 1).copied().unwrap_or(self.m_arcs);
        self.data[start..end].iter()
    }

    fn reverse_arcs<'a>(&'a self, node: &usize) -> impl Iterator<Item = &'a ArcInfo<usize, A>> + 'a
    where
        A: 'a,
    {
        let start: usize = self.rjump.get(*node).copied().unwrap_or(self.m_arcs);
        let end: usize = self.rjump.get(*node + 1).copied().unwrap_or(self.m_arcs);
        self.trace[start..end].iter().map(|&idx| &self.data[idx])
    }
}

impl<A> ForwardAndReverseStar<A>
where
    A: Default + std::fmt::Debug,
{
    /// Create a new `ForwadAndRevereStar` representation of the arcs in a network
    ///
    /// # Example
    /// ```rust
    /// use spokes::arc_storage::ForwardAndReverseStar;
    ///
    /// let fstar: ForwardAndReverseStar<()> = ForwardAndReverseStar::new();
    /// ```
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn insertion() {
        let mut fstar = ForwardAndReverseStar::new();

        fstar.add_arcs([
            ArcInfo::new(1, 3, ()),
            ArcInfo::new(0, 2, ()),
            ArcInfo::new(3, 2, ()),
            ArcInfo::new(3, 4, ()),
            ArcInfo::new(0, 1, ()),
            ArcInfo::new(4, 2, ()),
            ArcInfo::new(2, 1, ()),
            ArcInfo::new(4, 3, ()),
        ]);

        assert_eq!(fstar.m_arcs, 8);
        assert_eq!(
            fstar.data,
            vec![
                ArcInfo::new(0, 1, ()),
                ArcInfo::new(0, 2, ()),
                ArcInfo::new(1, 3, ()),
                ArcInfo::new(2, 1, ()),
                ArcInfo::new(3, 2, ()),
                ArcInfo::new(3, 4, ()),
                ArcInfo::new(4, 2, ()),
                ArcInfo::new(4, 3, ()),
            ]
        );
        assert_eq!(fstar.jump, vec![0, 2, 3, 4, 6]);
        assert_eq!(fstar.rjump, vec![0, 0, 2, 5, 7]);
        assert_eq!(fstar.trace, vec![0, 3, 1, 4, 6, 2, 7, 5]);
    }

    #[test]
    fn single_insertion_then_deletion() {
        let mut fstar = ForwardAndReverseStar::new();
        fstar.add_arc(ArcInfo::new(0, 1, ()));
        assert_eq!(fstar.remove_arc(&0, &1), Some(()));
        assert_eq!(fstar.m_arcs(), 0);
    }

    #[test]
    fn insertion_then_deletions() {
        let mut fstar = ForwardAndReverseStar::new();

        fstar.add_arcs([
            ArcInfo::new(1, 3, ()),
            ArcInfo::new(0, 2, ()),
            ArcInfo::new(3, 2, ()),
            ArcInfo::new(3, 4, ()),
            ArcInfo::new(0, 1, ()),
            ArcInfo::new(4, 2, ()),
            ArcInfo::new(2, 1, ()),
            ArcInfo::new(4, 3, ()),
        ]);

        dbg!(&fstar.data);
        dbg!(&fstar.jump);
        dbg!(&fstar.rjump);
        dbg!(&fstar.trace);

        assert_eq!(fstar.remove_arc(&4, &2), Some(()));
        assert_eq!(fstar.m_arcs, 7);

        dbg!(&fstar.data);
        dbg!(&fstar.jump);
        dbg!(&fstar.rjump);
        dbg!(&fstar.trace);

        assert_eq!(
            fstar.data,
            vec![
                ArcInfo::new(0, 1, ()), // 0
                ArcInfo::new(0, 2, ()), // 1
                ArcInfo::new(1, 3, ()), // 2
                ArcInfo::new(2, 1, ()), // 3
                ArcInfo::new(3, 2, ()), // 4
                ArcInfo::new(3, 4, ()), // 5
                ArcInfo::new(4, 3, ()), // 6
            ]
        );

        assert_eq!(fstar.trace, vec![0, 3, 1, 4, 2, 6, 5]);
        assert_eq!(fstar.jump, vec![0, 2, 3, 4, 6]);
        assert_eq!(fstar.rjump, vec![0, 0, 2, 4, 6]);

        assert_eq!(fstar.remove_arc(&1, &3), Some(()));
        assert_eq!(fstar.m_arcs, 6);

        dbg!(&fstar.data);
        dbg!(&fstar.jump);
        dbg!(&fstar.rjump);
        dbg!(&fstar.trace);

        assert_eq!(
            fstar.data,
            vec![
                ArcInfo::new(0, 1, ()), // 0
                ArcInfo::new(0, 2, ()), // 1
                ArcInfo::new(2, 1, ()), // 2
                ArcInfo::new(3, 2, ()), // 3
                ArcInfo::new(3, 4, ()), // 4
                ArcInfo::new(4, 3, ()), // 5
            ]
        );

        assert_eq!(fstar.trace, vec![0, 2, 1, 3, 5, 4]);
        assert_eq!(fstar.jump, vec![0, 2, 2, 3, 5]);
        assert_eq!(fstar.rjump, vec![0, 0, 2, 4, 5]);

        // Check forward arcs
        assert_eq!(
            fstar.forward_arcs(&0).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(0, 1, ()), &ArcInfo::new(0, 2, ()),]
        );

        assert_eq!(fstar.forward_arcs(&1).count(), 0);

        assert_eq!(
            fstar.forward_arcs(&2).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(2, 1, ()),]
        );

        assert_eq!(
            fstar.forward_arcs(&3).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(3, 2, ()), &ArcInfo::new(3, 4, ()),]
        );

        assert_eq!(
            fstar.forward_arcs(&4).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(4, 3, ())]
        );

        // Check reverse arcs
        assert_eq!(fstar.reverse_arcs(&0).count(), 0);

        assert_eq!(
            fstar.reverse_arcs(&1).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(0, 1, ()), &ArcInfo::new(2, 1, ()),]
        );

        assert_eq!(
            fstar.reverse_arcs(&2).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(0, 2, ()), &ArcInfo::new(3, 2, ()),]
        );

        assert_eq!(
            fstar.reverse_arcs(&3).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(4, 3, ()),]
        );

        assert_eq!(
            fstar.reverse_arcs(&4).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(3, 4, ()),]
        );
    }

    #[test]
    fn get_forward_arcs() {
        let mut fstar = ForwardAndReverseStar::with_capacity(5, 8);

        fstar.add_arcs([
            ArcInfo::new(1, 2, ()),
            ArcInfo::new(1, 3, ()),
            ArcInfo::new(2, 4, ()),
            ArcInfo::new(3, 2, ()),
            ArcInfo::new(4, 3, ()),
            ArcInfo::new(4, 5, ()),
            ArcInfo::new(5, 3, ()),
            ArcInfo::new(5, 4, ()),
        ]);

        assert_eq!(
            fstar.forward_arcs(&4).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(4, 3, ()), &ArcInfo::new(4, 5, ()),]
        );
    }

    #[test]
    fn get_reverse_arcs() {
        let mut fstar = ForwardAndReverseStar::with_capacity(5, 8);

        fstar.add_arcs([
            ArcInfo::new(1, 2, ()),
            ArcInfo::new(1, 3, ()),
            ArcInfo::new(2, 4, ()),
            ArcInfo::new(3, 2, ()),
            ArcInfo::new(4, 3, ()),
            ArcInfo::new(4, 5, ()),
            ArcInfo::new(5, 3, ()),
            ArcInfo::new(5, 4, ()),
        ]);

        assert_eq!(fstar.reverse_arcs(&1).count(), 0);
        assert_eq!(
            fstar.reverse_arcs(&2).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(1, 2, ()), &ArcInfo::new(3, 2, ()),]
        );
        assert_eq!(
            fstar.reverse_arcs(&3).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![
                &ArcInfo::new(1, 3, ()),
                &ArcInfo::new(4, 3, ()),
                &ArcInfo::new(5, 3, ()),
            ]
        );
        assert_eq!(
            fstar.reverse_arcs(&4).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(2, 4, ()), &ArcInfo::new(5, 4, ()),]
        );
        assert_eq!(
            fstar.reverse_arcs(&5).collect::<Vec<&ArcInfo<usize, ()>>>(),
            vec![&ArcInfo::new(4, 5, ())]
        );
    }
}