Skip to content

Commit fea923b

Browse files
molpopgenmolpopgen
authored
feat: Add lending iterators over table row "views" (#398)
1 parent 9acf22a commit fea923b

File tree

9 files changed

+364
-1
lines changed

9 files changed

+364
-1
lines changed

‎src/_macros.rs‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1077,6 +1077,33 @@ macro_rules! build_owned_table_type {
10771077
};
10781078
}
10791079

1080+
macro_rules! raw_metadata_getter_for_tables {
1081+
($idtype: ident) => {
1082+
fn raw_metadata(&self, row: $idtype) -> Option<&[u8]> {
1083+
$crate::metadata::char_column_to_slice(
1084+
self,
1085+
self.as_ref().metadata,
1086+
self.as_ref().metadata_offset,
1087+
row.0,
1088+
self.num_rows().into(),
1089+
self.as_ref().metadata_length,
1090+
)
1091+
}
1092+
};
1093+
}
1094+
1095+
macro_rules! row_lending_iterator_get {
1096+
() => {
1097+
fn get(&self) -> Option<&Self::Item> {
1098+
if crate::SizeType::try_from(self.id).ok()? < self.table.num_rows() {
1099+
Some(self)
1100+
} else {
1101+
None
1102+
}
1103+
}
1104+
};
1105+
}
1106+
10801107
#[cfg(test)]
10811108
mod test {
10821109
use crate::error::TskitError;

‎src/edge_table.rs‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,46 @@ impl Iterator for EdgeTableIterator {
6464
}
6565
}
6666

67+
/// Row of an [`EdgeTable`]
68+
pub struct EdgeTableRowView<'a> {
69+
table: &'a EdgeTable,
70+
pub id: EdgeId,
71+
pub left: Position,
72+
pub right: Position,
73+
pub parent: NodeId,
74+
pub child: NodeId,
75+
pub metadata: Option<&'a [u8]>,
76+
}
77+
78+
impl<'a> EdgeTableRowView<'a> {
79+
fn new(table: &'a EdgeTable) -> Self {
80+
Self {
81+
table,
82+
id: (-1).into(),
83+
left: f64::NAN.into(),
84+
right: f64::NAN.into(),
85+
parent: NodeId::NULL,
86+
child: NodeId::NULL,
87+
metadata: None,
88+
}
89+
}
90+
}
91+
92+
impl<'a> streaming_iterator::StreamingIterator for EdgeTableRowView<'a> {
93+
type Item = Self;
94+
95+
row_lending_iterator_get!();
96+
97+
fn advance(&mut self) {
98+
self.id = (i32::from(self.id) + 1).into();
99+
self.left = self.table.left(self.id).unwrap_or_else(|| f64::NAN.into());
100+
self.right = self.table.right(self.id).unwrap_or_else(|| f64::NAN.into());
101+
self.parent = self.table.parent(self.id).unwrap_or(NodeId::NULL);
102+
self.child = self.table.child(self.id).unwrap_or(NodeId::NULL);
103+
self.metadata = self.table.raw_metadata(self.id);
104+
}
105+
}
106+
67107
/// An immutable view of an edge table.
68108
///
69109
/// These are not created directly but are accessed
@@ -94,6 +134,8 @@ impl EdgeTable {
94134
self.as_ref().num_rows.into()
95135
}
96136

137+
raw_metadata_getter_for_tables!(EdgeId);
138+
97139
/// Return the ``parent`` value from row ``row`` of the table.
98140
///
99141
/// # Returns
@@ -193,6 +235,10 @@ impl EdgeTable {
193235
crate::table_iterator::make_table_iterator::<&EdgeTable>(self)
194236
}
195237

238+
pub fn lending_iter(&self) -> EdgeTableRowView {
239+
EdgeTableRowView::new(self)
240+
}
241+
196242
/// Return row `r` of the table.
197243
///
198244
/// # Parameters

‎src/individual_table.rs‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,42 @@ impl PartialEq for IndividualTableRow {
4545
}
4646
}
4747

