1 files changed, 136 insertions, 10 deletions
diff --git a/src/swarm/tree.rs b/src/swarm/tree.rs
index c6baff9..015d62c 100644
--- a/src/swarm/tree.rs
+++ b/src/swarm/tree.rs
@@ -174,6 +174,105 @@ pub(crate) fn whole(
     )
 }
 
+/// Whom we and our neighbours in `graph` reach, by pair key, as far as we know (our
+/// own, and what our neighbours said).
+async fn known_reaches<'g>(state: &AppState, graph: &'g Graph) -> HashMap<&'g str, Vec<u64>> {
+    let me = state.local_id();
+    let mut reaches: HashMap<&str, Vec<u64>> = HashMap::new();
+    if let Some((me, _)) = graph.get_key_value(me) {
+        reaches.insert(me.as_str(), reach::published(state));
+    }
+    for member in graph.get(me).into_iter().flatten() {
+        if let Some(peer_state) = state.get_peer_state(member).await
+            && let Some(keys) = reach::of(peer_state.lock().await.reaches.as_ref())
+        {
+            reaches.insert(member.as_str(), keys);
+        }
+    }
+    reaches
+}
+
+/// Our own tree, as our writes go along it (`Plan`, as every member computes its part).
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct Shape {
+    /// The other members in it.
+    pub(crate) members: usize,
+    /// The most hops from us to one of them.
+    pub(crate) depth: u32,
+    /// The hops to each, summed (for the mean).
+    pub(crate) hops: u64,
+    /// Members it doesn't reach: without a parent there (no member they reach comes
+    /// before them), below one without, or not connected to us in the member list's
+    /// graph. They get our writes by announcement.
+    pub(crate) off: usize,
+}
+
+/// The shape of our own tree; None while the member list isn't known, or we aren't in
+/// it. What members other than us and our neighbours reach is taken from the graph, as
+/// in `plan`.
+pub(crate) async fn own_shape(state: &AppState) -> Option<Shape> {
+    let me = state.local_id();
+    let (graph, distance) = distances(state, topology::pair_key(me))?;
+    let reaches = known_reaches(state, &graph).await;
+    let reached = |by: &str, member: &str| {
+        reaches
+            .get(by)
+            .is_none_or(|keys| keys.contains(&topology::pair_key(member)))
+    };
+    Some(shape(&graph, &distance, &reached, me))
+}
+
+/// `own_shape` over `graph`, from `me` at distance 0.
+fn shape(
+    graph: &Graph,
+    distance: &HashMap<String, u32>,
+    reached: &dyn Fn(&str, &str) -> bool,
+    me: &str,
+) -> Shape {
+    let parents: HashMap<&str, Option<&str>> = graph
+        .keys()
+        .filter(|member| distance.get(member.as_str()).is_some_and(|d| *d > 0))
+        .map(|member| (member.as_str(), parent(graph, distance, reached, member)))
+        .collect();
+    // Hops from us along parents; None where the chain ends without us. Parents come
+    // before their children, so chains end.
+    let mut hops: HashMap<&str, Option<u32>> = HashMap::from([(me, Some(0))]);
+    fn along<'g>(
+        member: &'g str,
+        parents: &HashMap<&'g str, Option<&'g str>>,
+        hops: &mut HashMap<&'g str, Option<u32>>,
+    ) -> Option<u32> {
+        if let Some(known) = hops.get(member) {
+            return *known;
+        }
+        let found = parents
+            .get(member)
+            .copied()
+            .flatten()
+            .and_then(|parent| along(parent, parents, hops))
+            .map(|hops| hops + 1);
+        hops.insert(member, found);
+        found
+    }
+    let mut shape = Shape {
+        members: 0,
+        depth: 0,
+        hops: 0,
+        off: graph.len() - 1 - parents.len(),
+    };
+    for member in parents.keys() {
+        match along(member, &parents, &mut hops) {
+            Some(h) => {
+                shape.members += 1;
+                shape.depth = shape.depth.max(h);
+                shape.hops += u64::from(h);
+            }
+            None => shape.off += 1,
+        }
+    }
+    shape
+}
+
 /// What `writer`'s tree means for us: our children in it, and our neighbours it leaves
 /// out (`uncovered`).
 pub(crate) struct Plan {
@@ -195,16 +294,7 @@ pub(crate) async fn plan(state: &AppState, writer: u64) -> Option<Plan> {
     if !distance.contains_key(me) {
         return None;
     }
-    // Whom we and our neighbours reach, by pair key.
-    let mut reaches: HashMap<&str, Vec<u64>> = HashMap::new();
-    reaches.insert(me, reach::published(state));
-    for member in graph.get(me).into_iter().flatten() {
-        if let Some(peer_state) = state.get_peer_state(member).await
-            && let Some(keys) = reach::of(peer_state.lock().await.reaches.as_ref())
-        {
-            reaches.insert(member.as_str(), keys);
-        }
-    }
+    let reaches = known_reaches(state, &graph).await;
     let reached = |by: &str, member: &str| {
         reaches
             .get(by)
@@ -277,3 +367,39 @@ pub(crate) async fn forward(state: &AppState, writer: u64, from: &str) -> Option
     }
     Some(plan)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Hops along parents from us; members without a parent, below one without, or
+    /// not connected count as off the tree.
+    #[test]
+    fn our_tree_in_numbers() {
+        let link = |graph: &mut Graph, a: &str, b: &str| {
+            graph.entry(a.into()).or_default().insert(b.into());
+            graph.entry(b.into()).or_default().insert(a.into());
+        };
+        let mut graph = Graph::new();
+        link(&mut graph, "me", "a");
+        link(&mut graph, "a", "b");
+        link(&mut graph, "me", "c");
+        link(&mut graph, "c", "e");
+        graph.insert("d".into(), BTreeSet::new());
+        let distance: HashMap<String, u32> = [("me", 0), ("a", 1), ("c", 1), ("b", 2), ("e", 2)]
+            .into_iter()
+            .map(|(member, d)| (member.to_owned(), d))
+            .collect();
+        // c doesn't reach us: no parent, nor e below it.
+        let reached = |by: &str, member: &str| !(by == "c" && member == "me");
+        assert_eq!(
+            shape(&graph, &distance, &reached, "me"),
+            Shape {
+                members: 2,
+                depth: 2,
+                hops: 3,
+                off: 3,
+            }
+        );
+    }
+}