identity_core/common/
ordered_set.rs

1// Copyright 2020-2022 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use core::convert::TryFrom;
5use core::fmt::Debug;
6use core::fmt::Formatter;
7use core::iter::FromIterator;
8use core::ops::Deref;
9use core::slice::Iter;
10use core::slice::IterMut;
11use std::vec::IntoIter;
12
13use serde::Deserialize;
14use serde::Serialize;
15
16use crate::common::KeyComparable;
17use crate::error::Error;
18use crate::error::Result;
19
20/// An ordered set backed by a `Vec<T>`.
21///
22/// Note: Ordering is based on insert order and **not** [`Ord`].
23///
24/// See the [Infra standard definition](https://infra.spec.whatwg.org/#ordered-set).
25#[derive(Clone, Hash, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize)]
26#[repr(transparent)]
27#[serde(bound(deserialize = "T: KeyComparable + Deserialize<'de>"), try_from = "Vec<T>")]
28pub struct OrderedSet<T>(Vec<T>);
29
30impl<T> OrderedSet<T> {
31  /// Creates a new `OrderedSet`.
32  #[inline]
33  pub const fn new() -> Self {
34    Self(Vec::new())
35  }
36
37  /// Creates a new `OrderedSet` with the specified capacity.
38  #[inline]
39  pub fn with_capacity(capacity: usize) -> Self {
40    Self(Vec::with_capacity(capacity))
41  }
42
43  /// Returns the number of elements in the `OrderedSet`.
44  #[inline]
45  pub fn len(&self) -> usize {
46    self.0.len()
47  }
48
49  /// Returns `true` if the `OrderedSet` contains no elements.
50  #[inline]
51  pub fn is_empty(&self) -> bool {
52    self.0.is_empty()
53  }
54
55  /// Returns an iterator over the slice of elements.
56  #[inline]
57  pub fn iter(&self) -> Iter<'_, T> {
58    self.0.iter()
59  }
60
61  /// Returns an iterator that allows modifying each value.
62  ///
63  /// WARNING: improper usage of this allows violating the key-uniqueness of the OrderedSet.
64  #[inline]
65  pub fn iter_mut_unchecked(&mut self) -> IterMut<'_, T> {
66    self.0.iter_mut()
67  }
68
69  /// Returns the first element in the set, or `None` if the set is empty.
70  #[inline]
71  pub fn head(&self) -> Option<&T> {
72    self.0.first()
73  }
74
75  /// Returns a mutable reference to the first element in the set, or `None` if
76  /// the set is empty.
77  ///
78  /// # Warning
79  /// Incorrect use of this can break the invariants of [`OrderedSet`].
80  #[inline]
81  pub fn head_mut(&mut self) -> Option<&mut T> {
82    self.0.first_mut()
83  }
84
85  /// Returns the last element in the set, or `None` if the set is empty.
86  #[inline]
87  pub fn tail(&self) -> Option<&T> {
88    self.0.last()
89  }
90
91  /// Returns a mutable reference the last element in the set, or `None` if the
92  /// set is empty.
93  ///
94  /// # Warning
95  /// Incorrect use of this can break the invariants of [`OrderedSet`].
96  #[inline]
97  pub fn tail_mut(&mut self) -> Option<&mut T> {
98    self.0.last_mut()
99  }
100
101  /// Returns a slice containing all elements in the `OrderedSet`.
102  #[inline]
103  pub fn as_slice(&self) -> &[T] {
104    &self.0
105  }
106
107  /// Consumes the `OrderedSet` and returns the elements as a `Vec<T>`.
108  #[inline]
109  pub fn into_vec(self) -> Vec<T> {
110    self.0
111  }
112
113  /// Clears the `OrderedSet`, removing all values.
114  #[inline]
115  pub fn clear(&mut self) {
116    self.0.clear();
117  }
118
119  /// Returns `true` if the `OrderedSet` contains the given value.
120  pub fn contains<U>(&self, item: &U) -> bool
121  where
122    T: KeyComparable,
123    U: KeyComparable<Key = T::Key> + ?Sized,
124  {
125    self.0.iter().any(|other| other.key() == item.key())
126  }
127
128  /// Adds a new value to the end of the `OrderedSet`; returns `true` if the
129  /// value was successfully added.
130  pub fn append(&mut self, item: T) -> bool
131  where
132    T: KeyComparable,
133  {
134    if self.contains(&item) {
135      false
136    } else {
137      self.0.push(item);
138      true
139    }
140  }
141
142  /// Adds a new value to the start of the `OrderedSet`; returns `true` if the
143  /// value was successfully added.
144  pub fn prepend(&mut self, item: T) -> bool
145  where
146    T: KeyComparable,
147  {
148    if self.contains(&item) {
149      false
150    } else {
151      self.0.insert(0, item);
152      true
153    }
154  }
155
156  /// Replaces a `current` value with the given `update` value; returns `true`
157  /// if the value was successfully replaced.
158  #[inline]
159  pub fn replace<U>(&mut self, current: &U, update: T) -> bool
160  where
161    T: KeyComparable,
162    U: KeyComparable<Key = T::Key>,
163  {
164    self.change(update, |item, update| {
165      item.key() == current.key() || item.key() == update.key()
166    })
167  }
168
169  /// Updates an existing value in the `OrderedSet`; returns `true` if the value
170  /// was successfully updated.
171  #[inline]
172  pub fn update(&mut self, update: T) -> bool
173  where
174    T: KeyComparable,
175  {
176    self.change(update, |item, update| item.key() == update.key())
177  }
178
179  /// Removes and returns the item with the matching key from the set, if it exists.
180  #[inline]
181  pub fn remove<U>(&mut self, item: &U) -> Option<T>
182  where
183    T: KeyComparable,
184    U: KeyComparable<Key = T::Key>,
185  {
186    self
187      .0
188      .iter()
189      .enumerate()
190      .find(|(_, entry)| entry.key() == item.key())
191      .map(|(idx, _)| idx)
192      .map(|idx| self.0.remove(idx))
193  }
194
195  fn change<F>(&mut self, data: T, f: F) -> bool
196  where
197    F: Fn(&T, &T) -> bool,
198  {
199    let index: Option<usize> = self.0.iter().position(|item| f(item, &data));
200
201    if let Some(index) = index {
202      let keep: Vec<T> = self.0.drain(index..).filter(|item| !f(item, &data)).collect();
203
204      self.0.extend(keep);
205      self.0.insert(index, data);
206    }
207
208    index.is_some()
209  }
210}
211
212impl<T> IntoIterator for OrderedSet<T> {
213  type Item = T;
214  type IntoIter = IntoIter<T>;
215
216  fn into_iter(self) -> Self::IntoIter {
217    self.0.into_iter()
218  }
219}
220
221impl<T> Debug for OrderedSet<T>
222where
223  T: Debug,
224{
225  #[inline]
226  fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
227    f.debug_set().entries(self.iter()).finish()
228  }
229}
230
231impl<T> Deref for OrderedSet<T> {
232  type Target = [T];
233
234  #[inline]
235  fn deref(&self) -> &Self::Target {
236    &self.0
237  }
238}
239
240impl<T: KeyComparable> Default for OrderedSet<T> {
241  #[inline]
242  fn default() -> Self {
243    Self::new()
244  }
245}
246
247impl<T: KeyComparable> FromIterator<T> for OrderedSet<T>
248where
249  T: KeyComparable,
250{
251  fn from_iter<I>(iter: I) -> Self
252  where
253    I: IntoIterator<Item = T>,
254  {
255    let iter = iter.into_iter();
256    let size: usize = iter.size_hint().1.unwrap_or(0);
257
258    let mut this: Self = Self::with_capacity(size);
259
260    // Ignore duplicates.
261    for item in iter {
262      this.append(item);
263    }
264
265    this
266  }
267}
268
269impl<T> TryFrom<Vec<T>> for OrderedSet<T>
270where
271  T: KeyComparable,
272{
273  type Error = Error;
274
275  fn try_from(other: Vec<T>) -> Result<Self, Self::Error> {
276    let mut this: Self = Self::with_capacity(other.len());
277
278    for item in other {
279      if !this.append(item) {
280        return Err(Error::OrderedSetDuplicate);
281      }
282    }
283
284    Ok(this)
285  }
286}
287
288#[cfg(test)]
289mod tests {
290  use std::collections::HashSet;
291  use std::hash::Hash;
292
293  use super::*;
294  use proptest::prelude::Rng;
295  use proptest::strategy::Strategy;
296  use proptest::test_runner::TestRng;
297  use proptest::*;
298
299  #[test]
300  fn test_ordered_set_works() {
301    let mut set = OrderedSet::new();
302
303    set.append("a");
304    set.append("b");
305    set.append("c");
306
307    assert_eq!(set.as_slice(), &["a", "b", "c"]);
308    assert_eq!(set.head(), Some(&"a"));
309    assert_eq!(set.tail(), Some(&"c"));
310
311    set.replace(&"a", "c");
312
313    assert_eq!(set.as_slice(), &["c", "b"]);
314
315    let mut set = OrderedSet::new();
316
317    set.prepend("a");
318    set.prepend("b");
319    set.prepend("c");
320
321    assert_eq!(set.as_slice(), &["c", "b", "a"]);
322    assert_eq!(set.head(), Some(&"c"));
323    assert_eq!(set.tail(), Some(&"a"));
324
325    set.replace(&"a", "c");
326
327    assert_eq!(set.as_slice(), &["c", "b"]);
328  }
329
330  #[test]
331  fn test_from_vec_valid() {
332    let source: Vec<u8> = vec![3, 1, 2, 0];
333    let set: OrderedSet<u8> = OrderedSet::try_from(source).unwrap();
334
335    assert_eq!(&*set, &[3, 1, 2, 0]);
336  }
337
338  #[test]
339  #[should_panic = "OrderedSetDuplicate"]
340  fn test_from_vec_invalid() {
341    let source: Vec<u8> = vec![1, 2, 2, 5];
342    let _: OrderedSet<u8> = OrderedSet::try_from(source).unwrap();
343  }
344
345  #[test]
346  fn test_collect() {
347    let source: Vec<u8> = vec![1, 2, 3, 3, 2, 4, 5, 1, 1];
348    let set: OrderedSet<u8> = source.into_iter().collect();
349
350    assert_eq!(&*set, &[1, 2, 3, 4, 5]);
351  }
352
353  #[test]
354  fn test_contains() {
355    let cs1 = ComparableStruct { key: 0, value: 10 };
356    let cs2 = ComparableStruct { key: 1, value: 20 };
357    let cs3 = ComparableStruct { key: 2, value: 10 };
358    let cs4 = ComparableStruct { key: 3, value: 20 };
359
360    let source: Vec<ComparableStruct> = vec![cs1, cs2];
361    let set: OrderedSet<ComparableStruct> = source.into_iter().collect();
362
363    assert!(set.contains(&cs1));
364    assert!(set.contains(&cs2));
365    assert!(!set.contains(&cs3));
366    assert!(!set.contains(&cs4));
367  }
368
369  #[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
370  struct ComparableStruct {
371    key: u8,
372    value: i32,
373  }
374
375  impl KeyComparable for ComparableStruct {
376    type Key = u8;
377
378    #[inline]
379    fn key(&self) -> &Self::Key {
380      &self.key
381    }
382  }
383
384  #[test]
385  fn test_ordered_set_replace() {
386    let mut set = OrderedSet::new();
387
388    // Create two structs with the same key.
389    let cs1 = ComparableStruct { key: 0, value: 10 };
390    let cs2 = ComparableStruct { key: 0, value: 20 };
391
392    // Try replace it with the second.
393    // This should succeed because the keys are equivalent.
394    assert!(set.append(cs1));
395    assert_eq!(set.len(), 1);
396
397    assert!(set.replace(&cs1, cs2));
398    assert_eq!(set.len(), 1);
399    assert_eq!(set.head().unwrap().key, cs2.key);
400    assert_eq!(set.head().unwrap().value, cs2.value);
401  }
402
403  #[test]
404  fn test_ordered_set_replace_all() {
405    let mut set = OrderedSet::new();
406    let cs1 = ComparableStruct { key: 0, value: 10 };
407    let cs2 = ComparableStruct { key: 1, value: 20 };
408    assert!(set.append(cs1));
409    assert!(set.append(cs2));
410    assert_eq!(set.len(), 2);
411
412    // Now replace cs1 with something that has the same key as cs2.
413    // This should replace BOTH cs1 AND cs2.
414    let cs3 = ComparableStruct { key: 1, value: 30 };
415    assert!(set.replace(&cs1, cs3));
416    assert_eq!(set.len(), 1);
417    assert_eq!(set.head().unwrap().key, cs3.key);
418    assert_eq!(set.head().unwrap().value, cs3.value);
419  }
420
421  #[test]
422  fn test_ordered_set_update() {
423    let mut set = OrderedSet::new();
424    let cs1 = ComparableStruct { key: 0, value: 10 };
425    assert!(set.append(cs1));
426    assert_eq!(set.len(), 1);
427
428    // This should update the value of cs1 since the keys are the same.
429    let cs2 = ComparableStruct { key: 0, value: 20 };
430    assert!(set.update(cs2));
431    assert_eq!(set.len(), 1);
432    assert_eq!(set.head().unwrap().key, cs2.key);
433    assert_eq!(set.head().unwrap().value, cs2.value);
434
435    // This should NOT update anything since the key does not match.
436    let cs3 = ComparableStruct { key: 1, value: 20 };
437    assert!(!set.update(cs3));
438    assert_eq!(set.len(), 1);
439    assert_eq!(set.head().unwrap().key, cs2.key);
440    assert_eq!(set.head().unwrap().value, cs2.value);
441  }
442
443  // ===========================================================================================================================
444  // Test key uniqueness invariant with randomly generated input
445  // ===========================================================================================================================
446
447  /// Produce a strategy for generating a pair of ordered sets of ComparableStruct
448  fn arbitrary_sets_comparable_struct(
449  ) -> impl Strategy<Value = (OrderedSet<ComparableStruct>, OrderedSet<ComparableStruct>)> {
450    proptest::arbitrary::any::<Vec<(u8, i32)>>().prop_map(|mut x_vec| {
451      let half = x_vec.len() / 2;
452      let y_vec = x_vec.split_off(half);
453      let mapper = |(key, value)| ComparableStruct { key, value };
454      (
455        x_vec.into_iter().map(mapper).collect(),
456        y_vec.into_iter().map(mapper).collect(),
457      )
458    })
459  }
460
461  /// Produce a strategy for generating a pair of ordered sets of u128
462  fn arbitrary_sets_u128() -> impl Strategy<Value = (OrderedSet<u128>, OrderedSet<u128>)> {
463    proptest::arbitrary::any::<Vec<u128>>().prop_map(|mut x_vec| {
464      let half = x_vec.len() / 2;
465      let y_vec = x_vec.split_off(half);
466      (x_vec.into_iter().collect(), y_vec.into_iter().collect())
467    })
468  }
469
470  /// Trait for replacing the key of a KeyComparable value
471  trait ReplaceKey: KeyComparable {
472    fn set_key(&mut self, key: Self::Key);
473  }
474
475  impl ReplaceKey for ComparableStruct {
476    fn set_key(&mut self, key: Self::Key) {
477      let ComparableStruct { key: current_key, .. } = self;
478      *current_key = key;
479    }
480  }
481
482  impl ReplaceKey for u128 {
483    fn set_key(&mut self, key: Self::Key) {
484      *self = key;
485    }
486  }
487
488  /// Produces a strategy for generating an ordered set together with two values according to the following algorithm:
489  /// 1. Call `f` to get a pair of sets (x,y).
490  /// 2. Toss a coin to decide whether to pick an element from x at random, or from y (if the chosen set is empty
491  ///    Default is called). 3. Repeat step 2 and let the two outcomes be denoted a and b.
492  /// 4. Toss a coin to decide whether to swap the keys of a and b.
493  /// 5. return (x,a,b)
494  fn set_with_values<F, T, U>(f: F) -> impl Strategy<Value = (OrderedSet<T>, T, T)>
495  where
496    T: KeyComparable + Default + Debug + Clone + ReplaceKey,
497    <T as KeyComparable>::Key: Clone,
498    U: Strategy<Value = (OrderedSet<T>, OrderedSet<T>)>,
499    F: Fn() -> U,
500  {
501    f().prop_perturb(|(x, y), mut rng| {
502      let sets = [&x, &y];
503
504      let sample = |generator: &mut TestRng| {
505        let set_idx = usize::from(generator.random_bool(0.5));
506        let set_range = if set_idx == 0 { 0..x.len() } else { 0..y.len() };
507        if set_range.is_empty() {
508          T::default()
509        } else {
510          let entry_idx = generator.random_range(set_range);
511          (sets[set_idx])[entry_idx].clone()
512        }
513      };
514
515      let (mut a, mut b) = (sample(&mut rng), sample(&mut rng));
516      if rng.random_bool(0.5) {
517        let key_a = a.key().clone();
518        let key_b = b.key().clone();
519        a.set_key(key_b);
520        b.set_key(key_a);
521      }
522
523      (x, a, b)
524    })
525  }
526
527  fn set_with_values_comparable_struct(
528  ) -> impl Strategy<Value = (OrderedSet<ComparableStruct>, ComparableStruct, ComparableStruct)> {
529    set_with_values(arbitrary_sets_comparable_struct)
530  }
531
532  fn set_with_values_u128() -> impl Strategy<Value = (OrderedSet<u128>, u128, u128)> {
533    set_with_values(arbitrary_sets_u128)
534  }
535
536  fn assert_operation_preserves_invariant<F, T, S>(mut set: OrderedSet<T>, operation: F, val: T)
537  where
538    T: KeyComparable,
539    T::Key: Hash + Eq,
540    F: Fn(&mut OrderedSet<T>, T) -> S,
541  {
542    operation(&mut set, val);
543    assert_unique_keys(set);
544  }
545
546  fn assert_unique_keys<T>(set: OrderedSet<T>)
547  where
548    T: KeyComparable,
549    T::Key: Hash + Eq,
550  {
551    let mut keys: HashSet<&T::Key> = HashSet::new();
552    for key in set.as_slice().iter().map(KeyComparable::key) {
553      assert!(keys.insert(key));
554    }
555  }
556
557  proptest! {
558    #[test]
559    fn preserves_invariant_append_comparable_struct((set, element, _) in set_with_values_comparable_struct()) {
560      assert_operation_preserves_invariant(set,OrderedSet::append, element);
561    }
562  }
563
564  proptest! {
565    #[test]
566    fn preserves_invariant_append_u128((set, element, _) in set_with_values_u128()) {
567      assert_operation_preserves_invariant(set,OrderedSet::append, element);
568    }
569  }
570
571  proptest! {
572    #[test]
573    fn preserves_invariant_prepend_comparable_struct((set, element, _) in set_with_values_comparable_struct()) {
574      assert_operation_preserves_invariant(set,OrderedSet::prepend, element);
575    }
576  }
577
578  proptest! {
579    #[test]
580    fn preserves_invariant_prepend_u128((set, element, _) in set_with_values_u128()) {
581      assert_operation_preserves_invariant(set,OrderedSet::prepend, element);
582    }
583  }
584
585  proptest! {
586    #[test]
587    fn preserves_invariant_update_comparable_struct((set, element, _) in set_with_values_comparable_struct()) {
588      assert_operation_preserves_invariant(set,OrderedSet::update, element);
589    }
590  }
591
592  proptest! {
593    #[test]
594    fn preserves_invariant_update_u128((set, element, _) in set_with_values_u128()) {
595      assert_operation_preserves_invariant(set,OrderedSet::update, element);
596    }
597  }
598
599  proptest! {
600    #[test]
601    fn preserves_invariant_replace_comparable_struct((mut set, current, update) in set_with_values_comparable_struct()) {
602      set.replace(current.key(), update);
603      assert_unique_keys(set);
604    }
605  }
606
607  proptest! {
608    #[test]
609    fn preserves_invariant_replace_u128((mut set, current, update) in set_with_values_u128()) {
610      set.replace(current.key(), update);
611      assert_unique_keys(set);
612    }
613  }
614
615  proptest! {
616    #[test]
617    fn preserves_invariant_remove_u128((mut set, element, _) in set_with_values_u128()) {
618      set.remove(element.key());
619      assert_unique_keys(set);
620    }
621  }
622
623  proptest! {
624    #[test]
625    fn preserves_invariant_remove_comparable_struct((mut set, element, _) in set_with_values_comparable_struct()) {
626      set.remove(element.key());
627      assert_unique_keys(set);
628    }
629  }
630}