Merge pull request #2308 from jeromekelleher/beta-segfault

Malloc Q rather than from stack in multi-merger
tskit-dev · Aug 5, 2024 · 7492287 · 7492287
2 parents 657fc8d + 6ac820a
commit 7492287
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 2 deletions.
diff --git a/lib/msprime.c b/lib/msprime.c
@@ -8010,7 +8010,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;
  lineage_t *x_lin, *y_lin;
@@ -8024,6 +8025,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);
@@ -8080,6 +8086,7 @@ msp_dirac_common_ancestor_event(msp_t *self, population_id_t pop_id, label_id_t
  }
  }
 out:
+ tsk_safe_free(Q);
  return ret;
 }
 
@@ -8229,7 +8236,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;
@@ -8241,6 +8249,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);
@@ -8300,6 +8313,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
@@ -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)
 {
@@ -4305,6 +4346,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 },