raphtory-core 0.18.5

//! A data structure for efficiently storing and querying the temporal adjacency set of a node in a temporal graph.
use raphtory_api::iter::BoxedLIter;
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, hash::Hash};

const SMALL_SET: usize = 1024;

/**
 * Temporal adjacency set can track when adding edge v -> u
 * does u exist already
 * and if it does what is the edge metadata
 *
 *  */
#[derive(Debug, Default, Serialize, Deserialize, PartialEq)]
pub enum AdjSet<K: Ord + Copy + Hash + Send + Sync, V: Into<usize> + Copy + Send + Sync> {
    #[default]
    Empty,
    One(K, V),
    Small {
        vs: Vec<K>,    // the neighbours
        edges: Vec<V>, // edge metadata
    },
    Large {
        vs: BTreeMap<K, V>, // this is equiv to vs and edges
    },
    // TODO: if we use BTreeSet<(K, Option<V>)> we could implement intersections and support edge label queries such as a && b
}

impl<K: Ord + Copy + Hash + Send + Sync, V: Into<usize> + Copy + Send + Sync> AdjSet<K, V> {
    pub fn len(&self) -> usize {
        match self {
            AdjSet::Empty => 0,
            AdjSet::One(_, _) => 1,
            AdjSet::Small { vs, .. } => vs.len(),
            AdjSet::Large { vs } => vs.len(),
        }
    }

    pub fn is_empty(&self) -> bool {
        match self {
            AdjSet::Empty => true,
            AdjSet::One(_, _) => false,
            AdjSet::Small { vs, .. } => vs.is_empty(),
            AdjSet::Large { vs } => vs.is_empty(),
        }
    }
    pub fn new(v: K, e: V) -> Self {
        Self::One(v, e)
    }

    /// Push a new node and edge into the adjacency set.
    ///
    /// If the node already exists, it will not be added again.
    /// Returns `true` if the node was added, `false` if it already existed
    pub fn push(&mut self, v: K, e: V) -> bool {
        match self {
            AdjSet::Empty => {
                *self = Self::new(v, e);
                true
            }
            AdjSet::One(vv, ee) => {
                if *vv < v {
                    *self = Self::Small {
                        vs: vec![*vv, v],
                        edges: vec![*ee, e],
                    };
                    true
                } else if *vv > v {
                    *self = Self::Small {
                        vs: vec![v, *vv],
                        edges: vec![e, *ee],
                    };
                    true
                } else {
                    // already exists
                    false
                }
            }
            AdjSet::Small { vs, edges } => match vs.binary_search(&v) {
                Ok(_) => false,
                Err(i) => {
                    if vs.len() < SMALL_SET {
                        vs.insert(i, v);
                        edges.insert(i, e);
                    } else {
                        let mut map =
                            BTreeMap::from_iter(vs.iter().copied().zip(edges.iter().copied()));
                        map.insert(v, e);
                        *self = Self::Large { vs: map }
                    }
                    true
                }
            },
            AdjSet::Large { vs } => vs.insert(v, e).is_none(),
        }
    }

