add tests

Author: molpopgen

src/tree_interface.rs

@@ -15,7 +15,6 @@ pub struct TreeInterface {
     num_nodes: tsk_size_t,
     array_len: tsk_size_t,
     flags: TreeFlags,
-    nodes: crate::NodeTable,
 }
 
 impl TreeInterface {
@@ -25,16 +24,11 @@ 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,
         }
     }
 
@@ -495,9 +489,9 @@ impl TreeInterface {
     /// (and the tree sequence from which it came).
     ///
     /// This is a convenience function for accessing node times, etc..
-    fn node_table(&self) -> &crate::NodeTable {
-        &self.nodes
-    }
+    //    fn node_table(&self) -> &crate::NodeTable {
+    //       &self.nodes
+    //  }
 
     /// Calculate the total length of the tree via a preorder traversal.
     ///
@@ -509,13 +503,19 @@ impl TreeInterface {
     ///
     /// [`TskitError`] may be returned if a node index is out of range.
     pub fn total_branch_length(&self, by_span: bool) -> Result<Time, TskitError> {
-        let nt = self.node_table();
+        let time: &[Time] = sys::generate_slice(
+            unsafe {
+                (*(*(*self.non_owned_pointer.as_ptr()).tree_sequence).tables)
+                    .nodes
+                    .time
+            },
+            self.num_nodes,
+        );
         let mut b = Time::from(0.);
         for n in self.traverse_nodes(NodeTraversalOrder::Preorder) {
             let p = self.parent(n).ok_or(TskitError::IndexError {})?;
             if p != NodeId::NULL {
-                b += nt.time(p).ok_or(TskitError::IndexError {})?
-                    - nt.time(n).ok_or(TskitError::IndexError {})?;
+                b += time[p.as_usize()] - time[n.as_usize()]
             }
         }
 

tests/book_trees.rs

@@ -149,6 +149,23 @@ fn initialize_from_table_collection() {
         total_branch_lengths.push(tree.total_branch_length(false).unwrap());
     }
 
+    let mut tree_iterator = treeseq.tree_iterator(TreeFlags::default()).unwrap();
+    let mut total_branch_lengths_ll = vec![];
+    let mut x = 0.0;
+    while let Some(tree) = tree_iterator.next() {
+        let l =
+            unsafe { tskit::bindings::tsk_tree_get_total_branch_length(tree.as_ptr(), -1, &mut x) };
+        assert!(l >= 0);
+        total_branch_lengths_ll.push(x);
+    }
+
+    for (i, j) in total_branch_lengths
+        .iter()
+        .zip(total_branch_lengths_ll.iter())
+    {
+        assert_eq!(i, j, "{} {}", i, j);
+    }
+
     // ANCHOR: iterate_edge_differences
     if let Ok(mut edge_diff_iterator) = treeseq.edge_differences_iter() {
         while let Some(diffs) = edge_diff_iterator.next() {