From d5d032170b98e037c55084c1eafe29a66e893878 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Wed, 28 Oct 2026 17:21:53 +0000
Subject: [PATCH 1/7] Clear likelihood caches with a stamp instead of reset
 walks

Profile (perf record, line-level annotation of tsls_hmm_find_path_matrix,
into which the site update, coalescing and walks are inlined) of chr20p
sample matching, as shares of that function's samples:

                                     map (2 queries)  constant r (3)
  allele walk (205-211) + clear (225)       ~30%          ~27%
  coalesce walk to inherited value (1343)   ~16%          ~12%
  coalesce cache writes (1354)               ~8%           ~10%
  cache reset walks (1021, 1412)            ~14%          ~15%
  traceback run (2269)                       ~2%            ~5%
  arithmetic (1615)                          <1%           <1%

tsls_hmm_update_tree is a further ~10% of the run under the map and ~25%
under constant r. The time is in walks up the tree, not in the values.
(An all-missing query is not a useful floor for the topology cost: it
takes 15 s, because nodes entering the tree at 0 never coalesce.)

The caches are now valid on nodes whose cache_stamp (in the node record,
which grows from 8 to 16 bytes so records don't straddle cache lines)
equals the HMM's, and clearing them all is an increment, so the reset
walks after edge removals, coalescing and sparse sites are gone. A node
with a likelihood ignores any stale cache.

Paired wall time (ab.py: alternating processes pinned to one core, best
of 3 per query, 6 samples), stamp / base: map 0.908, 0.935, 0.939;
constant r 0.957, 0.975, 0.980. Paths identical.

benchmarking/profile_samples.py runs and times tsls on chr20p samples
under a map or constant r; benchmarking/ab.py compares builds with it.
---
 benchmarking/ab.py              | 68 ++++++++++++++++++++++++++++
 benchmarking/profile_samples.py | 79 +++++++++++++++++++++++++++++++++
 tsls/lib/hmm.c                  | 79 ++++++++++++++-------------------
 tsls/lib/tsls.h                 | 20 ++++++---
 4 files changed, 194 insertions(+), 52 deletions(-)
 create mode 100644 benchmarking/ab.py
 create mode 100644 benchmarking/profile_samples.py

diff --git a/benchmarking/ab.py b/benchmarking/ab.py
new file mode 100644
index 0000000..5d405b4
--- /dev/null
+++ b/benchmarking/ab.py
@@ -0,0 +1,68 @@
+"""
+Compare the wall time of two or more builds of the tsls extension on chr20p
+sample matching, with paired runs.
+
+Each build is a directory holding a _tsls extension module, put first on
+PYTHONPATH ("." means the in-tree build). For each round, every build runs
+profile_samples.py in its own process, pinned to one core, in turn; the
+best time for each (build, query) over the rounds is kept. Reports the
+total time of each build over the first, and checks that every build finds
+the same paths.
+
+    uv run python ab.py base=/path/to/build new=. --map data/genetic_maps/... -r 3
+"""
+
+import collections
+import json
+import os
+import subprocess
+import sys
+
+import click
+
+
+def run_build(path, args, core):
+    env = dict(os.environ)
+    if path != ".":
+        env["PYTHONPATH"] = path
+    command = ["taskset", "-c", str(core), sys.executable, "profile_samples.py", *args]
+    out = subprocess.run(command, env=env, capture_output=True, text=True, check=True)
+    return [json.loads(line) for line in out.stdout.splitlines() if line.startswith("{")]
+
+
+@click.command()
+@click.argument("builds", nargs=-1, required=True)
+@click.option("--map", "genetic_map", default=None)
+@click.option("--mu", type=float, default=1e-20)
+@click.option("--queries", "query_ids", default="0,40,80,120,160,199")
+@click.option("-r", "--rounds", type=int, default=3)
+@click.option("--core", type=int, default=2)
+def main(builds, genetic_map, mu, query_ids, rounds, core):
+    named = [b.split("=", 1) for b in builds]
+    args = ["--mu", str(mu), "--queries", query_ids]
+    if genetic_map is not None:
+        args += ["--map", genetic_map]
+    best = collections.defaultdict(dict)
+    paths = collections.defaultdict(dict)
+    for round_ in range(rounds):
+        for name, path in named:
+            for row in run_build(path, args, core):
+                q = row["query"]
+                best[name][q] = min(best[name].get(q, float("inf")), row["time"])
+                paths[name][q] = row["path"]
+        totals = {name: sum(best[name].values()) for name, _ in named}
+        base = named[0][0]
+        click.echo(
+            f"round {round_ + 1}: "
+            + ", ".join(
+                f"{name} {totals[name]:.2f}s ({totals[name] / totals[base]:.3f})"
+                for name, _ in named
+            )
+        )
+    for name, _ in named[1:]:
+        same = paths[name] == paths[named[0][0]]
+        click.echo(f"{name}: paths {'identical' if same else 'DIFFER'}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarking/profile_samples.py b/benchmarking/profile_samples.py
new file mode 100644
index 0000000..406446d
--- /dev/null
+++ b/benchmarking/profile_samples.py
@@ -0,0 +1,79 @@
+"""
+Run tsls Viterbi on chr20p sample queries, for profiling and step counts.
+
+The reference is the samples workload's reference.trees with the vestigial
+root edges from node 0 removed (as engines.TslsEngine does), and the model
+is tsinfer's: r = 0.01 or per-site probabilities from a genetic map, with
+weight_by_n=True (tsls's model). Prints the time and the engine's
+statistics for each query, and the step counters if the extension was
+built with TSLS_STEP_COUNTERS.
+
+    uv run python profile_samples.py --map data/genetic_maps/chr20_HapMapII_GRCh38.txt
+"""
+
+import hashlib
+import json
+import time
+
+import click
+import numpy as np
+import tskit
+import tsls
+
+import queries
+
+WORKLOAD = "data/tgp_chr20p/samples"
+
+
+def load_reference(path):
+    ts = tskit.load(path)
+    tables = ts.dump_tables()
+    tables.edges.keep_rows(tables.edges.parent != 0)
+    return tables.tree_sequence()
+
+
+@click.command()
+@click.option("--map", "genetic_map", default=None, help="HapMap genetic map")
+@click.option("--mu", type=float, default=1e-20)
+@click.option("--queries", "query_ids", default="0,40,80,120,160,199")
+@click.option("--missing", is_flag=True, help="Use all-missing haplotypes")
+@click.option("--repeats", type=int, default=1)
+def main(genetic_map, mu, query_ids, missing, repeats):
+    ts = load_reference(f"{WORKLOAD}/reference.trees")
+    H, positions, _ = queries.load_queries(WORKLOAD)
+    m = ts.num_sites
+    if genetic_map is None:
+        r = np.full(m, 0.01)
+    else:
+        r = tsls.model.recombination_probabilities(genetic_map, positions)
+    reference = tsls.TreeSequence(ts)
+    for qid in [int(q) for q in query_ids.split(",")]:
+        h = H[qid].astype(np.int8)
+        start, end = queries.query_window(H, qid)
+        if missing:
+            h[start:end] = -1
+        for _ in range(repeats):
+            hmm = tsls.Hmm(reference, r, np.full(m, mu))
+            matched = np.zeros(m, dtype=np.int8)
+            before = time.perf_counter()
+            left, right, parent = hmm.find_path(h, start, end, matched)
+            seconds = time.perf_counter() - before
+            digest = hashlib.md5()
+            for array in [left, right, parent, matched]:
+                digest.update(np.ascontiguousarray(array).tobytes())
+            row = {
+                "query": qid,
+                "time": round(seconds, 4),
+                "segments": len(left),
+                "path": digest.hexdigest()[:12],
+                "mean_values": round(hmm.mean_num_values, 1),
+                "traceback": round(hmm.mean_traceback_size, 1),
+            }
+            counts = getattr(hmm, "step_counts", None)
+            if counts is not None:
+                row.update(counts)
+            print(json.dumps(row), flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c
index 1113bc7..131971b 100644
--- a/tsls/lib/hmm.c
+++ b/tsls/lib/hmm.c
@@ -616,7 +616,7 @@ tsls_hmm_check_state(const tsls_hmm_t *self)
         }
         assert(self->tree[u].allelic_state
                == (self->tree[u].parent == NULL_NODE ? ROOT_ALLELE : TSK_NULL));
-        assert(self->likelihood[u].cache == CACHE_UNSET);
+        assert(self->tree[u].cache_stamp != self->cache_stamp);
     }
     assert(num_likelihoods == self->num_likelihood_nodes);
     (void) num_likelihoods;
@@ -990,37 +990,39 @@ tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u)
 {
     tsls_likelihood_t *restrict L = self->likelihood;
     tsls_node_t *restrict tree = self->tree;
+    const uint32_t stamp = self->cache_stamp;
     double value;
     tsk_id_t v;
 
     v = u;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) {
+    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
         v = tree[v].parent;
     }
-    value = L[v].cache;
-    if (unlikely(value == CACHE_UNSET)) {
-        value = L[v].value;
-    }
+    /* A node with a likelihood may have a stale cache from before it was
+     * given one, so its likelihood takes precedence */
+    value = L[v].value >= 0 ? L[v].value : L[v].cache;
     v = u;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) {
+    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
         L[v].cache = value;
+        tree[v].cache_stamp = stamp;
         v = tree[v].parent;
     }
     return value;
 }
 
-/* Unset the cached values on the path up from u. Values are only cached on
- * nodes without likelihoods, and no root has a cached value (see
- * tsls_hmm_remove_edge), so the walk stops below a root. */
+/* Clear all the likelihood caches, which are only valid while the tree and
+ * the set of nodes with likelihoods are unchanged. */
 static inline void
-tsls_hmm_reset_likelihood_cache(tsls_hmm_t *self, tsk_id_t u)
+tsls_hmm_clear_likelihood_caches(tsls_hmm_t *self)
 {
-    tsls_likelihood_t *restrict L = self->likelihood;
-    tsls_node_t *restrict tree = self->tree;
+    size_t j;
 
-    while (L[u].cache != CACHE_UNSET) {
-        L[u].cache = CACHE_UNSET;
-        u = tree[u].parent;
+    self->cache_stamp++;
+    if (unlikely(self->cache_stamp == 0)) {
+        for (j = 0; j < self->num_nodes; j++) {
+            self->tree[j].cache_stamp = 0;
+        }
+        self->cache_stamp = 1;
     }
 }
 
@@ -1043,13 +1045,9 @@ tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c,
     }
     remove_edge(p, c, self->tree, self->left_child, self->right_child,
         self->left_sib, self->right_sib);
+    /* c is now a root, and so has a likelihood, which walks over the
+     * cache use in preference to any cached value */
     self->tree[c].allelic_state = ROOT_ALLELE;
-    /* c is now a root, and so has a likelihood. Its cached value (if any,
-     * from an earlier removal at this position) is the same as its
-     * likelihood, and clearing it means that no root has a cached value,
-     * so walks over the cache always stop below a root (see
-     * tsls_hmm_reset_likelihood_cache). */
-    self->likelihood[c].cache = CACHE_UNSET;
     if (update_likelihoods) {
         /* Each node can only become isolated once when removing edges, so
          * the isolated list has at most num_nodes entries. */
@@ -1108,9 +1106,7 @@ tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp)
         tsls_hmm_remove_edge(
             self, out[j].parent, out[j].child, true, isolated, &num_isolated);
     }
