Implement improved mrca.

c/tests/test_trees.c

@@ -4813,6 +4813,35 @@ test_internal_sample_diff_iter(void)
     CU_ASSERT_EQUAL(ret, 0);
 }
 
+static void
+test_multiroot_mrca(void)
+{
+    int ret;
+    tsk_treeseq_t ts;
+    tsk_tree_t tree;
+    tsk_id_t mrca;
+
+    tsk_treeseq_from_text(&ts, 10, multiroot_ex_nodes, multiroot_ex_edges, NULL, NULL,
+        NULL, NULL, NULL, 0);
+    ret = tsk_tree_init(&tree, &ts, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    ret = tsk_tree_first(&tree);
+    CU_ASSERT_EQUAL_FATAL(ret, 1);
+
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_mrca(&tree, 0, 0, &mrca), 0);
+    CU_ASSERT_EQUAL(mrca, 0);
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_mrca(&tree, 0, 1, &mrca), 0);
+    CU_ASSERT_EQUAL(mrca, 10);
+    /* MRCA of two nodes in different subtrees is TSK_NULL */
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_mrca(&tree, 0, 2, &mrca), 0);
+    CU_ASSERT_EQUAL(mrca, TSK_NULL);
+    CU_ASSERT_EQUAL_FATAL(tsk_tree_get_mrca(&tree, 2, 0, &mrca), 0);
+    CU_ASSERT_EQUAL(mrca, TSK_NULL);
+
+    tsk_tree_free(&tree);
+    tsk_treeseq_free(&ts);
+}
+
 static void
 test_multiroot_diff_iter(void)
 {
@@ -6235,12 +6264,15 @@ main(int argc, char **argv)
         { "test_gappy_multi_tree", test_gappy_multi_tree },
         { "test_tsk_treeseq_bad_records", test_tsk_treeseq_bad_records },
 
+        /* multiroot tests */
+        { "test_multiroot_mrca", test_multiroot_mrca },
+        { "test_multiroot_diff_iter", test_multiroot_diff_iter },
+
         /* Diff iter tests */
         { "test_simple_diff_iter", test_simple_diff_iter },
         { "test_nonbinary_diff_iter", test_nonbinary_diff_iter },
         { "test_unary_diff_iter", test_unary_diff_iter },
         { "test_internal_sample_diff_iter", test_internal_sample_diff_iter },
-        { "test_multiroot_diff_iter", test_multiroot_diff_iter },
         { "test_empty_diff_iter", test_empty_diff_iter },
 
         /* Sample sets */

c/tskit/trees.c

@@ -3490,20 +3490,9 @@ int TSK_WARN_UNUSED
 tsk_tree_get_mrca(const tsk_tree_t *self, tsk_id_t u, tsk_id_t v, tsk_id_t *mrca)
 {
     int ret = 0;
-    tsk_id_t w = 0;
-    tsk_id_t *s1 = NULL;
-    tsk_id_t *s2 = NULL;
-    tsk_id_t j;
-    int l1, l2;
-
-    /* We have a NULL value in these stacks, so need to be sure that
-     * we have an extra space allocated */
-    s1 = malloc((1 + self->num_nodes) * sizeof(*s1));
-    s2 = malloc((1 + self->num_nodes) * sizeof(*s2));
-    if (s1 == NULL || s2 == NULL) {
-        ret = TSK_ERR_NO_MEMORY;
-        goto out;
-    }
+    double tu, tv;
+    const tsk_id_t *restrict parent = self->parent;
+    const double *restrict time = self->tree_sequence->tables->nodes.time;
 
     ret = tsk_tree_check_node(self, u);
     if (ret != 0) {
@@ -3513,34 +3502,26 @@ tsk_tree_get_mrca(const tsk_tree_t *self, tsk_id_t u, tsk_id_t v, tsk_id_t *mrca
     if (ret != 0) {
         goto out;
     }
-    j = u;
-    l1 = 0;
-    while (j != TSK_NULL) {
-        tsk_bug_assert(l1 < (int) self->num_nodes);
-        s1[l1] = j;
-        l1++;
-        j = self->parent[j];
-    }
-    s1[l1] = TSK_NULL;
-    j = v;
-    l2 = 0;
-    while (j != TSK_NULL) {
-        tsk_bug_assert(l2 < (int) self->num_nodes);
-        s2[l2] = j;
-        l2++;
-        j = self->parent[j];
-    }
-    s2[l2] = TSK_NULL;
-    do {
-        w = s1[l1];
-        l1--;
-        l2--;
-    } while (l1 >= 0 && l2 >= 0 && s1[l1] == s2[l2]);
-    *mrca = w;
-    ret = 0;
+
+    tu = time[u];
+    tv = time[v];
+    while (u != v) {
+        if (tu < tv) {
+            u = parent[u];
+            if (u == TSK_NULL) {
+                break;
+            }
+            tu = time[u];
+        } else {
+            v = parent[v];
+            if (v == TSK_NULL) {
+                break;
+            }
+            tv = time[v];
+        }
+    }
+    *mrca = u == v ? u : TSK_NULL;
 out:
-    tsk_safe_free(s1);
-    tsk_safe_free(s2);
     return ret;
 }