implement reverting block state predictions on ack

3 files changed, 156 insertions, 43 deletions

diff --git a/azalea-client/src/plugins/interact.rs b/azalea-client/src/plugins/interact.rs
index 2e2b039f..31b5acf4 100644
--- a/azalea-client/src/plugins/interact.rs
+++ b/azalea-client/src/plugins/interact.rs
@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use azalea_block::BlockState;
 use azalea_core::{
     direction::Direction,
@@ -96,17 +98,95 @@ impl Client {
     }
 }
 
-/// A component that contains the number of changes this client has made to
-/// blocks.
-#[derive(Component, Copy, Clone, Debug, Default, Deref)]
-pub struct CurrentSequenceNumber(u32);
+/// A component that contains information about our local block state
+/// predictions.
+#[derive(Component, Clone, Debug, Default)]
+pub struct BlockStatePredictionHandler {
+    /// The total number of changes that this client has made to blocks.
+    seq: u32,
+    server_state: HashMap<BlockPos, ServerVerifiedState>,
+}
+#[derive(Clone, Debug)]
+struct ServerVerifiedState {
+    seq: u32,
+    block_state: BlockState,
+    /// Used for teleporting the player back if we're colliding with the block
+    /// that got placed back.
+    #[allow(unused)]
+    player_pos: Vec3,
+}
 
