more

Author: molpopgen

src/sys/tree.rs

@@ -209,6 +209,10 @@ impl<'treeseq> LLTree<'treeseq> {
     pub fn children(&self, u: NodeId) -> impl Iterator<Item = NodeId> + '_ {
         NodeIteratorAdapter(ChildIterator::new(self, u))
     }
+
+    pub fn parents(&self, u: NodeId) -> impl Iterator<Item = NodeId> + '_ {
+        NodeIteratorAdapter(ParentsIterator::new(self, u))
+    }
 }
 
 // Trait defining iteration over nodes.
@@ -450,3 +454,41 @@ impl NodeIterator for ChildIterator<'_> {
         self.current_child
     }
 }
+
+struct ParentsIterator<'a> {
+    current_node: Option<NodeId>,
+    next_node: NodeId,
+    tree: &'a LLTree<'a>,
+}
+
+impl<'a> ParentsIterator<'a> {
+    fn new(tree: &'a LLTree<'a>, u: NodeId) -> Self {
+        let u = match tsk_id_t::try_from(tree.treeseq.num_nodes_raw()) {
+            Ok(num_nodes) if u < num_nodes => u,
+            _ => NodeId::NULL,
+        };
+        ParentsIterator {
+            current_node: None,
+            next_node: u,
+            tree,
+        }
+    }
+}
+
+impl NodeIterator for ParentsIterator<'_> {
+    fn next_node(&mut self) {
+        self.current_node = match self.next_node {
+            NodeId::NULL => None,
+            r => {
+                assert!(r >= 0);
+                let cr = Some(r);
+                self.next_node = self.tree.parent(r).unwrap_or(NodeId::NULL);
+                cr
+            }
+        };
+    }
+
+    fn current_node(&mut self) -> Option<NodeId> {
+        self.current_node
+    }
+}

src/trees/tree.rs

@@ -114,7 +114,14 @@ impl<'treeseq> Tree<'treeseq> {
     ///
     /// Returns `None` if `u` is out of range.
     pub fn right_child<N: Into<NodeId> + Copy>(&self, u: N) -> Option<NodeId> {
-        self.inner.left_sib(u.into())
+        self.inner.right_child(u.into())
+    }
+
+    /// Get the left child of node `u`.
+    ///
+    /// Returns `None` if `u` is out of range.
+    pub fn left_child<N: Into<NodeId> + Copy>(&self, u: N) -> Option<NodeId> {
+        self.inner.left_child(u.into())
     }
 
     /// Get the number of samples below node `u`.
@@ -169,6 +176,15 @@ impl<'treeseq> Tree<'treeseq> {
     pub fn children<N: Into<NodeId> + Copy>(&self, u: N) -> impl Iterator<Item = NodeId> + '_ {
         self.inner.children(u.into())
     }
+
+    /// Return an [`Iterator`] over the parents of node `u`.
+    /// # Returns
+    ///
+    /// * `Some(iterator)` if `u` is valid
+    /// * `None` otherwise
+    pub fn parents<N: Into<NodeId> + Copy>(&self, u: N) -> impl Iterator<Item = NodeId> + '_ {
+        self.inner.parents(u.into())
+    }
 }
 
 impl<'ts> streaming_iterator::StreamingIterator for Tree<'ts> {