try to bring back branch length calcs

Author: molpopgen

src/tree_interface.rs

@@ -15,6 +15,7 @@ pub struct TreeInterface {
     num_nodes: tsk_size_t,
     array_len: tsk_size_t,
     flags: TreeFlags,
+    nodes: crate::NodeTable,
 }
 
 impl TreeInterface {
@@ -24,11 +25,16 @@ impl TreeInterface {
         array_len: tsk_size_t,
         flags: TreeFlags,
     ) -> Self {
+        let nodes = crate::NodeTable::new_from_table(
+            &mut unsafe { *(*non_owned_pointer.as_ref().tree_sequence).tables }.nodes,
+        )
+        .expect("null pointer to tsk_node_table_t");
         Self {
             non_owned_pointer,
             num_nodes,
             array_len,
             flags,
+            nodes,
         }
     }
 
@@ -340,7 +346,7 @@ impl TreeInterface {
     /// Get the parent of node `u`.
     ///
     /// Returns `None` if `u` is out of range.
-    pub fn parent<N: Into<NodeId> + Copy>(&self, u: N) -> Option<NodeId> {
+    pub fn parent<N: Into<NodeId> + Copy + std::fmt::Debug>(&self, u: N) -> Option<NodeId> {
         sys::tsk_column_access::<NodeId, _, _, _>(u.into(), self.as_ref().parent, self.array_len)
     }
 
@@ -489,11 +495,8 @@ impl TreeInterface {
     /// (and the tree sequence from which it came).
     ///
     /// This is a convenience function for accessing node times, etc..
-    pub fn node_table(&self) -> crate::NodeTable {
-        unimplemented!("this needs to return &NodeTable");
-        // crate::NodeTable::new_from_table(unsafe {
-        //     &(*(*(*self.as_ptr()).tree_sequence).tables).nodes
-        // })
+    fn node_table(&self) -> &crate::NodeTable {
+        &self.nodes
     }
 
     /// Calculate the total length of the tree via a preorder traversal.

tests/book_trees.rs

@@ -143,6 +143,12 @@ fn initialize_from_table_collection() {
     }
     // ANCHOR_END: iterate_node_siblings_via_array_getters
 
+    let mut tree_iterator = treeseq.tree_iterator(TreeFlags::default()).unwrap();
+    let mut total_branch_lengths = vec![];
+    while let Some(tree) = tree_iterator.next() {
+        total_branch_lengths.push(tree.total_branch_length(false).unwrap());
+    }
+
     // ANCHOR: iterate_edge_differences
     if let Ok(mut edge_diff_iterator) = treeseq.edge_differences_iter() {
         while let Some(diffs) = edge_diff_iterator.next() {