Add tsk_tree_size_bound to enable safe allocations

Closes #1725

Author: jeromekelleher

c/tskit/trees.c

@@ -3626,7 +3626,6 @@ tsk_tree_get_time(const tsk_tree_t *self, tsk_id_t u, double *t)
     if (u == self->virtual_root) {
         *t = INFINITY;
     } else {
-
         ret = tsk_treeseq_get_node(self->tree_sequence, u, &node);
         if (ret != 0) {
             goto out;
@@ -4259,16 +4258,48 @@ tsk_tree_clear(tsk_tree_t *self)
     return ret;
 }
 
+tsk_size_t
+tsk_tree_get_size_bound(const tsk_tree_t *self)
+{
+    tsk_size_t bound = 0;
+
+    if (self->tree_sequence != NULL) {
+        /* This is a safe upper bound which can be computed cheaply.
+         * We have at most n roots and each edge adds at most one new
+         * node to the tree. We also allow space for the virtual root,
+         * to simplify client code.
+         *
+         * In the common case of a binary tree with a single root, we have
+         * 2n - 1 nodes in total, and 2n - 2 edges. Therefore, we return
+         * 3n - 1, which is an over-estimate of 1/2 and we allocate
+         * 1.5 times as much memory as we need.
+         *
+         * Since tracking the exact number of nodes in the tree would require
+         * storing the number of nodes beneath every node and complicate
+         * the tree transition method, this seems like a good compromise
+         * and will result in less memory usage overall in nearly all cases.
+         */
+        bound = 1 + self->tree_sequence->num_samples + self->num_edges;
+    }
+    return bound;
+}
+
 /* Traversal orders */
 
+static tsk_id_t *
+tsk_tree_alloc_node_stack(const tsk_tree_t *self)
+{
+    return tsk_malloc(tsk_tree_get_size_bound(self) * sizeof(tsk_id_t));
+}
+
 int
 tsk_tree_preorder(
     const tsk_tree_t *self, tsk_id_t root, tsk_id_t *nodes, tsk_size_t *num_nodes_ret)
 {
     int ret = 0;
     const tsk_id_t *restrict right_child = self->right_child;
     const tsk_id_t *restrict left_sib = self->left_sib;
-    tsk_id_t *restrict stack = tsk_malloc((self->num_nodes + 1) * sizeof(*stack));
+    tsk_id_t *restrict stack = tsk_tree_alloc_node_stack(self);
     tsk_size_t num_nodes = 0;
     tsk_id_t u, v;
     int stack_top;
@@ -4320,7 +4351,7 @@ tsk_tree_postorder(
     const tsk_id_t *restrict right_child = self->right_child;
     const tsk_id_t *restrict left_sib = self->left_sib;
     const tsk_id_t *restrict parent = self->parent;
-    tsk_id_t *restrict stack = tsk_malloc((self->num_nodes + 1) * sizeof(*stack));
+    tsk_id_t *restrict stack = tsk_tree_alloc_node_stack(self);
     tsk_size_t num_nodes = 0;
     tsk_id_t u, v, postorder_parent;
     int stack_top;

c/tskit/trees.h

@@ -411,6 +411,31 @@ int tsk_tree_prev(tsk_tree_t *self);
 int tsk_tree_clear(tsk_tree_t *self);
 
 void tsk_tree_print_state(const tsk_tree_t *self, FILE *out);
+
+/**
+@brief Return an upper bound on the number of nodes reachable
+    from the roots of this tree.
+
+@rst
+This function provides an upper bound on the number of nodes that
+can be reached in tree traversals, and is intended to be used
+for memory allocation purposes. If ``num_nodes`` is the number
+of nodes visited in a tree traversal from the virtual root
+(e.g., ``tsk_tree_preorder(tree, tree->virtual_root, nodes,
+&num_nodes)``), the bound ``N`` returned here is guaranteed to
+be greater than or equal to ``num_nodes``.
+
+.. warning:: The precise value returned is not defined and should
+    not be depended on, as it may change from version-to-version.
+
+@endrst
+
+@param self A pointer to a tsk_tree_t object.
+@return An upper bound on the number nodes reachable from the roots
+    of this tree, or zero if this tree has not been initialised.
+*/
+tsk_size_t tsk_tree_get_size_bound(const tsk_tree_t *self);
+
 /** @} */
 
 int tsk_tree_set_root_threshold(tsk_tree_t *self, tsk_size_t root_threshold);

python/_tskitmodule.c

@@ -10070,7 +10070,7 @@ Tree_get_traversal_array(Tree *self, PyObject *args, tsk_traversal_func *func)
     if (!PyArg_ParseTuple(args, "i", &root)) {
         goto out;
     }
-    data = PyDataMem_NEW((self->tree->num_nodes + 1) * sizeof(*data));
+    data = PyDataMem_NEW(tsk_tree_get_size_bound(self->tree) * sizeof(*data));
     if (data == NULL) {
         ret = PyErr_NoMemory();
         goto out;

python/tests/test_highlevel.py

@@ -445,6 +445,24 @@ def get_samples(ts, time=None, population=None):
 
 
 class TestTreeTraversals:
+    def test_bad_traversal_order(self):
+        ts = msprime.sim_ancestry(2, random_seed=234)
+        tree = ts.first()
+        for bad_order in ["pre", "post", "preorderorder", ("x",), b"preorder"]:
+            with pytest.raises(ValueError, match="Traversal order"):
+                tree.nodes(order=bad_order)
+
+    @pytest.mark.parametrize("order", list(traversal_map.keys()))
+    def test_returned_types(self, order):
+        ts = msprime.sim_ancestry(2, random_seed=234)
+        tree = ts.first()
+        iterator = tree.nodes(order=order)
+        assert isinstance(iterator, collections.abc.Iterable)
+        lst = list(iterator)
+        assert len(lst) > 0
+        for u in lst:
+            assert isinstance(u, int)
+
     @pytest.mark.parametrize("ts", get_example_tree_sequences())
     @pytest.mark.parametrize("order", list(traversal_map.keys()))
     def test_traversals_virtual_root(self, ts, order):

python/tests/test_lowlevel.py

@@ -2714,7 +2714,7 @@ def test_cleared_tree(self):
         def check_tree(tree):
             assert tree.get_index() == -1
             assert tree.get_left_child(tree.get_virtual_root()) == samples[0]
-            assert tree.get_num_edge() == 0
+            assert tree.get_num_edges() == 0
             assert tree.get_mrca(0, 1) == _tskit.NULL
             for u in range(ts.get_num_nodes()):
                 assert tree.get_parent(u) == _tskit.NULL

python/tskit/trees.py

@@ -2127,9 +2127,11 @@ def timedesc(self, u=NULL):
         return self.timeasc(u)[::-1]
 
     def _preorder_traversal(self, root):
+        # Return Python integers for compatibility
         return map(int, self.preorder(root))
 
     def _postorder_traversal(self, root):
+        # Return Python integers for compatibility
         return map(int, self.postorder(root))
 
     def _inorder_traversal(self, root):
@@ -2166,13 +2168,13 @@ def _timeasc_traversal(self, root):
         """
         Sorts by increasing time but falls back to increasing ID for equal times.
         """
-        yield from self.timeasc(root)
+        return map(int, self.timeasc(root))
 
     def _timedesc_traversal(self, root):
         """
         The reverse of timeasc.
         """
-        yield from self.timedesc(root)
+        return map(int, self.timedesc(root))
 
     def _minlex_postorder_traversal(self, root):
         """