From fdf79c26b105d051559b973d9e3809e3c3d49e9e Mon Sep 17 00:00:00 2001
From: Jeremy Wall <jeremy@marzhillstudios.com>
Date: Mon, 5 Sep 2022 10:16:34 -0400
Subject: [PATCH] Find the next set of gap ids for a given set of nodes

---
 src/dag.rs  | 34 ++++++++++++++++++++++
 src/test.rs | 83 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+)

diff --git a/src/dag.rs b/src/dag.rs
index 4365a9a..396d50b 100644
--- a/src/dag.rs
+++ b/src/dag.rs
@@ -152,6 +152,40 @@ where
         })
     }
 
+    /// Find the immediate next non descendant nodes in this graph for the given `search_nodes`.
+    pub fn find_next_non_descendant_nodes(
+        &self,
+        search_nodes: &BTreeSet<Vec<u8>>,
+    ) -> Result<Vec<Node<HW>>> {
+        let mut stack: Vec<Vec<u8>> = dbg!(self.roots.iter().cloned().collect());
+        dbg!(search_nodes);
+        let mut ids = BTreeSet::new();
+        while !stack.is_empty() {
+            let node_id = dbg!(stack.pop().unwrap());
+            let node = self.get_node_by_id(node_id.as_slice())?.unwrap();
+            let deps = node.dependency_ids();
+            if dbg!(deps.len()) == 0 {
+                // This is a leaf node which means it's the beginning of a sub graph
+                // the search_nodes_are not part of.
+                ids.insert(node.id().to_owned());
+            }
+            for dep in deps {
+                // We found one of the search roots.
+                if dbg!(search_nodes.contains(dep.as_slice())) {
+                    // This means that the previous node is a parent of the search_roots.
+                    ids.insert(node.id().to_owned());
+                    continue;
+                }
+                stack.push(dep.to_owned())
+            }
+        }
+        let mut result = Vec::new();
+        for id in ids {
+            result.push(self.get_node_by_id(id.as_slice())?.unwrap());
+        }
+        Ok(result)
+    }
+
     fn search_graph(&self, root_id: &[u8], search_id: &[u8]) -> Result<bool> {
         if root_id == search_id {
             return Ok(true);
diff --git a/src/test.rs b/src/test.rs
index a4dfcd2..4040838 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -173,6 +173,89 @@ fn test_node_comparison_no_shared_graph() {
     );
 }
 
+#[test]
+fn test_find_next_missing_nodes_disjoint_graphs_no_deps() {
+    let mut dag1 = TestDag::new(BTreeMap::new());
+    let mut dag2 = TestDag::new(BTreeMap::new());
+    let quake_node_id = dag1.add_node("quake", BTreeSet::new()).unwrap();
+    let qualm_node_id = dag1.add_node("qualm", BTreeSet::new()).unwrap();
+    dag2.add_node("quell", BTreeSet::new()).unwrap();
+    let missing_nodes = dag1
+        .find_next_non_descendant_nodes(dag2.get_roots())
+        .unwrap();
+    assert_eq!(missing_nodes.len(), 2);
+    let mut found_quake = false;
+    let mut found_qualm = false;
+    for node in missing_nodes {
+        if node.id().to_owned() == quake_node_id {
+            found_quake = true;
+        }
+        if node.id().to_owned() == qualm_node_id {
+            found_qualm = true;
+        }
+    }
+    assert!(found_quake);
+    assert!(found_qualm);
+}
+
+#[test]
+fn test_find_next_missing_nodes_sub_graphs_one_degree_off() {
+    let mut dag1 = TestDag::new(BTreeMap::new());
+    let mut dag2 = TestDag::new(BTreeMap::new());
+    dag1.add_node("quake", BTreeSet::new()).unwrap();
+    let quake_node_id = dag2.add_node("quake", BTreeSet::new()).unwrap();
+
+    let mut deps = BTreeSet::new();
+    deps.insert(quake_node_id);
+    let qualm_node_id = dag1.add_node("qualm", deps).unwrap();
+
+    let missing_nodes = dag1
+        .find_next_non_descendant_nodes(dag2.get_roots())
+        .unwrap();
+    assert_eq!(missing_nodes.len(), 1);
+    let mut found_qualm = false;
+    for node in missing_nodes {
+        if node.id().to_owned() == qualm_node_id {
+            found_qualm = true;
+        }
+    }
+    assert!(found_qualm);
+}
+
+#[test]
+fn test_find_next_missing_nodes_sub_graphs_two_degree_off() {
+    let mut dag1 = TestDag::new(BTreeMap::new());
+    let mut dag2 = TestDag::new(BTreeMap::new());
+    dag1.add_node("quake", BTreeSet::new()).unwrap();
+    let quake_node_id = dag2.add_node("quake", BTreeSet::new()).unwrap();
+
+    let mut deps = BTreeSet::new();
+    deps.insert(quake_node_id.clone());
+    let qualm_node_id = dag1.add_node("qualm", deps).unwrap();
+
+    deps = BTreeSet::new();
+    deps.insert(quake_node_id.clone());
+    deps.insert(qualm_node_id.clone());
+    let quell_node_id = dag1.add_node("quell", deps).unwrap();
+
+    let missing_nodes = dag1
+        .find_next_non_descendant_nodes(dag2.get_roots())
+        .unwrap();
+    assert_eq!(missing_nodes.len(), 2);
+    let mut found_qualm = false;
+    let mut found_quell = false;
+    for node in missing_nodes {
+        if node.id().to_owned() == qualm_node_id {
+            found_qualm = true;
+        }
+        if node.id().to_owned() == quell_node_id {
+            found_quell = true;
+        }
+    }
+    assert!(found_qualm);
+    assert!(found_quell);
+}
+
 #[cfg(feature = "cbor")]
 mod cbor_serialization_tests {
     use super::TestDag;