feat: Add lending iterators over table row "views"

molpopgen · molpopgen · commit d835d6f93340 · 2022-11-07T15:04:43.000-08:00
diff --git 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
@@ -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
@@ -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::<MutationMetadata>(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
@@ -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
@@ -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
@@ -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
diff --git a/src/site_table.rs b/src/site_table.rs
diff --git a/tests/test_tables.rs b/tests/test_tables.rs
