Add PopulationId (#135)

Author: molpopgen

src/lib.rs

@@ -121,7 +121,7 @@ pub use bindings::tsk_size_t;
 /// assert_eq!(y, z);
 /// ```
 ///
-/// It is also possible to write functions accepting both the `NodeId` 
+/// It is also possible to write functions accepting both the `NodeId`
 /// and `tsk_id_t`:
 ///
 /// ```
@@ -151,8 +151,18 @@ pub struct NodeId(tsk_id_t);
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, std::hash::Hash)]
 pub struct IndividualId(tsk_id_t);
 
+/// A population ID
+///
+/// This is an integer referring to a row of an [``PopulationTable``].
+///
+/// The features for this type follow the same pattern as for [``NodeId``]
+#[repr(transparent)]
+#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, std::hash::Hash)]
+pub struct PopulationId(tsk_id_t);
+
 impl_id_traits!(NodeId);
 impl_id_traits!(IndividualId);
+impl_id_traits!(PopulationId);
 
 // tskit defines this via a type cast
 // in a macro. bindgen thus misses it.

src/node_table.rs

@@ -1,14 +1,14 @@
 use crate::bindings as ll_bindings;
 use crate::metadata;
-use crate::NodeId;
 use crate::{tsk_flags_t, tsk_id_t, TskitError};
+use crate::{NodeId, PopulationId};
 
 /// Row of a [`NodeTable`]
 pub struct NodeTableRow {
     pub id: NodeId,
     pub time: f64,
     pub flags: tsk_flags_t,
-    pub population: tsk_id_t,
+    pub population: PopulationId,
     pub individual: tsk_id_t,
     pub metadata: Option<Vec<u8>>,
 }
