From 0100ea7736481e0d9d350e86e4cde159f6c8feb7 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Fri, 26 Jul 2024 15:11:07 +0100
Subject: [PATCH] Other basic models passing in C

---
 lib/msprime.c | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/lib/msprime.c b/lib/msprime.c
index ff8f2d37f..ec2ed4a71 100644
--- a/lib/msprime.c
+++ b/lib/msprime.c
@@ -3025,6 +3025,7 @@ msp_dtwf_recombine(
                 ret = MSP_ERR_NO_MEMORY;
                 goto out;
             }
+            z->lineage = x->lineage;
             msp_set_segment_mass(self, z);
             tsk_bug_assert(z->left < z->right);
             if (x->next != NULL) {
@@ -3068,6 +3069,9 @@ msp_dtwf_recombine(
 
     for (j = 0; j < MSP_MAX_PED_PLOIDY; j++) {
         y = *rec_heads[j];
+        if (y == x_head) {
+            lineage_reset_segments(y->lineage);
+        }
         if (y != x_head && y != NULL) {
             lin = msp_alloc_lineage(self, y, label);
             if (lin == NULL) {
@@ -4083,10 +4087,6 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
                     ret = MSP_ERR_NO_MEMORY;
                     goto out;
                 }
-                ret = msp_insert_individual(self, new_lineage);
-                if (ret != 0) {
-                    goto out;
-                }
             } else {
                 if ((self->additional_nodes & MSP_NODE_IS_CA_EVENT)
                     || (!self->coalescing_segments_only && coalescence)) {
@@ -4099,10 +4099,21 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
                 z->next = alpha;
             }
             alpha->prev = z;
+            alpha->lineage = new_lineage;
             msp_set_segment_mass(self, alpha);
             z = alpha;
         }
     }
+    if (new_lineage != NULL) {
+        ret = msp_insert_individual(self, new_lineage);
+        if (ret != 0) {
+            goto out;
+        }
+        /* FIXME see note above about avoiding this by exausting
+         * the original chains */
+        lineage_reset_segments(new_lineage);
+    }
+
     if (coalescence) {
         if (!self->coalescing_segments_only) {
             ret = msp_store_arg_edges(self, z, new_node_id);
@@ -4129,6 +4140,7 @@ msp_merge_ancestors(msp_t *self, avl_tree_t *Q, population_id_t population_id,
             goto out;
         }
     }
+
     if (ret_merged_head != NULL) {
         *ret_merged_head = merged_head;
     }