RepeatN: use MaybeUninit

[fee1-dead]

Closes #130140. Closes #130141.

Use MaybeUninit instead of ManuallyDrop for soundness.

[rustbot]

r? @Mark-Simulacrum

rustbot has assigned @Mark-Simulacrum.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

[fee1-dead]

Note that this changes the Debug output, the element field will now be wrapped inside a Some if count is at least 1, and will display as None otherwise. Previously it just printed whatever is in there (which is unsound)

[lcnr]

also test clone please 🤔

r=me afterwards

[lcnr]

@bors r+ rollup=never

[bors]

📌 Commit 40b6f96 has been approved by lcnr

It is now in the queue for this repository.

[lcnr]

cc @rust-lang/libs

[clubby789]

Suggested change

	assert_eq(0, *x);
	assert_eq!(0, *x);

[fee1-dead]

oops

[fee1-dead]

@bors r-

[WaffleLapkin]

I think the use of ManuallyDrop was correct here (ignoring the UAF issue obviously), however before rust-lang/unsafe-code-guidelines#245 is fixed, it makes sense to use ManuallyDrop instead.

Can you leave a FIXME to replace it back with ManuallyDrop, once it doesn't cause issues?

[antonilol]

This iterator could be implemented using minimal to no unsafe code (I used none in this example) like this:

use core::num::NonZeroUsize;

#[derive(Debug, Clone)]
pub struct RepeatN<T>(Option<(T, NonZeroUsize)>);

impl<T> RepeatN<T> {
    pub fn new(value: T, count: usize) -> Self {
        Self(NonZeroUsize::new(count).map(|count| (value, count)))
    }

    // for use in last and maybe other methods
    fn take_element(&mut self) -> Option<T> {
        self.0.take().map(|(value, _count)| value)
    }
}

impl<T: Clone> Iterator for RepeatN<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        let (value, count) = self.0.as_mut()?;
        if let Some(new_count) = NonZeroUsize::new(count.get() - 1) {
            *count = new_count;
            Some(value.clone())
        } else {
            // this could be unwrap_unchecked
            Some(self.0.take().unwrap().0)
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let len = self.0.as_ref().map_or(0, |(_value, count)| count.get());
        (len, Some(len))
    }

    ...
}

[jwong101]

That also has the advantage of interacting with dropck a bit better(although you could also use may_dangle):

fn src() -> String {
    // fails today, but shouldn't if you use the option approach
    let x = "foo".to_owned();
    let mut y = std::iter::repeat_n(&*x, 1);
    // use y, but don't consume it entirely
    x
}

[antonilol]

Does this also preserve niches from T? MaybeUninit for sure not.

[scottmcm]

FWIW, it's intentional that this was ManuallyDrop and not MaybeUninit, because there's no case in which this is constructed without a value, and thus we can preserve niches for it.

Does this actually need to change to MaybeUninit? Feels to me like the fix is to have Clone do a bitwise copy if count == 0 -- which will have the nice behaviour of eliding the branch for T: Copy.

[WaffleLapkin]

Yes, ManuallyDrop currently doesn't work here, see #130141 and #130145 (review).

[scottmcm]

Does that mean that ManuallyDrop<Box<_>> is just never sound? Do we need #[lang = "manually_drop"] to remove that requirement?

[WaffleLapkin]

Does that mean that ManuallyDrop<Box<_>> is just never sound?

It's sound if you don't touch it in any way after running Box's drop I think (i.e. not even move). But generally yes.

