2 files changed, 80 insertions, 11 deletions

diff --git a/src/transport/mod.rs b/src/transport/mod.rs
index aee3ee3..d7e7135 100644
--- a/src/transport/mod.rs
+++ b/src/transport/mod.rs
@@ -12,14 +12,16 @@ pub use crate::document::replica::Epoch;
 pub mod firewall;
 pub mod mycelium;
 
-/// Most heads listed in a keepalive; more are sent as a digest. 16 heads are about
-/// 520 bytes, well within one Mycelium chunk (1300 bytes).
-pub const MAX_LISTED_HEADS: usize = 16;
+/// Most heads listed in a keepalive; more are sent as a digest. A head takes 34 bytes on
+/// the wire, so 32 of them keep a keepalive within one Mycelium chunk (1300 bytes; see
+/// the size test in `mycelium`). There are rarely more heads than nodes that wrote
+/// concurrently since the last merging write.
+pub const MAX_LISTED_HEADS: usize = 32;
 
 /// A document's heads, compactly: the heads themselves while there are few, otherwise a
-/// digest of them. Either way, two summaries are equal exactly when the heads are (the
-/// number of heads decides the form, so equal heads always take the same one).
-#[derive(Debug, Clone, PartialEq, Eq)]
+/// digest of them. Two summaries are equal exactly when the heads are, whatever form
+/// each took (nodes with different limits still agree).
+#[derive(Debug, Clone)]
 pub enum HeadsSummary {
     /// Sorted.
     Heads(Vec<ChangeHash>),
@@ -31,13 +33,18 @@ impl HeadsSummary {
     pub fn of(mut heads: Vec<ChangeHash>) -> Self {
         heads.sort_unstable();
         if heads.len() <= MAX_LISTED_HEADS {
-            return Self::Heads(heads);
+            Self::Heads(heads)
+        } else {
+            Self::Digest(digest_of(&heads))
         }
-        let mut hasher = Sha256::new();
-        for head in &heads {
-            hasher.update(head.0);
+    }
+
+    /// The digest of the heads, computed if they are listed.
+    pub fn digest(&self) -> [u8; 32] {
+        match self {
+            Self::Heads(heads) => digest_of(heads),
+            Self::Digest(digest) => *digest,
         }
-        Self::Digest(hasher.finalize().into())
     }
 
     /// The heads, if listed.
@@ -49,6 +56,26 @@ impl HeadsSummary {
     }
 }
 
+impl PartialEq for HeadsSummary {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::Heads(a), Self::Heads(b)) => a == b,
+            _ => self.digest() == other.digest(),
+        }
+    }
+}
+
+impl Eq for HeadsSummary {}
+
+/// SHA-256 over sorted heads.
+fn digest_of(sorted_heads: &[ChangeHash]) -> [u8; 32] {
+    let mut hasher = Sha256::new();
+    for head in sorted_heads {
+        hasher.update(head.0);
+    }
+    hasher.finalize().into()
+}
+
 /// PING and PONG. Addressed to one peer: `syncs_with_you` is about the recipient.
 #[derive(Debug, Clone)]
 pub struct Keepalive {
diff --git a/src/transport/mycelium.rs b/src/transport/mycelium.rs
index 9d1fb90..2683f82 100644
--- a/src/transport/mycelium.rs
+++ b/src/transport/mycelium.rs
@@ -973,3 +973,45 @@ pub enum TransportError {
         status: reqwest::StatusCode,
     },
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Mycelium splits messages into chunks of this size.
+    const MYCELIUM_CHUNK: usize = 1300;
+
+    #[test]
+    fn largest_keepalive_fits_in_one_chunk() {
+        let keepalive = Keepalive {
+            heads: HeadsSummary::Heads(
+                (0..MAX_LISTED_HEADS as u8)
+                    .map(|n| ChangeHash([n; 32]))
+                    .collect(),
+            ),
+            epoch: Epoch {
+                generation: u64::MAX,
+                id: [0xff; 32],
+            },
+            ready: true,
+            syncs_with_you: true,
+        };
+        let payload = keepalive_to_bytes(&keepalive).unwrap();
+        // Mycelium's message data: a length byte, the topic, then our payload.
+        let size = 1 + TOPIC_PONG.len().max(TOPIC_PING.len()) + payload.len();
+        assert!(
+            size + 50 <= MYCELIUM_CHUNK,
+            "keepalive with {MAX_LISTED_HEADS} heads takes {size} bytes"
+        );
+        assert_eq!(bytes_to_keepalive(&payload).unwrap().heads, keepalive.heads);
+    }
+
+    #[test]
+    fn listed_heads_equal_their_digest() {
+        let heads: Vec<_> = (0..3u8).map(|n| ChangeHash([n; 32])).collect();
+        let listed = HeadsSummary::of(heads.clone());
+        let digest = HeadsSummary::Digest(listed.digest());
+        assert_eq!(listed, digest);
+        assert_ne!(listed, HeadsSummary::of(heads[..2].to_vec()));
+    }
+}