Skip to main content

steel_registry/data_component_predicate/
collections.rs

1use super::{
2    ComponentHasher, DataComponentPredicateCodec, Debug, HashComponent, HashEntry, IntBounds,
3    NbtCompound, NbtList, NbtTag, decode_optional, hash_entries, push_hash_entry,
4};
5
6/// Generic collection predicate shared by container, firework, book, and attribute checks.
7#[derive(Debug, Clone, PartialEq)]
8pub struct CollectionPredicate<P> {
9    contains: Option<Vec<P>>,
10    counts: Option<Vec<CollectionCountPredicate<P>>>,
11    size: Option<IntBounds>,
12}
13
14impl<P> CollectionPredicate<P> {
15    #[must_use]
16    pub const fn new(
17        contains: Option<Vec<P>>,
18        counts: Option<Vec<CollectionCountPredicate<P>>>,
19        size: Option<IntBounds>,
20    ) -> Self {
21        Self {
22            contains,
23            counts,
24            size,
25        }
26    }
27
28    #[must_use]
29    pub const fn contains(&self) -> Option<&Vec<P>> {
30        self.contains.as_ref()
31    }
32
33    #[must_use]
34    pub const fn counts(&self) -> Option<&Vec<CollectionCountPredicate<P>>> {
35        self.counts.as_ref()
36    }
37
38    #[must_use]
39    pub const fn size(&self) -> Option<&IntBounds> {
40        self.size.as_ref()
41    }
42
43    pub(super) fn from_nbt_with(
44        tag: &NbtTag,
45        decode: impl Fn(&NbtTag) -> Option<P> + Copy,
46    ) -> Option<Self> {
47        let compound = tag.compound()?;
48        Some(Self::new(
49            decode_optional(compound, "contains", |tag| decode_list(tag, decode))?,
50            decode_optional(compound, "count", |tag| {
51                decode_list(tag, |tag| {
52                    CollectionCountPredicate::from_nbt_with(tag, decode)
53                })
54            })?,
55            decode_optional(compound, "size", IntBounds::from_owned_nbt)?,
56        ))
57    }
58
59    pub(super) fn to_nbt_with(&self, encode: impl Fn(&P) -> NbtTag + Copy) -> NbtTag {
60        let mut compound = NbtCompound::new();
61        if let Some(contains) = &self.contains {
62            compound.insert("contains", encode_list(contains, encode));
63        }
64        if let Some(counts) = &self.counts {
65            compound.insert(
66                "count",
67                encode_list(counts, |entry| entry.to_nbt_with(encode)),
68            );
69        }
70        if let Some(size) = &self.size {
71            compound.insert("size", size.as_nbt_tag());
72        }
73        NbtTag::Compound(compound)
74    }
75
76    pub(super) fn hash_with(&self, hasher: &mut ComponentHasher, hash: impl Fn(&P) -> i32 + Copy) {
77        let mut entries = Vec::new();
78        if let Some(contains) = &self.contains {
79            let mut value_hasher = ComponentHasher::new();
80            hash_list_with(contains, &mut value_hasher, hash);
81            crate::item_predicate::push_prehashed_entry(&mut entries, "contains", value_hasher);
82        }
83        if let Some(counts) = &self.counts {
84            let mut value_hasher = ComponentHasher::new();
85            value_hasher.start_list();
86            for entry in counts {
87                let mut entry_hasher = ComponentHasher::new();
88                entry.hash_with(&mut entry_hasher, hash);
89                value_hasher.put_raw_bytes(&(entry_hasher.finish() as u32).to_le_bytes());
90            }
91            value_hasher.end_list();
92            crate::item_predicate::push_prehashed_entry(&mut entries, "count", value_hasher);
93        }
94        if let Some(size) = &self.size {
95            push_hash_entry(&mut entries, "size", size);
96        }
97        hash_entries(hasher, &mut entries);
98    }
99}
100
101/// One element predicate and the accepted number of matching elements.
102#[derive(Debug, Clone, PartialEq)]
103pub struct CollectionCountPredicate<P> {
104    test: P,
105    count: IntBounds,
106}
107
108impl<P> CollectionCountPredicate<P> {
109    #[must_use]
110    pub const fn new(test: P, count: IntBounds) -> Self {
111        Self { test, count }
112    }
113
114    #[must_use]
115    pub const fn test(&self) -> &P {
116        &self.test
117    }
118
119    #[must_use]
120    pub const fn count(&self) -> IntBounds {
121        self.count
122    }
123
124    fn from_nbt_with(tag: &NbtTag, decode: impl Fn(&NbtTag) -> Option<P>) -> Option<Self> {
125        let compound = tag.compound()?;
126        Some(Self::new(
127            decode(compound.get("test")?)?,
128            IntBounds::from_owned_nbt(compound.get("count")?)?,
129        ))
130    }
131
132    fn to_nbt_with(&self, encode: impl Fn(&P) -> NbtTag) -> NbtTag {
133        let mut compound = NbtCompound::new();
134        compound.insert("test", encode(&self.test));
135        compound.insert("count", self.count.as_nbt_tag());
136        NbtTag::Compound(compound)
137    }
138
139    fn hash_with(&self, hasher: &mut ComponentHasher, hash: impl Fn(&P) -> i32) {
140        let mut entries = Vec::new();
141        let mut key_hasher = ComponentHasher::new();
142        "test".hash_component(&mut key_hasher);
143        entries.push(HashEntry::from_hashes(
144            key_hasher.finish() as u32,
145            hash(&self.test) as u32,
146        ));
147        push_hash_entry(&mut entries, "count", &self.count);
148        hash_entries(hasher, &mut entries);
149    }
150}
151
152pub(super) fn decode_list<T>(
153    tag: &NbtTag,
154    decode: impl Fn(&NbtTag) -> Option<T>,
155) -> Option<Vec<T>> {
156    tag.list()?.as_nbt_tags().iter().map(decode).collect()
157}
158
159pub(super) fn encode_list<T>(values: &[T], encode: impl Fn(&T) -> NbtTag) -> NbtTag {
160    NbtTag::List(NbtList::from(values.iter().map(encode).collect::<Vec<_>>()))
161}
162
163pub(super) fn hash_list_with<T>(
164    values: &[T],
165    hasher: &mut ComponentHasher,
166    hash: impl Fn(&T) -> i32,
167) {
168    hasher.start_list();
169    for value in values {
170        hasher.put_raw_bytes(&(hash(value) as u32).to_le_bytes());
171    }
172    hasher.end_list();
173}
174
175pub(super) fn collection_field_nbt<P>(
176    collection: Option<&CollectionPredicate<P>>,
177    name: &str,
178    encode: impl Fn(&P) -> NbtTag + Copy,
179) -> NbtTag {
180    let mut compound = NbtCompound::new();
181    if let Some(collection) = collection {
182        compound.insert(name, collection.to_nbt_with(encode));
183    }
184    NbtTag::Compound(compound)
185}
186
187pub(super) fn hash_optional_collection_field<P>(
188    collection: Option<&CollectionPredicate<P>>,
189    name: &str,
190    hasher: &mut ComponentHasher,
191    hash: impl Fn(&P) -> i32 + Copy,
192) {
193    let mut entries = Vec::new();
194    if let Some(collection) = collection {
195        let mut value_hasher = ComponentHasher::new();
196        collection.hash_with(&mut value_hasher, hash);
197        crate::item_predicate::push_prehashed_entry(&mut entries, name, value_hasher);
198    }
199    hash_entries(hasher, &mut entries);
200}
201
202pub(super) fn hash_nbt_codec<T: DataComponentPredicateCodec>(
203    value: &T,
204    hasher: &mut ComponentHasher,
205) {
206    value.to_nbt_value().hash_component(hasher);
207}
208
209pub(super) fn owned_string(tag: &NbtTag) -> Option<String> {
210    tag.string()?.to_owned().try_into_string().ok()
211}