orientdb 0.1.2

use std::collections::{HashMap, HashSet};

/// A simple in-memory graph with nodes and directed edges
pub struct SimpleGraph<T> {
    nodes: HashMap<usize, T>,
    edges: HashMap<usize, Vec<usize>>,
    next_id: usize,
}

impl<T> SimpleGraph<T> {
    /// Create a new empty graph
    pub fn new() -> Self {
        Self {
            nodes: HashMap::new(),
            edges: HashMap::new(),
            next_id: 0,
        }
    }

    /// Add a node to the graph, returns the node ID
    pub fn add_node(&mut self, value: T) -> usize {
        let id = self.next_id;
        self.nodes.insert(id, value);
        self.edges.insert(id, Vec::new());
        self.next_id += 1;
        id
    }

    /// Add a directed edge from source to target
    pub fn add_edge(&mut self, from: usize, to: usize) -> bool {
        if self.nodes.contains_key(&from) && self.nodes.contains_key(&to) {
            self.edges.get_mut(&from).unwrap().push(to);
            true
        } else {
            false
        }
    }

    /// Get a node value by ID
    pub fn get_node(&self, id: usize) -> Option<&T> {
        self.nodes.get(&id)
    }

    /// Get all neighbors of a node
    pub fn neighbors(&self, id: usize) -> Option<&Vec<usize>> {
        self.edges.get(&id)
    }

    /// Fast DFS tree search to find a node by predicate
    pub fn find<F>(&self, start: usize, predicate: F) -> Option<usize>
    where
        F: Fn(&T) -> bool,
    {
        let mut visited = HashSet::new();
        let mut stack = vec![start];

        while let Some(node_id) = stack.pop() {
            if visited.contains(&node_id) {
                continue;
            }
            visited.insert(node_id);

            if let Some(value) = self.nodes.get(&node_id) {
                if predicate(value) {
                    return Some(node_id);
                }

                if let Some(neighbors) = self.edges.get(&node_id) {
                    for &neighbor in neighbors.iter().rev() {
                        if !visited.contains(&neighbor) {
                            stack.push(neighbor);
                        }
                    }
                }
            }
        }
        None
    }

    /// Get total number of nodes
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }
}

impl<T> Default for SimpleGraph<T> {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_add_nodes() {
        let mut graph = SimpleGraph::new();
        let id1 = graph.add_node("A");
        let id2 = graph.add_node("B");

        assert_eq!(graph.node_count(), 2);
        assert_eq!(graph.get_node(id1), Some(&"A"));
        assert_eq!(graph.get_node(id2), Some(&"B"));
    }

    #[test]
    fn test_add_edges() {
        let mut graph = SimpleGraph::new();
        let a = graph.add_node("A");
        let b = graph.add_node("B");

        assert!(graph.add_edge(a, b));
        assert_eq!(graph.neighbors(a), Some(&vec![b]));
    }

    #[test]
    fn test_find() {
        let mut graph = SimpleGraph::new();
        let a = graph.add_node("A");
        let b = graph.add_node("B");
        let c = graph.add_node("C");

        graph.add_edge(a, b);
        graph.add_edge(b, c);

        let result = graph.find(a, |&v| v == "C");
        assert_eq!(result, Some(c));
    }
}