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