Merge pull request #33 from molpopgen/add_tree_sequence

Add TreeSequence

Author: molpopgen

src/_macros.rs

@@ -58,6 +58,35 @@ macro_rules! build_tskit_type {
     };
 }
 
+macro_rules! build_consuming_tskit_type {
+    ($name: ident, $ll_name: ty, $drop: ident, $consumed: ty) => {
+        impl Drop for $name {
+            fn drop(&mut self) {
+                let rv = unsafe { $drop(&mut *self.inner) };
+                panic_on_tskit_error!(rv);
+            }
+        }
+
+        impl crate::ffi::TskitConsumingType<$ll_name, $consumed> for $name {
+            fn wrap(consumed: $consumed) -> Self {
+                let temp: std::mem::MaybeUninit<$ll_name> = std::mem::MaybeUninit::uninit();
+                $name {
+                    consumed,
+                    inner: unsafe { Box::<$ll_name>::new(temp.assume_init()) },
+                }
+            }
+
+            fn as_ptr(&self) -> *const $ll_name {
+                &*self.inner
+            }
+
+            fn as_mut_ptr(&mut self) -> *mut $ll_name {
+                &mut *self.inner
+            }
+        }
+    };
+}
+
 macro_rules! metadata_to_vector {
     ($T: ty, $self: expr, $row: expr) => {
         crate::metadata::char_column_to_vector(

src/ffi.rs

@@ -15,6 +15,28 @@ pub trait TskitType<T> {
     fn as_mut_ptr(&mut self) -> *mut T;
 }
 
+/// Define what it means to wrap a tskit struct
+/// that contains another tskit type C (the
+/// "consumed" type).
+/// This trait models a type that takes another
+/// type as input for initialization and effectively
+/// owns it.
+/// A key example of such a type is a [`crate::TreeSequence`],
+/// which owns the underying [`crate::TableCollection`].
+/// In practice, one needs to implement Drop for
+/// test types, calling the tsk_foo_free() function
+/// corresponding to tsk_foo_t.
+pub trait TskitConsumingType<T, C> {
+    /// Encapsulate tsk_foo_t and return rust
+    /// object.  Best practices seem to
+    /// suggest using Box for this.
+    fn wrap(consumed: C) -> Self;
+    /// Return const pointer
+    fn as_ptr(&self) -> *const T;
+    /// Return mutable pointer
+    fn as_mut_ptr(&mut self) -> *mut T;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/lib.rs

@@ -16,6 +16,7 @@ mod node_table;
 mod population_table;
 mod site_table;
 mod table_collection;
+mod trees;
 pub mod types;
 
 // re-export fundamental constants that
@@ -42,6 +43,8 @@ pub use node_table::NodeTable;
 pub use population_table::PopulationTable;
 pub use site_table::SiteTable;
 pub use table_collection::TableCollection;
+pub use trees::TreeSequence;
+
 /// Handles return codes from low-level tskit functions.
 ///
 /// When an error from the tskit C API is detected,

src/table_collection.rs

@@ -375,11 +375,44 @@ impl TableCollection {
     pub fn equals(&self, other: &TableCollection, options: tsk_flags_t) -> bool {
         unsafe { ll_bindings::tsk_table_collection_equals(self.as_ptr(), other.as_ptr(), options) }
     }
+
+    /// Return a "deep" copy of the tables.
+    pub fn deepcopy(&self) -> Result<TableCollection, TskitError> {
+        let mut copy = TableCollection::new(1.)?;
+
+        let rv =
+            unsafe { ll_bindings::tsk_table_collection_copy(self.as_ptr(), copy.as_mut_ptr(), 0) };
+
+        if rv < 0 {
+            return Err(crate::error::TskitError::ErrorCode { code: rv });
+        }
+
+        Ok(copy)
+    }
+
+    /// Return a [`crate::TreeSequence`] based on the tables.
+    /// This function will raise errors if tables are not sorted,
+    /// not indexed, or invalid in any way.
+    pub fn tree_sequence(self) -> Result<crate::TreeSequence, TskitError> {
+        crate::TreeSequence::new(self)
+    }
 }
 
 #[cfg(test)]
 mod test {
     use super::*;
+    use crate::TSK_NULL;
+
+    fn make_small_table_collection() -> TableCollection {
+        let mut tables = TableCollection::new(1000.).unwrap();
+        tables.add_node(0, 1.0, TSK_NULL, TSK_NULL).unwrap();
+        tables.add_node(0, 0.0, TSK_NULL, TSK_NULL).unwrap();
+        tables.add_node(0, 0.0, TSK_NULL, TSK_NULL).unwrap();
+        tables.add_edge(0., 1000., 0, 1).unwrap();
+        tables.add_edge(0., 1000., 0, 2).unwrap();
+        tables.build_index(0).unwrap();
+        tables
+    }
 
     #[test]
     fn test_sequence_length() {
@@ -577,4 +610,11 @@ mod test {
         let mut tables = TableCollection::new(1000.).unwrap();
         tables.free().unwrap();
     }
+
+    #[test]
+    fn test_deepcopy() {
+        let tables = make_small_table_collection();
+        let dumps = tables.deepcopy().unwrap();
+        assert!(tables.equals(&dumps, 0));
+    }
 }

src/trees.rs

@@ -0,0 +1,85 @@
+use crate::bindings as ll_bindings;
+use crate::error::TskitError;
+use crate::ffi::{TskitConsumingType, TskitType};
+use crate::{tsk_flags_t, TableCollection, TSK_NULL};
+use ll_bindings::tsk_treeseq_free;
+
+/// A tree sequence.
+///
+/// This is a thin wrapper around the C type `tsk_treeseq_t`.
+///
+/// When created from a [`TableCollection`], the input tables are
+/// moved into the `TreeSequence` object.
+/// # Examples
+///
+/// ```
+/// let mut tables = tskit::TableCollection::new(1000.).unwrap();
+/// tables.add_node(0, 1.0, tskit::TSK_NULL, tskit::TSK_NULL).unwrap();
+/// tables.add_node(0, 0.0, tskit::TSK_NULL, tskit::TSK_NULL).unwrap();
+/// tables.add_node(0, 0.0, tskit::TSK_NULL, tskit::TSK_NULL).unwrap();
+/// tables.add_edge(0., 1000., 0, 1).unwrap();
+/// tables.add_edge(0., 1000., 0, 2).unwrap();
+///
+/// // tables gets moved into our treeseq variable:
+/// let treeseq = tables.tree_sequence();
+/// ```
+pub struct TreeSequence {
+    consumed: TableCollection,
+    inner: Box<ll_bindings::tsk_treeseq_t>,
+}
+
+build_consuming_tskit_type!(
+    TreeSequence,
+    ll_bindings::tsk_treeseq_t,
+    tsk_treeseq_free,
+    TableCollection
+);
+
+impl TreeSequence {
+    /// Create a tree sequence from a [`TableCollection`].
+    /// In general, [`TableCollection::tree_sequence`] may be preferred.
+    /// The table collection is moved/consumed.
+    pub fn new(tables: TableCollection) -> Result<Self, TskitError> {
+        let mut treeseq = Self::wrap(tables);
+        let rv = unsafe {
+            ll_bindings::tsk_treeseq_init(treeseq.as_mut_ptr(), treeseq.consumed.as_ptr(), 0)
+        };
+        if rv < 0 {
+            return Err(crate::error::TskitError::ErrorCode { code: rv });
+        }
+        Ok(treeseq)
+    }
+
+    /// Obtain a copy of the [`TableCollection`]
+    pub fn dump_tables(&self) -> Result<TableCollection, TskitError> {
+        self.consumed.deepcopy()
+    }
+}
+
+#[cfg(test)]
+mod test_trees {
+    use super::*;
+
+    fn make_small_table_collection() -> TableCollection {
+        let mut tables = TableCollection::new(1000.).unwrap();
+        tables.add_node(0, 1.0, TSK_NULL, TSK_NULL).unwrap();
+        tables.add_node(0, 0.0, TSK_NULL, TSK_NULL).unwrap();
+        tables.add_node(0, 0.0, TSK_NULL, TSK_NULL).unwrap();
+        tables.add_edge(0., 1000., 0, 1).unwrap();
+        tables.add_edge(0., 1000., 0, 2).unwrap();
+        tables.build_index(0).unwrap();
+        tables
+    }
+
+    #[test]
+    fn test_create_treeseq_new_from_tables() {
+        let tables = make_small_table_collection();
+        let _treeseq = TreeSequence::new(tables).unwrap();
+    }
+
+    #[test]
+    fn test_create_treeseq_from_tables() {
+        let tables = make_small_table_collection();
+        let _treeseq = tables.tree_sequence().unwrap();
+    }
+}