diff --git a/src/lib.rs b/src/lib.rs index e54c440..8b7c3ba 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -63,6 +63,12 @@ impl OneOrMany { fn is_empty(&self) -> bool { matches!(self, OneOrMany::Vec(v) if v.is_empty()) } + + /// Iterate over the values regardless of whether it contains one or many. + #[inline] + pub fn iter(&self) -> OneOrManyIter { + OneOrManyIter::new(self) + } } impl Default for OneOrMany { @@ -72,6 +78,37 @@ impl Default for OneOrMany { } } +pub struct OneOrManyIter<'a, T> { + one: Option<&'a T>, + many: Option>, +} + +impl<'a, T> OneOrManyIter<'a, T> { + fn new(inner: &'a OneOrMany) -> Self { + match inner { + OneOrMany::One(v) => Self { + one: Some(v), + many: None, + }, + OneOrMany::Vec(v) => Self { + one: None, + many: Some(v.iter()), + }, + } + } +} + +impl<'a, T> Iterator for OneOrManyIter<'a, T> { + type Item = &'a T; + + fn next(&mut self) -> Option { + if let Some(one) = self.one.take() { + return Some(one); + } + self.many.as_mut().and_then(|iter| iter.next()) + } +} + /// JWT Claims. #[serde_as] #[skip_serializing_none] @@ -502,6 +539,21 @@ mod tests { use super::*; + #[test] + fn one_or_many_iterator() { + let v = OneOrMany::Vec(vec![1, 2, 3]); + let mut iter = v.iter(); + assert_eq!(iter.next(), Some(&1)); + assert_eq!(iter.next(), Some(&2)); + assert_eq!(iter.next(), Some(&3)); + assert_eq!(iter.next(), None); + + let v = OneOrMany::One(1); + let mut iter = v.iter(); + assert_eq!(iter.next(), Some(&1)); + assert_eq!(iter.next(), None); + } + #[test] fn signing_and_verification() -> Result<()> { let mut claims = HeaderAndClaims::new_dynamic();