limbo_graph/src/lib.rs

//
// Copyright 2023 James Pace
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
//
// This Source Code Form is "Incompatible With Secondary Licenses", as
// defined by the Mozilla Public License, v. 2.0.
//
#![no_std]
extern crate alloc;
mod node;
mod graph_error;

use alloc::collections::vec_deque::VecDeque;
use alloc::vec;
use alloc::vec::Vec;
pub use node::*;
pub use graph_error::*;

/// A generic graph type holding values connected to other values.
/// Values can be added to the graph, but not removed.
pub struct Graph<NodeValueT: NodeValue> {
    nodes: Vec<Node<NodeValueT>>,
}

impl<NodeValueT: NodeValue> Graph<NodeValueT> {
    /// Make a new graph with a root node with value `root`.
    pub fn new(root: NodeValueT) -> Self {
        // Make root node from its value.
        let root_node = Node::new(root, None);
        // Make graph with root_node as the one value in the vec.
        Graph {
            nodes: vec![root_node],
        }
    }

    /// Add a child wth value `val` to the parent with Key `parent`.
    /// If the parent key is not in the graph, returns an error.
    /// Returns a result with the key of the new node or an error.
    pub fn add(&mut self, val: NodeValueT, parent: Key) -> Result<Key, GraphError> {
        // Make sure parent is valid.
        if parent >= self.nodes.len() {
            return Err(GraphError::from_msg("Parent node not in graph."));
        }
        // Add new node to graph, get it's key.
        let new_node = Node::new(val, Some(parent));
        self.nodes.push(new_node.clone());
        let new_node_key = self.nodes.len() - 1;
        // Add it's key to parent's children.
        self.nodes[parent].add_child(new_node_key);
        Ok(new_node_key)
    }

    /// Replace the value of `key` with value `value`.
    /// `key` must already exist in the graph, and no connections will be modified.
    pub fn replace_value_of(&mut self, key: &Key, value: NodeValueT) -> Result<(), GraphError> {
        if let Some(node) = self.nodes.get_mut(*key) {
            node.set_value(value);
            return Ok(());
        }
        Err(GraphError::from_msg("Can't set value of invalid key."))
    }

    /// Get the value of key `key` if the key is valid.
    pub fn value_of(&self, key: &Key) -> Result<NodeValueT, GraphError> {
        if key >= &self.nodes.len() {
            return Err(GraphError::from_msg("Can't get value of invalid key."));
        }
        Ok(self.nodes[*key].value())
    }

    /// Get the children (as a list of keys) of key `key` if the key is valid.
    pub fn children_of(&self, key: &Key) -> Result<Vec<Key>, GraphError> {
        if key >= &self.nodes.len() {
            return Err(GraphError::from_msg("Can't get children of invalid key."));
        }
        Ok(self.nodes[*key].children())
    }

    /// Get the parent of key `key` if the key is valid.
    /// Will return None if the node at `key` as no parent (i.e. is the root node).
    pub fn parent_of(&self, key: &Key) -> Result<Option<Key>, GraphError> {
        if key >= &self.nodes.len() {
            return Err(GraphError::from_msg("Can't get parent of invalid key."));
        }
        Ok(self.nodes[*key].parent())
    }

    /// Get the key for the root of the graph.
    pub fn root_key(&self) -> Key {
        // This is always 0.
        0
    }

    /// Find all nodes that are leaf nodes.
    /// A leaf node is one that doesn't have any children.
    pub fn find_leaf_keys(&self) -> Result<Vec<Key>, GraphError> {
        let mut leaf_keys = Vec::<Key>::new();

        for key in 0..self.nodes.len() {
            let children_of_node = self.children_of(&key)?;
            if children_of_node.len() == 0 {
                leaf_keys.push(key);
            }
        }

        Ok(leaf_keys)
    }

    /// Return all keys ordered like you were doing
    /// a depth first search.
    pub fn get_keys_by_depth(&self) -> Result<Vec<Key>, GraphError> {
        let mut visited_keys = Vec::<Key>::new();
        let mut stack = Vec::<Key>::new();

        stack.push(self.root_key());

        while stack.len() > 0 {
            let next_key = stack.pop().unwrap();
            if !visited_keys.contains(&next_key) {
                visited_keys.push(next_key);
                let children_of_next_key = self.children_of(&next_key)?;
                stack.extend(children_of_next_key);
            }
        }

        Ok(visited_keys)
    }

    /// Given a key, return the vec of keys connecting that key to the root.
    /// Order is root->key inclusive.
    pub fn backtrack_from_key(&self, key: &Key) -> Result<Vec<Key>, GraphError> {
        let mut curr_key = key.clone();
        let mut deque = VecDeque::<Key>::new();
        deque.push_front(curr_key.clone());

        while curr_key != self.root_key() {
            let new_parent = self.parent_of(&curr_key)?.unwrap();
            deque.push_front(new_parent.clone());
            curr_key = new_parent;
        }

        Ok(Vec::from(deque))
    }
}

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

    #[derive(Debug, Clone, PartialEq)]
    struct NodeType {
        pub x: f64,
    }
    impl NodeValue for NodeType {}
    impl NodeType {
        fn new(val: f64) -> Self {
            NodeType { x: val }
        }
    }

    #[test]
    fn node_manipulation() {
        let mut node = Node::new(NodeType::new(1.0), None);
        assert!(node.parent().is_none());
        assert!(node.value() == NodeType::new(1.0));
        let child_key: Key = 1;
        node.add_child(child_key.clone());
        assert!(node.children().len() == 1);
        assert!(node.children()[0] == child_key);
    }

    #[test]
    fn graph_manipulation() {
        let root_val = NodeType::new(1.0);
        let mut graph = Graph::new(root_val);

        let second_val = NodeType::new(2.0);
        let second_add_res = graph.add(second_val, 0);
        assert!(second_add_res.is_ok());
        assert!(second_add_res.unwrap() == 1);

        assert!(graph.value_of(&1).unwrap() == NodeType::new(2.0));
        assert!(graph.parent_of(&1).unwrap() == Some(0));
        assert!(graph.children_of(&1).unwrap().len() == 0);
    }

    #[test]
    fn graph_search() {
        // Build the graph.
        let root_val = NodeType::new(1.0);
        let mut graph = Graph::new(root_val);

        let second_val = NodeType::new(2.0);
        let second_add_res = graph.add(second_val, 0);
        assert!(second_add_res.is_ok());

        let third_val = NodeType::new(3.0);
        let third_add_res = graph.add(third_val, 0);
        assert!(third_add_res.is_ok());

        let fourth_val = NodeType::new(4.0);
        let fourth_add_res = graph.add(fourth_val, 1);
        assert!(fourth_add_res.is_ok());

        let fifth_val = NodeType::new(5.0);
        let fifth_add_res = graph.add(fifth_val, 2);
        assert!(fifth_add_res.is_ok());

        // Do the testing.
        let leaf_keys = graph.find_leaf_keys().unwrap();
        assert!(leaf_keys.len() == 2);

        let keys_by_depth = graph.get_keys_by_depth().unwrap();
        let expected_keys_by_depth = vec![0, 2, 4, 1, 3];
        assert!(keys_by_depth == expected_keys_by_depth);

        let backtracked_keys = graph.backtrack_from_key(&4).unwrap();
        let expected_backtracked_keys = vec![0, 2, 4];
        assert!(backtracked_keys == expected_backtracked_keys);
    }
}