Remove tree_pos iterator from KC distance and haplotype matching

Author: duncanMR

c/tskit/haplotype_matching.c

@@ -209,7 +209,6 @@ int
 tsk_ls_hmm_free(tsk_ls_hmm_t *self)
 {
     tsk_tree_free(&self->tree);
-    tsk_tree_position_free(&self->tree_pos);
     tsk_safe_free(self->recombination_rate);
     tsk_safe_free(self->mutation_rate);
     tsk_safe_free(self->recombination_rate);
@@ -247,11 +246,6 @@ tsk_ls_hmm_reset(tsk_ls_hmm_t *self, double value)
     tsk_memset(self->transition_parent, 0xff,
         self->max_transitions * sizeof(*self->transition_parent));
 
-    tsk_tree_position_free(&self->tree_pos);
-    ret = tsk_tree_position_init(&self->tree_pos, self->tree_sequence, 0);
-    if (ret != 0) {
-        goto out;
-    }
     samples = tsk_treeseq_get_samples(self->tree_sequence);
     for (j = 0; j < self->num_samples; j++) {
         u = samples[j];
@@ -260,7 +254,6 @@ tsk_ls_hmm_reset(tsk_ls_hmm_t *self, double value)
         self->transition_index[u] = (tsk_id_t) j;
     }
     self->num_transitions = self->num_samples;
-out:
     return ret;
 }
 
@@ -309,13 +302,11 @@ tsk_ls_hmm_update_tree(tsk_ls_hmm_t *self, int direction)
     tsk_value_transition_t *restrict T = self->transitions;
     tsk_id_t u, c, p, j, e;
     tsk_value_transition_t *vt;
+    tsk_tree_position_t tree_pos;
 
