details

Author: molpopgen

src/edge_differences.rs

@@ -76,37 +76,107 @@ pub type EdgeRemoval = EdgeDifference<Removal>;
 /// Manages iteration over trees to obtain
 /// edge differences.
 pub struct EdgeDifferencesIterator<'ts> {
-    edge_left: &'ts [Position],
-    edge_right: &'ts [Position],
-    edge_parent: &'ts [NodeId],
-    edge_child: &'ts [NodeId],
+    edges_left: &'ts [Position],
+    edges_right: &'ts [Position],
+    edges_parent: &'ts [NodeId],
+    edges_child: &'ts [NodeId],
     insertion_order: &'ts [EdgeId],
-    removal_orderr: &'ts [EdgeId],
+    removal_order: &'ts [EdgeId],
     left: f64,
+    sequence_length: f64,
+    insertion_index: usize,
+    removal_index: usize,
 }
 
 impl<'ts> EdgeDifferencesIterator<'ts> {
     pub(crate) fn new(treeseq: &'ts TreeSequence) -> Self {
         Self {
-            edge_left: treeseq.tables().edges().left_slice(),
-            edge_right: treeseq.tables().edges().right_slice(),
-            edge_parent: treeseq.tables().edges().parent_slice(),
-            edge_child: treeseq.tables().edges().child_slice(),
+            edges_left: treeseq.tables().edges().left_slice(),
+            edges_right: treeseq.tables().edges().right_slice(),
+            edges_parent: treeseq.tables().edges().parent_slice(),
+            edges_child: treeseq.tables().edges().child_slice(),
             insertion_order: treeseq.edge_insertion_order(),
-            removal_orderr: treeseq.edge_removal_order(),
+            removal_order: treeseq.edge_removal_order(),
             left: 0.,
+            sequence_length: treeseq.tables().sequence_length().into(),
+            insertion_index: 0,
+            removal_index: 0,
         }
     }
 }
 
 pub struct EdgeDifferences<'ts> {
-    marker: std::marker::PhantomData<&'ts ()>,
+    edges_left: &'ts [Position],
+    edges_right: &'ts [Position],
+    edges_parent: &'ts [NodeId],
+    edges_child: &'ts [NodeId],
+    removals: (usize, usize),
+    insertions: (usize, usize),
+    left: f64,
+    right: f64,
+}
+
+fn update_right(
+    right: f64,
+    index: usize,
+    position_slice: &[Position],
+    diff_slice: &[EdgeId],
+) -> f64 {
+    if index < diff_slice.len() {
+        let temp = position_slice[diff_slice[index].as_usize()];
+        if temp < right {
+            temp.into()
+        } else {
+            right
+        }
+    } else {
+        right
+    }
 }
 
 impl<'ts> Iterator for EdgeDifferencesIterator<'ts> {
     type Item = EdgeDifferences<'ts>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        todo!()
+        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.removal_index += 1;
+            }
+            let right = update_right(
+                self.sequence_length,
+                self.insertion_index,
+                self.edges_left,
+                self.insertion_order,
+            );
+            let right = update_right(
+                right,
+                self.removal_index,
+                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,
+                removals: (removals_start, self.removal_index),
+                insertions: (insertions_start, self.insertion_index),
+                left: self.left,
+                right,
+            };
+            self.left = right;
+            Some(diffs)
+        } else {
+            None
+        }
     }
 }