1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! Multiset is used as a set with multiple occurances of items.

use std::collections::hash_map::{HashMap, Iter as HashIter};
use std::hash::Hash;

/// A multiset or 'bag' is a set where items can be added multiple times
///
/// Example: '<< 1,1,2,4,4,4 >>'
///
/// Implementation detail: this set drops the redundant items and keeps a reference count
#[derive(Debug)]
pub struct MultiSet<T> where
    T: PartialEq + Eq + Hash
{
    data: HashMap<T, usize>,
}

impl<T> MultiSet<T> where T: PartialEq + Eq + Hash {
    pub fn new() -> Self {
        MultiSet { data: HashMap::new() }
    }

    pub fn with_capacity(capacity: usize) -> Self {
        MultiSet { data: HashMap::with_capacity(capacity) }
    }

    /// Returns the number of unique items the set can hold without allocating
    pub fn capacity(&self) -> usize {
        self.data.capacity()
    }

    /// Drop all items from the set
    pub fn clear(&mut self) {
        self.data.clear()
    }

    /// Returns true if no items exist in the set
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Returns true if the item occurs at least once
    pub fn contains(&self, value: &T) -> bool
    {
        self.data.get(value).map_or(false, |_size| true)
    }

    /// Returns the number of occurrences of the given item in the set
    pub fn count(&self, value: &T) -> usize
    {
        self.data.get(value).map_or(0, |size| *size)
    }

    /// Insert an item to the set.
    ///
    /// This method will drop the item if it already exists in the set, and increases the count
    /// instead, using reference counting instead of data duplication.
    pub fn insert(&mut self, value: T)
    {
        match self.data.get_mut(&value) {
            Some(size) => { *size = *size + 1; }
            None => { self.data.insert(value, 1); }
        };
    }

    /// Removes an item from the set, and returns wetter it was present or not.
    pub fn remove(&mut self, value: &T) -> bool
    {
        if match self.data.get_mut(value) {
            Some(size) => {
                *size = *size + 1;
                *size
            }
            None => { return false; }
        } == 0 {
            self.data.remove(value);
        }
        true
    }

    pub fn iter(&self) -> Iter<T> {
        Iter {
            it: self.data.iter(),
            remaining: 0,
            current: None,
        }
    }
}

/// Iterates over all the items in the multiset
pub struct Iter<'a, T: 'a> {
    it: HashIter<'a, T, usize>,
    remaining: usize,
    current: Option<&'a T>,
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        match self.current {
            Some(_t) => {
                self.remaining = self.remaining - 1;
                if self.remaining == 0 {
                    return self.current.take();
                } else {
                    self.current
                }
            }
            None => {
                if let Some((t, size)) = self.it.next() {
                    if *size > 1 {
                        self.remaining = *size - 1;
                        self.current = Some(t);
                    }
                    Some(t)
                } else {
                    None
                }
            }
        }
    }
}