1
use alloc::borrow::ToOwned;
2
use alloc::rc::Rc;
3
use alloc::string::String;
4
use alloc::vec;
5
use alloc::vec::Vec;
6
use core::borrow::Borrow;
7
use std::println;
8

            
9
use crate::map::{Entry, Field, Map};
10
use crate::Set;
11

            
12
1
#[test]
13
1
fn basics() {
14
1
    let mut map = Map::default();
15
1
    assert!(map.is_empty());
16
1
    assert!(map.insert("b", 1).is_none());
17
1
    assert_eq!(map.len(), 1);
18
1
    assert_eq!(map.insert("b", 2), Some(Field::new("b", 1)));
19
1
    assert_eq!(map.len(), 1);
20

            
21
1
    assert!(map.insert("a", 1).is_none());
22
1
    assert_eq!(map.len(), 2);
23

            
24
1
    assert!(map.contains(&"a"));
25
1
    assert_eq!(map.get(&"a"), Some(&1));
26
1
    assert!(map.contains(&"b"));
27
1
    assert_eq!(map.get(&"b"), Some(&2));
28
1
    assert!(!map.contains(&"c"));
29
1
    assert_eq!(map.get(&"c"), None);
30

            
31
1
    assert_eq!(*map.field(0).unwrap().key(), "a");
32
1
    assert_eq!(*map.field(1).unwrap().key(), "b");
33
1
    assert!(map.field(2).is_none());
34
1
    map.field_mut(1).unwrap().value += 1;
35
1

            
36
1
    // Various iteration.
37
1
    let mut iter = map.iter();
38
1
    assert_eq!(iter.next().unwrap(), &Field::new("a", 1));
39
1
    assert_eq!(iter.next().unwrap(), &Field::new("b", 3));
40
1
    let mut iter = map.values();
41
1
    assert_eq!(iter.next().unwrap(), &1);
42
1
    assert_eq!(iter.next().unwrap(), &3);
43
1
    let mut iter = map.clone().into_iter();
44
1
    assert_eq!(iter.next().unwrap(), Field::new("a", 1));
45
1
    assert_eq!(iter.next().unwrap(), Field::new("b", 3));
46
1
    let mut iter = map.clone().into_values();
47
1
    assert_eq!(iter.next().unwrap(), 1);
48
1
    assert_eq!(iter.next().unwrap(), 3);
49

            
50
    // Increment via iter_mut
51
    #[allow(clippy::explicit_iter_loop)] // for coverage
52
2
    for (_, value) in map.iter_mut() {
53
2
        *value += 1;
54
2
    }
55

            
56
    // Increment via values_mut
57
2
    for value in map.values_mut() {
58
2
        *value += 1;
59
2
    }
60

            
61
    // Removal
62
1
    assert_eq!(map.remove(&"a"), Some(Field::new("a", 3)));
63
1
    assert_eq!(map.remove(&"a"), None);
64

            
65
    // Drain
66
1
    let drained = map.drain();
67
1
    drop(drained);
68
1
    assert!(map.is_empty());
69
1
}
70

            
71
1
#[test]
72
1
fn clear_and_shrink() {
73
1
    let mut map = Map::<&'static str, usize>::with_capacity(10);
74
1
    map.insert("a", 1);
75
1
    assert_eq!(map.capacity(), 10);
76
1
    map.shrink_to(0);
77
1
    assert_eq!(map.capacity(), 1);
78
1
    map.clear();
79
1
    assert_eq!(map.capacity(), 1);
80
1
    assert!(map.is_empty());
81

            
82
1
    let mut map = Map::<&'static str, usize>::with_capacity(10);
83
1
    map.insert("a", 1);
84
1
    assert_eq!(map.capacity(), 10);
85
1
    map.shrink_to_fit();
86
1
    assert_eq!(map.capacity(), 1);
87

            
88
    // Set
89
1
    let mut map = Set::<&'static str>::with_capacity(10);
90
1
    map.insert("a");
91
1
    assert_eq!(map.capacity(), 10);
92
1
    map.shrink_to(0);
93
1
    assert_eq!(map.capacity(), 1);
94
1
    map.clear();
95
1
    assert_eq!(map.capacity(), 1);
96
1
    assert!(map.is_empty());
97

            
98
1
    let mut map = Set::<&'static str>::with_capacity(10);
99
1
    map.insert("a");
100
1
    assert_eq!(map.capacity(), 10);
101
1
    map.shrink_to_fit();
102
1
    assert_eq!(map.capacity(), 1);
103
1
}
104

            
105
1
#[test]
106
1
fn entry() {
107
1
    let mut map = Map::<String, usize>::new();
108
1
    let entry = map.entry("a").or_insert(1);
109
1
    assert_eq!(*entry, 1);
110
1
    let entry = map
111
1
        .entry(String::from("a"))
112
1
        .and_modify(|value| *value += 1)
113
1
        .or_insert_with(|| unreachable!());
114
1
    assert_eq!(*entry, 2);
115
1
    let entry = map
116
1
        .entry(&String::from("b"))
117
1
        .and_modify(|_| unreachable!())
118
1
        .or_insert_with(|| 1);
119
1
    assert_eq!(*entry, 1);
120

            
121
1
    let entry = map.entry("a").or_insert(0);
122
1
    assert_eq!(*entry, 2);
123

            
124
1
    let Entry::Occupied(entry) = map.entry("a") else {
125
        unreachable!()
126
    };
127
1
    assert_eq!(entry.key(), "a");
128
1
    assert_eq!(*entry, 2);
129
1
    let removed = entry.remove();
130
1
    assert_eq!(removed.key(), "a");
131
1
    assert_eq!(removed.value, 2);
132

            
133
1
    let Entry::Occupied(entry) = map.entry("b") else {
134
        unreachable!()
135
    };
136
1
    assert_eq!(entry.replace(2), 1);
137
1
    assert_eq!(map["b"], 2);
138

            
139
1
    assert_eq!(*map.entry("c").or_default(), 0);
140

            
141
    // Entry with [u8]/Vec<u8>
142
1
    let mut map = Map::<Vec<u8>, usize>::new();
143
1
    map.entry(vec![b'a']).or_insert(1);
144
1
    map.entry(&b"a"[..]).or_insert(1);
145
1

            
146
1
    let mut map = Map::<CustomType, usize>::new();
147
1
    let entry = map.entry(&CustomTypeBorrowed(1)).or_insert(42);
148
1
    assert_eq!(*entry, 42);
149
1
    let entry = map
150
1
        .entry(CustomType::new(1))
151
1
        .or_insert_with(|| unreachable!("key should be found"));
152
1
    assert_eq!(*entry, 42);
153
1
}
154

            
155
1
#[derive(Ord, PartialOrd, Eq, PartialEq, Debug)]
156
pub struct CustomType(CustomTypeBorrowed);
157

            
158
impl CustomType {
159
1
    pub fn new(value: usize) -> Self {
160
1
        Self(CustomTypeBorrowed(value))
161
1
    }
162
}
163

            
164
1
#[derive(Ord, PartialOrd, Eq, PartialEq, Debug)]
165
pub struct CustomTypeBorrowed(usize);
166

            
167
impl Borrow<CustomTypeBorrowed> for CustomType {
168
    fn borrow(&self) -> &CustomTypeBorrowed {
169
        &self.0
170
    }
171
}
172

            
173
impl ToOwned for CustomTypeBorrowed {
174
    type Owned = CustomType;
175

            
176
1
    fn to_owned(&self) -> Self::Owned {
177
1
        CustomType(CustomTypeBorrowed(self.0))
178
1
    }
179
}
180

            
181
impl ToOwned for CustomType {
182
    type Owned = Self;
183

            
184
    fn to_owned(&self) -> Self::Owned {
185
        CustomType(CustomTypeBorrowed(self.0 .0))
186
    }
187
}
188

            
189
impl PartialOrd<CustomTypeBorrowed> for CustomType {
190
    fn partial_cmp(&self, other: &CustomTypeBorrowed) -> Option<core::cmp::Ordering> {
191
        self.0.partial_cmp(other)
192
    }
193
}
194

            
195
impl PartialEq<CustomTypeBorrowed> for CustomType {
196
    fn eq(&self, other: &CustomTypeBorrowed) -> bool {
197
        self.0.eq(other)
198
    }
199
}
200

            
201
1
#[test]
202
1
fn binary_search_extremes() {
203
1
    // fill in 0..100 in two passes: first with evens, second with odds. This
204
1
    // should hit every possible combination of the binary search algorithm.
205
1
    let mut map = Map::new();
206
50
    for i in (0..100).step_by(2) {
207
50
        map.insert(i, i);
208
50
    }
209
50
    for i in (1..100).step_by(2) {
210
50
        map.insert(i, i);
211
50
    }
212
101
    for i in 0..100 {
213
100
        assert_eq!(map.get(&i), Some(&i));
214
    }
215
1
}
216

            
217
1
#[test]
218
1
fn merge() {
219
49
    let multiples_of_two = (2..100).step_by(2).map(|i| (i, 1)).collect::<Map<_, _>>();
220
33
    let multiples_of_three = (3..100).step_by(3).map(|i| (i, 1)).collect::<Map<_, _>>();
221
69
    let copy_if_not_five = |key: &usize, value: &usize| (*key % 5 != 0).then_some(*value);
222
1
    let multiples_of_2_and_3_but_not_5 = Map::new()
223
1
        .merged_with(
224
1
            &multiples_of_two,
225
1
            copy_if_not_five,
226
1
            |_key, _existing, _incoming| unreachable!(),
227
1
        )
228
1
        .merged_with(
229
1
            &multiples_of_three,
230
1
            copy_if_not_five,
231
13
            |_key, existing, incoming| *existing += *incoming,
232
1
        );
233
1
    println!(
234
1
        "All: {multiples_of_2_and_3_but_not_5:?}, {}",
235
1
        multiples_of_2_and_3_but_not_5.len()
236
1
    );
237
1
    assert_eq!(multiples_of_2_and_3_but_not_5.get(&2), Some(&1));
238
1
    assert_eq!(multiples_of_2_and_3_but_not_5.get(&3), Some(&1));
239
1
    assert_eq!(multiples_of_2_and_3_but_not_5.get(&6), Some(&2));
240
1
    assert_eq!(multiples_of_2_and_3_but_not_5.get(&30), None);
241
1
    assert_eq!(multiples_of_2_and_3_but_not_5.len(), 54);
242
1
}
243

            
244
1
#[test]
245
1
fn entry_to_owned_on_insert() {
246
1
    #[derive(Ord, PartialOrd, Eq, PartialEq)]
247
1
    struct NotCloneable;
248
1

            
249
1
    impl Clone for NotCloneable {
250
1
        fn clone(&self) -> Self {
251
            unreachable!()
252
1
        }
253
1
    }
254
1

            
255
1
    let rc = Rc::new(0);
256
1
    let mut map = Map::<Rc<usize>, ()>::new();
257
1
    map.entry(&rc);
258
1
    assert_eq!(Rc::strong_count(&rc), 1);
259
1
    map.entry(&rc).or_insert(());
260
1
    assert_eq!(Rc::strong_count(&rc), 2);
261

            
262
    // This final test proves that when passing in the owned copy, it is used
263
    // without being cloned.
264
1
    let mut map = Map::<NotCloneable, ()>::new();
265
1
    map.entry(NotCloneable).or_insert(());
266
1
    assert!(map.contains(&NotCloneable));
267
1
}
268

            
269
1
#[test]
270
1
fn capacity() {
271
1
    let mut map = Map::with_capacity(1);
272
1
    assert_eq!(map.capacity(), 1);
273
1
    map.insert(1, 1);
274
1
    assert_eq!(map.capacity(), 1);
275
1
    map.insert(2, 2);
276
1
    assert!(map.capacity() > 1);
277
1
}
278

            
279
1
#[test]
280
1
fn insert_with() {
281
1
    let mut map = Map::with_capacity(1);
282
1
    assert_eq!(map.insert_with("a", || 1), None);
283
1
    assert!(map.contains(&"a"));
284
1
    assert_eq!(map.insert_with("a", || unreachable!()), Some("a"));
285
1
}
286

            
287
1
#[test]
288
1
fn field() {
289
1
    let field = Field::new("a", 1);
290
1
    assert_eq!(field.key(), &"a");
291
1
    assert_eq!(field.value, 1);
292
1
    assert_eq!(field.into_key(), "a");
293
1
}
294

            
295
1
#[test]
296
1
fn vacant_entry_key() {
297
1
    let mut map = Map::<String, i32>::with_capacity(1);
298
1
    let borrowed = "a";
299
1
    let borrowed_ptr = borrowed.as_ptr();
300
1
    let Entry::Vacant(entry) = map.entry(borrowed) else {
301
        unreachable!()
302
    };
303
    // The key is still borrowed at this point.
304
1
    assert_eq!(entry.key().as_ptr(), borrowed_ptr);
305

            
306
    // This test just follows the path from taking the owned type to the
307
    // borrowed type.
308
1
    let Entry::Vacant(entry) = map.entry(String::from("a")) else {
309
        unreachable!()
310
    };
311
1
    assert_eq!(entry.key(), "a");
312
1
}
313

            
314
1
#[test]
315
1
fn union() {
316
1
    let mut a = Map::new();
317
1
    a.insert("a", 1);
318
1
    a.insert("b", 2);
319
1
    a.insert("c", 3);
320
1
    let mut b = Map::new();
321
1
    b.insert("b", 2);
322
1
    b.insert("d", 4);
323
1
    let merged = a
324
1
        .union(&b)
325
4
        .map(|unioned| unioned.map_both(|_key, a, b| *a + *b).into_owned())
326
1
        .collect::<Map<_, _>>();
327
1
    assert_eq!(merged.get(&"a"), Some(&1));
328
1
    assert_eq!(merged.get(&"b"), Some(&4));
329
1
    assert_eq!(merged.get(&"c"), Some(&3));
330
1
    assert_eq!(merged.get(&"d"), Some(&4));
331
1
    assert_eq!(merged.len(), 4);
332
1
}
333

            
334
1
#[test]
335
1
fn unioned_map_both_ref() {
336
1
    let mut a = Map::new();
337
1
    a.insert("a", 1);
338
1
    a.insert("b", 2);
339
1
    a.insert("c", 3);
340
1
    let mut b = Map::new();
341
1
    b.insert("b", 42); // This value will not make it to the result.
342
1
    b.insert("d", 4);
343
1
    let merged = a
344
1
        .union(&b)
345
4
        .map(|unioned| unioned.map_both(|_key, a, _b| a).into_owned())
346
1
        .collect::<Map<_, _>>();
347
1
    assert_eq!(merged.get(&"a"), Some(&1));
348
1
    assert_eq!(merged.get(&"b"), Some(&2));
349
1
    assert_eq!(merged.get(&"c"), Some(&3));
350
1
    assert_eq!(merged.get(&"d"), Some(&4));
351
1
    assert_eq!(merged.len(), 4);
352
1
}