Fix alloc error

tskit-dev · Aug 4, 2023 · ebb68bc · ebb68bc
1 parent f22f436
commit ebb68bc
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/c/tskit/trees.c b/c/tskit/trees.c
@@ -6679,18 +6679,20 @@ tsk_treeseq_divergence_matrix_site(const tsk_treeseq_t *self, tsk_size_t num_sam
     tsk_mutation_t mut;
     double left, right, span_left, span_right;
     double *restrict D;
+    const tsk_size_t num_nodes = self->tables->nodes.num_rows;
     int8_t *descending_bitset = tsk_malloc(num_samples * sizeof(*descending_bitset));
     tsk_id_t *descending_list = tsk_malloc(num_samples * sizeof(*descending_list));
     tsk_id_t *not_descending_list
         = tsk_malloc(num_samples * sizeof(*not_descending_list));
-    tsk_id_t *stack = NULL;
+    /* Do *not* use tsk_tree_get_size bound here because it gives a per-tree
+     * bound, not a global one! */
+    tsk_id_t *stack = tsk_malloc(num_nodes * sizeof(*stack));
 
     ret = tsk_tree_init(&tree, self, 0);
     if (ret != 0) {
         goto out;
     }
 
-    stack = tsk_malloc(tsk_tree_get_size_bound(&tree) * sizeof(*stack));
     if (descending_bitset == NULL || descending_list == NULL
         || not_descending_list == NULL || stack == NULL) {
         ret = TSK_ERR_NO_MEMORY;