From 0a420f21cb3e77b1fd779e564680b1e15dde1029 Mon Sep 17 00:00:00 2001
From: Hennadii Chernyshchyk <genaloner@gmail.com>
Date: Tue, 22 Apr 2025 22:47:34 +0300
Subject: [PATCH] Refactor `UpdateMessage` and `MutateMessage`

- Rename into `Updates` and `Mutations`. Shorter and these components
  aren't messages. We actually create messages from them.
- Use shorter methods in `Mutations`. Things like
  `mutations.add_mutated_entity` looks too redundant.
- In `Updates` I made few names more explicit since it accumulates not
  only changes.
- Use the variable name `changes` when referring to `ChangeRanges`.
---
 src/server.rs                                 |  84 +++++++-------
 src/server/replication_messages.rs            |   4 +-
 .../replication_messages/change_ranges.rs     |   4 +-
 .../{mutate_message.rs => mutations.rs}       | 103 +++++++++---------
 .../replication_messages/serialized_data.rs   |   4 +-
 .../{update_message.rs => updates.rs}         |  71 ++++++------
 6 files changed, 135 insertions(+), 135 deletions(-)
 rename src/server/replication_messages/{mutate_message.rs => mutations.rs} (70%)
 rename src/server/replication_messages/{update_message.rs => updates.rs} (86%)

diff --git a/src/server.rs b/src/server.rs
index cacc7381..69d8d05f 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -43,7 +43,7 @@ use client_visibility::{ClientVisibility, Visibility};
 use despawn_buffer::{DespawnBuffer, DespawnBufferPlugin};
 use removal_buffer::{RemovalBuffer, RemovalBufferPlugin};
 use replication_messages::{
-    mutate_message::MutateMessage, serialized_data::SerializedData, update_message::UpdateMessage,
+    mutations::Mutations, serialized_data::SerializedData, updates::Updates,
 };
 use server_tick::ServerTick;
 use server_world::{ReplicatedComponent, ServerWorld};