48+
pub struct IndividualTableRowView<'a> {
49+
table: &'a IndividualTable,
50+
pub id: IndividualId,
51+
pub flags: IndividualFlags,
52+
pub location: Option<&'a [Location]>,
53+
pub parents: Option<&'a [IndividualId]>,
54+
pub metadata: Option<&'a [u8]>,
55+
}
56+
57+
impl<'a> IndividualTableRowView<'a> {
58+
fn new(table: &'a IndividualTable) -> Self {
59+
Self {
60+
table,
61+
id: (-1_i32).into(),
62+
flags: 0.into(),
63+
location: None,
64+
parents: None,
65+
metadata: None,
66+
}
67+
}
68+
}
69+
70+
impl<'a> streaming_iterator::StreamingIterator for IndividualTableRowView<'a> {
71+
type Item = Self;
72+
73+
row_lending_iterator_get!();
74+
75+
fn advance(&mut self) {
76+
self.id = (i32::from(self.id) + 1).into();
77+
self.flags = self.table.flags(self.id).unwrap_or_else(|| 0.into());
78+
self.location = self.table.location(self.id);
79+
self.parents = self.table.parents(self.id);
80+
self.metadata = self.table.raw_metadata(self.id);
81+
}
82+
}
83+
4884
/// An immutable view of a individual table.
4985
///
5086
/// These are not created directly but are accessed
@@ -104,6 +140,8 @@ impl IndividualTable {
104140
unsafe { self.table_.as_ref() }
105141
}
106142

143+
raw_metadata_getter_for_tables!(IndividualId);
144+
107145
/// Return the number of rows
108146
pub fn num_rows(&self) -> crate::SizeType {
109147
self.as_ref().num_rows.into()
@@ -349,6 +387,10 @@ match tables.individuals().metadata::<MutationMetadata>(0.into())
349387
crate::table_iterator::make_table_iterator::<&IndividualTable>(self)
350388
}
351389

390+
pub fn lending_iter(&self) -> IndividualTableRowView {
391+
IndividualTableRowView::new(self)
392+
}
393+
352394
/// Return row `r` of the table.
353395
///
354396
/// # Parameters

‎src/migration_table.rs‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,51 @@ impl Iterator for MigrationTableIterator {
7373
}
7474
}
7575

76+
pub struct MigrationTableRowView<'a> {
77+
table: &'a MigrationTable,
78+
pub id: MigrationId,
79+
pub left: Position,
80+
pub right: Position,
81+
pub node: NodeId,
82+
pub source: PopulationId,
83+
pub dest: PopulationId,
84+
pub time: Time,
85+
pub metadata: Option<&'a [u8]>,
86+
}
87+
88+
impl<'a> MigrationTableRowView<'a> {
89+
fn new(table: &'a MigrationTable) -> Self {
90+
Self {
91+
table,
92+
id: MigrationId::NULL,
93+
left: Position::from(f64::NAN),
94+
right: Position::from(f64::NAN),
95+
node: NodeId::NULL,
96+
source: PopulationId::NULL,
97+
dest: PopulationId::NULL,
98+
time: Time::from(f64::NAN),
99+
metadata: None,
100+
}
101+
}
102+
}
103+
104+
impl<'a> streaming_iterator::StreamingIterator for MigrationTableRowView<'a> {
105+
type Item = Self;
106+
107+
row_lending_iterator_get!();
108+
109+
fn advance(&mut self) {
110+
self.id = (i32::from(self.id) + 1).into();
111+
self.left = self.table.left(self.id).unwrap_or_else(|| f64::NAN.into());
112+
self.right = self.table.right(self.id).unwrap_or_else(|| f64::NAN.into());
113+
self.node = self.table.node(self.id).unwrap_or(NodeId::NULL);
114+
self.source = self.table.source(self.id).unwrap_or(PopulationId::NULL);
115+
self.dest = self.table.dest(self.id).unwrap_or(PopulationId::NULL);
116+
self.time = self.table.time(self.id).unwrap_or_else(|| f64::NAN.into());
117+
self.metadata = self.table.raw_metadata(self.id);
118+
}
119+
}
120+
76121
/// An immutable view of a migration table.
77122
///
78123
/// These are not created directly but are accessed
@@ -102,6 +147,8 @@ impl MigrationTable {
102147
self.as_ref().num_rows.into()
103148
}
104149

150+
raw_metadata_getter_for_tables!(MigrationId);
151+
105152
/// Return the left coordinate for a given row.
106153
///
107154
/// # Returns
@@ -232,6 +279,10 @@ impl MigrationTable {
232279
crate::table_iterator::make_table_iterator::<&MigrationTable>(self)
233280
}
234281

282+
pub fn lending_iter(&self) -> MigrationTableRowView {
283+
MigrationTableRowView::new(self)
284+
}
285+
235286
/// Return row `r` of the table.
236287
///
237288
/// # Parameters

‎src/mutation_table.rs‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,48 @@ impl Iterator for MutationTableIterator {
7676
}
7777
}
7878

