Merge pull request #61 from molpopgen/add_individual_and_migration_tables

Add individual and migration tables

Author: molpopgen

src/_macros.rs

@@ -37,6 +37,32 @@ macro_rules! unsafe_tsk_column_access {
     }};
 }
 
+macro_rules! unsafe_tsk_ragged_column_access {
+    ($i: expr, $lo: expr, $hi: expr, $array: expr, $offset_array: expr, $offset_array_len: expr) => {{
+        if $i < $lo || $i >= ($hi as tsk_id_t) {
+            Err(TskitError::IndexError {})
+        } else if $offset_array_len == 0 {
+            Ok(None)
+        } else {
+            let start = unsafe { *$offset_array.offset($i as isize) };
+            let stop = if $i < ($hi as tsk_id_t) {
+                unsafe { *$offset_array.offset(($i + 1) as isize) }
+            } else {
+                $offset_array_len as tsk_size_t
+            };
+            if start == stop {
+                Ok(None)
+            } else {
+                let mut buffer = vec![];
+                for i in start..stop {
+                    buffer.push(unsafe { *$array.offset(i as isize) });
+                }
+                Ok(Some(buffer))
+            }
+        }
+    }};
+}
+
 macro_rules! drop_for_tskit_type {
     ($name: ident, $drop: ident) => {
         impl Drop for $name {

src/individual_table.rs

@@ -0,0 +1,176 @@
+use crate::bindings as ll_bindings;
+use crate::metadata;
+use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TskitError};
+
+/// Row of a [`IndividualTable`]
+pub struct IndividualTableRow {
+    pub flags: u32,
+    pub location: Option<Vec<f64>>,
+    pub parents: Option<Vec<tsk_id_t>>,
+    pub metadata: Option<Vec<u8>>,
+}
+
+impl PartialEq for IndividualTableRow {
+    fn eq(&self, other: &Self) -> bool {
+        self.flags == other.flags
+            && self.parents == other.parents
+            && self.metadata == other.metadata
+            && match &self.location {
+                None => other.location.is_none(),
+                Some(a) => match &other.location {
+                    None => false,
+                    Some(b) => {
+                        if a.len() != b.len() {
+                            false
+                        } else {
+                            for (i, j) in a.iter().enumerate() {
+                                if !crate::util::f64_partial_cmp_equal(&j, &b[i]) {
+                                    return false;
+                                }
+                            }
+                            true
+                        }
+                    }
+                },
+            }
+    }
+}
+
+/// An immutable view of a individual table.
+///
+/// These are not created directly.
+/// Instead, use [`TableCollection::individuals`](crate::TableCollection::individuals)
+/// to get a reference to an existing node table;
+pub struct IndividualTable<'a> {
+    table_: &'a ll_bindings::tsk_individual_table_t,
+}
+
+fn make_individual_table_row(
+    table: &IndividualTable,
+    pos: tsk_id_t,
+    decode_metadata: bool,
+) -> Option<IndividualTableRow> {
+    if pos < table.num_rows() as tsk_id_t {
+        let rv = IndividualTableRow {
+            flags: table.flags(pos).unwrap(),
+            location: table.location(pos).unwrap(),
+            parents: table.parents(pos).unwrap(),
+            metadata: match decode_metadata {
+                true => match metadata_to_vector!(table, pos).unwrap() {
+                    Some(x) => Some(x),
+                    None => None,
+                },
+                false => None,
+            },
+        };
+        Some(rv)
+    } else {
+        None
+    }
+}
+
+pub type IndividualTableRefIterator<'a> =
+    crate::table_iterator::TableIterator<&'a IndividualTable<'a>>;
+pub type IndividualTableIterator<'a> = crate::table_iterator::TableIterator<IndividualTable<'a>>;
+
+impl<'a> IndividualTable<'a> {
+    pub(crate) fn new_from_table(individuals: &'a ll_bindings::tsk_individual_table_t) -> Self {
+        IndividualTable {
+            table_: individuals,
+        }
+    }
+
+    /// Return the number of rows
+    pub fn num_rows(&'a self) -> ll_bindings::tsk_size_t {
+        self.table_.num_rows
+    }
+
+    /// Return the flags for a given row.
+    ///
+    /// # Errors
+    ///
+    /// * [`TskitError::IndexError`] if `row` is out of range.
+    pub fn flags(&self, row: tsk_id_t) -> Result<tsk_flags_t, TskitError> {
+        unsafe_tsk_column_access!(row, 0, self.num_rows(), self.table_.flags)
+    }
+
+    /// Return the locations for a given row.
+    ///
+    /// # Errors
+    ///
+    /// * [`TskitError::IndexError`] if `row` is out of range.
+    pub fn location(&self, row: tsk_id_t) -> Result<Option<Vec<f64>>, TskitError> {
+        unsafe_tsk_ragged_column_access!(
+            row,
+            0,
+            self.num_rows(),
+            self.table_.location,
+            self.table_.location_offset,
+            self.table_.location_length
+        )
+    }
+
+    /// Return the parents for a given row.
+    ///
+    /// # Errors
+    ///
+    /// * [`TskitError::IndexError`] if `row` is out of range.
+    pub fn parents(&self, row: tsk_id_t) -> Result<Option<Vec<tsk_id_t>>, TskitError> {
+        unsafe_tsk_ragged_column_access!(
+            row,
+            0,
+            self.num_rows(),
+            self.table_.parents,
+            self.table_.parents_offset,
+            self.table_.parents_length
+        )
+    }
+
+    /// Return the metadata for a given row.
+    ///
+    /// # Errors
+    ///
+    /// * [`TskitError::IndexError`] if `row` is out of range.
+    pub fn metadata<T: metadata::MetadataRoundtrip>(
+        &'a self,
+        row: tsk_id_t,
+    ) -> Result<Option<T>, TskitError> {
+        let buffer = metadata_to_vector!(self, row)?;
+        decode_metadata_row!(T, buffer)
+    }
+
+    /// Return an iterator over rows of the table.
+    /// The value of the iterator is [`IndividualTableRow`].
+    ///
+    /// # Parameters
+    ///
+    /// * `decode_metadata`: if `true`, then a *copy* of row metadata
+    ///    will be provided in [`IndividualTableRow::metadata`].
+    ///    The meta data are *not* decoded.
+    ///    Rows with no metadata will contain the value `None`.
+    ///
+    pub fn iter(&self, decode_metadata: bool) -> IndividualTableRefIterator {
+        crate::table_iterator::make_table_iterator::<&IndividualTable<'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 [`IndividualTableRow::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<IndividualTableRow, TskitError> {
+        table_row_access!(r, decode_metadata, self, make_individual_table_row)
+    }
+}

src/lib.rs

@@ -10,7 +10,9 @@ mod _macros; // Starts w/_ to be sorted at front by rustfmt!
 mod edge_table;
 pub mod error;
 pub mod ffi;
+mod individual_table;
 pub mod metadata;
+mod migration_table;
 mod mutation_table;
 mod node_table;
 mod population_table;
@@ -40,6 +42,8 @@ pub const TSK_NULL: tsk_id_t = -1;
 
 pub use edge_table::{EdgeTable, EdgeTableRow};
 pub use error::TskitError;
+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 population_table::{PopulationTable, PopulationTableRow};