Skip to content

Commit a95c8cc

Browse files
molpopgenmolpopgen
committed
Add IndividualId.
1 parent 8bf7d4b commit a95c8cc

File tree

4 files changed

+82
-36
lines changed

4 files changed

+82
-36
lines changed

‎src/_macros.rs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,30 @@ macro_rules! unsafe_tsk_ragged_column_access {
6868
}
6969
}
7070
}};
71+
72+
($i: expr, $lo: expr, $hi: expr, $array: expr, $offset_array: expr, $offset_array_len: expr, $output_id_type: expr) => {{
73+
if $i < $lo || $i >= ($hi as tsk_id_t) {
74+
Err(TskitError::IndexError {})
75+
} else if $offset_array_len == 0 {
76+
Ok(None)
77+
} else {
78+
let start = unsafe { *$offset_array.offset($i as isize) };
79+
let stop = if $i < ($hi as tsk_id_t) {
80+
unsafe { *$offset_array.offset(($i + 1) as isize) }
81+
} else {
82+
$offset_array_len as tsk_size_t
83+
};
84+
if start == stop {
85+
Ok(None)
86+
} else {
87+
let mut buffer = vec![];
88+
for i in start..stop {
89+
buffer.push($output_id_type(unsafe { *$array.offset(i as isize) }));
90+
}
91+
Ok(Some(buffer))
92+
}
93+
}
94+
}};
7195
}
7296

7397
// Allow this to be unused for default features

‎src/individual_table.rs‎

Lines changed: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
11
use crate::bindings as ll_bindings;
22
use crate::metadata;
3+
use crate::IndividualId;
34
use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TskitError};
45

56
/// Row of a [`IndividualTable`]
67
pub struct IndividualTableRow {
7-
pub id: tsk_id_t,
8+
pub id: IndividualId,
89
pub flags: tsk_flags_t,
910
pub location: Option<Vec<f64>>,
10-
pub parents: Option<Vec<tsk_id_t>>,
11+
pub parents: Option<Vec<IndividualId>>,
1112
pub metadata: Option<Vec<u8>>,
1213
}
1314

