Add TableAccess trait and implement for TableCollection and TreeSequence. (#66)

Author: molpopgen

src/lib.rs

@@ -19,6 +19,7 @@ mod population_table;
 mod site_table;
 mod table_collection;
 mod table_iterator;
+mod traits;
 mod trees;
 pub mod types;
 mod util;
@@ -49,6 +50,7 @@ pub use node_table::{NodeTable, NodeTableRow};
 pub use population_table::{PopulationTable, PopulationTableRow};
 pub use site_table::{SiteTable, SiteTableRow};
 pub use table_collection::TableCollection;
+pub use traits::TableAccess;
 pub use trees::{NodeIterator, NodeTraversalOrder, Tree, TreeFlags, TreeSequence};
 
 /// Handles return codes from low-level tskit functions.

src/table_collection.rs

@@ -10,6 +10,7 @@ use crate::MutationTable;
 use crate::NodeTable;
 use crate::PopulationTable;
 use crate::SiteTable;
+use crate::TableAccess;
 use crate::TskReturnValue;
 use crate::{tsk_flags_t, tsk_id_t, tsk_size_t};
 use ll_bindings::tsk_table_collection_free;
@@ -23,6 +24,7 @@ use ll_bindings::tsk_table_collection_free;
 /// # Examples
 ///
 /// ```
+/// use tskit::TableAccess;
 /// let mut tables = tskit::TableCollection::new(100.).unwrap();
 /// assert_eq!(tables.sequence_length(), 100.);
 ///
@@ -47,6 +49,7 @@ use ll_bindings::tsk_table_collection_free;
 ///
 /// ```
 /// use tskit;
+/// use tskit::TableAccess;
 /// use tskit::metadata::MetadataRoundtrip;
 ///
 /// // Define a type for metadata
@@ -156,130 +159,6 @@ impl TableCollection {
         unsafe { (*self.as_ptr()).sequence_length }
     }
 
-    /// Get reference to the [``EdgeTable``](crate::EdgeTable).
-    /// Lifetime of return value is tied to (this)
-    /// parent object.
-    pub fn edges<'a>(&'a self) -> EdgeTable<'a> {
-        EdgeTable::<'a>::new_from_table(&self.inner.edges)
-    }
-
-    /// Return an iterator over the edges.
-    /// See [`EdgeTable::iter`] for details.
-    pub fn edges_iter<'a>(
-        &'a self,
-        decode_metadata: bool,
-    ) -> crate::edge_table::EdgeTableIterator<'a> {
-        crate::table_iterator::make_table_iterator::<EdgeTable<'a>>(self.edges(), decode_metadata)
-    }
-
-    /// Get reference to the [``IndividualTable``](crate::IndividualTable).
-    /// Lifetime of return value is tied to (this)
-    /// parent object.
-    pub fn individuals<'a>(&'a self) -> IndividualTable<'a> {
-        IndividualTable::<'a>::new_from_table(&self.inner.individuals)
-    }
-
-    /// Return an iterator over the individuals.
-    /// See [`IndividualTable::iter`] for details.
-    pub fn individuals_iter<'a>(
-        &'a self,
-        decode_metadata: bool,
-    ) -> crate::individual_table::IndividualTableIterator<'a> {
-        crate::table_iterator::make_table_iterator::<IndividualTable<'a>>(
-            self.individuals(),
-            decode_metadata,
-        )
-    }
-
-    /// Get reference to the [``MigrationTable``](crate::MigrationTable).
-    /// Lifetime of return value is tied to (this)
-    /// parent object.
-    pub fn migrations<'a>(&'a self) -> MigrationTable<'a> {
-        MigrationTable::<'a>::new_from_table(&self.inner.migrations)
-    }
-
-    /// Return an iterator over the migration events.
-    /// See [`MigrationTable::iter`] for details.
-    pub fn migrations_iter<'a>(
-        &'a self,
-        decode_metadata: bool,
-    ) -> crate::migration_table::MigrationTableIterator<'a> {
-        crate::table_iterator::make_table_iterator::<MigrationTable<'a>>(
-            self.migrations(),
-            decode_metadata,
-        )
-    }
-
-    /// Get reference to the [``NodeTable``](crate::NodeTable).
-    /// Lifetime of return value is tied to (this)
-    /// parent object.
-    pub fn nodes<'a>(&'a self) -> NodeTable<'a> {
-        NodeTable::<'a>::new_from_table(&self.inner.nodes)
-    }
-
-    /// Return an iterator over the nodes.
-    /// See [`NodeTable::iter`] for details.
-    pub fn nodes_iter<'a>(
-        &'a self,
-        decode_metadata: bool,
-    ) -> crate::node_table::NodeTableIterator<'a> {
-        crate::table_iterator::make_table_iterator::<NodeTable<'a>>(self.nodes(), decode_metadata)
-    }
-
-    /// Get reference to the [``SiteTable``](crate::SiteTable).
-    /// Lifetime of return value is tied to (this)
-    /// parent object.
-    pub fn sites<'a>(&'a self) -> SiteTable<'a> {
-        SiteTable::<'a>::new_from_table(&self.inner.sites)
-    }
-
-    /// Return an iterator over the sites.
-    /// See [`SiteTable::iter`] for details.
-    pub fn sites_iter<'a>(
-        &'a self,
-        decode_metadata: bool,
-    ) -> crate::site_table::SiteTableIterator<'a> {
-        crate::table_iterator::make_table_iterator::<SiteTable<'a>>(self.sites(), decode_metadata)
-    }
-
-    /// Get reference to the [``MutationTable``](crate::MutationTable).
-    /// Lifetime of return value is tied to (this)
-    /// parent object.
-    pub fn mutations<'a>(&'a self) -> MutationTable<'a> {
-        MutationTable::<'a>::new_from_table(&self.inner.mutations)
-    }
-
-    /// Return an iterator over the mutations.
-    /// See [`MutationTable::iter`] for details.
-    pub fn mutations_iter<'a>(
-        &'a self,
-        decode_metadata: bool,
-    ) -> crate::mutation_table::MutationTableIterator<'a> {
-        crate::table_iterator::make_table_iterator::<MutationTable<'a>>(
-            self.mutations(),
-            decode_metadata,
-        )
-    }
-
-    /// Get reference to the [``PopulationTable``](crate::PopulationTable).
-    /// Lifetime of return value is tied to (this)
-    /// parent object.
-    pub fn populations<'a>(&'a self) -> PopulationTable<'a> {
-        PopulationTable::<'a>::new_from_table(&self.inner.populations)
-    }
-
-    /// Return an iterator over the populations.
-    /// See [`PopulationTable::iter`] for details.
-    pub fn populations_iter<'a>(
-        &'a self,
-        decode_metadata: bool,
-    ) -> crate::population_table::PopulationTableIterator<'a> {
-        crate::table_iterator::make_table_iterator::<PopulationTable<'a>>(
-            self.populations(),
-            decode_metadata,
-        )
-    }
-
     /// Add a row to the edge table
     pub fn add_edge(
         &mut self,
@@ -610,6 +489,36 @@ impl TableCollection {
     }
 }
 
