feat: Add tskit::OwnedNodeTable (#273)

Author: molpopgen

src/lib.rs

@@ -428,7 +428,7 @@ pub use flags::*;
 pub use individual_table::{IndividualTable, IndividualTableRow};
 pub use migration_table::{MigrationTable, MigrationTableRow};
 pub use mutation_table::{MutationTable, MutationTableRow};
-pub use node_table::{NodeTable, NodeTableRow};
+pub use node_table::{NodeTable, NodeTableRow, OwnedNodeTable};
 pub use population_table::{OwnedPopulationTable, PopulationTable, PopulationTableRow};
 pub use site_table::{SiteTable, SiteTableRow};
 pub use table_collection::TableCollection;

src/node_table.rs

@@ -5,6 +5,7 @@ use crate::SizeType;
 use crate::Time;
 use crate::{tsk_id_t, TskitError};
 use crate::{IndividualId, NodeId, PopulationId};
+use ll_bindings::{tsk_node_table_free, tsk_node_table_init};
 
 /// Row of a [`NodeTable`]
 pub struct NodeTableRow {
@@ -243,3 +244,121 @@ impl<'a> NodeTable<'a> {
         samples
     }
 }
+
+/// A standalone node table that owns its data.
+///
+/// # Examples
+///
+/// ```
+/// use tskit::OwnedNodeTable;
+///
+/// let mut nodes = OwnedNodeTable::default();
+/// let rowid = nodes.add_row(0, 1.1, -1, -1).unwrap();
+/// assert_eq!(rowid, 0);
+/// assert_eq!(nodes.num_rows(), 1);
+/// ```
+///
+/// An example with metadata.
+/// This requires the cargo feature `"derive"` for `tskit`.
+///
+/// ```
+/// # #[cfg(any(feature="doc", feature="derive"))] {
+/// use tskit::OwnedNodeTable;
+///
+/// #[derive(serde::Serialize,
+///          serde::Deserialize,
+///          tskit::metadata::NodeMetadata)]
+/// #[serializer("serde_json")]
+/// struct NodeMetadata {
+///     value: i32,
+/// }
+///
+/// let metadata = NodeMetadata{value: 42};
+///
+/// let mut nodes = OwnedNodeTable::default();
+///
+/// let rowid = nodes.add_row_with_metadata(0, 1., -1, -1, &metadata).unwrap();
+/// assert_eq!(rowid, 0);
+///
+/// if let Some(decoded) = nodes.metadata::<NodeMetadata>(rowid).unwrap() {
+///     assert_eq!(decoded.value, 42);
+/// } else {
+///     panic!("hmm...we expected some metadata!");
+/// }
+///
+/// # }
+/// ```
+pub struct OwnedNodeTable {
+    table: mbox::MBox<ll_bindings::tsk_node_table_t>,
+}
+
+impl OwnedNodeTable {
+    pub fn add_row(
+        &mut self,
+        flags: impl Into<NodeFlags>,
+        time: impl Into<Time>,
+        population: impl Into<PopulationId>,
+        individual: impl Into<IndividualId>,
+    ) -> Result<NodeId, TskitError> {
+        let rv = unsafe {
+            ll_bindings::tsk_node_table_add_row(
+                &mut (*self.table),
+                flags.into().bits(),
+                time.into().0,
+                population.into().0,
+                individual.into().0,
+                std::ptr::null(),
+                0,
+            )
+        };
+
+        handle_tsk_return_value!(rv, rv.into())
+    }
+
+    pub fn add_row_with_metadata<M>(
+        &mut self,
+        flags: impl Into<NodeFlags>,
+        time: impl Into<Time>,
+        population: impl Into<PopulationId>,
+        individual: impl Into<IndividualId>,
+        metadata: &M,
+    ) -> Result<NodeId, TskitError>
+    where
+        M: crate::metadata::NodeMetadata,
+    {
+        let md = crate::metadata::EncodedMetadata::new(metadata)?;
+        let rv = unsafe {
+            ll_bindings::tsk_node_table_add_row(
+                &mut (*self.table),
+                flags.into().bits(),
+                time.into().0,
+                population.into().0,
+                individual.into().0,
+                md.as_ptr(),
+                md.len().into(),
+            )
+        };
+        handle_tsk_return_value!(rv, rv.into())
+    }
+}
+
+build_owned_tables!(
+    OwnedNodeTable,
+    NodeTable,
+    ll_bindings::tsk_node_table_t,
+    tsk_node_table_init,
+    tsk_node_table_free
+);
+
+#[cfg(test)]
+mod test_owned_node_table {
+    use super::*;
+
+    #[test]
+    fn test_add_row() {
+        let mut nodes = OwnedNodeTable::default();
+        let rowid = nodes.add_row(0, 1.1, -1, -1).unwrap();
+        assert_eq!(rowid, 0);
+        assert_eq!(nodes.num_rows(), 1);
+    }
+}