To fix this we need to implement MaybeDangling (#118166) which would be a lang item, yes (and then wrap the only field of ManuallyDrop with MaybeDangling).

[scottmcm]

To @antonilol's point, the only reason this was using unsafe code was to preserve niches, and thus to allow the implementation for Copy types to always just read the payload without caring about the count.

If this is going to use MaybeUninit -- and thus not get noundef in LLVM and such -- then it should just use the safe implementation. The safe implementation is just as efficient as the MaybeUninit version.

[cuviper]

Option<(T, NonZeroUsize)> would have slightly less niche space than (ManuallyDrop<T>, usize), because the current layout implementation only propagates the largest niche.

• If T has no niche values, then the single niche in NonZeroUsize will be used by Option::None. This is the same as before with no niches available outside.
• If T has one niche value, then None may use that or the NonZeroUsize again, but the other will not be available. So this is a lost niche compared to before.
• If T has multiple niche values, then None will use one of those, and the rest will be available outside. The NonZeroUsize niche will not be available. This is one less niche compared to before.

[jwong101]

Tbh, it's probably best if the NonZero field is used as the niche for the option, as that allows just loading that field for stuff like size_hint.

use std::num::NonZero;

type T = Option<(String, NonZero<usize>)>;

#[no_mangle]
pub fn src(x: &T) -> usize {
    x.as_ref().map_or(0, |a| a.1.get())
}

type U = Option<(NonZero<usize>, String)>;

#[no_mangle]
pub fn dst(x: &U) -> usize {
    x.as_ref().map_or(0, |a| a.0.get())
}

#[no_mangle]
pub fn good(x: &Option<(NonZero<usize>, std::os::fd::OwnedFd)>) -> usize {
    x.as_ref().map_or(0, |a| a.0.get())
}

#[no_mangle]
pub fn bad(x: &Option<(std::os::fd::OwnedFd, NonZero<usize>)>) -> usize {
    x.as_ref().map_or(0, |a| a.1.get())
}

src:                                    # @src
# %bb.0:
	xor	eax, eax
	cmp	rax, qword ptr [rdi]
	jo	.LBB0_2
# %bb.1:
	mov	rax, qword ptr [rdi + 24]

.LBB0_2:
	ret
                                        # -- End function

dst:                                    # @dst
# %bb.0:
	xor	eax, eax
	cmp	rax, qword ptr [rdi + 8]
	jo	.LBB1_2
# %bb.1:
	mov	rax, qword ptr [rdi]

.LBB1_2:
	ret
                                        # -- End function

good:                                   # @good
# %bb.0:
	mov	rax, qword ptr [rdi]
	ret
                                        # -- End function

bad:                                    # @bad
# %bb.0:
	cmp	dword ptr [rdi], -1
	je	.LBB3_1
# %bb.2:
	mov	rax, qword ptr [rdi + 8]
	ret

.LBB3_1:
	xor	eax, eax
	ret

The loss of the other niches mainly hurts the size of it being wrapped in iterator adapters that would wrap it in an option like chain. Though I think the main regression would be losing noundef, and Rust potentially choosing a suboptimal niche that doesn't overlap with the NonZero<size>.

[antonilol]

Tbh, it's probably best if the NonZero field is used as the niche for the option, as that allows just loading that field for stuff like size_hint.

If the compiler doesn't know the original size, that could result in this extra branch, but if it does know, it is smart enough to figure it out (first 3 examples, compiler explorer link)

https://rust.godbolt.org/z/W8hqncT81

(How does this even produce 10x more code for 'len_5' compared to 'len_4'? That feels very counterintuitive.)

[theemathas]

Is this the same as returning Some(self.element.assume_init_read())?

[fee1-dead]

Hmm. not too sure about this.

[antonilol]

I think it is, but there could be a subtle difference between the two methods.

[zachs18]

Writing uninit() into self.element will likely be elided in release mode, but explicitly writing uninit() into self.element would let sanitizers (e.g. Miri) explicitly catch if the field is later read as initialized, even if T: Copy.

Compare these two functions: https://rust.godbolt.org/z/cK48E76z3 At -Copt-level=1, they are almost the same (notably the mem::replace(&mut self.element, uninit()) does not actually write anything to self.element), and at -Copt-level=2 or above, the two functions are compiled to the same code.

[fee1-dead]

I'd like a definite answer on what approach I should take (e.g. fixing the use-after-free first by not changing the field type and writing manual impls but keeping the miri reported UB due to using ManuallyDrop?)

[Noratrieb]

https://rust-lang.zulipchat.com/#narrow/stream/219381-t-libs/topic/repeat_n.20unsoundness

[workingjubilee]

It's fun to fuss over an unsoundness issue that is also an optimization problem, but it's better to actually ship the fix. This code seems adequate to quell the soundness issue, we can fuss over improvements later.

@bors r=lcnr,workingjubilee

[bors]

📌 Commit 4c8b84a has been approved by lcnr,workingjubilee

It is now in the queue for this repository.

[bors]

⌛ Testing commit 4c8b84a with merge 2e367d9...

[bors]

☀️ Test successful - checks-actions
Approved by: lcnr,workingjubilee
Pushing 2e367d9 to master...

[rust-timer]

Finished benchmarking commit (2e367d9): comparison URL.

Overall result: ❌ regressions - no action needed

@rustbot label: -perf-regression

Instruction count

This is a highly reliable metric that was used to determine the overall result at the top of this comment.

	mean	range	count
Regressions ❌ (primary)	-	-	0
Regressions ❌ (secondary)	0.3%	[0.3%, 0.3%]	1
Improvements ✅ (primary)	-	-	0
Improvements ✅ (secondary)	-	-	0
All ❌✅ (primary)	-	-	0

Max RSS (memory usage)

This benchmark run did not return any relevant results for this metric.

Cycles

This benchmark run did not return any relevant results for this metric.

Binary size

This benchmark run did not return any relevant results for this metric.

Bootstrap: 759.355s -> 759.624s (0.04%)
Artifact size: 341.37 MiB -> 341.31 MiB (-0.02%)

[cuviper]

Since this is stabilizing in 1.82:

@rustbot label beta-nominated