1use std::fmt::{self, Debug, Formatter};
4use std::io::{Cursor, Error, Result, Write};
5
6use rustc_hash::FxHashMap;
7use simdnbt::ToNbtTag as _;
8use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
9use steel_utils::codec::VarInt;
10use steel_utils::hash::{ComponentHasher, HashComponent, HashEntry, sort_map_entries};
11use steel_utils::nbt::NbtNumeric as _;
12use steel_utils::serial::{ReadFrom, WriteTo};
13use steel_utils::{Downcast as _, DowncastType, DowncastTypeKey, ErasedType, Identifier};
14
15use crate::mob_effect::MobEffect;
16use crate::mob_effect_instance::MobEffectInstance;
17use crate::sound_event::SoundEventHolder;
18use crate::{REGISTRY, RegistryEntry, RegistryExt, RegistryHolderSet};
19
20pub trait ConsumeEffectCodec:
22 DowncastType + Clone + Debug + PartialEq + Send + Sync + 'static
23{
24 fn read_fields(compound: &NbtCompound) -> Option<Self>;
25 fn write_fields(&self, compound: &mut NbtCompound);
26 fn read_network(data: &mut Cursor<&[u8]>) -> Result<Self>;
27 fn write_network(&self, writer: &mut Vec<u8>) -> Result<()>;
28 fn hash_fields(&self, entries: &mut Vec<HashEntry>);
29}
30
31trait ErasedConsumeEffect: ErasedType + Debug + Send + Sync {
32 fn clone_effect(&self) -> Box<dyn ErasedConsumeEffect>;
33 fn effect_eq(&self, other: &dyn ErasedConsumeEffect) -> bool;
34}
35
36impl<T: ConsumeEffectCodec> ErasedConsumeEffect for T {
37 fn clone_effect(&self) -> Box<dyn ErasedConsumeEffect> {
38 Box::new(self.clone())
39 }
40
41 fn effect_eq(&self, other: &dyn ErasedConsumeEffect) -> bool {
42 other.downcast_ref::<T>() == Some(self)
43 }
44}
45
46type PersistentReader = fn(&NbtCompound) -> Option<Box<dyn ErasedConsumeEffect>>;
47type PersistentWriter = fn(&dyn ErasedConsumeEffect, &mut NbtCompound);
48type NetworkReader = fn(&mut Cursor<&[u8]>) -> Result<Box<dyn ErasedConsumeEffect>>;
49type NetworkWriter = fn(&dyn ErasedConsumeEffect, &mut Vec<u8>) -> Result<()>;
50type FieldsHasher = fn(&dyn ErasedConsumeEffect, &mut Vec<HashEntry>);
51
52pub struct ConsumeEffectType {
54 pub key: Identifier,
55 expected_type_key: DowncastTypeKey,
56 persistent_reader: PersistentReader,
57 persistent_writer: PersistentWriter,
58 network_reader: NetworkReader,
59 network_writer: NetworkWriter,
60 fields_hasher: FieldsHasher,
61}
62
63impl ConsumeEffectType {
64 #[must_use]
65 pub const fn of<T: ConsumeEffectCodec>(key: Identifier) -> Self {
66 Self {
67 key,
68 expected_type_key: T::TYPE_KEY,
69 persistent_reader: read_persistent::<T>,
70 persistent_writer: write_persistent::<T>,
71 network_reader: read_network::<T>,
72 network_writer: write_network::<T>,
73 fields_hasher: hash_fields::<T>,
74 }
75 }
76}
77
78impl Debug for ConsumeEffectType {
79 fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
80 formatter
81 .debug_struct("ConsumeEffectType")
82 .field("key", &self.key)
83 .field("expected_type_key", &self.expected_type_key)
84 .finish_non_exhaustive()
85 }
86}
87
88pub type ConsumeEffectTypeRef = &'static ConsumeEffectType;
89
90pub struct ConsumeEffectData {
92 effect_type: ConsumeEffectTypeRef,
93 value: Box<dyn ErasedConsumeEffect>,
94}
95
96impl ConsumeEffectData {
97 #[must_use]
98 pub fn new<T: ConsumeEffectCodec>(effect_type: ConsumeEffectTypeRef, value: T) -> Self {
99 assert_eq!(
100 effect_type.expected_type_key,
101 T::TYPE_KEY,
102 "consume effect value does not match its registered type"
103 );
104 Self {
105 effect_type,
106 value: Box::new(value),
107 }
108 }
109
110 #[must_use]
111 pub const fn effect_type(&self) -> ConsumeEffectTypeRef {
112 self.effect_type
113 }
114
115 #[must_use]
116 pub fn downcast_ref<T: DowncastType>(&self) -> Option<&T> {
117 self.value.downcast_ref::<T>()
118 }
119
120 pub(crate) fn to_nbt_tag_ref(&self) -> NbtTag {
121 let mut compound = NbtCompound::new();
122 compound.insert("type", self.effect_type.key.to_string());
123 (self.effect_type.persistent_writer)(self.value.as_ref(), &mut compound);
124 NbtTag::Compound(compound)
125 }
126
127 pub(crate) fn from_owned_nbt(tag: &NbtTag) -> Option<Self> {
128 let compound = tag.compound()?;
129 let key = compound.get("type")?.string()?.to_string().parse().ok()?;
130 let effect_type = REGISTRY.consume_effect_types.by_key(&key)?;
131 let value = (effect_type.persistent_reader)(compound)?;
132 Some(Self { effect_type, value })
133 }
134}
135
136impl Clone for ConsumeEffectData {
137 fn clone(&self) -> Self {
138 Self {
139 effect_type: self.effect_type,
140 value: self.value.clone_effect(),
141 }
142 }
143}
144
145impl Debug for ConsumeEffectData {
146 fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
147 formatter
148 .debug_struct("ConsumeEffectData")
149 .field("effect_type", &self.effect_type.key)
150 .field("value", &self.value)
151 .finish()
152 }
153}
154
155impl PartialEq for ConsumeEffectData {
156 fn eq(&self, other: &Self) -> bool {
157 self.effect_type.key == other.effect_type.key && self.value.effect_eq(other.value.as_ref())
158 }
159}
160
161impl WriteTo for ConsumeEffectData {
162 fn write(&self, writer: &mut impl Write) -> Result<()> {
163 let id = self.effect_type.try_id().ok_or_else(|| {
164 Error::other(format!(
165 "Unknown consume effect type: {}",
166 self.effect_type.key
167 ))
168 })?;
169 let id = i32::try_from(id)
170 .map_err(|_| Error::other(format!("Consume effect type id out of range: {id}")))?;
171 VarInt(id).write(writer)?;
172 let mut payload = Vec::new();
173 (self.effect_type.network_writer)(self.value.as_ref(), &mut payload)?;
174 writer.write_all(&payload)
175 }
176}
177
178impl ReadFrom for ConsumeEffectData {
179 fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
180 let id = VarInt::read(data)?.0;
181 let id = usize::try_from(id)
182 .map_err(|_| Error::other(format!("Negative consume effect type id: {id}")))?;
183 let effect_type = REGISTRY
184 .consume_effect_types
185 .by_id(id)
186 .ok_or_else(|| Error::other(format!("Unknown consume effect type id: {id}")))?;
187 let value = (effect_type.network_reader)(data)?;
188 Ok(Self { effect_type, value })
189 }
190}
191
192impl HashComponent for ConsumeEffectData {
193 fn hash_component(&self, hasher: &mut ComponentHasher) {
194 let mut entries = Vec::new();
195 push_hash_entry(&mut entries, "type", &self.effect_type.key);
196 (self.effect_type.fields_hasher)(self.value.as_ref(), &mut entries);
197 hash_entries(hasher, &mut entries);
198 }
199}
200
201pub struct ConsumeEffectTypeRegistry {
202 types_by_id: Vec<ConsumeEffectTypeRef>,
203 types_by_key: FxHashMap<Identifier, usize>,
204 allows_registering: bool,
205}
206
207impl ConsumeEffectTypeRegistry {
208 #[must_use]
209 pub fn new() -> Self {
210 Self {
211 types_by_id: Vec::new(),
212 types_by_key: FxHashMap::default(),
213 allows_registering: true,
214 }
215 }
216}
217
218crate::impl_standard_methods!(
219 ConsumeEffectTypeRegistry,
220 ConsumeEffectTypeRef,
221 types_by_id,
222 types_by_key,
223 allows_registering
224);
225crate::impl_registry!(
226 ConsumeEffectTypeRegistry,
227 ConsumeEffectType,
228 types_by_id,
229 types_by_key,
230 consume_effect_types
231);
232
233#[derive(Debug, Clone)]
234pub struct ApplyStatusEffectsConsumeEffect {
235 effects: Vec<MobEffectInstance>,
236 probability: f32,
237}
238
239impl ApplyStatusEffectsConsumeEffect {
240 pub fn new(effects: Vec<MobEffectInstance>, probability: f32) -> Result<Self> {
241 if !is_float_in_unit_range(probability) {
242 return Err(Error::other("Consume-effect probability must be in 0..=1"));
243 }
244 Ok(Self {
245 effects,
246 probability,
247 })
248 }
249
250 pub(crate) const fn from_extracted(effects: Vec<MobEffectInstance>, probability: f32) -> Self {
251 assert!(
252 is_float_in_unit_range(probability),
253 "extracted consume-effect probability must be in 0..=1"
254 );
255 Self {
256 effects,
257 probability,
258 }
259 }
260
261 #[must_use]
262 pub fn effects(&self) -> &[MobEffectInstance] {
263 &self.effects
264 }
265
266 #[must_use]
267 pub const fn probability(&self) -> f32 {
268 self.probability
269 }
270}
271
272impl PartialEq for ApplyStatusEffectsConsumeEffect {
273 fn eq(&self, other: &Self) -> bool {
274 self.effects == other.effects && float_equals(self.probability, other.probability)
275 }
276}
277
278unsafe impl DowncastType for ApplyStatusEffectsConsumeEffect {
280 const TYPE_KEY: DowncastTypeKey =
281 DowncastTypeKey::new("steel:consume_effect/apply_status_effects");
282}
283
284impl ConsumeEffectCodec for ApplyStatusEffectsConsumeEffect {
285 fn read_fields(compound: &NbtCompound) -> Option<Self> {
286 let effects = compound
287 .get("effects")?
288 .list()?
289 .as_nbt_tags()
290 .iter()
291 .map(MobEffectInstance::from_owned_nbt)
292 .collect::<Option<Vec<_>>>()?;
293 Self::new(effects, optional_f32(compound.get("probability"), 1.0)?).ok()
294 }
295
296 fn write_fields(&self, compound: &mut NbtCompound) {
297 compound.insert(
298 "effects",
299 NbtList::Compound(
300 self.effects
301 .iter()
302 .map(|effect| match effect.to_nbt_tag_ref() {
303 NbtTag::Compound(compound) => compound,
304 _ => unreachable!("mob effect codec always produces a compound"),
305 })
306 .collect(),
307 ),
308 );
309 if !float_equals(self.probability, 1.0) {
310 compound.insert("probability", self.probability);
311 }
312 }
313
314 fn read_network(data: &mut Cursor<&[u8]>) -> Result<Self> {
315 let count = read_count(data, "mob effect")?;
316 let mut effects = Vec::with_capacity(count.min(65_536));
317 for _ in 0..count {
318 effects.push(MobEffectInstance::read(data)?);
319 }
320 Self::new(effects, f32::read(data)?)
321 }
322
323 fn write_network(&self, writer: &mut Vec<u8>) -> Result<()> {
324 write_count(self.effects.len(), writer, "mob effect")?;
325 for effect in &self.effects {
326 effect.write(writer)?;
327 }
328 self.probability.write(writer)
329 }
330
331 fn hash_fields(&self, entries: &mut Vec<HashEntry>) {
332 push_hash_entry(entries, "effects", &MobEffectList(&self.effects));
333 if !float_equals(self.probability, 1.0) {
334 push_hash_entry(entries, "probability", &self.probability);
335 }
336 }
337}
338
339#[derive(Debug, Clone, PartialEq)]
340pub struct RemoveStatusEffectsConsumeEffect {
341 effects: RegistryHolderSet<MobEffect>,
342}
343
344impl RemoveStatusEffectsConsumeEffect {
345 #[must_use]
346 pub const fn new(effects: RegistryHolderSet<MobEffect>) -> Self {
347 Self { effects }
348 }
349
350 #[must_use]
351 pub const fn effects(&self) -> &RegistryHolderSet<MobEffect> {
352 &self.effects
353 }
354}
355
356unsafe impl DowncastType for RemoveStatusEffectsConsumeEffect {
358 const TYPE_KEY: DowncastTypeKey =
359 DowncastTypeKey::new("steel:consume_effect/remove_status_effects");
360}
361
362impl ConsumeEffectCodec for RemoveStatusEffectsConsumeEffect {
363 fn read_fields(compound: &NbtCompound) -> Option<Self> {
364 Some(Self::new(RegistryHolderSet::from_owned_nbt(
365 compound.get("effects")?,
366 )?))
367 }
368
369 fn write_fields(&self, compound: &mut NbtCompound) {
370 compound.insert("effects", self.effects.clone().to_nbt_tag());
371 }
372
373 fn read_network(data: &mut Cursor<&[u8]>) -> Result<Self> {
374 Ok(Self::new(RegistryHolderSet::read(data)?))
375 }
376
377 fn write_network(&self, writer: &mut Vec<u8>) -> Result<()> {
378 self.effects.write(writer)
379 }
380
381 fn hash_fields(&self, entries: &mut Vec<HashEntry>) {
382 push_hash_entry(entries, "effects", &self.effects);
383 }
384}
385
386#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
387pub struct ClearAllStatusEffectsConsumeEffect;
388
389unsafe impl DowncastType for ClearAllStatusEffectsConsumeEffect {
391 const TYPE_KEY: DowncastTypeKey =
392 DowncastTypeKey::new("steel:consume_effect/clear_all_status_effects");
393}
394
395impl ConsumeEffectCodec for ClearAllStatusEffectsConsumeEffect {
396 fn read_fields(_compound: &NbtCompound) -> Option<Self> {
397 Some(Self)
398 }
399
400 fn write_fields(&self, _compound: &mut NbtCompound) {}
401
402 fn read_network(_data: &mut Cursor<&[u8]>) -> Result<Self> {
403 Ok(Self)
404 }
405
406 fn write_network(&self, _writer: &mut Vec<u8>) -> Result<()> {
407 Ok(())
408 }
409
410 fn hash_fields(&self, _entries: &mut Vec<HashEntry>) {}
411}
412
413#[derive(Debug, Clone, Copy)]
414pub struct TeleportRandomlyConsumeEffect {
415 diameter: f32,
416}
417
418impl TeleportRandomlyConsumeEffect {
419 pub const DEFAULT_DIAMETER: f32 = 16.0;
420
421 pub fn new(diameter: f32) -> Result<Self> {
422 if !is_positive_float(diameter) {
423 return Err(Error::other("Random teleport diameter must be positive"));
424 }
425 Ok(Self { diameter })
426 }
427
428 pub(crate) const fn from_extracted(diameter: f32) -> Self {
429 assert!(
430 is_positive_float(diameter),
431 "extracted random teleport diameter must be positive"
432 );
433 Self { diameter }
434 }
435
436 #[must_use]
437 pub const fn default_value() -> Self {
438 Self {
439 diameter: Self::DEFAULT_DIAMETER,
440 }
441 }
442
443 #[must_use]
444 pub const fn diameter(self) -> f32 {
445 self.diameter
446 }
447}
448
449impl PartialEq for TeleportRandomlyConsumeEffect {
450 fn eq(&self, other: &Self) -> bool {
451 float_equals(self.diameter, other.diameter)
452 }
453}
454
455unsafe impl DowncastType for TeleportRandomlyConsumeEffect {
457 const TYPE_KEY: DowncastTypeKey =
458 DowncastTypeKey::new("steel:consume_effect/teleport_randomly");
459}
460
461impl ConsumeEffectCodec for TeleportRandomlyConsumeEffect {
462 fn read_fields(compound: &NbtCompound) -> Option<Self> {
463 Self::new(optional_f32(
464 compound.get("diameter"),
465 Self::DEFAULT_DIAMETER,
466 )?)
467 .ok()
468 }
469
470 fn write_fields(&self, compound: &mut NbtCompound) {
471 if !float_equals(self.diameter, Self::DEFAULT_DIAMETER) {
472 compound.insert("diameter", self.diameter);
473 }
474 }
475
476 fn read_network(data: &mut Cursor<&[u8]>) -> Result<Self> {
477 Self::new(f32::read(data)?)
478 }
479
480 fn write_network(&self, writer: &mut Vec<u8>) -> Result<()> {
481 self.diameter.write(writer)
482 }
483
484 fn hash_fields(&self, entries: &mut Vec<HashEntry>) {
485 if !float_equals(self.diameter, Self::DEFAULT_DIAMETER) {
486 push_hash_entry(entries, "diameter", &self.diameter);
487 }
488 }
489}
490
491#[derive(Debug, Clone, PartialEq)]
492pub struct PlaySoundConsumeEffect {
493 sound: SoundEventHolder,
494}
495
496impl PlaySoundConsumeEffect {
497 #[must_use]
498 pub const fn new(sound: SoundEventHolder) -> Self {
499 Self { sound }
500 }
501
502 #[must_use]
503 pub const fn sound(&self) -> &SoundEventHolder {
504 &self.sound
505 }
506}
507
508unsafe impl DowncastType for PlaySoundConsumeEffect {
510 const TYPE_KEY: DowncastTypeKey = DowncastTypeKey::new("steel:consume_effect/play_sound");
511}
512
513impl ConsumeEffectCodec for PlaySoundConsumeEffect {
514 fn read_fields(compound: &NbtCompound) -> Option<Self> {
515 Some(Self::new(SoundEventHolder::from_owned_nbt(
516 compound.get("sound")?,
517 )?))
518 }
519
520 fn write_fields(&self, compound: &mut NbtCompound) {
521 compound.insert("sound", self.sound.clone().to_nbt_tag());
522 }
523
524 fn read_network(data: &mut Cursor<&[u8]>) -> Result<Self> {
525 Ok(Self::new(SoundEventHolder::read(data)?))
526 }
527
528 fn write_network(&self, writer: &mut Vec<u8>) -> Result<()> {
529 self.sound.write(writer)
530 }
531
532 fn hash_fields(&self, entries: &mut Vec<HashEntry>) {
533 push_hash_entry(entries, "sound", &self.sound);
534 }
535}
536
537pub mod vanilla_consume_effect_types {
538 use super::{
539 ApplyStatusEffectsConsumeEffect, ClearAllStatusEffectsConsumeEffect, ConsumeEffectType,
540 ConsumeEffectTypeRegistry, PlaySoundConsumeEffect, RemoveStatusEffectsConsumeEffect,
541 TeleportRandomlyConsumeEffect,
542 };
543 use steel_utils::Identifier;
544
545 pub static APPLY_EFFECTS: ConsumeEffectType =
546 ConsumeEffectType::of::<ApplyStatusEffectsConsumeEffect>(Identifier::vanilla_static(
547 "apply_effects",
548 ));
549 pub static REMOVE_EFFECTS: ConsumeEffectType =
550 ConsumeEffectType::of::<RemoveStatusEffectsConsumeEffect>(Identifier::vanilla_static(
551 "remove_effects",
552 ));
553 pub static CLEAR_ALL_EFFECTS: ConsumeEffectType =
554 ConsumeEffectType::of::<ClearAllStatusEffectsConsumeEffect>(Identifier::vanilla_static(
555 "clear_all_effects",
556 ));
557 pub static TELEPORT_RANDOMLY: ConsumeEffectType =
558 ConsumeEffectType::of::<TeleportRandomlyConsumeEffect>(Identifier::vanilla_static(
559 "teleport_randomly",
560 ));
561 pub static PLAY_SOUND: ConsumeEffectType =
562 ConsumeEffectType::of::<PlaySoundConsumeEffect>(Identifier::vanilla_static("play_sound"));
563
564 pub fn register_consume_effect_types(registry: &mut ConsumeEffectTypeRegistry) {
565 registry.register(&APPLY_EFFECTS);
566 registry.register(&REMOVE_EFFECTS);
567 registry.register(&CLEAR_ALL_EFFECTS);
568 registry.register(&TELEPORT_RANDOMLY);
569 registry.register(&PLAY_SOUND);
570 }
571}
572
573fn read_persistent<T: ConsumeEffectCodec>(
574 compound: &NbtCompound,
575) -> Option<Box<dyn ErasedConsumeEffect>> {
576 T::read_fields(compound).map(|value| Box::new(value) as Box<dyn ErasedConsumeEffect>)
577}
578
579fn write_persistent<T: ConsumeEffectCodec>(
580 value: &dyn ErasedConsumeEffect,
581 compound: &mut NbtCompound,
582) {
583 let Some(value) = value.downcast_ref::<T>() else {
584 panic!("registered consume effect payload type mismatch");
585 };
586 value.write_fields(compound);
587}
588
589fn read_network<T: ConsumeEffectCodec>(
590 data: &mut Cursor<&[u8]>,
591) -> Result<Box<dyn ErasedConsumeEffect>> {
592 Ok(Box::new(T::read_network(data)?))
593}
594
595fn write_network<T: ConsumeEffectCodec>(
596 value: &dyn ErasedConsumeEffect,
597 writer: &mut Vec<u8>,
598) -> Result<()> {
599 value
600 .downcast_ref::<T>()
601 .ok_or_else(|| Error::other("Consume effect payload type mismatch"))?
602 .write_network(writer)
603}
604
605fn hash_fields<T: ConsumeEffectCodec>(
606 value: &dyn ErasedConsumeEffect,
607 entries: &mut Vec<HashEntry>,
608) {
609 let Some(value) = value.downcast_ref::<T>() else {
610 panic!("registered consume effect payload type mismatch");
611 };
612 value.hash_fields(entries);
613}
614
615struct MobEffectList<'a>(&'a [MobEffectInstance]);
616
617impl HashComponent for MobEffectList<'_> {
618 fn hash_component(&self, hasher: &mut ComponentHasher) {
619 hasher.start_list();
620 for effect in self.0 {
621 hasher.put_component_hash(effect);
622 }
623 hasher.end_list();
624 }
625}
626
627const fn float_equals(left: f32, right: f32) -> bool {
628 (left.is_nan() && right.is_nan()) || left.to_bits() == right.to_bits()
629}
630
631const fn is_positive_float(value: f32) -> bool {
632 value > 0.0 && value <= f32::MAX
633}
634
635const fn is_float_in_unit_range(value: f32) -> bool {
636 value.is_finite() && !value.is_sign_negative() && value <= 1.0
637}
638
639fn optional_f32(tag: Option<&NbtTag>, default: f32) -> Option<f32> {
640 match tag {
641 Some(tag) => tag.codec_f32(),
642 None => Some(default),
643 }
644}
645
646fn write_count(count: usize, writer: &mut Vec<u8>, name: &str) -> Result<()> {
647 let count = i32::try_from(count).map_err(|_| Error::other(format!("{name} list too large")))?;
648 VarInt(count).write(writer)
649}
650
651fn read_count(data: &mut Cursor<&[u8]>, name: &str) -> Result<usize> {
652 let count = VarInt::read(data)?.0;
653 usize::try_from(count).map_err(|_| Error::other(format!("Negative {name} count: {count}")))
654}
655
656fn push_hash_entry<T: HashComponent + ?Sized>(entries: &mut Vec<HashEntry>, key: &str, value: &T) {
657 let mut key_hasher = ComponentHasher::new();
658 key_hasher.put_string(key);
659 let mut value_hasher = ComponentHasher::new();
660 value.hash_component(&mut value_hasher);
661 entries.push(HashEntry::new(key_hasher, value_hasher));
662}
663
664fn hash_entries(hasher: &mut ComponentHasher, entries: &mut [HashEntry]) {
665 sort_map_entries(entries);
666 hasher.start_map();
667 for entry in entries {
668 hasher.put_raw_bytes(&entry.key_bytes);
669 hasher.put_raw_bytes(&entry.value_bytes);
670 }
671 hasher.end_map();
672}