use bookmark for WF example program

Author: molpopgen

examples/haploid_wright_fisher.rs

@@ -8,6 +8,23 @@ use proptest::prelude::*;
 use rand::distributions::Distribution;
 use rand::SeedableRng;
 
+fn rotate_edge_table(rotation_point: usize, tables: &mut tskit::TableCollection) {
+    let num_edges = tables.edges().num_rows().as_usize();
+    // SAFETY: tables pointer is not null.
+    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 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) };
+    left.rotate_left(rotation_point);
+    right.rotate_left(rotation_point);
+    parent.rotate_left(rotation_point);
+    child.rotate_left(rotation_point);
+}
+
 // ANCHOR: haploid_wright_fisher
 fn simulate(
     seed: u64,
@@ -46,6 +63,7 @@ fn simulate(
     let parent_picker = rand::distributions::Uniform::new(0, popsize);
     let breakpoint_generator = rand::distributions::Uniform::new(0.0, 1.0);
     let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
+    let mut bookmark = tskit::types::Bookmark::new();
 
     for birth_time in (0..num_generations).rev() {
         for c in children.iter_mut() {
@@ -64,7 +82,10 @@ fn simulate(
         }
 
         if birth_time % simplify_interval == 0 {
-            tables.full_sort(tskit::TableSortOptions::default())?;
+            tables.sort(&bookmark, 0)?;
+            if bookmark.edges() > 0 {
+                rotate_edge_table(bookmark.edges().as_usize(), &mut tables);
+            }
             if let Some(idmap) =
                 tables.simplify(children, tskit::SimplificationOptions::default(), true)?
             {
@@ -73,6 +94,7 @@ fn simulate(
                     *o = idmap[usize::try_from(*o)?];
                 }
             }
+            bookmark.set_edges(tables.edges().num_rows());
         }
         std::mem::swap(&mut parents, &mut children);
     }