Skip to content

Instantly share code, notes, and snippets.

@jonhoo
Created May 27, 2020 22:07
Show Gist options
  • Save jonhoo/dd63b720fa4a220ea845a77e2d75831e to your computer and use it in GitHub Desktop.
Save jonhoo/dd63b720fa4a220ea845a77e2d75831e to your computer and use it in GitHub Desktop.
pub trait IteratorExt: Iterator + Sized {
fn our_flatten(self) -> Flatten<Self>
where
Self::Item: IntoIterator;
}
impl<T> IteratorExt for T
where
T: Iterator,
{
fn our_flatten(self) -> Flatten<Self>
where
Self::Item: IntoIterator,
{
flatten(self)
}
}
pub fn flatten<I>(iter: I) -> Flatten<I::IntoIter>
where
I: IntoIterator,
I::Item: IntoIterator,
{
Flatten::new(iter.into_iter())
}
pub struct Flatten<O>
where
O: Iterator,
O::Item: IntoIterator,
{
outer: O,
front_iter: Option<<O::Item as IntoIterator>::IntoIter>,
back_iter: Option<<O::Item as IntoIterator>::IntoIter>,
}
impl<O> Flatten<O>
where
O: Iterator,
O::Item: IntoIterator,
{
fn new(iter: O) -> Self {
Flatten {
outer: iter,
front_iter: None,
back_iter: None,
}
}
}
impl<O> Iterator for Flatten<O>
where
O: Iterator,
O::Item: IntoIterator,
{
type Item = <O::Item as IntoIterator>::Item;
fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(front_iter) = &mut self.front_iter {
if let Some(i) = front_iter.next() {
return Some(i);
}
self.front_iter = None;
}
if let Some(next_inner) = self.outer.next() {
self.front_iter = Some(next_inner.into_iter());
} else {
return self.back_iter.as_mut()?.next();
}
}
}
}
impl<O> DoubleEndedIterator for Flatten<O>
where
O: DoubleEndedIterator,
O::Item: IntoIterator,
<O::Item as IntoIterator>::IntoIter: DoubleEndedIterator,
{
fn next_back(&mut self) -> Option<Self::Item> {
loop {
if let Some(back_iter) = &mut self.back_iter {
if let Some(i) = back_iter.next_back() {
return Some(i);
}
self.back_iter = None;
}
if let Some(next_back_inner) = self.outer.next_back() {
self.back_iter = Some(next_back_inner.into_iter());
} else {
return self.front_iter.as_mut()?.next_back();
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty() {
assert_eq!(flatten(std::iter::empty::<Vec<()>>()).count(), 0);
}
#[test]
fn empty_wide() {
assert_eq!(flatten(vec![Vec::<()>::new(), vec![], vec![]]).count(), 0);
}
#[test]
fn one() {
assert_eq!(flatten(std::iter::once(vec!["a"])).count(), 1);
}
#[test]
fn two() {
assert_eq!(flatten(std::iter::once(vec!["a", "b"])).count(), 2);
}
#[test]
fn two_wide() {
assert_eq!(flatten(vec![vec!["a"], vec!["b"]]).count(), 2);
}
#[test]
fn reverse() {
assert_eq!(
flatten(std::iter::once(vec!["a", "b"]))
.rev()
.collect::<Vec<_>>(),
vec!["b", "a"]
);
}
#[test]
fn reverse_wide() {
assert_eq!(
flatten(vec![vec!["a"], vec!["b"]])
.rev()
.collect::<Vec<_>>(),
vec!["b", "a"]
);
}
#[test]
fn both_ends() {
let mut iter = flatten(vec![vec!["a1", "a2", "a3"], vec!["b1", "b2", "b3"]]);
assert_eq!(iter.next(), Some("a1"));
assert_eq!(iter.next_back(), Some("b3"));
assert_eq!(iter.next(), Some("a2"));
assert_eq!(iter.next_back(), Some("b2"));
assert_eq!(iter.next(), Some("a3"));
assert_eq!(iter.next_back(), Some("b1"));
assert_eq!(iter.next(), None);
assert_eq!(iter.next_back(), None);
}
#[test]
fn inf() {
let mut iter = flatten((0..).map(|i| 0..i));
// 0 => 0..0 => empty
// 1 => 0..1 => [0]
// 2 => 0..2 => [0, 1]
assert_eq!(iter.next(), Some(0));
assert_eq!(iter.next(), Some(0));
assert_eq!(iter.next(), Some(1));
}
#[test]
fn deep() {
assert_eq!(flatten(flatten(vec![vec![vec![0, 1]]])).count(), 2);
}
#[test]
fn ext() {
assert_eq!(vec![vec![0, 1]].into_iter().our_flatten().count(), 2);
}
}
@jonhoo
Copy link
Author

jonhoo commented Mar 10, 2024

self is mutable, and that's what we're accessing it via — there's no way to declare individual fields as mutable (or not mutable) in Rust

@1717chen
Copy link

thank you so much!

another question if I may,

when we do
if let Some(next_back_inner) = self.outer.next_back() {

does it involve transfer the ownership to next_back_inner variable?

thanks

@jonhoo
Copy link
Author

jonhoo commented Mar 30, 2024

That depends on the method signature of next_back. If it returns Some(&_), then we're transferring ownership of the reference (i.e., borrowing).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment