diff --git a/crates/valence/examples/random_equipment.rs b/crates/valence/examples/random_equipment.rs new file mode 100644 index 000000000..d60dd5d9e --- /dev/null +++ b/crates/valence/examples/random_equipment.rs @@ -0,0 +1,117 @@ +use bevy_app::CoreStage; +use rand::Rng; +use valence::client::despawn_disconnected_clients; +use valence::client::event::default_event_handler; +use valence::equipment::{Equipment, EquipmentSlot}; +use valence::prelude::*; + +const BOARD_MIN_X: i32 = -30; +const BOARD_MAX_X: i32 = 30; +const BOARD_MIN_Z: i32 = -30; +const BOARD_MAX_Z: i32 = 30; +const BOARD_Y: i32 = 64; + +const SPAWN_POS: DVec3 = DVec3::new( + (BOARD_MIN_X + BOARD_MAX_X) as f64 / 2.0, + BOARD_Y as f64 + 1.0, + (BOARD_MIN_Z + BOARD_MAX_Z) as f64 / 2.0, +); + +pub fn main() { + tracing_subscriber::fmt().init(); + + App::new() + .add_plugin(ServerPlugin::new(()).with_connection_mode(ConnectionMode::Offline)) + .add_system_to_stage(EventLoop, default_event_handler) + .add_system_set(PlayerList::default_system_set()) + .add_startup_system(setup) + .add_system(init_clients) + .add_system(despawn_disconnected_clients) + .add_system(randomize_equipment) + .run(); +} + +fn setup(world: &mut World) { + let mut instance = world + .resource::() + .new_instance(DimensionId::default()); + + for z in -10..10 { + for x in -10..10 { + instance.insert_chunk([x, z], Chunk::default()); + } + } + + for z in BOARD_MIN_Z..=BOARD_MAX_Z { + for x in BOARD_MIN_X..=BOARD_MAX_X { + instance.set_block([x, BOARD_Y, z], BlockState::DIRT); + } + } + + let instance = world.spawn(instance); + let instance_entity = instance.id(); + + let mut equipment = Equipment::default(); + equipment.set( + ItemStack::new(ItemKind::IronBoots, 1, None), + EquipmentSlot::Boots, + ); + + // Spawn armor stand + let mut armor_stand = world.spawn(( + McEntity::new(EntityKind::ArmorStand, instance_entity), + equipment, + )); + + if let Some(mut armor_stand) = armor_stand.get_mut::() { + let position = [SPAWN_POS.x, SPAWN_POS.y, SPAWN_POS.z + 3.0]; + armor_stand.set_position(position); + armor_stand.set_yaw(180.0); + } +} + +fn init_clients( + mut clients: Query<(&mut Client, Entity), Added>, + instances: Query>, + mut commands: Commands, +) { + let instance = instances.single(); + + for (mut client, entity) in &mut clients { + client.set_position(SPAWN_POS); + client.set_instance(instances.single()); + client.set_game_mode(GameMode::Creative); + + let equipment = Equipment::default(); + + commands.entity(entity).insert(( + equipment, + McEntity::with_uuid(EntityKind::Player, instance, client.uuid()), + )); + } +} + +fn randomize_equipment(mut query: Query<&mut Equipment>, server: Res) { + let ticks = server.current_tick(); + if ticks % server.tps() != 0 { + return; + } + + for mut equips in &mut query { + equips.clear(); + + let (slot, item_kind) = match rand::thread_rng().gen_range(0..=5) { + 0 => (EquipmentSlot::MainHand, ItemKind::DiamondSword), + 1 => (EquipmentSlot::OffHand, ItemKind::Torch), + 2 => (EquipmentSlot::Boots, ItemKind::IronBoots), + 3 => (EquipmentSlot::Leggings, ItemKind::DiamondLeggings), + 4 => (EquipmentSlot::Chestplate, ItemKind::ChainmailChestplate), + 5 => (EquipmentSlot::Helmet, ItemKind::LeatherHelmet), + _ => (EquipmentSlot::Boots, ItemKind::IronBoots), + }; + + let item = ItemStack::new(item_kind, 1, None); + + equips.set(item, slot); + } +} diff --git a/crates/valence/src/client.rs b/crates/valence/src/client.rs index 766990854..35e4f4db6 100644 --- a/crates/valence/src/client.rs +++ b/crates/valence/src/client.rs @@ -28,6 +28,7 @@ use valence_protocol::{ use crate::dimension::DimensionId; use crate::entity::data::Player; use crate::entity::{velocity_to_packet_units, EntityStatus, McEntity}; +use crate::equipment::Equipment; use crate::instance::Instance; use crate::packet::WritePacket; use crate::server::{NewClientInfo, Server}; @@ -626,6 +627,7 @@ pub(crate) fn update_clients( mut clients: Query<(Entity, &mut Client, Option<&McEntity>)>, instances: Query<&Instance>, entities: Query<&McEntity>, + equipment: Query<&Equipment>, ) { // TODO: what batch size to use? clients.par_for_each_mut(16, |(entity_id, mut client, self_entity)| { @@ -636,6 +638,7 @@ pub(crate) fn update_clients( entity_id, &instances, &entities, + &equipment, &server, ) { client.write_packet(&DisconnectPlay { @@ -662,6 +665,7 @@ fn update_one_client( _self_id: Entity, instances: &Query<&Instance>, entities: &Query<&McEntity>, + equipment: &Query<&Equipment>, server: &Server, ) -> anyhow::Result<()> { let Ok(instance) = instances.get(client.instance) else { @@ -804,6 +808,7 @@ fn update_one_client( &mut client.enc, entity.old_position(), &mut client.scratch, + equipment.get(id).ok(), ); } } @@ -881,6 +886,7 @@ fn update_one_client( &mut client.enc, entity.position(), &mut client.scratch, + equipment.get(id).ok(), ); } } @@ -934,6 +940,7 @@ fn update_one_client( &mut client.enc, entity.position(), &mut client.scratch, + equipment.get(id).ok(), ); } } diff --git a/crates/valence/src/entity.rs b/crates/valence/src/entity.rs index 4517e4d88..dcd62e277 100644 --- a/crates/valence/src/entity.rs +++ b/crates/valence/src/entity.rs @@ -15,9 +15,11 @@ use valence_protocol::packets::s2c::play::{ SetHeadRotation, SpawnEntity, SpawnExperienceOrb, SpawnPlayer, TeleportEntity, UpdateEntityPosition, UpdateEntityPositionAndRotation, UpdateEntityRotation, }; +use valence_protocol::packets::s2c::set_equipment::SetEquipment; use valence_protocol::{ByteAngle, RawBytes, VarInt}; use crate::config::DEFAULT_TPS; +use crate::equipment::Equipment; use crate::math::Aabb; use crate::packet::WritePacket; use crate::{Despawned, NULL_ENTITY}; @@ -611,9 +613,11 @@ impl McEntity { mut writer: impl WritePacket, position: DVec3, scratch: &mut Vec, + equipment: Option<&Equipment>, ) { + let entity_id = VarInt(self.protocol_id); let with_object_data = |data| SpawnEntity { - entity_id: VarInt(self.protocol_id), + entity_id, object_uuid: self.uuid, kind: VarInt(self.kind() as i32), position: position.to_array(), @@ -678,6 +682,16 @@ impl McEntity { metadata: RawBytes(scratch), }); } + + // If entity has equipment, send it to the client + if let Some(equipment) = equipment { + if !equipment.is_empty() { + writer.write_packet(&SetEquipment { + entity_id, + equipment: equipment.equipment().clone(), + }) + } + } } /// Writes the appropriate packets to update the entity (Position, tracked diff --git a/crates/valence/src/equipment.rs b/crates/valence/src/equipment.rs new file mode 100644 index 000000000..304501540 --- /dev/null +++ b/crates/valence/src/equipment.rs @@ -0,0 +1,374 @@ +use bevy_ecs::prelude::*; +use bevy_ecs::{query::Changed, system::Query}; +use valence_protocol::packets::s2c::set_equipment::SetEquipment; +use valence_protocol::VarInt; +use valence_protocol::{packets::s2c::set_equipment::EquipmentEntry, ItemStack}; + +use crate::prelude::*; +use crate::view::ChunkPos; + +/// ECS component to be added for entities with equipment. +/// +/// Equipment updates managed by `update_equipment`. +#[derive(Component, Default, PartialEq, Debug)] +pub struct Equipment { + equipment: Vec, + /// Bit set with the modified equipment slots + modified_slots: u8, +} + +#[derive(Copy, Clone)] +pub enum EquipmentSlot { + MainHand, + OffHand, + Boots, + Leggings, + Chestplate, + Helmet, +} + +impl Equipment { + pub fn new() -> Equipment { + Equipment::default() + } + + /// Set an equipment slot with an item stack + pub fn set(&mut self, item: ItemStack, slot: EquipmentSlot) { + if let Some(equip) = self.get_mut(slot) { + equip.item = Some(item); + } else { + self.equipment.push(EquipmentEntry { + item: Some(item), + slot: slot.into(), + }); + } + + self.set_modified_slot(slot); + } + + /// Remove all equipment + pub fn clear(&mut self) { + for equip in self.equipment.iter() { + self.modified_slots |= 1 << equip.slot as u8; + } + + self.equipment.clear(); + } + + /// Remove an equipment from a slot and return it if present + pub fn remove(&mut self, slot: EquipmentSlot) -> Option { + let slot_id: i8 = slot.into(); + + if let Some(idx) = self + .equipment + .iter() + .position(|equip| equip.slot == slot_id) + { + self.set_modified_slot(slot); + Some(self.equipment.remove(idx)) + } else { + None + } + } + + pub fn get_mut(&mut self, slot: EquipmentSlot) -> Option<&mut EquipmentEntry> { + let slot: i8 = slot.into(); + self.equipment.iter_mut().find(|equip| equip.slot == slot) + } + + pub fn get(&self, slot: EquipmentSlot) -> Option<&EquipmentEntry> { + let slot: i8 = slot.into(); + self.equipment.iter().find(|equip| equip.slot == slot) + } + + pub fn equipment(&self) -> &Vec { + &self.equipment + } + + pub fn is_empty(&self) -> bool { + self.equipment.is_empty() + } + + fn has_modified_slots(&self) -> bool { + self.modified_slots != 0 + } + + fn iter_modified_equipment(&self) -> impl Iterator + '_ { + self.iter_modified_slots().map(|slot| { + self.get(slot).cloned().unwrap_or_else(|| EquipmentEntry { + slot: slot.into(), + item: None, + }) + }) + } + + fn iter_modified_slots(&self) -> impl Iterator { + let modified_slots = self.modified_slots; + + (0..=5).filter_map(move |slot: i8| { + if modified_slots & (1 << slot) != 0 { + Some(EquipmentSlot::try_from(slot).unwrap()) + } else { + None + } + }) + } + + fn set_modified_slot(&mut self, slot: EquipmentSlot) { + let shifts: i8 = slot.into(); + self.modified_slots |= 1 << (shifts as u8); + } + + fn clear_modified_slot(&mut self) { + self.modified_slots = 0; + } +} + +impl TryFrom for EquipmentSlot { + type Error = &'static str; + + /// Convert from `id` according to + fn try_from(id: i8) -> Result { + let slot = match id { + 0 => EquipmentSlot::MainHand, + 1 => EquipmentSlot::OffHand, + 2 => EquipmentSlot::Boots, + 3 => EquipmentSlot::Leggings, + 4 => EquipmentSlot::Chestplate, + 5 => EquipmentSlot::Helmet, + _ => return Err("Invalid value"), + }; + + Ok(slot) + } +} + +impl From for i8 { + /// Convert to `id` according to + fn from(slot: EquipmentSlot) -> Self { + match slot { + EquipmentSlot::MainHand => 0, + EquipmentSlot::OffHand => 1, + EquipmentSlot::Boots => 2, + EquipmentSlot::Leggings => 3, + EquipmentSlot::Chestplate => 4, + EquipmentSlot::Helmet => 5, + } + } +} + +/// When a [Equipment] component is changed, send [SetEquipment] packet to all clients +/// that have the updated entity in their view distance. +/// +/// NOTE: [SetEquipment] packet only have cosmetic effect, which means it does not affect armor resistance or damage. +pub(crate) fn update_equipment( + mut equiped_entities: Query<(&McEntity, &mut Equipment), Changed>, + mut instances: Query<&mut Instance>, +) { + for (equiped_mc_entity, mut equips) in &mut equiped_entities { + if !equips.has_modified_slots() { + continue; + } + + let instance = equiped_mc_entity.instance(); + let chunk_pos = ChunkPos::from_dvec3(equiped_mc_entity.position()); + + if let Ok(mut instance) = instances.get_mut(instance) { + instance.write_packet_at( + &SetEquipment { + entity_id: VarInt(equiped_mc_entity.protocol_id()), + equipment: equips.iter_modified_equipment().collect(), + }, + chunk_pos, + ) + } + + equips.clear_modified_slot(); + } +} + +#[cfg(test)] +mod test { + use valence_protocol::packets::S2cPlayPacket; + + use crate::unit_test::util::scenario_single_client; + + use super::*; + + #[test] + fn test_set_boots_and_clear() { + let mut equipment = Equipment::default(); + assert_eq!( + equipment, + Equipment { + equipment: vec![], + modified_slots: 0 + } + ); + + let item = ItemStack::new(ItemKind::GreenWool, 1, None); + let slot = EquipmentSlot::Boots; + equipment.set(item.clone(), slot); + + assert_eq!( + equipment, + Equipment { + equipment: vec![EquipmentEntry { + slot: slot.into(), + item: Some(item) + }], + modified_slots: 0b100 + } + ); + + equipment.clear_modified_slot(); + equipment.clear(); + assert_eq!( + equipment, + Equipment { + equipment: vec![], + modified_slots: 0b100 + } + ); + assert_eq!( + equipment + .iter_modified_equipment() + .collect::>(), + vec![EquipmentEntry { + slot: slot.into(), + item: None + }] + ); + } + + #[test] + fn test_set_main_hand_and_remove_it() { + let mut equipment = Equipment::default(); + + let item = ItemStack::new(ItemKind::DiamondSword, 1, None); + let slot = EquipmentSlot::MainHand; + equipment.set(item.clone(), slot); + + assert_eq!( + equipment.remove(EquipmentSlot::MainHand), + Some(EquipmentEntry { + slot: slot.into(), + item: Some(item) + }) + ); + assert_eq!(equipment.remove(EquipmentSlot::Helmet), None); + assert_eq!( + equipment, + Equipment { + equipment: vec![], + modified_slots: 0b1 + } + ); + } + + #[test] + fn test_set_equipment_sent_packets() -> anyhow::Result<()> { + let mut app = App::new(); + let (client_ent, mut client_helper) = scenario_single_client(&mut app); + + // Setup server + let mut instance = app + .world + .resource::() + .new_instance(DimensionId::default()); + instance.insert_chunk([0, 0], Default::default()); + let instance = app.world.spawn(instance); + let instance_entity = instance.id(); + + // Setup client + let mut client = app.world.get_mut::(client_ent).unwrap(); + let uuid = client.uuid(); + client.set_position([0.0, 0.0, 0.0]); + client.set_instance(instance_entity); + let mut client_ent_mut = app + .world + .get_entity_mut(client_ent) + .expect("should have client component"); + client_ent_mut.insert(McEntity::with_uuid(EntityKind::Player, client_ent, uuid)); + + // Spawn armor stand + let equipment = Equipment::default(); + let mut mc_entity = McEntity::new(EntityKind::ArmorStand, instance_entity); + mc_entity.set_position([0.0, 0.0, 0.0]); + let armor_stand = app.world.spawn((mc_entity, equipment)); + let armor_stand = armor_stand.id(); + + // Process a tick to get past the "on join" logic. + app.update(); + client_helper.clear_sent(); + + // Set armor stand boots + let mut equipments = app + .world + .get_mut::(armor_stand) + .expect("should have Equipment component"); + let item = ItemStack::new(ItemKind::IronBoots, 1, None); + let slot = EquipmentSlot::Boots; + equipments.set(item.clone(), slot); + let armor_stand_entity_id = VarInt( + app.world + .get_mut::(armor_stand) + .expect("should have McEntity component") + .protocol_id(), + ); + + // Assert packet + app.update(); + let sent_packets = client_helper.collect_sent()?; + if let S2cPlayPacket::SetEquipment(packet) = &sent_packets[0] { + assert_eq!( + packet, + &SetEquipment { + entity_id: armor_stand_entity_id, + equipment: vec![EquipmentEntry { + slot: slot.into(), + item: Some(item) + }] + } + ) + } + + // Set up for next tick + app.update(); + client_helper.clear_sent(); + + // Remove boots and set main hand + let mut equipments = app + .world + .get_mut::(armor_stand) + .expect("should have Equipment component"); + equipments.remove(EquipmentSlot::Boots); + let item = ItemStack::new(ItemKind::DiamondSword, 1, None); + let slot = EquipmentSlot::MainHand; + equipments.set(item.clone(), slot); + + // Assert new packets + app.update(); + let sent_packets = client_helper.collect_sent()?; + if let S2cPlayPacket::SetEquipment(packet) = &sent_packets[0] { + assert_eq!( + packet, + &SetEquipment { + entity_id: armor_stand_entity_id, + equipment: vec![ + EquipmentEntry { + slot: slot.into(), + item: Some(item) + }, + EquipmentEntry { + slot: EquipmentSlot::Boots.into(), + item: None + }, + ] + } + ) + } + + Ok(()) + } +} diff --git a/crates/valence/src/lib.rs b/crates/valence/src/lib.rs index f3a704ce3..a8ac12cb4 100644 --- a/crates/valence/src/lib.rs +++ b/crates/valence/src/lib.rs @@ -36,6 +36,7 @@ pub mod client; pub mod config; pub mod dimension; pub mod entity; +pub mod equipment; pub mod instance; pub mod inventory; pub mod math; diff --git a/crates/valence/src/server.rs b/crates/valence/src/server.rs index 0d7e3a738..8447f635e 100644 --- a/crates/valence/src/server.rs +++ b/crates/valence/src/server.rs @@ -29,6 +29,7 @@ use crate::entity::{ check_entity_invariants, deinit_despawned_entities, init_entities, update_entities, McEntityManager, }; +use crate::equipment::update_equipment; use crate::instance::{ check_instance_invariants, update_instances_post_client, update_instances_pre_client, Instance, }; @@ -361,6 +362,7 @@ pub fn build_plugin( .with_system(handle_close_container) .with_system(update_client_on_close_inventory.after(update_open_inventories)) .with_system(update_player_inventories) + .with_system(update_equipment) .with_system( handle_click_container .before(update_open_inventories)