Add mutable access to NodeTable flag and time arrays.

Author: molpopgen

src/node_table.rs

@@ -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

src/table_collection.rs

@@ -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();