refactor: low level tables/tree sequences (#622)

* tree sequences prefer newtypes over raw C types
* improved mechanics for tables -> tree seq ownership transfer

Author: molpopgen

src/sys/table_collection.rs

@@ -105,4 +105,8 @@ impl TableCollection {
         // SAFETY: self pointer is not null
         unsafe { &mut (*self.as_mut_ptr()).sites }
     }
+
+    pub fn into_raw(self) -> *mut tsk_table_collection_t {
+        unsafe { self.0.into_raw() }
+    }
 }

src/sys/tree.rs

@@ -25,7 +25,7 @@ impl<'treeseq> LLTree<'treeseq> {
             let code = unsafe {
                 super::bindings::tsk_tree_set_tracked_samples(
                     inner.as_mut(),
-                    treeseq.num_samples(),
+                    treeseq.num_samples().into(),
                     (inner.as_mut()).samples,
                 )
             };

src/sys/treeseq.rs

@@ -11,9 +11,10 @@ pub struct TreeSequence(TskBox<bindings::tsk_treeseq_t>);
 
 impl TreeSequence {
     pub fn new(
-        tables: *mut bindings::tsk_table_collection_t,
+        tables: super::TableCollection,
         flags: super::flags::TreeSequenceFlags,
     ) -> Result<Self, TskitError> {
+        let tables = tables.into_raw();
         let inner = TskBox::new(|t: *mut bindings::tsk_treeseq_t| unsafe {
             tsk_treeseq_init(t, tables, flags.bits() | bindings::TSK_TAKE_OWNERSHIP)
         })?;
@@ -30,9 +31,9 @@ impl TreeSequence {
 
     pub fn simplify(
         &self,
-        samples: &[bindings::tsk_id_t],
+        samples: &[super::newtypes::NodeId],
         options: super::flags::SimplificationOptions,
-        idmap: *mut bindings::tsk_id_t,
+        idmap: Option<&mut [super::newtypes::NodeId]>,
     ) -> Result<Self, TskitError> {
         // The output is an UNINITIALIZED treeseq,
         // else we leak memory.
@@ -42,11 +43,15 @@ impl TreeSequence {
         let rv = unsafe {
             bindings::tsk_treeseq_simplify(
                 self.as_ref(),
-                samples.as_ptr(),
-                samples.len() as bindings::tsk_size_t,
+                // The cast is safe/sound b/c NodeId is repr(transparent)
+                samples.as_ptr().cast::<_>(),
+                samples.len().try_into().unwrap(),
                 options.bits(),
                 ts.as_mut_ptr(),
-                idmap,
+                match idmap {
+                    Some(s) => s.as_mut_ptr().cast::<_>(),
+                    None => std::ptr::null_mut(),
+                },
             )
         };
         if rv < 0 {
@@ -72,9 +77,9 @@ impl TreeSequence {
         }
     }
 
-    pub fn num_trees(&self) -> bindings::tsk_size_t {
+    pub fn num_trees(&self) -> super::newtypes::SizeType {
         // SAFETY: self pointer is not null
-        unsafe { bindings::tsk_treeseq_get_num_trees(self.as_ref()) }
+        unsafe { bindings::tsk_treeseq_get_num_trees(self.as_ref()) }.into()
     }
 
     pub fn num_nodes_raw(&self) -> bindings::tsk_size_t {
@@ -95,7 +100,7 @@ impl TreeSequence {
         }
     }
 
-    pub fn num_samples(&self) -> bindings::tsk_size_t {
-        unsafe { bindings::tsk_treeseq_get_num_samples(self.as_ref()) }
+    pub fn num_samples(&self) -> super::newtypes::SizeType {
+        unsafe { bindings::tsk_treeseq_get_num_samples(self.as_ref()) }.into()
     }
 }

src/table_collection.rs

@@ -109,13 +109,8 @@ impl TableCollection {
         })
     }
 
-    pub(crate) fn into_raw(self) -> Result<*mut ll_bindings::tsk_table_collection_t, TskitError> {
-        let mut tables = self;
-        let mut temp = crate::sys::TableCollection::new(1.)?;
-        std::mem::swap(&mut temp, &mut tables.inner);
-        let ptr = temp.as_mut_ptr();
-        std::mem::forget(temp);
-        handle_tsk_return_value!(0, ptr)
+    pub(crate) fn into_inner(self) -> crate::sys::TableCollection {
+        self.inner
     }
 
     /// Load a table collection from a file.

src/trees/treeseq.rs

@@ -9,7 +9,6 @@ use crate::TableOutputOptions;
 use crate::TreeFlags;
 use crate::TreeSequenceFlags;
 use crate::TskReturnValue;
-use ll_bindings::tsk_id_t;
 use sys::bindings as ll_bindings;
 
 use super::Tree;
@@ -113,7 +112,7 @@ impl TreeSequence {
         tables: TableCollection,
         flags: F,
     ) -> Result<Self, TskitError> {
-        let raw_tables_ptr = tables.into_raw()?;
+        let raw_tables_ptr = tables.into_inner();
         let mut inner = sys::TreeSequence::new(raw_tables_ptr, flags.into())?;
         let views = crate::table_views::TableViews::new_from_tree_sequence(inner.as_mut())?;
         Ok(Self { inner, views })
@@ -302,7 +301,7 @@ impl TreeSequence {
 
     /// Get the number of trees.
     pub fn num_trees(&self) -> SizeType {
-        self.inner.num_trees().into()
+        self.inner.num_trees()
     }
 
     /// Calculate the average Kendall-Colijn (`K-C`) distance between
@@ -322,7 +321,7 @@ impl TreeSequence {
 
     // FIXME: document
     pub fn num_samples(&self) -> SizeType {
-        self.inner.num_samples().into()
+        self.inner.num_samples()
     }
 
     /// Simplify tables and return a new tree sequence.
@@ -348,15 +347,12 @@ impl TreeSequence {
         if idmap {
             output_node_map.resize(usize::try_from(self.nodes().num_rows())?, NodeId::NULL);
         }
-        let llsamples = unsafe {
-            std::slice::from_raw_parts(samples.as_ptr().cast::<tsk_id_t>(), samples.len())
-        };
         let mut inner = self.inner.simplify(
-            llsamples,
+            samples,
             options.into(),
             match idmap {
-                true => output_node_map.as_mut_ptr().cast::<tsk_id_t>(),
-                false => std::ptr::null_mut(),
+                true => Some(&mut output_node_map),
+                false => None,
             },
         )?;
         let views = crate::table_views::TableViews::new_from_tree_sequence(inner.as_mut())?;