chunkedge_entity/
attributes.rs

1use std::collections::HashMap;
2
3use bevy_ecs::prelude::*;
4pub use chunkedge_generated::attributes::{EntityAttribute, EntityAttributeOperation};
5use chunkedge_protocol::packets::play::update_attributes_s2c::*;
6use chunkedge_protocol::{Ident, VarInt};
7use indexmap::IndexMap;
8
9/// An instance of an Entity Attribute.
10#[derive(Component, Clone, PartialEq, Debug)]
11pub struct EntityAttributeInstance {
12    /// The attribute.
13    attribute: EntityAttribute,
14    /// The base value of the attribute.
15    base_value: f64,
16    /// The add modifiers of the attribute.
17    add_modifiers: IndexMap<String, f64>,
18    /// The multiply base modifiers of the attribute.
19    multiply_base_modifiers: IndexMap<String, f64>,
20    /// The multiply total modifiers of the attribute.
21    multiply_total_modifiers: IndexMap<String, f64>,
22}
23
24impl EntityAttributeInstance {
25    /// Creates a new instance of an Entity Attribute.
26    pub fn new(attribute: EntityAttribute) -> Self {
27        Self {
28            attribute,
29            base_value: attribute.default_value(),
30            add_modifiers: IndexMap::new(),
31            multiply_base_modifiers: IndexMap::new(),
32            multiply_total_modifiers: IndexMap::new(),
33        }
34    }
35
36    /// Creates a new instance of an Entity Attribute with a value.
37    pub fn new_with_value(attribute: EntityAttribute, base_value: f64) -> Self {
38        Self {
39            attribute,
40            base_value,
41            add_modifiers: IndexMap::new(),
42            multiply_base_modifiers: IndexMap::new(),
43            multiply_total_modifiers: IndexMap::new(),
44        }
45    }
46
47    /// Gets the attribute.
48    pub fn attribute(&self) -> EntityAttribute {
49        self.attribute
50    }
51
52    /// Gets the base value of the attribute.
53    pub fn base_value(&self) -> f64 {
54        self.base_value
55    }
56
57    /// Gets the computed value of the attribute.
58    pub fn compute_value(&self) -> f64 {
59        let mut value = self.base_value;
60
61        // Increment value by modifier
62        for (_, modifier) in &self.add_modifiers {
63            value += modifier;
64        }
65
66        let v = value;
67
68        // Increment value by modifier * v
69        for (_, modifier) in &self.multiply_base_modifiers {
70            value += v * modifier;
71        }
72
73        // Increment value by modifier * value
74        for (_, modifier) in &self.multiply_total_modifiers {
75            value += value * modifier;
76        }
77
78        value.clamp(self.attribute.min_value(), self.attribute.max_value())
79    }
80
81    /// Sets an add modifier.
82    ///
83    /// If the modifier already exists, it will be overwritten.
84    ///
85    /// Returns a mutable reference to self.
86    pub fn with_add_modifier(&mut self, id: &str, modifier: f64) -> &mut Self {
87        self.add_modifiers.insert(id.to_owned(), modifier);
88        self
89    }
90
91    /// Sets a multiply base modifier.
92    ///
93    /// If the modifier already exists, it will be overwritten.
94    ///
95    /// Returns a mutable reference to self.
96    pub fn with_multiply_base_modifier(&mut self, id: &str, modifier: f64) -> &mut Self {
97        self.multiply_base_modifiers.insert(id.to_owned(), modifier);
98        self
99    }
100
101    /// Sets a multiply total modifier.
102    ///
103    /// If the modifier already exists, it will be overwritten.
104    ///
105    /// Returns a mutable reference to self.
106    pub fn with_multiply_total_modifier(&mut self, id: &str, modifier: f64) -> &mut Self {
107        self.multiply_total_modifiers
108            .insert(id.to_owned(), modifier);
109        self
110    }
111
112    /// Sets a value modifier based on the operation.
113    ///
114    /// If the modifier already exists, it will be overwritten.
115    ///
116    /// Returns a mutable reference to self.
117    pub fn with_modifier(
118        &mut self,
119        id: &str,
120        modifier: f64,
121        operation: EntityAttributeOperation,
122    ) -> &mut Self {
123        match operation {
124            EntityAttributeOperation::Add => self.with_add_modifier(id, modifier),
125            EntityAttributeOperation::MultiplyBase => {
126                self.with_multiply_base_modifier(id, modifier)
127            }
128            EntityAttributeOperation::MultiplyTotal => {
129                self.with_multiply_total_modifier(id, modifier)
130            }
131        }
132    }
133
134    /// Removes a modifier.
135    pub fn remove_modifier(&mut self, id: &String) {
136        self.add_modifiers.swap_remove(id);
137        self.multiply_base_modifiers.swap_remove(id);
138        self.multiply_total_modifiers.swap_remove(id);
139    }
140
141    /// Clears all modifiers.
142    pub fn clear_modifiers(&mut self) {
143        self.add_modifiers.clear();
144        self.multiply_base_modifiers.clear();
145        self.multiply_total_modifiers.clear();
146    }
147
148    /// Checks if a modifier exists.
149    pub fn has_modifier(&self, id: &String) -> bool {
150        self.add_modifiers.contains_key(id)
151            || self.multiply_base_modifiers.contains_key(id)
152            || self.multiply_total_modifiers.contains_key(id)
153    }
154
155    /// Converts to a `TrackedEntityProperty` for use in the
156    /// `EntityAttributesS2c` packet.
157    pub(crate) fn to_property(&self) -> TrackedEntityProperty {
158        TrackedEntityProperty {
159            id: self.attribute.get_id().into(),
160            value: self.base_value(),
161            modifiers: self
162                .add_modifiers
163                .iter()
164                .map(|(id, &amount)| TrackedAttributeModifier {
165                    id: id.clone(),
166                    amount,
167                    operation: 0,
168                })
169                .chain(self.multiply_base_modifiers.iter().map(|(id, &amount)| {
170                    TrackedAttributeModifier {
171                        id: id.clone(),
172                        amount,
173                        operation: 1,
174                    }
175                }))
176                .chain(self.multiply_total_modifiers.iter().map(|(id, &amount)| {
177                    TrackedAttributeModifier {
178                        id: id.clone(),
179                        amount,
180                        operation: 2,
181                    }
182                }))
183                .collect(),
184        }
185    }
186}
187
188/// The attributes of a Living Entity.
189#[derive(Component, Clone, PartialEq, Debug, Default)]
190pub struct EntityAttributes {
191    attributes: HashMap<EntityAttribute, EntityAttributeInstance>,
192    recently_changed: Vec<EntityAttribute>,
193}
194
195impl EntityAttributes {
196    /// Gets and clears the recently changed attributes.
197    pub(crate) fn take_recently_changed(&mut self) -> Vec<EntityAttribute> {
198        std::mem::take(&mut self.recently_changed)
199    }
200
201    /// Marks an attribute as recently changed.
202    pub(crate) fn mark_recently_changed(&mut self, attribute: EntityAttribute) {
203        if attribute.tracked() && !self.recently_changed.contains(&attribute) {
204            self.recently_changed.push(attribute);
205        }
206    }
207}
208
209impl EntityAttributes {
210    /// Creates a new instance of `EntityAttributes`.
211    pub fn new() -> Self {
212        Self {
213            attributes: HashMap::new(),
214            recently_changed: Vec::new(),
215        }
216    }
217
218    /// Gets the instance of an attribute.
219    pub fn get(&self, attribute: EntityAttribute) -> Option<&EntityAttributeInstance> {
220        self.attributes.get(&attribute)
221    }
222
223    /// Gets the base value of an attribute.
224    ///
225    /// Returns [`None`] if the attribute does not exist.
226    pub fn get_base_value(&self, attribute: EntityAttribute) -> Option<f64> {
227        self.get(attribute).map(|instance| instance.base_value())
228    }
229
230    /// Gets the computed value of an attribute.
231    ///
232    /// Returns [`None`] if the attribute does not exist.
233    pub fn get_compute_value(&self, attribute: EntityAttribute) -> Option<f64> {
234        self.get(attribute).map(|instance| instance.compute_value())
235    }
236
237    /// Checks if an attribute exists.
238    pub fn has_attribute(&self, attribute: EntityAttribute) -> bool {
239        self.attributes.contains_key(&attribute)
240    }
241
242    /// Creates an attribute if it does not exist.
243    pub fn create_attribute(&mut self, attribute: EntityAttribute) {
244        self.mark_recently_changed(attribute);
245        self.attributes
246            .entry(attribute)
247            .or_insert_with(|| EntityAttributeInstance::new(attribute));
248    }
249
250    /// Creates an attribute if it does not exist and sets its base value.
251    ///
252    /// Returns self.
253    ///
254    /// ## Note
255    ///
256    /// Only to be used in builder-like patterns.
257    pub(crate) fn with_attribute_and_value(
258        mut self,
259        attribute: EntityAttribute,
260        base_value: f64,
261    ) -> Self {
262        self.attributes
263            .entry(attribute)
264            .or_insert_with(|| EntityAttributeInstance::new_with_value(attribute, base_value))
265            .base_value = base_value;
266        self
267    }
268
269    /// Sets the base value of an attribute.
270    pub fn set_base_value(&mut self, attribute: EntityAttribute, value: f64) {
271        self.mark_recently_changed(attribute);
272        self.attributes
273            .entry(attribute)
274            .or_insert_with(|| EntityAttributeInstance::new(attribute))
275            .base_value = value;
276    }
277
278    /// Sets an add modifier of an attribute.
279    pub fn set_add_modifier(&mut self, attribute: EntityAttribute, id: &str, modifier: f64) {
280        self.mark_recently_changed(attribute);
281        self.attributes
282            .entry(attribute)
283            .or_insert_with(|| EntityAttributeInstance::new(attribute))
284            .with_add_modifier(id, modifier);
285    }
286
287    /// Sets a multiply base modifier of an attribute.
288    pub fn set_multiply_base_modifier(
289        &mut self,
290        attribute: EntityAttribute,
291        id: &str,
292        modifier: f64,
293    ) {
294        self.mark_recently_changed(attribute);
295        self.attributes
296            .entry(attribute)
297            .or_insert_with(|| EntityAttributeInstance::new(attribute))
298            .with_multiply_base_modifier(id, modifier);
299    }
300
301    /// Sets a multiply total modifier of an attribute.
302    pub fn set_multiply_total_modifier(
303        &mut self,
304        attribute: EntityAttribute,
305        id: &str,
306        modifier: f64,
307    ) {
308        self.mark_recently_changed(attribute);
309        self.attributes
310            .entry(attribute)
311            .or_insert_with(|| EntityAttributeInstance::new(attribute))
312            .with_multiply_total_modifier(id, modifier);
313    }
314
315    /// Sets a value modifier of an attribute based on the operation.
316    pub fn set_modifier<S: Into<String>>(
317        &mut self,
318        attribute: EntityAttribute,
319        id: S,
320        modifier: f64,
321        operation: EntityAttributeOperation,
322    ) {
323        self.mark_recently_changed(attribute);
324        self.attributes
325            .entry(attribute)
326            .or_insert_with(|| EntityAttributeInstance::new(attribute))
327            .with_modifier(&id.into(), modifier, operation);
328    }
329
330    /// Removes a modifier of an attribute.
331    pub fn remove_modifier<S: Into<String>>(&mut self, attribute: EntityAttribute, id: S) {
332        self.mark_recently_changed(attribute);
333        if let Some(instance) = self.attributes.get_mut(&attribute) {
334            instance.remove_modifier(&id.into());
335        }
336    }
337
338    /// Clears all modifiers of an attribute.
339    pub fn clear_modifiers(&mut self, attribute: EntityAttribute) {
340        self.mark_recently_changed(attribute);
341        if let Some(instance) = self.attributes.get_mut(&attribute) {
342            instance.clear_modifiers();
343        }
344    }
345
346    /// Checks if a modifier exists on an attribute.
347    pub fn has_modifier<S: Into<String>>(&self, attribute: EntityAttribute, id: S) -> bool {
348        self.attributes
349            .get(&attribute)
350            .is_some_and(|inst| inst.has_modifier(&id.into()))
351    }
352
353    /// **For internal use only.**
354    ///
355    /// Converts to a [`Vec`] of [`AttributeProperty`]s.
356    pub fn to_properties(&self) -> Vec<AttributeProperty<'_>> {
357        self.attributes
358            .iter()
359            .filter(|(_, instance)| instance.attribute().tracked())
360            .map(|(_, instance)| instance.to_property().to_property())
361            .collect()
362    }
363}
364
365/// Tracks the attributes of a Living Entity.
366#[derive(Component, Clone, Debug, Default)]
367pub struct TrackedEntityAttributes {
368    /// The attributes that have been modified.
369    modified: IndexMap<EntityAttribute, TrackedEntityProperty>,
370}
371
372#[derive(Clone, Debug)]
373pub(crate) struct TrackedEntityProperty {
374    id: i32,
375    value: f64,
376    modifiers: Vec<TrackedAttributeModifier>,
377}
378
379#[derive(Clone, Debug)]
380pub(crate) struct TrackedAttributeModifier {
381    id: String,
382    amount: f64,
383    operation: u8,
384}
385
386impl TrackedEntityProperty {
387    /// Converts to an [`AttributeProperty`]s.
388    fn to_property(&self) -> AttributeProperty<'static> {
389        AttributeProperty {
390            id: VarInt(self.id),
391            value: self.value,
392            modifiers: self
393                .modifiers
394                .iter()
395                .map(|modifier| AttributeModifier {
396                    id: Ident::new(modifier.id.clone()).unwrap(),
397                    amount: modifier.amount,
398                    operation: modifier.operation,
399                })
400                .collect(),
401        }
402    }
403}
404
405impl TrackedEntityAttributes {
406    /// Creates a new instance of [`TrackedEntityAttributes`].
407    pub fn new() -> Self {
408        Self {
409            modified: IndexMap::new(),
410        }
411    }
412
413    /// Marks an attribute as modified.
414    pub fn mark_modified(&mut self, attributes: &EntityAttributes, attribute: EntityAttribute) {
415        if let Some(instance) = attributes.get(attribute) {
416            self.modified.insert(attribute, instance.to_property());
417        }
418    }
419
420    /// Returns the properties turned into a [`Vec`] of [`AttributeProperty`]s.
421    pub fn get_properties(&self) -> Vec<AttributeProperty<'static>> {
422        self.modified
423            .iter()
424            .map(|(_, property)| property.to_property())
425            .collect()
426    }
427
428    /// Clears the modified attributes.
429    pub fn clear(&mut self) {
430        self.modified.clear();
431    }
432}
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437
438    #[test]
439    fn test_compute_value() {
440        let add_id = "my_attr".to_owned();
441        let mut attributes = EntityAttributes::new();
442        attributes.set_base_value(EntityAttribute::MaxHealth, 20.0);
443        attributes.set_add_modifier(EntityAttribute::MaxHealth, &add_id, 10.0);
444        attributes.set_multiply_base_modifier(EntityAttribute::MaxHealth, "1", 0.2);
445        attributes.set_multiply_base_modifier(EntityAttribute::MaxHealth, "2", 0.2);
446        attributes.set_multiply_total_modifier(EntityAttribute::MaxHealth, "3", 0.5);
447
448        assert_eq!(
449            attributes.get_compute_value(EntityAttribute::MaxHealth),
450            Some(63.0) // ((20 + 10) * (1 + 0.2 + 0.2)) * (1 + 0.5)
451        );
452
453        attributes.remove_modifier(EntityAttribute::MaxHealth, &add_id);
454
455        assert_eq!(
456            attributes.get_compute_value(EntityAttribute::MaxHealth),
457            Some(42.0) // ((20) * (1 + 0.2 + 0.2)) * (1 + 0.5)
458        );
459    }
460}