Skip to content

Implement efficient seeking from non-null trees using tree_pos #2911

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Aug 13, 2025
Merged

Implement efficient seeking from non-null trees using tree_pos #2911

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
dc3f3c4
Add tree seek_forward and seek_backward and revise seek linear
duncanMR Feb 23, 2024
f02781f
Restore seek_linear and add seek_skip as an option
duncanMR Mar 5, 2024
002b006
Refine API, add tests and fix implementation error
jeromekelleher Aug 6, 2025
c79fbd8
Update c/tskit/trees.h
jeromekelleher Aug 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 67 additions & 34 deletions c/tests/test_trees.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,7 +331,7 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
uint32_t mutation_index, site_index;
tsk_size_t k, l, tree_sites_length;
const tsk_site_t *sites = NULL;
tsk_tree_t tree;
tsk_tree_t tree, skip_tree;
tsk_size_t num_edges;
tsk_size_t num_nodes = tsk_treeseq_get_num_nodes(ts);
tsk_size_t num_sites = tsk_treeseq_get_num_sites(ts);
Expand All @@ -340,6 +340,7 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)

ret = tsk_tree_init(&tree, ts, 0);
CU_ASSERT_EQUAL(ret, 0);
ret = tsk_tree_init(&skip_tree, ts, 0);
CU_ASSERT_EQUAL(tsk_treeseq_get_num_trees(ts), num_trees);

CU_ASSERT_EQUAL(tree.index, -1);
Expand Down Expand Up @@ -372,6 +373,21 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
}
site_index++;
}
/* Check the skip tree */
ret = tsk_tree_first(&skip_tree);
CU_ASSERT_EQUAL(ret, TSK_TREE_OK);
ret = tsk_tree_seek(&skip_tree, breakpoints[j], TSK_SEEK_SKIP);
CU_ASSERT_EQUAL(ret, 0);
/* Calling print_state here also verifies the integrity of the tree */
tsk_tree_print_state(&skip_tree, _devnull);
check_trees_equal(&tree, &skip_tree);
ret = tsk_tree_last(&skip_tree);
CU_ASSERT_EQUAL(ret, TSK_TREE_OK);
ret = tsk_tree_seek(&skip_tree, breakpoints[j], TSK_SEEK_SKIP);
CU_ASSERT_EQUAL(ret, 0);
tsk_tree_print_state(&skip_tree, _devnull);
check_trees_equal(&tree, &skip_tree);

j++;
}
CU_ASSERT_EQUAL(ret, 0);
Expand All @@ -381,6 +397,7 @@ verify_trees(tsk_treeseq_t *ts, tsk_size_t num_trees, tsk_id_t *parents)
CU_ASSERT_EQUAL(tsk_treeseq_get_sequence_length(ts), breakpoints[j]);

tsk_tree_free(&tree);
tsk_tree_free(&skip_tree);

verify_tree_pos(ts, num_trees, parents);
}
Expand Down Expand Up @@ -811,7 +828,8 @@ typedef struct {
} sample_count_test_t;

