tf2_enum/
spell_set.rs

1//! Set for holding up to 2 spells.
2
3use crate::{
4    Attribute,
5    AttributeSet,
6    FootprintsSpell,
7    ItemAttribute,
8    PaintSpell,
9    Spell,
10    TryFromIntAttributeValue,
11};
12use crate::error::InsertError;
13use crate::serialize;
14use std::collections::HashSet;
15use std::fmt;
16use std::hash::{Hash, Hasher};
17use std::ops::{BitAnd, Sub};
18use serde::{Deserialize, Deserializer, Serialize, Serializer};
19use serde::de::{self, SeqAccess, Visitor};
20
21const SPELL_COUNT: usize = 2;
22
23/// Contains up to 2 spells. Although the underlying data structure is an array, this structure
24/// behaves like a set. Most methods mimic those of [`HashSet`](std::collections::HashSet),
25/// with a few differences.
26/// 
27/// This struct solves the following problems:
28/// - An item can only hold up to 2 spells.
29/// - An item cannot have duplicate spells or multiple spells of the same type.
30/// - Comparing spells for equality is order-agnostic.
31/// - Hashing is order-agnostic.
32/// - The type is `Copy`, allowing for cheap and easy duplication.
33/// 
34/// Most methods are implemented under the [`AttributeSet`] trait, make sure to import it to make
35/// use of them.
36/// 
37/// # Examples
38/// ```
39/// use tf2_enum::{SpellSet, Spell, AttributeSet};
40/// 
41/// // Create a set for spells with one spell.
42/// let mut spells = SpellSet::single(Spell::HeadlessHorseshoes);
43/// 
44/// // Check that spells contains Headless Horseshoes.
45/// assert!(spells.contains(&Spell::HeadlessHorseshoes));
46/// assert_eq!(spells.len(), 1);
47/// 
48/// // Add a spell.
49/// spells.insert(Spell::VoicesFromBelow);
50/// assert_eq!(spells.len(), 2);
51/// 
52/// // If a spell is added when spells are full, the insert will fail.
53/// assert!(!spells.insert(Spell::PumpkinBombs));
54/// assert!(!spells.contains(&Spell::PumpkinBombs));
55/// 
56/// // Iterate over spells.
57/// for spell in spells {
58///     println!("{}", spell);
59/// }
60/// ```
61#[derive(Default, Clone, Copy, Eq, PartialOrd, Ord)]
62pub struct SpellSet {
63    inner: [Option<Spell>; SPELL_COUNT]
64}
65
66impl SpellSet {
67    /// Creates a set for spells.
68    /// 
69    /// # Examples
70    /// ```
71    /// use tf2_enum::SpellSet;
72    /// 
73    /// let spells = SpellSet::new();
74    /// ```
75    pub fn new() -> Self {
76        Self::default()
77    }
78    
79    /// Creates a set for spells with one spell.
80    /// 
81    /// # Examples
82    /// ```
83    /// use tf2_enum::{SpellSet, Spell, AttributeSet};
84    /// 
85    /// let spells = SpellSet::single(Spell::HeadlessHorseshoes);
86    /// 
87    /// assert_eq!(spells.len(), 1);
88    /// ```
89    pub fn single(spell: Spell) -> Self {
90        Self {
91            inner: [
92                Some(spell),
93                None,
94            ],
95        }
96    }
97    
98    /// Creates a set for spells with two spells.
99    /// 
100    /// If the same spell is added multiple times, only one will be kept. This is also true for
101    /// spells of the same type. In cases of multiple spells of the same type, the first occuring
102    /// spell will be prioritized.
103    /// 
104    /// # Examples
105    /// ```
106    /// use tf2_enum::{SpellSet, Spell, AttributeSet};
107    /// 
108    /// let spells = SpellSet::double(Spell::HeadlessHorseshoes, Spell::VoicesFromBelow);
109    /// 
110    /// assert_eq!(spells.len(), 2);
111    /// 
112    /// let spells = SpellSet::double(Spell::HeadlessHorseshoes, Spell::TeamSpiritFootprints);
113    /// 
114    /// assert_eq!(spells.len(), 1);
115    /// assert_eq!(SpellSet::single(Spell::HeadlessHorseshoes), spells);
116    /// ```
117    pub fn double(spell1: Spell, spell2: Spell) -> Self {
118        Self::from([
119            Some(spell1),
120            Some(spell2),
121        ])
122    }
123}
124
125impl AttributeSet for SpellSet {
126    /// Max number of items.
127    const MAX_COUNT: usize = SPELL_COUNT;
128    /// An empty [`SpellSet`].
129    const NONE: Self = Self {
130        inner: [None, None],
131    };
132    /// The item type.
133    type Item = Spell;
134    
135    /// Clears the set, removing all spells.
136    /// 
137    /// # Examples
138    /// ```
139    /// use tf2_enum::{SpellSet, Spell, AttributeSet};
140    /// 
141    /// let mut spells = SpellSet::single(Spell::HeadlessHorseshoes);
142    /// 
143    /// spells.clear();
144    /// 
145    /// assert_eq!(spells.len(), 0);
146    /// ```
147    fn clear(&mut self) {
148        self.inner = [None, None];
149    }
150    
151    /// Adds a spell to the first available slot.
152    /// 
153    /// Returns `false` if:
154    /// - The spell is already in the set.
155    /// - The set is full.
156    /// 
157    /// # Examples
158    /// ```
159    /// use tf2_enum::{SpellSet, Spell, AttributeSet};
160    /// 
161    /// let mut spells = SpellSet::single(Spell::HeadlessHorseshoes);
162    /// 
163    /// assert_eq!(spells.len(), 1);
164    /// 
165    /// spells.insert(Spell::VoicesFromBelow);
166    /// 
167    /// assert_eq!(spells.len(), 2);
168    /// 
169    /// // Spells are full.
170    /// assert!(!spells.insert(Spell::PumpkinBombs));
171    /// ```
172    fn insert(&mut self, spell: Spell) -> bool {
173        self.try_insert(spell).is_ok()
174    }
175    
176    fn try_insert(&mut self, spell: Spell) -> Result<(), InsertError> {
177        let attribute_defindex = spell.attribute_defindex();
178        
179        if self.inner.iter().flatten().any(|s| s.attribute_defindex() == attribute_defindex) {
180            return Err(InsertError::Duplicate);
181        }
182        
183        if let Some(slot) = self.inner.iter_mut().find(|slot| slot.is_none()) {
184            *slot = Some(spell);
185            return Ok(());
186        }
187        
188        // full set, insertion failed
189        Err(InsertError::Full)
190    }
191    
192    fn insert_or_replace_last(&mut self, spell: Spell) -> bool {
193        if self.contains(&spell) {
194            return false;
195        }
196        
197        if let Some(slot) = self.inner.iter_mut().find(|slot| slot.is_none()) {
198            *slot = Some(spell);
199            return true;
200        }
201        
202        // replace the last item
203        self.inner[Self::MAX_COUNT - 1] = Some(spell);
204        true
205    }
206    
207    /// Removes a spell from the set. Returns whether the value was present in the set.
208    /// 
209    /// # Examples
210    /// ```
211    /// use tf2_enum::{SpellSet, Spell, AttributeSet};
212    /// 
213    /// let mut spells = SpellSet::single(Spell::HeadlessHorseshoes);
214    /// 
215    /// assert!(spells.remove(&Spell::HeadlessHorseshoes));
216    /// assert!(!spells.contains(&Spell::HeadlessHorseshoes));
217    /// ```
218    fn remove(&mut self, spell: &Spell) -> bool {
219        if self.inner[0] == Some(*spell) {
220            self.inner[0] = None;
221            true
222        } else if self.inner[1] == Some(*spell) {
223            self.inner[1] = None;
224            true
225        } else {
226            false
227        }
228    }
229    
230    /// Removes and returns the spell in the set, if any, that is equal to the given one.
231    fn take(&mut self, spell: &Spell) -> Option<Spell> {
232        if self.inner[0] == Some(*spell) {
233            self.inner[0] = None;
234            return Some(*spell);
235        } else if self.inner[1] == Some(*spell) {
236            self.inner[1] = None;
237            return Some(*spell);
238        }
239        
240        None
241    }
242    
243    /// Replaces a spell in the set with a new spell. `false` if the spell was not present.
244    fn replace(&mut self, spell: &Spell, new_spell: Spell) -> bool {
245        if !self.contains(spell) {
246            return false;
247        }
248        
249        for s in self.inner.iter_mut() {
250            if *s == Some(*spell) {
251                *s = Some(new_spell);
252                return true;
253            }
254        }
255        
256        false
257    }
258    
259    /// Converts each element to an [`ItemAttribute`]. 
260    fn iter_attributes(&self) -> impl Iterator<Item = ItemAttribute> {
261        self
262            .into_iter()
263            .map(ItemAttribute::from)
264    }
265    
266    /// Returns the inner storage as a slice.
267    fn as_slice(&self) -> &[Option<Spell>] {
268        &self.inner
269    }
270    
271    /// Returns the inner storage as a mutable slice.
272    fn as_mut_slice(&mut self) -> &mut [Option<Spell>] {
273        &mut self.inner
274    }
275}
276
277// Only Sub is implemented because Add wouldn't make much sense with spells being limited to 2.
278impl Sub for SpellSet {
279    type Output = Self;
280    
281    fn sub(self, other: Self) -> Self::Output {
282        self.difference(&other)
283    }
284}
285
286impl Sub for &SpellSet {
287    type Output = SpellSet;
288    
289    fn sub(self, other: &SpellSet) -> Self::Output {
290        self.difference(other)
291    }
292}
293
294impl BitAnd for SpellSet {
295    type Output = Self;
296    
297    fn bitand(self, other: Self) -> Self::Output {
298        self.intersection(&other)
299    }
300}
301
302impl BitAnd for &SpellSet {
303    type Output = SpellSet;
304    
305    fn bitand(self, other: &SpellSet) -> Self::Output {
306        self.intersection(other)
307    }
308}
309
310impl PartialEq<Self> for SpellSet {
311    fn eq(&self, other: &Self) -> bool {
312        (self.inner[0] == other.inner[0] && self.inner[1] == other.inner[1]) || 
313        (self.inner[0] == other.inner[1] && self.inner[1] == other.inner[0])
314    }
315}
316
317impl Hash for SpellSet {
318    fn hash<H: Hasher>(&self, state: &mut H) {
319        if self.inner[0] <= self.inner[1] {
320            self.inner[0].hash(state);
321            self.inner[1].hash(state);
322        } else {
323            self.inner[1].hash(state);
324            self.inner[0].hash(state);
325        }
326    }
327}
328
329impl From<[Option<Spell>; SPELL_COUNT]> for SpellSet {
330    fn from(mut inner: [Option<Spell>; SPELL_COUNT]) -> Self {
331        // remove duplicates
332        // since this only contains 2 spells it's not really necessary to do this using loops but
333        // the implementation is consistent with StrangePartSet
334        for i in 0..SPELL_COUNT {
335            if let Some(val_i) = inner[i] {
336                for j in 0..i {
337                    if let Some(val_j) = inner[j] {
338                        if val_i.attribute_defindex() == val_j.attribute_defindex() {
339                            inner[i] = None;
340                            break;
341                        } 
342                    }
343                }
344            }
345        }
346        
347        Self {
348            inner,
349        }
350    }
351}
352
353impl From<SpellSet> for Vec<Spell>{
354    fn from(spell_set: SpellSet) -> Self {
355        spell_set.into_iter().collect()
356    }
357}
358
359impl From<&SpellSet> for Vec<Spell> {
360    fn from(spell_set: &SpellSet) -> Self {
361        (*spell_set).into()
362    }
363}
364
365impl FromIterator<Spell> for SpellSet {
366    fn from_iter<I: IntoIterator<Item = Spell>>(iter: I) -> Self {
367        let mut spell_set = Self::new();
368        
369        for spell in iter {
370            spell_set.insert(spell);
371        }
372        
373        spell_set
374    }
375}
376
377impl<'a> FromIterator<&'a Spell> for SpellSet {
378    fn from_iter<I: IntoIterator<Item = &'a Spell>>(iter: I) -> Self {
379        let mut spell_set = Self::new();
380        
381        for spell in iter {
382            spell_set.insert(*spell);
383        }
384        
385        spell_set
386    }
387}
388
389impl FromIterator<Option<Spell>> for SpellSet {
390    fn from_iter<I: IntoIterator<Item = Option<Spell>>>(iter: I) -> Self {
391        let mut set = Self::new();
392        
393        for val in iter.into_iter().flatten() {
394            set.insert(val);
395        }
396        
397        set
398    }
399}
400
401impl<'a> FromIterator<&'a Option<Spell>> for SpellSet {
402    fn from_iter<I: IntoIterator<Item = &'a Option<Spell>>>(iter: I) -> Self {
403        let mut set = Self::new();
404        
405        for val in iter.into_iter().flatten() {
406            set.insert(*val);
407        }
408        
409        set
410    }
411}
412
413impl IntoIterator for SpellSet {
414    type Item = Spell;
415    type IntoIter = SpellSetIterator;
416    
417    fn into_iter(self) -> Self::IntoIter {
418        SpellSetIterator {
419            inner: self.inner.into_iter(),
420        }
421    }
422}
423
424impl IntoIterator for &SpellSet {
425    type Item = Spell;
426    type IntoIter = SpellSetIterator;
427    
428    fn into_iter(self) -> Self::IntoIter {
429        (*self).into_iter()
430    }
431}
432
433/// Iterator for spells.
434#[derive(Debug, Clone)]
435pub struct SpellSetIterator {
436    inner: std::array::IntoIter<Option<Spell>, SPELL_COUNT>,
437}
438
439impl Iterator for SpellSetIterator {
440    type Item = Spell;
441    
442    fn next(&mut self) -> Option<Self::Item> {
443        let iter = self.inner.by_ref();
444        
445        for opt in iter {
446            if opt.is_some() {
447                return opt;
448            }
449        }
450        
451        None
452    }
453}
454
455impl fmt::Display for SpellSet {
456    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
457        let mut iter = self.into_iter();
458        
459        if let Some(first) = iter.next() {
460            write!(f, "{first}")?;
461            
462            for s in iter {
463                write!(f, ", {s}")?;
464            }
465        }
466        
467        Ok(())
468    }
469}
470
471impl fmt::Debug for SpellSet {
472    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
473        f.write_str("{")?;
474        let mut first = true;
475        for part in self {
476            if !first {
477                f.write_str(", ")?;
478            }
479            write!(f, "{:?}", part as u32)?;
480            first = false;
481        }
482        f.write_str("}")
483    }
484}
485
486impl Serialize for SpellSet {
487    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
488    where
489        S: Serializer,
490    {
491        serialize::serialize_attribute_set(self, serializer)
492    }
493}
494
495impl<'de> Deserialize<'de> for SpellSet {
496    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
497    where
498        D: Deserializer<'de>,
499    {
500        struct StrangePartSetVisitor;
501        
502        impl<'de> Visitor<'de> for StrangePartSetVisitor {
503            type Value = SpellSet;
504
505            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
506                formatter.write_str("an array of maps with defindex, float_value")
507            }
508            
509            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
510            where
511                A: SeqAccess<'de>,
512            {
513                let mut set = Self::Value::new();
514                let mut defindex_map = HashSet::new();
515                
516                while let Some(map) = seq.next_element::<ItemAttribute>()? {
517                    if defindex_map.contains(&map.defindex) {
518                        // Skip if defindex is already in the set
519                        continue;
520                    }
521                    
522                    defindex_map.insert(map.defindex);
523                    
524                    match map.defindex {
525                        FootprintsSpell::DEFINDEX => {
526                            let float_value = map.float_value
527                                .ok_or_else(|| de::Error::missing_field("float_value"))?;
528                            let part = FootprintsSpell::try_from_attribute_float_value(
529                                    float_value
530                                )
531                                .ok_or_else(|| de::Error::custom(
532                                    "cannot convert from float_value"
533                                ))?;
534                            
535                            set.insert(part.into());
536                            continue;
537                        },
538                        PaintSpell::DEFINDEX => {
539                            let float_value = map.float_value
540                                .ok_or_else(|| de::Error::missing_field("float_value"))?;
541                            let part = PaintSpell::try_from_attribute_float_value(
542                                    float_value
543                                )
544                                .ok_or_else(|| de::Error::custom(
545                                    "cannot convert from float_value"
546                                ))?;
547                            
548                            set.insert(part.into());
549                            continue;
550                        },
551                        Spell::DEFINDEX_EXORCISM => {
552                            set.insert(Spell::Exorcism);
553                        },
554                        Spell::DEFINDEX_HALLOWEEN_FIRE => {
555                            set.insert(Spell::HalloweenFire);
556                        },
557                        Spell::DEFINDEX_VOICES_FROM_BELOW => {
558                            set.insert(Spell::VoicesFromBelow);
559                        },
560                        Spell::DEFINDEX_PUMPKIN_BOMBS => {
561                            set.insert(Spell::PumpkinBombs);
562                        },
563                        _ => continue,
564                    }
565                }
566                
567                Ok(set)
568            }
569        }
570        
571        deserializer.deserialize_seq(StrangePartSetVisitor)
572    }
573}
574    
575#[cfg(test)]
576mod tests {
577    use super::*;
578    use std::collections::HashSet;
579    
580    #[test]
581    fn serializes() {
582        let mut spell_set = SpellSet::new();
583        spell_set.insert(Spell::Exorcism);
584        spell_set.insert(Spell::HeadlessHorseshoes);
585        let serialized = serde_json::to_string(&spell_set).unwrap();
586                
587        assert_eq!(serialized, r#"[{"defindex":1009,"value":1065353216,"float_value":1},{"defindex":1005,"value":1073741824,"float_value":2}]"#);
588    }
589    
590    #[test]
591    fn base_methods() {
592        let mut spell_set = SpellSet::new();
593        spell_set.insert(Spell::Exorcism);
594        spell_set.insert(Spell::HalloweenFire);
595        assert_eq!(spell_set.len(), 2);
596        assert!(spell_set.contains(&Spell::Exorcism));
597        assert!(spell_set.contains(&Spell::HalloweenFire));
598        assert!(!spell_set.contains(&Spell::VoicesFromBelow));
599        assert!(spell_set.remove(&Spell::Exorcism));
600        assert!(!spell_set.contains(&Spell::Exorcism));
601        assert!(spell_set.take(&Spell::HalloweenFire).is_some());
602        assert!(!spell_set.contains(&Spell::HalloweenFire));
603        assert!(spell_set.is_empty());
604    }
605    
606    #[test]
607    fn spell_set_equals() {
608        assert_eq!(SpellSet::from([
609            Some(Spell::Exorcism),
610            Some(Spell::HalloweenFire),
611        ]), SpellSet::from([
612            Some(Spell::HalloweenFire),
613            Some(Spell::Exorcism),
614        ]));
615    }
616    
617    #[test]
618    fn spell_set_hashes() {
619        let mut set = HashSet::new();
620        
621        set.insert(SpellSet::from([
622            Some(Spell::Exorcism),
623            Some(Spell::HalloweenFire),
624        ]));
625        
626        assert!(set.contains(&SpellSet::from([
627            Some(Spell::HalloweenFire),
628            Some(Spell::Exorcism),
629        ])));
630    }
631    
632    #[test]
633    fn spell_set_no_duplicates() {
634        assert_eq!(SpellSet::from([
635            Some(Spell::Exorcism),
636            Some(Spell::Exorcism),
637        ]), SpellSet::from([
638            Some(Spell::Exorcism),
639            None,
640        ]));
641        
642        assert_eq!(SpellSet::from([
643            Some(Spell::TeamSpiritFootprints),
644            Some(Spell::HeadlessHorseshoes),
645        ]), SpellSet::from([
646            Some(Spell::TeamSpiritFootprints),
647            None,
648        ]));
649    }
650    
651    #[test]
652    fn iterates_spells() {
653        let spells = SpellSet::from([
654            Some(Spell::Exorcism),
655            Some(Spell::HalloweenFire),
656        ]);
657        let mut count = 0;
658        
659        for _spell in spells {
660            count += 1;
661        }
662        
663        assert_eq!(count, 2);
664        
665        let spells = spells.into_iter().collect::<Vec<_>>();
666        
667        assert_eq!(spells, vec![Spell::Exorcism, Spell::HalloweenFire]);
668        
669        let spells = SpellSet::from_iter(spells);
670        
671        assert_eq!(spells, SpellSet::from([
672            Some(Spell::Exorcism),
673            Some(Spell::HalloweenFire),
674        ]));
675    }
676    
677    #[test]
678    fn sub() {
679        let spells1 = SpellSet::from([
680            Some(Spell::HalloweenFire),
681            Some(Spell::Exorcism),
682        ]);
683        let spells2 = SpellSet::from([
684            Some(Spell::HalloweenFire),
685            Some(Spell::VoicesFromBelow),
686        ]);
687        
688        let difference = spells1 - spells2;
689        
690        assert_eq!(difference, SpellSet::from([
691            Some(Spell::Exorcism),
692            None,
693        ]));
694    }
695    
696    #[test]
697    fn stringify() {
698        let spells = SpellSet::from([
699            Some(Spell::Exorcism),
700            Some(Spell::HalloweenFire),
701        ]);
702        
703        assert_eq!(spells.to_string(), "Exorcism, Halloween Fire");
704    }
705}