Apply TickEvent to Correction (#928)

* simplify pop history

* fix

lightyear/src/client/prediction/correction.rs

@@ -45,6 +45,14 @@ pub struct Correction<C: Component> {
     pub current_correction: Option<C>,
 }
 
+impl <C: Component> Correction<C> {
+    /// In case of a TickEvent where the client tick is changed, we need to update the ticks in the buffer
+    pub(crate) fn update_ticks(&mut self, delta: i16) {
+        self.original_tick = self.original_tick + delta;
+        self.final_correction_tick = self.final_correction_tick + delta;
+    }
+}
+
 /// Perform the correction: we interpolate between the original (incorrect) prediction and the final confirmed value
 /// over a period of time. The intermediary state is called the Corrected state.
 pub(crate) fn get_corrected_state<C: SyncComponent>(

lightyear/src/client/prediction/predicted_history.rs

@@ -9,7 +9,7 @@ use crate::client::components::{ComponentSyncMode, Confirmed, SyncComponent};
 use crate::client::prediction::resource::PredictionManager;
 use crate::client::prediction::rollback::Rollback;
 use crate::client::prediction::Predicted;
-use crate::prelude::client::PredictionSet;
+use crate::prelude::client::{Correction, PredictionSet};
 use crate::prelude::{
     ComponentRegistry, HistoryBuffer, PrePredicted, PreSpawnedPlayerObject, TickManager,
 };
@@ -44,12 +44,15 @@ pub(crate) fn update_prediction_history<T: Component + PartialEq + Clone>(
 /// (i.e. X ticks in the past compared to the current tick)
 pub(crate) fn handle_tick_event_prediction_history<C: Component>(
     trigger: Trigger<TickEvent>,
-    mut query: Query<&mut PredictionHistory<C>>,
+    mut query: Query<(&mut PredictionHistory<C>, Option<&mut Correction<C>>)>,
 ) {
     match *trigger.event() {
         TickEvent::TickSnap { old_tick, new_tick } => {
-            for mut history in query.iter_mut() {
+            for (mut history, correction) in query.iter_mut() {
                 history.update_ticks(new_tick - old_tick);
+                if let Some(mut correction) = correction {
+                    correction.update_ticks(new_tick - old_tick);
+                }
             }
         }
     }

lightyear/src/client/prediction/rollback.rs

@@ -332,7 +332,8 @@ pub(crate) fn prepare_rollback<C: SyncComponent>(
         };
 
         // 2. we need to clear the history so we can write a new one
-        let original_predicted_value = predicted_history.pop_until_tick_and_clear(rollback_tick);
+        let original_predicted_value = predicted_history.pop_until_tick(rollback_tick);
+        predicted_history.clear();
 
         // if rollback is disabled, we will restore the component to its past value from the prediction history
         let correct_value = if disable_rollback {

lightyear/src/utils/history_buffer.rs

@@ -191,27 +191,6 @@ impl<R: Clone> HistoryBuffer<R> {
         };
         res
     }
-
-    /// Return the value at the specified tick and then clear the history
-    ///
-    /// We create this separate function to avoid cloning the value.
-    pub fn pop_until_tick_and_clear(&mut self, tick: Tick) -> Option<HistoryState<R>> {
-        // self.buffer[partition] is the first element where the buffer_tick > tick
-        let partition = self
-            .buffer
-            .partition_point(|(buffer_tick, _)| buffer_tick <= &tick);
-        // all elements are strictly more recent than the tick
-        if partition == 0 {
-            return None;
-        }
-        // remove all elements strictly older than the tick. We need to keep the element at index `partition-1`
-        // because that is the value at tick `tick`
-        self.buffer.drain(0..(partition - 1));
-        let res = self.buffer.pop_front().map(|(_, state)| state);
-
-        self.buffer.clear();
-        res
-    }
 }
 
 #[cfg(test)]