The implementation of InPlaceIterable for Flatten&FlatMap is unsound

[steffahn]

const S: String = String::new();
fn main() {
    let v = vec![[S, "Hello World!".into()], [S, S]];
    let mut i = v.into_iter().flatten();
    let _ = i.next();
    let result: Vec<String> = i.clone().collect();
    println!("{result:?}");
}

(playground)

["\0\0\0\0\0\0\0\0rld!", "\0\0\0\0\0\0\0\0rld!", ""]
free(): invalid pointer
[1]    1131073 IOT instruction (core dumped)

The above code is analogous to #85322 but applied to Flatten instead of Peekable: cloning the whole iterator doesn't preserve capacity in the inner vec::IntoIter. (This also applies to FlatMap.)

Introduced in 1.76

cc @the8472, #110353

@rustbot label T-libs, I-unsound, A-iterators

[the8472]

Oof... looks like I didn't learn from #85322 :(
I can prepare a backportable PR that disables the specialization for flatten/flatmap without ripping out all the changes and then do more cleanup afterwards if I can't figure out a way to recover this... it'd likely require cloning the full capacity or extending the clone interface somehow...

Edit: PR at #135104

[steffahn]

Right… either change something about the clone; or add some way to detect the after-clone state.

---

I'm thinking the latter might be quite approachable. If there's a method added to … InPlaceIterable I suppose … that asks “do you have extra elements outside of the underlying vec::IntoIter?” The most common case of wanting the in-place behavior if the consumption of the iterator starts on a freshly constructed stack of iterators would be fully covered.

So that's one method like

/// Conservaticely returns true whenever the iterator *may* contain extra items on the side
/// E.g. if the peeked element is `Some` in `Peekable`,
/// or either `frontiter` or `backiter` is `Some` for `Flatten`/`FlatMap`
fn has_extra_items(&self) -> bool

E.g. as an easy first extra step, there can be a const (or specialization marker?¹) that says

/// For optimization, `true`, if `has_extra_items` *never* returns `true`
const NEVER_HAS_EXTRA_ITEMS: bool;

---

One can potentially get successively even more precise later… e.g. we consider if it's worth adding an additional field to Peekable and Flatten/FlatMap that tracks whether or not the local “extra items” are truly “potentially additional items”. The Clone impl would ensure to set this field to “yes, they are potentially additional items” in the clone (unless the value being cloned contained only None-valued extra-item-carrying fields [peeked, frontiter, backiter]]); but the normal constructor via .peekable() or .flatten() etc… would always be able to set these fields to (false) (and the relevant fields are None anyway).

Footnotes

1. what exactly is better about those? IMHO the whole of in-place-iteration could be reduced town entirely to const items that are part of Iterator anyway… we don't need any specialization for this at all; this could also achieve an "or" logic for determining a source in Zip ↩

[the8472]

I'm thinking the latter might be quite approachable.

Hrrm, you mean with fallback to the default collect implementation? Yeah, that would work but I'm concerned that it'd lead to code bloat. Currently the path is chosen at compile time, a runtime fallback would require both paths to exist.

I'm thinking about specializing clone calls to pass down an upper bound of how many empty slots it'd have to preserve (if any), similar to the EXPAND_BY tracking. But if that gets too hairy then collect fallback may be the better choice.

IMHO the whole of in-place-iteration could be reduced town entirely to const items that are part of Interator anyway…

That could work for InPlaceIterable but we'd still need SourceIter and idk how we could select that at constant time? Offsets and field accesses maybe, but that'd be wildly unsafe.

[steffahn]

Yeah, that would work but I'm concerned that it'd lead to code bloat. Currently the path is chosen at compile time, a runtime fallback would require both paths to exist.

I guess that's a valid concern.

That could work for InPlaceIterable but we'd still need SourceIter and idk how we could select that at constant time?

Something like a const containing an Option<fn(&mut Self) -> &mut Source>; but I don't know how well that works with codegen and stuff…

[the8472]

Not certain, but it might work.
Currently the SourceIter trait is designed to support other in-place-collect impls by having different associated Source types, e.g. we could have a separate one for VecDeque sources that'd put in some extra work for rearranging the items or even one for reusing the HashSet allocation. But that flexibility is unused so for constification it could be stripped out.

