Refactor to remove iteration

Author: benjeffery

python/tests/test_jit.py

@@ -24,16 +24,17 @@ def test_correct_trees_forward(ts):
     numba_ts = jit_numba.numba_tree_sequence(ts)
     in_index = ts.indexes_edge_insertion_order
     out_index = ts.indexes_edge_removal_order
-    for numba_edge_diff, edge_diff in itertools.zip_longest(
-        numba_ts.edge_diffs(), ts.edge_diffs()
-    ):
-        assert edge_diff.interval == numba_edge_diff.interval
+    tree_pos = numba_ts.tree_position()
+    ts_edge_diffs = ts.edge_diffs()
+    while tree_pos.next():
+        edge_diff = next(ts_edge_diffs)
+        assert edge_diff.interval == tree_pos.interval
         for edge_in_index, edge in itertools.zip_longest(
-            range(*numba_edge_diff.edges_in_index_range), edge_diff.edges_in
+            range(*tree_pos.edges_in_index_range), edge_diff.edges_in
         ):
             assert edge.id == in_index[edge_in_index]
         for edge_out_index, edge in itertools.zip_longest(
-            range(*numba_edge_diff.edges_out_index_range), edge_diff.edges_out
+            range(*tree_pos.edges_out_index_range), edge_diff.edges_out
         ):
             assert edge.id == out_index[edge_out_index]
 
@@ -52,7 +53,8 @@ def test_using_from_jit_function():
     def _coalescent_nodes_numba(numba_ts, num_nodes, edges_parent):
         is_coalescent = np.zeros(num_nodes, dtype=np.int8)
         num_children = np.zeros(num_nodes, dtype=np.int64)
-        for tree_pos in numba_ts.edge_diffs():
+        tree_pos = numba_ts.tree_position()
+        while tree_pos.next():
             for j in range(*tree_pos.edges_out_index_range):
                 e = numba_ts.indexes_edge_removal_order[j]
                 num_children[edges_parent[e]] -= 1

python/tskit/jit/numba.py

@@ -30,13 +30,6 @@ def jitdataclass(cls):
     return numba.experimental.jitclass(dc_cls)
 
 
-@jitdataclass
-class NumbaEdgeDiff:
-    interval: numba.types.UniTuple(numba.float64, 2)
-    edges_in_index_range: numba.types.UniTuple(numba.int32, 2)
-    edges_out_index_range: numba.types.UniTuple(numba.int32, 2)
-
-
 @jitdataclass
 class NumbaTreeSequence:
     num_edges: numba.int64
@@ -46,39 +39,44 @@ class NumbaTreeSequence:
     indexes_edge_insertion_order: numba.int32[:]
     indexes_edge_removal_order: numba.int32[:]
 
-    def edge_diffs(self, include_terminal=False):
-        left = 0.0
-        j = 0
-        k = 0
-        edges_left = self.edges_left
-        edges_right = self.edges_right
-        in_order = self.indexes_edge_insertion_order
-        out_order = self.indexes_edge_removal_order
-
-        while j < self.num_edges or left < self.sequence_length:
-            in_start = j
-            out_start = k
-
-            while k < self.num_edges and edges_right[out_order[k]] == left:
-                k += 1
-            while j < self.num_edges and edges_left[in_order[j]] == left:
-                j += 1
-            in_end = j
-            out_end = k
-
-            right = self.sequence_length
-            if j < self.num_edges:
-                right = min(right, edges_left[in_order[j]])
-            if k < self.num_edges:
-                right = min(right, edges_right[out_order[k]])
-
-            yield NumbaEdgeDiff((left, right), (in_start, in_end), (out_start, out_end))
-
-            left = right
-
-        # Handle remaining edges that haven't been processed
-        if include_terminal:
-            yield NumbaEdgeDiff((left, right), (j, j), (k, self.num_edges))
+    def tree_position(self):
+        return NumbaTreePosition(self, (0, 0), (0, 0), (0, 0))
+
+
+@jitdataclass
+class NumbaTreePosition:
+    ts: NumbaTreeSequence
+    interval: numba.types.UniTuple(numba.float64, 2)
+    edges_in_index_range: numba.types.UniTuple(numba.int32, 2)
+    edges_out_index_range: numba.types.UniTuple(numba.int32, 2)
+
+    def next(self):  # noqa: A003
+        M = self.ts.num_edges
+        edges_left = self.ts.edges_left
+        edges_right = self.ts.edges_right
+        in_order = self.ts.indexes_edge_insertion_order
+        out_order = self.ts.indexes_edge_removal_order
+
+        left = self.interval[1]
+        j = self.edges_in_index_range[1]
+        k = self.edges_out_index_range[1]
+
+        while k < M and edges_right[out_order[k]] == left:
+            k += 1
+        while j < M and edges_left[in_order[j]] == left:
+            j += 1
+
+        self.edges_in_index_range = (self.edges_in_index_range[1], j)
+        self.edges_out_index_range = (self.edges_out_index_range[1], k)
+
+        right = self.ts.sequence_length
+        if j < M:
+            right = min(right, edges_left[in_order[j]])
+        if k < M:
+            right = min(right, edges_right[out_order[k]])
+
+        self.interval = (left, right)
+        return j < M or left < self.ts.sequence_length
 
 
 def numba_tree_sequence(ts):