C adjustments

Author: petrelharp

c/tests/test_stats.c

@@ -1277,6 +1277,7 @@ verify_afs(tsk_treeseq_t *ts)
     int ret;
     tsk_size_t n = tsk_treeseq_get_num_samples(ts);
     tsk_size_t sample_set_sizes[2];
+    double time_windows[] = { 0, 1 };
     const tsk_id_t *samples = tsk_treeseq_get_samples(ts);
     double *result = tsk_malloc(n * n * sizeof(*result));
 
@@ -1305,6 +1306,10 @@ verify_afs(tsk_treeseq_t *ts)
         NULL, 0, NULL, TSK_STAT_BRANCH | TSK_STAT_SPAN_NORMALISE, result);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
 
+    ret = tsk_treeseq_allele_frequency_spectrum(ts, 2, sample_set_sizes, samples, 0,
+        NULL, 1, time_windows, TSK_STAT_BRANCH | TSK_STAT_SPAN_NORMALISE, result);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
     free(result);
 }
 

c/tskit/trees.c

@@ -3539,30 +3539,23 @@ tsk_treeseq_site_allele_frequency_spectrum(const tsk_treeseq_t *self,
     return ret;
 }
 
-static int TSK_WARN_UNUSED
+static void
 tsk_treeseq_update_branch_afs(const tsk_treeseq_t *self, tsk_id_t u, double right,
     double *restrict last_update, const double *restrict time, tsk_id_t *restrict parent,
-    const double *counts, tsk_size_t num_sample_sets, tsk_size_t num_time_windows,
-    const double *time_windows, tsk_size_t window_index, const tsk_size_t *result_dims,
-    tsk_flags_t options, double *result)
+    tsk_size_t *restrict coordinate, const double *counts, tsk_size_t num_sample_sets,
+    tsk_size_t num_time_windows, const double *time_windows, tsk_size_t window_index,
+    const tsk_size_t *result_dims, tsk_flags_t options, double *result)
 {
-    int ret = 0;
     tsk_size_t afs_size;
     tsk_size_t k;
     tsk_size_t time_window_index;
     double *afs;
-    // note: moving this malloc outside this function doesn't speed things up
-    tsk_size_t *coordinate = tsk_malloc(num_sample_sets * sizeof(*coordinate));
     bool polarised = !!(options & TSK_STAT_POLARISED);
     const double *count_row = GET_2D_ROW(counts, num_sample_sets + 1, u);
     double x = 0;
     double t_u, t_v;
     double tw_branch_length = 0;
     const tsk_size_t all_samples = (tsk_size_t) count_row[num_sample_sets];
-    if (coordinate == NULL) {
-        ret = tsk_trace_error(TSK_ERR_NO_MEMORY);
-        goto out;
-    }
     if (parent[u] != TSK_NULL) {
         t_u = time[u];
         t_v = time[parent[u]];
@@ -3589,9 +3582,6 @@ tsk_treeseq_update_branch_afs(const tsk_treeseq_t *self, tsk_id_t u, double righ
         }
     }
     last_update[u] = right;
-out:
-    tsk_safe_free(coordinate);
-    return ret;
 }
 
 static int
@@ -3616,6 +3606,7 @@ tsk_treeseq_branch_allele_frequency_spectrum(const tsk_treeseq_t *self,
     tsk_id_t *restrict parent = tsk_malloc(num_nodes * sizeof(*parent));
     double *restrict last_update = tsk_calloc(num_nodes, sizeof(*last_update));
     double *restrict branch_length = tsk_calloc(num_nodes, sizeof(*branch_length));
+    tsk_size_t *restrict coordinate = tsk_malloc(num_sample_sets * sizeof(*coordinate));
     tsk_id_t tj, tk, h;
     double t_left, t_right, w_right;
     const tsk_size_t K = num_sample_sets + 1;
@@ -3625,7 +3616,7 @@ tsk_treeseq_branch_allele_frequency_spectrum(const tsk_treeseq_t *self,
         goto out;
     }
 
-    if (parent == NULL || last_update == NULL) {
+    if (parent == NULL || last_update == NULL || coordinate == NULL) {
         ret = tsk_trace_error(TSK_ERR_NO_MEMORY);
         goto out;
     }
@@ -3643,19 +3634,13 @@ tsk_treeseq_branch_allele_frequency_spectrum(const tsk_treeseq_t *self,
             tk++;
             u = edge_child[h];
             v = edge_parent[h];
-            ret = tsk_treeseq_update_branch_afs(self, u, t_left, last_update, node_time,
-                parent, counts, num_sample_sets, num_time_windows, time_windows,
-                window_index, result_dims, options, result);
-            if (ret != 0) {
-                goto out;
-            }
+            tsk_treeseq_update_branch_afs(self, u, t_left, last_update, node_time,
+                parent, coordinate, counts, num_sample_sets, num_time_windows,
+                time_windows, window_index, result_dims, options, result);
             while (v != TSK_NULL) {
-                ret = tsk_treeseq_update_branch_afs(self, v, t_left, last_update,
-                    node_time, parent, counts, num_sample_sets, num_time_windows,
+                tsk_treeseq_update_branch_afs(self, v, t_left, last_update, node_time,
+                    parent, coordinate, counts, num_sample_sets, num_time_windows,
                     time_windows, window_index, result_dims, options, result);
-                if (ret != 0) {
-                    goto out;
-                }
                 update_state(counts, K, v, u, -1);
                 v = parent[v];
             }
@@ -3671,12 +3656,9 @@ tsk_treeseq_branch_allele_frequency_spectrum(const tsk_treeseq_t *self,
             parent[u] = v;
             branch_length[u] = node_time[v] - node_time[u];
             while (v != TSK_NULL) {
-                ret = tsk_treeseq_update_branch_afs(self, v, t_left, last_update,
-                    node_time, parent, counts, num_sample_sets, num_time_windows,
+                tsk_treeseq_update_branch_afs(self, v, t_left, last_update, node_time,
+                    parent, coordinate, counts, num_sample_sets, num_time_windows,
                     time_windows, window_index, result_dims, options, result);
-                if (ret != 0) {
-                    goto out;
-                }
                 update_state(counts, K, v, u, +1);
                 v = parent[v];
             }
@@ -3695,12 +3677,9 @@ tsk_treeseq_branch_allele_frequency_spectrum(const tsk_treeseq_t *self,
             /* Flush the contributions of all nodes to the current window */
             for (u = 0; u < (tsk_id_t) num_nodes; u++) {
                 tsk_bug_assert(last_update[u] < w_right);
-                ret = tsk_treeseq_update_branch_afs(self, u, w_right, last_update,
-                    node_time, parent, counts, num_sample_sets, num_time_windows,
+                tsk_treeseq_update_branch_afs(self, u, w_right, last_update, node_time,
+                    parent, coordinate, counts, num_sample_sets, num_time_windows,
                     time_windows, window_index, result_dims, options, result);
-                if (ret != 0) {
-                    goto out;
-                }
             }
             window_index++;
         }