feature: Tree::total_branch_length now works. (#429)

* Replace unimplemented! macro with working code.
* Remove Tree::node_table.
  This is not a breaking change due to
  use of unimplemented!.

Author: molpopgen

examples/haploid_wright_fisher.rs

@@ -114,7 +114,25 @@ proptest! {
 #[test]
     fn test_simulate_proptest(seed in any::<u64>(),
                               num_generations in 50..100i32,
-                              simplify_interval in 1..100i32 ) {
-        let _ = simulate(seed, 100, num_generations, simplify_interval).unwrap();
+                              simplify_interval in 1..100i32) {
+        let ts = simulate(seed, 100, num_generations, simplify_interval).unwrap();
+
+        // stress test the branch length fn b/c it is not a trivial
+        // wrapper around the C API.
+        {
+            use streaming_iterator::StreamingIterator;
+            let mut x = f64::NAN;
+            if let Ok(mut tree_iter) = ts.tree_iterator(0) {
+                // We will only do the first tree to save time.
+                if let Some(tree) = tree_iter.next() {
+                    let b = tree.total_branch_length(false).unwrap();
+                    let b2 = unsafe {
+                        tskit::bindings::tsk_tree_get_total_branch_length(tree.as_ptr(), -1, &mut x)
+                    };
+                    assert!(b2 >= 0, "{}", b2);
+                    assert!(f64::from(b) - x <= 1e-8);
+                }
+            }
+        }
     }
 }

src/tree_interface.rs

@@ -340,7 +340,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,12 +489,9 @@ 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.
     ///
@@ -506,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

@@ -143,6 +143,29 @@ 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());
+    }
+
+    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() {