Implement FromIterator for OneOrMany

identity-core/src/common/one_or_many.rs

@@ -74,7 +74,11 @@ impl<T> OneOrMany<T> {
         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<T> From<OneOrMany<T>> for Vec<T> {
   }
 }
 
+impl<T> FromIterator<T> for OneOrMany<T> {
+  fn from_iter<I: IntoIterator<Item = T>>(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::<Vec<T>>().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::<u32>::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::<u32>::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::<u32>::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<i32> = (0..42).collect();
+    let mut collection = OneOrMany::Many(v);
+    collection.push(42);
+    assert_eq!(collection, OneOrMany::Many((0..=42).collect()));
+  }
+}