Skip to main content

steel_registry/data_components/registry/
patch_persistence.rs

1use super::{
2    BorrowedNbtTag, Component, ComponentData, ComponentHasher, ComponentPatchEntry,
3    DataComponentMap, DataComponentPatch, DataComponentType, DowncastType, EmbeddedNbtCodec,
4    FromNbtTag, HashComponent, HashEntry, Identifier, NbtCompound, OwnedNbtTag, Result, ToNbtTag,
5    sort_map_entries,
6};
7
8impl DataComponentPatch {
9    /// Computes the persistent map hash for one build-time validated component
10    /// without consulting the not-yet-published global registry.
11    pub(crate) fn compute_single_extracted_hash<T: HashComponent>(
12        key: &Identifier,
13        value: &T,
14    ) -> i32 {
15        let entry = hash_entry(key.to_string().compute_hash(), value.compute_hash());
16        let mut hasher = ComponentHasher::new();
17        hasher.start_map();
18        hasher.put_raw_bytes(&entry.key_bytes);
19        hasher.put_raw_bytes(&entry.value_bytes);
20        hasher.end_map();
21        hasher.finish()
22    }
23
24    /// Computes Vanilla's `HashOps` value for the persistent patch codec.
25    pub fn compute_persistent_hash(&self) -> Result<i32> {
26        use crate::{REGISTRY, RegistryExt};
27
28        let mut entries = Vec::new();
29        for (key, patch_entry) in &self.entries {
30            let Some(component) = REGISTRY.data_components.by_key(key) else {
31                continue;
32            };
33            if !component.is_persistent() {
34                continue;
35            }
36
37            let (encoded_key, value_hash) = match patch_entry {
38                ComponentPatchEntry::Set(data) => (key.to_string(), component.compute_hash(data)?),
39                ComponentPatchEntry::Removed => (format!("!{key}"), ().compute_hash()),
40            };
41            entries.push(hash_entry(encoded_key.compute_hash(), value_hash));
42        }
43        sort_map_entries(&mut entries);
44
45        let mut hasher = ComponentHasher::new();
46        hasher.start_map();
47        for entry in &entries {
48            hasher.put_raw_bytes(&entry.key_bytes);
49            hasher.put_raw_bytes(&entry.value_bytes);
50        }
51        hasher.end_map();
52        Ok(hasher.finish())
53    }
54
55    /// Iterates over removed component keys.
56    pub fn iter_removed(&self) -> impl Iterator<Item = &Identifier> {
57        self.entries.iter().filter_map(|(k, v)| {
58            if matches!(v, ComponentPatchEntry::Removed) {
59                Some(k)
60            } else {
61                None
62            }
63        })
64    }
65
66    fn encode_nbt(&self, validate: bool) -> (OwnedNbtTag, Vec<std::io::Error>) {
67        use crate::{REGISTRY, RegistryExt};
68
69        let mut compound = NbtCompound::new();
70        let mut errors = Vec::new();
71
72        for (key, entry) in &self.entries {
73            let Some(component) = REGISTRY.data_components.by_key(key) else {
74                continue;
75            };
76            if !component.is_persistent() {
77                continue;
78            }
79            match entry {
80                ComponentPatchEntry::Set(data) => {
81                    let encoded = if validate {
82                        component.validate_persistent_encoding(data)
83                    } else {
84                        component.write_nbt(data)
85                    };
86                    match encoded {
87                        Ok(nbt) => {
88                            compound.insert(key.to_string(), nbt);
89                        }
90                        Err(error) => errors.push(std::io::Error::other(format!(
91                            "failed to encode component {key}: {error}"
92                        ))),
93                    }
94                }
95                ComponentPatchEntry::Removed => {
96                    compound.insert(format!("!{key}"), NbtCompound::new());
97                }
98            }
99        }
100
101        (OwnedNbtTag::Compound(compound), errors)
102    }
103
104    /// Strictly encodes this component patch through its persistent codecs.
105    ///
106    /// This is the equivalent of Vanilla encoding an untrusted stack through
107    /// `ItemStack.CODEC` before accepting it into server state.
108    pub fn try_to_nbt_tag_ref(&self) -> Result<OwnedNbtTag> {
109        let (tag, errors) = self.encode_nbt(true);
110        match errors.into_iter().next() {
111            Some(error) => Err(error),
112            None => Ok(tag),
113        }
114    }
115
116    /// Converts this component patch to NBT without consuming it.
117    ///
118    /// Save-time encoding mirrors Vanilla's `TagValueOutput`: invalid fields
119    /// are reported and omitted from the partial result rather than aborting
120    /// the owner save.
121    #[must_use]
122    pub fn to_nbt_tag_ref(&self) -> OwnedNbtTag {
123        let (tag, errors) = self.encode_nbt(false);
124        for error in errors {
125            log::warn!("Data component serialization error: {error}");
126        }
127        tag
128    }
129}
130
131impl DataComponentMap {
132    fn as_patch(&self) -> DataComponentPatch {
133        DataComponentPatch {
134            entries: self
135                .map
136                .iter()
137                .map(|(key, value)| (key.clone(), ComponentPatchEntry::Set(value.clone())))
138                .collect(),
139        }
140    }
141
142    /// Strictly encodes this component map through its persistent codecs.
143    pub fn try_to_nbt_tag_ref(&self) -> Result<OwnedNbtTag> {
144        self.as_patch().try_to_nbt_tag_ref()
145    }
146
147    /// Encodes this component map while reporting and omitting invalid values.
148    #[must_use]
149    pub fn to_nbt_tag_ref(&self) -> OwnedNbtTag {
150        self.as_patch().to_nbt_tag_ref()
151    }
152}
153
154pub(super) fn hash_entry(key_hash: i32, value_hash: i32) -> HashEntry {
155    let key_hash = key_hash as u32;
156    let value_hash = value_hash as u32;
157    HashEntry {
158        key_hash: i64::from(key_hash),
159        value_hash: i64::from(value_hash),
160        key_bytes: key_hash.to_le_bytes(),
161        value_bytes: value_hash.to_le_bytes(),
162    }
163}
164impl ToNbtTag for DataComponentPatch {
165    fn to_nbt_tag(self) -> OwnedNbtTag {
166        self.to_nbt_tag_ref()
167    }
168}
169
170impl EmbeddedNbtCodec for &DataComponentPatch {
171    type Error = std::io::Error;
172
173    fn encode_embedded_nbt(self) -> Result<OwnedNbtTag> {
174        self.try_to_nbt_tag_ref()
175    }
176}
177
178impl FromNbtTag for DataComponentPatch {
179    fn from_nbt_tag(tag: BorrowedNbtTag) -> Option<Self> {
180        use crate::{REGISTRY, RegistryExt};
181
182        let compound = tag.compound()?;
183        let mut patch = Self::new();
184
185        for (key, value) in compound.iter() {
186            let key_str = key.to_str();
187
188            if let Some(stripped) = key_str.strip_prefix('!') {
189                let id = stripped.parse::<Identifier>().ok()?;
190                let entry = REGISTRY.data_components.by_key(&id)?;
191                if !entry.is_persistent() || value.compound().is_none() {
192                    return None;
193                }
194                patch.entries.insert(id, ComponentPatchEntry::Removed);
195            } else {
196                let id = key_str.parse::<Identifier>().ok()?;
197                let entry = REGISTRY.data_components.by_key(&id)?;
198                if !entry.is_persistent() {
199                    return None;
200                }
201                let component_data = entry.read_nbt(value)?;
202                patch
203                    .entries
204                    .insert(id, ComponentPatchEntry::Set(component_data));
205            }
206        }
207
208        Some(patch)
209    }
210}
211
212impl FromNbtTag for DataComponentMap {
213    fn from_nbt_tag(tag: BorrowedNbtTag) -> Option<Self> {
214        let patch = DataComponentPatch::from_nbt_tag(tag)?;
215        let mut map = rustc_hash::FxHashMap::default();
216        for (key, entry) in patch.entries {
217            let ComponentPatchEntry::Set(value) = entry else {
218                return None;
219            };
220            map.insert(key, value);
221        }
222        Some(Self { map })
223    }
224}
225
226/// Attempts to extract a typed component from `ComponentData`.
227#[must_use]
228pub fn component_try_into<T: Component + DowncastType>(
229    data: &ComponentData,
230    _component: DataComponentType<T>,
231) -> Option<&T> {
232    data.downcast_ref::<T>()
233}