streamline the postorder iterator

Author: molpopgen

src/trees.rs

@@ -686,87 +686,63 @@ impl NodeIterator for PreorderNodeIterator<'_> {
 iterator_for_nodeiterator!(PreorderNodeIterator<'_>);
 
 struct PostorderNodeIterator<'a> {
-    //root_stack: Vec<NodeId>,
     nodes: Vec<NodeId>,
-    tree: &'a Tree,
     current_node_index: usize,
     num_nodes_current_tree: usize,
+    // Make the lifetime checker happy.
+    tree: std::marker::PhantomData<&'a Tree>,
 }
 
 impl<'a> PostorderNodeIterator<'a> {
     fn new(tree: &'a Tree) -> Self {
-        let mut rv = PostorderNodeIterator {
-            //root_stack: tree.roots_to_vec(),
-            nodes: vec![
+        let mut num_nodes_current_tree: usize = 0;
+        let ptr = std::ptr::addr_of_mut!(num_nodes_current_tree);
+        let mut nodes = vec![
                 NodeId::NULL;
                 // NOTE: this fn does not return error codes
                 unsafe { ll_bindings::tsk_tree_get_size_bound(tree.inner) } as usize
-            ],
-            tree,
-            current_node_index: 0,
-            num_nodes_current_tree: 0,
-        };
-        //rv.root_stack.reverse();
-        // if let Some(root) = rv.root_stack.pop() {
-        //     // hmm--iterator setup shouldn't be
-        //     // able to fail...
-        //     // The error codes can be:
-        //     // * out of memory
-        //     // * node out of bounds, which "can't be" possible here?
-        //     let _ret = rv.update_nodes(NodeId::NULL);
-        // }
-        let _ret = rv.update_nodes(NodeId::NULL);
-        rv
-    }
+            ];
 
-    fn update_nodes(&mut self, u: NodeId) -> i32 {
-        let ptr = std::ptr::addr_of_mut!(self.num_nodes_current_tree);
-
-        unsafe {
+        let rv = unsafe {
             ll_bindings::tsk_tree_postorder(
-                self.tree.inner,
-                u.into(),
-                self.nodes.as_mut_ptr() as *mut tsk_id_t,
+                tree.inner,
+                NodeId::NULL.into(), // start from virtual root
+                nodes.as_mut_ptr() as *mut tsk_id_t,
                 ptr as *mut tsk_size_t,
             )
+        };
+
+        // This is either out of memory
+        // or node out of range.
+        // The former is fatal, and the latter
+        // not relevant (for now), as we start at roots.
+        if rv < 0 {
+            panic!("fatal error calculating postoder node list");
         }
-    }
-}
 
-impl NodeIterator for PostorderNodeIterator<'_> {
-    fn next_node(&mut self) {
-        //match self.current_node_index < self.num_nodes_current_tree {
-        //    true => (),
-        //    false => {
-        //        match self.root_stack.pop() {
-        //            Some(root) => {
-        //                self.current_node_index = 0;
-        //                //FIXME: failure should not be an option.
-        //                let _ret = self.update_nodes(root);
-        //            }
-        //            None => {
-        //                // should be sufficient to terminate iteration
-        //                self.current_node_index = 0;
-        //                self.num_nodes_current_tree = 0;
-        //            }
-        //        }
-        //    }
-        //}
+        Self {
+            nodes,
+            current_node_index: 0,
+            num_nodes_current_tree,
+            tree: std::marker::PhantomData,
+        }
     }
+}
 
-    fn current_node(&mut self) -> Option<NodeId> {
-        if self.current_node_index < self.num_nodes_current_tree {
-            let rv = Some(self.nodes[self.current_node_index]);
-            self.current_node_index += 1;
-            rv
-        } else {
-            None
+impl<'a> Iterator for PostorderNodeIterator<'a> {
+    type Item = NodeId;
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.current_node_index < self.num_nodes_current_tree {
+            true => {
+                let rv = self.nodes[self.current_node_index];
+                self.current_node_index += 1;
+                Some(rv)
+            }
+            false => None,
         }
     }
 }
 
-iterator_for_nodeiterator!(PostorderNodeIterator<'_>);
-
 struct RootIterator<'a> {
     current_root: Option<NodeId>,
     next_root: NodeId,