1 files changed, 53 insertions, 20 deletions

diff --git a/src/swarm/tree.rs b/src/swarm/tree.rs
index d362f74..ee46271 100644
--- a/src/swarm/tree.rs
+++ b/src/swarm/tree.rs
@@ -95,13 +95,30 @@ fn distances(state: &AppState, writer: u64) -> Option<(Arc<Graph>, Arc<HashMap<S
 /// others take them as parents. What a member reaches we know of ourselves and our
 /// neighbours, as they said; of others, the graph is taken as right (if we're wrong,
 /// a change goes twice or by announcement).
+#[cfg(test)]
 pub(crate) async fn children(state: &AppState, writer: u64) -> Vec<String> {
-    let Some((graph, distance)) = distances(state, writer) else {
-        return Vec::new();
-    };
+    plan(state, writer)
+        .await
+        .map(|plan| plan.children)
+        .unwrap_or_default()
+}
+
+/// What `writer`'s tree means for us: our children in it, and our neighbours it leaves
+/// out (`uncovered`).
+pub(crate) struct Plan {
+    pub(crate) children: Vec<String>,
+    /// Our neighbours that no member they reach comes before (the tree reaches them only
+    /// by announcement), and those not in the member list's graph (whom no tree has).
+    pub(crate) uncovered: Vec<String>,
+}
+
+/// Our part in `writer`'s tree (`children`, `Plan`); None while the member list isn't
+/// known, or the writer or we aren't in it.
+pub(crate) async fn plan(state: &AppState, writer: u64) -> Option<Plan> {
+    let (graph, distance) = distances(state, writer)?;
     let me = state.local_id();
     if !distance.contains_key(me) {
-        return Vec::new();
+        return None;
     }
     // Whom we and our neighbours reach, by pair key.
     let mut reaches: HashMap<&str, Vec<u64>> = HashMap::new();
@@ -131,15 +148,14 @@ pub(crate) async fn children(state: &AppState, writer: u64) -> Vec<String> {
             .get(member)
             .map(|d| (*d, !closer(member), topology::pair_key(member)))
     };
-    let Some(my_rank) = rank(me) else {
-        return Vec::new();
-    };
+    let my_rank = rank(me)?;
     let mut children = Vec::new();
+    let mut uncovered = Vec::new();
     for member in graph.get(me).into_iter().flatten() {
         let Some(its_rank) = rank(member) else {
             continue;
         };
-        if its_rank <= my_rank {
+        if distance[member.as_str()] == 0 {
             continue;
         }
         let parent = graph
@@ -152,22 +168,39 @@ pub(crate) async fn children(state: &AppState, writer: u64) -> Vec<String> {
                 let (d, far, _) = rank(candidate).expect("ranked");
                 (d, far, topology::pair_score(candidate, member))
             });
-        if parent.is_some_and(|parent| parent == me) && topology::reaches(state, member).await {
-            children.push(member.clone());
+        match parent {
+            None => uncovered.push(member.clone()),
+            Some(parent)
+                if parent == me && its_rank > my_rank && topology::reaches(state, member).await =>
+            {
+                children.push(member.clone());
+            }
+            Some(_) => {}
+        }
+    }
+    // Neighbours the member list's graph doesn't give us (not listed yet, or only
+    // connecting out): no tree has them.
+    for (peer, _) in state.peers_snapshot().await {
+        if topology::is_neighbour(state, &peer)
+            && !graph.get(me).is_some_and(|ours| ours.contains(&peer))
+        {
+            uncovered.push(peer);
         }
     }
-    children
+    Some(Plan {
+        children,
+        uncovered,
+    })
 }
 
 /// We applied a change of `writer`'s that `from` passed on: to our children in its
-/// tree (but `from`), a sync message each, saying whose. Returns them.
-pub(crate) async fn forward(state: &AppState, writer: u64, from: &str) -> Vec<String> {
-    let children: Vec<String> = children(state, writer)
-        .await
-        .into_iter()
-        .filter(|child| child != from)
-        .collect();
-    for child in &children {
+/// tree (but `from`), a sync message each, saying whose. Returns our part in the tree
+/// (`Plan`), None if we have none (announcements then do it all).
+pub(crate) async fn forward(state: &AppState, writer: u64, from: &str) -> Option<Plan> {
+    let mut plan = plan(state, writer).await?;
+    plan.children.retain(|child| child != from);
+    let children = &plan.children;
+    for child in children {
         if let Some(peer_state) = state.get_peer_state(child).await {
             peer_state.lock().await.tree_writer = Some(writer);
         }
@@ -181,5 +214,5 @@ pub(crate) async fn forward(state: &AppState, writer: u64, from: &str) -> Vec<St
             "Passing a change on along the writer's tree"
         );
     }
-    children
+    Some(plan)
 }