one more test

python/tests/test_extend_edges.py

@@ -461,12 +461,11 @@ def get_simple_ex(self, samples=None):
             (6, 8, 0, 10),
         ]
         tables = tskit.TableCollection(sequence_length=10)
-        nodes = tables.nodes
         if samples is None:
             samples = [0, 1, 2, 3]
         for n, t in node_times.items():
             flags = tskit.NODE_IS_SAMPLE if n in samples else 0
-            nodes.add_row(time=t, flags=flags)
+            tables.nodes.add_row(time=t, flags=flags)
         for p, c, l, r in edges:
             tables.edges.add_row(parent=p, child=c, left=l, right=r)
         ts = tables.tree_sequence()
@@ -524,6 +523,51 @@ def test_internal_samples(self):
         # (and it would be a pain to make it work)
         # self.verify_extend_edges(ts)
 
+    def test_iterative_example(self):
+        # Here is the full tree; extend edges should be able to
+        # recover all unary nodes after simplification:
+        #
+        #       9         9         9          9
+        #     +-+-+    +--+--+  +---+---+  +-+-+--+
+        #     8   |    8     |  8   |   |  8 | |  |
+        #     |   |  +-+-+   |  |   |   |  | | |  |
+        #     7   |  |   7   |  |   7   |  | | |  7
+        #   +-+-+ |  | +-++  |  | +-++  |  | | |  |
+        #   6   | |  | |  6  |  | |  6  |  | | |  6
+        # +-++  | |  | |  |  |  | |  |  |  | | |  |
+        # 1  0  2 3  1 2  0  3  1 2  0  3  1 2 3  0
+        #   +++          +++        +++          +++
+        #   4 5          4 5        4 5          4 5
+        #
+        samples = [0, 1, 2, 3, 4, 5]
+        node_times = [1, 1, 1, 1, 0, 0, 2, 3, 4, 5]
+        # (p, c, l, r)
+        edges = [
+            (0, 4, 0, 10),
+            (0, 5, 0, 10),
+            (6, 0, 0, 10),
+            (6, 1, 0, 3),
+            (7, 2, 0, 7),
+            (7, 6, 0, 10),
+            (8, 1, 3, 10),
+            (8, 7, 0, 5),
+            (9, 2, 7, 10),
+            (9, 3, 0, 10),
+            (9, 7, 5, 10),
+            (9, 8, 0, 10),
+        ]
+        tables = tskit.TableCollection(sequence_length=10)
+        for n, t in enumerate(node_times):
+            flags = tskit.NODE_IS_SAMPLE if n in samples else 0
+            tables.nodes.add_row(time=t, flags=flags)
+        for p, c, l, r in edges:
+            tables.edges.add_row(parent=p, child=c, left=l, right=r)
+        ts = tables.tree_sequence()
+        sts = ts.simplify()
+        assert ts.num_edges == 12
+        assert sts.num_edges == 16
+        tables.assert_equals(sts.extend_edges().tables, ignore_provenance=True)
+
     def test_wright_fisher(self):
         tables = wf.wf_sim(N=5, ngens=20, num_loci=100, deep_history=False, seed=3)
         tables.sort()