@@ -242,8 +242,8 @@ pub(super) fn send_replication(
     world: ServerWorld,
     mut clients: Query<(
         Entity,
-        &mut UpdateMessage,
-        &mut MutateMessage,
+        &mut Updates,
+        &mut Mutations,
         &mut ConnectedClient,
         &mut ClientEntityMap,
         &mut ClientTicks,
@@ -259,9 +259,9 @@ pub(super) fn send_replication(
     server_tick: Res<ServerTick>,
     time: Res<Time>,
 ) -> Result<()> {
-    for (_, mut mutate_message, mut update_message, ..) in &mut clients {
-        update_message.clear();
-        mutate_message.clear();
+    for (_, mut mutations, mut updates, ..) in &mut clients {
+        updates.clear();
+        mutations.clear();
     }
 
     collect_mappings(&mut serialized, &mut clients)?;
@@ -310,8 +310,8 @@ fn reset(
 fn send_messages(
     clients: &mut Query<(
         Entity,
-        &mut UpdateMessage,
-        &mut MutateMessage,
+        &mut Updates,
+        &mut Mutations,
         &mut ConnectedClient,
         &mut ClientEntityMap,
         &mut ClientTicks,
@@ -326,23 +326,21 @@ fn send_messages(
     time: &Time,
 ) -> Result<()> {
     let mut server_tick_range = None;
-    for (client_entity, update_message, mut mutate_message, client, .., mut ticks, visibility) in
-        clients
-    {
-        if !update_message.is_empty() {
+    for (client_entity, updates, mut mutations, client, .., mut ticks, visibility) in clients {
+        if !updates.is_empty() {
             ticks.set_update_tick(server_tick);
             let server_tick = write_tick_cached(&mut server_tick_range, serialized, server_tick)?;
 
             trace!("sending update message to client `{client_entity}`");
-            update_message.send(server, client_entity, serialized, server_tick)?;
+            updates.send(server, client_entity, serialized, server_tick)?;
         } else {
             trace!("no updates to send for client `{client_entity}`");
         }
 
-        if !mutate_message.is_empty() || track_mutate_messages {
+        if !mutations.is_empty() || track_mutate_messages {
             let server_tick = write_tick_cached(&mut server_tick_range, serialized, server_tick)?;
 
-            let messages_count = mutate_message.send(
+            let messages_count = mutations.send(
                 server,
                 client_entity,
                 &mut ticks,
@@ -372,8 +370,8 @@ fn collect_mappings(
     serialized: &mut SerializedData,
     clients: &mut Query<(
         Entity,
-        &mut UpdateMessage,
-        &mut MutateMessage,
+        &mut Updates,
+        &mut Mutations,
         &mut ConnectedClient,
         &mut ClientEntityMap,
         &mut ClientTicks,
@@ -394,8 +392,8 @@ fn collect_despawns(
     serialized: &mut SerializedData,
     clients: &mut Query<(
         Entity,
-        &mut UpdateMessage,
-        &mut MutateMessage,
+        &mut Updates,
+        &mut Mutations,
         &mut ConnectedClient,
         &mut ClientEntityMap,
         &mut ClientTicks,
@@ -436,8 +434,8 @@ fn collect_removals(
     serialized: &mut SerializedData,
     clients: &mut Query<(
         Entity,
-        &mut UpdateMessage,
-        &mut MutateMessage,
+        &mut Updates,
+        &mut Mutations,
         &mut ConnectedClient,
         &mut ClientEntityMap,
         &mut ClientTicks,
@@ -464,8 +462,8 @@ fn collect_changes(
     serialized: &mut SerializedData,
     clients: &mut Query<(
         Entity,
-        &mut UpdateMessage,
-        &mut MutateMessage,
+        &mut Updates,
+        &mut Mutations,
         &mut ConnectedClient,
         &mut ClientEntityMap,
         &mut ClientTicks,
@@ -481,12 +479,12 @@ fn collect_changes(
     for (archetype, replicated_archetype) in world.iter_archetypes() {
         for entity in archetype.entities() {
             let mut entity_range = None;
-            for (_, mut update_message, mut mutate_message, .., visibility) in &mut *clients {
+            for (_, mut updates, mut mutations, .., visibility) in &mut *clients {
                 let visibility = visibility
                     .map(|v| v.state(entity.id()))
                     .unwrap_or(Visibility::Visible);
-                update_message.start_entity_changes(visibility);
-                mutate_message.start_entity_mutations();
+                updates.start_entity_changes(visibility);
+                mutations.start_entity();
             }
 
             // SAFETY: all replicated archetypes have marker component with table storage.
@@ -524,27 +522,25 @@ fn collect_changes(
                     type_registry,
                 };
                 let mut component_range = None;
-                for (_, mut update_message, mut mutate_message, .., client_ticks, _) in
-                    &mut *clients
-                {
-                    if update_message.entity_visibility() == Visibility::Hidden {
+                for (_, mut updates, mut mutations, .., client_ticks, _) in &mut *clients {
+                    if updates.entity_visibility() == Visibility::Hidden {
                         continue;
                     }
 
                     if let Some(tick) = client_ticks
                         .mutation_tick(entity.id())
                         .filter(|_| !marker_added)
-                        .filter(|_| update_message.entity_visibility() != Visibility::Gained)
+                        .filter(|_| updates.entity_visibility() != Visibility::Gained)
                         .filter(|_| !ticks.is_added(change_tick.last_run(), change_tick.this_run()))
                     {
                         if ticks.is_changed(tick, change_tick.this_run()) {
-                            if !mutate_message.mutations_written() {
+                            if !mutations.entity_added() {
                                 let entity_range = write_entity_cached(
                                     &mut entity_range,
                                     serialized,
                                     entity.id(),
                                 )?;
-                                mutate_message.add_mutated_entity(entity.id(), entity_range);
+                                mutations.add_entity(entity.id(), entity_range);
                             }
                             let component_range = write_component_cached(
                                 &mut component_range,
@@ -555,13 +551,13 @@ fn collect_changes(
                                 replicated_component,
                                 component,
                             )?;
-                            mutate_message.add_mutated_component(component_range);
+                            mutations.add_component(component_range);
                         }
                     } else {
-                        if !update_message.entity_written() {
+                        if !updates.changed_entity_added() {
                             let entity_range =
                                 write_entity_cached(&mut entity_range, serialized, entity.id())?;
-                            update_message.add_changed_entity(entity_range);
+                            updates.add_changed_entity(entity_range);
                         }
                         let component_range = write_component_cached(
                             &mut component_range,
@@ -572,33 +568,33 @@ fn collect_changes(
                             replicated_component,
                             component,
                         )?;
-                        update_message.add_inserted_component(component_range);
+                        updates.add_inserted_component(component_range);
                     }
                 }
             }
 
-            for (_, mut update_message, mut mutate_message, .., mut ticks, _) in &mut *clients {
-                let visibility = update_message.entity_visibility();
+            for (_, mut updates, mut mutations, .., mut ticks, _) in &mut *clients {
+                let visibility = updates.entity_visibility();
                 if visibility == Visibility::Hidden {
                     continue;
                 }
 
                 let new_entity = marker_added || visibility == Visibility::Gained;
                 if new_entity
-                    || update_message.entity_written()
+                    || updates.changed_entity_added()
                     || removal_buffer.contains_key(&entity.id())
                 {
                     // If there is any insertion, removal, or it's a new entity for a client, include all mutations
                     // into update message and bump the last acknowledged tick to keep entity updates atomic.
-                    update_message.take_mutations(&mut mutate_message);
+                    updates.take_added_entity(&mut mutations);
                     ticks.set_mutation_tick(entity.id(), change_tick.this_run());
                 }
 
-                if new_entity && !update_message.entity_written() {
+                if new_entity && !updates.changed_entity_added() {
                     // Force-write new entity even if it doesn't have any components.
                     let entity_range =
                         write_entity_cached(&mut entity_range, serialized, entity.id())?;
-                    update_message.add_changed_entity(entity_range);
+                    updates.add_changed_entity(entity_range);
                 }
             }
         }
@@ -719,7 +715,7 @@ pub enum TickPolicy {
 ///
 /// Pausing replication by temporarily removing this component is not supported.
 #[derive(Component, Default)]
-#[require(ClientTicks, ClientEntityMap, UpdateMessage, MutateMessage)]
+#[require(ClientTicks, ClientEntityMap, Updates, Mutations)]
 pub struct ReplicatedClient;
 
 /// Controls how visibility will be managed via [`ClientVisibility`].
diff --git a/src/server/replication_messages.rs b/src/server/replication_messages.rs
index 76a3911f..421adb3a 100644
--- a/src/server/replication_messages.rs
+++ b/src/server/replication_messages.rs
@@ -1,4 +1,4 @@
 mod change_ranges;
-pub(super) mod mutate_message;
+pub(super) mod mutations;
 pub(super) mod serialized_data;
-pub(super) mod update_message;
+pub(super) mod updates;
diff --git a/src/server/replication_messages/change_ranges.rs b/src/server/replication_messages/change_ranges.rs
index 32a75906..bf9bd434 100644
--- a/src/server/replication_messages/change_ranges.rs
+++ b/src/server/replication_messages/change_ranges.rs
@@ -7,8 +7,8 @@ use postcard::experimental::serialized_size;
 /// Component insertions or mutations for an entity in form of serialized ranges
 /// from [`SerializedData`](super::serialized_data::SerializedData).
 ///
-/// Used inside [`UpdateMessage`](super::update_message::UpdateMessage) and
-/// [`MutateMessage`](super::mutate_message::MutateMessage).
+/// Used inside [`Updates`](super::updates::Updates) and
+/// [`Mutations`](super::mutations::Mutations).
 pub(super) struct ChangeRanges {
     pub(super) entity: Range<usize>,
     pub(super) components_len: usize,
diff --git a/src/server/replication_messages/mutate_message.rs b/src/server/replication_messages/mutations.rs
similarity index 70%
rename from src/server/replication_messages/mutate_message.rs
rename to src/server/replication_messages/mutations.rs
index 1c6a0324..75aa1a89 100644
--- a/src/server/replication_messages/mutate_message.rs
+++ b/src/server/replication_messages/mutations.rs
@@ -14,24 +14,14 @@ use crate::shared::{
     replicon_tick::RepliconTick,
 };
 
-/// A message with replicated component mutations.
-///
-/// Contains update tick, current tick, mutate index and component mutations since
-/// the last acknowledged tick for each entity.
-///
-/// Cannot be applied on the client until the update message matching this message's update tick
-/// has been applied to the client world.
-/// The message will be manually split into packets up to max size, and each packet will be applied
-/// independently on the client.
-/// Message splits only happen per-entity to avoid weird behavior from partial entity mutations.
+/// Component mutations for the current tick.
 ///
 /// The data is serialized manually and stored in the form of ranges
 /// from [`SerializedData`].
 ///
-/// Sent over the [`ReplicationChannel::Mutations`] channel. If the message gets lost, we try to resend it manually,
-/// using the last up-to-date mutations to avoid re-sending old values.
+/// Can be packed into messages using [`Self::send`].
 #[derive(Default, Component)]
-pub(crate) struct MutateMessage {
+pub(crate) struct Mutations {
     /// List of entity values for [`Self::mutations`].
     ///
     /// Used to associate entities with the mutate index that the client
@@ -44,15 +34,15 @@ pub(crate) struct MutateMessage {
     /// Components are stored in multiple chunks because some clients may acknowledge mutations,
     /// while others may not.
     ///
-    /// Unlike [`UpdateMessage`](super::update_message::UpdateMessage), we serialize the number
+    /// Unlike [`Updates`](super::updates::Updates), we serialize the number
     /// of chunk bytes instead of the number of components. This is because, during deserialization,
     /// some entities may be skipped if they have already been updated (as mutations are sent until
     /// the client acknowledges them).
     mutations: Vec<ChangeRanges>,
 
     /// Indicates that an entity has been written since the
-    /// last call of [`Self::start_entity_mutations`].
-    mutations_written: bool,
+    /// last call of [`Self::start_entity`].
+    entity_added: bool,
 
     /// Intermediate buffer to reuse allocated memory from [`Self::mutations`].
     buffer: Vec<Vec<Range<usize>>>,
@@ -63,23 +53,23 @@ pub(crate) struct MutateMessage {
     messages: Vec<(MutateIndex, usize, Range<usize>)>,
 }
 
-impl MutateMessage {
+impl Mutations {
     /// Updates internal state to start writing mutated components for an entity.
     ///
     /// Entities and their data written lazily during the iteration.
-    /// See [`Self::add_mutated_entity`] and [`Self::add_mutated_component`].
-    pub(crate) fn start_entity_mutations(&mut self) {
-        self.mutations_written = false;
+    /// See [`Self::add_entity`] and [`Self::add_component`].
+    pub(crate) fn start_entity(&mut self) {
+        self.entity_added = false;
     }
 
-    /// Returns `true` if [`Self::add_mutated_entity`] were called since the last
-    /// call of [`Self::start_entity_mutations`].
-    pub(crate) fn mutations_written(&mut self) -> bool {
-        self.mutations_written
+    /// Returns `true` if [`Self::add_entity`] were called since the last
+    /// call of [`Self::start_entity`].
+    pub(crate) fn entity_added(&mut self) -> bool {
+        self.entity_added
     }
 
     /// Adds an entity chunk.
-    pub(crate) fn add_mutated_entity(&mut self, entity: Entity, entity_range: Range<usize>) {
+    pub(crate) fn add_entity(&mut self, entity: Entity, entity_range: Range<usize>) {
         let components = self.buffer.pop().unwrap_or_default();
         self.mutations.push(ChangeRanges {
             entity: entity_range,
@@ -87,30 +77,32 @@ impl MutateMessage {
             components,
         });
         self.entities.push(entity);
-        self.mutations_written = true;
+        self.entity_added = true;
     }
 
-    /// Adds a component chunk to the last added entity from [`Self::add_mutated_entity`].
-    pub(crate) fn add_mutated_component(&mut self, component: Range<usize>) {
-        let mutations = self
+    /// Adds a component chunk to the last added entity from [`Self::add_entity`].
+    pub(crate) fn add_component(&mut self, component: Range<usize>) {
+        let changes = self
             .mutations
             .last_mut()
             .expect("entity should be written before adding components");
 
-        mutations.add_component(component);
+        changes.add_component(component);
     }
 
-    /// Returns written mutations for the last entity from [`Self::add_mutated_entity`].
-    pub(super) fn last_mutations(&mut self) -> Option<&ChangeRanges> {
+    /// Returns written mutations for the last entity from [`Self::add_entity`].
+    pub(super) fn last(&mut self) -> Option<&ChangeRanges> {
         self.mutations.last()
     }
 
-    /// Removes last added entity from [`Self::add_mutated_entity`] with associated components.
-    pub(super) fn pop_mutations(&mut self) {
+    /// Removes last added entity from [`Self::add_entity`] with associated components.
+    ///
+    /// keeps allocated memory for reuse.
+    pub(super) fn pop(&mut self) {
         self.entities.pop();
-        if let Some(mut mutations) = self.mutations.pop() {
-            mutations.components.clear();
-            self.buffer.push(mutations.components);
+        if let Some(mut changes) = self.mutations.pop() {
+            changes.components.clear();
+            self.buffer.push(changes.components);
         }
     }
 
@@ -118,6 +110,19 @@ impl MutateMessage {
         self.mutations.is_empty()
     }
 
+    /// Packs mutations into messages.
+    ///
+    /// Contains update tick, current tick, mutate index and component mutations since
+    /// the last acknowledged tick for each entity.
+    ///
+    /// Cannot be applied on the client until the update message matching this message's update tick
+    /// has been applied to the client world.
+    /// The message will be manually split into packets up to max size, and each packet will be applied
+    /// independently on the client.
+    /// Message splits only happen per-entity to avoid weird behavior from partial entity mutations.
+    ///
+    /// Sent over the [`ReplicationChannel::Mutations`] channel. If the message gets lost, we try to resend it manually,
+    /// using the last up-to-date mutations to avoid re-sending old values.
     pub(crate) fn send(
         &mut self,
         server: &mut RepliconServer,
@@ -146,13 +151,13 @@ impl MutateMessage {
         let mut header_size = metadata_size + serialized_size(&mutate_index)?;
         let mut body_size = 0;
         let mut mutations_range = Range::<usize>::default();
-        for (entity, mutations) in self.entities.iter().zip(&self.mutations) {
-            let mutations_size = mutations.size_with_components_size()?;
+        for (entity, changes) in self.entities.iter().zip(&self.mutations) {
+            let changes_size = changes.size_with_components_size()?;
 
             // Try to pack back first, then try to pack forward.
             if body_size != 0
-                && !can_pack(header_size + body_size, mutations_size, max_size)
-                && !can_pack(header_size + mutations_size, body_size, max_size)
+                && !can_pack(header_size + body_size, changes_size, max_size)
+                && !can_pack(header_size + changes_size, body_size, max_size)
             {
                 self.messages.push((
                     mutate_index,
@@ -169,7 +174,7 @@ impl MutateMessage {
 
             entities.push(*entity);
             mutations_range.end += 1;
-            body_size += mutations_size;
+            body_size += changes_size;
         }
         if !mutations_range.is_empty() || track_mutate_messages {
             // When the loop ends, pack all leftovers into a message.
@@ -195,10 +200,10 @@ impl MutateMessage {
                 postcard_utils::to_extend_mut(&messages_count, &mut message)?;
             }
             postcard_utils::to_extend_mut(&mutate_index, &mut message)?;
-            for mutations in &self.mutations[mutations_range.clone()] {
-                message.extend_from_slice(&serialized[mutations.entity.clone()]);
-                postcard_utils::to_extend_mut(&mutations.components_size(), &mut message)?;
-                for component in &mutations.components {
+            for changes in &self.mutations[mutations_range.clone()] {
+                message.extend_from_slice(&serialized[changes.entity.clone()]);
+                postcard_utils::to_extend_mut(&changes.components_size(), &mut message)?;
+                for component in &changes.components {
                     message.extend_from_slice(&serialized[component.clone()]);
                 }
             }
@@ -217,9 +222,9 @@ impl MutateMessage {
     pub(crate) fn clear(&mut self) {
         self.entities.clear();
         self.buffer
-            .extend(self.mutations.drain(..).map(|mut mutations| {
-                mutations.components.clear();
-                mutations.components
+            .extend(self.mutations.drain(..).map(|mut changes| {
+                changes.components.clear();
+                changes.components
             }));
     }
 }
diff --git a/src/server/replication_messages/serialized_data.rs b/src/server/replication_messages/serialized_data.rs
index 650a3ca6..0216a415 100644
--- a/src/server/replication_messages/serialized_data.rs
+++ b/src/server/replication_messages/serialized_data.rs
@@ -12,8 +12,8 @@ use crate::shared::{
 
 /// Single continuous buffer that stores serialized data for messages.
 ///
-/// See [`UpdateMessage`](super::update_message::UpdateMessage) and
-/// [`MutateMessage`](super::mutate_message::MutateMessage).
+/// See [`Updates`](super::updates::Updates) and
+/// [`MutateMessage`](super::mutations::MutateMessage).
 #[derive(Default, Deref, DerefMut)]
 pub(crate) struct SerializedData(Vec<u8>);
 
diff --git a/src/server/replication_messages/update_message.rs b/src/server/replication_messages/updates.rs
similarity index 86%
rename from src/server/replication_messages/update_message.rs
rename to src/server/replication_messages/updates.rs
index c6b2686d..068b4456 100644
--- a/src/server/replication_messages/update_message.rs
+++ b/src/server/replication_messages/updates.rs
@@ -3,9 +3,7 @@ use core::ops::Range;
 use bevy::prelude::*;
 use postcard::experimental::serialized_size;
 
-use super::{
-    change_ranges::ChangeRanges, mutate_message::MutateMessage, serialized_data::SerializedData,
-};
+use super::{change_ranges::ChangeRanges, mutations::Mutations, serialized_data::SerializedData};
 use crate::server::client_visibility::Visibility;
 use crate::shared::{
     backend::{replicon_channels::ReplicationChannel, replicon_server::RepliconServer},
@@ -13,28 +11,14 @@ use crate::shared::{
     replication::update_message_flags::UpdateMessageFlags,
 };
 
-/// A message with replicated data.
-///
-/// Contains tick, mappings, insertions, removals, and despawns that
-/// happened in this tick.
+/// Entity updates for the current tick.
 ///
 /// The data is serialized manually and stored in the form of ranges
 /// from [`SerializedData`].
 ///
-/// Sent over [`ReplicationChannel::Updates`] channel.
-///
-/// Some data is optional, and their presence is encoded in the [`UpdateMessageFlags`] bitset.
-///
-/// To know how much data array takes, we serialize it's length. We use `usize`,
-/// but we use variable integer encoding, so they are correctly deserialized even
-/// on a client with a different pointer size. However, if the server sends a value
-/// larger than what a client can fit into `usize` (which is very unlikely), the client will panic.
-/// This is expected, as the client can't have an array of such a size anyway.
-///
-/// Additionally, we don't serialize the size for the last array and
-/// on deserialization just consume all remaining bytes.
+/// Can be packed into a message using [`Self::send`].
 #[derive(Default, Component)]
-pub(crate) struct UpdateMessage {
+pub(crate) struct Updates {
     /// Mappings for client's pre-spawned entities.
     ///
     /// Serialized as single continuous chunk of entity pairs.
@@ -82,13 +66,13 @@ pub(crate) struct UpdateMessage {
 
     /// Indicates that an entity has been written since the
     /// last call of [`Self::start_entity_changes`].
-    entity_written: bool,
+    changed_entity_added: bool,
 
     /// Intermediate buffer to reuse allocated memory from [`Self::changes`].
     buffer: Vec<Vec<Range<usize>>>,
 }
 
-impl UpdateMessage {
+impl Updates {
     pub(crate) fn set_mappings(&mut self, mappings: Range<usize>, len: usize) {
         self.mappings = mappings;
         self.mappings_len = len;
@@ -125,7 +109,7 @@ impl UpdateMessage {
     /// See [`Self::add_changed_entity`] and [`Self::add_changed_component`].
     pub(crate) fn start_entity_changes(&mut self, visibility: Visibility) {
         self.entity_visibility = visibility;
-        self.entity_written = false;
+        self.changed_entity_added = false;
     }
 
     /// Visibility from the last call of [`Self::start_entity_changes`].
@@ -135,8 +119,8 @@ impl UpdateMessage {
 
     /// Returns `true` if [`Self::add_changed_entity`] were called since the last
     /// call of [`Self::start_entity_changes`].
-    pub(crate) fn entity_written(&mut self) -> bool {
-        self.entity_written
+    pub(crate) fn changed_entity_added(&mut self) -> bool {
+        self.changed_entity_added
     }
 
     /// Adds an entity chunk.
@@ -147,7 +131,7 @@ impl UpdateMessage {
             components_len: 0,
             components,
         });
-        self.entity_written = true;
+        self.changed_entity_added = true;
     }
 
     /// Adds a component chunk to the last added entity from [`Self::add_changed_entity`].
@@ -161,29 +145,27 @@ impl UpdateMessage {
     }
 
     /// Takes last mutated entity with its component chunks from the mutate message.
-    pub(crate) fn take_mutations(&mut self, mutate_message: &mut MutateMessage) {
-        if !mutate_message.mutations_written() {
+    pub(crate) fn take_added_entity(&mut self, mutations: &mut Mutations) {
+        if !mutations.entity_added() {
             return;
         }
 
-        let mutations = mutate_message
-            .last_mutations()
-            .expect("entity should be written");
+        let last_changes = mutations.last().expect("entity should be written");
 
-        if !self.entity_written {
+        if !self.changed_entity_added {
             let components = self.buffer.pop().unwrap_or_default();
             let changes = ChangeRanges {
-                entity: mutations.entity.clone(),
+                entity: last_changes.entity.clone(),
                 components_len: 0,
                 components,
             };
             self.changes.push(changes);
         }
         let changes = self.changes.last_mut().unwrap();
-        debug_assert_eq!(mutations.entity, changes.entity);
-        changes.extend(mutations);
+        debug_assert_eq!(last_changes.entity, changes.entity);
+        changes.extend(last_changes);
 
-        mutate_message.pop_mutations();
+        mutations.pop();
     }
 
     pub(crate) fn is_empty(&self) -> bool {
@@ -193,6 +175,23 @@ impl UpdateMessage {
             && self.mappings.is_empty()
     }
 
+    /// Packs updates into a message.
+    ///
+    /// Contains tick, mappings, insertions, removals, and despawns that
+    /// happened in this tick.
+    ///
+    /// Sent over [`ReplicationChannel::Updates`] channel.
+    ///
+    /// Some data is optional, and their presence is encoded in the [`UpdateMessageFlags`] bitset.
+    ///
+    /// To know how much data array takes, we serialize it's length. We use `usize`,
+    /// but we use variable integer encoding, so they are correctly deserialized even
+    /// on a client with a different pointer size. However, if the server sends a value
+    /// larger than what a client can fit into `usize` (which is very unlikely), the client will panic.
+    /// This is expected, as the client can't have an array of such a size anyway.
+    ///
+    /// Additionally, we don't serialize the size for the last array and
+    /// on deserialization just consume all remaining bytes.
     pub(crate) fn send(
         &self,
         server: &mut RepliconServer,