Merge pull request #58 from molpopgen/table_row_access

Add table row API

Author: molpopgen

src/_macros.rs

@@ -160,6 +160,23 @@ macro_rules! err_if_not_tracking_samples {
     };
 }
 
+// This macro assumes that table row access helper
+// functions have a standard interface.
+// Here, we convert the None type to an Error,
+// as it applies $row is out of range.
+macro_rules! table_row_access {
+    ($row: expr, $decode_metadata: expr, $table: expr, $row_fn: ident) => {
+        if $row < 0 {
+            Err(TskitError::IndexError)
+        } else {
+            match $row_fn($table, $row, $decode_metadata) {
+                Some(x) => Ok(x),
+                None => Err(TskitError::IndexError),
+            }
+        }
+    };
+}
+
 #[cfg(test)]
 mod test {
     use crate::error::TskitError;

src/edge_table.rs

@@ -2,6 +2,7 @@ use crate::bindings as ll_bindings;
 use crate::metadata;
 use crate::{tsk_id_t, tsk_size_t, TskitError};
 
+/// Row of an [`EdgeTable`]
 pub struct EdgeTableRow {
     pub left: f64,
     pub right: f64,
@@ -10,6 +11,16 @@ pub struct EdgeTableRow {
     pub metadata: Option<Vec<u8>>,
 }
 
+impl PartialEq for EdgeTableRow {
+    fn eq(&self, other: &Self) -> bool {
+        self.parent == other.parent
+            && self.child == other.child
+            && crate::util::f64_partial_cmp_equal(&self.left, &other.left)
+            && crate::util::f64_partial_cmp_equal(&self.right, &other.right)
+            && crate::util::metadata_like_are_equal(&self.metadata, &other.metadata)
+    }
+}
+
 fn make_edge_table_row(
     table: &EdgeTable,
     pos: tsk_id_t,
@@ -138,4 +149,21 @@ impl<'a> EdgeTable<'a> {
     pub fn iter(&self, decode_metadata: bool) -> EdgeTableRefIterator {
         crate::table_iterator::make_table_iterator::<&EdgeTable<'a>>(&self, decode_metadata)
     }
+
+    /// Return row `r` of the table.
+    ///
+    /// # Parameters
+    ///
+    /// * `r`: the row id.
+    /// * `decode_metadata`: if `true`, then a *copy* of row metadata
+    ///    will be provided in [`EdgeTableRow::metadata`].
+    ///    The meta data are *not* decoded.
+    ///    Rows with no metadata will contain the value `None`.
+    ///
+    /// # Errors
+    ///
+    /// [`TskitError::IndexError`] if `r` is out of range.
+    pub fn row(&self, r: tsk_id_t, decode_metadata: bool) -> Result<EdgeTableRow, TskitError> {
+        table_row_access!(r, decode_metadata, self, make_edge_table_row)
+    }
 }

src/lib.rs

@@ -19,6 +19,7 @@ mod table_collection;
 mod table_iterator;
 mod trees;
 pub mod types;
+mod util;
 
 // re-export fundamental constants that
 // we can't live without

src/mutation_table.rs

@@ -12,6 +12,17 @@ pub struct MutationTableRow {
     pub metadata: Option<Vec<u8>>,
 }
 
+impl PartialEq for MutationTableRow {
+    fn eq(&self, other: &Self) -> bool {
+        self.site == other.site
+            && self.node == other.node
+            && self.parent == other.parent
+            && crate::util::f64_partial_cmp_equal(&self.time, &other.time)
+            && crate::util::metadata_like_are_equal(&self.derived_state, &other.derived_state)
+            && crate::util::metadata_like_are_equal(&self.metadata, &other.metadata)
+    }
+}
+
 fn make_mutation_table_row(
     table: &MutationTable,
     pos: tsk_id_t,
@@ -160,4 +171,21 @@ impl<'a> MutationTable<'a> {
     pub fn iter(&self, decode_metadata: bool) -> MutationTableRefIterator {
         crate::table_iterator::make_table_iterator::<&MutationTable<'a>>(&self, decode_metadata)
     }
+
+    /// Return row `r` of the table.
+    ///
+    /// # Parameters
+    ///
+    /// * `r`: the row id.
+    /// * `decode_metadata`: if `true`, then a *copy* of row metadata
+    ///    will be provided in [`MutationTableRow::metadata`].
+    ///    The meta data are *not* decoded.
+    ///    Rows with no metadata will contain the value `None`.
+    ///
+    /// # Errors
+    ///
+    /// [`TskitError::IndexError`] if `r` is out of range.
+    pub fn row(&self, r: tsk_id_t, decode_metadata: bool) -> Result<MutationTableRow, TskitError> {
+        table_row_access!(r, decode_metadata, self, make_mutation_table_row)
+    }
 }

src/node_table.rs

@@ -11,6 +11,16 @@ pub struct NodeTableRow {
     pub metadata: Option<Vec<u8>>,
 }
 
+impl PartialEq for NodeTableRow {
+    fn eq(&self, other: &Self) -> bool {
+        self.flags == other.flags
+            && self.population == other.population
+            && self.individual == other.individual
+            && crate::util::metadata_like_are_equal(&self.metadata, &other.metadata)
+            && crate::util::f64_partial_cmp_equal(&self.time, &other.time)
+    }
+}
+
 fn make_node_table_row(
     table: &NodeTable,
     pos: tsk_id_t,
@@ -148,4 +158,21 @@ impl<'a> NodeTable<'a> {
     pub fn iter(&self, decode_metadata: bool) -> NodeTableRefIterator {
         crate::table_iterator::make_table_iterator::<&NodeTable<'a>>(&self, decode_metadata)
     }
+
+    /// Return row `r` of the table.
+    ///
+    /// # Parameters
+    ///
+    /// * `r`: the row id.
+    /// * `decode_metadata`: if `true`, then a *copy* of row metadata
+    ///    will be provided in [`NodeTableRow::metadata`].
+    ///    The meta data are *not* decoded.
+    ///    Rows with no metadata will contain the value `None`.
+    ///
+    /// # Errors
+    ///
+    /// [`TskitError::IndexError`] if `r` is out of range.
+    pub fn row(&self, r: tsk_id_t, decode_metadata: bool) -> Result<NodeTableRow, TskitError> {
+        table_row_access!(r, decode_metadata, self, make_node_table_row)
+    }
 }

src/population_table.rs

@@ -4,10 +4,17 @@ use crate::TskitError;
 use crate::{tsk_id_t, tsk_size_t};
 
 /// Row of a [`PopulationTable`]
+#[derive(Eq)]
 pub struct PopulationTableRow {
     pub metadata: Option<Vec<u8>>,
 }
 
+impl PartialEq for PopulationTableRow {
+    fn eq(&self, other: &Self) -> bool {
+        crate::util::metadata_like_are_equal(&self.metadata, &other.metadata)
+    }
+}
+
 fn make_population_table_row(
     table: &PopulationTable,
     pos: tsk_id_t,
@@ -93,4 +100,25 @@ impl<'a> PopulationTable<'a> {
     pub fn iter(&self, decode_metadata: bool) -> PopulationTableRefIterator {
         crate::table_iterator::make_table_iterator::<&PopulationTable<'a>>(&self, decode_metadata)
     }
+
+    /// Return row `r` of the table.
+    ///
+    /// # Parameters
+    ///
+    /// * `r`: the row id.
+    /// * `decode_metadata`: if `true`, then a *copy* of row metadata
+    ///    will be provided in [`PopulationTableRow::metadata`].
+    ///    The meta data are *not* decoded.
+    ///    Rows with no metadata will contain the value `None`.
+    ///
+    /// # Errors
+    ///
+    /// [`TskitError::IndexError`] if `r` is out of range.
+    pub fn row(
+        &self,
+        r: tsk_id_t,
+        decode_metadata: bool,
+    ) -> Result<PopulationTableRow, TskitError> {
+        table_row_access!(r, decode_metadata, self, make_population_table_row)
+    }
 }

src/site_table.rs

@@ -10,6 +10,14 @@ pub struct SiteTableRow {
     pub metadata: Option<Vec<u8>>,
 }
 
+impl PartialEq for SiteTableRow {
+    fn eq(&self, other: &Self) -> bool {
+        crate::util::f64_partial_cmp_equal(&self.position, &other.position)
+            && crate::util::metadata_like_are_equal(&self.ancestral_state, &other.ancestral_state)
+            && crate::util::metadata_like_are_equal(&self.metadata, &other.metadata)
+    }
+}
+
 fn make_site_table_row(
     table: &SiteTable,
     pos: tsk_id_t,
@@ -126,4 +134,21 @@ impl<'a> SiteTable<'a> {
     pub fn iter(&self, decode_metadata: bool) -> SiteTableRefIterator {
         crate::table_iterator::make_table_iterator::<&SiteTable<'a>>(&self, decode_metadata)
     }
+
+    /// Return row `r` of the table.
+    ///
+    /// # Parameters
+    ///
+    /// * `r`: the row id.
+    /// * `decode_metadata`: if `true`, then a *copy* of row metadata
+    ///    will be provided in [`SiteTableRow::metadata`].
+    ///    The meta data are *not* decoded.
+    ///    Rows with no metadata will contain the value `None`.
+    ///
+    /// # Errors
+    ///
+    /// [`TskitError::IndexError`] if `r` is out of range.
+    pub fn row(&self, r: tsk_id_t, decode_metadata: bool) -> Result<SiteTableRow, TskitError> {
+        table_row_access!(r, decode_metadata, self, make_site_table_row)
+    }
 }

src/table_collection.rs

@@ -865,4 +865,27 @@ mod test {
         let dumps = tables.deepcopy().unwrap();
         assert!(tables.equals(&dumps, 0));
     }
+
+    #[test]
+    fn test_edge_table_row_equality() {
+        let tables = make_small_table_collection();
+        for (i, row) in tables.edges_iter(true).enumerate() {
+            assert!(row == tables.edges().row(i as tsk_id_t, true).unwrap());
+            assert!(!(row != tables.edges().row(i as tsk_id_t, true).unwrap()));
+            if i > 0 {
+                assert!(row != tables.edges().row(i as tsk_id_t - 1, true).unwrap());
+            }
+        }
+    }
+
+    #[test]
+    fn test_node_table_row_equality() {
+        let tables = make_small_table_collection();
+        for (i, row) in tables.nodes_iter(true).enumerate() {
+            assert!(row == tables.nodes().row(i as tsk_id_t, true).unwrap());
+            assert!(!(row != tables.nodes().row(i as tsk_id_t, true).unwrap()));
+        }
+        assert!(tables.nodes().row(0, true) != tables.nodes().row(1, true));
+        assert!(tables.nodes().row(1, true) == tables.nodes().row(2, true));
+    }
 }

src/util.rs

@@ -0,0 +1,58 @@
+pub(crate) fn metadata_like_are_equal(a: &Option<Vec<u8>>, b: &Option<Vec<u8>>) -> bool {
+    match a {
+        Some(x) => match b {
+            Some(y) => x == y,
+            None => false,
+        },
+        None => match b.is_none() {
+            true => true,
+            false => false,
+        },
+    }
+}
+
+pub(crate) fn f64_partial_cmp_equal(a: &f64, b: &f64) -> bool {
+    match a.partial_cmp(b) {
+        Some(std::cmp::Ordering::Equal) => true,
+        Some(std::cmp::Ordering::Less) => false,
+        Some(std::cmp::Ordering::Greater) => false,
+        None => false,
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::metadata_like_are_equal;
+
+    #[test]
+    fn compare_some_to_none() {
+        let v: Vec<u8> = vec![1, 2, 3];
+        assert!(!metadata_like_are_equal(&Some(v), &None));
+    }
+
+    #[test]
+    fn compare_none_to_some() {
+        let v: Vec<u8> = vec![1, 2, 3];
+        assert!(!metadata_like_are_equal(&None, &Some(v)));
+    }
+
+    #[test]
+    fn compare_none_to_none() {
+        assert!(metadata_like_are_equal(&None, &None));
+    }
+
+    #[test]
+    fn compare_some_to_some_are_equal() {
+        let v: Vec<u8> = vec![1, 2, 3];
+        let vc = v.clone();
+        assert!(metadata_like_are_equal(&Some(v), &Some(vc)));
+    }
+
+    #[test]
+    fn compare_some_to_some_are_not_equal() {
+        let v: Vec<u8> = vec![1, 2, 3];
+        let mut vc = v.clone();
+        vc.push(11);
+        assert!(!metadata_like_are_equal(&Some(v), &Some(vc)));
+    }
+}