-    for (j = tp->out.start; j != tp->out.stop; j += direction) {
-        tsls_hmm_reset_likelihood_cache(self, out[j].parent);
-    }
+    tsls_hmm_clear_likelihood_caches(self);
     for (j = tp->in.start; j != tp->in.stop; j += direction) {
         tsls_hmm_insert_edge(self, in[j].parent, in[j].child, true);
     }
@@ -1151,7 +1147,7 @@ tsls_hmm_init_sweep(
     memset(self->right_sib, 0xff, N * sizeof(*self->right_sib));
     for (k = 0; k < N; k++) {
         /* Every node starts as a root. The likelihoods are set below. */
-        self->likelihood[k].cache = CACHE_UNSET;
+        self->tree[k].cache_stamp = 0;
         self->tree[k].parent = NULL_NODE;
         self->tree[k].allelic_state = ROOT_ALLELE;
         self->tree[k].recombination_required = -1;
@@ -1159,6 +1155,7 @@ tsls_hmm_init_sweep(
         self->tree[k].traceback_flag = -1;
         self->num_below[k] = 1;
     }
+    self->cache_stamp = 1;
     self->tree_numbered = false;
     self->all_changed = true;
     self->num_changed = 0;
@@ -1335,15 +1332,16 @@ tsls_hmm_get_inherited_likelihood_and_flag(
 {
     tsls_likelihood_t *restrict L = self->likelihood;
     tsls_node_t *restrict tree = self->tree;
+    const uint32_t stamp = self->cache_stamp;
     double value;
     int8_t r;
     tsk_id_t v;
 
     v = u;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) {
+    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
         v = tree[v].parent;
     }
-    if (L[v].cache == CACHE_UNSET) {
+    if (L[v].value >= 0) {
         value = L[v].value;
         r = tree[v].recombination_required;
     } else {
@@ -1351,9 +1349,10 @@ tsls_hmm_get_inherited_likelihood_and_flag(
         r = tree[v].recombination_required_cache;
     }
     v = u;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(L[v].cache == CACHE_UNSET)) {
+    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
         L[v].cache = value;
         tree[v].recombination_required_cache = r;
+        tree[v].cache_stamp = stamp;
         v = tree[v].parent;
     }
     *flag = r;
@@ -1374,19 +1373,15 @@ tsls_hmm_coalesce_likelihoods(
     double L_p;
     int8_t R_p;
     tsk_id_t u, p;
-    tsk_id_t *restrict cached_paths = self->likelihood_nodes_tmp;
     const int old_num_likelihood_nodes = self->num_likelihood_nodes;
     tsk_id_t *restrict L_nodes = self->likelihood_nodes;
-    int j, num_cached_paths, num_likelihood_nodes;
+    int j, num_likelihood_nodes;
 
-    num_cached_paths = 0;
     num_likelihood_nodes = 0;
     for (j = 0; j < old_num_likelihood_nodes; j++) {
         u = L_nodes[j];
         p = tree[u].parent;
         if (p != NULL_NODE) {
-            cached_paths[num_cached_paths] = p;
-            num_cached_paths++;
             if (check_recombination_required) {
                 /* A removed node had the inherited value and flag, so the
                  * cached values below it stay valid. */
@@ -1409,9 +1404,7 @@ tsls_hmm_coalesce_likelihoods(
         }
     }
     self->num_likelihood_nodes = num_likelihood_nodes;
-    for (j = 0; j < num_cached_paths; j++) {
-        tsls_hmm_reset_likelihood_cache(self, cached_paths[j]);
-    }
+    tsls_hmm_clear_likelihood_caches(self);
 }
 
 static inline double
@@ -1766,7 +1759,6 @@ tsls_hmm_process_site_sparse(tsls_hmm_t *self, const tsk_id_t site, const allele
     const int32_t *restrict subtree_end = self->subtree_end;
     const tsk_id_t *restrict L_nodes = self->likelihood_nodes;
     tsk_id_t *restrict nodes = self->site_nodes;
-    tsk_id_t *restrict cached_paths = self->likelihood_nodes_tmp;
     tsk_id_t *restrict changed = self->changed_next;
     const double rho = self->recombination_rate[site];
     const double mu = self->mismatch_rate[site];
@@ -1784,7 +1776,7 @@ tsls_hmm_process_site_sparse(tsls_hmm_t *self, const tsk_id_t site, const allele
     int8_t old_flag, R_p;
     allele_t allele;
     bool mismatch;
-    int j, k, num_ranges, num_bounds, num_nodes, num_cached, num_changed;
+    int j, k, num_ranges, num_bounds, num_nodes, num_changed;
     int32_t x;
     tsk_size_t i;
     tsk_id_t u, p;
@@ -1833,7 +1825,7 @@ tsls_hmm_process_site_sparse(tsls_hmm_t *self, const tsk_id_t site, const allele
                         ranges, num_ranges, preorder[p], ancestral_state)
                     != state)) {
                 value = tsls_hmm_get_inherited_likelihood_and_flag(self, p, &R_p);
-                tsls_hmm_reset_likelihood_cache(self, p);
+                tsls_hmm_clear_likelihood_caches(self);
                 tsls_hmm_add_likelihood_sorted(self, u, value);
                 tree[u].recombination_required = R_p;
                 tsls_hmm_add_site_node(self, u, &num_nodes);
@@ -1904,22 +1896,17 @@ tsls_hmm_process_site_sparse(tsls_hmm_t *self, const tsk_id_t site, const allele
     self->all_changed = false;
 
     /* Coalesce the recomputed nodes, as tsls_hmm_coalesce_likelihoods does */
-    num_cached = 0;
     for (j = 0; j < num_nodes; j++) {
         u = nodes[j];
         p = tree[u].parent;
         if (p != NULL_NODE) {
-            cached_paths[num_cached] = p;
-            num_cached++;
             L_p = tsls_hmm_get_inherited_likelihood_and_flag(self, p, &R_p);
             if (L[u].value == L_p && tree[u].recombination_required == R_p) {
                 tsls_hmm_remove_likelihood_sorted(self, u);
             }
         }
     }
-    for (j = 0; j < num_cached; j++) {
-        tsls_hmm_reset_likelihood_cache(self, cached_paths[j]);
-    }
+    tsls_hmm_clear_likelihood_caches(self);
     ret = tsls_hmm_store_traceback(self, site, nodes, num_nodes, NULL, 0);
     if (ret != 0) {
         goto out;
diff --git a/tsls/lib/tsls.h b/tsls/lib/tsls.h
index a84305a..16c6741 100644
--- a/tsls/lib/tsls.h
+++ b/tsls/lib/tsls.h
@@ -30,8 +30,7 @@
 #define NULL_LIKELIHOOD (-1)
 #define NOT_IN_TREE     (-2)
 
-#define NULL_NODE   (-1)
-#define CACHE_UNSET (-1)
+#define NULL_NODE (-1)
 
 /* Matcher flags */
 #define TSLS_EXTENDED_CHECKS     (1 << 0)
@@ -130,7 +129,8 @@ typedef struct {
 } tsls_index_t;
 
 /* The state of a node that the walks up the tree read, kept together (in
- * 8 bytes) so that each step of a walk reads one record. */
+ * 16 bytes, so that records don't straddle cache lines) so that each step
+ * of a walk reads one record. */
 typedef struct {
     tsk_id_t parent;
     /* The allelic state at the current site, if the node has a mutation
@@ -138,18 +138,23 @@ typedef struct {
     allele_t allelic_state;
     int8_t recombination_required;
     /* The recombination_required flag inherited by the node, valid where
-     * the likelihood cache is set. Only filled when coalescing Viterbi. */
+     * the likelihood cache is. Only filled when coalescing Viterbi. */
     int8_t recombination_required_cache;
     /* The node's flag in the traceback as of the last site stored, or -1
      * if it is not in the traceback's set */
     int8_t traceback_flag;
+    /* The likelihood cache of the node (and its inherited flag) is valid
+     * if this equals the HMM's cache_stamp */
+    uint32_t cache_stamp;
+    uint32_t unused;
 } tsls_node_t;
 
-/* A node's likelihood, with the cache used when coalescing. */
+/* A node's likelihood, with the cache used by the walks up the tree. */
 typedef struct {
     /* The node's likelihood, or NULL_LIKELIHOOD or NOT_IN_TREE */
     double value;
-    /* The likelihood the node inherits, cached during coalescing */
+    /* The likelihood the node inherits, valid if the node's cache_stamp
+     * equals the HMM's (see tsls_hmm_get_inherited_likelihood) */
     double cache;
 } tsls_likelihood_t;
 
@@ -178,6 +183,9 @@ typedef struct {
      * tree, which are completed by the arrays below */
     tsls_node_t *tree;
     tsls_likelihood_t *likelihood;
+    /* Likelihood caches are valid on nodes whose cache_stamp equals this;
+     * incrementing it clears them all */
+    uint32_t cache_stamp;
     tsk_id_t *left_child;
     tsk_id_t *right_child;
     tsk_id_t *left_sib;

From 8ac6acf8c97990c853159c21dc7da74689df7f05 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Wed, 28 Oct 2026 17:31:10 +0000
Subject: [PATCH 2/7] Find Viterbi alleles and coalesce through the list's own
 tree

Every node with a mutation at a site is in the likelihood list
(tsls_hmm_insert_mutation_likelihoods), so a listed node without a
mutation has the allele of its nearest listed ancestor. At the start of
each Viterbi site, one memoised walk finds each listed node's nearest
listed proper ancestor (cached on the unlisted nodes passed, with the
likelihood cache's stamp, in the node record's spare field). Alleles are
then resolved by walking up the list's own tree, and coalescing compares
each node with its listed ancestor's value and flag, so the allele walks
up to a mutation or root, and the separate coalescing walks, are gone.
Removals are applied after all the comparisons, since a removed node had
the same value and flag as its own ancestor. Forward and backward are
unchanged.

Paired wall time (ab.py, 6 chr20p samples, best of 3), over the build
before the stamps:

  map:        stamp 0.925-0.936, listed 0.719-0.726
  constant r: stamp 0.991-0.997, listed 0.849-0.916

Paths identical.
---
 tsls/lib/hmm.c  | 115 +++++++++++++++++++++++++++++++++++++++++++++++-
 tsls/lib/tsls.h |   6 ++-
 2 files changed, 118 insertions(+), 3 deletions(-)

diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c
index 131971b..2f9a57e 100644
--- a/tsls/lib/hmm.c
+++ b/tsls/lib/hmm.c
@@ -1359,6 +1359,111 @@ tsls_hmm_get_inherited_likelihood_and_flag(
     return value;
 }
 
+/* Returns the nearest proper ancestor of u that has a likelihood, or
+ * NULL_NODE if u is a root. The answer is cached on the nodes without
+ * likelihoods on the path, using the likelihood cache's stamp, so that
+ * finding the ancestors of all the nodes in the likelihood list costs the
+ * size of the union of their paths up to the nearest node with a
+ * likelihood. These are much shorter than the paths up to a mutation or
+ * root. The caches must be cleared before (they share the stamp with the
+ * likelihood caches) and after, once the likelihood list changes. */
+static ALWAYS_INLINE tsk_id_t
+tsls_hmm_get_listed_ancestor(tsls_hmm_t *self, tsk_id_t u)
+{
+    const tsls_likelihood_t *restrict L = self->likelihood;
+    tsls_node_t *restrict tree = self->tree;
+    const uint32_t stamp = self->cache_stamp;
+    const tsk_id_t p = tree[u].parent;
+    tsk_id_t v, a;
+
+    if (p == NULL_NODE) {
+        return NULL_NODE;
+    }
+    v = p;
+    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
+        v = tree[v].parent;
+    }
+    a = L[v].value >= 0 ? v : tree[v].listed_ancestor;
+    v = p;
+    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
+        tree[v].listed_ancestor = a;
+        tree[v].cache_stamp = stamp;
+        v = tree[v].parent;
+    }
+    return a;
+}
+
+/* Returns the allelic state of the node u with a likelihood, given that
+ * every node in the likelihood list has its listed_ancestor set, and that
+ * the nodes with mutations at the site have their allelic states set and
+ * the roots are marked with ROOT_ALLELE. Because every node with a mutation
+ * at the site has a likelihood (tsls_hmm_insert_mutation_likelihoods), a
+ * node without a mutation has the state of its nearest ancestor with a
+ * likelihood, so the walk goes up the list's own tree rather than the
+ * whole tree. States found are memoised on the list nodes passed, as in
+ * get_memoised_allelic_state. */
+static ALWAYS_INLINE allele_t
+get_listed_allelic_state(tsls_node_t *restrict tree, tsk_id_t u,
+    allele_t ancestral_state, tsk_id_t *restrict touched, int *num_touched)
+{
+    allele_t state;
+    tsk_id_t v = u;
+    int n = *num_touched;
+
+    while (tree[v].allelic_state == TSK_NULL) {
+        v = tree[v].listed_ancestor;
+    }
+    state = tree[v].allelic_state;
+    while (u != v) {
+        tree[u].allelic_state = state;
+        touched[n] = u;
+        n++;
+        u = tree[u].listed_ancestor;
+    }
+    *num_touched = n;
+    return likely(state != ROOT_ALLELE) ? state : ancestral_state;
+}
+
+/* Remove the Viterbi likelihoods that are equal to those of the nearest
+ * ancestor with a likelihood, with the same recombination_required flag, so
+ * that the traceback is unchanged (see tsls_hmm_coalesce_likelihoods). The
+ * listed_ancestor of every node in the list must be up to date. A node
+ * whose ancestor is itself removed here had the same value and flag as that
+ * ancestor, which had the same as its own ancestor, so the nodes are only
+ * removed after all the comparisons. */
+static void
+tsls_hmm_coalesce_listed_likelihoods(tsls_hmm_t *self)
+{
+    tsls_likelihood_t *restrict L = self->likelihood;
+    const tsls_node_t *restrict tree = self->tree;
+    tsk_id_t *restrict L_nodes = self->likelihood_nodes;
+    tsk_id_t *restrict removed = self->likelihood_nodes_tmp;
+    const int old_num_likelihood_nodes = self->num_likelihood_nodes;
+    int j, num_removed, num_likelihood_nodes;
+    tsk_id_t u, a;
+
+    num_removed = 0;
+    num_likelihood_nodes = 0;
+    for (j = 0; j < old_num_likelihood_nodes; j++) {
+        u = L_nodes[j];
+        a = tree[u].listed_ancestor;
+        if (a != NULL_NODE && L[u].value == L[a].value
+            && tree[u].recombination_required == tree[a].recombination_required) {
+            removed[num_removed] = u;
+            num_removed++;
+        } else {
+            L_nodes[num_likelihood_nodes] = u;
+            num_likelihood_nodes++;
+        }
+    }
+    for (j = 0; j < num_removed; j++) {
+        L[removed[j]].value = NULL_LIKELIHOOD;
+        tsls_hmm_record_removed(self, removed[j]);
+    }
+    self->num_likelihood_nodes = num_likelihood_nodes;
+    tsls_hmm_clear_likelihood_caches(self);
+}
+
 /* Remove likelihoods that are equal to the likelihood inherited from the
  * parent. If check_recombination_required is true, also require the
  * recombination_required flag to equal the one inherited from the parent,
@@ -1594,6 +1699,12 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site,
     tsls_hmm_insert_mutation_likelihoods(self, site);
     tsls_hmm_set_allelic_state(self, site);
     self->site_num_updated = (size_t) self->num_likelihood_nodes;
+    /* Find the list's own tree, for the alleles and coalescing */
+    tsls_hmm_clear_likelihood_caches(self);
+    for (j = 0; j < self->num_likelihood_nodes; j++) {
+        u = L_nodes[j];
+        tree[u].listed_ancestor = tsls_hmm_get_listed_ancestor(self, u);
+    }
     max_L = -1;
     max_L_node = NULL_NODE;
     for (j = 0; j < self->num_likelihood_nodes; j++) {
@@ -1611,7 +1722,7 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site,
             p_t = p_recomb;
             tree[u].recombination_required = true;
         }
-        allele = get_memoised_allelic_state(
+        allele = get_listed_allelic_state(
             tree, u, ancestral_state, touched, &num_touched);
         p_e = emission_probability(allele, state, mu, num_alleles);
         if (track_changes && state != TSLS_MISSING_DATA && allele != state) {
@@ -1664,7 +1775,7 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site,
     /* Remove redundant nodes before storing the traceback, keeping those
      * whose recombination_required flag differs from the inherited one.
      * Those are left for the next site, where their flags are recomputed. */
-    tsls_hmm_coalesce_likelihoods(self, true);
+    tsls_hmm_coalesce_listed_likelihoods(self);
     ret = tsls_hmm_store_traceback(self, site, self->removed_nodes,
         self->num_removed_nodes, L_nodes, self->num_likelihood_nodes);
     if (ret != 0) {
diff --git a/tsls/lib/tsls.h b/tsls/lib/tsls.h
index 16c6741..0b8e860 100644
--- a/tsls/lib/tsls.h
+++ b/tsls/lib/tsls.h
@@ -146,7 +146,11 @@ typedef struct {
     /* The likelihood cache of the node (and its inherited flag) is valid
      * if this equals the HMM's cache_stamp */
     uint32_t cache_stamp;
-    uint32_t unused;
+    /* In Viterbi, the nearest proper ancestor with a likelihood: for a node
+     * with a likelihood, as computed at the start of each site, and for a
+     * node without one, cached when cache_stamp is valid (see
+     * tsls_hmm_get_listed_ancestor) */
+    tsk_id_t listed_ancestor;
 } tsls_node_t;
 
 /* A node's likelihood, with the cache used by the walks up the tree. */

From 77277b0fdad80a907e09625dd6e1febf9f572915 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Wed, 28 Oct 2026 17:37:52 +0000
Subject: [PATCH 3/7] Walk to listed ancestors reading only the node records

The walk to each listed node's nearest listed ancestor was ~50% of
tsls_hmm_find_path_matrix after the previous change (perf, chr20p samples
under the map), and each step read two records: the likelihood, to see
whether the node is listed, and the node record, for the parent and
cache. Each listed node is now first marked with the cache stamp and
itself as its listed_ancestor, so the walk stops at a marked node and
reads only the node records. The proper ancestors are collected in a
scratch array and written back after all the walks.

Paired wall time (ab.py, 6 chr20p samples, best of 3):

  map:        0.793, 0.858, 0.864 of the previous build
  constant r: 0.748-0.765 of the build before the stamps
              (previous build 0.809-0.828)

Paths identical.
---
 tsls/lib/hmm.c | 75 +++++++++++++++++++++++++++++---------------------
 1 file changed, 44 insertions(+), 31 deletions(-)

diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c
index 2f9a57e..687d279 100644
--- a/tsls/lib/hmm.c
+++ b/tsls/lib/hmm.c
@@ -1359,38 +1359,55 @@ tsls_hmm_get_inherited_likelihood_and_flag(
     return value;
 }
 
-/* Returns the nearest proper ancestor of u that has a likelihood, or
- * NULL_NODE if u is a root. The answer is cached on the nodes without
- * likelihoods on the path, using the likelihood cache's stamp, so that
- * finding the ancestors of all the nodes in the likelihood list costs the
- * size of the union of their paths up to the nearest node with a
- * likelihood. These are much shorter than the paths up to a mutation or
- * root. The caches must be cleared before (they share the stamp with the
- * likelihood caches) and after, once the likelihood list changes. */
-static ALWAYS_INLINE tsk_id_t
-tsls_hmm_get_listed_ancestor(tsls_hmm_t *self, tsk_id_t u)
+/* Set the listed_ancestor of every node in the likelihood list to its
+ * nearest proper ancestor with a likelihood (NULL_NODE for a root). Each
+ * listed node is first marked, with the stamp of the likelihood caches and
+ * itself as its listed_ancestor, so that the walks up the tree only read
+ * the node records. The answers are cached on the nodes without
+ * likelihoods on the paths, so the cost is the size of the union of the
+ * paths up to the nearest nodes with likelihoods, which are much shorter
+ * than the paths up to a mutation or root. Clears the likelihood caches,
+ * which share the stamp, before; they must be cleared again once the
+ * likelihood list changes. */
+static void
+tsls_hmm_set_listed_ancestors(tsls_hmm_t *self)
 {
-    const tsls_likelihood_t *restrict L = self->likelihood;
     tsls_node_t *restrict tree = self->tree;
-    const uint32_t stamp = self->cache_stamp;
-    const tsk_id_t p = tree[u].parent;
-    tsk_id_t v, a;
+    const tsk_id_t *restrict L_nodes = self->likelihood_nodes;
+    tsk_id_t *restrict ancestor = self->likelihood_nodes_tmp;
+    const int num_likelihood_nodes = self->num_likelihood_nodes;
+    uint32_t stamp;
+    tsk_id_t u, v, a;
+    int j;
 
-    if (p == NULL_NODE) {
-        return NULL_NODE;
+    tsls_hmm_clear_likelihood_caches(self);
+    stamp = self->cache_stamp;
+    for (j = 0; j < num_likelihood_nodes; j++) {
+        u = L_nodes[j];
+        tree[u].cache_stamp = stamp;
+        tree[u].listed_ancestor = u;
     }
-    v = p;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
-        v = tree[v].parent;
+    for (j = 0; j < num_likelihood_nodes; j++) {
+        u = L_nodes[j];
+        a = NULL_NODE;
+        if (tree[u].parent != NULL_NODE) {
+            v = tree[u].parent;
+            while (tree[v].cache_stamp != stamp) {
+                v = tree[v].parent;
+            }
+            a = tree[v].listed_ancestor;
+            v = tree[u].parent;
+            while (tree[v].cache_stamp != stamp) {
+                tree[v].listed_ancestor = a;
+                tree[v].cache_stamp = stamp;
+                v = tree[v].parent;
+            }
+        }
+        ancestor[j] = a;
     }
-    a = L[v].value >= 0 ? v : tree[v].listed_ancestor;
-    v = p;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
-        tree[v].listed_ancestor = a;
-        tree[v].cache_stamp = stamp;
-        v = tree[v].parent;
+    for (j = 0; j < num_likelihood_nodes; j++) {
+        tree[L_nodes[j]].listed_ancestor = ancestor[j];
     }
-    return a;
 }
 
 /* Returns the allelic state of the node u with a likelihood, given that
@@ -1700,11 +1717,7 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site,
     tsls_hmm_set_allelic_state(self, site);
     self->site_num_updated = (size_t) self->num_likelihood_nodes;
     /* Find the list's own tree, for the alleles and coalescing */
-    tsls_hmm_clear_likelihood_caches(self);
-    for (j = 0; j < self->num_likelihood_nodes; j++) {
-        u = L_nodes[j];
-        tree[u].listed_ancestor = tsls_hmm_get_listed_ancestor(self, u);
-    }
+    tsls_hmm_set_listed_ancestors(self);
     max_L = -1;
     max_L_node = NULL_NODE;
     for (j = 0; j < self->num_likelihood_nodes; j++) {

From 5c1d397ef5ae2db4afe9f3aab91a92d3066f27d2 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Wed, 28 Oct 2026 17:58:07 +0000
Subject: [PATCH 4/7] Keep the tree as parent links and child counts

The left_child, right_child, left_sib and right_sib arrays were only
used by the in-tree test and by tsls_hmm_number_tree (sparse sites), but
every edge insertion and removal updated them, touching several cache
lines. The tree is now the parent links plus a count of children per
node. tsls_hmm_number_tree lists the children with a counting sort over
the parent array when it numbers a tree, which is O(N) like the
numbering itself; children are now visited in node ID order rather than
insertion order, which can only change sparse sites' choice between
equally likely nodes.

Measurements before this change (perf, chr20p samples under the map):
the walk to listed ancestors was ~40% of tsls_hmm_find_path_matrix
(~31% of the run), and tsls_hmm_update_tree plus tree position stepping
~7%. Walks averaged 271-280 steps per site under the map (109-123 listed
nodes, 56-58% of them with a listed parent) and 120-129 under constant r
(14-28 listed); 7.2 edges out and 7.2 in per site.

Paired wall time (ab.py, 6 chr20p samples, best of 3), over the previous
build: map 0.994, 0.956, 0.957; constant r 0.943, 0.936, 0.929. Paths
identical. (A guard-rail run on sc2ts was invalidated by rebuilding the
extension mid-run; sc2ts is checked after the next change.)
---
 tsls/lib/hmm.c  | 147 +++++++++++++++++++++---------------------------
 tsls/lib/tsls.h |  14 +++--
 2 files changed, 73 insertions(+), 88 deletions(-)

diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c
index 687d279..05083d6 100644
--- a/tsls/lib/hmm.c
+++ b/tsls/lib/hmm.c
@@ -133,47 +133,23 @@ tsls_large_free(void *p, size_t size)
 #define MODE_FORWARD  1
 #define MODE_BACKWARD 2
 
+/* The tree is kept as parent links and numbers of children: the walks only
+ * go up the tree, and a node is in the tree if it has a parent or a child.
+ * Children are only listed when a tree is numbered (tsls_hmm_number_tree). */
 static inline void
-remove_edge(tsk_id_t p, tsk_id_t c, tsls_node_t *restrict tree,
-    tsk_id_t *restrict left_child, tsk_id_t *restrict right_child,
-    tsk_id_t *restrict left_sib, tsk_id_t *restrict right_sib)
+remove_edge(
+    tsk_id_t p, tsk_id_t c, tsls_node_t *restrict tree, int32_t *restrict num_children)
 {
-    const tsk_id_t lsib = left_sib[c];
-    const tsk_id_t rsib = right_sib[c];
-
-    if (lsib == NULL_NODE) {
-        left_child[p] = rsib;
-    } else {
-        right_sib[lsib] = rsib;
-    }
-    if (rsib == NULL_NODE) {
-        right_child[p] = lsib;
-    } else {
-        left_sib[rsib] = lsib;
-    }
     tree[c].parent = NULL_NODE;
-    left_sib[c] = NULL_NODE;
-    right_sib[c] = NULL_NODE;
+    num_children[p]--;
 }
 
 static inline void
-insert_edge(tsk_id_t p, tsk_id_t c, tsls_node_t *restrict tree,
-    tsk_id_t *restrict left_child, tsk_id_t *restrict right_child,
-    tsk_id_t *restrict left_sib, tsk_id_t *restrict right_sib)
+insert_edge(
+    tsk_id_t p, tsk_id_t c, tsls_node_t *restrict tree, int32_t *restrict num_children)
 {
-    const tsk_id_t u = right_child[p];
-
     tree[c].parent = p;
-    if (u == NULL_NODE) {
-        left_child[p] = c;
-        left_sib[c] = NULL_NODE;
-        right_sib[c] = NULL_NODE;
-    } else {
-        right_sib[u] = c;
-        left_sib[c] = u;
-        right_sib[c] = NULL_NODE;
-    }
-    right_child[p] = c;
+    num_children[p]++;
 }
 
 /* During the forward and backward sweeps, the allelic state of every node
@@ -587,7 +563,7 @@ tsls_index_free(tsls_index_t *self)
 static inline bool
 tsls_hmm_in_tree(const tsls_hmm_t *self, tsk_id_t u)
 {
-    return self->tree[u].parent != NULL_NODE || self->left_child[u] != NULL_NODE;
+    return self->tree[u].parent != NULL_NODE || self->num_children[u] > 0;
 }
 
 static void
@@ -636,11 +612,10 @@ tsls_hmm_print_state(tsls_hmm_t *self, FILE *out)
             (int) self->index->num_alleles[j], (int) self->index->ancestral_state[j]);
     }
     fprintf(out, "tree = \n");
-    fprintf(out, "id\tparent\tlchild\trchild\tlsib\trsib\tlikelihood\n");
+    fprintf(out, "id\tparent\tnum_children\tlikelihood\n");
     for (j = 0; j < (int) self->num_nodes; j++) {
-        fprintf(out, "%d\t%d\t%d\t%d\t%d\t%d\t%f\n", (int) j, self->tree[j].parent,
-            self->left_child[j], self->right_child[j], self->left_sib[j],
-            self->right_sib[j], self->likelihood[j].value);
+        fprintf(out, "%d\t%d\t%d\t%f\n", (int) j, self->tree[j].parent,
+            self->num_children[j], self->likelihood[j].value);
     }
     fprintf(out, "likelihood nodes\n");
     for (j = 0; j < self->num_likelihood_nodes; j++) {
@@ -686,10 +661,9 @@ tsls_hmm_alloc(tsls_hmm_t *self, const tsls_index_t *index,
 
     self->tree = tsls_large_alloc((N + 1) * sizeof(*self->tree));
     self->likelihood = tsls_large_alloc((N + 1) * sizeof(*self->likelihood));
-    self->left_child = tsls_large_alloc((N + 1) * sizeof(*self->left_child));
-    self->right_child = tsls_large_alloc((N + 1) * sizeof(*self->right_child));
-    self->left_sib = tsls_large_alloc((N + 1) * sizeof(*self->left_sib));
-    self->right_sib = tsls_large_alloc((N + 1) * sizeof(*self->right_sib));
+    self->num_children = tsls_large_alloc((N + 1) * sizeof(*self->num_children));
+    self->child_offset = tsls_large_alloc((N + 1) * sizeof(*self->child_offset));
+    self->children = tsls_large_alloc((N + 1) * sizeof(*self->children));
     self->num_below = tsls_large_alloc((N + 1) * sizeof(*self->num_below));
     self->num_inheriting = tsls_large_alloc((N + 1) * sizeof(*self->num_inheriting));
     self->likelihood_nodes = tsls_large_alloc((N + 1) * sizeof(*self->likelihood_nodes));
@@ -709,8 +683,8 @@ tsls_hmm_alloc(tsls_hmm_t *self, const tsls_index_t *index,
         || self->mismatch_rate == NULL || self->traceback == NULL
         || self->max_likelihood_node == NULL || self->tree == NULL
         || self->likelihood == NULL
-        || self->left_child == NULL || self->right_child == NULL
-        || self->left_sib == NULL || self->right_sib == NULL || self->num_below == NULL
+        || self->num_children == NULL || self->child_offset == NULL
+        || self->children == NULL || self->num_below == NULL
         || self->num_inheriting == NULL || self->likelihood_nodes == NULL
         || self->likelihood_nodes_tmp == NULL || self->values_tmp == NULL
         || self->traceback_flags_tmp == NULL || self->removed_nodes == NULL
@@ -750,10 +724,9 @@ tsls_hmm_free(tsls_hmm_t *self)
     tsi_safe_free(self->mismatch_rate);
     tsls_large_free_array(self->tree, self->num_nodes);
     tsls_large_free_array(self->likelihood, self->num_nodes);
-    tsls_large_free_array(self->left_child, self->num_nodes);
-    tsls_large_free_array(self->right_child, self->num_nodes);
-    tsls_large_free_array(self->left_sib, self->num_nodes);
-    tsls_large_free_array(self->right_sib, self->num_nodes);
+    tsls_large_free_array(self->num_children, self->num_nodes);
+    tsls_large_free_array(self->child_offset, self->num_nodes);
+    tsls_large_free_array(self->children, self->num_nodes);
     tsls_large_free_array(self->num_below, self->num_nodes);
     tsls_large_free_array(self->num_inheriting, self->num_nodes);
     tsls_large_free_array(self->likelihood_nodes, self->num_nodes);
@@ -1043,8 +1016,7 @@ tsls_hmm_remove_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c,
     if (update_likelihoods && L[c].value == NULL_LIKELIHOOD) {
         tsls_hmm_add_likelihood(self, c, tsls_hmm_get_inherited_likelihood(self, p));
     }
-    remove_edge(p, c, self->tree, self->left_child, self->right_child,
-        self->left_sib, self->right_sib);
+    remove_edge(p, c, self->tree, self->num_children);
     /* c is now a root, and so has a likelihood, which walks over the
      * cache use in preference to any cached value */
     self->tree[c].allelic_state = ROOT_ALLELE;
@@ -1081,8 +1053,7 @@ tsls_hmm_insert_edge(tsls_hmm_t *self, tsk_id_t p, tsk_id_t c, bool update_likel
         /* Otherwise the child was a root and keeps its likelihood */
         assert(L[c].value >= 0);
     }
-    insert_edge(p, c, self->tree, self->left_child, self->right_child,
-        self->left_sib, self->right_sib);
+    insert_edge(p, c, self->tree, self->num_children);
     self->tree[c].allelic_state = TSK_NULL;
     if (self->track_num_below) {
         tsls_hmm_update_num_below(self, p, self->num_below[c]);
@@ -1141,10 +1112,7 @@ tsls_hmm_init_sweep(
     tsk_id_t u, j;
     size_t k;
 
-    memset(self->left_child, 0xff, N * sizeof(*self->left_child));
-    memset(self->right_child, 0xff, N * sizeof(*self->right_child));
-    memset(self->left_sib, 0xff, N * sizeof(*self->left_sib));
-    memset(self->right_sib, 0xff, N * sizeof(*self->right_sib));
+    memset(self->num_children, 0, N * sizeof(*self->num_children));
     for (k = 0; k < N; k++) {
         /* Every node starts as a root. The likelihoods are set below. */
         self->tree[k].cache_stamp = 0;
@@ -1174,7 +1142,7 @@ tsls_hmm_init_sweep(
     for (u = 0; u < (tsk_id_t) N; u++) {
         if (self->tree[u].parent != NULL_NODE) {
             self->likelihood[u].value = NULL_LIKELIHOOD;
-        } else if (self->left_child[u] != NULL_NODE) {
+        } else if (self->num_children[u] > 0) {
             tsls_hmm_add_likelihood(self, u, initial_value);
         } else {
             self->likelihood[u].value = NOT_IN_TREE;
@@ -1191,48 +1159,63 @@ tsls_hmm_number_tree(tsls_hmm_t *self)
 {
     const tsk_id_t N = (tsk_id_t) self->num_nodes;
     const tsls_node_t *restrict tree = self->tree;
-    const tsk_id_t *restrict left_child = self->left_child;
-    const tsk_id_t *restrict right_sib = self->right_sib;
+    const int32_t *restrict num_children = self->num_children;
+    int32_t *restrict child_offset = self->child_offset;
+    tsk_id_t *restrict children = self->children;
     const tsls_likelihood_t *restrict L = self->likelihood;
     int32_t *restrict preorder = self->preorder;
     int32_t *restrict subtree_end = self->subtree_end;
     tsk_id_t *restrict L_nodes = self->likelihood_nodes;
     tsk_id_t *restrict order = self->likelihood_nodes_tmp;
-    tsk_id_t root, u;
-    int32_t k, j;
+    /* The position in each node's children of the next child to visit */
+    int32_t *restrict next_child = subtree_end;
+    tsk_id_t root, u, p;
+    int32_t k, j, total;
     int n;
-    bool done;
 
+    /* List the children of each node, in order of node ID */
+    total = 0;
     for (u = 0; u < N; u++) {
+        child_offset[u] = total;
+        next_child[u] = total;
+        total += num_children[u];
+    }
+    child_offset[N] = total;
+    for (u = 0; u < N; u++) {
+        p = tree[u].parent;
+        if (p != NULL_NODE) {
+            children[next_child[p]] = u;
+            next_child[p]++;
+        }
         preorder[u] = -1;
     }
+    for (u = 0; u < N; u++) {
+        next_child[u] = child_offset[u];
+    }
     k = 0;
     for (root = 0; root < N; root++) {
-        if (tree[root].parent != NULL_NODE || left_child[root] == NULL_NODE) {
+        if (tree[root].parent != NULL_NODE || num_children[root] == 0) {
             continue;
         }
+        preorder[root] = k;
+        order[k] = root;
+        k++;
         u = root;
-        done = false;
-        while (!done) {
-            preorder[u] = k;
-            order[k] = u;
-            k++;
-            if (left_child[u] != NULL_NODE) {
-                u = left_child[u];
+        while (true) {
+            if (next_child[u] < child_offset[u + 1]) {
+                /* Visit u's next child */
+                u = children[next_child[u]++];
+                preorder[u] = k;
+                order[k] = u;
+                k++;
             } else {
-                /* Leave u, and every ancestor whose last child we've left */
-                while (true) {
-                    subtree_end[u] = k;
-                    if (u == root) {
-                        done = true;
-                        break;
-                    }
-                    if (right_sib[u] != NULL_NODE) {
-                        u = right_sib[u];
-                        break;
-                    }
-                    u = tree[u].parent;
+                /* Leave u, whose subtree has been visited (this overwrites
+                 * u's next_child, which is no longer needed) */
+                subtree_end[u] = k;
+                if (u == root) {
+                    break;
                 }
+                u = tree[u].parent;
             }
         }
     }
diff --git a/tsls/lib/tsls.h b/tsls/lib/tsls.h
index 0b8e860..53b5346 100644
--- a/tsls/lib/tsls.h
+++ b/tsls/lib/tsls.h
@@ -183,17 +183,19 @@ typedef struct {
     double prune_threshold;
     double *recombination_rate;
     double *mismatch_rate;
-    /* The per-node state, with the parent links of the quintuply linked
-     * tree, which are completed by the arrays below */
+    /* The per-node state, including the tree's parent links */
     tsls_node_t *tree;
     tsls_likelihood_t *likelihood;
     /* Likelihood caches are valid on nodes whose cache_stamp equals this;
      * incrementing it clears them all */
     uint32_t cache_stamp;
-    tsk_id_t *left_child;
-    tsk_id_t *right_child;
-    tsk_id_t *left_sib;
-    tsk_id_t *right_sib;
+    /* The number of children of each node in the current tree */
+    int32_t *num_children;
+    /* The children of each node, as ranges child_offset[u] to
+     * child_offset[u + 1] of children, built when a tree is numbered (see
+     * tsls_hmm_number_tree) */
+    int32_t *child_offset;
+    tsk_id_t *children;
     /* The number of nodes in the subtree below each node, including itself.
      * Only maintained for the forward and backward algorithms. */
     int64_t *num_below;

From 099cec67a28e4c97d7f04d9ff0b74b3e4a678578 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Wed, 28 Oct 2026 17:58:30 +0000
Subject: [PATCH 5/7] Keep the walk caches across sites

Every walk up the tree now goes to the nearest node with a likelihood
(tsls_hmm_find_listed) and caches the answer, as a node ID, on the nodes
without likelihoods that it passes. Nodes with likelihoods are always
marked (the cache stamp, and themselves as listed_ancestor), so walks
read only node records. The caches are no longer cleared after each
phase: a cached answer stays right until a node on a cached path gets a
likelihood or a new parent. So:

- a node given a likelihood that has a cached walk clears all the caches
  (and the list is marked again), since walks from below may have passed
  through it;
- an edge removal gives the child a likelihood first, which does the same;
- a node losing its likelihood (coalescing) becomes a cached node pointing
  to its old nearest listed ancestor, and walks follow such pointers until
  a node with a likelihood;
- nodes leaving the tree are isolated, so no walk ends there.

The value caches (L.cache, recombination_required_cache) are gone: the
inherited value and flag are read from the listed ancestor. In Viterbi,
each listed node's proper listed ancestor is kept in its likelihood
record, for the alleles and coalescing. tsls_hmm_check_state checks that
every cached walk leads to the nearest listed ancestor.

Measured before implementing, by checking the previous site's stamps: a
full clear would be needed at 24% of sites under the map (mutation node on
a cached path 9%, an edge-removal child on one 19%) and 14-15% under
constant r.

Paired wall time (ab.py, 6 chr20p samples, best of 3), over the previous
build: map 0.732, 0.719, 0.716; constant r 0.668, 0.698, 0.697. Paths
identical.
---
 tsls/lib/hmm.c  | 288 +++++++++++++++++++++++++-----------------------
 tsls/lib/tsls.h |  30 ++---
 2 files changed, 163 insertions(+), 155 deletions(-)

diff --git a/tsls/lib/hmm.c b/tsls/lib/hmm.c
index 05083d6..b470915 100644
--- a/tsls/lib/hmm.c
+++ b/tsls/lib/hmm.c
@@ -571,7 +571,7 @@ tsls_hmm_check_state(const tsls_hmm_t *self)
 {
     int num_likelihoods = 0;
     int j;
-    tsk_id_t u;
+    tsk_id_t u, v, a;
 
     for (j = 0; j < self->num_likelihood_nodes; j++) {
         u = self->likelihood_nodes[j];
@@ -592,7 +592,22 @@ tsls_hmm_check_state(const tsls_hmm_t *self)
         }
         assert(self->tree[u].allelic_state
                == (self->tree[u].parent == NULL_NODE ? ROOT_ALLELE : TSK_NULL));
-        assert(self->tree[u].cache_stamp != self->cache_stamp);
+        if (self->likelihood[u].value >= 0) {
+            /* Nodes with likelihoods are marked for the walk caches */
+            assert(self->tree[u].cache_stamp == self->cache_stamp);
+            assert(self->tree[u].listed_ancestor == u);
+        } else if (self->tree[u].cache_stamp == self->cache_stamp) {
+            /* A cached walk leads to the nearest ancestor with a likelihood */
+            v = self->tree[u].parent;
+            while (self->likelihood[v].value < 0) {
+                v = self->tree[v].parent;
+            }
+            a = self->tree[u].listed_ancestor;
+            while (self->tree[a].listed_ancestor != a) {
+                a = self->tree[a].listed_ancestor;
+            }
+            assert(a == v);
+        }
     }
     assert(num_likelihoods == self->num_likelihood_nodes);
     (void) num_likelihoods;
@@ -851,9 +866,62 @@ tsls_hmm_count_distinct(tsls_hmm_t *self)
     self->total_distinct_values += num_distinct;
 }
 
+/* Clear all the walk caches (see tsls_node_t), marking the nodes with
+ * likelihoods again. The caches stay valid as long as no node on a cached
+ * path gets a likelihood or a new parent, so this is only needed when one
+ * does (see tsls_hmm_mark_listed and tsls_hmm_remove_edge). */
+static void
+tsls_hmm_clear_walk_caches(tsls_hmm_t *self)
+{
+    tsls_node_t *restrict tree = self->tree;
+    const tsk_id_t *restrict L_nodes = self->likelihood_nodes;
+    tsk_id_t u;
+    size_t k;
+    int j;
+
+    self->cache_stamp++;
+    if (unlikely(self->cache_stamp == 0)) {
+        for (k = 0; k < self->num_nodes; k++) {
+            tree[k].cache_stamp = 0;
+        }
+        self->cache_stamp = 1;
+    }
+    for (j = 0; j < self->num_likelihood_nodes; j++) {
+        u = L_nodes[j];
+        tree[u].cache_stamp = self->cache_stamp;
+        tree[u].listed_ancestor = u;
+    }
+}
+
+/* Mark u, which is being given a likelihood, as a node with a likelihood
+ * for the walk caches. If u has a cached walk, walks from below may have
+ * passed through it to an ancestor that is no longer the nearest, so the
+ * caches are cleared first. The caller must not have walked from u itself
+ * since it last changed. */
+static inline void
+tsls_hmm_mark_listed(tsls_hmm_t *self, tsk_id_t u)
+{
+    tsls_node_t *restrict tree = self->tree;
+
+    if (unlikely(tree[u].cache_stamp == self->cache_stamp)) {
+        tsls_hmm_clear_walk_caches(self);
+    }
+    tree[u].cache_stamp = self->cache_stamp;
+    tree[u].listed_ancestor = u;
+}
+
+/* Mark u, which has just lost its likelihood, as a node without one whose
+ * nearest ancestor with a likelihood is ancestor (via listed_ancestor). */
+static inline void
+tsls_hmm_mark_unlisted(tsls_hmm_t *self, tsk_id_t u, tsk_id_t ancestor)
+{
+    self->tree[u].listed_ancestor = ancestor;
+}
+
 static inline void
 tsls_hmm_add_likelihood(tsls_hmm_t *self, tsk_id_t u, double value)
 {
+    tsls_hmm_mark_listed(self, u);
     self->likelihood[u].value = value;
     self->likelihood_nodes[self->num_likelihood_nodes] = u;
     self->num_likelihood_nodes++;
@@ -888,6 +956,7 @@ tsls_hmm_add_likelihood_sorted(tsls_hmm_t *self, tsk_id_t u, double value)
     tsk_id_t *L_nodes = self->likelihood_nodes;
     const int j = tsls_hmm_lower_bound(self, self->preorder[u]);
 
+    tsls_hmm_mark_listed(self, u);
     memmove(L_nodes + j + 1, L_nodes + j,
         (size_t) (self->num_likelihood_nodes - j) * sizeof(*L_nodes));
     L_nodes[j] = u;
@@ -895,9 +964,10 @@ tsls_hmm_add_likelihood_sorted(tsls_hmm_t *self, tsk_id_t u, double value)
     self->likelihood[u].value = value;
 }
 
-/* Remove u from the likelihood list, which is sorted by preorder. */
+/* Remove u from the likelihood list, which is sorted by preorder, given
+ * its nearest ancestor with a likelihood. */
 static void
-tsls_hmm_remove_likelihood_sorted(tsls_hmm_t *self, tsk_id_t u)
+tsls_hmm_remove_likelihood_sorted(tsls_hmm_t *self, tsk_id_t u, tsk_id_t ancestor)
 {
     tsk_id_t *L_nodes = self->likelihood_nodes;
     const int j = tsls_hmm_lower_bound(self, self->preorder[u]);
@@ -907,6 +977,7 @@ tsls_hmm_remove_likelihood_sorted(tsls_hmm_t *self, tsk_id_t u)
         (size_t) (self->num_likelihood_nodes - j - 1) * sizeof(*L_nodes));
     self->num_likelihood_nodes--;
     self->likelihood[u].value = NULL_LIKELIHOOD;
+    tsls_hmm_mark_unlisted(self, u, ancestor);
 }
 
 /* Record that u has been removed from the likelihood list at the current
@@ -953,50 +1024,45 @@ tsls_hmm_update_num_below(tsls_hmm_t *self, tsk_id_t u, int64_t delta)
     }
 }
 
-/* Returns the likelihood inherited by node u, filling in the likelihood
- * cache along the path so that later queries are fast. The walks up the
- * tree here and in tsls_hmm_get_inherited_likelihood_and_flag are inner
- * loops, so they have no asserts (every root has a likelihood, which
- * tsls_hmm_check_state checks). */
-static inline double
-tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u)
+/* Returns the nearest ancestor of u (or u itself) with a likelihood. The
+ * answer is cached on the nodes without likelihoods on the path (see
+ * tsls_node_t), so that later walks through them stop there; the caches
+ * stay valid across sites until a node on a cached path gets a likelihood
+ * or a new parent. A walk stops at a node with the current stamp, which is
+ * either a node with a likelihood (its own listed_ancestor) or a cached
+ * node, whose listed_ancestor may be a node that has since lost its
+ * likelihood, and so is followed until a node with a likelihood. This is
+ * the innermost loop of the HMM, so it has no asserts (every root has a
+ * likelihood, which tsls_hmm_check_state checks). */
+static ALWAYS_INLINE tsk_id_t
+tsls_hmm_find_listed(tsls_hmm_t *self, tsk_id_t u)
 {
-    tsls_likelihood_t *restrict L = self->likelihood;
     tsls_node_t *restrict tree = self->tree;
     const uint32_t stamp = self->cache_stamp;
-    double value;
-    tsk_id_t v;
+    tsk_id_t v, a;
 
     v = u;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
+    while (tree[v].cache_stamp != stamp) {
         v = tree[v].parent;
     }
-    /* A node with a likelihood may have a stale cache from before it was
-     * given one, so its likelihood takes precedence */
-    value = L[v].value >= 0 ? L[v].value : L[v].cache;
+    a = tree[v].listed_ancestor;
+    while (unlikely(tree[a].listed_ancestor != a)) {
+        a = tree[a].listed_ancestor;
+    }
     v = u;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
-        L[v].cache = value;
+    while (tree[v].cache_stamp != stamp) {
+        tree[v].listed_ancestor = a;
         tree[v].cache_stamp = stamp;
         v = tree[v].parent;
     }
-    return value;
+    return a;
 }
 
-/* Clear all the likelihood caches, which are only valid while the tree and
- * the set of nodes with likelihoods are unchanged. */
-static inline void
-tsls_hmm_clear_likelihood_caches(tsls_hmm_t *self)
+/* Returns the likelihood inherited by node u. */
+static inline double
+tsls_hmm_get_inherited_likelihood(tsls_hmm_t *self, tsk_id_t u)
 {
-    size_t j;
-
-    self->cache_stamp++;
-    if (unlikely(self->cache_stamp == 0)) {
-        for (j = 0; j < self->num_nodes; j++) {
-            self->tree[j].cache_stamp = 0;
-        }
-        self->cache_stamp = 1;
-    }
+    return self->likelihood[tsls_hmm_find_listed(self, u)].value;
 }
 
 /* Remove the edge from p to c from the tree, and if update_likelihoods is
@@ -1077,7 +1143,6 @@ tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp)
         tsls_hmm_remove_edge(
             self, out[j].parent, out[j].child, true, isolated, &num_isolated);
     }
-    tsls_hmm_clear_likelihood_caches(self);
     for (j = tp->in.start; j != tp->in.stop; j += direction) {
         tsls_hmm_insert_edge(self, in[j].parent, in[j].child, true);
     }
@@ -1089,7 +1154,9 @@ tsls_hmm_update_tree(tsls_hmm_t *self, const tsk_tree_position_t *tp)
     for (k = 0; k < num_isolated; k++) {
         u = isolated[k];
         if (!tsls_hmm_in_tree(self, u)) {
+            /* An isolated node has no descendants, so no walk ends here */
             L[u].value = NOT_IN_TREE;
+            self->tree[u].cache_stamp = 0;
             tsls_hmm_record_removed(self, u);
             num_left++;
         }
@@ -1119,7 +1186,6 @@ tsls_hmm_init_sweep(
         self->tree[k].parent = NULL_NODE;
         self->tree[k].allelic_state = ROOT_ALLELE;
         self->tree[k].recombination_required = -1;
-        self->tree[k].recombination_required_cache = -1;
         self->tree[k].traceback_flag = -1;
         self->num_below[k] = 1;
     }
@@ -1293,10 +1359,9 @@ tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site)
     for (j = index->site_mutations_offset[site]; j < stop; j++) {
         node = index->mutation_node[j];
         if (L[node].value == NULL_LIKELIHOOD) {
-            u = node;
-            while (L[u].value == NULL_LIKELIHOOD) {
-                u = tree[u].parent;
-            }
+            /* Walk from the parent, so that the node's own cache (if any)
+             * is seen when it's marked (see tsls_hmm_mark_listed) */
+            u = tsls_hmm_find_listed(self, tree[node].parent);
             if (self->tree_numbered) {
                 tsls_hmm_add_likelihood_sorted(self, node, L[u].value);
             } else {
@@ -1307,89 +1372,36 @@ tsls_hmm_insert_mutation_likelihoods(tsls_hmm_t *self, const tsk_id_t site)
 }
 
 /* As tsls_hmm_get_inherited_likelihood, but also return the
- * recombination_required flag that u inherits, caching it along the path
- * with the likelihood. */
+ * recombination_required flag that u inherits. */
 static ALWAYS_INLINE double
 tsls_hmm_get_inherited_likelihood_and_flag(
     tsls_hmm_t *self, tsk_id_t u, int8_t *flag)
 {
-    tsls_likelihood_t *restrict L = self->likelihood;
-    tsls_node_t *restrict tree = self->tree;
-    const uint32_t stamp = self->cache_stamp;
-    double value;
-    int8_t r;
-    tsk_id_t v;
+    const tsk_id_t a = tsls_hmm_find_listed(self, u);
 
-    v = u;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
-        v = tree[v].parent;
-    }
-    if (L[v].value >= 0) {
-        value = L[v].value;
-        r = tree[v].recombination_required;
-    } else {
-        value = L[v].cache;
-        r = tree[v].recombination_required_cache;
-    }
-    v = u;
-    while (likely(L[v].value == NULL_LIKELIHOOD) && likely(tree[v].cache_stamp != stamp)) {
-        L[v].cache = value;
-        tree[v].recombination_required_cache = r;
-        tree[v].cache_stamp = stamp;
-        v = tree[v].parent;
-    }
-    *flag = r;
-    return value;
+    *flag = self->tree[a].recombination_required;
+    return self->likelihood[a].value;
 }
 
-/* Set the listed_ancestor of every node in the likelihood list to its
- * nearest proper ancestor with a likelihood (NULL_NODE for a root). Each
- * listed node is first marked, with the stamp of the likelihood caches and
- * itself as its listed_ancestor, so that the walks up the tree only read
- * the node records. The answers are cached on the nodes without
- * likelihoods on the paths, so the cost is the size of the union of the
- * paths up to the nearest nodes with likelihoods, which are much shorter
- * than the paths up to a mutation or root. Clears the likelihood caches,
- * which share the stamp, before; they must be cleared again once the
- * likelihood list changes. */
+/* Set the ancestor of every node in the likelihood list to its nearest
+ * proper ancestor with a likelihood (NULL_NODE for a root). The walks are
+ * cached (tsls_hmm_find_listed), so the cost is the size of the union of
+ * the paths not already cached, up to the nearest nodes with likelihoods,
+ * which are much shorter than the paths up to a mutation or root. */
 static void
 tsls_hmm_set_listed_ancestors(tsls_hmm_t *self)
 {
-    tsls_node_t *restrict tree = self->tree;
+    tsls_likelihood_t *restrict L = self->likelihood;
+    const tsls_node_t *restrict tree = self->tree;
     const tsk_id_t *restrict L_nodes = self->likelihood_nodes;
-    tsk_id_t *restrict ancestor = self->likelihood_nodes_tmp;
     const int num_likelihood_nodes = self->num_likelihood_nodes;
-    uint32_t stamp;
-    tsk_id_t u, v, a;
+    tsk_id_t u, p;
     int j;
 
-    tsls_hmm_clear_likelihood_caches(self);
-    stamp = self->cache_stamp;
     for (j = 0; j < num_likelihood_nodes; j++) {
         u = L_nodes[j];
-        tree[u].cache_stamp = stamp;
-        tree[u].listed_ancestor = u;
-    }
-    for (j = 0; j < num_likelihood_nodes; j++) {
-        u = L_nodes[j];
-        a = NULL_NODE;
-        if (tree[u].parent != NULL_NODE) {
-            v = tree[u].parent;
-            while (tree[v].cache_stamp != stamp) {
-                v = tree[v].parent;
-            }
-            a = tree[v].listed_ancestor;
-            v = tree[u].parent;
-            while (tree[v].cache_stamp != stamp) {
-                tree[v].listed_ancestor = a;
-                tree[v].cache_stamp = stamp;
-                v = tree[v].parent;
-            }
-        }
-        ancestor[j] = a;
-    }
-    for (j = 0; j < num_likelihood_nodes; j++) {
-        tree[L_nodes[j]].listed_ancestor = ancestor[j];
+        p = tree[u].parent;
+        L[u].ancestor = p == NULL_NODE ? NULL_NODE : tsls_hmm_find_listed(self, p);
     }
 }
 
@@ -1403,22 +1415,23 @@ tsls_hmm_set_listed_ancestors(tsls_hmm_t *self)
  * whole tree. States found are memoised on the list nodes passed, as in
  * get_memoised_allelic_state. */
 static ALWAYS_INLINE allele_t
-get_listed_allelic_state(tsls_node_t *restrict tree, tsk_id_t u,
-    allele_t ancestral_state, tsk_id_t *restrict touched, int *num_touched)
+get_listed_allelic_state(tsls_node_t *restrict tree,
+    const tsls_likelihood_t *restrict L, tsk_id_t u, allele_t ancestral_state,
+    tsk_id_t *restrict touched, int *num_touched)
 {
     allele_t state;
     tsk_id_t v = u;
     int n = *num_touched;
 
     while (tree[v].allelic_state == TSK_NULL) {
-        v = tree[v].listed_ancestor;
+        v = L[v].ancestor;
     }
     state = tree[v].allelic_state;
     while (u != v) {
         tree[u].allelic_state = state;
         touched[n] = u;
         n++;
-        u = tree[u].listed_ancestor;
+        u = L[u].ancestor;
     }
     *num_touched = n;
     return likely(state != ROOT_ALLELE) ? state : ancestral_state;
@@ -1427,7 +1440,7 @@ get_listed_allelic_state(tsls_node_t *restrict tree, tsk_id_t u,
 /* Remove the Viterbi likelihoods that are equal to those of the nearest
  * ancestor with a likelihood, with the same recombination_required flag, so
  * that the traceback is unchanged (see tsls_hmm_coalesce_likelihoods). The
- * listed_ancestor of every node in the list must be up to date. A node
+ * ancestor of every node in the list must be up to date. A node
  * whose ancestor is itself removed here had the same value and flag as that
  * ancestor, which had the same as its own ancestor, so the nodes are only
  * removed after all the comparisons. */
@@ -1446,7 +1459,7 @@ tsls_hmm_coalesce_listed_likelihoods(tsls_hmm_t *self)
     num_likelihood_nodes = 0;
     for (j = 0; j < old_num_likelihood_nodes; j++) {
         u = L_nodes[j];
-        a = tree[u].listed_ancestor;
+        a = L[u].ancestor;
         if (a != NULL_NODE && L[u].value == L[a].value
             && tree[u].recombination_required == tree[a].recombination_required) {
             removed[num_removed] = u;
@@ -1457,11 +1470,12 @@ tsls_hmm_coalesce_listed_likelihoods(tsls_hmm_t *self)
         }
     }
     for (j = 0; j < num_removed; j++) {
-        L[removed[j]].value = NULL_LIKELIHOOD;
-        tsls_hmm_record_removed(self, removed[j]);
+        u = removed[j];
+        L[u].value = NULL_LIKELIHOOD;
+        tsls_hmm_mark_unlisted(self, u, L[u].ancestor);
+        tsls_hmm_record_removed(self, u);
     }
     self->num_likelihood_nodes = num_likelihood_nodes;
-    tsls_hmm_clear_likelihood_caches(self);
 }
 
 /* Remove likelihoods that are equal to the likelihood inherited from the
@@ -1475,9 +1489,7 @@ tsls_hmm_coalesce_likelihoods(
 {
     tsls_likelihood_t *restrict L = self->likelihood;
     tsls_node_t *restrict tree = self->tree;
-    double L_p;
-    int8_t R_p;
-    tsk_id_t u, p;
+    tsk_id_t u, p, a;
     const int old_num_likelihood_nodes = self->num_likelihood_nodes;
     tsk_id_t *restrict L_nodes = self->likelihood_nodes;
     int j, num_likelihood_nodes;
@@ -1487,20 +1499,18 @@ tsls_hmm_coalesce_likelihoods(
         u = L_nodes[j];
         p = tree[u].parent;
         if (p != NULL_NODE) {
-            if (check_recombination_required) {
-                /* A removed node had the inherited value and flag, so the
-                 * cached values below it stay valid. */
-                L_p = tsls_hmm_get_inherited_likelihood_and_flag(
-                    self, p, &R_p);
-                if (L[u].value == L_p && tree[u].recombination_required == R_p) {
-                    L[u].value = NULL_LIKELIHOOD;
+            /* A removed node had the inherited value (and flag), so walks
+             * through it to its own ancestor give the same answer. */
+            a = tsls_hmm_find_listed(self, p);
+            if (L[u].value == L[a].value
+                && (!check_recombination_required
+                    || tree[u].recombination_required
+                           == tree[a].recombination_required)) {
+                L[u].value = NULL_LIKELIHOOD;
+                tsls_hmm_mark_unlisted(self, u, a);
+                if (check_recombination_required) {
                     tsls_hmm_record_removed(self, u);
                 }
-            } else {
-                L_p = tsls_hmm_get_inherited_likelihood(self, p);
-                if (L[u].value == L_p) {
-                    L[u].value = NULL_LIKELIHOOD;
-                }
             }
         }
         if (L[u].value >= 0) {
@@ -1509,7 +1519,6 @@ tsls_hmm_coalesce_likelihoods(
         }
     }
     self->num_likelihood_nodes = num_likelihood_nodes;
-    tsls_hmm_clear_likelihood_caches(self);
 }
 
 static inline double
@@ -1719,7 +1728,7 @@ tsls_hmm_process_site_viterbi(tsls_hmm_t *self, const tsk_id_t site,
             tree[u].recombination_required = true;
         }
         allele = get_listed_allelic_state(
-            tree, u, ancestral_state, touched, &num_touched);
+            tree, L, u, ancestral_state, touched, &num_touched);
         p_e = emission_probability(allele, state, mu, num_alleles);
         if (track_changes && state != TSLS_MISSING_DATA && allele != state) {
             old_flag[j] = (int8_t) (old_flag[j] + 4);
@@ -1879,14 +1888,14 @@ tsls_hmm_process_site_sparse(tsls_hmm_t *self, const tsk_id_t site, const allele
     const double p_recomb = rho / n;
     mutation_range_t ranges[SPARSE_MAX_MUTATIONS];
     int32_t bounds[2 * SPARSE_MAX_MUTATIONS + 2];
-    double value, old_value, p_no_recomb, p_t, L_p;
+    double value, old_value, p_no_recomb, p_t;
     int8_t old_flag, R_p;
     allele_t allele;
     bool mismatch;
     int j, k, num_ranges, num_bounds, num_nodes, num_changed;
     int32_t x;
     tsk_size_t i;
-    tsk_id_t u, p;
+    tsk_id_t u, p, a;
 
     if (mutations_stop - mutations_start > SPARSE_MAX_MUTATIONS) {
         goto out;
@@ -1932,7 +1941,6 @@ tsls_hmm_process_site_sparse(tsls_hmm_t *self, const tsk_id_t site, const allele
                         ranges, num_ranges, preorder[p], ancestral_state)
                     != state)) {
                 value = tsls_hmm_get_inherited_likelihood_and_flag(self, p, &R_p);
-                tsls_hmm_clear_likelihood_caches(self);
                 tsls_hmm_add_likelihood_sorted(self, u, value);
                 tree[u].recombination_required = R_p;
                 tsls_hmm_add_site_node(self, u, &num_nodes);
@@ -2007,13 +2015,13 @@ tsls_hmm_process_site_sparse(tsls_hmm_t *self, const tsk_id_t site, const allele
         u = nodes[j];
         p = tree[u].parent;
         if (p != NULL_NODE) {
-            L_p = tsls_hmm_get_inherited_likelihood_and_flag(self, p, &R_p);
-            if (L[u].value == L_p && tree[u].recombination_required == R_p) {
-                tsls_hmm_remove_likelihood_sorted(self, u);
+            a = tsls_hmm_find_listed(self, p);
+            if (L[u].value == L[a].value
+                && tree[u].recombination_required == tree[a].recombination_required) {
+                tsls_hmm_remove_likelihood_sorted(self, u, a);
             }
         }
     }
-    tsls_hmm_clear_likelihood_caches(self);
     ret = tsls_hmm_store_traceback(self, site, nodes, num_nodes, NULL, 0);
     if (ret != 0) {
         goto out;
diff --git a/tsls/lib/tsls.h b/tsls/lib/tsls.h
index 53b5346..16d37c7 100644
--- a/tsls/lib/tsls.h
+++ b/tsls/lib/tsls.h
@@ -137,29 +137,29 @@ typedef struct {
      * there (or is a root, during the sweeps) */
     allele_t allelic_state;
     int8_t recombination_required;
-    /* The recombination_required flag inherited by the node, valid where
-     * the likelihood cache is. Only filled when coalescing Viterbi. */
-    int8_t recombination_required_cache;
+    int8_t unused;
     /* The node's flag in the traceback as of the last site stored, or -1
      * if it is not in the traceback's set */
     int8_t traceback_flag;
-    /* The likelihood cache of the node (and its inherited flag) is valid
-     * if this equals the HMM's cache_stamp */
+    /* The walks up the tree to the nearest node with a likelihood are
+     * cached (see tsls_hmm_find_listed). A node whose cache_stamp equals the
+     * HMM's is either a node with a likelihood, whose listed_ancestor is
+     * itself, or a node without one whose listed_ancestor leads, by
+     * following listed_ancestor, to its nearest ancestor with a likelihood.
+     * Nodes with likelihoods always have the current stamp. */
     uint32_t cache_stamp;
-    /* In Viterbi, the nearest proper ancestor with a likelihood: for a node
-     * with a likelihood, as computed at the start of each site, and for a
-     * node without one, cached when cache_stamp is valid (see
-     * tsls_hmm_get_listed_ancestor) */
     tsk_id_t listed_ancestor;
 } tsls_node_t;
 
-/* A node's likelihood, with the cache used by the walks up the tree. */
+/* A node's likelihood. */
 typedef struct {
     /* The node's likelihood, or NULL_LIKELIHOOD or NOT_IN_TREE */
     double value;
-    /* The likelihood the node inherits, valid if the node's cache_stamp
-     * equals the HMM's (see tsls_hmm_get_inherited_likelihood) */
-    double cache;
+    /* For a node with a likelihood, in Viterbi, its nearest proper ancestor
+     * with a likelihood (NULL_NODE for a root), set at the start of each
+     * site (see tsls_hmm_set_listed_ancestors) */
+    tsk_id_t ancestor;
+    int32_t unused;
 } tsls_likelihood_t;
 
 /* The Li and Stephens HMM on a tree sequence. The index is borrowed, and
@@ -186,8 +186,8 @@ typedef struct {
     /* The per-node state, including the tree's parent links */
     tsls_node_t *tree;
     tsls_likelihood_t *likelihood;
-    /* Likelihood caches are valid on nodes whose cache_stamp equals this;
-     * incrementing it clears them all */
+    /* The walk caches are valid on nodes whose cache_stamp equals this;
+     * incrementing it clears them all (see tsls_hmm_clear_walk_caches) */
     uint32_t cache_stamp;
     /* The number of children of each node in the current tree */
     int32_t *num_children;

From 9b0112be62aee8d7f45fc3f5d224fe4770a18cc8 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Wed, 28 Oct 2026 18:14:11 +0000
Subject: [PATCH 6/7] Document the walk cache optimisations and where the time
 goes

Add a benchmarking README section on this round of algorithmic work on
sample matching under a recombination map: the profile it started from,
each change and its paired timing, the end-to-end result (0.45-0.47 of
the earlier build under the map, 0.50-0.52 under constant r, identical
paths), the guard-rail workloads (ancestor groups 84 and 1485, sc2ts
100k; all faster, path classes unchanged), and the dead ends and
remaining ideas, each estimated under 10%.

Dead end recorded here (not committed): reusing each listed node's own
ancestor across sites, stamped and recomputed when stale, gave 1.00-1.02
of the previous build under the map, since checking the stored ancestor
costs a record load, as the one-step walk it replaces does.
---
 benchmarking/README.md | 108 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 108 insertions(+)

diff --git a/benchmarking/README.md b/benchmarking/README.md
index 33665e0..3540cfa 100644
--- a/benchmarking/README.md
+++ b/benchmarking/README.md
@@ -745,3 +745,111 @@ time in proportion to all 3.4M nodes, so most trees aren't numbered.
 Numbering 16 times more eagerly made more sites sparse (6,490 and 9,780
 for queries 11 and 22) but was slower (35 and 40 s), so numbering whole
 trees is the limit at this scale.
+
+## Walk caches: sample matching under a recombination map (October 2026)
+
+Under the HapMap II chr20 map, Viterbi on a chr20p sample took about
+4.1 s against 2.2 s with tsinfer's constant r, because values from
+different sites' switch odds rarely coalesce, so the likelihood list is
+about 8 times longer (`experiments/coarsening`). This round assessed
+algorithmic strategies, with sample matching under the map as the main
+target, and stopped when every remaining idea's estimated gain was under
+10%. Each step's commit records its measurements.
+
+### Measuring
+
+- `profile_samples.py` runs tsls on chr20p samples under a map or
+  constant r, printing times and a hash of each path.
+- `ab.py` compares builds of the extension (each a directory holding
+  `_tsls*.so`, put first on `PYTHONPATH`): each round runs every build in
+  its own process, pinned to one core, in turn, keeping the best time per
+  query over the rounds, and checks that the paths are identical.
+- Where the time goes was found with `sudo perf record` and
+  `perf annotate -l` on `tsls_hmm_find_path_matrix`, into which the site
+  update, coalescing and walks are all inlined, so source lines rather
+  than functions give the split. Step counters (walk lengths, edges per
+  site, how often caches could survive) were added temporarily and not
+  committed.
+
+### Where the time went
+
+At the start, the time was in walks up the tree, not in the values:
+
+| | map | constant r |
+|---|---|---|
+| allele walks up to a mutation or root, and clearing their marks | ~30% | ~27% |
+| coalescing walks to the inherited value, and cache writes | ~24% | ~22% |
+| walks resetting the caches | ~14% | ~15% |
+| arithmetic on the values | <1% | <1% |
+
+(Shares of `tsls_hmm_find_path_matrix`, which was about 80% of the run
+under the map and 63% under constant r; `tsls_hmm_update_tree` was most
+of the rest.) There were about 7.2 edges out and 7.2 in per site.
+
+### Changes
+
+| commit | change | map | constant r |
+|---|---|---|---|
+| `009f587` | clear the caches with a stamp, not reset walks | 0.93 | 0.98 |
+| `a1d957f` | get alleles, and coalesce, through each listed node's nearest listed ancestor | 0.78 | 0.86 |
+| `f00bcc0` | mark listed nodes, so walks read only node records | 0.86 | 0.92 |
+| `39fd01e` | keep the tree as parent links and child counts | 0.96 | 0.93 |
+| `7914bc2` | keep the walk caches across sites | 0.72 | 0.70 |
+
+Ratios are paired wall times over the previous build (6 samples, best of
+3). The two main ideas:
+
+- **The list's own tree.** Every node with a mutation at a site is in
+  the likelihood list, so a listed node's allele is that of its nearest
+  listed ancestor. One walk per listed node, to its nearest listed
+  ancestor, serves both the alleles and coalescing, and these walks are
+  much shorter than walks to a mutation or root.
+- **Persistent walk caches.** Each walk caches its answer, a node ID, on
+  the unlisted nodes it passes. The answer stays right until a node on a
+  cached path gets a likelihood or a new parent, which happened at 24% of
+  sites under the map and 14% under constant r. A node that loses its
+  likelihood becomes a cached node pointing to its old listed ancestor.
+  `tsls_hmm_check_state` checks that every cached walk leads to the
+  nearest listed ancestor.
+
+Overall, against the build before these changes, with identical paths:
+
+| workload | time |
+|---|---|
+| samples, map, mu = 1e-20 | 0.45-0.47 |
+| samples, map, mu = 1e-8 | 0.46 |
+| samples, constant r | 0.50-0.52 |
+
+`paired.py` (tsls over tsinfer-hmm; the path classes are unchanged):
+
+| workload | before | after |
+|---|---|---|
+| ancestors, group 84 | 0.773 | 0.674 |
+| ancestors, group 1485 | 0.714 | 0.528 |
+| sc2ts 100k, k = 1 | 0.020 | 0.018 |
+| sc2ts 100k, exact | 0.013 | 0.012 |
+
+A sample now takes about 2.0-2.3 s under the map and 0.8-0.9 s under
+constant r.
+
+### Dead ends and remaining ideas
+
+- **Reusing listed nodes' own ancestors between sites** (stamped, and
+  recomputed when stale): 1.00-1.02. Checking a stored ancestor costs a
+  record load, as the one-step walk it replaces does.
+- **Batching queries** in one sweep: only the tree position stepping and
+  the structural part of the edge updates are shared (about 7% of the run
+  after the changes above, and at most about 15% before its own
+  overheads), since each query's list, walks and caches are its own.
+  Not pursued.
+- **Remaining profile** (map, shares of the run): walks to listed
+  ancestors ~22%, tree updates ~13%, tree position stepping ~7%, the
+  site update's arithmetic ~7%, the traceback's walk up from the path
+  node to find the matched allele ~6%, storing the traceback ~5%, the
+  allele walk over the list ~4%, inserting mutation nodes ~4%. The
+  candidates each have a ceiling under 10%: precomputing tree transitions
+  in the index (~5-7%), bounding the traceback's allele walk by the
+  mutation's time (exact, since parents are older than children; ~6%),
+  and fusing traceback storage into the site update (~5%).
+- Approximations were not needed; earlier thresholding and coarsening
+  work found them costly in accuracy for little time.

From 0bdfb81c281c01d6478e3a2c9436cd4d2ae08019 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Thu, 29 Oct 2026 11:16:11 +0000
Subject: [PATCH 7/7] Time samples under the HapMap map in paired.py, and
 record results

paired.py gains samples_map and samples_map_mu1e-8 workloads (6 chr20p
samples with per-site rho from the HapMap II chr20 map). Two full rounds
at 9b0112b, tsls / tsinfer-hmm: ancestors 0.39-0.67, samples 0.30,
samples under the map 0.16-0.17, sc2ts 0.013-0.044. Paths identical on
all tsinfer workloads; sc2ts path classes unchanged.
---
 benchmarking/README.md | 25 +++++++++++++++++++++++++
 benchmarking/paired.py | 32 ++++++++++++++++++++++++++------
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/benchmarking/README.md b/benchmarking/README.md
index 3540cfa..9f817d5 100644
--- a/benchmarking/README.md
+++ b/benchmarking/README.md
@@ -853,3 +853,28 @@ constant r.
   and fusing traceback storage into the site update (~5%).
 - Approximations were not needed; earlier thresholding and coarsening
   work found them costly in accuracy for little time.
+
+### Against tsinfer-hmm after the walk caches
+
+Two full rounds of `paired.py` (`9b0112b`, pinned to one core with
+`taskset`, best of 3 per query), tsls / tsinfer-hmm. `samples_map` and
+`samples_map_mu1e-8` are the 6 samples under the HapMap II chr20 map.
+
+| workload | queries | round 1 | round 2 | memory (MiB), tsls / tsinfer-hmm | paths same/tie/better/worse |
+|---|---|---|---|---|---|
+| ancestors, group 84 | 50 | 0.64 | 0.67 | 4 / 64 | 50/0/0/0 |
+| ancestors, group 1000 | 50 | 0.51 | 0.52 | 4 / 64 | 50/0/0/0 |
+| ancestors, group 1483 | 100 | 0.46 | 0.49 | 4 / 64 | 100/0/0/0 |
+| ancestors, group 1484 | 50 | 0.39 | 0.47 | 8 / 64 | 50/0/0/0 |
+| ancestors, group 1485 | 100 | 0.51 | 0.50 | 4 / 64 | 100/0/0/0 |
+| samples, constant r | 6 | 0.30 | 0.30 | 12 / 128 | 6/0/0/0 |
+| samples, map, mu = 1e-20 | 6 | 0.16 | 0.17 | 36 / 384 | 6/0/0/0 |
+| samples, map, mu = 1e-8 | 6 | 0.17 | 0.16 | 40 / 448 | 6/0/0/0 |
+| sc2ts 100k, k = 1 | 10 | 0.018 | 0.018 | 4 / 640 | 4/2/3/1 |
+| sc2ts 100k, exact | 10 | 0.013 | 0.013 | 4 / 960 | 8/2/0/0 |
+| sc2ts 300k, exact | 4 | 0.044 | 0.043 | 4 / 1,792 | 3/1/0/0 |
+
+A chr20p sample takes about 0.8 s with tsls against 2.7 s with
+tsinfer-hmm under constant r, and 1.7-1.9 s against 10-12 s under the
+map. The one worse sc2ts k = 1 path is the tie in floored values
+described in `c81ac3b`.
diff --git a/benchmarking/paired.py b/benchmarking/paired.py
index e8e5aac..bea3fa8 100644
--- a/benchmarking/paired.py
+++ b/benchmarking/paired.py
@@ -11,7 +11,9 @@
 There are two kinds of workload:
 
 - tsinfer: the chr20p ancestor groups and samples made by tsinfer_data.py,
-  with the model in their model.json.
+  with the model in their model.json. The "samples_map" workloads replace
+  its constant rho with per-site probabilities from the HapMap II chr20 map
+  (genetic_map.py), and "_mu1e-8" also replaces its mu.
 - sc2ts: queries from data/queries_scaling matched against the sc2ts
   sub-ARGs, with the sc2ts model (num_mismatches = 3, weight_by_n=False).
   sc2ts matches every sample first with a likelihood threshold of k = 0
@@ -32,6 +34,7 @@
 
 import click
 import numpy as np
+import tsls
 
 import callgrind
 import engines
@@ -53,11 +56,18 @@ class Workload:
     # None for tsinfer workloads, which use their model.json
     mismatch_threshold: float = None
     sc2ts: bool = False
+    # A genetic map (relative to the data directory) and mu to use in place
+    # of model.json's rho and mu
+    genetic_map: str = None
+    mu: float = None
 
 
-def tsinfer_workload(name, num_queries):
-    path = f"tgp_chr20p/{name}"
-    return Workload(name, f"{path}/reference.trees", path, num_queries)
+GENETIC_MAP = "genetic_maps/chr20_HapMapII_GRCh38.txt"
+
+
+def tsinfer_workload(name, num_queries, directory=None, **kwargs):
+    path = f"tgp_chr20p/{name if directory is None else directory}"
+    return Workload(name, f"{path}/reference.trees", path, num_queries, **kwargs)
 
 
 def sc2ts_workload(name, num_samples, num_queries, mismatch_threshold=None):
@@ -78,6 +88,10 @@ def sc2ts_workload(name, num_samples, num_queries, mismatch_threshold=None):
     tsinfer_workload("ancestors_g1484", 50),
     tsinfer_workload("ancestors_g1485", 100),
     tsinfer_workload("samples", 6),
+    tsinfer_workload("samples_map", 6, "samples", genetic_map=GENETIC_MAP),
+    tsinfer_workload(
+        "samples_map_mu1e-8", 6, "samples", genetic_map=GENETIC_MAP, mu=1e-8
+    ),
     sc2ts_workload("sc2ts_n100k_k1", 100000, 10, mismatch_threshold=1),
     sc2ts_workload("sc2ts_n100k", 100000, 10),
     sc2ts_workload("sc2ts_n300k", 300000, 4),
@@ -107,6 +121,12 @@ def load_workload(data_dir, workload):
         query_ids = SC2TS_QUERIES[: workload.num_queries]
     else:
         model = queries.load_model(data_dir / workload.query_dir)
+        if workload.genetic_map is not None:
+            model["rho"] = tsls.model.recombination_probabilities(
+                data_dir / workload.genetic_map, ts.sites_position
+            )
+        if workload.mu is not None:
+            model["mu"] = workload.mu
         query_ids = list(df.query_id[: workload.num_queries])
     return ts, H, query_ids, model
 
@@ -130,8 +150,8 @@ def time_workload(data_dir, workload, repeats):
     n = ts.num_nodes if model["weight_by_n"] else 1
     num_alleles = engines.site_num_alleles(ts, model["alleles"])
     rates = {
-        "rho": np.full(ts.num_sites, model["rho"]),
-        "mu": np.full(ts.num_sites, model["mu"]),
+        "rho": np.broadcast_to(model["rho"], ts.num_sites),
+        "mu": np.broadcast_to(model["mu"], ts.num_sites),
     }
     best = {name: np.full(len(query_ids), np.inf) for name in ENGINES}
     max_memory = {name: 0 for name in ENGINES}