From 3b8ee27f74a0b65f1dfa9620ef5e5243a718658a Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Sat, 8 Jul 2023 19:49:40 +0100
Subject: [PATCH] Test for correct edge ordering

---
 ...tree_iters.py => test_tree_positioning.py} | 60 +++++++++++++++----
 1 file changed, 50 insertions(+), 10 deletions(-)
 rename python/tests/{test_tree_iters.py => test_tree_positioning.py} (81%)

diff --git a/python/tests/test_tree_iters.py b/python/tests/test_tree_positioning.py
similarity index 81%
rename from python/tests/test_tree_iters.py
rename to python/tests/test_tree_positioning.py
index cb0fa6145d..82e628d74d 100644
--- a/python/tests/test_tree_iters.py
+++ b/python/tests/test_tree_positioning.py
@@ -179,23 +179,47 @@ def prev(self):
 
 def check_iters_forward(ts):
     alg_t_output = tsutil.algorithm_T(ts)
-
+    lib_tree = tskit.Tree(ts)
     tree_pos = TreePosition(ts)
+    sample_count = np.zeros(ts.num_nodes, dtype=int)
+    sample_count[ts.samples()] = 1
     parent1 = [-1 for _ in range(ts.num_nodes)]
+    i = 0
+    lib_tree.next()
     while tree_pos.next():
+        out_times = []
         for j in range(*tree_pos.out_range):
             e = tree_pos.out_order[j]
             c = ts.edges_child[e]
+            p = ts.edges_parent[e]
+            out_times.append(ts.nodes_time[p])
             parent1[c] = -1
+        in_times = []
         for j in range(*tree_pos.in_range):
             e = tree_pos.in_order[j]
             c = ts.edges_child[e]
             p = ts.edges_parent[e]
+            in_times.append(ts.nodes_time[p])
             parent1[c] = p
+        # We must visit the edges in *increasing* time order on the way in,
+        # and *decreasing* order on the way out. Otherwise we get quadratic
+        # behaviour for algorithms that need to propagate changes up to the
+        # root.
+        assert out_times == sorted(out_times, reverse=True)
+        assert in_times == sorted(in_times)
 
         interval, parent2 = next(alg_t_output)
         assert list(interval) == tree_pos.interval
         assert parent1 == parent2
+
+        assert lib_tree.index == i
+        assert list(lib_tree.interval) == list(interval)
+        assert list(lib_tree.parent_array[:-1]) == parent1
+
+        lib_tree.next()
+        i += 1
+    assert i == ts.num_trees
+    assert lib_tree.index == -1
     assert next(alg_t_output, None) is None
 
 
@@ -205,32 +229,48 @@ def check_iters_back(ts):
     ]
     i = len(alg_t_output) - 1
 
-    tree = tskit.Tree(ts)
-    while tree.prev():
-        assert alg_t_output[i][0] == list(tree.interval)
-        assert alg_t_output[i][1] == list(tree.parent_array[:-1])
-        i -= 1
-    assert i == -1
-
-    i = len(alg_t_output) - 1
-
+    lib_tree = tskit.Tree(ts)
     tree_pos = TreePosition(ts)
     parent1 = [-1 for _ in range(ts.num_nodes)]
+
+    lib_tree.last()
+
     while tree_pos.prev():
+        # print(tree_pos.out_range)
+        out_times = []
         for j in range(*tree_pos.out_range, -1):
             e = tree_pos.out_order[j]
             c = ts.edges_child[e]
+            p = ts.edges_parent[e]
+            out_times.append(ts.nodes_time[p])
             parent1[c] = -1
+        in_times = []
         for j in range(*tree_pos.in_range, -1):
             e = tree_pos.in_order[j]
             c = ts.edges_child[e]
             p = ts.edges_parent[e]
+            in_times.append(ts.nodes_time[p])
             parent1[c] = p
 
+        # We must visit the edges in *increasing* time order on the way in,
+        # and *decreasing* order on the way out. Otherwise we get quadratic
+        # behaviour for algorithms that need to propagate changes up to the
+        # root.
+        assert out_times == sorted(out_times, reverse=True)
+        assert in_times == sorted(in_times)
+
         interval, parent2 = alg_t_output[i]
         assert list(interval) == tree_pos.interval
         assert parent1 == parent2
+
+        assert lib_tree.index == i
+        assert list(lib_tree.interval) == interval
+        assert list(lib_tree.parent_array[:-1]) == parent1
+
+        lib_tree.prev()
         i -= 1
+
+    assert lib_tree.index == -1
     assert i == -1