Skip to main content

steel_core/entity/
attribute.rs

1//! Runtime entity attribute system.
2//!
3use core::iter;
4
5use simdnbt::owned::{NbtCompound, NbtList};
6use steel_protocol::packets::game::{AttributeModifierData, AttributeSnapshot};
7pub use steel_registry::attribute::AttributeModifierOperation;
8use steel_registry::attribute::AttributeRef;
9use steel_registry::entity_type::EntityTypeRef;
10use steel_registry::{REGISTRY, RegistryEntry, RegistryExt};
11use steel_utils::Identifier;
12
13/// Growable bitmask for tracking dirty attribute IDs.
14pub struct DirtySet {
15    chunks: Vec<u64>,
16}
17
18impl Default for DirtySet {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24impl DirtySet {
25    /// Creates an empty dirty set
26    #[must_use]
27    pub const fn new() -> Self {
28        Self { chunks: Vec::new() }
29    }
30
31    /// Marks an attribute ID as dirty
32    pub fn mark(&mut self, id: u16) {
33        let chunk = id as usize / 64;
34        let bit = id as usize % 64;
35        if chunk >= self.chunks.len() {
36            self.chunks.resize(chunk + 1, 0);
37        }
38        self.chunks[chunk] |= 1 << bit;
39    }
40
41    /// Returns `true` if no attributes are dirty
42    #[must_use]
43    pub fn is_empty(&self) -> bool {
44        self.chunks.iter().all(|&c| c == 0)
45    }
46
47    /// Drains all marked IDs, clearing the set
48    pub fn drain(&mut self) -> impl Iterator<Item = u16> + '_ {
49        self.chunks.iter_mut().enumerate().flat_map(|(i, chunk)| {
50            let mut bits = *chunk;
51            *chunk = 0;
52            iter::from_fn(move || {
53                if bits == 0 {
54                    return None;
55                }
56                let bit = bits.trailing_zeros() as u16;
57                bits &= bits - 1;
58                Some(i as u16 * 64 + bit)
59            })
60        })
61    }
62}
63
64/// A modifier applied to an attribute instance
65#[derive(Clone, Debug)]
66pub struct AttributeModifier {
67    /// Unique identifier (e.g. `minecraft:sprinting`)
68    pub id: Identifier,
69    /// The modifier value
70    pub amount: f64,
71    /// How the modifier is applied during calculation
72    pub operation: AttributeModifierOperation,
73}
74
75/// Runtime state for a single attribute on an entity
76pub struct AttributeInstance {
77    attribute: AttributeRef,
78    base_value: f64,
79    modifiers: Vec<AttributeModifier>,
80    /// Parallel to `modifiers` `true` means the modifier survives serialization
81    persistent: Vec<bool>,
82    cached_value: f64,
83}
84
85impl AttributeInstance {
86    const fn new(attribute: AttributeRef, base_value: f64) -> Self {
87        let cached = attribute.sanitize_value(base_value);
88        Self {
89            attribute,
90            base_value,
91            modifiers: Vec::new(),
92            persistent: Vec::new(),
93            cached_value: cached,
94        }
95    }
96
97    /// Returns the attribute definition
98    #[must_use]
99    pub const fn attribute(&self) -> AttributeRef {
100        self.attribute
101    }
102
103    /// Returns the base value before modifiers
104    #[must_use]
105    pub const fn base_value(&self) -> f64 {
106        self.base_value
107    }
108
109    /// Sets the base value and recalculates. Returns `true` if changed
110    #[expect(
111        clippy::float_cmp,
112        reason = "vanilla uses exact base value equality for dirty checks"
113    )]
114    pub fn set_base_value(&mut self, value: f64) -> bool {
115        if self.base_value == value {
116            return false;
117        }
118        self.base_value = value;
119        self.recalculate();
120        true
121    }
122
123    /// Returns the final calculated value (base + modifiers, clamped)
124    #[must_use]
125    pub const fn value(&self) -> f64 {
126        self.cached_value
127    }
128
129    /// Adds a modifier. Returns `false` if a modifier with this ID already exists
130    pub fn add_modifier(&mut self, modifier: AttributeModifier, persistent: bool) -> bool {
131        if self.modifiers.iter().any(|m| m.id == modifier.id) {
132            return false;
133        }
134        self.modifiers.push(modifier);
135        self.persistent.push(persistent);
136        self.recalculate();
137        true
138    }
139
140    /// Returns whether a modifier with the given ID exists.
141    #[must_use]
142    pub fn has_modifier(&self, id: &Identifier) -> bool {
143        self.modifiers.iter().any(|modifier| modifier.id == *id)
144    }
145
146    /// Adds or replaces a modifier. Returns `true` if the value actually changed.
147    #[expect(
148        clippy::float_cmp,
149        reason = "exact equality is intentional — we want to skip recalculation when the modifier is identical"
150    )]
151    pub fn set_modifier(&mut self, modifier: AttributeModifier, persistent: bool) -> bool {
152        if let Some(idx) = self.modifiers.iter().position(|m| m.id == modifier.id) {
153            let existing = &self.modifiers[idx];
154            if existing.amount == modifier.amount
155                && existing.operation == modifier.operation
156                && self.persistent[idx] == persistent
157            {
158                return false;
159            }
160            self.modifiers[idx] = modifier;
161            self.persistent[idx] = persistent;
162        } else {
163            self.modifiers.push(modifier);
164            self.persistent.push(persistent);
165        }
166        self.recalculate();
167        true
168    }
169
170    /// Removes a modifier by ID, Returns `true` if it existed
171    pub fn remove_modifier(&mut self, id: &Identifier) -> bool {
172        let Some(idx) = self.modifiers.iter().position(|m| m.id == *id) else {
173            return false;
174        };
175        self.modifiers.swap_remove(idx);
176        self.persistent.swap_remove(idx);
177        self.recalculate();
178        true
179    }
180
181    /// Returns an iterator over permanent modifiers (for serialization)
182    pub fn permanent_modifiers(&self) -> impl Iterator<Item = &AttributeModifier> {
183        self.modifiers
184            .iter()
185            .zip(self.persistent.iter())
186            .filter(|&(_, p)| *p)
187            .map(|(m, _)| m)
188    }
189
190    /// Removes all transient (non-persistent) modifiers. Returns `true` if any were removed
191    fn remove_transient_modifiers(&mut self) -> bool {
192        if !self.persistent.iter().any(|&p| !p) {
193            return false;
194        }
195        let mut i = 0;
196        while i < self.modifiers.len() {
197            if self.persistent[i] {
198                i += 1;
199            } else {
200                self.modifiers.swap_remove(i);
201                self.persistent.swap_remove(i);
202            }
203        }
204        self.recalculate();
205        true
206    }
207
208    /// Three-phase vanilla calculation:
209    /// 1. `ADD_VALUE`
210    /// 2. `ADD_MULTIPLIED_BASE`
211    /// 3. `ADD_MULTIPLIED_TOTAL`
212    fn recalculate(&mut self) {
213        let mut base = self.base_value;
214        for m in &self.modifiers {
215            if m.operation == AttributeModifierOperation::AddValue {
216                base += m.amount;
217            }
218        }
219
220        let mut result = base;
221        for m in &self.modifiers {
222            if m.operation == AttributeModifierOperation::AddMultipliedBase {
223                result += base * m.amount;
224            }
225        }
226        for m in &self.modifiers {
227            if m.operation == AttributeModifierOperation::AddMultipliedTotal {
228                result *= 1.0 + m.amount;
229            }
230        }
231
232        self.cached_value = self.attribute.sanitize_value(result);
233    }
234
235    /// Builds a network snapshot for `CUpdateAttributes`
236    fn to_snapshot(&self, attribute_id: i32) -> AttributeSnapshot {
237        AttributeSnapshot {
238            attribute_id,
239            base_value: self.base_value,
240            modifiers: self
241                .modifiers
242                .iter()
243                .map(|m| AttributeModifierData {
244                    id: m.id.clone(),
245                    amount: m.amount,
246                    operation: m.operation,
247                })
248                .collect(),
249        }
250    }
251}
252
253/// Per-entity container for attribute instances with dirty tracking
254///
255/// Indexed by attribute registry ID. Two dirty sets match vanilla's design:
256/// - `to_update`: all dirty attributes, drained for server-side effects
257/// - `to_sync`: syncable dirty attributes, drained for network packets
258pub struct AttributeMap {
259    instances: Vec<Option<AttributeInstance>>,
260    to_update: DirtySet,
261    to_sync: DirtySet,
262}
263
264impl AttributeMap {
265    /// Creates an `AttributeMap` from an entity type's default attributes
266    ///
267    /// # Panics
268    /// Panics if generated entity default attributes reference an attribute
269    /// that is missing from the generated vanilla registry.
270    // TODO: Add AttributeSupplier for lazy instantiation when mob entities are implemented
271    #[must_use]
272    pub fn new_for_entity(entity_type: EntityTypeRef) -> Self {
273        let attr_count = REGISTRY.attributes.len();
274        let mut instances = Vec::with_capacity(attr_count);
275        instances.resize_with(attr_count, || None);
276
277        for &(attr_name, base_value) in entity_type.default_attributes {
278            let key = Identifier::vanilla_static(attr_name);
279            let Some(id) = REGISTRY.attributes.id_from_key(&key) else {
280                panic!(
281                    "default attributes for entity type {} reference unregistered attribute {key}",
282                    entity_type.key
283                );
284            };
285            let Some(attr) = REGISTRY.attributes.by_id(id) else {
286                panic!("attribute registry id {id} for default attribute {key} is not registered");
287            };
288            instances[id] = Some(AttributeInstance::new(attr, base_value));
289        }
290
291        Self {
292            instances,
293            to_update: DirtySet::new(),
294            to_sync: DirtySet::new(),
295        }
296    }
297
298    /// Returns `true` if the entity has this attribute registered
299    #[must_use]
300    pub fn has_attribute(&self, attribute: AttributeRef) -> bool {
301        attribute
302            .try_id()
303            .and_then(|id| self.instances.get(id))
304            .is_some_and(Option::is_some)
305    }
306
307    /// Gets the calculated value of an attribute
308    #[must_use]
309    pub fn get_value(&self, attribute: AttributeRef) -> Option<f64> {
310        let id = attribute.try_id()?;
311        self.instances
312            .get(id)?
313            .as_ref()
314            .map(AttributeInstance::value)
315    }
316
317    /// Gets the calculated value of an attribute that must exist on this entity.
318    ///
319    /// # Panics
320    ///
321    /// Panics if the entity type was not constructed with the requested
322    /// attribute. Vanilla's `AttributeSupplier.getValue` is a hard failure for
323    /// missing attributes; using this keeps required living attributes from
324    /// silently falling back to unrelated defaults.
325    #[must_use]
326    pub fn required_value(&self, attribute: AttributeRef) -> f64 {
327        let Some(value) = self.get_value(attribute) else {
328            panic!("required attribute {} is missing", attribute.key);
329        };
330        value
331    }
332
333    /// Gets the base value of an attribute
334    #[must_use]
335    pub fn get_base_value(&self, attribute: AttributeRef) -> Option<f64> {
336        let id = attribute.try_id()?;
337        self.instances
338            .get(id)?
339            .as_ref()
340            .map(AttributeInstance::base_value)
341    }
342
343    /// Gets a reference to an attribute instance
344    #[must_use]
345    pub fn get_instance(&self, attribute: AttributeRef) -> Option<&AttributeInstance> {
346        let id = attribute.try_id()?;
347        self.instances.get(id)?.as_ref()
348    }
349
350    /// Serializes the vanilla `LivingEntity.attributes` list.
351    #[must_use]
352    pub(crate) fn to_vanilla_nbt(&self) -> NbtList {
353        let attributes = self
354            .instances
355            .iter()
356            .flatten()
357            .map(|instance| {
358                let mut attribute = NbtCompound::new();
359                attribute.insert("id", instance.attribute().key.to_string());
360                attribute.insert("base", instance.base_value());
361
362                let modifiers = instance
363                    .permanent_modifiers()
364                    .map(|modifier| {
365                        let mut packed = NbtCompound::new();
366                        packed.insert("id", modifier.id.to_string());
367                        packed.insert("amount", modifier.amount);
368                        packed.insert("operation", modifier.operation.name());
369                        packed
370                    })
371                    .collect::<Vec<_>>();
372                if !modifiers.is_empty() {
373                    attribute.insert("modifiers", NbtList::Compound(modifiers));
374                }
375
376                attribute
377            })
378            .collect();
379        NbtList::Compound(attributes)
380    }
381
382    /// Returns whether an attribute has a modifier with the given ID.
383    #[must_use]
384    pub fn has_modifier(&self, attribute: AttributeRef, modifier_id: &Identifier) -> bool {
385        self.get_instance(attribute)
386            .is_some_and(|instance| instance.has_modifier(modifier_id))
387    }
388
389    /// Sets the base value of an attribute
390    pub fn set_base_value(&mut self, attribute: AttributeRef, value: f64) {
391        let Some(id) = attribute.try_id() else { return };
392        let Some(Some(instance)) = self.instances.get_mut(id) else {
393            return;
394        };
395        if instance.set_base_value(value) {
396            self.mark_dirty(id, attribute);
397        }
398    }
399
400    /// Adds a modifier to an attribute. Returns `false` if the modifier ID already exists
401    pub fn add_modifier(
402        &mut self,
403        attribute: AttributeRef,
404        modifier: AttributeModifier,
405        persistent: bool,
406    ) -> bool {
407        let Some(id) = attribute.try_id() else {
408            return false;
409        };
410        let Some(Some(instance)) = self.instances.get_mut(id) else {
411            return false;
412        };
413        if instance.add_modifier(modifier, persistent) {
414            self.mark_dirty(id, attribute);
415            true
416        } else {
417            false
418        }
419    }
420
421    /// Adds or replaces a modifier on an attribute
422    pub fn set_modifier(
423        &mut self,
424        attribute: AttributeRef,
425        modifier: AttributeModifier,
426        persistent: bool,
427    ) {
428        let Some(id) = attribute.try_id() else { return };
429        let Some(Some(instance)) = self.instances.get_mut(id) else {
430            return;
431        };
432        if instance.set_modifier(modifier, persistent) {
433            self.mark_dirty(id, attribute);
434        }
435    }
436
437    /// Removes a modifier from an attribute. Returns `true` if it existed
438    pub fn remove_modifier(&mut self, attribute: AttributeRef, modifier_id: &Identifier) -> bool {
439        let Some(id) = attribute.try_id() else {
440            return false;
441        };
442        let Some(Some(instance)) = self.instances.get_mut(id) else {
443            return false;
444        };
445        if instance.remove_modifier(modifier_id) {
446            self.mark_dirty(id, attribute);
447            true
448        } else {
449            false
450        }
451    }
452
453    fn mark_dirty(&mut self, id: usize, attribute: AttributeRef) {
454        let id = id as u16;
455        self.to_update.mark(id);
456        if attribute.syncable {
457            self.to_sync.mark(id);
458        }
459    }
460
461    /// Returns `true` if there are dirty attributes needing server-side effects
462    #[must_use]
463    pub fn has_dirty_updates(&self) -> bool {
464        !self.to_update.is_empty()
465    }
466
467    /// Returns `true` if there are syncable dirty attributes needing network send
468    #[must_use]
469    pub fn has_dirty_sync(&self) -> bool {
470        !self.to_sync.is_empty()
471    }
472
473    /// Drains `to_update` and returns the dirty attribute refs
474    pub fn drain_dirty_updates(&mut self) -> Vec<AttributeRef> {
475        self.to_update
476            .drain()
477            .filter_map(|id| {
478                self.instances
479                    .get(id as usize)?
480                    .as_ref()
481                    .map(|inst| inst.attribute)
482            })
483            .collect()
484    }
485
486    /// Drains `to_sync` and builds `AttributeSnapshot`s for `CUpdateAttributes`
487    pub fn drain_dirty_sync(&mut self) -> Vec<AttributeSnapshot> {
488        let ids: Vec<u16> = self.to_sync.drain().collect();
489        let mut snapshots = Vec::with_capacity(ids.len());
490        for id in ids {
491            if let Some(Some(instance)) = self.instances.get(id as usize) {
492                snapshots.push(instance.to_snapshot(i32::from(id)));
493            }
494        }
495        snapshots
496    }
497
498    /// Returns snapshots for ALL syncable attributes (initial tracking sync)
499    #[must_use]
500    pub fn syncable_snapshots(&self) -> Vec<AttributeSnapshot> {
501        let mut snapshots = Vec::new();
502        for (id, slot) in self.instances.iter().enumerate() {
503            if let Some(inst) = slot
504                && inst.attribute.syncable
505            {
506                snapshots.push(inst.to_snapshot(id as i32));
507            }
508        }
509        snapshots
510    }
511
512    /// Removes all transient modifiers (e.g. on respawn)
513    pub fn remove_all_transient(&mut self) {
514        let Self {
515            instances,
516            to_update,
517            to_sync,
518        } = self;
519        for (id, slot) in instances.iter_mut().enumerate() {
520            if let Some(inst) = slot
521                && inst.remove_transient_modifiers()
522            {
523                to_update.mark(id as u16);
524                if inst.attribute.syncable {
525                    to_sync.mark(id as u16);
526                }
527            }
528        }
529    }
530}
531
532#[cfg(test)]
533mod tests {
534    use simdnbt::owned::NbtTag;
535    use steel_registry::{REGISTRY, init_vanilla_registry, vanilla_attributes, vanilla_entities};
536
537    use super::*;
538
539    #[test]
540    fn all_generated_entity_default_attributes_resolve() {
541        init_vanilla_registry();
542
543        for (_, entity_type) in REGISTRY.entity_types.iter() {
544            let _ = AttributeMap::new_for_entity(entity_type);
545        }
546    }
547
548    #[test]
549    fn player_gravity_is_initialized_from_default_attributes() {
550        init_vanilla_registry();
551
552        let attributes = AttributeMap::new_for_entity(&vanilla_entities::PLAYER);
553
554        assert_eq!(
555            attributes
556                .required_value(vanilla_attributes::GRAVITY)
557                .to_bits(),
558            vanilla_attributes::GRAVITY.default_value.to_bits()
559        );
560    }
561
562    #[test]
563    fn vanilla_nbt_only_contains_permanent_modifiers() {
564        init_vanilla_registry();
565        let mut attributes = AttributeMap::new_for_entity(&vanilla_entities::PLAYER);
566        attributes.add_modifier(
567            vanilla_attributes::MAX_HEALTH,
568            AttributeModifier {
569                id: Identifier::vanilla_static("transient_test"),
570                amount: 1.0,
571                operation: AttributeModifierOperation::AddValue,
572            },
573            false,
574        );
575        attributes.add_modifier(
576            vanilla_attributes::MAX_HEALTH,
577            AttributeModifier {
578                id: Identifier::vanilla_static("permanent_test"),
579                amount: 2.0,
580                operation: AttributeModifierOperation::AddMultipliedBase,
581            },
582            true,
583        );
584
585        let NbtList::Compound(packed_attributes) = attributes.to_vanilla_nbt() else {
586            panic!("attributes should serialize as a compound list");
587        };
588        let max_health = packed_attributes
589            .iter()
590            .find(|attribute| {
591                attribute.string("id").is_some_and(|id| {
592                    id.to_str().as_ref() == vanilla_attributes::MAX_HEALTH.key.to_string()
593                })
594            })
595            .unwrap_or_else(|| panic!("max health should be serialized"));
596        let Some(NbtTag::List(NbtList::Compound(modifiers))) = max_health.get("modifiers") else {
597            panic!("permanent modifier should be serialized");
598        };
599
600        assert_eq!(modifiers.len(), 1);
601        assert_eq!(
602            modifiers[0].string("id").map(ToString::to_string),
603            Some("minecraft:permanent_test".to_owned())
604        );
605        assert_eq!(modifiers[0].double("amount"), Some(2.0));
606        assert_eq!(
607            modifiers[0].string("operation").map(ToString::to_string),
608            Some("add_multiplied_base".to_owned())
609        );
610    }
611}