More tests for seek_from_null

Also minor style tweaks

Author: jeromekelleher

c/tskit/trees.c

@@ -4581,8 +4581,8 @@ tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(optio
     const tsk_table_collection_t *tables = treeseq->tables;
     const tsk_id_t *restrict edge_parent = tables->edges.parent;
     const tsk_id_t *restrict edge_child = tables->edges.child;
-    const tsk_size_t num_edges = tables->edges.num_rows,
-                     num_trees = self->tree_sequence->num_trees;
+    const tsk_size_t num_edges = tables->edges.num_rows;
+    const tsk_size_t num_trees = self->tree_sequence->num_trees;
     const double *restrict edge_left = tables->edges.left;
     const double *restrict edge_right = tables->edges.right;
     const double *restrict breakpoints = treeseq->breakpoints;
@@ -4595,7 +4595,7 @@ tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(optio
     // of trees.
     j = -1;
     if (x <= L / 2.0) {
-        for (edge = 0; edge < num_edges; ++edge) {
+        for (edge = 0; edge < num_edges; edge++) {
             e = insertion[edge];
             if (edge_left[e] > x) {
                 j = (tsk_id_t) edge;
@@ -4608,12 +4608,12 @@ tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(optio
             }
         }
     } else {
-        for (edge = 0; edge < num_edges; ++edge) {
+        for (edge = 0; edge < num_edges; edge++) {
             e = removal[num_edges - edge - 1];
             if (edge_right[e] < x) {
                 j = (tsk_id_t)(num_edges - edge - 1);
                 while (j < (tsk_id_t) num_edges && edge_left[insertion[j]] <= x) {
-                    j += 1;
+                    j++;
                 }
                 break;
             }
@@ -4628,12 +4628,12 @@ tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(optio
     if (j == -1) {
         j = 0;
         while (j < (tsk_id_t) num_edges && edge_left[insertion[j]] <= x) {
-            j += 1;
+            j++;
         }
     }
     k = 0;
     while (k < (tsk_id_t) num_edges && edge_right[removal[k]] <= x) {
-        k += 1;
+        k--;
     }
 
     /* NOTE: tsk_search_sorted finds the first the first
@@ -4683,6 +4683,7 @@ tsk_tree_seek_linear(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(options)
     const double t_r = self->interval.right;
     int ret = 0;
     double distance_left, distance_right;
+
     if (x < t_l) {
         /* |-----|-----|========|---------| */
         /* 0     x    t_l      t_r        L */

python/tests/test_highlevel.py

@@ -4598,10 +4598,13 @@ def test_index_from_different_directions(self, index):
             t2.prev()
         assert_same_tree_different_order(t1, t2)
 
-    def test_seek_0_from_null(self):
+    @pytest.mark.parametrize("position", [0, 1, 2, 3])
+    def test_seek_from_null(self, position):
         t1, t2 = self.setup()
-        t1.first()
-        t2.seek(0)
+        t1.clear()
+        t1.seek(position)
+        t2.first()
+        t2.seek(position)
         assert_trees_identical(t1, t2)
 
     @pytest.mark.parametrize("index", range(3))
@@ -4654,6 +4657,14 @@ def test_seek_3_from_null(self):
         t2.seek(3)
         assert_trees_identical(t1, t2)
 
+    def test_seek_3_from_null_prev(self):
+        t1, t2 = self.setup()
+        t1.last()
+        t1.prev()
+        t2.seek(3)
+        t2.prev()
+        assert_trees_identical(t1, t2)
+
     def test_seek_3_from_0(self):
         t1, t2 = self.setup()
         t1.last()
@@ -4669,6 +4680,37 @@ def test_seek_0_from_3(self):
         t2.seek(0)
         assert_trees_identical(t1, t2)
 
+    @pytest.mark.parametrize("ts", get_example_tree_sequences())
+    def test_seek_mid_null_and_middle(self, ts):
+        breakpoints = ts.breakpoints(as_array=True)
+        mid = breakpoints[:-1] + np.diff(breakpoints) / 2
+        for index, x in enumerate(mid[:-1]):
+            t1 = tskit.Tree(ts)
+            t1.seek(x)
+            # Also seek to this point manually to make sure we're not
+            # reusing the seek from null under the hood.
+            t2 = tskit.Tree(ts)
+            if index <= ts.num_trees / 2:
+                while t2.index != index:
+                    t2.next()
+            else:
+                while t2.index != index:
+                    t2.prev()
+            assert t1.index == t2.index
+            assert np.all(t1.parent_array == t2.parent_array)
+
+    @pytest.mark.parametrize("ts", get_example_tree_sequences())
+    def test_seek_last_then_prev(self, ts):
+        t1 = tskit.Tree(ts)
+        t1.seek(ts.sequence_length - 0.00001)
+        assert t1.index == ts.num_trees - 1
+        t2 = tskit.Tree(ts)
+        t2.prev()
+        assert_trees_identical(t1, t2)
+        t1.prev()
+        t2.prev()
+        assert_trees_identical(t1, t2)
+
 
 class TestSeek:
     @pytest.mark.parametrize("ts", get_example_tree_sequences())