valence_entity/
attributes.rs

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