Skip to content

Commit

Permalink
Add mutable access to NodeTable flag and time arrays. (#125)
Browse files Browse the repository at this point in the history
Closes #121
  • Loading branch information
molpopgen committed May 19, 2021
1 parent 634df56 commit edd36fa
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 0 deletions.
10 changes: 10 additions & 0 deletions src/node_table.rs
Expand Up @@ -104,6 +104,16 @@ impl<'a> NodeTable<'a> {
unsafe_tsk_column_access!(row, 0, self.num_rows(), self.table_.flags)
}

/// Mutable access to node flags.
pub fn flags_array_mut(&mut self) -> &mut [tsk_flags_t] {
unsafe { std::slice::from_raw_parts_mut(self.table_.flags, self.table_.num_rows as usize) }
}

/// Mutable access to node times.
pub fn time_array_mut(&mut self) -> &mut [f64] {
unsafe { std::slice::from_raw_parts_mut(self.table_.time, self.table_.num_rows as usize) }
}

/// Return the ``population`` value from row ``row`` of the table.
///
/// # Errors
Expand Down
19 changes: 19 additions & 0 deletions src/table_collection.rs
Expand Up @@ -696,6 +696,25 @@ mod test {
}
}

#[test]
fn test_mutable_node_access() {
let tables = TableCollection::new(1000.).unwrap();
let mut nodes = tables.nodes();
let f = nodes.flags_array_mut();
for i in f {
*i = 11;
}

for t in nodes.time_array_mut() {
*t = -33.0;
}

for i in tables.nodes_iter(true) {
assert_eq!(i.flags, 11);
assert_eq!(i.time as i64, -33);
}
}

#[test]
fn test_node_iteration() {
let tables = make_small_table_collection();
Expand Down

0 comments on commit edd36fa

Please sign in to comment.