@@ -50,7 +51,7 @@ pub struct IndividualTable<'a> {
5051
fn make_individual_table_row(table: &IndividualTable, pos: tsk_id_t) -> Option<IndividualTableRow> {
5152
if pos < table.num_rows() as tsk_id_t {
5253
let rv = IndividualTableRow {
53-
id: pos,
54+
id: pos.into(),
5455
flags: table.flags(pos).unwrap(),
5556
location: table.location(pos).unwrap(),
5657
parents: table.parents(pos).unwrap(),
@@ -83,18 +84,21 @@ impl<'a> IndividualTable<'a> {
8384
/// # Errors
8485
///
8586
/// * [`TskitError::IndexError`] if `row` is out of range.
86-
pub fn flags(&self, row: tsk_id_t) -> Result<tsk_flags_t, TskitError> {
87-
unsafe_tsk_column_access!(row, 0, self.num_rows(), self.table_.flags)
87+
pub fn flags<I: Into<IndividualId> + Copy>(&self, row: I) -> Result<tsk_flags_t, TskitError> {
88+
unsafe_tsk_column_access!(row.into().0, 0, self.num_rows(), self.table_.flags)
8889
}
8990

9091
/// Return the locations for a given row.
9192
///
9293
/// # Errors
9394
///
9495
/// * [`TskitError::IndexError`] if `row` is out of range.
95-
pub fn location(&self, row: tsk_id_t) -> Result<Option<Vec<f64>>, TskitError> {
96+
pub fn location<I: Into<IndividualId> + Copy>(
97+
&self,
98+
row: I,
99+
) -> Result<Option<Vec<f64>>, TskitError> {
96100
unsafe_tsk_ragged_column_access!(
97-
row,
101+
row.into().0,
98102
0,
99103
self.num_rows(),
100104
self.table_.location,
@@ -108,14 +112,18 @@ impl<'a> IndividualTable<'a> {
108112
/// # Errors
109113
///
110114
/// * [`TskitError::IndexError`] if `row` is out of range.
111-
pub fn parents(&self, row: tsk_id_t) -> Result<Option<Vec<tsk_id_t>>, TskitError> {
115+
pub fn parents<I: Into<IndividualId> + Copy>(
116+
&self,
117+
row: I,
118+
) -> Result<Option<Vec<IndividualId>>, TskitError> {
112119
unsafe_tsk_ragged_column_access!(
113-
row,
120+
row.into().0,
114121
0,
115122
self.num_rows(),
116123
self.table_.parents,
117124
self.table_.parents_offset,
118-
self.table_.parents_length
125+
self.table_.parents_length,
126+
IndividualId
119127
)
120128
}
121129

@@ -124,11 +132,11 @@ impl<'a> IndividualTable<'a> {
124132
/// # Errors
125133
///
126134
/// * [`TskitError::IndexError`] if `row` is out of range.
127-
pub fn metadata<T: metadata::MetadataRoundtrip>(
135+
pub fn metadata<I: Into<IndividualId>, T: metadata::MetadataRoundtrip>(
128136
&'a self,
129-
row: tsk_id_t,
137+
row: I,
130138
) -> Result<Option<T>, TskitError> {
131-
let buffer = metadata_to_vector!(self, row)?;
139+
let buffer = metadata_to_vector!(self, row.into().0)?;
132140
decode_metadata_row!(T, buffer)
133141
}
134142

@@ -148,7 +156,10 @@ impl<'a> IndividualTable<'a> {
148156
/// # Errors
149157
///
150158
/// [`TskitError::IndexError`] if `r` is out of range.
151-
pub fn row(&self, r: tsk_id_t) -> Result<IndividualTableRow, TskitError> {
152-
table_row_access!(r, self, make_individual_table_row)
159+
pub fn row<I: Into<IndividualId> + Copy>(
160+
&self,
161+
r: I,
162+
) -> Result<IndividualTableRow, TskitError> {
163+
table_row_access!(r.into().0, self, make_individual_table_row)
153164
}
154165
}

‎src/lib.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,17 @@ pub use bindings::tsk_size_t;
142142
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, std::hash::Hash)]
143143
pub struct NodeId(tsk_id_t);
144144

145+
/// An individual ID
146+
///
147+
/// This is an integer referring to a row of an [``IndividualTable``].
148+
///
149+
/// The features for this type follow the same pattern as for [``NodeId``]
150+
#[repr(transparent)]
151+
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, std::hash::Hash)]
152+
pub struct IndividualId(tsk_id_t);
153+
145154
impl_id_traits!(NodeId);
155+
impl_id_traits!(IndividualId);
146156

147157
// tskit defines this via a type cast
148158
// in a macro. bindgen thus misses it.

‎src/table_collection.rs‎

Lines changed: 22 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use crate::TreeSequenceFlags;
2020
use crate::TskReturnValue;
2121
use crate::TskitTypeAccess;
2222
use crate::{tsk_flags_t, tsk_id_t, tsk_size_t, TSK_NULL};
23+
use crate::{IndividualId, NodeId};
2324
use ll_bindings::tsk_table_collection_free;
2425

2526
/// A table collection.
@@ -165,7 +166,7 @@ impl TableCollection {
165166
}
166167

167168
/// Add a row to the edge table
168-
pub fn add_edge<P: Into<crate::NodeId>, C: Into<crate::NodeId>>(
169+
pub fn add_edge<P: Into<NodeId>, C: Into<NodeId>>(
169170
&mut self,
170171
left: f64,
171172
right: f64,
@@ -176,7 +177,7 @@ impl TableCollection {
176177
}
177178

178179
/// Add a row with metadata to the edge table
179-
pub fn add_edge_with_metadata<P: Into<crate::NodeId>, C: Into<crate::NodeId>>(
180+
pub fn add_edge_with_metadata<P: Into<NodeId>, C: Into<NodeId>>(
180181
&mut self,
181182
left: f64,
182183
right: f64,
@@ -201,23 +202,23 @@ impl TableCollection {
201202
}
202203

203204
/// Add a row to the individual table
204-
pub fn add_individual<N: Into<crate::NodeId>>(
205+
pub fn add_individual<N: Into<NodeId>>(
205206
&mut self,
206207
flags: tsk_flags_t,
207208
location: &[f64],
208209
parents: &[N],
209-
) -> TskReturnValue {
210+
) -> Result<IndividualId, TskitError> {
210211
self.add_individual_with_metadata(flags, location, parents, None)
211212
}
212213

213214
/// Add a row with metadata to the individual table
214-
pub fn add_individual_with_metadata<N: Into<crate::NodeId>>(
215+
pub fn add_individual_with_metadata<N: Into<NodeId>>(
215216
&mut self,
216217
flags: tsk_flags_t,
217218
location: &[f64],
218219
parents: &[N],
219220
metadata: Option<&dyn MetadataRoundtrip>,
220-
) -> TskReturnValue {
221+
) -> Result<IndividualId, TskitError> {
221222
let md = EncodedMetadata::new(metadata)?;
222223
let rv = unsafe {
223224
ll_bindings::tsk_individual_table_add_row(
@@ -231,7 +232,7 @@ impl TableCollection {
231232
md.len(),
232233
)
233234
};
234-
handle_tsk_return_value!(rv)
235+
handle_tsk_return_value!(rv, IndividualId::from(rv))
235236
}
236237

237238
/// Add a row to the migration table
@@ -240,7 +241,7 @@ impl TableCollection {
240241
///
241242
/// Migration tables are not currently supported
242243
/// by tree sequence simplification.
243-
pub fn add_migration<N: Into<crate::NodeId>>(
244+
pub fn add_migration<N: Into<NodeId>>(
244245
&mut self,
245246
span: (f64, f64),
246247
node: N,
@@ -256,7 +257,7 @@ impl TableCollection {
256257
///
257258
/// Migration tables are not currently supported
258259
/// by tree sequence simplification.
259-
pub fn add_migration_with_metadata<N: Into<crate::NodeId>>(
260+
pub fn add_migration_with_metadata<N: Into<NodeId>>(
260261
&mut self,
261262
span: (f64, f64),
262263
node: N,
@@ -282,33 +283,33 @@ impl TableCollection {
282283
}
283284

284285
/// Add a row to the node table
285-
pub fn add_node(
286+
pub fn add_node<I: Into<IndividualId>>(
286287
&mut self,
287288
flags: ll_bindings::tsk_flags_t,
288289
time: f64,
289290
population: tsk_id_t,
290-
individual: tsk_id_t,
291-
) -> Result<crate::NodeId, TskitError> {
291+
individual: I,
292+
) -> Result<NodeId, TskitError> {
292293
self.add_node_with_metadata(flags, time, population, individual, None)
293294
}
294295

295296
/// Add a row with metadata to the node table
296-
pub fn add_node_with_metadata(
297+
pub fn add_node_with_metadata<I: Into<IndividualId>>(
297298
&mut self,
298299
flags: ll_bindings::tsk_flags_t,
299300
time: f64,
300301
population: tsk_id_t,
301-
individual: tsk_id_t,
302+
individual: I,
302303
metadata: Option<&dyn MetadataRoundtrip>,
303-
) -> Result<crate::NodeId, TskitError> {
304+
) -> Result<NodeId, TskitError> {
304305
let md = EncodedMetadata::new(metadata)?;
305306
let rv = unsafe {
306307
ll_bindings::tsk_node_table_add_row(
307308
&mut (*self.as_mut_ptr()).nodes,
308309
flags,
309310
time,
310311
population,
311-
individual,
312+
individual.into().0,
312313
md.as_ptr(),
313314
md.len(),
314315
)
@@ -347,7 +348,7 @@ impl TableCollection {
347348
}
348349

349350
/// Add a row to the mutation table.
350-
pub fn add_mutation<N: Into<crate::NodeId>>(
351+
pub fn add_mutation<N: Into<NodeId>>(
351352
&mut self,
352353
site: tsk_id_t,
353354
node: N,
@@ -359,7 +360,7 @@ impl TableCollection {
359360
}
360361

361362
/// Add a row with metadata to the mutation table.
362-
pub fn add_mutation_with_metadata<N: Into<crate::NodeId>>(
363+
pub fn add_mutation_with_metadata<N: Into<NodeId>>(
363364
&mut self,
364365
site: tsk_id_t,
365366
node: N,
@@ -558,13 +559,13 @@ impl TableCollection {
558559
/// in length to the input node table. For each input node,
559560
/// this vector either contains the node's new index or [`TSK_NULL`]
560561
/// if the input node is not part of the simplified history.
561-
pub fn simplify<N: Into<crate::NodeId>>(
562+
pub fn simplify<N: Into<NodeId>>(
562563
&mut self,
563564
samples: &[N],
564565
options: SimplificationOptions,
565566
idmap: bool,
566-
) -> Result<Option<Vec<crate::NodeId>>, TskitError> {
567-
let mut output_node_map: Vec<crate::NodeId> = vec![];
567+
) -> Result<Option<Vec<NodeId>>, TskitError> {
568+
let mut output_node_map: Vec<NodeId> = vec![];
568569
if idmap {
569570
output_node_map.resize(self.nodes().num_rows() as usize, TSK_NULL.into());
570571
}

0 commit comments

Comments
 (0)