There are times where I need to iterate through some items and the first element in iterator has some special handling.

A simple example is joining a slice of str. One wants to go from:

["a", "b", "c"]

and get:

"a + b + c"

I understand that for this specific example one could use the join() method on slice:

let joined = ["a", "b", "c"].join(" + ");

The intent isn’t to implement a better join, but to focus on the special handling around the first element of an iterator. It’s also possible that someone is in a no alloc environment, in these instances they wouldn’t have join() available to them as it returns a String, which requires alloc.

I’ll walk through 4 ways this could be implemented (I’m sure there are many more):

1. Using a boolean flag to know if we’ve seen the first element
2. Keying off of the index
3. Using a zip iterator
4. Using once() with default alternative

Boolean flag

Using the boolean flag method we’ll initialize a boolean flag to one state. We’ll condition on the state of the flag to determine if we need to perform extra logic and then we’ll set the flag to the alternate state.

pub fn boolean_flag_join(slice: &[&str]) -> String {
    let mut add_separator = false;
    let mut output = String::new();
    for item in slice {
        if add_separator {
            output.push_str(" + ");
        }
        add_separator = true;
        output.push_str(item);
    }
    output
}

For the above example I chose a positive variable name, so even though the odd state is for the initial item in the iterator of str I named it based on what is being done in the logic, adding the separator between successive elements.

This is a common approach and is often the way one does this type of thing in C.

Keying off of the index

Keying off the index uses the index in the iterator to determine when to do the special logic.

pub fn enumerator_index_join(slice: &[&str]) -> String {
    let mut output = String::new();
    for (i, item) in slice.iter().enumerate() {
        if i != 0 {
            output.push_str(" + ");
        }
        output.push_str(item);
    }
    output
}

While the code example is provided a slice and could actually navigate the slice by index. It’s more idiomatic to treat the slice as an iterator. This then requires using enumerate() to get to the index.

Zip Iterator

To use a zip iterator to solve this, an iterator representing the separator strings is zipped up with the slice items, always pushing a separator onto the String.

pub fn zip_join(slice: &[&str]) -> String {
    let separators = core::iter::once("").chain(core::iter::repeat(" + "));
    let mut output = String::new();
    separators.zip(slice).for_each(|(separator, item)| {
        output.push_str(separator);
        output.push_str(item);
    });
    output
}

For this problem the zipped iterator starts with a once() iterator which will provide the specified item once and then be exhausted. The initial item is an empty string so that the push for the first element results in no change to the output. The once() iterator is chained with a repeat() iterator of the real separator. Using a repeat() ensures we have a separator for however many items may be in the slice.

This approach removes the condtion in the for loop, but it can take a bit of a mental shift if one is used to the if conditions that’s needed in C for loops.

once() with default alternative

Similar to the zip iterator, an iterator could be created from once(). However instead of chaining this to a repeat() iterator we exhaust the iterator and fall back to an alternate value.

pub fn once_join(slice: &[&str]) -> String {
    let mut separator_iter = core::iter::once("");
    let mut output = String::new();
    for item in slice {
        output.push_str(separator_iter.next().unwrap_or_else(|| " + "));
        output.push_str(item);
    }
    output
}

This example creates an iterator of one itme being an empty string. The iterator will be exausted after the first time it’s called and then the unwrap_or_else will provide the value we’ll use in all other instances.

Similar to the zip iterator this may seem a bit foreign to those coming from C or similar where a condition in a for loop is often used. This can also feel a bit wasteful as it’s creating more or less a throw away iterator just to avoid a condition in the code.

Comparison

For this simple example one shouldn’t worry too much about performance and should lean toward readability. At the same time I think it’s good to know or understand how the solutions may differ.

I put all of the approaches in one main file using opt-level 3 on compiler explorer:

Looking at the assembly, the boolean_flag_join() and the enumerator_index_join() were compiled into the same function. The zip_join() and the once_join() remain separate functions. The assembly for once_join() and boolean_flag_join() are very similar. The zip_join() deviates quite a bit and looks to have more assembly instructions.

I decided to use criterion to see how these compare performance wise:

once_join               time:   [34.016 ns 34.068 ns 34.134 ns]

zip_join                time:   [45.314 ns 45.369 ns 45.442 ns]

enumerator_join         time:   [34.014 ns 34.152 ns 34.309 ns]

boolean_flag_join       time:   [33.965 ns 34.047 ns 34.154 ns]

As the extra assembly implied, zip_join() is slower. These benchmarks were using the example slice of ["a", "b", "c"]. To get a better idea of how much the zip_join() deviates I ran the benches again with ["a", "b", "c", "d", "e", "f", "g", "h"]:

once_join               time:   [73.588 ns 73.787 ns 73.992 ns]

zip_join                time:   [102.24 ns 102.55 ns 102.87 ns]

enumerator_join         time:   [74.523 ns 75.011 ns 75.541 ns]

boolean_flag_join       time:   [77.397 ns 77.928 ns 78.417 ns]

Since the zip_join() delta increased it seems to be a per iteration cost. As opposed to a one time setup cost.

Summary

For some reason seeing if statements inside of for loops feels like a code smell. I understand at times it may be the only way to do things, especially for more complex cases than the example. It may be that having a condition makes one pause and think:

Under what circumstances does this for loop behave differently?

Coming from C I’ve had a bit harder time becoming fluent in iterators and what they’re doing at times, but I’ve been trying to lean on them more.

I was a bit surprised and disappointed that the zip_join() was ~33% slower in these simple tests. I wouldn’t avoid it as a solution unless there is a performance concern.