remove code duplication

tskit-dev · Aug 15, 2022 · a3fa336 · a3fa336
1 parent dbb8224
commit a3fa336
Showing 1 changed file with 20 additions and 35 deletions.
diff --git a/examples/forward_simulation.rs b/examples/forward_simulation.rs
@@ -320,6 +320,24 @@ fn births(
     }
 }
 
+fn rotate_edge_table(mid: usize, tables: &mut tskit::TableCollection) {
+    // NOTE: using unsafe here because we don't have
+    // a rust API yet.
+    let num_edges: usize = tables.edges().num_rows().try_into().unwrap();
+    let parent =
+        unsafe { std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.parent, num_edges) };
+    let child =
+        unsafe { std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.child, num_edges) };
+    let left =
+        unsafe { std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.left, num_edges) };
+    let right =
+        unsafe { std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.right, num_edges) };
+    parent.rotate_left(mid);
+    child.rotate_left(mid);
+    left.rotate_left(mid);
+    right.rotate_left(mid);
+}
+
 fn simplify(
     bookmark: &tskit::types::Bookmark,
     alive: &mut [Diploid],
@@ -338,25 +356,8 @@ fn simplify(
     }
 
     if bookmark.offsets.edges > 0 {
-        // NOTE: there should be API support for this operation
-        unimplemented!("fix code duplication");
-        let num_edges: usize = tables.edges().num_rows().try_into().unwrap();
-        let parent = unsafe {
-            std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.parent, num_edges)
-        };
-        let child = unsafe {
-            std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.child, num_edges)
-        };
-        let left =
-            unsafe { std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.left, num_edges) };
-        let right = unsafe {
-            std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.right, num_edges)
-        };
         let mid: usize = bookmark.offsets.edges.try_into().unwrap();
-        parent.rotate_left(mid);
-        child.rotate_left(mid);
-        left.rotate_left(mid);
-        right.rotate_left(mid);
+        rotate_edge_table(mid, tables);
     }
 
     match tables.simplify(
@@ -431,24 +432,8 @@ fn update_bookmark(
             .find(|(_index, time)| **time > most_ancient_birth_time)
         {
             Some((index, _time)) => {
-                unimplemented!("fix code duplication");
+                rotate_edge_table(index, tables);
                 let num_edges: usize = tables.edges().num_rows().try_into().unwrap();
-                let parent = unsafe {
-                    std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.parent, num_edges)
-                };
-                let child = unsafe {
-                    std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.child, num_edges)
-                };
-                let left = unsafe {
-                    std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.left, num_edges)
-                };
-                let right = unsafe {
-                    std::slice::from_raw_parts_mut((*tables.as_mut_ptr()).edges.right, num_edges)
-                };
-                parent.rotate_left(index);
-                child.rotate_left(index);
-                left.rotate_left(index);
-                right.rotate_left(index);
                 bookmark.offsets.edges = (num_edges - index).try_into().unwrap();
             }
             None => bookmark.offsets.edges = 0,