Skip to main content

bevy_reflect/
set.rs

1//! A trait used to power [set-like] operations via reflection.
2//!
3//! [set-like]: https://doc.rust-lang.org/stable/std/collections/struct.HashSet.html
4use alloc::{boxed::Box, vec::Vec};
5use core::fmt::{Debug, Formatter};
6
7use bevy_platform::collections::{hash_table::OccupiedEntry as HashTableOccupiedEntry, HashTable};
8use bevy_reflect_derive::impl_type_path;
9
10use crate::{
11    generics::impl_generic_info_methods, hash_error, ty::impl_type_methods, ApplyError, Generics,
12    PartialReflect, Reflect, ReflectCloneError, ReflectKind, ReflectMut, ReflectOwned, ReflectRef,
13    Type, TypeInfo, TypePath,
14};
15
16/// A trait used to power [set-like] operations via [reflection].
17///
18/// Sets contain zero or more entries of a fixed type, and correspond to types
19/// like [`HashSet`] and [`BTreeSet`].
20/// The order of these entries is not guaranteed by this trait.
21///
22/// # Hashing and equality
23///
24/// All values are expected to return a valid hash value from [`PartialReflect::reflect_hash`] and be
25/// comparable using [`PartialReflect::reflect_partial_eq`].
26/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding
27/// `#[reflect(Hash, PartialEq)]` to the entire struct or enum.
28/// The ordering is expected to be total, that is as if the reflected type implements the [`Eq`] trait.
29/// This is true even for manual implementors who do not hash or compare values,
30/// as it is still relied on by [`DynamicSet`].
31///
32/// # Example
33///
34/// ```
35/// use bevy_reflect::{PartialReflect, set::Set};
36/// use std::collections::HashSet;
37///
38///
39/// let foo: &mut dyn Set = &mut HashSet::<u32>::new();
40/// foo.insert_boxed(Box::new(123_u32));
41/// assert_eq!(foo.len(), 1);
42///
43/// let field: &dyn PartialReflect = foo.get(&123_u32).unwrap();
44/// assert_eq!(field.try_downcast_ref::<u32>(), Some(&123_u32));
45/// ```
46///
47/// [`HashSet`]: std::collections::HashSet
48/// [`BTreeSet`]: alloc::collections::BTreeSet
49/// [set-like]: https://doc.rust-lang.org/stable/std/collections/struct.HashSet.html
50/// [reflection]: crate
51// Prevents unexpectedly importing this trait when trying to call, for example, `HashSet::iter`
52#[rust_analyzer::completions(ignore_flyimport_methods)]
53pub trait Set: PartialReflect {
54    /// Returns a reference to the value.
55    ///
56    /// If no value is contained, returns `None`.
57    fn get(&self, value: &dyn PartialReflect) -> Option<&dyn PartialReflect>;
58
59    /// Returns the number of elements in the set.
60    fn len(&self) -> usize;
61
62    /// Returns `true` if the list contains no elements.
63    fn is_empty(&self) -> bool {
64        self.len() == 0
65    }
66
67    /// Returns an iterator over the values of the set.
68    fn iter(&self) -> Box<dyn Iterator<Item = &dyn PartialReflect> + '_>;
69
70    /// Drain the values of this set to get a vector of owned values.
71    ///
72    /// After calling this function, `self` will be empty.
73    fn drain(&mut self) -> Vec<Box<dyn PartialReflect>>;
74
75    /// Retain only the elements specified by the predicate.
76    ///
77    /// In other words, remove all elements `e` for which `f(&e)` returns `false`.
78    fn retain(&mut self, f: &mut dyn FnMut(&dyn PartialReflect) -> bool);
79
80    /// Creates a new [`DynamicSet`] from this set.
81    ///
82    /// Returns an error if any value cannot be converted via [`PartialReflect::to_dynamic`].
83    fn to_dynamic_set(&self) -> Result<DynamicSet, ReflectCloneError> {
84        let mut set = DynamicSet::default();
85        set.set_represented_type(self.get_represented_type_info());
86        for value in self.iter() {
87            set.insert_boxed(value.to_dynamic()?);
88        }
89        Ok(set)
90    }
91
92    /// Inserts a value into the set.
93    ///
94    /// If the set did not have this value present, `true` is returned.
95    /// If the set did have this value present, `false` is returned.
96    fn insert_boxed(&mut self, value: Box<dyn PartialReflect>) -> bool;
97
98    /// Removes a value from the set.
99    ///
100    /// If the set did have this value present, `true` is returned.
101    /// If the set did not have this value present, `false` is returned.
102    fn remove(&mut self, value: &dyn PartialReflect) -> bool;
103
104    /// Checks if the given value is contained in the set
105    fn contains(&self, value: &dyn PartialReflect) -> bool;
106}
107
108/// A container for compile-time set info.
109#[derive(Clone, Debug)]
110pub struct SetInfo {
111    ty: Type,
112    generics: Generics,
113    value_ty: Type,
114    #[cfg(feature = "reflect_documentation")]
115    docs: Option<&'static str>,
116}
117
118impl SetInfo {
119    /// Create a new [`SetInfo`].
120    pub fn new<TSet: Set + TypePath, TValue: Reflect + TypePath>() -> Self {
121        Self {
122            ty: Type::of::<TSet>(),
123            generics: Generics::new(),
124            value_ty: Type::of::<TValue>(),
125            #[cfg(feature = "reflect_documentation")]
126            docs: None,
127        }
128    }
129
130    /// Sets the docstring for this set.
131    #[cfg(feature = "reflect_documentation")]
132    pub fn with_docs(self, docs: Option<&'static str>) -> Self {
133        Self { docs, ..self }
134    }
135
136    impl_type_methods!(ty);
137
138    /// The [type] of the value.
139    ///
140    /// [type]: Type
141    pub fn value_ty(&self) -> Type {
142        self.value_ty
143    }
144
145    /// The docstring of this set, if any.
146    #[cfg(feature = "reflect_documentation")]
147    pub fn docs(&self) -> Option<&'static str> {
148        self.docs
149    }
150
151    impl_generic_info_methods!(generics);
152}
153
154/// An unordered set of reflected values.
155#[derive(Default)]
156pub struct DynamicSet {
157    represented_type: Option<&'static TypeInfo>,
158    hash_table: HashTable<Box<dyn PartialReflect>>,
159}
160
161impl DynamicSet {
162    /// Sets the [type] to be represented by this `DynamicSet`.
163    ///
164    /// # Panics
165    ///
166    /// Panics if the given [type] is not a [`TypeInfo::Set`].
167    ///
168    /// [type]: TypeInfo
169    pub fn set_represented_type(&mut self, represented_type: Option<&'static TypeInfo>) {
170        if let Some(represented_type) = represented_type {
171            assert!(
172                matches!(represented_type, TypeInfo::Set(_)),
173                "expected TypeInfo::Set but received: {represented_type:?}"
174            );
175        }
176
177        self.represented_type = represented_type;
178    }
179
180    /// Inserts a typed value into the set.
181    pub fn insert<V: Reflect>(&mut self, value: V) {
182        self.insert_boxed(Box::new(value));
183    }
184
185    fn internal_hash(value: &dyn PartialReflect) -> u64 {
186        value.reflect_hash().expect(&hash_error!(value))
187    }
188
189    fn internal_eq(
190        value: &dyn PartialReflect,
191    ) -> impl FnMut(&Box<dyn PartialReflect>) -> bool + '_ {
192        |other| {
193            value
194                .reflect_partial_eq(&**other)
195                .expect("Underlying type does not reflect `PartialEq` and hence doesn't support equality checks")
196        }
197    }
198}
199
200impl Set for DynamicSet {
201    fn get(&self, value: &dyn PartialReflect) -> Option<&dyn PartialReflect> {
202        self.hash_table
203            .find(Self::internal_hash(value), Self::internal_eq(value))
204            .map(|value| &**value)
205    }
206
207    fn len(&self) -> usize {
208        self.hash_table.len()
209    }
210
211    fn iter(&self) -> Box<dyn Iterator<Item = &dyn PartialReflect> + '_> {
212        let iter = self.hash_table.iter().map(|v| &**v);
213        Box::new(iter)
214    }
215
216    fn drain(&mut self) -> Vec<Box<dyn PartialReflect>> {
217        self.hash_table.drain().collect::<Vec<_>>()
218    }
219
220    fn retain(&mut self, f: &mut dyn FnMut(&dyn PartialReflect) -> bool) {
221        self.hash_table.retain(move |value| f(&**value));
222    }
223
224    fn insert_boxed(&mut self, value: Box<dyn PartialReflect>) -> bool {
225        assert_eq!(
226            value.reflect_partial_eq(&*value),
227            Some(true),
228            "Values inserted in `Set` like types are expected to reflect `PartialEq`"
229        );
230        match self
231            .hash_table
232            .find_mut(Self::internal_hash(&*value), Self::internal_eq(&*value))
233        {
234            Some(old) => {
235                *old = value;
236                false
237            }
238            None => {
239                self.hash_table.insert_unique(
240                    Self::internal_hash(value.as_ref()),
241                    value,
242                    |boxed| Self::internal_hash(boxed.as_ref()),
243                );
244                true
245            }
246        }
247    }
248
249    fn remove(&mut self, value: &dyn PartialReflect) -> bool {
250        self.hash_table
251            .find_entry(Self::internal_hash(value), Self::internal_eq(value))
252            .map(HashTableOccupiedEntry::remove)
253            .is_ok()
254    }
255
256    fn contains(&self, value: &dyn PartialReflect) -> bool {
257        self.hash_table
258            .find(Self::internal_hash(value), Self::internal_eq(value))
259            .is_some()
260    }
261}
262
263impl PartialReflect for DynamicSet {
264    #[inline]
265    fn get_represented_type_info(&self) -> Option<&'static TypeInfo> {
266        self.represented_type
267    }
268
269    #[inline]
270    fn into_partial_reflect(self: Box<Self>) -> Box<dyn PartialReflect> {
271        self
272    }
273
274    #[inline]
275    fn as_partial_reflect(&self) -> &dyn PartialReflect {
276        self
277    }
278
279    #[inline]
280    fn as_partial_reflect_mut(&mut self) -> &mut dyn PartialReflect {
281        self
282    }
283
284    #[inline]
285    fn try_into_reflect(self: Box<Self>) -> Result<Box<dyn Reflect>, Box<dyn PartialReflect>> {
286        Err(self)
287    }
288
289    #[inline]
290    fn try_as_reflect(&self) -> Option<&dyn Reflect> {
291        None
292    }
293
294    #[inline]
295    fn try_as_reflect_mut(&mut self) -> Option<&mut dyn Reflect> {
296        None
297    }
298
299    fn apply(&mut self, value: &dyn PartialReflect) {
300        set_apply(self, value);
301    }
302
303    fn try_apply(&mut self, value: &dyn PartialReflect) -> Result<(), ApplyError> {
304        set_try_apply(self, value)
305    }
306
307    fn reflect_kind(&self) -> ReflectKind {
308        ReflectKind::Set
309    }
310
311    fn reflect_ref(&self) -> ReflectRef<'_> {
312        ReflectRef::Set(self)
313    }
314
315    fn reflect_mut(&mut self) -> ReflectMut<'_> {
316        ReflectMut::Set(self)
317    }
318
319    fn reflect_owned(self: Box<Self>) -> ReflectOwned {
320        ReflectOwned::Set(self)
321    }
322
323    fn reflect_partial_eq(&self, value: &dyn PartialReflect) -> Option<bool> {
324        set_partial_eq(self, value)
325    }
326
327    fn debug(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
328        write!(f, "DynamicSet(")?;
329        set_debug(self, f)?;
330        write!(f, ")")
331    }
332
333    #[inline]
334    fn is_dynamic(&self) -> bool {
335        true
336    }
337}
338
339impl_type_path!((in bevy_reflect) DynamicSet);
340
341impl Debug for DynamicSet {
342    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
343        self.debug(f)
344    }
345}
346
347impl FromIterator<Box<dyn PartialReflect>> for DynamicSet {
348    fn from_iter<I: IntoIterator<Item = Box<dyn PartialReflect>>>(values: I) -> Self {
349        let mut this = Self {
350            represented_type: None,
351            hash_table: HashTable::new(),
352        };
353
354        for value in values {
355            this.insert_boxed(value);
356        }
357
358        this
359    }
360}
361
362impl<T: Reflect> FromIterator<T> for DynamicSet {
363    fn from_iter<I: IntoIterator<Item = T>>(values: I) -> Self {
364        let mut this = Self {
365            represented_type: None,
366            hash_table: HashTable::new(),
367        };
368
369        for value in values {
370            this.insert(value);
371        }
372
373        this
374    }
375}
376
377impl IntoIterator for DynamicSet {
378    type Item = Box<dyn PartialReflect>;
379    type IntoIter = bevy_platform::collections::hash_table::IntoIter<Self::Item>;
380
381    fn into_iter(self) -> Self::IntoIter {
382        self.hash_table.into_iter()
383    }
384}
385
386impl<'a> IntoIterator for &'a DynamicSet {
387    type Item = &'a dyn PartialReflect;
388    type IntoIter = core::iter::Map<
389        bevy_platform::collections::hash_table::Iter<'a, Box<dyn PartialReflect>>,
390        fn(&'a Box<dyn PartialReflect>) -> Self::Item,
391    >;
392
393    fn into_iter(self) -> Self::IntoIter {
394        self.hash_table.iter().map(|v| v.as_ref())
395    }
396}
397
398/// Compares a [`Set`] with a [`PartialReflect`] value.
399///
400/// Returns true if and only if all of the following are true:
401/// - `b` is a set;
402/// - `b` is the same length as `a`;
403/// - For each value pair in `a`, `b` contains the value too,
404///   and [`PartialReflect::reflect_partial_eq`] returns `Some(true)` for the two values.
405///
406/// Returns [`None`] if the comparison couldn't even be performed.
407#[inline]
408pub fn set_partial_eq<M: Set>(a: &M, b: &dyn PartialReflect) -> Option<bool> {
409    let ReflectRef::Set(set) = b.reflect_ref() else {
410        return Some(false);
411    };
412
413    if a.len() != set.len() {
414        return Some(false);
415    }
416
417    for value in a.iter() {
418        if let Some(set_value) = set.get(value) {
419            let eq_result = value.reflect_partial_eq(set_value);
420            if let failed @ (Some(false) | None) = eq_result {
421                return failed;
422            }
423        } else {
424            return Some(false);
425        }
426    }
427
428    Some(true)
429}
430
431/// The default debug formatter for [`Set`] types.
432///
433/// # Example
434/// ```
435/// # use std::collections::HashSet;
436/// use bevy_reflect::Reflect;
437///
438/// let mut my_set = HashSet::new();
439/// my_set.insert(String::from("Hello"));
440/// println!("{:#?}", &my_set as &dyn Reflect);
441///
442/// // Output:
443///
444/// // {
445/// //   "Hello",
446/// // }
447/// ```
448#[inline]
449pub fn set_debug(dyn_set: &dyn Set, f: &mut Formatter<'_>) -> core::fmt::Result {
450    let mut debug = f.debug_set();
451    for value in dyn_set.iter() {
452        debug.entry(&value as &dyn Debug);
453    }
454    debug.finish()
455}
456
457/// Applies the elements of reflected set `b` to the corresponding elements of set `a`.
458///
459/// If a value from `b` does not exist in `a`, the value is cloned and inserted.
460/// If a value from `a` does not exist in `b`, the value is removed.
461///
462/// # Panics
463///
464/// This function panics if `b` is not a reflected set.
465#[inline]
466pub fn set_apply<M: Set>(a: &mut M, b: &dyn PartialReflect) {
467    if let Err(err) = set_try_apply(a, b) {
468        panic!("{err}");
469    }
470}
471
472/// Tries to apply the elements of reflected set `b` to the corresponding elements of set `a`
473/// and returns a Result.
474///
475/// If a value from `b` does not exist in `a`, the value is cloned and inserted.
476/// If a value from `a` does not exist in `b`, the value is removed.
477///
478/// # Errors
479///
480/// This function returns an [`ApplyError::MismatchedKinds`] if `b` is not a reflected set or if
481/// applying elements to each other fails.
482#[inline]
483pub fn set_try_apply<S: Set>(a: &mut S, b: &dyn PartialReflect) -> Result<(), ApplyError> {
484    let set_value = b.reflect_ref().as_set()?;
485
486    for b_value in set_value.iter() {
487        if a.get(b_value).is_none() {
488            a.insert_boxed(b_value.to_dynamic()?);
489        }
490    }
491    a.retain(&mut |value| set_value.get(value).is_some());
492
493    Ok(())
494}
495
496#[cfg(test)]
497mod tests {
498    use crate::{set::Set, PartialReflect};
499
500    use super::DynamicSet;
501    use alloc::string::{String, ToString};
502
503    #[test]
504    fn test_into_iter() {
505        let expected = ["foo", "bar", "baz"];
506
507        let mut set = DynamicSet::default();
508        set.insert(expected[0].to_string());
509        set.insert(expected[1].to_string());
510        set.insert(expected[2].to_string());
511
512        for item in set.into_iter() {
513            let value = item
514                .try_take::<String>()
515                .expect("couldn't downcast to String");
516            let index = expected
517                .iter()
518                .position(|i| *i == value.as_str())
519                .expect("Element found in expected array");
520            assert_eq!(expected[index], value);
521        }
522    }
523
524    #[test]
525    fn apply() {
526        let mut map_a = DynamicSet::default();
527        map_a.insert(0);
528        map_a.insert(1);
529
530        let mut map_b = DynamicSet::default();
531        map_b.insert(1);
532        map_b.insert(2);
533
534        map_a.apply(&map_b);
535
536        assert!(map_a.get(&0).is_none());
537        assert_eq!(map_a.get(&1).unwrap().try_downcast_ref(), Some(&1));
538        assert_eq!(map_a.get(&2).unwrap().try_downcast_ref(), Some(&2));
539    }
540}