How to implement IntoIterator for an enum of iterable variants?

user_1171710 0
Viewed 1073

I have an enum that has different variant possibilities like so:

pub enum Enum {
  Empty,
  Single(Struct),
  Multi(Vec<Struct>),
}

And I want to create the typical three iteration functions to get Structs for Enum: one .iter() for an iterator of references, and two flavors of .into_iter() to get iterators over references and over values. How does one do this?

Here is my failed attempt:

use std::{iter, slice, vec};

#[derive(Debug)]
pub struct Struct(u32);

#[derive(Debug)]
pub enum Enum {
  Empty,
  Single(Struct),
  Multi(Vec<Struct>),
}

impl Enum {
  fn iter(&self) -> slice::Iter<'_, Struct> {
    match self {
      Enum::Empty => iter::empty(),
      Enum::Single(s) => iter::once(s),
      Enum::Multi(v) => v.iter(),
    }
  }
}

impl IntoIterator for Enum {
  type Item = Struct;
  type IntoIter = vec::IntoIter<Struct>;

  fn into_iter(self) -> Self::IntoIter {
    match self {
      Enum::Empty => iter::empty(),
      Enum::Single(s) => iter::once(s),
      Enum::Multi(v) => v.into_iter(),
    }
  }
}

impl<'a> IntoIterator for &'a Enum {
  type Item = &'a Struct;
  type IntoIter = slice::Iter<'a, Struct>;

  fn into_iter(self) -> Self::IntoIter {
    self.iter()
  }
}

fn main() {
  let enums = vec![
    Enum::Empty,
    Enum::Single(Struct(1)),
    Enum::Multi(vec![Struct(2), Struct(3)])];
  for e in enums {
    for s in e.iter() {
      println!(".iter() over refs: {:?}", s);
    }
    for s in &e {  // leverages IntoIterator for &'a Enum
      println!(".into_iter() over refs: {:?}", s);
    }
    for s in e {  // leverages IntoIterator for Enum
      println!(".into_iter() moves: {:?}", s);
    }
  }
}

Playground

I'm failing to see how to coerce the iter::empty and iter::once into the slice::Iter and vec::IntoIter, or is that even the right thing to want to do?

rust enums iterator
1 Answers

In this case, all these variants can be represented by slices or easily converted into Vecs. So you can do that and then use their iterators directly.

impl IntoIterator for Enum {
    type Item = Struct;
    type IntoIter = vec::IntoIter<Struct>;
  
    fn into_iter(self) -> vec::IntoIter<Struct> {
        let vec = match self {
            Enum::Empty => Vec::new(),
            Enum::Single(single) => vec![single],
            Enum::Multi(vec) => vec,
        };
        
        vec.into_iter()
    }
}

impl<'a> IntoIterator for &'a Enum {
  type Item = &'a Struct;
  type IntoIter = slice::Iter<'a, Struct>;

    fn into_iter(self) -> slice::Iter<'a, Struct> {
        let slice = match self {
            Enum::Empty => &[],
            Enum::Single(single) => std::slice::from_ref(single),
            Enum::Multi(vec) => vec.as_slice(),
        };
        
        slice.iter()
    }
}

See it on the playground.

If you actually have different enough variants that you can't coerce them into an existing iterator, you can implement the iterator as an enum of other iterators. Here's what that would look like:

pub enum EnumIntoIter {
    Empty(std::iter::Empty<Struct>),
    Single(std::iter::Once<Struct>),
    Multi(std::vec::IntoIter<Struct>),
}

impl Iterator for EnumIntoIter {
    type Item = Struct;
    
    fn next(&mut self) -> Option<Struct> {
        match self {
            EnumIntoIter::Empty(iter) => iter.next(),
            EnumIntoIter::Single(iter) => iter.next(),
            EnumIntoIter::Multi(iter) => iter.next(),
        }
    }
}

impl IntoIterator for Enum {
    type Item = Struct;
    type IntoIter = EnumIntoIter;
  
    fn into_iter(self) -> EnumIntoIter {
        match self {
            Enum::Empty => EnumIntoIter::Empty(std::iter::empty()),
            Enum::Single(single) => EnumIntoIter::Single(std::iter::once(single)),
            Enum::Multi(vec) => EnumIntoIter::Multi(vec.into_iter()),
        }
    }
}

pub enum EnumIter<'a> {
    Empty(std::iter::Empty<&'a Struct>),
    Single(std::iter::Once<&'a Struct>),
    Multi(std::slice::Iter<'a, Struct>),
}

impl<'a> Iterator for EnumIter<'a> {
    type Item = &'a Struct;
    
    fn next(&mut self) -> Option<&'a Struct> {
        match self {
            EnumIter::Empty(iter) => iter.next(),
            EnumIter::Single(iter) => iter.next(),
            EnumIter::Multi(iter) => iter.next(),
        }
    }
}

impl<'a> IntoIterator for &'a Enum {
    type Item = &'a Struct;
    type IntoIter = EnumIter<'a>;
  
    fn into_iter(self) -> EnumIter<'a> {
        match self {
            Enum::Empty => EnumIter::Empty(std::iter::empty()),
            Enum::Single(single) => EnumIter::Single(std::iter::once(single)),
            Enum::Multi(vec) => EnumIter::Multi(vec.iter()),
        }
    }
}

See it on the playground

user_21891300
Related