diff --git a/identity-core/src/common/one_or_many.rs b/identity-core/src/common/one_or_many.rs index c2c0b2c1e8..97619c53f6 100644 --- a/identity-core/src/common/one_or_many.rs +++ b/identity-core/src/common/one_or_many.rs @@ -74,7 +74,11 @@ impl OneOrMany { Self::Many(_) => unreachable!(), }, Self::Many(ref mut inner) => { - inner.push(value); + if inner.is_empty() { + *self = Self::One(value); + } else { + inner.push(value); + } } } } @@ -168,6 +172,24 @@ impl From> for Vec { } } +impl FromIterator for OneOrMany { + fn from_iter>(iter: I) -> Self { + let mut iter = iter.into_iter(); + // if the iterator contains one element or less and a correct size hint is provided we can save an allocation + let size_hint = iter.size_hint(); + if size_hint.1.is_some() && (1, Some(1)) >= size_hint { + let mut this = iter.next().map(Self::One).unwrap_or_else(|| Self::Many(Vec::new())); + // if the hinted upper bound was incorrect we need to correct for it + for next in iter.by_ref() { + this.push(next); + } + this + } else { + iter.into_iter().collect::>().into() + } + } +} + // ============================================================================= // Iterator // ============================================================================= @@ -191,3 +213,60 @@ impl<'a, T> Iterator for OneOrManyIter<'a, T> { self.inner.get(self.index - 1) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn from_iterator_empty() { + let empty_vec = Vec::::new(); + assert_eq!(OneOrMany::from_iter(empty_vec.clone()), OneOrMany::Many(empty_vec)); + } + + #[test] + fn from_iterator_single() { + let single_item = [1]; + assert_eq!(OneOrMany::from_iter(single_item), OneOrMany::One(1)); + } + + #[test] + fn from_iterator_many() { + let letters = ["a", "b", "c", "d"]; + assert_eq!(OneOrMany::from_iter(letters), OneOrMany::Many(vec!["a", "b", "c", "d"])); + } + + #[test] + fn from_iterator_iter() { + let none = OneOrMany::Many(Vec::::new()); + assert_eq!(OneOrMany::from_iter(none.iter()), OneOrMany::Many(Vec::<&u32>::new())); + + let one = OneOrMany::One(42); + assert_eq!(OneOrMany::from_iter(one.iter()), OneOrMany::One(&42)); + + let two = OneOrMany::Many(vec![0, 1]); + assert_eq!(OneOrMany::from_iter(two.iter()), OneOrMany::Many(vec![&0, &1])); + } + + #[test] + fn push_from_zero_elements() { + let mut collection = OneOrMany::Many(Vec::::new()); + collection.push(42); + assert_eq!(collection, OneOrMany::One(42)); + } + + #[test] + fn push_one_element() { + let mut collection = OneOrMany::One(42); + collection.push(42); + assert_eq!(collection, OneOrMany::Many(vec![42, 42])); + } + + #[test] + fn push_many_elements() { + let v: Vec = (0..42).collect(); + let mut collection = OneOrMany::Many(v); + collection.push(42); + assert_eq!(collection, OneOrMany::Many((0..=42).collect())); + } +}