-    tsk_tree_position_step(&self->tree_pos, direction);
-    tsk_bug_assert(self->tree_pos.index != -1);
-    tsk_bug_assert(self->tree_pos.index == self->tree.index);
-
-    for (j = self->tree_pos.out.start; j != self->tree_pos.out.stop; j += direction) {
-        e = self->tree_pos.out.order[j];
+    tree_pos = self->tree.tree_pos;
+    for (j = tree_pos.out.start; j != tree_pos.out.stop; j += direction) {
+        e = tree_pos.out.order[j];
         c = edges_child[e];
         u = c;
         if (T_index[u] == TSK_NULL) {
@@ -333,8 +324,8 @@ tsk_ls_hmm_update_tree(tsk_ls_hmm_t *self, int direction)
         parent[c] = TSK_NULL;
     }
 
-    for (j = self->tree_pos.in.start; j != self->tree_pos.in.stop; j += direction) {
-        e = self->tree_pos.in.order[j];
+    for (j = tree_pos.in.start; j != tree_pos.in.stop; j += direction) {
+        e = tree_pos.in.order[j];
         c = edges_child[e];
         p = edges_parent[e];
         parent[c] = p;

c/tskit/haplotype_matching.h

@@ -1,7 +1,7 @@
 /*
  * MIT License
  *
- * Copyright (c) 2019-2023 Tskit Developers
+ * Copyright (c) 2019-2024 Tskit Developers
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -98,10 +98,6 @@ typedef struct _tsk_ls_hmm_t {
     tsk_size_t num_nodes;
     /* state */
     tsk_tree_t tree;
-    /* NOTE: this tree_position will be redundant once we integrate the top-level
-     * tree class with this.
-     */
-    tsk_tree_position_t tree_pos;
     tsk_id_t *parent;
     /* The probability value transitions on the tree */
     tsk_value_transition_t *transitions;

c/tskit/trees.c

@@ -4554,19 +4554,6 @@ tsk_tree_position_prev(tsk_tree_position_t *self)
     return self->index != -1;
 }
 
-bool
-tsk_tree_position_step(tsk_tree_position_t *self, int direction)
-{
-    bool ret = false;
-
-    if (direction == TSK_DIR_FORWARD) {
-        ret = tsk_tree_position_next(self);
-    } else if (direction == TSK_DIR_REVERSE) {
-        ret = tsk_tree_position_prev(self);
-    }
-    return ret;
-}
-
 int TSK_WARN_UNUSED
 tsk_tree_position_seek_forward(tsk_tree_position_t *self, tsk_id_t index)
 {
@@ -5674,18 +5661,20 @@ tsk_tree_next(tsk_tree_t *self)
     const tsk_id_t *restrict edge_parent = tables->edges.parent;
     const tsk_id_t *restrict edge_child = tables->edges.child;
     tsk_id_t j, e;
+    tsk_tree_position_t tree_pos;
     bool valid;
 
     valid = tsk_tree_position_next(&self->tree_pos);
+    tree_pos = self->tree_pos;
 
     if (valid) {
-        for (j = self->tree_pos.out.start; j != self->tree_pos.out.stop; j++) {
-            e = self->tree_pos.out.order[j];
+        for (j = tree_pos.out.start; j != tree_pos.out.stop; j++) {
+            e = tree_pos.out.order[j];
             tsk_tree_remove_edge(self, edge_parent[e], edge_child[e]);
         }
 
-        for (j = self->tree_pos.in.start; j != self->tree_pos.in.stop; j++) {
-            e = self->tree_pos.in.order[j];
+        for (j = tree_pos.in.start; j != tree_pos.in.stop; j++) {
+            e = tree_pos.in.order[j];
             tsk_tree_insert_edge(self, edge_parent[e], edge_child[e], e);
         }
         ret = TSK_TREE_OK;
@@ -5704,18 +5693,20 @@ tsk_tree_prev(tsk_tree_t *self)
     const tsk_id_t *restrict edge_parent = tables->edges.parent;
     const tsk_id_t *restrict edge_child = tables->edges.child;
     tsk_id_t j, e;
+    tsk_tree_position_t tree_pos;
     bool valid;
 
     valid = tsk_tree_position_prev(&self->tree_pos);
+    tree_pos = self->tree_pos;
 
     if (valid) {
-        for (j = self->tree_pos.out.start; j != self->tree_pos.out.stop; j--) {
-            e = self->tree_pos.out.order[j];
+        for (j = tree_pos.out.start; j != tree_pos.out.stop; j--) {
+            e = tree_pos.out.order[j];
             tsk_tree_remove_edge(self, edge_parent[e], edge_child[e]);
         }
 
-        for (j = self->tree_pos.in.start; j != self->tree_pos.in.stop; j--) {
-            e = self->tree_pos.in.order[j];
+        for (j = tree_pos.in.start; j != tree_pos.in.stop; j--) {
+            e = tree_pos.in.order[j];
             tsk_tree_insert_edge(self, edge_parent[e], edge_child[e], e);
         }
         ret = TSK_TREE_OK;
@@ -5741,12 +5732,12 @@ tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(optio
     const tsk_id_t *restrict edge_child = tables->edges.child;
     const double *restrict edge_left = tables->edges.left;
     const double *restrict edge_right = tables->edges.right;
-    double interval_left = self->interval.left;
-    double interval_right = self->interval.right;
+    double interval_left, interval_right;
     const double *restrict breakpoints = self->tree_sequence->breakpoints;
     const tsk_size_t num_trees = self->tree_sequence->num_trees;
     const double L = tsk_treeseq_get_sequence_length(self->tree_sequence);
     tsk_id_t j, e, index;
+    tsk_tree_position_t tree_pos;
 
     index = (tsk_id_t) tsk_search_sorted(breakpoints, num_trees + 1, x);
     if (breakpoints[index] > x) {
@@ -5758,9 +5749,11 @@ tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(optio
         if (ret != 0) {
             goto out;
         }
-        interval_left = self->tree_pos.interval.left;
-        for (j = self->tree_pos.in.start; j != self->tree_pos.in.stop; j++) {
-            e = self->tree_pos.in.order[j];
+        // Since we are seeking from null, there are no edges to remove
+        tree_pos = self->tree_pos;
+        interval_left = tree_pos.interval.left;
+        for (j = tree_pos.in.start; j != tree_pos.in.stop; j++) {
+            e = tree_pos.in.order[j];
             if (edge_left[e] <= interval_left && interval_left < edge_right[e]) {
                 tsk_tree_insert_edge(self, edge_parent[e], edge_child[e], e);
             }
@@ -5770,9 +5763,10 @@ tsk_tree_seek_from_null(tsk_tree_t *self, double x, tsk_flags_t TSK_UNUSED(optio
         if (ret != 0) {
             goto out;
         }
-        interval_right = self->tree_pos.interval.right;
-        for (j = self->tree_pos.in.start; j != self->tree_pos.in.stop; j--) {
-            e = self->tree_pos.in.order[j];
+        tree_pos = self->tree_pos;
+        interval_right = tree_pos.interval.right;
+        for (j = tree_pos.in.start; j != tree_pos.in.stop; j--) {
+            e = tree_pos.in.order[j];
             if (edge_right[e] >= interval_right && interval_right > edge_left[e]) {
                 tsk_tree_insert_edge(self, edge_parent[e], edge_child[e], e);
             }
@@ -6962,23 +6956,19 @@ update_kc_subtree_state(
 }
 
 static int
-update_kc_incremental(
-    tsk_tree_t *tree, kc_vectors *kc, tsk_tree_position_t *tree_pos, tsk_size_t *depths)
+update_kc_incremental(tsk_tree_t *tree, kc_vectors *kc, tsk_size_t *depths)
 {
     int ret = 0;
     tsk_id_t u, v, e, j;
     double root_time, time;
     const double *restrict times = tree->tree_sequence->tables->nodes.time;
     const tsk_id_t *restrict edges_child = tree->tree_sequence->tables->edges.child;
     const tsk_id_t *restrict edges_parent = tree->tree_sequence->tables->edges.parent;
-
-    tsk_bug_assert(tree_pos->index == tree->index);
-    tsk_bug_assert(tree_pos->interval.left == tree->interval.left);
-    tsk_bug_assert(tree_pos->interval.right == tree->interval.right);
+    tsk_tree_position_t tree_pos = tree->tree_pos;
 
     /* Update state of detached subtrees */
-    for (j = tree_pos->out.stop - 1; j >= tree_pos->out.start; j--) {
-        e = tree_pos->out.order[j];
+    for (j = tree_pos.out.stop - 1; j >= tree_pos.out.start; j--) {
+        e = tree_pos.out.order[j];
         u = edges_child[e];
         depths[u] = 0;
 
@@ -6992,8 +6982,8 @@ update_kc_incremental(
     }
 
     /* Propagate state change down into reattached subtrees. */
-    for (j = tree_pos->in.stop - 1; j >= tree_pos->in.start; j--) {
-        e = tree_pos->in.order[j];
+    for (j = tree_pos.in.stop - 1; j >= tree_pos.in.start; j--) {
+        e = tree_pos.in.order[j];
         u = edges_child[e];
         v = edges_parent[e];
 
@@ -7026,17 +7016,11 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
     const tsk_treeseq_t *treeseqs[2] = { self, other };
     tsk_tree_t trees[2];
     kc_vectors kcs[2];
-    /* TODO the tree_pos here is redundant because we should be using this interally
-     * in the trees to do the advancing. Once we have converted the tree over to using
-     * tree_pos internally, we can get rid of these tree_pos variables and use
-     * the values stored in the trees themselves */
-    tsk_tree_position_t tree_pos[2];
     tsk_size_t *depths[2];
     int ret = 0;
 
     for (i = 0; i < 2; i++) {
         tsk_memset(&trees[i], 0, sizeof(trees[i]));
-        tsk_memset(&tree_pos[i], 0, sizeof(tree_pos[i]));
         tsk_memset(&kcs[i], 0, sizeof(kcs[i]));
         depths[i] = NULL;
     }
@@ -7052,10 +7036,6 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
         if (ret != 0) {
             goto out;
         }
-        ret = tsk_tree_position_init(&tree_pos[i], treeseqs[i], 0);
-        if (ret != 0) {
-            goto out;
-        }
         ret = kc_vectors_alloc(&kcs[i], n);
         if (ret != 0) {
             goto out;
@@ -7079,10 +7059,8 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
     if (ret != 0) {
         goto out;
     }
-    tsk_tree_position_next(&tree_pos[0]);
-    tsk_bug_assert(tree_pos[0].index == 0);
 
-    ret = update_kc_incremental(&trees[0], &kcs[0], &tree_pos[0], depths[0]);
+    ret = update_kc_incremental(&trees[0], &kcs[0], depths[0]);
     if (ret != 0) {
         goto out;
     }
@@ -7091,17 +7069,11 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
         if (ret != 0) {
             goto out;
         }
-        tsk_tree_position_next(&tree_pos[1]);
-        tsk_bug_assert(tree_pos[1].index != -1);
 
-        ret = update_kc_incremental(&trees[1], &kcs[1], &tree_pos[1], depths[1]);
+        ret = update_kc_incremental(&trees[1], &kcs[1], depths[1]);
         if (ret != 0) {
             goto out;
         }
-        tsk_bug_assert(trees[0].interval.left == tree_pos[0].interval.left);
-        tsk_bug_assert(trees[0].interval.right == tree_pos[0].interval.right);
-        tsk_bug_assert(trees[1].interval.left == tree_pos[1].interval.left);
-        tsk_bug_assert(trees[1].interval.right == tree_pos[1].interval.right);
         while (trees[0].interval.right < trees[1].interval.right) {
             span = trees[0].interval.right - left;
             total += norm_kc_vectors(&kcs[0], &kcs[1], lambda_) * span;
@@ -7113,9 +7085,7 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
             if (ret != 0) {
                 goto out;
             }
-            tsk_tree_position_next(&tree_pos[0]);
-            tsk_bug_assert(tree_pos[0].index != -1);
-            ret = update_kc_incremental(&trees[0], &kcs[0], &tree_pos[0], depths[0]);
+            ret = update_kc_incremental(&trees[0], &kcs[0], depths[0]);
             if (ret != 0) {
                 goto out;
             }
@@ -7132,7 +7102,6 @@ tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
 out:
     for (i = 0; i < 2; i++) {
         tsk_tree_free(&trees[i]);
-        tsk_tree_position_free(&tree_pos[i]);
         kc_vectors_free(&kcs[i]);
         tsk_safe_free(depths[i]);
     }

c/tskit/trees.h

@@ -1850,10 +1850,8 @@ int tsk_tree_position_free(tsk_tree_position_t *self);
 int tsk_tree_position_print_state(const tsk_tree_position_t *self, FILE *out);
 bool tsk_tree_position_next(tsk_tree_position_t *self);
 bool tsk_tree_position_prev(tsk_tree_position_t *self);
-bool tsk_tree_position_step(tsk_tree_position_t *self, int direction);
 int tsk_tree_position_seek_forward(tsk_tree_position_t *self, tsk_id_t index);
 int tsk_tree_position_seek_backward(tsk_tree_position_t *self, tsk_id_t index);
-int tsk_tree_position_seek(tsk_tree_position_t *self, double position);
 
 #ifdef __cplusplus
 }