-impl CurrentSequenceNumber {
+impl BlockStatePredictionHandler {
     /// Get the next sequence number that we're going to use and increment the
     /// value.
-    pub fn get_next(&mut self) -> u32 {
-        self.0 += 1;
-        self.0
+    pub fn start_predicting(&mut self) -> u32 {
+        self.seq += 1;
+        self.seq
+    }
+
+    /// Should be called right before the client updates a block with its
+    /// prediction.
+    ///
+    /// This is used to make sure that we can rollback to this state if the
+    /// server acknowledges the sequence number (with
+    /// [`ClientboundBlockChangedAck`]) without having sent a block update.
+    ///
+    /// [`ClientboundBlockChangedAck`]: azalea_protocol::packets::game::ClientboundBlockChangedAck
+    pub fn retain_known_server_state(
+        &mut self,
+        pos: BlockPos,
+        old_state: BlockState,
+        player_pos: Vec3,
+    ) {
+        self.server_state
+            .entry(pos)
+            .and_modify(|s| s.seq = self.seq)
+            .or_insert(ServerVerifiedState {
+                seq: self.seq,
+                block_state: old_state,
+                player_pos: player_pos,
+            });
+    }
+
+    /// Save this update as the correct server state so when the server sends a
+    /// [`ClientboundBlockChangedAck`] we don't roll back this new update.
+    ///
+    /// This should be used when we receive a block update from the server.
+    ///
+    /// [`ClientboundBlockChangedAck`]: azalea_protocol::packets::game::ClientboundBlockChangedAck
+    pub fn update_known_server_state(&mut self, pos: BlockPos, state: BlockState) -> bool {
+        if let Some(s) = self.server_state.get_mut(&pos) {
+            s.block_state = state;
+            true
+        } else {
+            false
+        }
+    }
+
+    pub fn end_prediction_up_to(&mut self, seq: u32, world: &Instance) {
+        let mut to_remove = Vec::new();
+        for (pos, state) in &self.server_state {
+            if state.seq > seq {
+                continue;
+            }
+            to_remove.push(*pos);
+
+            // syncBlockState
+            let client_block_state = world.get_block_state(*pos).unwrap_or_default();
+            let server_block_state = state.block_state;
+            if client_block_state == server_block_state {
+                continue;
+            }
+            world.set_block_state(*pos, server_block_state);
+            // TODO: implement these two functions
+            // if is_colliding(player, *pos, server_block_state) {
+            //     abs_snap_to(state.player_pos);
+            // }
+        }
+
+        for pos in to_remove {
+            self.server_state.remove(&pos);
+        }
     }
 }
 
@@ -163,13 +243,15 @@ pub fn handle_start_use_item_queued(
     query: Query<(
         Entity,
         &StartUseItemQueued,
-        &mut CurrentSequenceNumber,
+        &mut BlockStatePredictionHandler,
         &HitResultComponent,
         &LookDirection,
         Option<&Mining>,
     )>,
 ) {
-    for (entity, start_use_item, mut sequence_number, hit_result, look_direction, mining) in query {
+    for (entity, start_use_item, mut prediction_handler, hit_result, look_direction, mining) in
+        query
+    {
         commands.entity(entity).remove::<StartUseItemQueued>();
 
         if mining.is_some() {
@@ -203,12 +285,13 @@ pub fn handle_start_use_item_queued(
 
         match &hit_result {
             HitResult::Block(block_hit_result) => {
+                let seq = prediction_handler.start_predicting();
                 if block_hit_result.miss {
                     commands.trigger(SendPacketEvent::new(
                         entity,
                         ServerboundUseItem {
                             hand: start_use_item.hand,
-                            sequence: sequence_number.get_next(),
+                            seq,
                             x_rot: look_direction.x_rot,
                             y_rot: look_direction.y_rot,
                         },
@@ -219,7 +302,7 @@ pub fn handle_start_use_item_queued(
                         ServerboundUseItemOn {
                             hand: start_use_item.hand,
                             block_hit: block_hit_result.into(),
-                            sequence: sequence_number.get_next(),
+                            seq,
                         },
                     ));
                     // TODO: depending on the result of useItemOn, this might
diff --git a/azalea-client/src/plugins/mining.rs b/azalea-client/src/plugins/mining.rs
index 51bb5529..88bd3be8 100644
--- a/azalea-client/src/plugins/mining.rs
+++ b/azalea-client/src/plugins/mining.rs
@@ -1,6 +1,6 @@
 use azalea_block::{BlockState, BlockTrait, fluid_state::FluidState};
 use azalea_core::{direction::Direction, game_type::GameMode, position::BlockPos, tick::GameTick};
-use azalea_entity::{FluidOnEyes, Physics, mining::get_mine_progress};
+use azalea_entity::{FluidOnEyes, Physics, Position, mining::get_mine_progress};
 use azalea_inventory::ItemStack;
 use azalea_physics::{PhysicsSet, collision::BlockWithShape};
 use azalea_protocol::packets::game::s_player_action::{self, ServerboundPlayerAction};
@@ -13,7 +13,7 @@ use tracing::trace;
 use crate::{
     Client,
     interact::{
-        CurrentSequenceNumber, HitResultComponent, SwingArmEvent, can_use_game_master_blocks,
+        BlockStatePredictionHandler, HitResultComponent, SwingArmEvent, can_use_game_master_blocks,
         check_is_interaction_restricted,
     },
     inventory::{Inventory, InventorySet},
@@ -216,7 +216,7 @@ fn handle_mining_queued(
         &FluidOnEyes,
         &Physics,
         Option<&Mining>,
-        &mut CurrentSequenceNumber,
+        &mut BlockStatePredictionHandler,
         &mut MineDelay,
         &mut MineProgress,
         &mut MineTicks,
@@ -280,7 +280,7 @@ fn handle_mining_queued(
                         pos: current_mining_pos
                             .expect("IsMining is true so MineBlockPos must be present"),
                         direction: mining_queued.direction,
-                        sequence: 0,
+                        seq: 0,
                     },
                 ));
             }
@@ -345,7 +345,7 @@ fn handle_mining_queued(
                     action: s_player_action::Action::StartDestroyBlock,
                     pos: mining_queued.position,
                     direction: mining_queued.direction,
-                    sequence: sequence_number.get_next(),
+                    seq: sequence_number.start_predicting(),
                 },
             ));
             // vanilla really does send two swing arm packets
@@ -440,14 +440,22 @@ pub fn handle_finish_mining_block_observer(
         &Inventory,
         &PlayerAbilities,
         &PermissionLevel,
-        &mut CurrentSequenceNumber,
+        &Position,
+        &mut BlockStatePredictionHandler,
     )>,
     instances: Res<InstanceContainer>,
 ) {
     let event = trigger.event();
 
-    let (instance_name, game_mode, inventory, abilities, permission_level, _sequence_number) =
-        query.get_mut(trigger.target()).unwrap();
+    let (
+        instance_name,
+        game_mode,
+        inventory,
+        abilities,
+        permission_level,
+        player_pos,
+        mut prediction_handler,
+    ) = query.get_mut(trigger.target()).unwrap();
     let instance_lock = instances.get(instance_name).unwrap();
     let instance = instance_lock.read();
     if check_is_interaction_restricted(&instance, event.position, &game_mode.current, inventory) {
@@ -469,7 +477,8 @@ pub fn handle_finish_mining_block_observer(
         return;
     };
 
-    let registry_block = Box::<dyn BlockTrait>::from(block_state).as_registry_block();
+    let registry_block: azalea_registry::Block =
+        Box::<dyn BlockTrait>::from(block_state).as_registry_block();
     if !can_use_game_master_blocks(abilities, permission_level)
         && matches!(
             registry_block,
@@ -485,7 +494,10 @@ pub fn handle_finish_mining_block_observer(
     // when we break a waterlogged block we want to keep the water there
     let fluid_state = FluidState::from(block_state);
     let block_state_for_fluid = BlockState::from(fluid_state);
-    instance.set_block_state(event.position, block_state_for_fluid);
+    let old_state = instance
+        .set_block_state(event.position, block_state_for_fluid)
+        .unwrap_or_default();
+    prediction_handler.retain_known_server_state(event.position, old_state, **player_pos);
 }
 
 /// Abort mining a block.
@@ -510,7 +522,7 @@ pub fn handle_stop_mining_block_event(
                 action: s_player_action::Action::AbortDestroyBlock,
                 pos: mine_block_pos,
                 direction: Direction::Down,
-                sequence: 0,
+                seq: 0,
             },
         ));
         commands.entity(event.entity).remove::<Mining>();
@@ -538,7 +550,7 @@ pub fn continue_mining_block(
         &mut MineDelay,
         &mut MineProgress,
         &mut MineTicks,
-        &mut CurrentSequenceNumber,
+        &mut BlockStatePredictionHandler,
     )>,
     mut commands: Commands,
     mut mine_block_progress_events: EventWriter<MineBlockProgressEvent>,
@@ -557,7 +569,7 @@ pub fn continue_mining_block(
         mut mine_delay,
         mut mine_progress,
         mut mine_ticks,
-        mut sequence_number,
+        mut prediction_handler,
     ) in query.iter_mut()
     {
         if **mine_delay > 0 {
@@ -580,7 +592,7 @@ pub fn continue_mining_block(
                     action: s_player_action::Action::StartDestroyBlock,
                     pos: mining.pos,
                     direction: mining.dir,
-                    sequence: sequence_number.get_next(),
+                    seq: prediction_handler.start_predicting(),
                 },
             ));
             commands.trigger(SwingArmEvent { entity });
@@ -634,7 +646,7 @@ pub fn continue_mining_block(
                         action: s_player_action::Action::StopDestroyBlock,
                         pos: mining.pos,
                         direction: mining.dir,
-                        sequence: sequence_number.get_next(),
+                        seq: prediction_handler.start_predicting(),
                     },
                 ));
                 **mine_progress = 0.;
diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index a49a0209..b2a4abc4 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -25,6 +25,7 @@ use crate::{
     connection::RawConnection,
     declare_packet_handlers,
     disconnect::DisconnectEvent,
+    interact::BlockStatePredictionHandler,
     inventory::{
         ClientSideCloseContainerEvent, Inventory, MenuOpenedEvent, SetContainerContentEvent,
     },
@@ -1061,13 +1062,17 @@ impl GamePacketHandler<'_> {
     pub fn block_update(&mut self, p: &ClientboundBlockUpdate) {
         debug!("Got block update packet {p:?}");
 
-        as_system::<Query<&InstanceHolder>>(self.ecs, |mut query| {
-            let local_player = query.get_mut(self.player).unwrap();
-
-            let world = local_player.instance.write();
+        as_system::<Query<(&InstanceHolder, &mut BlockStatePredictionHandler)>>(
+            self.ecs,
+            |mut query| {
+                let (local_player, mut prediction_handler) = query.get_mut(self.player).unwrap();
 
-            world.chunks.set_block_state(p.pos, p.block_state);
-        });
+                let world = local_player.instance.read();
+                if !prediction_handler.update_known_server_state(p.pos, p.block_state) {
+                    world.chunks.set_block_state(p.pos, p.block_state);
+                }
+            },
+        );
     }
 
     pub fn animate(&mut self, p: &ClientboundAnimate) {
@@ -1077,15 +1082,19 @@ impl GamePacketHandler<'_> {
     pub fn section_blocks_update(&mut self, p: &ClientboundSectionBlocksUpdate) {
         debug!("Got section blocks update packet {p:?}");
 
-        as_system::<Query<&InstanceHolder>>(self.ecs, |mut query| {
-            let local_player = query.get_mut(self.player).unwrap();
-            let world = local_player.instance.write();
-            for state in &p.states {
-                world
-                    .chunks
-                    .set_block_state(p.section_pos + state.pos, state.state);
-            }
-        });
+        as_system::<Query<(&InstanceHolder, &mut BlockStatePredictionHandler)>>(
+            self.ecs,
+            |mut query| {
+                let (local_player, mut prediction_handler) = query.get_mut(self.player).unwrap();
+                let world = local_player.instance.read();
+                for new_state in &p.states {
+                    let pos = p.section_pos + new_state.pos;
+                    if !prediction_handler.update_known_server_state(pos, new_state.state) {
+                        world.chunks.set_block_state(pos, new_state.state);
+                    }
+                }
+            },
+        );
     }
 
     pub fn game_event(&mut self, p: &ClientboundGameEvent) {
@@ -1125,7 +1134,16 @@ impl GamePacketHandler<'_> {
 
     pub fn award_stats(&mut self, _p: &ClientboundAwardStats) {}
 
-    pub fn block_changed_ack(&mut self, _p: &ClientboundBlockChangedAck) {}
+    pub fn block_changed_ack(&mut self, p: &ClientboundBlockChangedAck) {
+        as_system::<Query<(&InstanceHolder, &mut BlockStatePredictionHandler)>>(
+            self.ecs,
+            |mut query| {
+                let (local_player, mut prediction_handler) = query.get_mut(self.player).unwrap();
+                let world = local_player.instance.read();
+                prediction_handler.end_prediction_up_to(p.seq, &world);
+            },
+        );
+    }
 
     pub fn block_destruction(&mut self, _p: &ClientboundBlockDestruction) {}