+impl TableAccess for TableCollection {
+    fn edges(&self) -> EdgeTable {
+        EdgeTable::new_from_table(&self.inner.edges)
+    }
+
+    fn individuals(&self) -> IndividualTable {
+        IndividualTable::new_from_table(&self.inner.individuals)
+    }
+
+    fn migrations(&self) -> MigrationTable {
+        MigrationTable::new_from_table(&self.inner.migrations)
+    }
+
+    fn nodes(&self) -> NodeTable {
+        NodeTable::new_from_table(&self.inner.nodes)
+    }
+
+    fn sites(&self) -> SiteTable {
+        SiteTable::new_from_table(&self.inner.sites)
+    }
+
+    fn mutations(&self) -> MutationTable {
+        MutationTable::new_from_table(&self.inner.mutations)
+    }
+
+    fn populations(&self) -> PopulationTable {
+        PopulationTable::new_from_table(&self.inner.populations)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

src/traits.rs

@@ -0,0 +1,113 @@
+//! Traits related to user-facing types
+
+use crate::edge_table::EdgeTableIterator;
+use crate::individual_table::IndividualTableIterator;
+use crate::migration_table::MigrationTableIterator;
+use crate::mutation_table::MutationTableIterator;
+use crate::node_table::NodeTableIterator;
+use crate::population_table::PopulationTableIterator;
+use crate::site_table::SiteTableIterator;
+use crate::table_iterator::make_table_iterator;
+use crate::EdgeTable;
+use crate::IndividualTable;
+use crate::MigrationTable;
+use crate::MutationTable;
+use crate::NodeTable;
+use crate::PopulationTable;
+use crate::SiteTable;
+
+/// Immutable access to tables.
+///
+/// For objects that contain the full suite of tables,
+/// implementing this trait provides immutable access
+/// to their data.
+///
+/// For most types, the provided implementations of `_iter`
+/// methods should do.
+///
+/// # Examples
+///
+/// ```
+/// use tskit::TableAccess;
+///
+/// let mut tables = tskit::TableCollection::new(1.).unwrap();
+/// // The borrows are immuatble, so we can
+/// // take multiple table references from the same object.
+/// let e = tables.edges();
+/// let n = tables.nodes();
+/// ```
+///
+/// The borrow checker will keep you from getting in trouble:
+///
+/// ```compile_fail
+/// use tskit::TableAccess;
+///
+/// let mut tables = tskit::TableCollection::new(1.).unwrap();
+/// let e = tables.edges();
+/// tables.add_edge(0.0, 1.0, 0, 1).unwrap();
+/// let p = e.parent(0).unwrap();   // FAIL!
+/// ```
+pub trait TableAccess {
+    /// Get reference to the [``EdgeTable``](crate::EdgeTable).
+    fn edges(&self) -> EdgeTable;
+
+    /// Return an iterator over the edges.
+    /// See [`EdgeTable::iter`] for details.
+    fn edges_iter(&self, decode_metadata: bool) -> EdgeTableIterator {
+        make_table_iterator::<EdgeTable>(self.edges(), decode_metadata)
+    }
+
+    /// Get reference to the [``NodeTable``](crate::NodeTable).
+    fn nodes(&self) -> NodeTable;
+
+    /// Return an iterator over the nodes.
+    /// See [`NodeTable::iter`] for details.
+    fn nodes_iter(&self, decode_metadata: bool) -> NodeTableIterator {
+        make_table_iterator::<NodeTable>(self.nodes(), decode_metadata)
+    }
+
+    /// Get reference to the [``MutationTable``](crate::MutationTable).
+    fn mutations(&self) -> MutationTable;
+
+    /// Return an iterator over the mutations.
+    /// See [`MutationTable::iter`] for details.
+    fn mutations_iter(&self, decode_metadata: bool) -> MutationTableIterator {
+        make_table_iterator::<MutationTable>(self.mutations(), decode_metadata)
+    }
+
+    /// Get reference to the [``SiteTable``](crate::SiteTable).
+    fn sites(&self) -> SiteTable;
+
+    /// Return an iterator over the sites.
+    /// See [`SiteTable::iter`] for details.
+    fn sites_iter(&self, decode_metadata: bool) -> SiteTableIterator {
+        make_table_iterator::<SiteTable>(self.sites(), decode_metadata)
+    }
+
+    /// Get reference to the [``PopulationTable``](crate::PopulationTable).
+    fn populations(&self) -> PopulationTable;
+
+    /// Return an iterator over the populations.
+    /// See [`PopulationTable::iter`] for details.
+    fn populations_iter(&self, decode_metadata: bool) -> PopulationTableIterator {
+        make_table_iterator::<PopulationTable>(self.populations(), decode_metadata)
+    }
+
+    /// Get reference to the [``MigrationTable``](crate::MigrationTable).
+    fn migrations(&self) -> MigrationTable;
+
+    /// Return an iterator over the migration events.
+    /// See [`MigrationTable::iter`] for details.
+    fn migrations_iter(&self, decode_metadata: bool) -> MigrationTableIterator {
+        make_table_iterator::<MigrationTable>(self.migrations(), decode_metadata)
+    }
+
+    /// Get reference to the [``IndividualTable``](crate::IndividualTable).
+    fn individuals(&self) -> IndividualTable;
+
+    /// Return an iterator over the individuals.
+    /// See [`IndividualTable::iter`] for details.
+    fn individuals_iter(&self, decode_metadata: bool) -> IndividualTableIterator {
+        make_table_iterator::<IndividualTable>(self.individuals(), decode_metadata)
+    }
+}

src/trees.rs

@@ -1,6 +1,14 @@
 use crate::bindings as ll_bindings;
 use crate::error::TskitError;
 use crate::ffi::{TskitTypeAccess, WrapTskitType};
+use crate::EdgeTable;
+use crate::IndividualTable;
+use crate::MigrationTable;
+use crate::MutationTable;
+use crate::NodeTable;
+use crate::PopulationTable;
+use crate::SiteTable;
+use crate::TableAccess;
 use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TableCollection, TSK_NULL};
 use bitflags::bitflags;
 use ll_bindings::{tsk_tree_free, tsk_treeseq_free};
@@ -825,6 +833,36 @@ impl TreeSequence {
     }
 }
 
+impl TableAccess for TreeSequence {
+    fn edges(&self) -> EdgeTable {
+        EdgeTable::new_from_table(unsafe { &(*self.inner.tables).edges })
+    }
+
+    fn individuals(&self) -> IndividualTable {
+        IndividualTable::new_from_table(unsafe { &(*self.inner.tables).individuals })
+    }
+
+    fn migrations(&self) -> MigrationTable {
+        MigrationTable::new_from_table(unsafe { &(*self.inner.tables).migrations })
+    }
+
+    fn nodes(&self) -> NodeTable {
+        NodeTable::new_from_table(unsafe { &(*self.inner.tables).nodes })
+    }
+
+    fn sites(&self) -> SiteTable {
+        SiteTable::new_from_table(unsafe { &(*self.inner.tables).sites })
+    }
+
+    fn mutations(&self) -> MutationTable {
+        MutationTable::new_from_table(unsafe { &(*self.inner.tables).mutations })
+    }
+
+    fn populations(&self) -> PopulationTable {
+        PopulationTable::new_from_table(unsafe { &(*self.inner.tables).populations })
+    }
+}
+
 #[cfg(test)]
 mod test_trees {
     use super::*;