diff --git a/c/tskit/trees.c b/c/tskit/trees.c
index e5286b6123..3d4f54c766 100644
--- a/c/tskit/trees.c
+++ b/c/tskit/trees.c
@@ -2205,16 +2205,15 @@ get_allele_samples(const tsk_site_t *site, const tsk_bit_array_t *state,
 }
 
 static int
-get_mutation_samples(const tsk_tree_t *self, const tsk_size_t *site_offsets,
+get_mutation_samples(const tsk_tree_t *tree, const tsk_size_t *site_offsets,
     tsk_bit_array_t *allele_samples, tsk_size_t **num_alleles)
 {
     int ret = 0;
 
-    const tsk_treeseq_t *restrict ts = self->tree_sequence;
+    const tsk_treeseq_t *restrict ts = tree->tree_sequence;
 
     const tsk_size_t num_samples = tsk_treeseq_get_num_samples(ts);
-    const tsk_size_t num_tree_sites = ts->tree_sites_length[self->index];
-    const tsk_site_t *restrict tree_sites = ts->tree_sites[self->index];
+    const tsk_site_t *restrict tree_sites = ts->tree_sites[tree->index];
     const tsk_flags_t *restrict flags = ts->tables->nodes.flags;
     const tsk_id_t *restrict mut_nodes = ts->tables->mutations.node;
 
@@ -2224,11 +2223,11 @@ get_mutation_samples(const tsk_tree_t *self, const tsk_size_t *site_offsets,
     tsk_size_t num_mutations, mut_samples_offset, mut_offset;
     tsk_size_t s, m, n;
 
-    tsk_memset(&mut_samples, 0, sizeof(mut_samples));
-
     tsk_id_t node;
     tsk_size_t num_nodes;
-    tsk_id_t *nodes = tsk_malloc(tsk_tree_get_size_bound(self) * sizeof(*nodes));
+    tsk_id_t *nodes = tsk_malloc(tsk_tree_get_size_bound(tree) * sizeof(*nodes));
+
+    tsk_memset(&mut_samples, 0, sizeof(mut_samples));
 
     if (nodes == NULL) {
         ret = TSK_ERR_NO_MEMORY;
@@ -2237,9 +2236,8 @@ get_mutation_samples(const tsk_tree_t *self, const tsk_size_t *site_offsets,
 
     // Get the number of mutations
     num_mutations = 0;
-    for (s = 0; s < num_tree_sites; s++) {
-        site = &tree_sites[s];
-        num_mutations += site->mutations_length;
+    for (s = 0; s < tree->sites_length; s++) {
+        num_mutations += tree_sites[s].mutations_length;
     }
 
     ret = tsk_bit_array_init(&mut_samples, num_samples, num_mutations);
@@ -2248,7 +2246,7 @@ get_mutation_samples(const tsk_tree_t *self, const tsk_size_t *site_offsets,
     }
 
     // Current position minus num ancestral alleles to this point, if there are any sites
-    mut_offset = num_tree_sites
+    mut_offset = tree->sites_length
                      ? (site_offsets[tree_sites[0].id] - (tsk_size_t) tree_sites[0].id)
                      : 0;
 
@@ -2259,7 +2257,7 @@ get_mutation_samples(const tsk_tree_t *self, const tsk_size_t *site_offsets,
     for (m = 0; m < num_mutations; m++) {
         tsk_bit_array_get_row(&mut_samples, m, &mut_samples_row);
 
-        ret = tsk_tree_preorder_from(self, mut_nodes[m + mut_offset], nodes, &num_nodes);
+        ret = tsk_tree_preorder_from(tree, mut_nodes[m + mut_offset], nodes, &num_nodes);
         if (ret != 0) {
             goto out;
         }
@@ -2273,7 +2271,7 @@ get_mutation_samples(const tsk_tree_t *self, const tsk_size_t *site_offsets,
     }
 
     mut_samples_offset = 0;
-    for (s = 0; s < num_tree_sites; s++) {
+    for (s = 0; s < tree->sites_length; s++) {
         site = &tree_sites[s];
 
         tsk_bit_array_get_row(&mut_samples, mut_samples_offset, &mut_samples_row);
@@ -2351,11 +2349,6 @@ compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state,
     tsk_bit_array_t ss_row;
     tsk_bit_array_t ss_A_samples, ss_B_samples, ss_AB_samples, AB_samples;
 
-    tsk_memset(&ss_A_samples, 0, sizeof(ss_A_samples));
-    tsk_memset(&ss_B_samples, 0, sizeof(ss_B_samples));
-    tsk_memset(&ss_AB_samples, 0, sizeof(ss_AB_samples));
-    tsk_memset(&AB_samples, 0, sizeof(AB_samples));
-
     // Sample sets and b sites are rows, a sites are columns
     //       b1           b2           b3
     // a1   [s1, s2, s3] [s1, s2, s3] [s1, s2, s3]
@@ -2373,6 +2366,11 @@ compute_general_two_site_stat_result(const tsk_bit_array_t *site_a_state,
     double *norm = tsk_malloc(state_dim * sizeof(*norm));
     double *result_tmp = tsk_malloc(row_len * num_a_alleles * sizeof(*result_tmp));
 
+    tsk_memset(&ss_A_samples, 0, sizeof(ss_A_samples));
+    tsk_memset(&ss_B_samples, 0, sizeof(ss_B_samples));
+    tsk_memset(&ss_AB_samples, 0, sizeof(ss_AB_samples));
+    tsk_memset(&AB_samples, 0, sizeof(AB_samples));
+
     if (weights == NULL || norm == NULL || result_tmp == NULL) {
         ret = TSK_ERR_NO_MEMORY;
         goto out;
@@ -2484,15 +2482,15 @@ tsk_treeseq_two_site_count_stat(const tsk_treeseq_t *self, tsk_size_t state_dim,
     tsk_size_t site_a, site_b;
     bool polarised = false;
 
-    tsk_memset(&all_samples_bits, 0, sizeof(all_samples_bits));
-    tsk_memset(&allele_samples, 0, sizeof(allele_samples));
-
     const tsk_size_t num_sites = self->tables->sites.num_rows;
     const tsk_size_t num_samples = self->num_samples;
     const tsk_size_t max_alleles = self->tables->mutations.num_rows + num_sites;
     tsk_size_t *restrict site_offsets = tsk_malloc(num_sites * sizeof(*site_offsets));
     tsk_size_t *restrict num_alleles = tsk_malloc(num_sites * sizeof(*num_alleles));
 
+    tsk_memset(&all_samples_bits, 0, sizeof(all_samples_bits));
+    tsk_memset(&allele_samples, 0, sizeof(allele_samples));
+
     if (site_offsets == NULL || num_alleles == NULL) {
         ret = TSK_ERR_NO_MEMORY;
         goto out;
@@ -2620,7 +2618,6 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl
     int ret = 0;
 
     tsk_bit_array_t sample_sets_bits;
-    tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits));
 
     bool stat_site = !!(options & TSK_STAT_SITE);
     bool stat_branch = !!(options & TSK_STAT_BRANCH);
@@ -2632,6 +2629,8 @@ tsk_treeseq_two_locus_count_stat(const tsk_treeseq_t *self, tsk_size_t num_sampl
         .sample_set_sizes = sample_set_sizes,
         .set_indexes = set_indexes };
 
+    tsk_memset(&sample_sets_bits, 0, sizeof(sample_sets_bits));
+
     ret = tsk_treeseq_check_sample_sets(
         self, num_sample_sets, sample_set_sizes, sample_sets);
     if (ret != 0) {