Use a dedicated type for the index

Previously, we used `u64`, which is unlikely to ever overflow.
Even at 60 Hz, it would take ~9.7 billion years.

But the code used checked operations, even though they would not help.
I was going to switch to regular operations, but realized that if we
handle overflows properly, we could save some traffic.

By default, integers in postcard are serialized with varint encoding.
> 128 takes 2 bytes, > 16384 takes 3 bytes, and so on. But we can switch
  to `u16` with overflow handling and fixint encoding, which always
  takes 2
  bytes.

This also simplifies the logic inside `batches_after`.
I also had to replace `BTreeMap` with `HashMap`, but we don't need to
preserve the ordering since we iterate over indices, so it's fine.

Author: Shatur

src/lib.rs

@@ -750,7 +750,7 @@ pub mod prelude {
             protocol::{ProtocolHash, ProtocolHasher, ProtocolMismatch},
             replication::{
                 Replicated,
-                diff::{DiffEntityExt, Diffable, PatchHistory, PatchIndex},
+                diff::{DiffEntityExt, Diffable, PatchHistory, patch_index::PatchIndex},
                 receive_markers::AppMarkerExt,
                 registry::rule_fns::RuleFns,
                 rules::{AppRuleExt, component::ReplicationMode},

src/server/replication_messages/mutations.rs

@@ -13,7 +13,7 @@ use crate::{
         backend::channels::ServerChannel,
         replication::{
             client_ticks::{ClientTicks, MutateInfo, MutatedEntityInfo, PatchCursors},
-            diff::PatchIndex,
+            diff::patch_index::PatchIndex,
             mutate_index::MutateIndex,
             registry::{ComponentIndex, component_mask::ComponentMask},
         },

src/server/replication_messages/serialized_data.rs

@@ -6,7 +6,7 @@ use crate::{
     postcard_utils,
     prelude::*,
     shared::replication::{
-        diff::{DiffFns, PatchIndex},
+        diff::{DiffFns, patch_index::PatchIndex},
         registry::{FnsId, ctx::SerializeCtx, serde_fns::SerdeFns},
     },
 };

src/shared/replication/client_ticks.rs

@@ -12,7 +12,7 @@ use super::mutate_index::MutateIndex;
 use crate::{
     prelude::*,
     shared::replication::{
-        diff::PatchIndex,
+        diff::patch_index::PatchIndex,
         registry::{ComponentIndex, component_mask::ComponentMask},
     },
 };
@@ -183,15 +183,17 @@ impl EntityTicks {
     /// Advances the acknowledged diff patch cursor for `component`.
     ///
     /// Mutation ACKs can arrive late relative to newer mutation ACKs. Keep the
-    /// greatest cursor so an older ACK cannot make the sender resend patches
+    /// newest cursor so an older ACK cannot make the sender resend patches
     /// already acknowledged by a newer packet.
     fn set_patch_cursor(&mut self, component: ComponentIndex, cursor: PatchIndex) {
         if let Some((_, existing)) = self
             .patch_cursors
             .iter_mut()
             .find(|(index, _)| *index == component)
         {
-            *existing = (*existing).max(cursor);
+            if cursor.is_newer_than(*existing) {
+                *existing = cursor;
+            }
         } else {
             self.patch_cursors.push((component, cursor));
         }

src/shared/replication/diff.rs

@@ -2,10 +2,12 @@
 //!
 //! See [`Diffable`] for the main user-facing API and example.
 
+pub mod patch_index;
+
 use core::{iter, mem};
 
 use alloc::{
-    collections::{BTreeMap, VecDeque, vec_deque},
+    collections::{VecDeque, vec_deque},
     format,
     vec::Vec,
 };
@@ -16,6 +18,7 @@ use bevy::{
         system::EntityCommands,
         world::EntityWorldMut,
     },
+    platform::collections::HashMap,
     prelude::*,
     ptr::Ptr,
 };
@@ -33,9 +36,7 @@ use crate::{
         },
     },
 };
-
-/// Monotonic index assigned to a sent diff batch.
-pub type PatchIndex = u64;
+use patch_index::PatchIndex;
 
 /// Component whose mutations can be represented as an ordered history of patches.
 ///
@@ -159,9 +160,7 @@ impl<C: Diffable> PatchHistory<C> {
             return self.last_index;
         }
 
-        let index = self
-            .last_index
-            .map_or(0, |last_index| last_index.saturating_add(1));
+        let index = self.last_index.map_or(PatchIndex::new(0), |i| i + 1);
         self.last_index = Some(index);
         self.batches.push_back(mem::take(&mut self.pending));
         self.prune_to_limit();
@@ -178,20 +177,22 @@ impl<C: Diffable> PatchHistory<C> {
             return None;
         }
 
-        let first_index = last_index - (self.batches.len() as PatchIndex - 1);
-        let next_index = cursor + 1;
-        if next_index < first_index {
-            // Cursor is outside of the history window.
+        let missing_count = last_index.distance_after(cursor) as usize;
+        if missing_count == 0 {
+            // Client is already at the latest cursor.
+            // The component was mutated directly.
             return None;
         }
 
-        if cursor >= last_index {
+        if missing_count > self.batches.len() {
+            // Client cursor is outside the history window.
             return None;
         }
 
-        let start = (next_index - first_index) as usize;
+        let start = self.batches.len() - missing_count;
+
         Some(BatchSlice {
-            first_index: next_index,
+            first_index: cursor + 1,
             batches: self.batches.range(start..),
         })
     }
@@ -233,7 +234,7 @@ impl<C: Diffable> Default for PatchHistory<C> {
 #[derive(Component, Debug)]
 pub struct PatchBuffer<C: Diffable> {
     last_applied: Option<PatchIndex>,
-    pending: BTreeMap<PatchIndex, Vec<C::Patch>>,
+    pending: HashMap<PatchIndex, Vec<C::Patch>>,
 }
 
 impl<C: Diffable> PatchBuffer<C> {
@@ -260,29 +261,22 @@ impl<C: Diffable> PatchBuffer<C> {
         batches: Vec<Vec<C::Patch>>,
     ) -> impl Iterator<Item = Vec<C::Patch>> + '_ {
         for (offset, batch) in batches.into_iter().enumerate() {
-            let index = first_index + offset as PatchIndex;
+            let index = first_index + offset as u16;
             if self
                 .last_applied
-                .is_none_or(|last_applied| index > last_applied)
+                .is_none_or(|last_applied| index.is_newer_than(last_applied))
             {
                 self.pending.entry(index).or_insert(batch);
             }
         }
 
         iter::from_fn(move || {
-            let next_index = self.next_patch_index()?;
+            let next_index = self.last_applied.map_or(PatchIndex::new(0), |i| i + 1);
             let batch = self.pending.remove(&next_index)?;
             self.last_applied = Some(next_index);
             Some(batch)
         })
     }
-
-    fn next_patch_index(&self) -> Option<PatchIndex> {
-        match self.last_applied {
-            Some(index) => index.checked_add(1),
-            None => Some(0),
-        }
-    }
 }
 
 impl<C: Diffable> Default for PatchBuffer<C> {
@@ -569,8 +563,8 @@ mod tests {
         history.record(2);
         history.finish_pending();
 
-        let slice = history.batches_after(0).unwrap();
-        assert_eq!(slice.first_index, 1);
+        let slice = history.batches_after(PatchIndex::new(0)).unwrap();
+        assert_eq!(slice.first_index.get(), 1);
         assert_ne!(slice.batches.len(), 0);
     }
 }

src/shared/replication/diff/patch_index.rs

@@ -0,0 +1,66 @@
+use core::ops::{Add, Sub};
+
+use serde::{Deserialize, Serialize};
+
+/// Monotonic index assigned to a sent diff batch.
+///
+/// All operations on it are wrapping.
+#[derive(Debug, Default, Serialize, Deserialize, Eq, PartialEq, Hash, Clone, Copy)]
+pub struct PatchIndex(#[serde(with = "postcard::fixint::le")] u16);
+
+impl PatchIndex {
+    /// The maximum wrapping distance at which an index is considered newer.
+    pub const MAX_NEWER_DISTANCE: u16 = u16::MAX / 2;
+
+    #[inline]
+    pub fn new(value: u16) -> Self {
+        Self(value)
+    }
+
+    #[inline]
+    pub fn get(self) -> u16 {
+        self.0
+    }
+
+    /// Returns `true` if `self` is newer than `other`.
+    ///
+    /// The value is considered newer if it is ahead of the other value
+    /// by less than [`PatchIndex::MAX_NEWER_DISTANCE`].
+    pub fn is_newer_than(self, other: Self) -> bool {
+        let distance = self.distance_after(other);
+        distance != 0 && distance <= Self::MAX_NEWER_DISTANCE
+    }
+
+    /// Returns the wrapping distance from `base` to `self`.
+    #[inline]
+    pub fn distance_after(self, base: Self) -> u16 {
+        self.0.wrapping_sub(base.0)
+    }
+}
+
+impl Add<u16> for PatchIndex {
+    type Output = Self;
+
+    #[inline]
+    fn add(self, rhs: u16) -> Self::Output {
+        Self(self.0.wrapping_add(rhs))
+    }
+}
+
+impl Sub<u16> for PatchIndex {
+    type Output = Self;
+
+    #[inline]
+    fn sub(self, rhs: u16) -> Self::Output {
+        Self(self.0.wrapping_sub(rhs))
+    }
+}
+
+impl Sub for PatchIndex {
+    type Output = u16;
+
+    #[inline]
+    fn sub(self, rhs: Self) -> Self::Output {
+        self.0.wrapping_sub(rhs.0)
+    }
+}

tests/diff.rs

@@ -157,7 +157,7 @@ fn multiple_patches_in_same_send_share_patch_cursor() {
     replicate_and_ack(&mut server_app, &mut client_app);
 
     assert_client_point_values(&mut client_app, 0..=2);
-    assert_diff_cursor(&server_app, server_entity, Some(0));
+    assert_diff_cursor(&server_app, server_entity, Some(PatchIndex::new(0)));
 
     server_app
         .world_mut()
@@ -167,7 +167,7 @@ fn multiple_patches_in_same_send_share_patch_cursor() {
     replicate_and_ack(&mut server_app, &mut client_app);
 
     assert_client_point_values(&mut client_app, 0..=3);
-    assert_diff_cursor(&server_app, server_entity, Some(1));
+    assert_diff_cursor(&server_app, server_entity, Some(PatchIndex::new(1)));
 }
 
 #[test]
@@ -334,15 +334,15 @@ fn duplicate_patches_are_ignored_by_receiver() {
     );
     entity.apply_write(
         wire(DiffWire::Patches {
-            first_index: 0,
+            first_index: PatchIndex::new(0),
             patches: vec![vec![PointPatch::PushBack(Vec2::new(2.0, 2.0))]],
         }),
         fns_id,
         RepliconTick::default(),
     );
     entity.apply_write(
         wire(DiffWire::Patches {
-            first_index: 0,
+            first_index: PatchIndex::new(0),
             patches: vec![vec![PointPatch::PushBack(Vec2::new(2.0, 2.0))]],
         }),
         fns_id,
@@ -368,7 +368,7 @@ fn out_of_order_patches_wait_for_missing_predecessor() {
     );
     entity.apply_write(
         wire(DiffWire::Patches {
-            first_index: 1,
+            first_index: PatchIndex::new(1),
             patches: vec![vec![PointPatch::PushBack(Vec2::new(3.0, 3.0))]],
         }),
         fns_id,
@@ -378,7 +378,7 @@ fn out_of_order_patches_wait_for_missing_predecessor() {
 
     entity.apply_write(
         wire(DiffWire::Patches {
-            first_index: 0,
+            first_index: PatchIndex::new(0),
             patches: vec![vec![PointPatch::PushBack(Vec2::new(2.0, 2.0))]],
         }),
         fns_id,
@@ -396,7 +396,7 @@ fn patches_before_snapshot_are_rejected() {
 
     entity.apply_write(
         wire(DiffWire::Patches {
-            first_index: 0,
+            first_index: PatchIndex::new(0),
             patches: vec![vec![PointPatch::PushBack(Vec2::new(1.0, 1.0))]],
         }),
         fns_id,
@@ -520,8 +520,8 @@ fn write_point_history(
             // Batch N transforms state cursor N - 1 into cursor N. Batch 0
             // transforms the pre-patch base, represented by `None`, into
             // cursor `Some(0)`.
-            let base_cursor = first_index.checked_sub(1);
-            let cursor = Some(first_index + patches.len() as PatchIndex - 1);
+            let base_cursor = (first_index != PatchIndex::new(0)).then_some(first_index - 1);
+            let cursor = Some(first_index + patches.len() as u16 - 1);
             // The base must come from a confirmed value in the history: consumers
             // like prediction/interpolation may locally mutate the live component,
             // so it can never be used as a patch base.