Implement efficient initialization of null trees.

Author: molpopgen

c/tskit/trees.c

@@ -4551,8 +4551,103 @@ tsk_tree_position_in_interval(const tsk_tree_t *self, double x)
     return self->interval.left <= x && x < self->interval.right;
 }
 
+static int
+tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(options))
+{
+    int ret = 0;
+    tsk_size_t edge, num_edges = self->tree_sequence->tables->edges.num_rows,
+                     num_trees = self->tree_sequence->num_trees;
+    tsk_id_t p, c, e, j, k, tree_index;
+    const double L = tsk_treeseq_get_sequence_length(self->tree_sequence);
+    const tsk_id_t *restrict edge_parent = self->tree_sequence->tables->edges.parent;
+    const tsk_id_t *restrict edge_child = self->tree_sequence->tables->edges.child;
+    const double *restrict edge_left = self->tree_sequence->tables->edges.left;
+    const double *restrict edge_right = self->tree_sequence->tables->edges.right;
+    const double *restrict breakpoints = self->tree_sequence->breakpoints;
+    const tsk_id_t *restrict insertion
+        = self->tree_sequence->tables->indexes.edge_insertion_order;
+    const tsk_id_t *restrict removal
+        = self->tree_sequence->tables->indexes.edge_removal_order;
+
+    /* NOTE: this differ's from JK's proposal by not using
+     * edge diffs explicitly. However:
+     * 1. It benchmarked faster (in rust).
+     * 2. It is a lot less code and doesn't have
+     * to do "edge diffs, but backwards".
+     */
+
+    // NOTE: it may be better to get the
+    // index first and then ask if we are
+    // searching in the first or last 1/2
+    // of trees.
+    j = -1;
+    if (x <= L / 2.0) {
+        for (edge = 0; edge < num_edges; ++edge) {
+            e = insertion[edge];
+            if (edge_left[e] > x) {
+                j = (tsk_id_t) edge;
+                break;
+            }
+            if (x >= edge_left[e] && x < edge_right[e]) {
+                p = edge_parent[e];
+                c = edge_child[e];
+                tsk_tree_insert_edge(self, p, c, e);
+            }
+        }
+    } else {
+        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;
+                }
+                break;
+            }
+            if (x >= edge_left[e] && x < edge_right[e]) {
+                p = edge_parent[e];
+                c = edge_child[e];
+                tsk_tree_insert_edge(self, p, c, e);
+            }
+        }
+    }
+
+    if (j == -1) {
+        j = 0;
+        while (j < (tsk_id_t) num_edges && edge_left[insertion[j]] <= x) {
+            j += 1;
+        }
+    }
+    k = 0;
+    while (k < (tsk_id_t) num_edges && edge_right[removal[k]] <= x) {
+        k += 1;
+    }
+
+    /* NOTE: tsk_search_sorted finds the first the first
+     * insertion locatiom >= the query point, which
+     * finds a RIGHT value for queries not at the left edge.
+     */
+    tree_index = (tsk_id_t) tsk_search_sorted(breakpoints, num_trees + 1, x);
+    if (breakpoints[tree_index] > x) {
+        tree_index -= 1;
+    }
+    self->index = tree_index;
+    self->interval.left = breakpoints[tree_index];
+    self->interval.right = breakpoints[tree_index + 1];
+    self->left_index = j;
+    self->right_index = k;
+    self->direction = TSK_DIR_FORWARD;
+    self->num_nodes = self->tree_sequence->tables->nodes.num_rows;
+    if (self->tree_sequence->tables->sites.num_rows > 0) {
+        self->sites = self->tree_sequence->tree_sites[self->index];
+        self->sites_length = self->tree_sequence->tree_sites_length[self->index];
+    }
+
+    return ret;
+}
+
 int TSK_WARN_UNUSED
-tsk_tree_seek(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(options))
+tsk_tree_seek(tsk_tree_t *self, double x, tsk_flags_t options)
 {
     int ret = 0;
     const double L = tsk_treeseq_get_sequence_length(self->tree_sequence);
@@ -4565,6 +4660,13 @@ tsk_tree_seek(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(options))
         goto out;
     }
 
+    if (self->index == -1) {
+        ret = tsk_tree_seek_from_null(self, x, options);
+        if (ret < 0) {
+            goto out;
+        }
+    }
+
     if (x < t_l) {
         /* |-----|-----|========|---------| */
         /* 0     x    t_l      t_r        L */
@@ -4592,6 +4694,7 @@ tsk_tree_seek(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(options))
         }
     }
     ret = 0;
+
 out:
     return ret;
 }