Skip to main content

steel_registry/data_component_predicate/
core.rs

1use super::*;
2
3/// Typed payload behavior for a registered partial component predicate.
4pub trait DataComponentPredicateCodec:
5    DowncastType + Clone + Debug + PartialEq + HashComponent + Send + Sync + 'static
6{
7    fn from_nbt_value(tag: &NbtTag) -> Option<Self>;
8    fn to_nbt_value(&self) -> NbtTag;
9}
10
11trait ErasedDataComponentPredicate: ErasedType + Debug + Send + Sync {
12    fn clone_predicate(&self) -> Box<dyn ErasedDataComponentPredicate>;
13    fn predicate_eq(&self, other: &dyn ErasedDataComponentPredicate) -> bool;
14}
15
16impl<T: DataComponentPredicateCodec> ErasedDataComponentPredicate for T {
17    fn clone_predicate(&self) -> Box<dyn ErasedDataComponentPredicate> {
18        Box::new(self.clone())
19    }
20
21    fn predicate_eq(&self, other: &dyn ErasedDataComponentPredicate) -> bool {
22        other.downcast_ref::<T>() == Some(self)
23    }
24}
25
26type PredicateReader = fn(&NbtTag) -> Option<Box<dyn ErasedDataComponentPredicate>>;
27type PredicateWriter = fn(&dyn ErasedDataComponentPredicate) -> NbtTag;
28type PredicateHasher = fn(&dyn ErasedDataComponentPredicate, &mut ComponentHasher);
29
30/// Registered discriminator and codecs for one concrete predicate value.
31pub struct DataComponentPredicateType {
32    pub key: Identifier,
33    expected_type_key: DowncastTypeKey,
34    reader: PredicateReader,
35    writer: PredicateWriter,
36    hasher: PredicateHasher,
37}
38
39impl DataComponentPredicateType {
40    #[must_use]
41    pub const fn of<T: DataComponentPredicateCodec>(key: Identifier) -> Self {
42        Self {
43            key,
44            expected_type_key: T::TYPE_KEY,
45            reader: read_predicate::<T>,
46            writer: write_predicate::<T>,
47            hasher: hash_predicate::<T>,
48        }
49    }
50}
51
52impl Debug for DataComponentPredicateType {
53    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
54        formatter
55            .debug_struct("DataComponentPredicateType")
56            .field("key", &self.key)
57            .field("expected_type_key", &self.expected_type_key)
58            .finish_non_exhaustive()
59    }
60}
61
62pub type DataComponentPredicateTypeRef = &'static DataComponentPredicateType;
63
64#[derive(Clone, Copy, PartialEq, Eq)]
65enum PredicateDiscriminator {
66    Concrete(DataComponentPredicateTypeRef),
67    Any(ComponentEntryRef),
68}
69
70impl PredicateDiscriminator {
71    const fn key(&self) -> &Identifier {
72        match *self {
73            Self::Concrete(predicate_type) => &predicate_type.key,
74            Self::Any(component) => &component.key,
75        }
76    }
77}
78
79impl Debug for PredicateDiscriminator {
80    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
81        formatter
82            .debug_tuple("PredicateDiscriminator")
83            .field(self.key())
84            .finish()
85    }
86}
87
88/// Type-erased predicate value retaining Steel's deterministic concrete key.
89pub struct DataComponentPredicateData {
90    discriminator: PredicateDiscriminator,
91    value: Option<Box<dyn ErasedDataComponentPredicate>>,
92}
93
94impl DataComponentPredicateData {
95    #[must_use]
96    pub fn new<T: DataComponentPredicateCodec>(
97        predicate_type: DataComponentPredicateTypeRef,
98        value: T,
99    ) -> Self {
100        assert_eq!(
101            predicate_type.expected_type_key,
102            T::TYPE_KEY,
103            "component predicate value does not match its registered type"
104        );
105        Self {
106            discriminator: PredicateDiscriminator::Concrete(predicate_type),
107            value: Some(Box::new(value)),
108        }
109    }
110
111    #[must_use]
112    pub const fn any(component: ComponentEntryRef) -> Self {
113        Self {
114            discriminator: PredicateDiscriminator::Any(component),
115            value: None,
116        }
117    }
118
119    #[must_use]
120    pub const fn predicate_type(&self) -> Option<DataComponentPredicateTypeRef> {
121        match self.discriminator {
122            PredicateDiscriminator::Concrete(predicate_type) => Some(predicate_type),
123            PredicateDiscriminator::Any(_) => None,
124        }
125    }
126
127    #[must_use]
128    pub const fn any_component(&self) -> Option<ComponentEntryRef> {
129        match self.discriminator {
130            PredicateDiscriminator::Concrete(_) => None,
131            PredicateDiscriminator::Any(component) => Some(component),
132        }
133    }
134
135    #[must_use]
136    pub fn downcast_ref<T: DowncastType>(&self) -> Option<&T> {
137        self.value.as_deref()?.downcast_ref::<T>()
138    }
139
140    #[must_use]
141    pub const fn key(&self) -> &Identifier {
142        self.discriminator.key()
143    }
144
145    pub(super) fn from_persistent_entry(key: &Identifier, tag: &NbtTag) -> Option<Self> {
146        if let Some(predicate_type) = REGISTRY.data_component_predicate_types.by_key(key) {
147            return Some(Self {
148                discriminator: PredicateDiscriminator::Concrete(predicate_type),
149                value: Some((predicate_type.reader)(tag)?),
150            });
151        }
152        let component = REGISTRY.data_components.by_key(key)?;
153        tag.compound()?;
154        Some(Self::any(component))
155    }
156
157    fn to_nbt_value(&self) -> NbtTag {
158        match (self.discriminator, self.value.as_deref()) {
159            (PredicateDiscriminator::Concrete(predicate_type), Some(value)) => {
160                (predicate_type.writer)(value)
161            }
162            (PredicateDiscriminator::Any(_), None) => NbtTag::Compound(NbtCompound::new()),
163            _ => panic!("component predicate discriminator and value disagree"),
164        }
165    }
166
167    fn hash_value(&self, hasher: &mut ComponentHasher) {
168        match (self.discriminator, self.value.as_deref()) {
169            (PredicateDiscriminator::Concrete(predicate_type), Some(value)) => {
170                (predicate_type.hasher)(value, hasher);
171            }
172            (PredicateDiscriminator::Any(_), None) => {
173                hasher.start_map();
174                hasher.end_map();
175            }
176            _ => panic!("component predicate discriminator and value disagree"),
177        }
178    }
179}
180
181impl Clone for DataComponentPredicateData {
182    fn clone(&self) -> Self {
183        Self {
184            discriminator: self.discriminator,
185            value: self.value.as_ref().map(|value| value.clone_predicate()),
186        }
187    }
188}
189
190impl Debug for DataComponentPredicateData {
191    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
192        formatter
193            .debug_struct("DataComponentPredicateData")
194            .field("key", self.key())
195            .field("value", &self.value)
196            .finish()
197    }
198}
199
200impl PartialEq for DataComponentPredicateData {
201    fn eq(&self, other: &Self) -> bool {
202        if self.discriminator != other.discriminator {
203            return false;
204        }
205        match (self.value.as_deref(), other.value.as_deref()) {
206            (Some(left), Some(right)) => left.predicate_eq(right),
207            (None, None) => true,
208            _ => false,
209        }
210    }
211}
212
213/// Registry of concrete partial component predicate types.
214pub struct DataComponentPredicateTypeRegistry {
215    types_by_id: Vec<DataComponentPredicateTypeRef>,
216    types_by_key: FxHashMap<Identifier, usize>,
217    allows_registering: bool,
218}
219
220impl DataComponentPredicateTypeRegistry {
221    #[must_use]
222    pub fn new() -> Self {
223        Self {
224            types_by_id: Vec::new(),
225            types_by_key: FxHashMap::default(),
226            allows_registering: true,
227        }
228    }
229}
230
231crate::impl_standard_methods!(
232    DataComponentPredicateTypeRegistry,
233    DataComponentPredicateTypeRef,
234    types_by_id,
235    types_by_key,
236    allows_registering
237);
238crate::impl_registry!(
239    DataComponentPredicateTypeRegistry,
240    DataComponentPredicateType,
241    types_by_id,
242    types_by_key,
243    data_component_predicate_types
244);
245
246fn read_predicate<T: DataComponentPredicateCodec>(
247    tag: &NbtTag,
248) -> Option<Box<dyn ErasedDataComponentPredicate>> {
249    T::from_nbt_value(tag).map(|value| Box::new(value) as Box<dyn ErasedDataComponentPredicate>)
250}
251
252fn write_predicate<T: DataComponentPredicateCodec>(
253    value: &dyn ErasedDataComponentPredicate,
254) -> NbtTag {
255    let Some(value) = value.downcast_ref::<T>() else {
256        panic!("registered component predicate writer received the wrong concrete type");
257    };
258    value.to_nbt_value()
259}
260
261fn hash_predicate<T: DataComponentPredicateCodec>(
262    value: &dyn ErasedDataComponentPredicate,
263    hasher: &mut ComponentHasher,
264) {
265    let Some(value) = value.downcast_ref::<T>() else {
266        panic!("registered component predicate hasher received the wrong concrete type");
267    };
268    value.hash_component(hasher);
269}
270
271/// Exact component values required by a component matcher.
272#[derive(Clone, PartialEq)]
273pub struct DataComponentExactPredicate {
274    values: Vec<(ComponentEntryRef, ComponentData)>,
275}
276
277impl Debug for DataComponentExactPredicate {
278    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
279        formatter
280            .debug_list()
281            .entries(self.values.iter().map(|(entry, value)| (&entry.key, value)))
282            .finish()
283    }
284}
285
286impl DataComponentExactPredicate {
287    pub const EMPTY: Self = Self { values: Vec::new() };
288
289    /// Creates an exact predicate only when every persistent value can round-trip.
290    ///
291    /// Vanilla's direct stream codec can admit values rejected by a component's
292    /// persistent codec. Steel rejects those here because exact predicates are
293    /// nested in item components and must not make their containing stack unsavable.
294    #[must_use]
295    pub fn new(values: Vec<(ComponentEntryRef, ComponentData)>) -> Option<Self> {
296        let mut keys = FxHashSet::default();
297        values
298            .iter()
299            .all(|(entry, value)| {
300                entry.validates(value)
301                    && keys.insert(entry.key.clone())
302                    && (!entry.is_persistent() || entry.validate_persistent_encoding(value).is_ok())
303            })
304            .then_some(Self { values })
305    }
306
307    #[must_use]
308    pub fn all_of(components: &DataComponentMap) -> Option<Self> {
309        let values = components
310            .keys()
311            .map(|key| {
312                Some((
313                    REGISTRY.data_components.by_key(key)?,
314                    components.get_raw(key)?.clone(),
315                ))
316            })
317            .collect::<Option<Vec<_>>>()?;
318        Self::new(values)
319    }
320
321    #[must_use]
322    pub const fn is_empty(&self) -> bool {
323        self.values.is_empty()
324    }
325
326    #[must_use]
327    pub fn values(&self) -> &[(ComponentEntryRef, ComponentData)] {
328        &self.values
329    }
330
331    fn from_owned_nbt(tag: &NbtTag) -> Option<Self> {
332        let compound = tag.compound()?;
333        let mut values = Vec::with_capacity(compound.len());
334        for (key, value) in compound.iter() {
335            let key = key.to_owned().try_into_string().ok()?.parse().ok()?;
336            let entry = REGISTRY.data_components.by_key(&key)?;
337            if !entry.is_persistent() {
338                return None;
339            }
340            values.push((entry, entry.read_nbt_owned(value)?));
341        }
342        Self::new(values)
343    }
344
345    fn to_nbt_value(&self) -> NbtTag {
346        let mut compound = NbtCompound::new();
347        for (entry, value) in &self.values {
348            if !entry.is_persistent() {
349                continue;
350            }
351            let Ok(value) = entry.write_nbt(value) else {
352                panic!("validated exact component predicate failed to encode");
353            };
354            compound.insert(entry.key.to_string(), value);
355        }
356        NbtTag::Compound(compound)
357    }
358}
359
360impl WriteTo for DataComponentExactPredicate {
361    fn write(&self, writer: &mut impl Write) -> Result<()> {
362        write_len(self.values.len(), writer)?;
363        for (entry, value) in &self.values {
364            write_registry_id(*entry, writer, "data component")?;
365            let mut encoded = Vec::new();
366            entry.write_network(value, &mut encoded)?;
367            writer.write_all(&encoded)?;
368        }
369        Ok(())
370    }
371}
372
373impl ReadFrom for DataComponentExactPredicate {
374    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
375        let count = read_len(data, usize::MAX, "exact component predicate")?;
376        let mut values = Vec::with_capacity(count.min(1024));
377        for _ in 0..count {
378            let entry = read_component_entry(data)?;
379            values.push((entry, entry.read_network(data)?));
380        }
381        Self::new(values)
382            .ok_or_else(|| Error::other("duplicate or mismatched exact component predicate"))
383    }
384}
385
386impl HashComponent for DataComponentExactPredicate {
387    fn hash_component(&self, hasher: &mut ComponentHasher) {
388        let mut entries = Vec::new();
389        for (entry, value) in &self.values {
390            if !entry.is_persistent() {
391                continue;
392            }
393            let Ok(value_hash) = entry.compute_hash(value) else {
394                panic!("validated exact component predicate failed to hash");
395            };
396            entries.push(HashEntry::from_hashes(
397                entry.key.compute_hash() as u32,
398                value_hash as u32,
399            ));
400        }
401        hash_entries(hasher, &mut entries);
402    }
403}
404
405/// Exact and partial component conditions flattened into block/item predicates.
406#[derive(Debug, Clone, PartialEq)]
407pub struct DataComponentMatchers {
408    exact: DataComponentExactPredicate,
409    partial: Vec<DataComponentPredicateData>,
410}
411
412impl DataComponentMatchers {
413    pub const ANY: Self = Self {
414        exact: DataComponentExactPredicate::EMPTY,
415        partial: Vec::new(),
416    };
417
418    #[must_use]
419    pub fn new(
420        exact: DataComponentExactPredicate,
421        partial: Vec<DataComponentPredicateData>,
422    ) -> Option<Self> {
423        let mut keys = FxHashSet::default();
424        partial
425            .iter()
426            .all(|predicate| keys.insert(predicate.key().clone()))
427            .then_some(Self { exact, partial })
428    }
429
430    #[must_use]
431    pub const fn is_empty(&self) -> bool {
432        self.exact.is_empty() && self.partial.is_empty()
433    }
434
435    #[must_use]
436    pub const fn exact(&self) -> &DataComponentExactPredicate {
437        &self.exact
438    }
439
440    #[must_use]
441    pub fn partial(&self) -> &[DataComponentPredicateData] {
442        &self.partial
443    }
444
445    pub(crate) fn from_fields(compound: &NbtCompound) -> Option<Self> {
446        let exact = compound.get("components").map_or(
447            Some(DataComponentExactPredicate::EMPTY),
448            DataComponentExactPredicate::from_owned_nbt,
449        )?;
450        let partial = if let Some(tag) = compound.get("predicates") {
451            let values = tag.compound()?;
452            let mut predicates = Vec::with_capacity(values.len());
453            for (key, value) in values.iter() {
454                let key = key.to_owned().try_into_string().ok()?.parse().ok()?;
455                predicates.push(DataComponentPredicateData::from_persistent_entry(
456                    &key, value,
457                )?);
458            }
459            predicates
460        } else {
461            Vec::new()
462        };
463        Self::new(exact, partial)
464    }
465
466    pub(crate) fn write_fields(&self, compound: &mut NbtCompound) {
467        if !self.exact.is_empty() {
468            compound.insert("components", self.exact.to_nbt_value());
469        }
470        if !self.partial.is_empty() {
471            let mut predicates = NbtCompound::new();
472            for predicate in &self.partial {
473                predicates.insert(predicate.key().to_string(), predicate.to_nbt_value());
474            }
475            compound.insert("predicates", predicates);
476        }
477    }
478
479    pub(crate) fn hash_fields(&self, entries: &mut Vec<HashEntry>) {
480        if !self.exact.is_empty() {
481            push_hash_entry(entries, "components", &self.exact);
482        }
483        if !self.partial.is_empty() {
484            let mut value_hasher = ComponentHasher::new();
485            self.hash_partial(&mut value_hasher);
486            crate::item_predicate::push_prehashed_entry(entries, "predicates", value_hasher);
487        }
488    }
489
490    fn hash_partial(&self, hasher: &mut ComponentHasher) {
491        let mut entries = self
492            .partial
493            .iter()
494            .map(|predicate| {
495                let mut key_hasher = ComponentHasher::new();
496                predicate.key().hash_component(&mut key_hasher);
497                let mut value_hasher = ComponentHasher::new();
498                predicate.hash_value(&mut value_hasher);
499                HashEntry::new(key_hasher, value_hasher)
500            })
501            .collect::<Vec<_>>();
502        hash_entries(hasher, &mut entries);
503    }
504}
505
506impl WriteTo for DataComponentMatchers {
507    fn write(&self, writer: &mut impl Write) -> Result<()> {
508        self.exact.write(writer)?;
509        if self.partial.len() > 64 {
510            return Err(Error::other("partial component predicate count exceeds 64"));
511        }
512        write_len(self.partial.len(), writer)?;
513        for predicate in &self.partial {
514            match predicate.discriminator {
515                PredicateDiscriminator::Concrete(predicate_type) => {
516                    true.write(writer)?;
517                    write_registry_id(predicate_type, writer, "component predicate type")?;
518                }
519                PredicateDiscriminator::Any(component) => {
520                    false.write(writer)?;
521                    write_registry_id(component, writer, "data component")?;
522                }
523            }
524            let mut encoded = Vec::new();
525            predicate.to_nbt_value().write(&mut encoded);
526            writer.write_all(&encoded)?;
527        }
528        Ok(())
529    }
530}
531
532impl ReadFrom for DataComponentMatchers {
533    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
534        let exact = DataComponentExactPredicate::read(data)?;
535        let count = read_len(data, 64, "partial component predicate")?;
536        let mut partial = Vec::with_capacity(count);
537        for _ in 0..count {
538            let discriminator = if bool::read(data)? {
539                let id = read_registry_id(data, "component predicate type")?;
540                PredicateDiscriminator::Concrete(
541                    REGISTRY
542                        .data_component_predicate_types
543                        .by_id(id)
544                        .ok_or_else(|| {
545                            Error::other(format!("unknown component predicate type id: {id}"))
546                        })?,
547                )
548            } else {
549                PredicateDiscriminator::Any(read_component_entry(data)?)
550            };
551            let tag = read_network_nbt(data)?;
552            let value = match discriminator {
553                PredicateDiscriminator::Concrete(predicate_type) => Some(
554                    (predicate_type.reader)(&tag)
555                        .ok_or_else(|| Error::other("invalid component predicate payload"))?,
556                ),
557                PredicateDiscriminator::Any(_) => {
558                    if tag.compound().is_none() {
559                        return Err(Error::other(
560                            "any-value predicate payload is not a compound",
561                        ));
562                    }
563                    None
564                }
565            };
566            partial.push(DataComponentPredicateData {
567                discriminator,
568                value,
569            });
570        }
571        Self::new(exact, partial)
572            .ok_or_else(|| Error::other("duplicate partial component predicate"))
573    }
574}
575
576impl HashComponent for DataComponentMatchers {
577    fn hash_component(&self, hasher: &mut ComponentHasher) {
578        let mut entries = Vec::new();
579        self.hash_fields(&mut entries);
580        hash_entries(hasher, &mut entries);
581    }
582}
583
584pub(super) fn write_registry_id(
585    entry: &impl RegistryEntry,
586    writer: &mut impl Write,
587    name: &str,
588) -> Result<()> {
589    let id = entry
590        .try_id()
591        .ok_or_else(|| Error::other(format!("unknown {name}: {}", entry.key())))?;
592    let id = i32::try_from(id).map_err(|_| Error::other(format!("{name} id out of range")))?;
593    VarInt(id).write(writer)
594}
595
596fn read_registry_id(data: &mut Cursor<&[u8]>, name: &str) -> Result<usize> {
597    let id = VarInt::read(data)?.0;
598    usize::try_from(id).map_err(|_| Error::other(format!("negative {name} id: {id}")))
599}
600
601fn read_component_entry(data: &mut Cursor<&[u8]>) -> Result<ComponentEntryRef> {
602    let id = read_registry_id(data, "data component")?;
603    REGISTRY
604        .data_components
605        .by_id(id)
606        .ok_or_else(|| Error::other(format!("unknown data component id: {id}")))
607}