static void
verify_sample_counts(tsk_treeseq_t *ts, tsk_size_t num_tests, sample_count_test_t *tests)
verify_sample_counts(tsk_treeseq_t *ts, tsk_size_t num_tests, sample_count_test_t *tests,
tsk_flags_t seek_options)
{
int ret;
tsk_size_t j, num_samples, n, k;
Expand All @@ -826,13 +844,9 @@ verify_sample_counts(tsk_treeseq_t *ts, tsk_size_t num_tests, sample_count_test_

ret = tsk_tree_init(&tree, ts, TSK_NO_SAMPLE_COUNTS);
CU_ASSERT_EQUAL(ret, 0);
ret = tsk_tree_first(&tree);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
for (j = 0; j < num_tests; j++) {
while (tree.index < tests[j].tree_index) {
ret = tsk_tree_next(&tree);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
}
ret = tsk_tree_seek_index(&tree, tests[j].tree_index, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = tsk_tree_get_num_samples(&tree, tests[j].node, &num_samples);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(tests[j].count, num_samples);
Expand All @@ -850,10 +864,8 @@ verify_sample_counts(tsk_treeseq_t *ts, tsk_size_t num_tests, sample_count_test_
ret = tsk_tree_first(&tree);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
for (j = 0; j < num_tests; j++) {
while (tree.index < tests[j].tree_index) {
ret = tsk_tree_next(&tree);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
}
ret = tsk_tree_seek_index(&tree, tests[j].tree_index, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = tsk_tree_get_num_samples(&tree, tests[j].node, &num_samples);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(tests[j].count, num_samples);
Expand All @@ -872,10 +884,8 @@ verify_sample_counts(tsk_treeseq_t *ts, tsk_size_t num_tests, sample_count_test_
ret = tsk_tree_first(&tree);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
for (j = 0; j < num_tests; j++) {
while (tree.index < tests[j].tree_index) {
ret = tsk_tree_next(&tree);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
}
ret = tsk_tree_seek_index(&tree, tests[j].tree_index, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = tsk_tree_get_num_samples(&tree, tests[j].node, &num_samples);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(tests[j].count, num_samples);
Expand Down Expand Up @@ -908,10 +918,8 @@ verify_sample_counts(tsk_treeseq_t *ts, tsk_size_t num_tests, sample_count_test_
ret = tsk_tree_first(&tree);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
for (j = 0; j < num_tests; j++) {
while (tree.index < tests[j].tree_index) {
ret = tsk_tree_next(&tree);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
}
ret = tsk_tree_seek_index(&tree, tests[j].tree_index, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = tsk_tree_get_num_samples(&tree, tests[j].node, &num_samples);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL(tests[j].count, num_samples);
Expand Down Expand Up @@ -1008,6 +1016,7 @@ verify_sample_sets(tsk_treeseq_t *ts)
{
int ret;
tsk_tree_t t;
tsk_id_t j;

ret = tsk_tree_init(&t, ts, TSK_SAMPLE_LISTS);
CU_ASSERT_EQUAL(ret, 0);
Expand All @@ -1021,6 +1030,20 @@ verify_sample_sets(tsk_treeseq_t *ts)
}
CU_ASSERT_EQUAL_FATAL(ret, 0);

for (j = 0; j < (tsk_id_t) tsk_treeseq_get_num_trees(ts); j++) {
ret = tsk_tree_first(&t);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
ret = tsk_tree_seek_index(&t, j, TSK_SEEK_SKIP);
CU_ASSERT_EQUAL_FATAL(ret, 0);
verify_sample_sets_for_tree(&t);

ret = tsk_tree_last(&t);
CU_ASSERT_EQUAL_FATAL(ret, TSK_TREE_OK);
ret = tsk_tree_seek_index(&t, j, TSK_SEEK_SKIP);
CU_ASSERT_EQUAL_FATAL(ret, 0);
verify_sample_sets_for_tree(&t);
}

tsk_tree_free(&t);
}

Expand Down Expand Up @@ -5870,7 +5893,8 @@ test_simple_sample_sets(void)

tsk_treeseq_from_text(&ts, 10, paper_ex_nodes, paper_ex_edges, NULL, NULL, NULL,
paper_ex_individuals, NULL, 0);
verify_sample_counts(&ts, num_tests, tests);
verify_sample_counts(&ts, num_tests, tests, 0);
verify_sample_counts(&ts, num_tests, tests, TSK_SEEK_SKIP);
verify_sample_sets(&ts);

tsk_treeseq_free(&ts);
Expand All @@ -5889,7 +5913,8 @@ test_nonbinary_sample_sets(void)

tsk_treeseq_from_text(&ts, 100, nonbinary_ex_nodes, nonbinary_ex_edges, NULL, NULL,
NULL, NULL, NULL, 0);
verify_sample_counts(&ts, num_tests, tests);
verify_sample_counts(&ts, num_tests, tests, 0);
verify_sample_counts(&ts, num_tests, tests, TSK_SEEK_SKIP);
verify_sample_sets(&ts);

tsk_treeseq_free(&ts);
Expand All @@ -5909,7 +5934,8 @@ test_internal_sample_sample_sets(void)

tsk_treeseq_from_text(&ts, 10, internal_sample_ex_nodes, internal_sample_ex_edges,
NULL, NULL, NULL, NULL, NULL, 0);
verify_sample_counts(&ts, num_tests, tests);
verify_sample_counts(&ts, num_tests, tests, 0);
verify_sample_counts(&ts, num_tests, tests, TSK_SEEK_SKIP);
verify_sample_sets(&ts);

tsk_treeseq_free(&ts);
Expand Down Expand Up @@ -6283,7 +6309,7 @@ test_multiroot_tree_traversal(void)
}

static void
test_seek_multi_tree(void)
verify_seek_multi_tree(tsk_flags_t seek_options)
{
int ret;
tsk_treeseq_t ts;
Expand All @@ -6299,29 +6325,29 @@ test_seek_multi_tree(void)
CU_ASSERT_EQUAL_FATAL(ret, 0);

for (j = 0; j < num_trees; j++) {
ret = tsk_tree_seek(&t, breakpoints[j], 0);
ret = tsk_tree_seek(&t, breakpoints[j], seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, j);
ret = tsk_tree_seek_index(&t, j, 0);
ret = tsk_tree_seek_index(&t, j, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, j);
for (k = 0; k < num_trees; k++) {
ret = tsk_tree_seek(&t, breakpoints[k], 0);
ret = tsk_tree_seek(&t, breakpoints[k], seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, k);
ret = tsk_tree_seek_index(&t, k, 0);
ret = tsk_tree_seek_index(&t, k, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, k);
}
}

ret = tsk_tree_seek(&t, 1.99999, 0);
ret = tsk_tree_seek(&t, 1.99999, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, 0);
ret = tsk_tree_seek(&t, 6.99999, 0);
ret = tsk_tree_seek(&t, 6.99999, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, 1);
ret = tsk_tree_seek(&t, 9.99999, 0);
ret = tsk_tree_seek(&t, 9.99999, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, 2);

Expand All @@ -6331,7 +6357,7 @@ test_seek_multi_tree(void)
for (j = 0; j < num_trees; j++) {
ret = tsk_tree_init(&t, &ts, 0);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = tsk_tree_seek(&t, breakpoints[j], 0);
ret = tsk_tree_seek(&t, breakpoints[j], seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, j);
tsk_tree_free(&t);
Expand All @@ -6341,12 +6367,12 @@ test_seek_multi_tree(void)
ret = tsk_tree_init(&t, &ts, 0);
CU_ASSERT_EQUAL_FATAL(ret, 0);
for (j = 0; j < num_trees; j++) {
ret = tsk_tree_seek(&t, 0, 0);
ret = tsk_tree_seek(&t, 0, seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
ret = tsk_tree_prev(&t);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, -1);
ret = tsk_tree_seek(&t, breakpoints[j], 0);
ret = tsk_tree_seek(&t, breakpoints[j], seek_options);
CU_ASSERT_EQUAL_FATAL(ret, 0);
CU_ASSERT_EQUAL_FATAL(t.index, j);
}
Expand All @@ -6355,6 +6381,13 @@ test_seek_multi_tree(void)
tsk_treeseq_free(&ts);
}

static void
test_seek_multi_tree(void)
{
verify_seek_multi_tree(0);
verify_seek_multi_tree(TSK_SEEK_SKIP);
}

static void
test_seek_errors(void)
{
Expand Down
Loading
Loading