Seeking alternative to `array::IntoIter::new_unchecked` on stable rust

My function needs to return a variable number of elements, depending on the branch executed internally. It's also important that the calling side doesn't require the result in the form of an array, so it can be an iterator.

For example:

fn foo(cond: bool) -> ... {
    if cond {
        [1, 2]
    } else {
        [1, 2, 3, 4]
    }
}

The maximum number of elements is small and strictly limited; in the example above, it's 4. Thus, we want to use stack memory.

However, since in Rust we need to have a unified output type, I had to consider several possible approaches.

Options that came to mind:

  1. Use arrayvec or smallvec. But adding a dependency seems excessive for a single use case in the code.

  2. Use an enum to enumerate all possible lengths and implement Iterator. However, each call to .next() would require checking which variant it belongs to. These are extra runtime checks that we'd like to avoid, although they might be optimized by the compiler.

  3. Construct array::IntoIter<T, N>::new_unchecked. But this is only available in nightly Rust.

  4. Create a custom wrapper ArrayIter.

The idea behind the ArrayIter wrapper is simple:
Transform [T; 2] into std::array::IntoIter<T, 4>, but with a shifted internal index start: {start: 0, end: 4} -> {start: 0, end: 2} (alternatively, we could shift the start index: start: 0 -> start: 2).
The function output would be std::array::IntoIter<T, 4> with its internal state correctly pre-modified.

I'd appreciate your insights on the following aspects:

  1. Overall approach
  2. Safety of the unsafe code blocks

I've commented out some alternative unsafe implementations. Would these be preferable, and if so, why?

gist: Code shared from the Rust Playground · GitHub
playground: Rust Playground

Thank you in advance for your time and expertise.

1 Like

I think this is the way to go. If you don't want a dependency for some reason, reimplement the data structure yourself.

2 Likes

Are there any safety issues or potential undefined behavior in my current implementation?
That's the main question I have at the moment.

Why so you need a custom type?

let mut iter= [1, 2, 0, 0].into_iter();
iter.nth_back(1);
iter

This should be nicely optimizable.

4 Likes

Well, it passes Miri. I don't see any issues in current code (haven't checked alternatives in comments).

I'd still suggest to use something simpler. Would T: Default be sufficient for your use case?

EDIT: By the way, you don't need to use unsafe in Drop. You can just iterate and drop as usual. This way your only unsafe block is in Iterator::next, and it's seems easy to check that all elements in the iterator were initialized.

1 Like

Of what type?

It's better to avoid placeholder values and use Option types instead of null values. But here I use MaybeUninit instead of Option, because it seems that I can enforce memory invariants. This is the central issue that can be found in the code I provided.

I deliberately didn't specify this requirement in my original question, assuming that others shared my view on null values.

I used integers ([1, 2], [1, 2, 3, 4]) as a simple example to illustrate the concept, but the type could be anything accepted by std::array, including non-Copy or resource-heavy types. This is why I left the trait bounds unconstrained.

My focus is on identifying any significant misunderstandings I might have about working with uninitialized memory.

1 Like

Minor improvement: If you use a const block for the assertion in new, you get a compiler error instead of a runtime error when N is too small:

        const {
            assert!(LEN <= N);
        }
3 Likes

I know you don't want this, but just in case someone reads this thread and would rather add an arrayvec dependency than add new unsafe code, here's how it might work with arrayvec.

use arrayvec::ArrayVec;
use core::iter::IntoIterator;

// use a type alias to make this less verbose
pub type ArrayIter<T, const CAP: usize> = <ArrayVec<T, CAP> as IntoIterator>::IntoIter;

// use a fn to add a compile time check for LEN <= CAP
pub fn array_iter<T, const LEN: usize, const CAP: usize>(a: [T; LEN]) -> ArrayIter<T, CAP> {
    const {
        assert!(LEN <= CAP);
    }
    ArrayVec::<T, CAP>::from_iter(a).into_iter()
}

fn main() {
    fn check(cond: bool) -> impl Iterator<Item = i32> {
        let iter: ArrayIter<_, 4> = if cond {
            array_iter([1, 2])
        } else {
            array_iter([1, 2, 3, 4])
        };
        iter
    }
    assert!(check(true).eq([1, 2].into_iter()));
    assert!(check(false).eq([1, 2, 3, 4].into_iter()));
}

playground