79+
pub struct MutationTableRowView<'a> {
80+
table: &'a MutationTable,
81+
pub id: MutationId,
82+
pub site: SiteId,
83+
pub node: NodeId,
84+
pub parent: MutationId,
85+
pub time: Time,
86+
pub derived_state: Option<&'a [u8]>,
87+
pub metadata: Option<&'a [u8]>,
88+
}
89+
90+
impl<'a> MutationTableRowView<'a> {
91+
fn new(table: &'a MutationTable) -> Self {
92+
Self {
93+
table,
94+
id: MutationId::NULL,
95+
site: SiteId::NULL,
96+
node: NodeId::NULL,
97+
parent: MutationId::NULL,
98+
time: f64::NAN.into(),
99+
derived_state: None,
100+
metadata: None,
101+
}
102+
}
103+
}
104+
105+
impl<'a> streaming_iterator::StreamingIterator for MutationTableRowView<'a> {
106+
type Item = Self;
107+
108+
row_lending_iterator_get!();
109+
110+
fn advance(&mut self) {
111+
self.id = (i32::from(self.id) + 1).into();
112+
self.site = self.table.site(self.id).unwrap_or(SiteId::NULL);
113+
self.node = self.table.node(self.id).unwrap_or(NodeId::NULL);
114+
self.parent = self.table.parent(self.id).unwrap_or(MutationId::NULL);
115+
self.time = self.table.time(self.id).unwrap_or_else(|| f64::NAN.into());
116+
self.derived_state = self.table.derived_state(self.id);
117+
self.metadata = self.table.raw_metadata(self.id);
118+
}
119+
}
120+
79121
/// An immutable view of site table.
80122
///
81123
/// These are not created directly but are accessed
@@ -105,6 +147,8 @@ impl MutationTable {
105147
self.as_ref().num_rows.into()
106148
}
107149

150+
raw_metadata_getter_for_tables!(MutationId);
151+
108152
/// Return the ``site`` value from row ``row`` of the table.
109153
///
110154
/// # Errors
@@ -218,6 +262,10 @@ impl MutationTable {
218262
crate::table_iterator::make_table_iterator::<&MutationTable>(self)
219263
}
220264

265+
pub fn lending_iter(&self) -> MutationTableRowView {
266+
MutationTableRowView::new(self)
267+
}
268+
221269
/// Return row `r` of the table.
222270
///
223271
/// # Parameters

‎src/node_table.rs‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,45 @@ impl Iterator for NodeTableIterator {
6666
}
6767
}
6868

69+
pub struct NodeTableRowView<'a> {
70+
table: &'a NodeTable,
71+
pub id: NodeId,
72+
pub time: Time,
73+
pub flags: NodeFlags,
74+
pub population: PopulationId,
75+
pub individual: IndividualId,
76+
pub metadata: Option<&'a [u8]>,
77+
}
78+
79+
impl<'a> NodeTableRowView<'a> {
80+
fn new(table: &'a NodeTable) -> Self {
81+
Self {
82+
table,
83+
id: NodeId::NULL,
84+
time: f64::NAN.into(),
85+
flags: 0.into(),
86+
population: PopulationId::NULL,
87+
individual: IndividualId::NULL,
88+
metadata: None,
89+
}
90+
}
91+
}
92+
93+
impl<'a> streaming_iterator::StreamingIterator for NodeTableRowView<'a> {
94+
type Item = Self;
95+
96+
row_lending_iterator_get!();
97+
98+
fn advance(&mut self) {
99+
self.id = (i32::from(self.id) + 1).into();
100+
self.time = self.table.time(self.id).unwrap_or_else(|| f64::NAN.into());
101+
self.flags = self.table.flags(self.id).unwrap_or_else(|| 0.into());
102+
self.population = self.table.population(self.id).unwrap_or(PopulationId::NULL);
103+
self.individual = self.table.individual(self.id).unwrap_or(IndividualId::NULL);
104+
self.metadata = self.table.raw_metadata(self.id);
105+
}
106+
}
107+
69108
/// An immtable view of a node table.
70109
///
71110
/// These are not created directly but are accessed
@@ -95,6 +134,8 @@ impl NodeTable {
95134
self.as_ref().num_rows.into()
96135
}
97136

137+
raw_metadata_getter_for_tables!(NodeId);
138+
98139
/// Return the ``time`` value from row ``row`` of the table.
99140
///
100141
/// # Returns
@@ -386,6 +427,10 @@ impl NodeTable {
386427
crate::table_iterator::make_table_iterator::<&NodeTable>(self)
387428
}
388429

430+
pub fn lending_iter(&self) -> NodeTableRowView {
431+
NodeTableRowView::new(self)
432+
}
433+
389434
/// Return row `r` of the table.
390435
///
391436
/// # Parameters

0 commit comments

Comments
 (0)