1use 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
13pub 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 #[must_use]
27 pub const fn new() -> Self {
28 Self { chunks: Vec::new() }
29 }
30
31 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 #[must_use]
43 pub fn is_empty(&self) -> bool {
44 self.chunks.iter().all(|&c| c == 0)
45 }
46
47 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#[derive(Clone, Debug)]
66pub struct AttributeModifier {
67 pub id: Identifier,
69 pub amount: f64,
71 pub operation: AttributeModifierOperation,
73}
74
75pub struct AttributeInstance {
77 attribute: AttributeRef,
78 base_value: f64,
79 modifiers: Vec<AttributeModifier>,
80 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 #[must_use]
99 pub const fn attribute(&self) -> AttributeRef {
100 self.attribute
101 }
102
103 #[must_use]
105 pub const fn base_value(&self) -> f64 {
106 self.base_value
107 }
108
109 #[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 #[must_use]
125 pub const fn value(&self) -> f64 {
126 self.cached_value
127 }
128
129 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 #[must_use]
142 pub fn has_modifier(&self, id: &Identifier) -> bool {
143 self.modifiers.iter().any(|modifier| modifier.id == *id)
144 }
145
146 #[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 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 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 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 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 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
253pub struct AttributeMap {
259 instances: Vec<Option<AttributeInstance>>,
260 to_update: DirtySet,
261 to_sync: DirtySet,
262}
263
264impl AttributeMap {
265 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 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 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 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 #[must_use]
463 pub fn has_dirty_updates(&self) -> bool {
464 !self.to_update.is_empty()
465 }
466
467 #[must_use]
469 pub fn has_dirty_sync(&self) -> bool {
470 !self.to_sync.is_empty()
471 }
472
473 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 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 #[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 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}