Restore seek_linear and add seek_skip as an option

Author: duncanMR

c/tskit/trees.c

@@ -6633,7 +6633,47 @@ tsk_tree_seek_index(tsk_tree_t *self, tsk_id_t tree, tsk_flags_t options)
 }
 
 static int TSK_WARN_UNUSED
-tsk_tree_seek_linear(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(options))
+tsk_tree_seek_linear(tsk_tree_t *self, double x)
+{
+    const double L = tsk_treeseq_get_sequence_length(self->tree_sequence);
+    const double t_l = self->interval.left;
+    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 */
+        distance_left = t_l - x;
+        distance_right = L - t_r + x;
+    } else {
+        /* |------|========|------|-------| */
+        /* 0     t_l      t_r     x       L */
+        distance_right = x - t_r;
+        distance_left = t_l + L - x;
+    }
+    if (distance_right <= distance_left) {
+        while (!tsk_tree_position_in_interval(self, x)) {
+            ret = tsk_tree_next(self);
+            if (ret < 0) {
+                goto out;
+            }
+        }
+    } else {
+        while (!tsk_tree_position_in_interval(self, x)) {
+            ret = tsk_tree_prev(self);
+            if (ret < 0) {
+                goto out;
+            }
+        }
+    }
+    ret = 0;
+out:
+    return ret;
+}
+
+static int TSK_WARN_UNUSED
+tsk_tree_seek_skip(tsk_tree_t *self, double x)
 {
     const double t_l = self->interval.left;
     int ret = 0;
@@ -6669,7 +6709,11 @@ tsk_tree_seek(tsk_tree_t *self, double x, tsk_flags_t options)
     if (self->index == -1) {
         ret = tsk_tree_seek_from_null(self, x, options);
     } else {
-        ret = tsk_tree_seek_linear(self, x, options);
+        if (options & TSK_TREE_SEEK_ENABLE_SKIPPING) {
+            ret = tsk_tree_seek_skip(self, x);
+        } else {
+            ret = tsk_tree_seek_linear(self, x);
+        }
     }
 
 out:

c/tskit/trees.h

@@ -1270,6 +1270,13 @@ int tsk_tree_copy(const tsk_tree_t *self, tsk_tree_t *dest, tsk_flags_t options)
 @{
 */
 
+/** @brief Option to seek by skipping to the target tree, adding and removing as few
+   edges as possible. If not specified, a linear time algorithm is used instead.
+
+    @ingroup TREE_API_SEEKING_GROUP
+*/
+#define TSK_TREE_SEEK_ENABLE_SKIPPING (1 << 0)
+
 /**
 @brief Seek to the first tree in the sequence.
 
@@ -1279,7 +1286,7 @@ tree sequence.
 @endrst
 
 @param self A pointer to an initialised tsk_tree_t object.
-@return Return TSK_TREE_OK on success; or a negative value if an error occurs.
+@return Return on success; or a negative value if an error occurs.
 */
 int tsk_tree_first(tsk_tree_t *self);
 
@@ -1375,12 +1382,22 @@ we will have ``position < tree.interval.right``.
 
 Seeking to a position currently covered by the tree is
 a constant time operation.
+
+Seeking to a position from a non-null tree use a linear time
+algorithm by default, unless the option :c:macro:`TSK_TREE_SEEK_ENABLE_SKIPPING`
+is specified. In this case, a faster algorithm is employed which skips
+to the target tree by removing and adding the minimal number of edges
+possible. However, this approach does not guarantee that edges are
+inserted and removed in time-sorted order.
+
+.. warning:: Using the :c:macro:`TSK_TREE_SEEK_ENABLE_SKIPPING` option
+    may lead to edges not being inserted or removed in time-sorted order.
+
 @endrst
 
 @param self A pointer to an initialised tsk_tree_t object.
 @param position The position in genome coordinates
-@param options Seek options. Currently unused. Set to 0 for compatibility
-    with future versions of tskit.
+@param options Seek options. See the notes above for details.
 @return Return 0 on success or a negative value on failure.
 */
 int tsk_tree_seek(tsk_tree_t *self, double position, tsk_flags_t options);

python/_tskitmodule.c

@@ -12103,16 +12103,22 @@ static PyObject *
 Tree_seek(Tree *self, PyObject *args)
 {
     PyObject *ret = NULL;
+    tsk_flags_t options = 0;
+    int enable_skipping = true;
     double position;
     int err;
 
+    if (enable_skipping) {
+        options |= TSK_TREE_SEEK_ENABLE_SKIPPING;
+    }
+
     if (Tree_check_state(self) != 0) {
         goto out;
     }
     if (!PyArg_ParseTuple(args, "d", &position)) {
         goto out;
     }
-    err = tsk_tree_seek(self->tree, position, 0);
+    err = tsk_tree_seek(self->tree, position, options);
     if (err != 0) {
         handle_library_error(err);
         goto out;