azalea_core/registry_holder/
enchantment.rs

1use std::{any::Any, fmt::Debug, str::FromStr};
2
3use azalea_registry::builtin::EnchantmentEffectComponentKind;
4use indexmap::IndexMap;
5use simdnbt::{DeserializeError, borrow::NbtCompound};
6
7use crate::registry_holder::components::{
8    EffectComponentTrait, EffectComponentUnion, EffectNbtTag, ResolvedEffectComponent,
9};
10
11pub struct EnchantmentData {
12    // TODO: make these two deserializable
13    // pub description: TextComponent,
14    // pub exclusive_set: HolderSet<Enchantment, ResourceLocation>,
15    effects: IndexMap<EnchantmentEffectComponentKind, Vec<EffectComponentUnion>>,
16}
17
18impl EnchantmentData {
19    pub fn get<T: EffectComponentTrait>(&self) -> Option<Vec<&T>> {
20        let components = self.get_kind(T::KIND)?;
21        let components_any = components
22            .into_iter()
23            .map(|c| (c as &dyn Any).downcast_ref::<T>())
24            .collect::<Option<_>>()?;
25        Some(components_any)
26    }
27
28    pub fn get_kind(
29        &self,
30        kind: EnchantmentEffectComponentKind,
31    ) -> Option<Vec<&dyn ResolvedEffectComponent>> {
32        self.effects.get(&kind).map(|c| {
33            c.iter()
34                .map(|c| {
35                    // SAFETY: we just got the component from the map, so it must be the correct
36                    // kind
37                    unsafe { c.as_kind(kind) }
38                })
39                .collect()
40        })
41    }
42}
43
44impl simdnbt::Deserialize for EnchantmentData {
45    fn from_compound(nbt: NbtCompound) -> Result<Self, DeserializeError> {
46        let mut effects: IndexMap<EnchantmentEffectComponentKind, Vec<EffectComponentUnion>> =
47            IndexMap::new();
48
49        if let Some(effects_tag) = nbt.compound("effects") {
50            for (key, list) in effects_tag.iter() {
51                let kind = EnchantmentEffectComponentKind::from_str(&key.to_str())
52                    .map_err(|_| DeserializeError::UnknownField(key.to_string()))?;
53
54                let mut components = Vec::new();
55                if let Some(tag) = list.compound() {
56                    if !tag.is_empty() {
57                        let value = EffectComponentUnion::from_effect_nbt_tag_as(
58                            kind,
59                            EffectNbtTag::Compound(tag),
60                        )?;
61                        components.push(value);
62                    }
63                } else {
64                    let list = list.list().ok_or_else(|| {
65                        DeserializeError::MismatchedFieldType("effects".to_owned())
66                    })?;
67
68                    if let Some(tags) = list.compounds() {
69                        for tag in tags {
70                            let value = EffectComponentUnion::from_effect_nbt_tag_as(
71                                kind,
72                                EffectNbtTag::Compound(tag),
73                            )?;
74                            components.push(value);
75                        }
76                    } else {
77                        let value = EffectComponentUnion::from_effect_nbt_tag_as(
78                            kind,
79                            EffectNbtTag::List(list),
80                        )?;
81                        components.push(value);
82                    }
83                }
84
85                effects.insert(kind, components);
86            }
87        }
88
89        let value = Self { effects };
90        Ok(value)
91    }
92}
93
94impl Debug for EnchantmentData {
95    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
96        f.debug_struct("EnchantmentData")
97            .field("effects", &self.effects.keys())
98            .finish()
99    }
100}
101
102impl Clone for EnchantmentData {
103    fn clone(&self) -> Self {
104        let mut effects = IndexMap::with_capacity(self.effects.len());
105        for (kind, effect) in &self.effects {
106            effects.insert(
107                *kind,
108                effect
109                    .iter()
110                    .map(|e| unsafe { e.clone_as(*kind) })
111                    .collect(),
112            );
113        }
114        EnchantmentData { effects }
115    }
116}