Snapshot interpolation plugin for bevy_replicon

snaphot_interpolation.rs

struct SnapshotInterpolationPlugin {
    max_tick_rate: u16,
}
impl Plugin for SnapshotInterpolationPlugin {
    fn build(&self, app: &mut App) {
        app.replicate::<Interpolated>()
            .insert_resource(InterpolationConfig {
                max_tick_rate: self.max_tick_rate,
            });
    }
}

#[derive(Resource)]
struct InterpolationConfig {
    max_tick_rate: u16,
}

#[derive(Component, Deserialize, Serialize)]
struct Interpolated;

trait Interpolate {
    fn interpolate(&self, other: Self, t: f32) -> Self;
}

#[derive(Component, Deserialize, Serialize)]
struct SnapshotBuffer<T: Component + Interpolate + Clone> {
    buffer: VecDeque<T>,
    time_since_last_snapshot: f32,
}
impl<T: Component + Interpolate + Clone> SnapshotBuffer<T> {
    pub fn new() -> Self {
        Self {
            buffer: VecDeque::new(),
            time_since_last_snapshot: 0.0,
        }
    }
    pub fn insert(&mut self, element: T) {
        if self.buffer.len() > 1 {
            self.buffer.pop_front();
        }
        self.buffer.push_back(element);
        self.time_since_last_snapshot = 0.0;
    }
}

fn snapshot_buffer_init_system<T: Component + Interpolate + Clone>(
    q_interpolated: Query<(Entity, &T), Added<Interpolated>>,
    mut commands: Commands,
) {
    for (e, initial_value) in q_interpolated.iter() {
        let mut buffer = SnapshotBuffer::new();
        buffer.insert(initial_value.clone());
        commands.entity(e).insert(buffer);
        info!("intialized buffer");
    }
}
fn snapshot_interpolation_system<T: Component + Interpolate + Clone + Debug>(
    mut q: Query<(Entity, &mut T, &mut SnapshotBuffer<T>), With<Interpolated>>,
    time: Res<Time>,
    config: Res<InterpolationConfig>,
) {
    for (e, mut component, mut snapshot_buffer) in q.iter_mut() {
        let buffer = &snapshot_buffer.buffer;
        let elapsed = snapshot_buffer.time_since_last_snapshot;
        if buffer.len() < 2 {
            continue;
        }

        let tick_duration = 1.0 / (config.max_tick_rate as f32);

        if elapsed > tick_duration + time.delta_seconds() {
            continue;
        }

        let t = (elapsed / tick_duration).clamp(0., 1.);
        info!(
            "Snapshot 0: {:?} ,Snapshot 1: {:?}, t: {:?}",
            buffer[0].clone(),
            buffer[1].clone(),
            t
        );
        *component = buffer[0].interpolate(buffer[1].clone(), t);
        snapshot_buffer.time_since_last_snapshot += time.delta_seconds();
    }
}

trait AppInterpolationExt {
    /// TODO: Add docs
    fn replicate_with_interpolation<C>(
        &mut self,
        serialize: SerializeFn,
        deserialize: DeserializeFn,
        remove: RemoveComponentFn,
    ) -> &mut Self
    where
        C: Component + Interpolate + Clone + Debug;
}
impl AppInterpolationExt for App {
    fn replicate_with_interpolation<T>(
        &mut self,
        serialize: SerializeFn,
        deserialize: DeserializeFn,
        remove: RemoveComponentFn,
    ) -> &mut Self
    where
        T: Component + Interpolate + Clone + Debug,
    {
        self.add_systems(
            PreUpdate,
            (
                snapshot_buffer_init_system::<T>,
                snapshot_interpolation_system::<T>,
            )
                .after(ClientSet::Receive)
                .run_if(resource_exists::<RenetClient>()),
        );
        self.replicate_with::<T>(serialize, deserialize, remove)
    }
}