[Solved] Possible to access current iterator state from within map?

I have a tree traversal iterator that maintains the depth of most recently tree entry yielded by the iterator. I also sometimes want to map that maintained depth state.

For example here is how I can do it imperatively:

let mut iterator = tree.iterator();
let depths = Vec::new();
for _ in iterator {
  depths.push(iterator.depth());
}

But what I can't figure out is how to map this state functionally. To do that I think that I need access to the current iterator from within the map function; something like:

let mut iterator = tree.iterator();
let depths = iterator.map(|_| {
  iterator.depth()
}).collect();

But of course that's problematic with the borrowing rules. Is there some way to do this? Or if not a good workaround pattern?

What I have been doing is creating new wrapping iterator for cases like this. In this case I would wrap my tree iterator in a new TreeDepthIterator that would then drive the tree iterator and yield the depth values. This works but I would rather not have to create these wrapper types.

I think the normal solution here is to be imperative.

Related: Rust: Looping on a member variable without mutably borrowing self | Stephan Sokolow's Blog

Ok thanks.

An alternative would be for the iterator to yield tuples of (depth, item), and you consume those similarly to how you would downstream of an enumerate().

1 Like

This could be made into a nice iterator adapter if it is used often I think

That faces the same challenge as the original question: access to the internal state of the iterator downstream in the adaptor. Enumerate works because counting things is easy within the internal state of the adaptor, regardless of source.

Ah, I think you are right. I haven't tried yet, but I think I can make a generic adapter that works like this:

let mut iterator = tree.iterator();
let depths = iterator.map_with_self(|tree_iter, _| {
  tree_iter.depth()
}).collect();

I'll report back.

I created a wrapper that passes in the original iterator while performing a map. It seems to be working well for me. It's based on a rather ill informed copy/paste/modify of the std::itr::Map implementation. Seems to be working well for me:

use std::fmt;
use std::iter::FusedIterator;

#[derive(Clone)]
pub struct MapWithSelf<I, F> {
  iter: I,
  f: F,
}

impl<I: fmt::Debug, F> fmt::Debug for MapWithSelf<I, F> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    f.debug_struct("MapWithSelf")
      .field("iter", &self.iter)
      .finish()
  }
}

impl<B, I: Iterator, F> Iterator for MapWithSelf<I, F>
where
  F: FnMut(&I, I::Item) -> B,
{
  type Item = B;

  #[inline]
  fn next(&mut self) -> Option<B> {
    self.iter.next().map(|n| (self.f)(&self.iter, n))
  }

  #[inline]
  fn size_hint(&self) -> (usize, Option<usize>) {
    self.iter.size_hint()
  }
}

impl<B, I: DoubleEndedIterator, F> DoubleEndedIterator for MapWithSelf<I, F>
where
  F: FnMut(&I, I::Item) -> B,
{
  #[inline]
  fn next_back(&mut self) -> Option<B> {
    self.iter.next_back().map(|n| (self.f)(&self.iter, n))
  }
}

impl<B, I: ExactSizeIterator, F> ExactSizeIterator for MapWithSelf<I, F>
where
  F: FnMut(&I, I::Item) -> B,
{
  fn len(&self) -> usize {
    self.iter.len()
  }
}

impl<B, I: FusedIterator, F> FusedIterator for MapWithSelf<I, F> where F: FnMut(&I, I::Item) -> B {}

pub trait MapWithSelfExt: Iterator {
  fn map_with_self<F, B>(self, f: F) -> MapWithSelf<Self, F>
  where
    Self: Sized,
    F: FnMut(&Self, Self::Item) -> B,
  {
    MapWithSelf { iter: self, f: f }
  }
}

impl<I: Iterator> MapWithSelfExt for I {}

(Playground)

2 Likes