Behavior of `map` when there is an `nth` available

help
1

It seems like the map adapter doesn't take advantage of the nth implementation provided by the underlying iterator. This makes calling nth on the results of a map potentially inefficient.

I'd expect this code to never invoke my next implementation, since I was expecting the map adapter to use my nth implementation instead of the mixin from Iterator. But the actual behavior is different.

struct Custom { count: usize, max: usize }

impl Iterator for Custom {
    type Item = usize;
    
    fn next(&mut self) -> Option<usize> {
        println!("Call to next!");
        if self.count >= self.max { return None };
        let to_return = self.count;
        self.count += 1;
        return Some(to_return);
    }
    
    fn nth(&mut self, val: usize) -> Option<usize> {
        self.count += val;
        if self.count >= self.max { return None; }
        return Some(self.count);
    }
}

fn main() {
    let mut c = Custom { count: 0, max: 10000 };
    println!("5th value: {}", c.nth(5).unwrap());
    
    c = Custom { count: 0, max: 10000 };
    println!("5th mapped value: {}", c.map(|x| x * 2).nth(5).unwrap());
}

(Playground)

Output:

5th value: 5
Call to next!
Call to next!
Call to next!
Call to next!
Call to next!
Call to next!
5th mapped value: 10

Errors:

   Compiling playground v0.0.1 (/playground)
    Finished dev [unoptimized + debuginfo] target(s) in 0.78s
     Running `target/debug/playground`

After asking on the Rust Discord, I now understand that this is to ensure the side effects of the closure given to map are applied to every item the iterator passes over.

Is there some external crate or way to use map that will take advantage of more efficient implementations of things like nth?

2

When you know your iterator is random-access and you want to skip the first N elements, you could just exchange mapping and skipping, like this: c.skip(5).next().map(|x| x * 2).

3

Ah, thanks, that's a nice workaround!

Outside of my toy example above, I often want to return a mapped iterator, and sometimes I want to do the nth call in the thing consuming the mapped iterator. This won't work in that case, but will solve my immediate problem.

I think it would still be nice for the nth call to a map adapter to use the provided nth implementation, or perhaps something like pure_map that assumes no side effects...

4

If you have full control over the full Iterator chain, you can also "just" do the nth before the map.

5

I would consider it a mistake for any iterator to implement nth in a way that it has different behavior than the default implementation in Iterator. slice::Iter has an nth method that can skip ahead, but doing so doesn't have any side effects, so that's all good. Map is a whole different story.

Suppose I wrote a function like this...

/// Returns an iterator of differences between consecutive items.
fn deltas(arg: &[i32]) -> impl '_ + Iterator<Item = i32> {
    let mut last = arg[0];
    arg[1..].iter().map(move |&a| {
        let ret = a - last;
        last = a;
        ret
    })
}

Handy, right?

    let mut i = deltas(&[0, 10, 12, 20, 25, 30]);
    assert_eq!(i.next(), Some(10));
    assert_eq!(i.next(), Some(2));

Clearly if you call .nth on it you should get the difference between elements n and n + 1, and that's what the default implementation in Iterator does. But if Map worked the way you think it should, it would give you the difference between element 0 and element n + 1. Worse, this "skipping" behavior would seem to mysteriously disappear if you put something like .flat_map(Some) in between deltas and nth. I think this would just be a huge footgun, making certain iterators work differently than others in the name of performance.

