diff --git a/lib/msprime.c b/lib/msprime.c index 474f7ca7e..063a2d61d 100644 --- a/lib/msprime.c +++ b/lib/msprime.c @@ -6824,7 +6824,8 @@ msp_dirac_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t { int ret = 0; uint32_t j, n, num_participants, num_parental_copies; - avl_tree_t *ancestors, Q[4]; /* MSVC won't let us use num_pots here */ + avl_tree_t *ancestors; + avl_tree_t *Q = NULL; avl_node_t *x_node, *y_node; segment_t *x, *y; double nC2, p; @@ -6837,6 +6838,11 @@ msp_dirac_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t } else { num_parental_copies = 2 * self->ploidy; } + Q = tsk_malloc(num_parental_copies * sizeof(*Q)); + if (Q == NULL) { + ret = MSP_ERR_NO_MEMORY; + goto out; + } ancestors = &self->populations[pop_id].ancestors[label]; n = avl_count(ancestors); @@ -6889,6 +6895,7 @@ msp_dirac_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t } } out: + tsk_safe_free(Q); return ret; } @@ -7035,7 +7042,8 @@ msp_beta_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t l { int ret = 0; uint32_t j, n, num_participants, num_parental_copies; - avl_tree_t *ancestors, Q[4]; /* MSVC won't let us use num_pots here */ + avl_tree_t *ancestors; + avl_tree_t *Q = NULL; double alpha = self->model.params.beta_coalescent.alpha; double truncation_point = beta_compute_truncation(self); double beta_x, u, increment; @@ -7047,6 +7055,11 @@ msp_beta_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t l } else { num_parental_copies = 2 * self->ploidy; } + Q = tsk_malloc(num_parental_copies * sizeof(*Q)); + if (Q == NULL) { + ret = MSP_ERR_NO_MEMORY; + goto out; + } for (j = 0; j < num_parental_copies; j++) { avl_init_tree(&Q[j], cmp_segment_queue, NULL); @@ -7106,6 +7119,7 @@ msp_beta_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t l } out: + tsk_safe_free(Q); return ret; } diff --git a/lib/tests/test_ancestry.c b/lib/tests/test_ancestry.c index 9974758a1..26a5a7be4 100644 --- a/lib/tests/test_ancestry.c +++ b/lib/tests/test_ancestry.c @@ -1515,6 +1515,47 @@ test_multiple_mergers_simulation(void) gsl_rng_free(rng); } +static void +test_multiple_mergers_ploidy(void) +{ + int ret; + size_t j, ploidy; + uint32_t n = 10; + long seed = 10; + msp_t msp; + gsl_rng *rng = safe_rng_alloc(); + tsk_table_collection_t tables; + + for (j = 0; j < 2; j++) { + gsl_rng_set(rng, seed); + for (ploidy = 1; ploidy < 12; ploidy++) { + ret = build_sim(&msp, &tables, rng, 10, 1, NULL, n); + CU_ASSERT_EQUAL(ret, 0); + CU_ASSERT_EQUAL_FATAL(msp_set_recombination_rate(&msp, 0.1), 0); + if (j == 0) { + ret = msp_set_simulation_model_dirac(&msp, 0.9, 10); + } else { + ret = msp_set_simulation_model_beta(&msp, 1.8, 1); + } + CU_ASSERT_EQUAL(ret, 0); + + msp_set_ploidy(&msp, ploidy); + ret = msp_initialise(&msp); + + ret = msp_run(&msp, DBL_MAX, ULONG_MAX); + CU_ASSERT_EQUAL_FATAL(ret, 0); + CU_ASSERT_TRUE(msp_is_completed(&msp)); + CU_ASSERT_TRUE(msp.time > 0); + msp_verify(&msp, 0); + + ret = msp_free(&msp); + CU_ASSERT_EQUAL(ret, 0); + tsk_table_collection_free(&tables); + } + } + gsl_rng_free(rng); +} + static void test_multiple_mergers_growth_rate(void) { @@ -3875,6 +3916,7 @@ main(int argc, char **argv) { "test_gc_rates", test_gc_rates }, { "test_multiple_mergers_simulation", test_multiple_mergers_simulation }, + { "test_multiple_mergers_ploidy", test_multiple_mergers_ploidy }, { "test_multiple_mergers_growth_rate", test_multiple_mergers_growth_rate }, { "test_dirac_coalescent_bad_parameters", test_dirac_coalescent_bad_parameters }, { "test_beta_coalescent_bad_parameters", test_beta_coalescent_bad_parameters },