[steffahn]

Looks like const might be a problem with dyn-safety.

[the8472]

And doing it through a helper trait runs into the issue that it'll require a blanket impl and then specializing that, which requires full specialization for const items, not min_specialization.

[steffahn]

@the8472 I came up with a possible way to attach a const without any helper traits. I haven't tested it beyond checking that this minimal prototype compiles at all.

[notably this should work for supporting OR-like logic for Zip. Even better: being able to make this decision with the final desired alignment as a parameter]

This approach is a bit verbose, but I found a way to get by with minimal usage of a macro, to make it less ugly.

It works without even requiring any nightly features at all.

It does restrict the setting to only vec::IntoIter(-based) sources for now. (You can't define a type that depends on consts, and it's super useful to make an informed decision in Zip based on info such as alignment, sizes, expansion/merge factors, …)

use std::{mem, num::NonZeroUsize, ptr::NonNull};

trait Iterator {
    // (implementation elided here for brevity)
    inplace_iterable_defaults!();

    type Item;
    fn next(&mut self) -> Option<Self::Item>;
}

// ####################################################
// #         normal iterators still work              #
// ####################################################

impl<I: ?Sized + Iterator> Iterator for Box<I> {
    type Item = I::Item;
    fn next(&mut self) -> Option<Self::Item> {
        (**self).next()
    }
}

fn objects_work(mut x: Box<dyn Iterator<Item = i32>>) {
    while let Some(_i) = x.next() {
        // TODO
    }
}

// ####################################################
// #        relatively small relevant API             #
// ####################################################

// type-erased struct containing pointers from a `vec::IntoIter`
// only for the non-ZST case. Values like `cap` are NOT byte-counts,
// but understood in terms of an item-size that is communicated *separately*
// (in `struct InplaceIterable`)
struct Source {
    buf: NonNull<()>,
    cap: usize,
    // TODO: add more fields as needed
}

// automatically implemented for all iterators
trait GetInplaceIterable<const ALIGN: usize>: Sized {
    const INPLACE_ITERABLE: Option<InplaceIterable<ALIGN, Self>>;
}

struct InplaceIterable<const ALIGN: usize, This> {
    get_source: fn(&This) -> Source,
    size: NonZeroUsize,
    expand: NonZeroUsize,
    merge: NonZeroUsize,
}

// and a macro
// inplace_iterable!(Type[Params][Constraints] => ident);

// ####################################################
// #             EXAMPLE 1: Map<F, I>                 #
// ####################################################

struct Map<F, I>(F, I);

impl<F: Fn(I::Item) -> R, R, I: Iterator> Iterator for Map<F, I> {
    inplace_iterable!(Map[F I][I: Iterator] => inplace_iterable_map);

    type Item = R;
    fn next(&mut self) -> Option<Self::Item> {
        todo!()
    }
}

const fn inplace_iterable_map<const ALIGN: usize, F, I>(
) -> Option<InplaceIterable<ALIGN, Map<F, I>>>
where
    I: GetInplaceIterable<ALIGN>,
{
    let Some(i) = I::INPLACE_ITERABLE else {
        return None;
    };
    Some(InplaceIterable {
        get_source: |this| (I::INPLACE_ITERABLE.unwrap().get_source)(&this.1),
        size: i.size,
        expand: i.expand,
        merge: i.merge,
    })
}

// ####################################################
// #           EXAMPLE 2: vec::IntoIter               #
// ####################################################

trait Allocator {}
struct VecIntoIter<T, A: Allocator>(T, A);

impl<T, A: Allocator> Iterator for VecIntoIter<T, A> {
    inplace_iterable!(VecIntoIter[T A][A: Allocator] => inplace_iterable_vec);

    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        todo!()
    }
}

const fn inplace_iterable_vec<const ALIGN: usize, T, A: Allocator>(
) -> Option<InplaceIterable<ALIGN, VecIntoIter<T, A>>> {
    if mem::align_of::<T>() != ALIGN {
        return None;
    }
    let Some(size) = NonZeroUsize::new(mem::size_of::<T>()) else {
        return None;
    };
    Some(InplaceIterable {
        get_source: |_this| Source {
            buf: NonNull::dangling(), // TODO
            cap: 0,                   // TODO
        },
        size,
        expand: NonZeroUsize::MIN,
        merge: NonZeroUsize::MIN,
    })
}

