Add support for mob effects (#269)

* Add support for mob effects

* Remove Option

* MobEffectFlags

* jump_boost_power f32

8 files changed, 242 insertions, 45 deletions

diff --git a/azalea-client/src/plugins/mining.rs b/azalea-client/src/plugins/mining.rs
index 101cd4df..b3880c00 100644
--- a/azalea-client/src/plugins/mining.rs
+++ b/azalea-client/src/plugins/mining.rs
@@ -1,6 +1,8 @@
 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, PlayerAbilities, Position, mining::get_mine_progress};
+use azalea_entity::{
+    ActiveEffects, FluidOnEyes, Physics, PlayerAbilities, Position, mining::get_mine_progress,
+};
 use azalea_inventory::ItemStack;
 use azalea_physics::{PhysicsSystems, collision::BlockWithShape};
 use azalea_protocol::packets::game::s_player_action::{self, ServerboundPlayerAction};
@@ -243,6 +245,7 @@ pub fn handle_mining_queued(
         &InstanceHolder,
         &LocalGameMode,
         &Inventory,
+        &ActiveEffects,
         &FluidOnEyes,
         &Physics,
         Option<&mut Mining>,
@@ -260,6 +263,7 @@ pub fn handle_mining_queued(
         instance_holder,
         game_mode,
         inventory,
+        active_effects,
         fluid_on_eyes,
         physics,
         mut mining,
@@ -359,6 +363,7 @@ pub fn handle_mining_queued(
                     &inventory.inventory_menu,
                     fluid_on_eyes,
                     physics,
+                    active_effects,
                 ) >= 1.
             {
                 // block was broken instantly (instamined)
@@ -593,6 +598,7 @@ pub fn continue_mining_block(
         &Inventory,
         &MineBlockPos,
         &MineItem,
+        &ActiveEffects,
         &FluidOnEyes,
         &Physics,
         &Mining,
@@ -612,6 +618,7 @@ pub fn continue_mining_block(
         inventory,
         current_mining_pos,
         current_mining_item,
+        active_effects,
         fluid_on_eyes,
         physics,
         mining,
@@ -669,6 +676,7 @@ pub fn continue_mining_block(
                 &inventory.inventory_menu,
                 fluid_on_eyes,
                 physics,
+                active_effects,
             );
 
             if **mine_ticks % 4. == 0. {
diff --git a/azalea-client/src/plugins/packet/game/mod.rs b/azalea-client/src/plugins/packet/game/mod.rs
index 46cfd531..40446cee 100644
--- a/azalea-client/src/plugins/packet/game/mod.rs
+++ b/azalea-client/src/plugins/packet/game/mod.rs
@@ -7,8 +7,8 @@ use azalea_core::{
     position::{ChunkPos, Vec3},
 };
 use azalea_entity::{
-    Dead, EntityBundle, EntityKindComponent, HasClientLoaded, LoadedBy, LocalEntity, LookDirection,
-    Physics, PlayerAbilities, Position, RelativeEntityUpdate,
+    ActiveEffects, Dead, EntityBundle, EntityKindComponent, HasClientLoaded, LoadedBy, LocalEntity,
+    LookDirection, MobEffectData, Physics, PlayerAbilities, Position, RelativeEntityUpdate,
     indexing::{EntityIdIndex, EntityUuidIndex},
     metadata::{Health, apply_metadata},
 };
@@ -1106,6 +1106,46 @@ impl GamePacketHandler<'_> {
 
     pub fn update_mob_effect(&mut self, p: &ClientboundUpdateMobEffect) {
         debug!("Got update mob effect packet {p:?}");
+
+        let mob_effect = p.mob_effect;
+        let effect_data = MobEffectData::new(
+            p.effect_amplifier,
+            p.effect_duration_ticks,
+            p.flags.ambient,
+            p.flags.show_particles,
+            p.flags.show_icon,
+        );
+
+        as_system::<(Commands, Query<(&EntityIdIndex, &InstanceHolder)>)>(
+            self.ecs,
+            |(mut commands, query)| {
+                let (entity_id_index, instance_holder) = query.get(self.player).unwrap();
+
+                let Some(entity) = entity_id_index.get_by_minecraft_entity(p.entity_id) else {
+                    debug!(
+                        "Got update mob effect packet for unknown entity id {}",
+                        p.entity_id
+                    );
+                    return;
+                };
+
+                let partial_instance = instance_holder.partial_instance.clone();
+                let mob_effect = mob_effect;
+                let effect_data = effect_data.clone();
+                commands.entity(entity).queue(RelativeEntityUpdate::new(
+                    partial_instance,
+                    move |entity| {
+                        if let Some(mut active_effects) = entity.get_mut::<ActiveEffects>() {
+                            active_effects.insert(mob_effect, effect_data.clone());
+                        } else {
+                            let mut active_effects = ActiveEffects::default();
+                            active_effects.insert(mob_effect, effect_data.clone());
+                            entity.insert(active_effects);
+                        }
+                    },
+                ));
+            },
+        );
     }
 
     pub fn award_stats(&mut self, _p: &ClientboundAwardStats) {}
@@ -1314,7 +1354,36 @@ impl GamePacketHandler<'_> {
 
     pub fn player_look_at(&mut self, _p: &ClientboundPlayerLookAt) {}
 
-    pub fn remove_mob_effect(&mut self, _p: &ClientboundRemoveMobEffect) {}
+    pub fn remove_mob_effect(&mut self, p: &ClientboundRemoveMobEffect) {
+        debug!("Got remove mob effect packet {p:?}");
+
+        let mob_effect = p.effect;
+
+        as_system::<(Commands, Query<(&EntityIdIndex, &InstanceHolder)>)>(
+            self.ecs,
+            |(mut commands, query)| {
+                let (entity_id_index, instance_holder) = query.get(self.player).unwrap();
+
+                let Some(entity) = entity_id_index.get_by_minecraft_entity(p.entity_id) else {
+                    debug!(
+                        "Got remove mob effect packet for unknown entity id {}",
+                        p.entity_id
+                    );
+                    return;
+                };
+
+                let partial_instance = instance_holder.partial_instance.clone();
+                commands.entity(entity).queue(RelativeEntityUpdate::new(
+                    partial_instance,
+                    move |entity| {
+                        if let Some(mut active_effects) = entity.get_mut::<ActiveEffects>() {
+                            active_effects.remove(mob_effect);
+                        }
+                    },
+                ));
+            },
+        );
+    }
 
     pub fn resource_pack_push(&mut self, p: &ClientboundResourcePackPush) {
         debug!("Got resource pack packet {p:?}");
diff --git a/azalea-entity/src/effects.rs b/azalea-entity/src/effects.rs
index 9cc750e5..d905414d 100644
--- a/azalea-entity/src/effects.rs
+++ b/azalea-entity/src/effects.rs
@@ -1,21 +1,80 @@
-// TODO
+use std::collections::HashMap;
 
-// pub struct ActiveEffects(HashMap<azalea_registry::MobEffect, MobEffectData>);
+use azalea_registry::MobEffect;
+use bevy_ecs::component::Component;
 
-/// Returns the level of the given effect, or `None` if the effect is not
-/// active. The lowest level is 0.
-pub fn get_effect(_effect: azalea_registry::MobEffect) -> Option<u32> {
-    // TODO
-    None
+/// Data about an active mob effect that the client knows about.
+#[derive(Clone, Debug, Default)]
+pub struct MobEffectData {
+    pub amplifier: u32,
+    pub duration_ticks: u32,
+    pub ambient: bool,
+    pub show_particles: bool,
+    pub show_icon: bool,
+}
+impl MobEffectData {
+    pub fn new(
+        amplifier: u32,
+        duration_ticks: u32,
+        ambient: bool,
+        show_particles: bool,
+        show_icon: bool,
+    ) -> Self {
+        Self {
+            amplifier,
+            duration_ticks,
+            ambient,
+            show_particles,
+            show_icon,
+        }
+    }
+
+    pub fn is_ambient(&self) -> bool {
+        self.ambient
+    }
+    pub fn should_show_particles(&self) -> bool {
+        self.show_particles
+    }
+    pub fn should_show_icon(&self) -> bool {
+        self.show_icon
+    }
+}
+
+/// Component storing the active mob effects on an entity.
+#[derive(Component, Clone, Debug, Default)]
+pub struct ActiveEffects(pub HashMap<MobEffect, MobEffectData>);
+impl ActiveEffects {
+    pub fn insert(&mut self, effect: MobEffect, data: MobEffectData) {
+        self.0.insert(effect, data);
+    }
+
+    pub fn remove(&mut self, effect: MobEffect) -> Option<MobEffectData> {
+        self.0.remove(&effect)
+    }
+
+    pub fn get_level(&self, effect: MobEffect) -> Option<u32> {
+        self.0.get(&effect).map(|data| data.amplifier)
+    }
+
+    pub fn get(&self, effect: MobEffect) -> Option<&MobEffectData> {
+        self.0.get(&effect)
+    }
+}
+
+/// Returns the level (amplifier) of the given effect, or `None` if the effect
+/// is not active. The lowest level is 0.
+pub fn get_effect(active_effects: &ActiveEffects, effect: MobEffect) -> Option<u32> {
+    active_effects.get_level(effect)
 }
 
-pub fn get_dig_speed_amplifier() -> Option<u32> {
+/// Returns the amplifier for dig speed (haste / conduit power), if present.
+pub fn get_dig_speed_amplifier(active_effects: &ActiveEffects) -> Option<u32> {
     let effect_plus_one = u32::max(
-        get_effect(azalea_registry::MobEffect::Haste)
-            .map(|x| x + 1)
+        get_effect(active_effects, MobEffect::Haste)
+            .map(|level| level + 1)
             .unwrap_or_default(),
-        get_effect(azalea_registry::MobEffect::ConduitPower)
-            .map(|x| x + 1)
+        get_effect(active_effects, MobEffect::ConduitPower)
+            .map(|level| level + 1)
             .unwrap_or_default(),
     );
     if effect_plus_one > 0 {
diff --git a/azalea-entity/src/lib.rs b/azalea-entity/src/lib.rs
index 4645e4ba..1f9e426e 100644
--- a/azalea-entity/src/lib.rs
+++ b/azalea-entity/src/lib.rs
@@ -31,6 +31,7 @@ use azalea_world::{ChunkStorage, InstanceName};
 use bevy_ecs::{bundle::Bundle, component::Component};
 pub use data::*;
 use derive_more::{Deref, DerefMut};
+pub use effects::{ActiveEffects, MobEffectData};
 use plugin::indexing::EntityChunkPos;
 use uuid::Uuid;
 use vec_delta_codec::VecDeltaCodec;
@@ -487,6 +488,7 @@ pub struct EntityBundle {
     pub crouching: Crouching,
     pub fluid_on_eyes: FluidOnEyes,
     pub on_climbable: OnClimbable,
+    pub active_effects: ActiveEffects,
 }
 
 impl EntityBundle {
@@ -515,6 +517,7 @@ impl EntityBundle {
             crouching: Crouching(false),
             fluid_on_eyes: FluidOnEyes(FluidKind::Empty),
             on_climbable: OnClimbable(false),
+            active_effects: ActiveEffects::default(),
         }
     }
 }
diff --git a/azalea-entity/src/mining.rs b/azalea-entity/src/mining.rs
index cd526799..7c142020 100644
--- a/azalea-entity/src/mining.rs
+++ b/azalea-entity/src/mining.rs
@@ -2,7 +2,7 @@ use azalea_block::{BlockBehavior, BlockTrait};
 use azalea_core::tier::get_item_tier;
 use azalea_registry as registry;
 
-use crate::{FluidOnEyes, Physics, effects};
+use crate::{ActiveEffects, FluidOnEyes, Physics, effects};
 
 /// How much progress is made towards mining the block per tick, as a
 /// percentage.
@@ -20,6 +20,7 @@ pub fn get_mine_progress(
     player_inventory: &azalea_inventory::Menu,
     fluid_on_eyes: &FluidOnEyes,
     physics: &Physics,
+    active_effects: &ActiveEffects,
 ) -> f32 {
     let block_behavior: BlockBehavior = block.behavior();
 
@@ -39,6 +40,7 @@ pub fn get_mine_progress(
         player_inventory,
         fluid_on_eyes,
         physics,
+        active_effects,
     );
     (base_destroy_speed / destroy_time) / divisor as f32
 }
@@ -82,6 +84,7 @@ fn destroy_speed(
     _player_inventory: &azalea_inventory::Menu,
     _fluid_on_eyes: &FluidOnEyes,
     physics: &Physics,
+    active_effects: &ActiveEffects,
 ) -> f32 {
     let mut base_destroy_speed = base_destroy_speed(block, tool);
 
@@ -95,11 +98,13 @@ fn destroy_speed(
     // efficiency_level + 1) as f32;     }
     // }
 
-    if let Some(dig_speed_amplifier) = effects::get_dig_speed_amplifier() {
+    if let Some(dig_speed_amplifier) = effects::get_dig_speed_amplifier(active_effects) {
         base_destroy_speed *= 1. + (dig_speed_amplifier + 1) as f32 * 0.2;
     }
 
-    if let Some(dig_slowdown) = effects::get_effect(registry::MobEffect::MiningFatigue) {
+    if let Some(dig_slowdown) =
+        effects::get_effect(active_effects, registry::MobEffect::MiningFatigue)
+    {
         let multiplier = match dig_slowdown {
             0 => 0.3,
             1 => 0.09,
diff --git a/azalea-physics/src/lib.rs b/azalea-physics/src/lib.rs
index ab9c2084..e3b95484 100644
--- a/azalea-physics/src/lib.rs
+++ b/azalea-physics/src/lib.rs
@@ -16,11 +16,11 @@ use azalea_core::{
     tick::GameTick,
 };
 use azalea_entity::{
-    Attributes, EntityKindComponent, HasClientLoaded, Jumping, LocalEntity, LookDirection,
-    OnClimbable, Physics, Pose, Position, dimensions::EntityDimensions, metadata::Sprinting,
-    move_relative,
+    ActiveEffects, Attributes, EntityKindComponent, HasClientLoaded, Jumping, LocalEntity,
+    LookDirection, OnClimbable, Physics, Pose, Position, dimensions::EntityDimensions,
+    metadata::Sprinting, move_relative,
 };
-use azalea_registry::{Block, EntityKind};
+use azalea_registry::{Block, EntityKind, MobEffect};
 use azalea_world::{Instance, InstanceContainer, InstanceName};
 use bevy_app::{App, Plugin, Update};
 use bevy_ecs::prelude::*;
@@ -70,6 +70,7 @@ pub fn ai_step(
             &Position,
             &LookDirection,
             &Sprinting,
+            &ActiveEffects,
             &InstanceName,
             &EntityKindComponent,
         ),
@@ -77,8 +78,16 @@ pub fn ai_step(
     >,
     instance_container: Res<InstanceContainer>,
 ) {
-    for (mut physics, jumping, position, look_direction, sprinting, instance_name, entity_kind) in
-        &mut query
+    for (
+        mut physics,
+        jumping,
+        position,
+        look_direction,
+        sprinting,
+        active_effects,
+        instance_name,
+        entity_kind,
+    ) in &mut query
     {
         let is_player = **entity_kind == EntityKind::Player;
 
@@ -140,6 +149,7 @@ pub fn ai_step(
                             *sprinting,
                             instance_name,
                             &instance_container,
+                            active_effects,
                         );
                         physics.no_jump_delay = 10;
                     }
@@ -331,17 +341,19 @@ pub fn jump_from_ground(
     sprinting: Sprinting,
     instance_name: &InstanceName,
     instance_container: &InstanceContainer,
+    active_effects: &ActiveEffects,
 ) {
     let world_lock = instance_container
         .get(instance_name)
         .expect("All entities should be in a valid world");
     let world = world_lock.read();
 
-    let jump_power: f64 = jump_power(&world, position) as f64 + jump_boost_power();
+    let base_jump = jump_power(&world, position);
+    let jump_power = base_jump + jump_boost_power(active_effects);
     let old_delta_movement = physics.velocity;
     physics.velocity = Vec3 {
         x: old_delta_movement.x,
-        y: jump_power,
+        y: f64::from(jump_power),
         z: old_delta_movement.z,
     };
     if *sprinting {
@@ -504,16 +516,9 @@ fn jump_power(world: &Instance, position: Position) -> f32 {
     0.42 * block_jump_factor(world, position)
 }
 
-fn jump_boost_power() -> f64 {
-    // TODO: potion effects
-    // if let Some(effects) = entity.effects() {
-    //     if let Some(jump_effect) = effects.get(&Effect::Jump) {
-    //         0.1 * (jump_effect.amplifier + 1) as f32
-    //     } else {
-    //         0.
-    //     }
-    // } else {
-    //     0.
-    // }
-    0.
+fn jump_boost_power(active_effects: &ActiveEffects) -> f32 {
+    active_effects
+        .get_level(MobEffect::JumpBoost)
+        .map(|level| 0.1 * (level + 1) as f32)
+        .unwrap_or(0.)
 }
diff --git a/azalea-protocol/src/packets/game/c_update_mob_effect.rs b/azalea-protocol/src/packets/game/c_update_mob_effect.rs
index 201589fb..896d547b 100644
--- a/azalea-protocol/src/packets/game/c_update_mob_effect.rs
+++ b/azalea-protocol/src/packets/game/c_update_mob_effect.rs
@@ -1,4 +1,6 @@
-use azalea_buf::AzBuf;
+use std::io::{Cursor, Write};
+
+use azalea_buf::{AzBuf, AzaleaRead, AzaleaWrite, BufReadError};
 use azalea_protocol_macros::ClientboundGamePacket;
 use azalea_registry::MobEffect;
 use azalea_world::MinecraftEntityId;
@@ -12,5 +14,39 @@ pub struct ClientboundUpdateMobEffect {
     pub effect_amplifier: u32,
     #[var]
     pub effect_duration_ticks: u32,
-    pub flags: u8,
+    pub flags: MobEffectFlags,
+}
+
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct MobEffectFlags {
+    pub ambient: bool,
+    pub show_particles: bool,
+    pub show_icon: bool,
+}
+
+impl AzaleaRead for MobEffectFlags {
+    fn azalea_read(buf: &mut Cursor<&[u8]>) -> Result<Self, BufReadError> {
+        let bits = u8::azalea_read(buf)?;
+        Ok(MobEffectFlags {
+            ambient: bits & 0x01 != 0,
+            show_particles: bits & 0x02 != 0,
+            show_icon: bits & 0x04 != 0,
+        })
+    }
+}
+
+impl AzaleaWrite for MobEffectFlags {
+    fn azalea_write(&self, buf: &mut impl Write) -> std::io::Result<()> {
+        let mut bits = 0;
+        if self.ambient {
+            bits |= 0x01;
+        }
+        if self.show_particles {
+            bits |= 0x02;
+        }
+        if self.show_icon {
+            bits |= 0x04;
+        }
+        bits.azalea_write(buf)
+    }
 }
diff --git a/azalea/src/auto_tool.rs b/azalea/src/auto_tool.rs
index e7eb614d..af3d4352 100644
--- a/azalea/src/auto_tool.rs
+++ b/azalea/src/auto_tool.rs
@@ -1,7 +1,7 @@
 use azalea_block::{BlockState, BlockTrait, fluid_state::FluidKind};
 use azalea_client::{Client, inventory::Inventory};
 use azalea_core::position::BlockPos;
-use azalea_entity::{FluidOnEyes, Physics};
+use azalea_entity::{ActiveEffects, FluidOnEyes, Physics};
 use azalea_inventory::{ItemStack, Menu, components};
 
 use crate::bot::BotClientExt;
@@ -19,10 +19,16 @@ pub trait AutoToolClientExt {
 
 impl AutoToolClientExt for Client {
     fn best_tool_in_hotbar_for_block(&self, block: BlockState) -> BestToolResult {
-        self.query_self::<(&Inventory, &Physics, &FluidOnEyes), _>(
-            |(inventory, physics, fluid_on_eyes)| {
+        self.query_self::<(&Inventory, &Physics, &FluidOnEyes, &ActiveEffects), _>(
+            |(inventory, physics, fluid_on_eyes, active_effects)| {
                 let menu = &inventory.inventory_menu;
-                accurate_best_tool_in_hotbar_for_block(block, menu, physics, fluid_on_eyes)
+                accurate_best_tool_in_hotbar_for_block(
+                    block,
+                    menu,
+                    physics,
+                    fluid_on_eyes,
+                    active_effects,
+                )
             },
         )
     }
@@ -48,11 +54,13 @@ pub fn best_tool_in_hotbar_for_block(block: BlockState, menu: &Menu) -> BestTool
     let mut physics = Physics::default();
     physics.set_on_ground(true);
 
+    let inactive_effects = ActiveEffects::default();
     accurate_best_tool_in_hotbar_for_block(
         block,
         menu,
         &physics,
         &FluidOnEyes::new(FluidKind::Empty),
+        &inactive_effects,
     )
 }
 
@@ -61,6 +69,7 @@ pub fn accurate_best_tool_in_hotbar_for_block(
     menu: &Menu,
     physics: &Physics,
     fluid_on_eyes: &FluidOnEyes,
+    active_effects: &ActiveEffects,
 ) -> BestToolResult {
     let hotbar_slots = &menu.slots()[menu.hotbar_slots_range()];
 
@@ -92,6 +101,7 @@ pub fn accurate_best_tool_in_hotbar_for_block(
                     menu,
                     fluid_on_eyes,
                     physics,
+                    active_effects,
                 ));
             }
             ItemStack::Present(item_stack) => {
@@ -104,6 +114,7 @@ pub fn accurate_best_tool_in_hotbar_for_block(
                         menu,
                         fluid_on_eyes,
                         physics,
+                        active_effects,
                     ));
                 } else {
                     this_item_speed = None;
@@ -127,6 +138,7 @@ pub fn accurate_best_tool_in_hotbar_for_block(
                 menu,
                 fluid_on_eyes,
                 physics,
+                active_effects,
             );
             if this_item_speed > best_speed {
                 best_slot = Some(i);