    pub fn iter(&self) -> BoxedLIter<'_, (K, V)> {
        match self {
            AdjSet::Empty => Box::new(std::iter::empty()),
            AdjSet::One(v, e) => Box::new(std::iter::once((*v, *e))),
            AdjSet::Small { vs, edges } => Box::new(vs.iter().copied().zip(edges.iter().copied())),
            AdjSet::Large { vs } => Box::new(vs.iter().map(|(k, v)| (*k, *v))),
        }
    }

    pub fn nodes(&self) -> Box<dyn Iterator<Item = K> + Send + '_> {
        match self {
            AdjSet::Empty => Box::new(std::iter::empty()),
            AdjSet::One(v, ..) => Box::new(std::iter::once(*v)),
            AdjSet::Small { vs, .. } => Box::new(vs.iter().copied()),
            AdjSet::Large { vs } => Box::new(vs.keys().copied()),
        }
    }

    pub fn find(&self, v: K) -> Option<V> {
        match self {
            AdjSet::Empty => None,
            AdjSet::One(vv, e) => (*vv == v).then_some(*e),
            AdjSet::Small { vs, edges } => vs.binary_search(&v).ok().map(|i| edges[i]),
            AdjSet::Large { vs } => vs.get(&v).copied(),
        }
    }

    /// puts elements into page and returns number of returned elements
    pub fn fill_page<const P: usize>(&self, last: Option<K>, page: &mut [(K, V); P]) -> usize {
        match self {
            AdjSet::Empty => 0,
            AdjSet::One(v, i) => {
                if let Some(l) = last {
                    if l < *v {
                        page[0] = (*v, *i);
                        1
                    } else {
                        0
                    }
                } else {
                    page[0] = (*v, *i);
                    1
                }
            }
            AdjSet::Small { vs, edges } => {
                if let Some(l) = last {
                    let i = match vs.binary_search(&l) {
                        Ok(i) => i + 1,
                        Err(i) => i,
                    };

                    if i >= vs.len() {
                        return 0;
                    }

                    let mut index = 0;
                    vs[i..]
                        .iter()
                        .zip(edges[i..].iter())
                        .take(P)
                        .for_each(|(a, b)| {
                            page[index] = (*a, *b);
                            index += 1;
                        });
                    index
                } else {
                    let mut index = 0;
                    vs.iter().zip(edges.iter()).take(P).for_each(|(a, b)| {
                        page[index] = (*a, *b);
                        index += 1;
                    });
                    index
                }
            }
            AdjSet::Large { vs } => {
                if let Some(l) = last {
                    let mut index = 0;
                    vs.range(l..).skip(1).take(P).for_each(|(a, b)| {
                        page[index] = (*a, *b);
                        index += 1;
                    });
                    index
                } else {
                    let mut index = 0;
                    vs.iter().take(P).for_each(|(a, b)| {
                        page[index] = (*a, *b);
                        index += 1
                    });
                    index
                }
            }
        }
    }
    pub fn get_page_vec(&self, last: Option<K>, page_size: usize) -> Vec<(K, V)> {
        match self {
            AdjSet::Empty => vec![],
            AdjSet::One(v, i) => {
                if let Some(l) = last {
                    if l < *v {
                        vec![(*v, *i)]
                    } else {
                        vec![]
                    }
                } else {
                    vec![(*v, *i)]
                }
            }
            AdjSet::Small { vs, edges } => {
                if let Some(l) = last {
                    let i = match vs.binary_search(&l) {
                        Ok(i) => i + 1,
                        Err(i) => i,
                    };

                    if i >= vs.len() {
                        return vec![];
                    }

                    vs[i..]
                        .iter()
                        .zip(edges[i..].iter())
                        .take(page_size)
                        .map(|(a, b)| (*a, *b))
                        .collect()
                } else {
                    vs.iter()
                        .zip(edges.iter())
                        .take(page_size)
                        .map(|(a, b)| (*a, *b))
                        .collect()
                }
            }
            AdjSet::Large { vs } => {
                if let Some(l) = last {
                    vs.range(l..)
                        .skip(1)
                        .take(page_size)
                        .map(|(a, b)| (*a, *b))
                        .collect()
                } else {
                    vs.iter().take(page_size).map(|(a, b)| (*a, *b)).collect()
                }
            }
        }
    }
}

#[cfg(test)]
mod tadjset_tests {
    use super::*;
    use proptest::{arbitrary::any, prop_assert, proptest};

    #[test]
    fn insert_fuzz_proptest() {
        proptest!(|(input in any::<Vec<usize>>())| {
            let mut ts: AdjSet<usize, usize> = AdjSet::default();
            for (e, i) in input.iter().enumerate() {
                ts.push(*i, e);
            }
            prop_assert!(input.iter().all(|i| ts.find(*i).is_some()));
        })
    }

    #[test]
    fn insert() {
        let mut ts: AdjSet<usize, usize> = AdjSet::default();

        ts.push(7, 5);
        let actual = ts.iter().collect::<Vec<_>>();
        let expected: Vec<(usize, usize)> = vec![(7, 5)];
        assert_eq!(actual, expected)
    }

    #[test]
    fn insert_large() {
        let mut ts: AdjSet<usize, usize> = AdjSet::default();

        for i in 0..SMALL_SET + 2 {
            ts.push(i, i);
        }

        for i in 0..SMALL_SET + 2 {
            assert_eq!(ts.find(i), Some(i));
        }
    }

    #[test]
    fn insert_twice() {
        let mut ts: AdjSet<usize, usize> = AdjSet::default();

        ts.push(7, 9);
        ts.push(7, 9);

        let actual = ts.iter().collect::<Vec<_>>();
        let expected: Vec<(usize, usize)> = vec![(7, 9)];
        assert_eq!(actual, expected);
    }

    #[test]
    fn insert_two_different() {
        let mut ts: AdjSet<usize, usize> = AdjSet::default();

        ts.push(1, 0);
        ts.push(7, 1);

        let actual = ts.iter().collect::<Vec<_>>();
        let expected: Vec<(usize, usize)> = vec![(1, 0), (7, 1)];
        assert_eq!(actual, expected);
    }
}