From d835d6f93340e3fdf1869dbe91c27c227c08c3fe Mon Sep 17 00:00:00 2001 From: "Kevin R. Thornton" Date: Mon, 7 Nov 2022 10:48:50 -0800 Subject: [PATCH] feat: Add lending iterators over table row "views" --- src/_macros.rs | 27 ++++++++++++++++++++++ src/edge_table.rs | 46 +++++++++++++++++++++++++++++++++++++ src/individual_table.rs | 42 +++++++++++++++++++++++++++++++++ src/migration_table.rs | 51 +++++++++++++++++++++++++++++++++++++++++ src/mutation_table.rs | 48 ++++++++++++++++++++++++++++++++++++++ src/node_table.rs | 45 ++++++++++++++++++++++++++++++++++++ src/population_table.rs | 33 ++++++++++++++++++++++++++ src/site_table.rs | 42 +++++++++++++++++++++++++++++++++ tests/test_tables.rs | 31 ++++++++++++++++++++++++- 9 files changed, 364 insertions(+), 1 deletion(-) diff --git a/src/_macros.rs b/src/_macros.rs index 1942875cc..ee4967881 100644 --- a/src/_macros.rs +++ b/src/_macros.rs @@ -1077,6 +1077,33 @@ macro_rules! build_owned_table_type { }; } +macro_rules! raw_metadata_getter_for_tables { + ($idtype: ident) => { + fn raw_metadata(&self, row: $idtype) -> Option<&[u8]> { + $crate::metadata::char_column_to_slice( + self, + self.as_ref().metadata, + self.as_ref().metadata_offset, + row.0, + self.num_rows().into(), + self.as_ref().metadata_length, + ) + } + }; +} + +macro_rules! row_lending_iterator_get { + () => { + fn get(&self) -> Option<&Self::Item> { + if crate::SizeType::try_from(self.id).ok()? < self.table.num_rows() { + Some(self) + } else { + None + } + } + }; +} + #[cfg(test)] mod test { use crate::error::TskitError; diff --git a/src/edge_table.rs b/src/edge_table.rs index b660d24ea..f18b2893f 100644 --- a/src/edge_table.rs +++ b/src/edge_table.rs @@ -64,6 +64,46 @@ impl Iterator for EdgeTableIterator { } } +/// Row of an [`EdgeTable`] +pub struct EdgeTableRowView<'a> { + table: &'a EdgeTable, + pub id: EdgeId, + pub left: Position, + pub right: Position, + pub parent: NodeId, + pub child: NodeId, + pub metadata: Option<&'a [u8]>, +} + +impl<'a> EdgeTableRowView<'a> { + fn new(table: &'a EdgeTable) -> Self { + Self { + table, + id: (-1).into(), + left: f64::NAN.into(), + right: f64::NAN.into(), + parent: NodeId::NULL, + child: NodeId::NULL, + metadata: None, + } + } +} + +impl<'a> streaming_iterator::StreamingIterator for EdgeTableRowView<'a> { + type Item = Self; + + row_lending_iterator_get!(); + + fn advance(&mut self) { + self.id = (i32::from(self.id) + 1).into(); + self.left = self.table.left(self.id).unwrap_or_else(|| f64::NAN.into()); + self.right = self.table.right(self.id).unwrap_or_else(|| f64::NAN.into()); + self.parent = self.table.parent(self.id).unwrap_or(NodeId::NULL); + self.child = self.table.child(self.id).unwrap_or(NodeId::NULL); + self.metadata = self.table.raw_metadata(self.id); + } +} + /// An immutable view of an edge table. /// /// These are not created directly but are accessed @@ -94,6 +134,8 @@ impl EdgeTable { self.as_ref().num_rows.into() } + raw_metadata_getter_for_tables!(EdgeId); + /// Return the ``parent`` value from row ``row`` of the table. /// /// # Returns @@ -193,6 +235,10 @@ impl EdgeTable { crate::table_iterator::make_table_iterator::<&EdgeTable>(self) } + pub fn lending_iter(&self) -> EdgeTableRowView { + EdgeTableRowView::new(self) + } + /// Return row `r` of the table. /// /// # Parameters diff --git a/src/individual_table.rs b/src/individual_table.rs index 9014f42d5..bc782b81b 100644 --- a/src/individual_table.rs +++ b/src/individual_table.rs @@ -45,6 +45,42 @@ impl PartialEq for IndividualTableRow { } } +pub struct IndividualTableRowView<'a> { + table: &'a IndividualTable, + pub id: IndividualId, + pub flags: IndividualFlags, + pub location: Option<&'a [Location]>, + pub parents: Option<&'a [IndividualId]>, + pub metadata: Option<&'a [u8]>, +} + +impl<'a> IndividualTableRowView<'a> { + fn new(table: &'a IndividualTable) -> Self { + Self { + table, + id: (-1_i32).into(), + flags: 0.into(), + location: None, + parents: None, + metadata: None, + } + } +} + +impl<'a> streaming_iterator::StreamingIterator for IndividualTableRowView<'a> { + type Item = Self; + + row_lending_iterator_get!(); + + fn advance(&mut self) { + self.id = (i32::from(self.id) + 1).into(); + self.flags = self.table.flags(self.id).unwrap_or_else(|| 0.into()); + self.location = self.table.location(self.id); + self.parents = self.table.parents(self.id); + self.metadata = self.table.raw_metadata(self.id); + } +} + /// An immutable view of a individual table. /// /// These are not created directly but are accessed @@ -104,6 +140,8 @@ impl IndividualTable { unsafe { self.table_.as_ref() } } + raw_metadata_getter_for_tables!(IndividualId); + /// Return the number of rows pub fn num_rows(&self) -> crate::SizeType { self.as_ref().num_rows.into() @@ -349,6 +387,10 @@ match tables.individuals().metadata::(0.into()) crate::table_iterator::make_table_iterator::<&IndividualTable>(self) } + pub fn lending_iter(&self) -> IndividualTableRowView { + IndividualTableRowView::new(self) + } + /// Return row `r` of the table. /// /// # Parameters diff --git a/src/migration_table.rs b/src/migration_table.rs index 1adbbd240..18f668607 100644 --- a/src/migration_table.rs +++ b/src/migration_table.rs @@ -73,6 +73,51 @@ impl Iterator for MigrationTableIterator { } } +pub struct MigrationTableRowView<'a> { + table: &'a MigrationTable, + pub id: MigrationId, + pub left: Position, + pub right: Position, + pub node: NodeId, + pub source: PopulationId, + pub dest: PopulationId, + pub time: Time, + pub metadata: Option<&'a [u8]>, +} + +impl<'a> MigrationTableRowView<'a> { + fn new(table: &'a MigrationTable) -> Self { + Self { + table, + id: MigrationId::NULL, + left: Position::from(f64::NAN), + right: Position::from(f64::NAN), + node: NodeId::NULL, + source: PopulationId::NULL, + dest: PopulationId::NULL, + time: Time::from(f64::NAN), + metadata: None, + } + } +} + +impl<'a> streaming_iterator::StreamingIterator for MigrationTableRowView<'a> { + type Item = Self; + + row_lending_iterator_get!(); + + fn advance(&mut self) { + self.id = (i32::from(self.id) + 1).into(); + self.left = self.table.left(self.id).unwrap_or_else(|| f64::NAN.into()); + self.right = self.table.right(self.id).unwrap_or_else(|| f64::NAN.into()); + self.node = self.table.node(self.id).unwrap_or(NodeId::NULL); + self.source = self.table.source(self.id).unwrap_or(PopulationId::NULL); + self.dest = self.table.dest(self.id).unwrap_or(PopulationId::NULL); + self.time = self.table.time(self.id).unwrap_or_else(|| f64::NAN.into()); + self.metadata = self.table.raw_metadata(self.id); + } +} + /// An immutable view of a migration table. /// /// These are not created directly but are accessed @@ -102,6 +147,8 @@ impl MigrationTable { self.as_ref().num_rows.into() } + raw_metadata_getter_for_tables!(MigrationId); + /// Return the left coordinate for a given row. /// /// # Returns @@ -232,6 +279,10 @@ impl MigrationTable { crate::table_iterator::make_table_iterator::<&MigrationTable>(self) } + pub fn lending_iter(&self) -> MigrationTableRowView { + MigrationTableRowView::new(self) + } + /// Return row `r` of the table. /// /// # Parameters diff --git a/src/mutation_table.rs b/src/mutation_table.rs index 14b0c54de..6de31ef66 100644 --- a/src/mutation_table.rs +++ b/src/mutation_table.rs @@ -76,6 +76,48 @@ impl Iterator for MutationTableIterator { } } +pub struct MutationTableRowView<'a> { + table: &'a MutationTable, + pub id: MutationId, + pub site: SiteId, + pub node: NodeId, + pub parent: MutationId, + pub time: Time, + pub derived_state: Option<&'a [u8]>, + pub metadata: Option<&'a [u8]>, +} + +impl<'a> MutationTableRowView<'a> { + fn new(table: &'a MutationTable) -> Self { + Self { + table, + id: MutationId::NULL, + site: SiteId::NULL, + node: NodeId::NULL, + parent: MutationId::NULL, + time: f64::NAN.into(), + derived_state: None, + metadata: None, + } + } +} + +impl<'a> streaming_iterator::StreamingIterator for MutationTableRowView<'a> { + type Item = Self; + + row_lending_iterator_get!(); + + fn advance(&mut self) { + self.id = (i32::from(self.id) + 1).into(); + self.site = self.table.site(self.id).unwrap_or(SiteId::NULL); + self.node = self.table.node(self.id).unwrap_or(NodeId::NULL); + self.parent = self.table.parent(self.id).unwrap_or(MutationId::NULL); + self.time = self.table.time(self.id).unwrap_or_else(|| f64::NAN.into()); + self.derived_state = self.table.derived_state(self.id); + self.metadata = self.table.raw_metadata(self.id); + } +} + /// An immutable view of site table. /// /// These are not created directly but are accessed @@ -105,6 +147,8 @@ impl MutationTable { self.as_ref().num_rows.into() } + raw_metadata_getter_for_tables!(MutationId); + /// Return the ``site`` value from row ``row`` of the table. /// /// # Errors @@ -218,6 +262,10 @@ impl MutationTable { crate::table_iterator::make_table_iterator::<&MutationTable>(self) } + pub fn lending_iter(&self) -> MutationTableRowView { + MutationTableRowView::new(self) + } + /// Return row `r` of the table. /// /// # Parameters diff --git a/src/node_table.rs b/src/node_table.rs index e3652e08d..272fd8298 100644 --- a/src/node_table.rs +++ b/src/node_table.rs @@ -66,6 +66,45 @@ impl Iterator for NodeTableIterator { } } +pub struct NodeTableRowView<'a> { + table: &'a NodeTable, + pub id: NodeId, + pub time: Time, + pub flags: NodeFlags, + pub population: PopulationId, + pub individual: IndividualId, + pub metadata: Option<&'a [u8]>, +} + +impl<'a> NodeTableRowView<'a> { + fn new(table: &'a NodeTable) -> Self { + Self { + table, + id: NodeId::NULL, + time: f64::NAN.into(), + flags: 0.into(), + population: PopulationId::NULL, + individual: IndividualId::NULL, + metadata: None, + } + } +} + +impl<'a> streaming_iterator::StreamingIterator for NodeTableRowView<'a> { + type Item = Self; + + row_lending_iterator_get!(); + + fn advance(&mut self) { + self.id = (i32::from(self.id) + 1).into(); + self.time = self.table.time(self.id).unwrap_or_else(|| f64::NAN.into()); + self.flags = self.table.flags(self.id).unwrap_or_else(|| 0.into()); + self.population = self.table.population(self.id).unwrap_or(PopulationId::NULL); + self.individual = self.table.individual(self.id).unwrap_or(IndividualId::NULL); + self.metadata = self.table.raw_metadata(self.id); + } +} + /// An immtable view of a node table. /// /// These are not created directly but are accessed @@ -95,6 +134,8 @@ impl NodeTable { self.as_ref().num_rows.into() } + raw_metadata_getter_for_tables!(NodeId); + /// Return the ``time`` value from row ``row`` of the table. /// /// # Returns @@ -386,6 +427,10 @@ impl NodeTable { crate::table_iterator::make_table_iterator::<&NodeTable>(self) } + pub fn lending_iter(&self) -> NodeTableRowView { + NodeTableRowView::new(self) + } + /// Return row `r` of the table. /// /// # Parameters diff --git a/src/population_table.rs b/src/population_table.rs index 84390205e..d8b25a285 100644 --- a/src/population_table.rs +++ b/src/population_table.rs @@ -61,6 +61,33 @@ impl Iterator for PopulationTableIterator { } } +pub struct PopulationTableRowView<'a> { + table: &'a PopulationTable, + pub id: PopulationId, + pub metadata: Option<&'a [u8]>, +} + +impl<'a> PopulationTableRowView<'a> { + fn new(table: &'a PopulationTable) -> Self { + Self { + table, + id: PopulationId::NULL, + metadata: None, + } + } +} + +impl<'a> streaming_iterator::StreamingIterator for PopulationTableRowView<'a> { + type Item = Self; + + row_lending_iterator_get!(); + + fn advance(&mut self) { + self.id = (i32::from(self.id) + 1).into(); + self.metadata = self.table.raw_metadata(self.id); + } +} + /// An immutable view of site table. /// /// These are not created directly but are accessed @@ -86,6 +113,8 @@ impl PopulationTable { unsafe { self.table_.as_ref() } } + raw_metadata_getter_for_tables!(PopulationId); + /// Return the number of rows. pub fn num_rows(&self) -> SizeType { self.as_ref().num_rows.into() @@ -122,6 +151,10 @@ impl PopulationTable { crate::table_iterator::make_table_iterator::<&PopulationTable>(self) } + pub fn lending_iter(&self) -> PopulationTableRowView { + PopulationTableRowView::new(self) + } + /// Return row `r` of the table. /// /// # Parameters diff --git a/src/site_table.rs b/src/site_table.rs index f7ee4c75b..97f7da738 100644 --- a/src/site_table.rs +++ b/src/site_table.rs @@ -61,6 +61,42 @@ impl Iterator for SiteTableIterator { } } +pub struct SiteTableRowView<'a> { + table: &'a SiteTable, + pub id: SiteId, + pub position: Position, + pub ancestral_state: Option<&'a [u8]>, + pub metadata: Option<&'a [u8]>, +} + +impl<'a> SiteTableRowView<'a> { + fn new(table: &'a SiteTable) -> Self { + Self { + table, + id: SiteId::NULL, + position: f64::NAN.into(), + ancestral_state: None, + metadata: None, + } + } +} + +impl<'a> streaming_iterator::StreamingIterator for SiteTableRowView<'a> { + type Item = Self; + + row_lending_iterator_get!(); + + fn advance(&mut self) { + self.id = (i32::from(self.id) + 1).into(); + self.position = self + .table + .position(self.id) + .unwrap_or_else(|| f64::NAN.into()); + self.ancestral_state = self.table.ancestral_state(self.id); + self.metadata = self.table.raw_metadata(self.id); + } +} + /// An immutable view of site table. /// /// These are not created directly but are accessed @@ -85,6 +121,8 @@ impl SiteTable { unsafe { self.table_.as_ref() } } + raw_metadata_getter_for_tables!(SiteId); + /// Return the number of rows pub fn num_rows(&self) -> SizeType { self.as_ref().num_rows.into() @@ -155,6 +193,10 @@ impl SiteTable { crate::table_iterator::make_table_iterator::<&SiteTable>(self) } + pub fn lending_iter(&self) -> SiteTableRowView { + SiteTableRowView::new(self) + } + /// Return row `r` of the table. /// /// # Parameters diff --git a/tests/test_tables.rs b/tests/test_tables.rs index 647ace8f1..315dcdce7 100644 --- a/tests/test_tables.rs +++ b/tests/test_tables.rs @@ -65,6 +65,7 @@ mod test_adding_rows_without_metadata { }, Err(e) => panic!("Err from tables.{}: {:?}", stringify!(adder), e) } + assert_eq!(tables.$table().iter().count(), 2); } }}; } @@ -314,20 +315,48 @@ mod test_metadata_round_trips { macro_rules! add_row_with_metadata { ($table: ident, $adder: ident, $md: ident) => {{ { + use tskit::prelude::*; + use tskit::metadata::MetadataRoundtrip; build_metadata_types!($md); let mut tables = tskit::TableCollection::new(10.).unwrap(); let md = MyMetadata::new(); let row = tables.$adder(&md); - match_block_impl!(tables, $table, $adder, row, md) + match_block_impl!(tables, $table, $adder, row, md); + let mut lending_iter = tables.$table().lending_iter(); + let mut iter = tables.$table().iter(); + while let Some(row) = lending_iter.next() { + if let Some(row_from_iter) = iter.next() { + assert_eq!(row.id, row_from_iter.id); + } + if let Some(metadata) = row.metadata { + assert_eq!(MyMetadata::decode(metadata).unwrap(), md); + }else { + panic!("expected Some(metadata)"); + } + } } }}; ($table: ident, $adder: ident, $md: ident $(,$payload: expr) + ) => {{ { + use tskit::prelude::*; + use tskit::metadata::MetadataRoundtrip; build_metadata_types!($md); let mut tables = tskit::TableCollection::new(10.).unwrap(); let md = MyMetadata::new(); let row = tables.$adder($($payload ), *, &md); match_block_impl!(tables, $table, $adder, row, md); + let mut lending_iter = tables.$table().lending_iter(); + let mut iter = tables.$table().iter(); + while let Some(row) = lending_iter.next() { + if let Some(row_from_iter) = iter.next() { + assert_eq!(row.id, row_from_iter.id); + } + if let Some(metadata) = row.metadata { + assert_eq!(MyMetadata::decode(metadata).unwrap(), md); + }else { + panic!("expected Some(metadata)"); + } + } } }}; }