// ####################################################
// #               EXAMPLE 3: Zip                     #
// ####################################################

struct Zip<A, B>(A, B);

impl<A: Iterator, B: Iterator> Iterator for Zip<A, B> {
    inplace_iterable!(Zip[A B][A: Iterator, B: Iterator] => inplace_iterable_zip);

    type Item = (A::Item, B::Item);
    fn next(&mut self) -> Option<Self::Item> {
        todo!()
    }
}

const fn inplace_iterable_zip<const ALIGN: usize, A, B>(
) -> Option<InplaceIterable<ALIGN, Zip<A, B>>>
where
    A: GetInplaceIterable<ALIGN>,
    B: GetInplaceIterable<ALIGN>,
{
    let ab = (A::INPLACE_ITERABLE, B::INPLACE_ITERABLE);
    let mut choose_left = true;
    'choice: {
        // tie-breaker when both iterators support in-place iteration
        // with an appropriate alignment
        //
        // the value of `choose_left` will be ignored in the `match` below
        // if `ab` isn't `Some(_), Some(_)`,
        // but this way of structuring the code avoids repeating `match` arms
        let (Some(a), Some(b)) = ab else {
            break 'choice;
        };
        // when the multiplications fail, values have gotten very large
        // we fall back to arbitrarily choosing the left side, though
        // in-place iteration might be unusable down the line anyway,
        // and this case is uncommon to begin with
        let Some(pro_a) = a.size.checked_mul(a.merge) else {
            break 'choice;
        };
        let Some(pro_b) = b.size.checked_mul(b.merge) else {
            break 'choice;
        };
        let Some(a_advantage) = pro_a.checked_mul(b.expand) else {
            break 'choice;
        };
        let Some(b_advantage) = pro_b.checked_mul(a.expand) else {
            break 'choice;
        };
        // choose whichever side promises more free space per item consumed
        if b_advantage.get() > b_advantage.get() {
            choose_left = false;
        } else if a_advantage.get() == b_advantage.get() {
            // on a tie:
            // larger item size makes it more likely the written-back
            // values divide the existing capacity without needing re-alloc
            choose_left = a.size.get() >= b.size.get();
        }
        // else we default to the left side
    };
    Some(match (ab, choose_left) {
        ((None, None), _) => return None,
        ((Some(a), None), _) | ((Some(a), _), true) => InplaceIterable {
            get_source: |this| (A::INPLACE_ITERABLE.unwrap().get_source)(&this.0),
            size: a.size,
            expand: a.expand,
            merge: a.merge,
        },
        ((_, Some(b)), _) => InplaceIterable {
            get_source: |this| (B::INPLACE_ITERABLE.unwrap().get_source)(&this.1),
            size: b.size,
            expand: b.expand,
            merge: b.merge,
        },
    })
}

(playground see the full implementation here)

---

Rough idea how it works:

Since neither const nor associated types admit defaults, and const don't admit Self: Sized bounds, a workaround is needed. The thing that does support defaults and Self: Sized is a method. But nowadays we can have a method have an impl Trait return-type.

So the method effectively does define an associated type with a default: its return-type! And that return type can then carry the const as an item of Trait in the impl Trait.

[the8472]

Awesome. That's even on the more readable side of type system crime spectrum, with the comments at least 😆
Though it would be nice if the language had nicer syntax for such things that are already kinda supported by the type system.

As-is Source wouldn't work, we need &mut access to it to steal the allocation from the source. Maybe that's fixable by returning an adapter that points to the pointers... or could InplaceIterable have an associated type that defaults to !?

It does restrict the setting to only vec::IntoIter(-based) sources for now. (You can't define a type that depends on consts,

Hrrm, with __inplace_iterable<SourceTy> the caller could specify what it's looking for... but the source impls in the const tower would then need const fn specialization to choose None / Self.

---

This would be a fairly big refactoring and I'm not even sure yet it'd work (but dealing with Zip is a good sign, that's one of the more complicated adapters). I'll work a narrower fix for the clone problem before trying this.