okay

Author: molpopgen

src/sys/traits.rs

@@ -1,3 +1,5 @@
+use super::bindings::tsk_id_t;
+
 /// For a type `tsk_foo_t`, this trait abstracts
 /// out the functionality of `tsk_foo_free`
 ///

src/sys/tree.rs

@@ -1,3 +1,5 @@
+use crate::bindings;
+
 use super::bindings::tsk_id_t;
 use super::bindings::tsk_size_t;
 use super::bindings::tsk_tree_t;
@@ -74,4 +76,114 @@ impl<'treeseq> LLTree<'treeseq> {
     pub fn as_ref(&self) -> &tsk_tree_t {
         self.inner.as_ref()
     }
+
+    pub fn left_sib<N: Into<tsk_id_t> + Copy>(&self, u: N) -> Option<tsk_id_t> {
+        todo!("doc SAFETY");
+        super::tsk_column_access::<tsk_id_t, _, _, _>(
+            u.into(),
+            self.as_ref().left_sib,
+            unsafe {
+                self.as_ref()
+                    .tree_sequence
+                    .as_ref()
+                    .unwrap()
+                    .tables
+                    .as_ref()
+            }
+            .unwrap()
+            .nodes
+            .num_rows,
+        )
+    }
+
+    pub fn right_child<N: Into<tsk_id_t> + Copy>(&self, u: N) -> Option<tsk_id_t> {
+        todo!("doc SAFETY");
+        super::tsk_column_access::<tsk_id_t, _, _, _>(
+            u.into(),
+            self.as_ref().right_child,
+            unsafe {
+                self.as_ref()
+                    .tree_sequence
+                    .as_ref()
+                    .unwrap()
+                    .tables
+                    .as_ref()
+            }
+            .unwrap()
+            .nodes
+            .num_rows,
+        )
+    }
+}
+
+// Trait defining iteration over nodes.
+pub trait NodeIterator {
+    fn next_node(&mut self);
+    fn current_node(&mut self) -> Option<tsk_id_t>;
+}
+
+struct NodeIteratorAdapter<T>
+where
+    T: NodeIterator,
+{
+    ni: T,
+}
+
+impl<T> Iterator for NodeIteratorAdapter<T>
+where
+    T: NodeIterator,
+{
+    type Item = tsk_id_t;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.ni.next_node();
+        self.ni.current_node()
+    }
+}
+
+struct PreorderNodeIterator<'a> {
+    current_root: tsk_id_t,
+    node_stack: Vec<tsk_id_t>,
+    tree: &'a LLTree<'a>,
+    current_node_: Option<tsk_id_t>,
+}
+
+impl<'a> PreorderNodeIterator<'a> {
+    fn new(tree: &'a LLTree) -> Self {
+        debug_assert!(tree.right_child(tree.virtual_root()).is_some());
+        let mut rv = PreorderNodeIterator {
+            current_root: tree.right_child(tree.virtual_root()).unwrap_or(-1),
+            node_stack: vec![],
+            tree,
+            current_node_: None,
+        };
+        let mut c = rv.current_root;
+        while c != -1 {
+            rv.node_stack.push(c);
+            debug_assert!(rv.tree.left_sib(c).is_some());
+            c = rv.tree.left_sib(c).unwrap_or(-1);
+        }
+        rv
+    }
+}
+
+impl NodeIterator for PreorderNodeIterator<'_> {
+    fn next_node(&mut self) {
+        self.current_node_ = self.node_stack.pop();
+        if let Some(u) = self.current_node_ {
+            // NOTE: process children right-to-left
+            // because we later pop them from a steck
+            // to generate the expected left-to-right ordering.
+            debug_assert!(self.tree.right_child(u).is_some());
+            let mut c = self.tree.right_child(u).unwrap_or(-1);
+            while c != NodeId::NULL {
+                self.node_stack.push(c);
+                debug_assert!(self.tree.right_child(c).is_some());
+                c = self.tree.left_sib(c).unwrap_or(-1);
+            }
+        };
+    }
+
+    fn current_node(&mut self) -> Option<tsk_id_t> {
+        self.current_node_
+    }
 }