tests pass

Author: molpopgen

src/edge_differences.rs

@@ -111,6 +111,8 @@ pub struct EdgeDifferences<'ts> {
     edges_right: &'ts [Position],
     edges_parent: &'ts [NodeId],
     edges_child: &'ts [NodeId],
+    insertion_order: &'ts [EdgeId],
+    removal_order: &'ts [EdgeId],
     removals: (usize, usize),
     insertions: (usize, usize),
     left: f64,
@@ -122,6 +124,7 @@ pub struct X<'ts, T: EdgeDifferenceIteration> {
     edges_right: &'ts [Position],
     edges_parent: &'ts [NodeId],
     edges_child: &'ts [NodeId],
+    order: &'ts [EdgeId],
     range: (usize, usize),
     marker: std::marker::PhantomData<T>,
 }
@@ -133,10 +136,10 @@ impl<'ts, T: EdgeDifferenceIteration> Iterator for X<'ts, T> {
             let index = self.range.0;
             self.range.0 += 1;
             Some(Self::Item::new(
-                self.edges_left[index],
-                self.edges_right[index],
-                self.edges_parent[index],
-                self.edges_child[index],
+                self.edges_left[self.order[index].as_usize()],
+                self.edges_right[self.order[index].as_usize()],
+                self.edges_parent[self.order[index].as_usize()],
+                self.edges_child[self.order[index].as_usize()],
             ))
         } else {
             None
@@ -151,6 +154,7 @@ impl<'ts> EdgeDifferences<'ts> {
             edges_right: self.edges_right,
             edges_parent: self.edges_parent,
             edges_child: self.edges_child,
+            order: self.removal_order,
             range: self.removals,
             marker: std::marker::PhantomData::<Removal> {},
         }
@@ -162,6 +166,7 @@ impl<'ts> EdgeDifferences<'ts> {
             edges_right: self.edges_right,
             edges_parent: self.edges_parent,
             edges_child: self.edges_child,
+            order: self.insertion_order,
             range: self.insertions,
             marker: std::marker::PhantomData::<Insertion> {},
         }
@@ -193,10 +198,17 @@ impl<'ts> Iterator for EdgeDifferencesIterator<'ts> {
         if self.insertion_index < self.insertion_order.len() && self.left < self.sequence_length {
             let removals_start = self.removal_index;
             while self.removal_index < self.removal_order.len()
-                && self.edges_right[self.removal_index] == self.left
+                && self.edges_right[self.removal_order[self.removal_index].as_usize()] == self.left
             {
                 self.removal_index += 1;
             }
+            let insertions_start = self.insertion_index;
+            while self.insertion_index < self.insertion_order.len()
+                && self.edges_left[self.insertion_order[self.insertion_index].as_usize()]
+                    == self.left
+            {
+                self.insertion_index += 1;
+            }
             let right = update_right(
                 self.sequence_length,
                 self.insertion_index,
@@ -209,17 +221,13 @@ impl<'ts> Iterator for EdgeDifferencesIterator<'ts> {
                 self.edges_right,
                 self.removal_order,
             );
-            let insertions_start = self.insertion_index;
-            while self.insertion_index < self.insertion_order.len()
-                && self.edges_left[self.insertion_index] == self.left
-            {
-                self.insertion_index += 1;
-            }
             let diffs = EdgeDifferences {
                 edges_left: self.edges_left,
                 edges_right: self.edges_right,
                 edges_parent: self.edges_parent,
                 edges_child: self.edges_child,
+                insertion_order: self.insertion_order,
+                removal_order: self.removal_order,
                 removals: (removals_start, self.removal_index),
                 insertions: (insertions_start, self.insertion_index),
                 left: self.left,

tests/book_trees.rs

@@ -170,10 +170,10 @@ fn initialize_from_table_collection() {
     // ANCHOR: iterate_edge_differences
     for diffs in treeseq.edge_differences_iter() {
         for edge_removal in diffs.removals() {
-            println!("{}", edge_removal);
+            println!("edge removal: {}", edge_removal);
         }
         for edge_insertion in diffs.insertions() {
-            println!("{}", edge_insertion);
+            println!("edge insertion: {}", edge_insertion);
         }
     }
     // ANCHOR_END: iterate_edge_differences