Make simplify thread-safe in no-filter case

jeromekelleher · jeromekelleher · commit 4f49c76adc75 · 2022-12-14T16:01:05.000Z
diff --git a/c/tskit/tables.c b/c/tskit/tables.c
@@ -9542,71 +9542,115 @@ simplifier_finalise_site_references(
 }
 
 static int TSK_WARN_UNUSED
-simplifier_finalise_population_references(
-    simplifier_t *self, const bool *population_referenced, tsk_id_t *population_id_map)
+simplifier_finalise_population_references(simplifier_t *self)
 {
     int ret = 0;
-    tsk_id_t ret_id;
     tsk_size_t j;
+    tsk_id_t pop_id, ret_id;
     tsk_population_t pop;
+    tsk_id_t *node_population = self->tables->nodes.population;
+    const tsk_size_t num_nodes = self->tables->nodes.num_rows;
     const tsk_size_t num_populations = self->input_tables.populations.num_rows;
+    bool *population_referenced
+        = tsk_calloc(num_populations, sizeof(*population_referenced));
+    tsk_id_t *population_id_map
+        = tsk_malloc(num_populations * sizeof(*population_id_map));
 
-    if (self->options & TSK_SIMPLIFY_FILTER_POPULATIONS) {
-        for (j = 0; j < num_populations; j++) {
-            tsk_population_table_get_row_unsafe(
-                &self->input_tables.populations, (tsk_id_t) j, &pop);
-            population_id_map[j] = TSK_NULL;
-            if (population_referenced[j]) {
-                ret_id = tsk_population_table_add_row(
-                    &self->tables->populations, pop.metadata, pop.metadata_length);
-                if (ret_id < 0) {
-                    ret = (int) ret_id;
-                    goto out;
-                }
-                population_id_map[j] = ret_id;
+    tsk_bug_assert(self->options & TSK_SIMPLIFY_FILTER_POPULATIONS);
+
+    if (population_referenced == NULL || population_id_map == NULL) {
+        ret = TSK_ERR_NO_MEMORY;
+        goto out;
+    }
+
+    for (j = 0; j < num_nodes; j++) {
+        pop_id = node_population[j];
+        if (pop_id != TSK_NULL) {
+            population_referenced[pop_id] = true;
+        }
+    }
+
+    for (j = 0; j < num_populations; j++) {
+        tsk_population_table_get_row_unsafe(
+            &self->input_tables.populations, (tsk_id_t) j, &pop);
+        population_id_map[j] = TSK_NULL;
+        if (population_referenced[j]) {
+            ret_id = tsk_population_table_add_row(
+                &self->tables->populations, pop.metadata, pop.metadata_length);
+            if (ret_id < 0) {
+                ret = (int) ret_id;
+                goto out;
             }
+            population_id_map[j] = ret_id;
         }
-    } else {
-        tsk_bug_assert(self->tables->populations.num_rows == num_populations);
-        for (j = 0; j < num_populations; j++) {
-            population_id_map[j] = (tsk_id_t) j;
+    }
+
+    /* Remap the IDs in the node table */
+    for (j = 0; j < num_nodes; j++) {
+        pop_id = node_population[j];
+        if (pop_id != TSK_NULL) {
+            node_population[j] = population_id_map[pop_id];
         }
     }
 out:
+    tsk_safe_free(population_id_map);
+    tsk_safe_free(population_referenced);
     return ret;
 }
 
 static int TSK_WARN_UNUSED
-simplifier_finalise_individual_references(
-    simplifier_t *self, const bool *individual_referenced, tsk_id_t *individual_id_map)
+simplifier_finalise_individual_references(simplifier_t *self)
 {
     int ret = 0;
-    tsk_id_t ret_id;
     tsk_size_t j;
+    tsk_id_t pop_id, ret_id;
     tsk_individual_t ind;
-    const tsk_size_t num_individuals = self->input_tables.individuals.num_rows;
+    tsk_id_t *node_individual = self->tables->nodes.individual;
     tsk_id_t *parents;
+    const tsk_size_t num_nodes = self->tables->nodes.num_rows;
+    const tsk_size_t num_individuals = self->input_tables.individuals.num_rows;
+    bool *individual_referenced
+        = tsk_calloc(num_individuals, sizeof(*individual_referenced));
+    tsk_id_t *individual_id_map
+        = tsk_malloc(num_individuals * sizeof(*individual_id_map));
 
-    if (self->options & TSK_SIMPLIFY_FILTER_INDIVIDUALS) {
-        for (j = 0; j < num_individuals; j++) {
-            tsk_individual_table_get_row_unsafe(
-                &self->input_tables.individuals, (tsk_id_t) j, &ind);
-            individual_id_map[j] = TSK_NULL;
-            if (individual_referenced[j]) {
-                ret_id = tsk_individual_table_add_row(&self->tables->individuals,
-                    ind.flags, ind.location, ind.location_length, ind.parents,
-                    ind.parents_length, ind.metadata, ind.metadata_length);
-                if (ret_id < 0) {
-                    ret = (int) ret_id;
-                    goto out;
-                }
-                individual_id_map[j] = ret_id;
+    tsk_bug_assert(self->options & TSK_SIMPLIFY_FILTER_INDIVIDUALS);
+
+    if (individual_referenced == NULL || individual_id_map == NULL) {
+        ret = TSK_ERR_NO_MEMORY;
+        goto out;
+    }
+
+    for (j = 0; j < num_nodes; j++) {
+        pop_id = node_individual[j];
+        if (pop_id != TSK_NULL) {
+            individual_referenced[pop_id] = true;
+        }
+    }
+
+    for (j = 0; j < num_individuals; j++) {
+        tsk_individual_table_get_row_unsafe(
+            &self->input_tables.individuals, (tsk_id_t) j, &ind);
+        individual_id_map[j] = TSK_NULL;
+        if (individual_referenced[j]) {
+            /* Can't remap the parents inline here because we have no
+             * guarantees about sortedness */
+            ret_id = tsk_individual_table_add_row(&self->tables->individuals, ind.flags,
+                ind.location, ind.location_length, ind.parents, ind.parents_length,
+                ind.metadata, ind.metadata_length);
+            if (ret_id < 0) {
+                ret = (int) ret_id;
+                goto out;
             }
+            individual_id_map[j] = ret_id;
         }
-    } else {
-        tsk_bug_assert(self->tables->individuals.num_rows == num_individuals);
-        for (j = 0; j < num_individuals; j++) {
-            individual_id_map[j] = (tsk_id_t) j;
+    }
+
+    /* Remap the IDs in the node table */
+    for (j = 0; j < num_nodes; j++) {
+        pop_id = node_individual[j];
+        if (pop_id != TSK_NULL) {
+            node_individual[j] = individual_id_map[pop_id];
         }
     }
 
@@ -9619,7 +9663,10 @@ simplifier_finalise_individual_references(
             parents[j] = individual_id_map[parents[j]];
         }
     }
+
 out:
+    tsk_safe_free(individual_id_map);
+    tsk_safe_free(individual_referenced);
     return ret;
 }
 
@@ -9687,26 +9734,6 @@ static int TSK_WARN_UNUSED
 simplifier_flush_output(simplifier_t *self)
 {
     int ret = 0;
-    const tsk_size_t num_nodes = self->tables->nodes.num_rows;
-    const tsk_size_t num_populations = self->input_tables.populations.num_rows;
-    const tsk_size_t num_individuals = self->input_tables.individuals.num_rows;
-    tsk_id_t *node_population = self->tables->nodes.population;
-    bool *population_referenced
-        = tsk_calloc(num_populations, sizeof(*population_referenced));
-    tsk_id_t *population_id_map
-        = tsk_malloc(num_populations * sizeof(*population_id_map));
-    tsk_id_t *node_individual = self->tables->nodes.individual;
-    bool *individual_referenced
-        = tsk_calloc(num_individuals, sizeof(*individual_referenced));
-    tsk_id_t *individual_id_map
-        = tsk_malloc(num_individuals * sizeof(*individual_id_map));
-    tsk_size_t j;
-    tsk_id_t pop_id, ind_id;
-
-    if (population_referenced == NULL || population_id_map == NULL
-        || individual_referenced == NULL || individual_id_map == NULL) {
-        goto out;
-    }
 
     /* TODO Migrations fit reasonably neatly into the pattern that we have here. We
      * can consider references to populations from migration objects in the same way
@@ -9727,44 +9754,20 @@ simplifier_flush_output(simplifier_t *self)
         goto out;
     }
 
-    for (j = 0; j < num_nodes; j++) {
-        pop_id = node_population[j];
-        if (pop_id != TSK_NULL) {
-            population_referenced[pop_id] = true;
-        }
-        ind_id = node_individual[j];
-        if (ind_id != TSK_NULL) {
-            individual_referenced[ind_id] = true;
+    if (self->options & TSK_SIMPLIFY_FILTER_POPULATIONS) {
+        ret = simplifier_finalise_population_references(self);
+        if (ret != 0) {
+            goto out;
         }
     }
-
-    ret = simplifier_finalise_population_references(
-        self, population_referenced, population_id_map);
-    if (ret != 0) {
-        goto out;
-    }
-    ret = simplifier_finalise_individual_references(
-        self, individual_referenced, individual_id_map);
-    if (ret != 0) {
-        goto out;
-    }
-
-    /* Remap node IDs referencing the above */
-    for (j = 0; j < num_nodes; j++) {
-        pop_id = node_population[j];
-        if (pop_id != TSK_NULL) {
-            node_population[j] = population_id_map[pop_id];
-        }
-        ind_id = node_individual[j];
-        if (ind_id != TSK_NULL) {
-            node_individual[j] = individual_id_map[ind_id];
+    if (self->options & TSK_SIMPLIFY_FILTER_INDIVIDUALS) {
+        ret = simplifier_finalise_individual_references(self);
+        if (ret != 0) {
+            goto out;
         }
     }
+
 out:
-    tsk_safe_free(population_referenced);
-    tsk_safe_free(population_id_map);
-    tsk_safe_free(individual_referenced);
-    tsk_safe_free(individual_id_map);
     return ret;
 }
 
diff --git a/c/tskit/tables.h b/c/tskit/tables.h
@@ -686,6 +686,9 @@ reference them. */
 #define TSK_SIMPLIFY_FILTER_POPULATIONS (1 << 1)
 /** Remove individuals from the output if there are no nodes that reference them.*/
 #define TSK_SIMPLIFY_FILTER_INDIVIDUALS (1 << 2)
+/** Do not remove nodes from the output if there are no edges that reference
+them. */
+#define TSK_SIMPLIFY_NO_FILTER_NODES (1 << 7)
 /**
 Reduce the topological information in the tables to the minimum necessary to
 represent the trees that contain sites. If there are zero sites this will
@@ -719,7 +722,6 @@ flag). It keeps unary nodes, but only if the unary node is referenced from an in
 this flag is negated compare to other filtering options because the default
 behaviour is remove unreferenced nodes.
 */
-#define TSK_SIMPLIFY_NO_FILTER_NODES (1 << 7)
 /** @} */
 
 /**
@@ -3923,6 +3925,21 @@ input tables, but not in the specified list of sample IDs (if provided). The
 ``node_map[u] == u`` for all nodes. Note also that the order of the list of
 samples is not important in this case.
 
+When table is not filtered (i.e., if the `TSK_SIMPLIFY_NO_FILTER_NODES`
+option is provided or the `TSK_SIMPLIFY_FILTER_SITES`,
+`TSK_SIMPLIFY_FILTER_POPULATIONS` or `TSK_SIMPLIFY_FILTER_INDIVIDUALS`
+options are *not* provided) the corresponding table is modified as
+little as possible, and all pointers are guaranteed to remain valid
+after simplification. The only changes made to an unfiltered table are
+to update any references to tables that may have changed (for example,
+remapping population IDs in the node table if
+`TSK_SIMPLIFY_FILTER_POPULATIONS` was specified).
+
+.. note:: It is possible for populations and individuals to be filtered
+   even if `TSK_SIMPLIFY_NO_FILTER_NODES` is specified because there
+   may be entirely unreferences entities in the input tables, which
+   are not affected by whether we filter nodes or not.
+
 The table collection will always be unindexed after simplify successfully completes.
 
 .. note:: Migrations are currently not supported by simplify, and an error will
@@ -3937,11 +3954,11 @@ Options can be specified by providing one or more of the following bitwise
 - :c:macro:`TSK_SIMPLIFY_FILTER_SITES`
 - :c:macro:`TSK_SIMPLIFY_FILTER_POPULATIONS`
 - :c:macro:`TSK_SIMPLIFY_FILTER_INDIVIDUALS`
+- :c:macro:`TSK_SIMPLIFY_NO_FILTER_NODES`
 - :c:macro:`TSK_SIMPLIFY_REDUCE_TO_SITE_TOPOLOGY`
 - :c:macro:`TSK_SIMPLIFY_KEEP_UNARY`
 - :c:macro:`TSK_SIMPLIFY_KEEP_INPUT_ROOTS`
 - :c:macro:`TSK_SIMPLIFY_KEEP_UNARY_IN_INDIVIDUALS`
-- :c:macro:`TSK_SIMPLIFY_NO_FILTER_NODES`
 @endrst
 
 @param self A pointer to a tsk_table_collection_t object.
diff --git a/python/tests/test_topology.py b/python/tests/test_topology.py
@@ -5655,6 +5655,96 @@ def test_many_trees_recurrent_mutations_internal_samples(self):
                         self.verify_simplify_haplotypes(ts, samples, keep_unary=keep)
 
 
+class TestSimplifyUnreferencedPopulations:
+    def example(self):
+        tables = tskit.TableCollection(1)
+        tables.populations.add_row()
+        tables.populations.add_row()
+        # No references to population 0
+        tables.nodes.add_row(time=0, population=1, flags=1)
+        tables.nodes.add_row(time=0, population=1, flags=1)
+        tables.nodes.add_row(time=1, population=1, flags=0)
+        # Unreference node
+        tables.nodes.add_row(time=1, population=1, flags=0)
+        tables.edges.add_row(0, 1, parent=2, child=0)
+        tables.edges.add_row(0, 1, parent=2, child=1)
+        tables.sort()
+        return tables
+
+    def test_no_filter_populations(self):
+        tables = self.example()
+        tables.simplify(filter_populations=False)
+        assert len(tables.populations) == 2
+        assert len(tables.nodes) == 3
+        assert np.all(tables.nodes.population == 1)
+
+    def test_no_filter_populations_nodes(self):
+        tables = self.example()
+        tables.simplify(filter_populations=False, filter_nodes=False)
+        assert len(tables.populations) == 2
+        assert len(tables.nodes) == 4
+        assert np.all(tables.nodes.population == 1)
+
+    def test_filter_populations_no_filter_nodes(self):
+        tables = self.example()
+        tables.simplify(filter_populations=True, filter_nodes=False)
+        assert len(tables.populations) == 1
+        assert len(tables.nodes) == 4
+        assert np.all(tables.nodes.population == 0)
+
+    def test_remapped_default(self):
+        tables = self.example()
+        tables.simplify()
+        assert len(tables.populations) == 1
+        assert len(tables.nodes) == 3
+        assert np.all(tables.nodes.population == 0)
+
+
+class TestSimplifyUnreferencedIndividuals:
+    def example(self):
+        tables = tskit.TableCollection(1)
+        tables.individuals.add_row()
+        tables.individuals.add_row()
+        # No references to individual 0
+        tables.nodes.add_row(time=0, individual=1, flags=1)
+        tables.nodes.add_row(time=0, individual=1, flags=1)
+        tables.nodes.add_row(time=1, individual=1, flags=0)
+        # Unreference node
+        tables.nodes.add_row(time=1, individual=1, flags=0)
+        tables.edges.add_row(0, 1, parent=2, child=0)
+        tables.edges.add_row(0, 1, parent=2, child=1)
+        tables.sort()
+        return tables
+
+    def test_no_filter_individuals(self):
+        tables = self.example()
+        tables.simplify(filter_individuals=False)
+        assert len(tables.individuals) == 2
+        assert len(tables.nodes) == 3
+        assert np.all(tables.nodes.individual == 1)
+
+    def test_no_filter_individuals_nodes(self):
+        tables = self.example()
+        tables.simplify(filter_individuals=False, filter_nodes=False)
+        assert len(tables.individuals) == 2
+        assert len(tables.nodes) == 4
+        assert np.all(tables.nodes.individual == 1)
+
+    def test_filter_individuals_no_filter_nodes(self):
+        tables = self.example()
+        tables.simplify(filter_individuals=True, filter_nodes=False)
+        assert len(tables.individuals) == 1
+        assert len(tables.nodes) == 4
+        assert np.all(tables.nodes.individual == 0)
+
+    def test_remapped_default(self):
+        tables = self.example()
+        tables.simplify()
+        assert len(tables.individuals) == 1
+        assert len(tables.nodes) == 3
+        assert np.all(tables.nodes.individual == 0)
+
+
 class TestSimplifyKeepInputRoots(SimplifyTestBase, ExampleTopologyMixin):
     """
     Tests for the keep_input_roots option to simplify.