@@ -117,8 +117,17 @@ impl<'a> NodeTable<'a> {
     ///
     /// Will return [``IndexError``](crate::TskitError::IndexError)
     /// if ``row`` is out of range.
-    pub fn population<N: Into<NodeId> + Copy>(&'a self, row: N) -> Result<tsk_id_t, TskitError> {
-        unsafe_tsk_column_access!(row.into().0, 0, self.num_rows(), self.table_.population)
+    pub fn population<N: Into<NodeId> + Copy>(
+        &'a self,
+        row: N,
+    ) -> Result<PopulationId, TskitError> {
+        unsafe_tsk_column_access!(
+            row.into().0,
+            0,
+            self.num_rows(),
+            self.table_.population,
+            PopulationId
+        )
     }
 
     /// Return the ``population`` value from row ``row`` of the table.
@@ -127,7 +136,7 @@ impl<'a> NodeTable<'a> {
     ///
     /// Will return [``IndexError``](crate::TskitError::IndexError)
     /// if ``row`` is out of range.
-    pub fn deme<N: Into<NodeId> + Copy>(&'a self, row: N) -> Result<tsk_id_t, TskitError> {
+    pub fn deme<N: Into<NodeId> + Copy>(&'a self, row: N) -> Result<PopulationId, TskitError> {
         self.population(row)
     }
 

src/population_table.rs

@@ -1,12 +1,13 @@
 use crate::bindings as ll_bindings;
 use crate::metadata;
+use crate::PopulationId;
 use crate::TskitError;
 use crate::{tsk_id_t, tsk_size_t};
 
 /// Row of a [`PopulationTable`]
 #[derive(Eq)]
 pub struct PopulationTableRow {
-    pub id: tsk_id_t,
+    pub id: PopulationId,
     pub metadata: Option<Vec<u8>>,
 }
 
@@ -19,7 +20,7 @@ impl PartialEq for PopulationTableRow {
 fn make_population_table_row(table: &PopulationTable, pos: tsk_id_t) -> Option<PopulationTableRow> {
     if pos < table.num_rows() as tsk_id_t {
         let rv = PopulationTableRow {
-            id: pos,
+            id: pos.into(),
             metadata: table_row_decode_metadata!(table, pos),
         };
         Some(rv)
@@ -71,11 +72,11 @@ impl<'a> PopulationTable<'a> {
         self.table_.num_rows
     }
 
-    pub fn metadata<T: metadata::MetadataRoundtrip>(
+    pub fn metadata<P: Into<PopulationId>, T: metadata::MetadataRoundtrip>(
         &'a self,
-        row: tsk_id_t,
+        row: P,
     ) -> Result<Option<T>, TskitError> {
-        let buffer = metadata_to_vector!(self, row)?;
+        let buffer = metadata_to_vector!(self, row.into().0)?;
         decode_metadata_row!(T, buffer)
     }
 
@@ -94,7 +95,10 @@ impl<'a> PopulationTable<'a> {
     /// # Errors
     ///
     /// [`TskitError::IndexError`] if `r` is out of range.
-    pub fn row(&self, r: tsk_id_t) -> Result<PopulationTableRow, TskitError> {
-        table_row_access!(r, self, make_population_table_row)
+    pub fn row<P: Into<PopulationId> + Copy>(
+        &self,
+        r: P,
+    ) -> Result<PopulationTableRow, TskitError> {
+        table_row_access!(r.into().0, self, make_population_table_row)
     }
 }

src/table_collection.rs

@@ -20,7 +20,7 @@ use crate::TreeSequenceFlags;
 use crate::TskReturnValue;
 use crate::TskitTypeAccess;
 use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TSK_NULL};
-use crate::{IndividualId, NodeId};
+use crate::{IndividualId, NodeId, PopulationId};
 use ll_bindings::tsk_table_collection_free;
 
 /// A table collection.
@@ -241,11 +241,11 @@ impl TableCollection {
     ///
     /// Migration tables are not currently supported
     /// by tree sequence simplification.
-    pub fn add_migration<N: Into<NodeId>>(
+    pub fn add_migration<N: Into<NodeId>, SOURCE: Into<PopulationId>, DEST: Into<PopulationId>>(
         &mut self,
         span: (f64, f64),
         node: N,
-        source_dest: (tsk_id_t, tsk_id_t),
+        source_dest: (SOURCE, DEST),
         time: f64,
     ) -> TskReturnValue {
         self.add_migration_with_metadata(span, node, source_dest, time, None)
@@ -257,11 +257,15 @@ impl TableCollection {
     ///
     /// Migration tables are not currently supported
     /// by tree sequence simplification.
-    pub fn add_migration_with_metadata<N: Into<NodeId>>(
+    pub fn add_migration_with_metadata<
+        N: Into<NodeId>,
+        SOURCE: Into<PopulationId>,
+        DEST: Into<PopulationId>,
+    >(
         &mut self,
         span: (f64, f64),
         node: N,
-        source_dest: (tsk_id_t, tsk_id_t),
+        source_dest: (SOURCE, DEST),
         time: f64,
         metadata: Option<&dyn MetadataRoundtrip>,
     ) -> TskReturnValue {
@@ -272,8 +276,8 @@ impl TableCollection {
                 span.0,
                 span.1,
                 node.into().0,
-                source_dest.0,
-                source_dest.1,
+                source_dest.0.into().0,
+                source_dest.1.into().0,
                 time,
                 md.as_ptr(),
                 md.len(),
@@ -283,22 +287,22 @@ impl TableCollection {
     }
 
     /// Add a row to the node table
-    pub fn add_node<I: Into<IndividualId>>(
+    pub fn add_node<I: Into<IndividualId>, POP: Into<PopulationId>>(
         &mut self,
         flags: ll_bindings::tsk_flags_t,
         time: f64,
-        population: tsk_id_t,
+        population: POP,
         individual: I,
     ) -> Result<NodeId, TskitError> {
         self.add_node_with_metadata(flags, time, population, individual, None)
     }
 
     /// Add a row with metadata to the node table
-    pub fn add_node_with_metadata<I: Into<IndividualId>>(
+    pub fn add_node_with_metadata<I: Into<IndividualId>, POP: Into<PopulationId>>(
         &mut self,
         flags: ll_bindings::tsk_flags_t,
         time: f64,
-        population: tsk_id_t,
+        population: POP,
         individual: I,
         metadata: Option<&dyn MetadataRoundtrip>,
     ) -> Result<NodeId, TskitError> {
@@ -308,7 +312,7 @@ impl TableCollection {
                 &mut (*self.as_mut_ptr()).nodes,
                 flags,
                 time,
-                population,
+                population.into().0,
                 individual.into().0,
                 md.as_ptr(),
                 md.len(),
@@ -390,15 +394,15 @@ impl TableCollection {
     }
 
     /// Add a row to the population_table
-    pub fn add_population(&mut self) -> TskReturnValue {
+    pub fn add_population(&mut self) -> Result<PopulationId, TskitError> {
         self.add_population_with_metadata(None)
     }
 
     /// Add a row with metadata to the population_table
     pub fn add_population_with_metadata(
         &mut self,
         metadata: Option<&dyn MetadataRoundtrip>,
-    ) -> TskReturnValue {
+    ) -> Result<PopulationId, TskitError> {
         let md = EncodedMetadata::new(metadata)?;
         let rv = unsafe {
             ll_bindings::tsk_population_table_add_row(
@@ -408,7 +412,7 @@ impl TableCollection {
             )
         };
 
-        handle_tsk_return_value!(rv)
+        handle_tsk_return_value!(rv, PopulationId::from(rv))
     }
 
     /// Build the "input" and "output"
@@ -1030,30 +1034,33 @@ mod test {
     #[test]
     fn test_add_population() {
         let mut tables = TableCollection::new(1000.).unwrap();
-        tables.add_population().unwrap();
+        let pop_id = tables.add_population().unwrap();
+        assert_eq!(pop_id, 0);
         assert_eq!(tables.populations().num_rows(), 1);
+
+        tables
+            .add_node(crate::TSK_NODE_IS_SAMPLE, 0.0, pop_id, crate::TSK_NULL)
+            .unwrap();
+
+        match tables.nodes().row(NodeId::from(0)) {
+            Ok(x) => match x.population {
+                PopulationId(0) => (),
+                _ => panic!("expected PopulationId(0)"),
+            },
+            Err(_) => panic!("expected Ok(_)"),
+        };
     }
 
     #[test]
     fn test_dump_tables() {
         let treefile = "trees.trees";
         let mut tables = TableCollection::new(1000.).unwrap();
-        tables.add_population().unwrap();
+        let pop_id = tables.add_population().unwrap();
         tables
-            .add_node(
-                crate::TSK_NODE_IS_SAMPLE,
-                0.0,
-                crate::TSK_NULL,
-                crate::TSK_NULL,
-            )
+            .add_node(crate::TSK_NODE_IS_SAMPLE, 0.0, pop_id, crate::TSK_NULL)
             .unwrap();
         tables
-            .add_node(
-                crate::TSK_NODE_IS_SAMPLE,
-                1.0,
-                crate::TSK_NULL,
-                crate::TSK_NULL,
-            )
+            .add_node(crate::TSK_NODE_IS_SAMPLE, 1.0, pop_id, crate::TSK_NULL)
             .unwrap();
         tables.add_edge(0., tables.sequence_length(), 1, 0).unwrap();
         tables