I wonder if your background is in C++? C++ has the concept of a "random access iterator" (which is, frankly, an oxymoron) that is used instead of a generic collections API. Rust's iterators are more focused around, well, iteration rather than whatever-you-might-do-with-any-generic-collection. Unfortunately Rust also doesn't have a proper generic collections API (there's some work to do on GATs before it can be supported in the most useful and general way), so you might find iterators a little limiting in Rust compared to what C++ offers.

2 Likes
6

This is a great example of taking advantage of a side effect of the function you are passing to map, as I mentioned in my original post.

My background is more on the functional side, which is why I was surprised that an operation called "map" could have side effects at all. Of course, your example could also be implemented in a functional style, by first grouping together overlapping pairs of items.

I agree that maintaining such side effects is important to avoid a "foot gun" to users who do not think of map in a functional paradigm. Since map can take a FnMut`, I totally agree that the current behavior is the only sane choice.

My problem is that I have a function producing an iterator (that is itself a map over another iterator), and then I consume that iterator in a number of different places. In some of the those places, I want to skip some items without calling my (pure) closure for performance. Perhaps a pure_map could be implemented in the standard library or an external crate?

I think it might also be the case that the Rust type system already has enough information to know if a more efficient nth could be used: whether or not the map closure is a Fn or FnMut. It would be nice if, in cases where the closure canb e proven side-effect-free, like [1, 2, 3, 4, 5].iter().map(|x| x * 2).nth(3), only required a single multiply-by-two.

7

I've never used a map that couldn't have side effects; Python, Perl and Lisp (at least the ones I'm aware of) all allow them.

The closest type level analog to "pure functions" in Rust is const fn; all non-const functions may have side effects, even Fn closures. But I don't think there's any way to write a pure_map that only accepts const fns.

I also prefer to rely on type level guarantees over optimizations, but this is the kind of thing that I would expect the optimizer to do pretty well with most of the time. Have you profiled / looked at the assembly to see whether it is actually calling the closure a bunch of times?

1 Like
8

I would assume that pure_map wouldn't check purity. That is impossible in an impure turing complete language. Instead, it would be a hint that the mapping operation is pure, so it is ok to skip, and this would allow the Iterator::nth optimization.

3 Likes
9

@trentj -- no need to profile the assembly, I believe my example from the original post demonstrates that such optimizations are both not present, and, as you pointed out, not desirable. :slight_smile: Thanks for pointing out that Fn and FnMut weren't exactly what I thought they were, I missed that.

@RustyYato -- yes, this makes sense. I think the behavior of pure_map might even be desirable in some cases when the closure even does have side effects. At least you could construct them, not sure if they actually exist.

10

I have used a language (VHDL) that distinguishes between pure and impure functions. Of course, you can't write every pure function, but much like const fn there are some things the compiler just doesn't allow in pure functions and you deal with it. So I don't think it's a compelling argument to say that you can't check purity.

Not at all. The next in the original post has a side effect (calling println!), so the compiler can't eliminate it. If you remove the println! so that next is actually side-effect-less, it's all removed.

4 Likes
11

Ok, I should ammend my statement to: you can'tcheck purity in an impure turing complete language without annotations. Checking for purity is non-trivial without annotations. Given that Rust used to have an annotation for purity, and got rid of it, I don't think we're going to get annotations. But that's beside the point.

I would find a pure_map combinator nice, but it should be in itertools. I think it will cause more confusion than it is worth.

2 Likes
12

Maybe call it something like lazy_map, since purity isn't really required.

1 Like
13

Good point! Maybe there's some trick, but I can't seem to get the optimizer to get rid of the loop even for a pretty simple example: Compiler Explorer

Maybe I have the wrong flags?

1 Like
14

It can possibly require that the mapper function is Fn and not just FnMut - this at least rules out the explicit environment mutation.

15

Requiring Fn instead of FnMut can be a hint to users, but we can't optimize the code based on it since we can still call the RefCell::borrow_mut() within the Fn closures.

1 Like
16

No, that is a needlessly restrictive api. I think I like @kornel's lazy_map, since purity isn't the important bit, it's that the mapping operation is only evaluated as needed (which sounds like laxy evaluation).

I don't want a needlessly restrictive api, because it will likely run into issues later on. For example: Vec::retain could give you an exclusive reference in the callback parameter, but it gives you a shared reference. This means that you can't easily mutate things inside the Vec as you remove them, which is a useful thing to do without significant regressions (via shared mutability, or looping twice). All because the api was needlessly restrictive. (Yes, drain_filter will fix this specific problem, but I'm illustrating a point)

1 Like
17

Hmm, it seems you are right. Oddly, although it doesn't get rid of the loop entirely, it does lift the indexing and x * 2 out of it:

        mov     rax, -5001
.LBB2_4:
        add     rax, 3                        // keep adding 3 until rax=0
        jne     .LBB2_4
        mov     eax, dword ptr [rbx + 20000]  // load v[5000]
        add     eax, eax                      // multiply it by 2

The last two lines are not inside the loop. Which is... better than no optimization at all, but worse than I'd like. Curiouser still, if you change 5000 to a different number you get loops that are unrolled by different amounts, or even skipped entirely.

I'm guessing there is some flag being set by add that LLVM thinks might be getting used later.

18

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.