more

Author: molpopgen

src/sys/tree.rs

@@ -96,6 +96,42 @@ impl<'treeseq> LLTree<'treeseq> {
         )
     }
 
+    pub fn right_sib(&self, u: NodeId) -> Option<NodeId> {
+        super::tsk_column_access::<NodeId, _, _, _>(
+            u,
+            self.as_ref().right_sib,
+            unsafe {
+                self.as_ref()
+                    .tree_sequence
+                    .as_ref()
+                    .unwrap()
+                    .tables
+                    .as_ref()
+            }
+            .unwrap()
+            .nodes
+            .num_rows,
+        )
+    }
+
+    pub fn left_child(&self, u: NodeId) -> Option<NodeId> {
+        super::tsk_column_access::<NodeId, _, _, _>(
+            u,
+            self.as_ref().left_child,
+            unsafe {
+                self.as_ref()
+                    .tree_sequence
+                    .as_ref()
+                    .unwrap()
+                    .tables
+                    .as_ref()
+            }
+            .unwrap()
+            .nodes
+            .num_rows,
+        )
+    }
+
     pub fn right_child(&self, u: NodeId) -> Option<NodeId> {
         super::tsk_column_access::<NodeId, _, _, _>(
             u,
@@ -169,6 +205,10 @@ impl<'treeseq> LLTree<'treeseq> {
             }
         }
     }
+
+    pub fn children(&self, u: NodeId) -> impl Iterator<Item = NodeId> + '_ {
+        NodeIteratorAdapter(ChildIterator::new(self, u))
+    }
 }
 
 // Trait defining iteration over nodes.

src/trees/tree.rs

@@ -160,6 +160,15 @@ impl<'treeseq> Tree<'treeseq> {
     ) -> Box<dyn Iterator<Item = NodeId> + '_> {
         self.inner.traverse_nodes(order)
     }
+
+    /// Return an [`Iterator`] over the children of node `u`.
+    /// # Returns
+    ///
+    /// * `Some(iterator)` if `u` is valid
+    /// * `None` otherwise
+    pub fn children<N: Into<NodeId> + Copy>(&self, u: N) -> impl Iterator<Item = NodeId> + '_ {
+        self.inner.children(u.into())
+    }
 }
 
 impl<'ts> streaming_iterator::StreamingIterator for Tree<'ts> {