1#![doc = include_str!("../README.md")]
2
3use bevy_app::prelude::*;
4use bevy_ecs::prelude::*;
5mod interaction_broadcast;
6pub use interaction_broadcast::EquipmentInteractionBroadcast;
7mod inventory_sync;
8pub use inventory_sync::EquipmentInventorySync;
9use valence_server::client::{Client, FlushPacketsSet, LoadEntityForClientEvent};
10use valence_server::entity::living::LivingEntity;
11use valence_server::entity::{EntityId, EntityLayerId, Position};
12use valence_server::protocol::packets::play::entity_equipment_update_s2c::EquipmentEntry;
13use valence_server::protocol::packets::play::EntityEquipmentUpdateS2c;
14use valence_server::protocol::WritePacket;
15use valence_server::{EntityLayer, ItemStack, Layer};
16
17pub struct EquipmentPlugin;
18
19impl Plugin for EquipmentPlugin {
20 fn build(&self, app: &mut App) {
21 app.add_systems(
22 PreUpdate,
23 (
24 on_entity_init,
25 interaction_broadcast::start_interaction,
26 interaction_broadcast::stop_interaction,
27 inventory_sync::on_attach_inventory_sync,
28 inventory_sync::equipment_inventory_sync,
29 inventory_sync::equipment_held_item_sync_from_client,
30 ),
31 )
32 .add_systems(
33 PostUpdate,
34 (
35 update_equipment.before(FlushPacketsSet),
36 on_entity_load.before(FlushPacketsSet),
37 ),
38 )
39 .add_event::<EquipmentChangeEvent>();
40 }
41}
42
43#[derive(Debug, Default, Clone, Component)]
51pub struct Equipment {
52 equipment: [ItemStack; Self::SLOT_COUNT],
53 #[doc(hidden)]
55 pub(crate) changed: u8,
56}
57
58impl Equipment {
59 pub const SLOT_COUNT: usize = 6;
60
61 pub const MAIN_HAND_IDX: u8 = 0;
62 pub const OFF_HAND_IDX: u8 = 1;
63 pub const FEET_IDX: u8 = 2;
64 pub const LEGS_IDX: u8 = 3;
65 pub const CHEST_IDX: u8 = 4;
66 pub const HEAD_IDX: u8 = 5;
67
68 pub fn new(
69 main_hand: ItemStack,
70 off_hand: ItemStack,
71 boots: ItemStack,
72 leggings: ItemStack,
73 chestplate: ItemStack,
74 helmet: ItemStack,
75 ) -> Self {
76 Self {
77 equipment: [main_hand, off_hand, boots, leggings, chestplate, helmet],
78 changed: 0,
79 }
80 }
81
82 pub fn slot(&self, idx: u8) -> &ItemStack {
83 &self.equipment[idx as usize]
84 }
85
86 pub fn set_slot(&mut self, idx: u8, item: ItemStack) {
87 assert!(
88 idx < Self::SLOT_COUNT as u8,
89 "slot index of {idx} out of bounds"
90 );
91 if self.equipment[idx as usize] != item {
92 self.equipment[idx as usize] = item;
93 self.changed |= 1 << idx;
94 }
95 }
96
97 pub fn main_hand(&self) -> &ItemStack {
98 self.slot(Self::MAIN_HAND_IDX)
99 }
100
101 pub fn off_hand(&self) -> &ItemStack {
102 self.slot(Self::OFF_HAND_IDX)
103 }
104
105 pub fn feet(&self) -> &ItemStack {
106 self.slot(Self::FEET_IDX)
107 }
108
109 pub fn legs(&self) -> &ItemStack {
110 self.slot(Self::LEGS_IDX)
111 }
112
113 pub fn chest(&self) -> &ItemStack {
114 self.slot(Self::CHEST_IDX)
115 }
116
117 pub fn head(&self) -> &ItemStack {
118 self.slot(Self::HEAD_IDX)
119 }
120
121 pub fn set_main_hand(&mut self, item: ItemStack) {
122 self.set_slot(Self::MAIN_HAND_IDX, item);
123 }
124
125 pub fn set_off_hand(&mut self, item: ItemStack) {
126 self.set_slot(Self::OFF_HAND_IDX, item);
127 }
128
129 pub fn set_feet(&mut self, item: ItemStack) {
130 self.set_slot(Self::FEET_IDX, item);
131 }
132
133 pub fn set_legs(&mut self, item: ItemStack) {
134 self.set_slot(Self::LEGS_IDX, item);
135 }
136
137 pub fn set_chest(&mut self, item: ItemStack) {
138 self.set_slot(Self::CHEST_IDX, item);
139 }
140
141 pub fn set_head(&mut self, item: ItemStack) {
142 self.set_slot(Self::HEAD_IDX, item);
143 }
144
145 pub fn clear(&mut self) {
146 for slot in 0..Self::SLOT_COUNT as u8 {
147 self.set_slot(slot, ItemStack::EMPTY);
148 }
149 }
150
151 pub fn is_default(&self) -> bool {
152 self.equipment.iter().all(|item| item.is_empty())
153 }
154}
155
156#[derive(Debug, Clone)]
157pub struct EquipmentSlotChange {
158 idx: u8,
159 stack: ItemStack,
160}
161
162#[derive(Debug, Clone, Event)]
163pub struct EquipmentChangeEvent {
164 pub client: Entity,
165 pub changed: Vec<EquipmentSlotChange>,
166}
167
168fn update_equipment(
169 mut clients: Query<
170 (Entity, &EntityId, &EntityLayerId, &Position, &mut Equipment),
171 Changed<Equipment>,
172 >,
173 mut event_writer: EventWriter<EquipmentChangeEvent>,
174 mut entity_layer: Query<&mut EntityLayer>,
175) {
176 for (entity, entity_id, entity_layer_id, position, mut equipment) in &mut clients {
177 let Ok(mut entity_layer) = entity_layer.get_mut(entity_layer_id.0) else {
178 continue;
179 };
180
181 if equipment.changed != 0 {
182 let mut slots_changed: Vec<EquipmentSlotChange> =
183 Vec::with_capacity(Equipment::SLOT_COUNT);
184
185 for slot in 0..Equipment::SLOT_COUNT {
186 if equipment.changed & (1 << slot) != 0 {
187 slots_changed.push(EquipmentSlotChange {
188 idx: slot as u8,
189 stack: equipment.equipment[slot].clone(),
190 });
191 }
192 }
193
194 entity_layer
195 .view_except_writer(position.0, entity)
196 .write_packet(&EntityEquipmentUpdateS2c {
197 entity_id: entity_id.get().into(),
198 equipment: slots_changed
199 .iter()
200 .map(|change| EquipmentEntry {
201 slot: change.idx as i8,
202 item: change.stack.clone(),
203 })
204 .collect(),
205 });
206
207 event_writer.send(EquipmentChangeEvent {
208 client: entity,
209 changed: slots_changed,
210 });
211
212 equipment.changed = 0;
213 }
214 }
215}
216
217fn on_entity_load(
220 mut clients: Query<&mut Client>,
221 entities: Query<(&EntityId, &Equipment)>,
222 mut events: EventReader<LoadEntityForClientEvent>,
223) {
224 for event in events.read() {
225 let Ok(mut client) = clients.get_mut(event.client) else {
226 continue;
227 };
228
229 let Ok((entity_id, equipment)) = entities.get(event.entity_loaded) else {
230 continue;
231 };
232
233 if equipment.is_default() {
234 continue;
235 }
236
237 let mut entries: Vec<EquipmentEntry> = Vec::with_capacity(Equipment::SLOT_COUNT);
238 for (idx, stack) in equipment.equipment.iter().enumerate() {
239 entries.push(EquipmentEntry {
240 slot: idx as i8,
241 item: stack.clone(),
242 });
243 }
244
245 client.write_packet(&EntityEquipmentUpdateS2c {
246 entity_id: entity_id.get().into(),
247 equipment: entries,
248 });
249 }
250}
251
252fn on_entity_init(
255 mut commands: Commands,
256 mut entities: Query<Entity, (Added<LivingEntity>, Without<Equipment>)>,
257) {
258 for entity in &mut entities {
259 commands.entity(entity).insert(Equipment::default());
260 }
261}