[asder8215]

View all comments

This implements the Vec::retain_unstable method accepted as an ACP in here. Here's the tracking issue as well.

The first approach for Vec::retain_unstable is taken from what was discussed in this Zulip thread (namely @cuviper was the one who came up with this algorithm): #t-libs/api-changes > ✔ swap_retain for Vec @ 💬?

The first approach did the following:

1. Scan forward for something we want to remove (predicate returns false).
2. Scan backward for something we want to keep (predicate returns true).
   • In the reverse search process, we drop elements that the predicate
     evaluated as false.
   • In the event that no elements in the reverse search evaluated to true
     by the predicate, then the vector is simply truncated to remove the
     current and remaining elements.
3. Overwrite the forward element with the backward element.
4. Rinse and repeat from 1) until there are no elements we can iterate through.

I then changed the implementation of Vec::retain_unstable to use a different approach. This one was also discussed in the same Zulip, where I took this approach from @quaternic who proposed the idea here: #t-libs/api-changes > ✔ swap_retain for Vec @ 💬.

This approach (basically the 6th approach I tried) essentially does the following:

1. Search backward until you find an element you want to remove (implicitly that means all elements after the one you want to remove are elements you want to keep)
2. Swap this element with an element that was kept in the back (or if there are no remaining elements, simply truncate the vec)
3. Repeat from 1) until there's no element to iterate through.

It uses significantly less unsafe code, doesn't have the shenanigans of using a guard struct/drop handler, and it optimizes itself really well (performing relatively the same or better than the first approach). Vec::pop and mem::replace seems to do a really good job here to do the work of Vec::swap_remove; they both seem to have similar performance, but I do notice they each perform worse than each other in certain cases (unsure why). In any case, I opted to use Vec::pop and mem::replace.

Regarding AI Usage:
No, I have not used an LLM to write this code and algorithm. Prior to opening the PR, I took the last 3 days (probably ~20 hours), with the last two days spending time overnight to do this all by hand (as I write this now, it's 7 AM, and I've been working on this since 1 AM due to realizing what caused the double drop/repeated "unaligned tcache chunk detected" errors were coming from and just cleaning up everything). I took time to read through the ACP and Zulip thread, understood how Vec::retain_mut worked underneath the hood:

rust/library/alloc/src/vec/mod.rs

Lines 2502 to 2610 in 7211ea5

	pub fn retain_mut<F>(&mut self, mut f: F)
	where
	F: FnMut(&mut T) -> bool,
	{
	let original_len = self.len();
	if original_len == 0 {
	// Empty case: explicit return allows better optimization, vs letting compiler infer it
	return;
	}
	#[cfg(all(target_arch = "aarch64", target_feature = "sve"))]
	{
	let long_enough = match mem::size_of::<T>() {
	1 => original_len >= sve_retain::MIN_SVE_SIZE_1,
	2 => original_len >= sve_retain::MIN_SVE_SIZE_2,
	4 => original_len >= sve_retain::MIN_SVE_SIZE_4,
	8 => original_len >= sve_retain::MIN_SVE_SIZE_8,
	_ => false,
	};
	if long_enough {
	// SAFETY: size_of::<T>() is 1, 2, 4 or 8, matching
	// the kernel lane widths.
	return unsafe { sve_retain::chunked_retain(self, f) };
	}
	}
	// Vec: [Kept, Kept, Hole, Hole, Hole, Hole, Unchecked, Unchecked]
	// | ^- write ^- read |
	// |<- original_len ->|
	// Kept: Elements which predicate returns true on.
	// Hole: Moved or dropped element slot.
	// Unchecked: Unchecked valid elements.
	//
	// This drop guard will be invoked when predicate or `drop` of element panicked.
	// It shifts unchecked elements to cover holes and `set_len` to the correct length.
	// In cases when predicate and `drop` never panick, it will be optimized out.
	struct PanicGuard<'a, T, A: Allocator> {
	v: &'a mut Vec<T, A>,
	read: usize,
	write: usize,
	original_len: usize,
	}
	impl<T, A: Allocator> Drop for PanicGuard<'_, T, A> {
	#[cold]
	fn drop(&mut self) {
	let remaining = self.original_len - self.read;
	// SAFETY: Trailing unchecked items must be valid since we never touch them.
	unsafe {
	ptr::copy(
	self.v.as_ptr().add(self.read),
	self.v.as_mut_ptr().add(self.write),
	remaining,
	);
	}
	// SAFETY: After filling holes, all items are in contiguous memory.
	unsafe {
	self.v.set_len(self.write + remaining);
	}
	}
	}
	let mut read = 0;
	loop {
	// SAFETY: read < original_len
	let cur = unsafe { self.get_unchecked_mut(read) };
	if hint::unlikely(!f(cur)) {
	break;
	}
	read += 1;
	if read == original_len {
	// All elements are kept, return early.
	return;
	}
	}
	// Critical section starts here and at least one element is going to be removed.
	// Advance `g.read` early to avoid double drop if `drop_in_place` panicked.
	let mut g = PanicGuard { v: self, read: read + 1, write: read, original_len };
	// SAFETY: previous `read` is always less than original_len.
	unsafe { ptr::drop_in_place(&mut *g.v.as_mut_ptr().add(read)) };
	while g.read < g.original_len {
	// SAFETY: `read` is always less than original_len.
	let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.read) };
	if !f(cur) {
	// Advance `read` early to avoid double drop if `drop_in_place` panicked.
	g.read += 1;
	// SAFETY: We never touch this element again after dropped.
	unsafe { ptr::drop_in_place(cur) };
	} else {
	// SAFETY: `read` > `write`, so the slots don't overlap.
	// We use copy for move, and never touch the source element again.
	unsafe {
	let hole = g.v.as_mut_ptr().add(g.write);
	ptr::copy_nonoverlapping(cur, hole, 1);
	}
	g.write += 1;
	g.read += 1;
	}
	}
	// We are leaving the critical section and no panic happened,
	// Commit the length change and forget the guard.
	// SAFETY: `write` is always less than or equal to original_len.
	unsafe { g.v.set_len(g.write) };
	mem::forget(g);
	}

and traded some of my sanity to debug and fix double drop/panic errors (be it from the predicate end or the drop of the element); you may notice that the earlier implementation for this follows a similar structure to the Vec::retain_mut implementation. The tests and benchmarks are templated off the ones made for Vec::retain (if you look earlier/later in the file, you'll see them). I did add few more test cases to cover for the case on what happens if the predicate returned false on all elements (or vice versa true on all elements).

I have, however, used Claude Sonnet 5.5 to create markdown formatted table of the data I collected. All the raw data collection and benchmarking code were manually collected via ./x bench library/alloc --test-args=bench_retain and pasted into txt files. I made sure to verify that the markdown-formatted tables parsed the data correctly from the txt files I pasted them in. I hope this is an acceptable usage of AI here.

To further note, I am only showing the first Vec::retain_unstable implementation/6th + 7th approach (the pure reverse search iteration, one with Vec::pop + mem::replace and other with Vec::swap_remove); all the other approaches in between performed pretty badly or didn't exceed the performance of approach 1.

The following is benchmarking data on Vec::retain and Vec::retain_unstable (the latter of which only shows the two main approaches that were used in commit history):

Benchmark	Vec::retain	Vec::retain_unstable Approach 1	Vec::retain_unstable Approach 6	Vec::retain_unstable Approach 7
bench_retain_100000	37476.58 ± 8850.81	36563.43 ± 619.68	36798.82 ± 239.89	56966.76 ± 392.28
bench_retain_none_100000	1.24 ± 0.01	1.24 ± 0.01	1.24 ± 0.00	1.24 ± 0.00
bench_retain_quarter_false_100000	34050.51 ± 117.18	42963.59 ± 2798.66	27929.89 ± 183.32	27748.10 ± 3916.34
bench_retain_100000_alternating_with_alignment	32711.64 ± 518.39	33780.88 ± 358.60	33282.16 ± 202.27	30596.01 ± 400.60
bench_retain_100000_quarter_false_with_alignment	48943.07 ± 583.16	49239.34 ± 5107.02	49642.47 ± 389.36	48160.21 ± 1277.09
bench_retain_whole_100000	15616.17 ± 82.73	15662.04 ± 135.53	15558.17 ± 53.02	15575.18 ± 148.82
bench_retain_whole_u8_000004	1.67 ± 0.02	1.86 ± 0.01	1.65 ± 0.00	1.66 ± 0.01
bench_retain_whole_u8_000008	2.29 ± 0.03	2.28 ± 0.01	2.07 ± 0.01	2.28 ± 0.03
bench_retain_whole_u8_000016	3.54 ± 0.07	3.52 ± 0.01	3.18 ± 0.09	3.31 ± 0.11
bench_retain_whole_u8_000032	6.07 ± 1.28	6.00 ± 0.09	5.75 ± 0.14	5.92 ± 0.11
bench_retain_whole_u8_000064	10.99 ± 0.10	10.97 ± 0.11	10.75 ± 0.11	10.86 ± 0.21
bench_retain_whole_u8_000128	20.94 ± 1.52	21.12 ± 0.23	20.69 ± 0.05	20.89 ± 0.04
bench_retain_whole_u8_000256	41.42 ± 0.81	40.95 ± 0.04	44.47 ± 0.10	44.68 ± 0.87
bench_retain_whole_u8_000512	84.63 ± 0.52	84.38 ± 0.10	84.17 ± 0.06	84.39 ± 0.51
bench_retain_whole_u8_001000	163.09 ± 4.57	160.71 ± 0.26	159.89 ± 0.22	160.09 ± 1.70
bench_retain_whole_u8_010000	1563.25 ± 170.51	1556.72 ± 2.64	1555.94 ± 3.38	1556.07 ± 2.92
bench_retain_whole_u8_100000	15584.98 ± 3022.73	15519.22 ± 37.93	15519.27 ± 34.22	15520.80 ± 48.78
bench_retain_whole_u16_000004	1.67 ± 0.01	1.66 ± 0.01	1.45 ± 0.01	1.66 ± 0.01
bench_retain_whole_u16_000008	2.29 ± 0.04	2.09 ± 0.01	1.87 ± 0.01	2.27 ± 0.00
bench_retain_whole_u16_000016	3.54 ± 0.05	3.33 ± 0.14	3.19 ± 0.21	3.43 ± 0.08
bench_retain_whole_u16_000032	6.04 ± 0.05	5.89 ± 0.09	5.73 ± 0.21	5.91 ± 0.07
bench_retain_whole_u16_000064	11.02 ± 0.10	10.86 ± 0.67	10.76 ± 0.22	10.92 ± 0.14
bench_retain_whole_u16_000128	21.00 ± 0.08	20.88 ± 0.39	20.69 ± 0.35	20.86 ± 0.21
bench_retain_whole_u16_000256	40.97 ± 0.38	40.72 ± 0.32	44.74 ± 0.62	44.70 ± 3.02
bench_retain_whole_u16_000512	87.09 ± 15.45	84.59 ± 0.18	83.96 ± 0.09	84.40 ± 7.96
bench_retain_whole_u16_001000	164.42 ± 3.27	160.71 ± 0.54	159.66 ± 0.17	160.08 ± 0.41
bench_retain_whole_u16_010000	1594.39 ± 146.76	1556.66 ± 2.00	1555.61 ± 4.25	1556.03 ± 5.08
bench_retain_whole_u16_100000	15848.05 ± 3849.79	15517.92 ± 19.14	15522.59 ± 26.40	15524.23 ± 15.75
bench_retain_whole_u32_000004	1.88 ± 0.10	1.66 ± 0.01	1.45 ± 0.29	1.66 ± 0.01
bench_retain_whole_u32_000008	2.29 ± 0.06	2.08 ± 0.01	1.88 ± 0.03	2.27 ± 0.01
bench_retain_whole_u32_000016	3.55 ± 0.05	3.42 ± 0.08	3.20 ± 0.17	3.43 ± 0.21
bench_retain_whole_u32_000032	6.05 ± 0.08	5.93 ± 0.09	5.71 ± 0.25	5.93 ± 0.14
bench_retain_whole_u32_000064	11.04 ± 0.14	10.95 ± 0.18	10.75 ± 0.30	10.93 ± 0.18
bench_retain_whole_u32_000128	21.20 ± 0.30	20.86 ± 0.21	20.68 ± 0.24	20.88 ± 0.09
bench_retain_whole_u32_000256	41.20 ± 0.49	45.23 ± 5.01	44.92 ± 0.57	44.96 ± 0.28
bench_retain_whole_u32_000512	85.30 ± 1.42	84.58 ± 0.12	83.97 ± 0.16	84.59 ± 0.15
bench_retain_whole_u32_001000	161.23 ± 2.12	160.74 ± 0.52	159.68 ± 0.31	160.31 ± 0.43
bench_retain_whole_u32_010000	1570.97 ± 15.19	1566.32 ± 5.97	1559.63 ± 4.83	1560.98 ± 3.39
bench_retain_whole_u32_100000	15646.12 ± 150.06	15626.45 ± 58.76	15532.71 ± 37.13	15567.46 ± 84.21
bench_retain_whole_u64_000004	1.66 ± 0.02	1.66 ± 0.00	1.45 ± 0.00	1.86 ± 0.00
bench_retain_whole_u64_000008	2.30 ± 0.06	2.24 ± 0.20	1.87 ± 0.01	2.27 ± 0.00
bench_retain_whole_u64_000016	3.54 ± 0.04	3.53 ± 0.19	3.16 ± 0.06	3.40 ± 0.06
bench_retain_whole_u64_000032	6.05 ± 0.10	6.03 ± 0.22	5.71 ± 0.09	5.88 ± 0.15
bench_retain_whole_u64_000064	11.04 ± 0.10	11.00 ± 0.21	10.74 ± 0.14	10.89 ± 0.18
bench_retain_whole_u64_000128	20.99 ± 0.13	20.90 ± 0.45	20.68 ± 0.13	20.90 ± 0.07
bench_retain_whole_u64_000256	41.02 ± 0.45	40.77 ± 0.66	44.67 ± 0.12	45.70 ± 5.03
bench_retain_whole_u64_000512	85.22 ± 0.85	85.29 ± 1.03	84.38 ± 0.08	84.59 ± 0.33
bench_retain_whole_u64_001000	161.32 ± 1.46	160.48 ± 1.51	160.09 ± 2.16	160.27 ± 0.19
bench_retain_whole_u64_010000	1583.79 ± 23.06	1558.87 ± 16.21	1560.58 ± 3.89	1560.95 ± 3.16
bench_retain_whole_u64_100000	15924.35 ± 276.94	15710.06 ± 383.21	15707.73 ± 88.41	15646.32 ± 1949.35
bench_retain_u8_000004	6.70 ± 0.22	6.27 ± 0.06	6.84 ± 0.04	6.65 ± 0.22
bench_retain_u8_000008	7.29 ± 0.05	6.97 ± 0.11	8.96 ± 1.12	8.22 ± 0.13
bench_retain_u8_000016	9.98 ± 0.06	9.29 ± 2.05	13.86 ± 0.09	12.02 ± 0.54
bench_retain_u8_000032	15.54 ± 0.22	16.57 ± 1.01	22.92 ± 2.46	18.31 ± 0.81
bench_retain_u8_000064	26.20 ± 0.72	27.53 ± 4.85	41.98 ± 1.49	33.25 ± 3.08
bench_retain_u8_000128	49.54 ± 1.24	50.94 ± 2.90	76.62 ± 0.97	59.92 ± 8.54
bench_retain_u8_000256	91.23 ± 0.71	92.18 ± 4.19	146.77 ± 1.15	116.72 ± 3.17
bench_retain_u8_000512	174.23 ± 1.87	181.76 ± 2.15	286.72 ± 2.87	229.30 ± 9.38
bench_retain_u8_001000	332.92 ± 1.09	334.72 ± 4.55	553.55 ± 4.15	442.96 ± 7.01
bench_retain_u8_010000	3251.47 ± 75.14	3260.39 ± 28.63	5478.62 ± 42.90	4414.76 ± 25.27
bench_retain_u8_100000	32657.05 ± 2751.67	32454.35 ± 115.06	54724.37 ± 112.44	43629.83 ± 169.57
bench_retain_u16_000004	6.10 ± 0.03	5.79 ± 0.15	6.72 ± 0.04	6.70 ± 0.18
bench_retain_u16_000008	7.36 ± 0.06	6.96 ± 0.03	8.42 ± 0.04	8.35 ± 0.30
bench_retain_u16_000016	9.27 ± 0.60	9.90 ± 0.25	12.37 ± 0.31	11.70 ± 0.53
bench_retain_u16_000032	13.16 ± 1.47	15.05 ± 1.80	19.48 ± 0.86	16.50 ± 2.08
bench_retain_u16_000064	20.20 ± 2.33	26.86 ± 3.79	36.86 ± 2.86	27.13 ± 6.02
bench_retain_u16_000128	38.05 ± 4.88	50.21 ± 1.83	65.07 ± 2.02	51.31 ± 9.55
bench_retain_u16_000256	66.12 ± 15.22	91.89 ± 3.33	121.84 ± 2.32	94.34 ± 11.84
bench_retain_u16_000512	123.60 ± 0.66	175.47 ± 2.48	236.31 ± 11.70	181.07 ± 29.55
bench_retain_u16_001000	241.17 ± 34.55	341.28 ± 3.05	453.58 ± 0.83	346.32 ± 4.42
bench_retain_u16_010000	2376.12 ± 350.65	3369.49 ± 38.55	4453.28 ± 34.60	3385.89 ± 76.81
bench_retain_u16_100000	23800.32 ± 1808.66	33784.22 ± 150.95	44440.85 ± 59.13	33771.11 ± 155.56
bench_retain_u32_000004	6.67 ± 0.34	5.86 ± 0.03	6.63 ± 0.23	7.04 ± 0.02
bench_retain_u32_000008	7.38 ± 0.16	6.55 ± 0.03	7.48 ± 0.35	8.74 ± 0.06
bench_retain_u32_000016	10.14 ± 1.22	9.06 ± 0.05	10.60 ± 0.67	13.70 ± 3.53
bench_retain_u32_000032	16.34 ± 1.35	15.13 ± 0.91	15.62 ± 2.41	22.55 ± 0.17
bench_retain_u32_000064	27.31 ± 0.43	26.00 ± 0.41	26.94 ± 4.25	42.43 ± 0.12
bench_retain_u32_000128	51.99 ± 0.39	52.05 ± 4.56	51.69 ± 12.33	77.16 ± 0.45
bench_retain_u32_000256	96.51 ± 2.01	94.58 ± 0.67	97.03 ± 17.27	147.31 ± 0.71
bench_retain_u32_000512	188.00 ± 1.40	188.41 ± 2.63	189.28 ± 0.98	292.59 ± 2.53
bench_retain_u32_001000	369.73 ± 6.73	366.19 ± 3.19	366.53 ± 3.08	572.00 ± 4.90
bench_retain_u32_010000	3670.92 ± 52.18	3637.99 ± 10.76	3642.67 ± 234.14	5695.97 ± 22.42
bench_retain_u32_100000	36795.72 ± 626.05	36411.14 ± 247.22	36417.13 ± 112.45	56958.37 ± 308.11
bench_retain_u64_000004	6.41 ± 0.20	5.71 ± 0.08	6.24 ± 0.22	6.47 ± 0.04
bench_retain_u64_000008	7.91 ± 0.55	6.70 ± 0.13	8.49 ± 1.02	10.13 ± 0.08
bench_retain_u64_000016	10.69 ± 0.22	9.63 ± 0.09	10.85 ± 2.67	13.08 ± 1.25
bench_retain_u64_000032	17.20 ± 0.51	16.05 ± 0.14	17.76 ± 1.65	18.26 ± 3.04
bench_retain_u64_000064	28.55 ± 0.73	27.87 ± 0.17	30.98 ± 0.90	30.85 ± 6.37
bench_retain_u64_000128	55.61 ± 7.67	54.88 ± 0.23	61.73 ± 9.24	63.07 ± 4.98
bench_retain_u64_000256	107.25 ± 2.72	108.75 ± 1.52	114.94 ± 1.65	114.59 ± 20.70
bench_retain_u64_000512	215.64 ± 12.76	214.78 ± 1.37	224.26 ± 1.21	223.68 ± 65.43
bench_retain_u64_001000	417.88 ± 8.98	416.98 ± 2.12	427.78 ± 9.27	427.92 ± 92.07
bench_retain_u64_010000	4173.49 ± 59.43	4164.23 ± 21.13	4171.13 ± 2053.29	4159.72 ± 528.47
bench_retain_u64_100000	45586.61 ± 5794.73	41787.79 ± 701.76	41832.71 ± 196.84	41790.43 ± 373.08

The raw data can be found in this repository I made.

I was working with the current benchmarking cases that Vec::retain had and trying to see if I could optimize Vec::retain_unstable to at least be on par with Vec::retain for certain cases and then excel better than it in additional cases like filtering 25% of the data (though this was in a contiguous block) or with filtering elements aligned to 64 bytes (filling up a cache line).

There are likely more cases I could've benchmarked like how it does when we're filtering 5%, 10%, 15%, 20% of the vector, whether those filtered elements came in contiguous portions or more randomly spread apart, how it does with bigger sized elements (128 bytes, 256 bytes, etc.), etc.. If you want me to add benchmark cases for those, I could do that.

Let me know if there's anything else I could improve in this PR.

[rustbot]

Any special-casing of Miri in the standard library requires review.

cc @rust-lang/miri

[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

Why was this reviewer chosen?

The reviewer was selected based on:

• Owners of files modified in this PR: libs
• libs expanded to 12 candidates
• Random selection from JohnTitor, Mark-Simulacrum, clarfonthey

[asder8215]

Oh and to note: I did not write any SVE-accelerated code for this method. That's outside of stuff that I know currently (though I'd be curious to learn how to write SVE-accelerated code). If this method does get merged, I think @ywxt can probably optimize it further with SVE-accelerated code.

[ywxt]

SVE-accelerated code is an extra work and I will try it later. I'm not sure if it deserves tho.

[asder8215]

I'll update the PR description with the benchmark results in a formatted md table later this week. For now, this is what I have from what ./x bench library/alloc --test-args=bench_retain on the latest commit (9ad832a74c9e27e7002165e10aaad29b55c34503).

retain_benchmarks_base.txt

retain_unstable_benchmarks_9ad832a.txt

I want to say the performance between retain_unstable and current retain are similar? The cases where retain_unstable perform worse on (e.g. quarter_false_100000) makes sense because the first 25000 elements will cause the predicate to return false and the last 75000 elements will cause the predicate to return true. For retain, that just means it'll sequentially drop the elements and then shift the last 75000 elements to cover the holes (linear order) while retain_unstable will repeatedly swap the Holes with an element that can be retained from the back (does this 25000 times before just linearly retaining the rest of the items). I don't fully understand the big performance difference between bench_retain_u16_* vs bench_retain_unstable_u16_*.

[asder8215]

I know ptr::copy_nonoverlapping cannot be used when memory regions between the the src and dst pointers are overlapping, but is it safe to use it on overlapping src/dst pointers when we perform no copying (count is 0)?

I wasn't sure from reading the docs and seeing this part:

Note that even if the effectively copied size (count * size_of::<T>()) is 0, the pointers must be properly aligned.

Does that mean src == dst when count = 0 is okay here? Vec elements should be properly aligned. If it's safe, I can just update this quickly to use copy_nonoverlapping instead.

View changes since the review

[scottmcm]

Zero-length memory regions do not overlap, so with a count of zero the copy is always non-overlapping.

(See also how two distinct Box::new([]) calls will give the same address despite being unique because when the size_of_val is zero they fundamentally cannot overlap. Or similarly using [()]::split_at_mut gives two distinct mutable slices at the same address, etc.)

[scottmcm]

ymmv: I would suggest moving this to a // comment in the implementation, not a /// public comment. If they care at all what order things are in after calling this, they shouldn't be using it.

(Since it wouldn't change the complexity class, I'd like to keep open the possibility of things like "in debug if we remove anything we shuffle everything" to help people make sure they're not depending on it. Which, yes, this still does because it just says "current", but I think it's more distracting than helpful to have it in the rustdocs at all.)

View changes since the review

[asder8215]

I'm okay with not mentioning what the current implementation does. I only had this here to be in line with what Iterator::partition_in_place does in being transparent about how the swapping occurs.

I don't think it needs to be a user-facing comment since this could be subject to change. I actually have a slightly different approach I want to try out to see if it'll perform better than the current one; the approach I had in mind is don't do the reverse search to find an element which the predicate returns true on, but rather just keep swap removing the first element that you see the predicate return false on until it becomes true (and if the predicate evaluates true on an element, just iterate forward).

Basically:
If we want to retain x == 3

[1,               2,               3,               4,               5]
 ^read                                                               ^read_back

The predicate would return false on read, and just swap 1 with 5 truncating the length to 4, evaluate at the same index again and rinse and repeat the swap_remove process until its an element we can retain (which is 3, and it ends up being [3]). I don't know if this approach would be more optimal than the reverse search approach, but I want to say it's more cache-local than the current approach (since any elements that the predicate evaluates to false will be dropped on the front side rather than the front from forward search and possibly multiple back elements from the reverse search)?

[scottmcm]

suggestion: emphasize the "might be in any order" here with something like

Suggested change

	/// assert_eq!(vec, [2, 4, 3]);
	/// vec.sort(); // it might be in any order, so normalize before asserting
	/// assert_eq!(vec, [2, 3, 4]);

View changes since the review

[scottmcm]

(The tests in alloctests being tied to the exact implementation is fine, especially if that helps emphasize exactly the goal of the particular test, but in something shown to users we shouldn't show an implementation detail if we don't want people to start depending on it.)

[scottmcm]

(You've obviously spent more time thinking about this than I have, so definitely push back if what I'm saying just doesn't work in some way I missed.)

I'm really wondering how much it's worth doing all the unsafe for things here. It's very worth maintaining the hole and having drop guards for retain since we can't fill the hole every step without going quadratic.

But here we can put exactly the right thing into the hole every time without needing any extra data movement for retain_unstable. So do we really need all the extra infrastructure?

As a proof of concept, I whipped up this, which despite being fully safe code I think still does the minimal data movement and even happens to produce the same result for the doctest: https://play.rust-lang.org/?version=nightly&mode=debug&edition=2024&gist=9f4f031efc574eeaea3044232a14ee36

fn retain_unstable<T>(vec: &mut Vec<T>, mut pred: impl FnMut(&mut T) -> bool) {
    let mut front = 0;

    loop {
        // Ensure that the last element is something we want to keep.
        loop {
            let Some(mut last) = vec.peek_mut() else {
                // empty vec; we're done
                return;
            };
            if pred(&mut *last) {
                break;
            } else {
                PeekMut::pop(last);
            }
        }

        // Now that we have something to keep, move it over something we don't want.
        loop {
            // We don't get here for empty vecs, so the subtraction can't overflow.
            if vec.len() - 1 <= front {
                // there's nowhere to move it, so we're done.
                return;
            }

            if pred(&mut vec[front]) {
                front += 1;
            } else {
                vec[front] = vec.pop().unwrap();
                front += 1;
                // we used up the one we found, time to scan the end again
                break;
            }
        }
    }
}

Thanks to PeekMut that first sub-loop clearly has no bounds checks. It's just lowering the length and dropping items, exactly as we need. (And if a predicate panics, all those "removed some but not everything" states are perfectly consistent without any extra handling.)

The second subloop at the rustc level would probably emit 3 (the two indexings and the pop-unwrap), but all three are unreachable because of the "we're done" guard, which I thus would expect LLVM to optimize away. And the vec[front] = vec.pop().unwrap(); line is exactly the data movement we need, including also automatically dropping the previous one without needing to write code for it.

View changes since the review

[asder8215]

Thanks to PeekMut that first sub-loop clearly has no bounds checks. It's just lowering the length and dropping items, exactly as we need.

I tried pasting your code to see if it'll work and pass the tests, and it says that the type parameter A: core::alloc::AllocatorNightly is not satisfied. This method currently resides in impl<T, A: Allocator> Vec<T, A>, which is also where Vec::retain resides in.

[asder8215]

Looking at AllocatorNightly trait, it says in the documentation:

rust/library/core/src/alloc/mod.rs

Lines 764 to 777 in 2d37576

	/// Allocator marker trait that is only implemented stably for `Global`.
	///
	/// This trait is used as a bound whenever a function constructing
	/// a type with an `#[unstable] A: Allocator = Global` parameter
	/// may be callable for `A != Global`.
	///
	/// This is to prevent stable code from e.g. constructing `Arc<T, NotGlobal>`
	/// using the `From<Box<T, A>> for Arc<T, A>` impl.
	///
	/// Note that this trait *will* be removed in the future, but downstream
	/// allocators should still feel free to implement it on nightly, to
	/// allow use of all containers with unstable allocator support.
	#[unstable(feature = "allocator_ext", issue = "163177", implied_by = "allocator_api")]
	pub trait AllocatorNightly: Allocator {}

I guess we're not supposed to use Vec::peek_mut in stable library code currently?

Note: to be clear, I would be happier to use safer code for this method. I wrote this the way it was because I assumed doing it in a similar fashion as Vec::retain_mut squeezed out the most performance as possible. The different approach I'm trying out, which I mentioned in #164155 (comment), I hesitate to use Vec::swap_remove because it uses ptr::copy underneath the hood when we could easily use ptr::copy_nonoverlapping` for this code.

[scottmcm]

Yeah, I was contemplating swap_remove as well but ended up with vec[front] = vec.pop().unwrap(); because I don't think the bounds checks are materially different (the same check is there in swap_remove, just hidden) and as you say it can be nonoverlapping.

I guess we're not supposed to use Vec::peek_mut in stable library code currently?

I have no idea why it's written that way, since peek_mut and PeekMut are themselves unstable. (In general, using unstable alloc features inside alloc is completely fine and done all over the place.)

Could make another PR to make peek_mut usable, or just use a last+pop loop for now with a FIXME -- it'll probably optimize well anyway (if last exists then pop().unwrap() can typically optimize out because the last is a non-empty proof) and if not that's fixable later.

---

That said, I'm not at all tied to the specific code I wrote there. Feel free to change it as much as you see fit, or take inspiration from the idea and ignore the code entirely.

Some unsafe would be fine by me too -- if you show that the vec[front] = vec.pop().unwrap(); isn't good enough and should be done with unsafe, for example, that sounds entirely plausible to me. But that's a very local piece of unsafe that could be much simpler and not need raw pointers everywhere, for example. (And wouldn't need to keep anything across loop iterations.)

[asder8215]

I made slight changes to the code and here's what I tried out:

fn retain_unstable<T>(vec: &mut Vec<T>, mut pred: impl FnMut(&mut T) -> bool) {
        let mut front = 0;

        loop {
            // Ensure that the last element is something we want to keep.
            loop {
                let Some(_last) = self.last() else {
                    // empty vec; we're done
                    return;
                };

                let back = self.len() - 1;
                if f(&mut self[back]) {
                    break;
                } else {
                    self.pop().unwrap();
                }
            }

            // Now that we have something to keep, move it over something we don't want.
            loop {
                // We don't get here for empty vecs, so the subtraction can't overflow.
                if self.len() - 1 <= front {
                    // there's nowhere to move it, so we're done.
                    return;
                }

                if f(&mut self[front]) {
                    front += 1;
                } else {
                    self[front] = self.pop().unwrap();
                    front += 1;
                    // we used up the one we found, time to scan the end again
                    break;
                }
            }
        }
}

I ran the benchmark cases that I put in benches/vec.rs (although I realize I forgot to update this branch with a couple of benchmark cases, and I'll add them later today after cleaning up my local code and commits) and here is what I got. When I compare this to the current PR's timings, it's kind of worse altogether except for the case where we have the first 25000 elements in a vec returning false by the predicate and the last 75000 returns true. For context on bench_retain_*_quarter_false_100000, it's just this:

#[bench]
fn bench_retain_quarter_false_100000(b: &mut Bencher) {
    let mut v = Vec::with_capacity(100000);

    b.iter(|| {
        v.clear();
        v.extend(black_box(1..=100000));
        v.retain(|x| if *x < 25000 { false } else { true }) // change to v.retain_unstable for the unstable counterpart
    });
}

[asder8215]

Also, just putting it out there, I'll try doing the same algorithm you suggested using unsafe code (seems like it's just swap around of the predicate logic and read and read_back conditions in the current code here) to see if that improves performance.

[scottmcm]

Note that I had the steps in that order because it was convenient for the safe implementation, not because I have any particular strong feeling that said rearrangement is better. Starting with a scan forward also seems entirely reasonable to me.

---

To uplevel this thread a bit: I saw the note about this PR getting opened, said "oh, great, I'll go review that", then saw that there was way more code to review than I was going to be able to easily understand in ≈10 minutes, so I'm just hoping that there's a way -- even if it still uses unsafe code -- that I can more easily look at and think "oh, yeah, I'm easily confident in it".

[quaternic]

@scottmcm
Your code has an issue; the first inner loop needs to stop before re-testing the last element that was already tested and kept in the second inner loop. Crucially, the peek_mut succeeds exactly when len > 0, but we actually want len > front. Easy to miss since it doesn't matter for the typical case of pure predicates.

[asder8215]

@scottmcm should be easier to read and reason about the code now since it uses significantly less unsafe. I took the approach @quaternic suggested earlier in the Zulip thread with doing a reverse sweep in Vec::retain_unstable and some sort of Vec::swap_remove. This approach is a lot cleaner to do than doing a forward then backward search (or vice versa) and optimizes just as well as the unsafe version.

[rustbot]

This PR was rebased onto a different main commit. Here's a range-diff highlighting what actually changed.

Rebasing is a normal part of keeping PRs up to date, so no action is needed—this note is just to help reviewers.

[Mark-Simulacrum]

Is it good for us to hint this is unlikely? I guess the rationale is that this branch is true at most once, but that seems like it's inferable by the compiler and it's not obvious the annotation here adds value?

(It seems like it may well be likely that e.g. the last element should be the one removed for a particular call site / T?)

View changes since the review

[asder8215]

I took this from what Vec::retain_mut does currently. It seems like it made a difference in the code there, but it also uses a lot of unsafe code, so unsure if some things just weren't easy for the compiler to optimize.

I can also try benchmarking again without this hint.

[asder8215]

So after applying the Vec::truncate suggestion you made and benchmarking it with/without hint::unlikely, it seems like the hint::unlikely annotation matters here. Without it, there are several cases where I see lower performance than the one applies hint::unlikely.

I gathered the results here.

[Mark-Simulacrum]

I'm unsurprised that we can find example scenarios where unlikely helps -- e.g. if the element to be found is commonly in the middle -- but do we have a benchmark where we are specifically trying to make this condition likely and still finding that unlikely is helpful?

[asder8215]

The only benchmarks I can where this would matter if this was likely is if we find an element we want to filter out early towards the end of the Vec, which technically we do have that with bench_retain_unstable_none (which returns false on all elements) or in the bench_retain_unstable_(u8/u16/u32/u64) case which makes the predicate alternate between returning true or false on the elements (there's an early false occurring at the end of the Vec).

Neither case showed that hint::unlikely had an impact on performance here (from the linked data above), which I think makes sense because, like you said earlier, we want this branch to be true at most once. A large bulk of filtering usually occurs in the loop below, which doesn't use hint::unlikely.

[Mark-Simulacrum]

I think my expectation is that we would want to swap_remove here, such that we avoid the to-stack and from-stack copies of the last element? But seems like that's mostly an optimization so doesn't necessarily need to land in this PR...

View changes since the review

[asder8215]

I'm okay with using Vec::swap_remove, but I'll admit that I'm not a fan of how it'll use ptr::copy from its implementation:

rust/library/alloc/src/vec/mod.rs

Lines 2170 to 2193 in 7d0f9b4

	pub fn swap_remove(&mut self, index: usize) -> T {
	#[cold]
	#[cfg_attr(not(panic = "immediate-abort"), inline(never))]
	#[optimize(size)]
	fn assert_failed(index: usize, len: usize) -> ! {
	panic!("swap_remove index (is {index}) should be < len (is {len})");
	}
	let len = self.len();
	if index >= len {
	assert_failed(index, len);
	}
	// ignore-tidy-undocumented-unsafe
	unsafe {
	// We replace self[index] with the last element. Note that if the
	// bounds check above succeeds there must be a last element (which
	// can be self[index] itself).
	let value = ptr::read(self.as_ptr().add(index));
	let base_ptr = self.as_mut_ptr();
	ptr::copy(base_ptr.add(len - 1), base_ptr.add(index), 1);
	self.set_len(len - 1);
	value
	}
	}

This is a scenario where ptr::copy_nonoverlapping seems like it should be possible here since we can easily fill the gap with the item at the back of the vec (which should be non-overlapping), truncate the vec, and then drop the value that was there in the hole. However, when I tried copying Vec::swap_remove code in place replacing ptr::copy with ptr::copy_nonoverlapping:

if read_back < self.len {
         let _ = unsafe { ptr::read(self.as_ptr().add(read_back - 1)) }; 
         let base_ptr = self.as_mut_ptr(); 
         unsafe { 
             ptr::copy_nonoverlapping(base_ptr.add(self.len - 1), base_ptr.add(read_back - 1), 1); 
             self.set_len(self.len - 1);
         };
}

This caused a panic issue with test_retain_unstable_drop_panic, which I didn't really understand why since these memory regions should be non-overlapping and we don't drop the element until the end of the conditional block.

I know mem::replace uses both ptr::read_via_copy and ptr::write_via_move:

rust/library/core/src/mem/mod.rs

Lines 963 to 981 in 7d0f9b4

	pub const fn replace<T>(dest: &mut T, src: T) -> T {
	// It may be tempting to use `swap` to avoid `unsafe` here. Don't!
	// The compiler optimizes the implementation below to two `memcpy`s
	// while `swap` would require at least three. See PR#83022 for details.
	// SAFETY: We read from `dest` but directly write `src` into it afterwards,
	// such that the old value is not duplicated. Nothing is dropped and
	// nothing here can panic.
	unsafe {
	// Ideally we wouldn't use the intrinsics here, but going through the
	// `ptr` methods introduces two unnecessary UbChecks, so until we can
	// remove those for pointers that come from references, this uses the
	// intrinsics instead so this stays very cheap in MIR (and debug).
	let result = crate::intrinsics::read_via_copy(dest);
	crate::intrinsics::write_via_move(dest, src);
	result
	}
	}

So between ptr::copy and what mem::replace does, I was hoping mem::replace would perform faster than ptr::copy or do something akin to ptr::copy_nonoverlapping (note, I don't know too much the performance diff of these intrinsics, which I'd like to hear if it's just better to use ptr::copy directly than a ptr::read_via_copy and ptr::write_via_move)

I do have a separate implementation with a PanicGuard that takes the same approach we have here and works with ptr::copy_nonoverlapping, but I was hoping to use less unsafe code.

[rustbot]

Reminder, once the PR becomes ready for a review, use @rustbot ready.

[Mark-Simulacrum]

r=me with nits fixed, or if you want to pursue the swap_remove here that's fine with me too (but maybe better to land and iterate in a subsequent PR)

View changes since this review

[Mark-Simulacrum]

Should we say something about this being generally faster than retain_mut due to the less-strict guarantee? I think retain_mut is also O(n)?

[asder8215]

I'm not opposed to saying that Vec::retain_unstable is generally faster than Vec::retain_mut if ordering doesn't matter to you, but I'm wondering if we should leave this to people to benchmark their usecase for this method?

[Mark-Simulacrum]

I think guidance is useful here. We give it for sort vs. sort_unstable:

It is typically faster than stable sorting, except in a few special cases, e.g., when the slice is partially sorted.

[asder8215]

I'll make a separate PR to adjust documentation once this is merged.

[asder8215]

I only changed mem::replace line with *hole = swap, so I doubt this would affect CI. Marking this as ready.

@rustbot ready

[Mark-Simulacrum]

@bors r+

[rust-bors]

📌 Commit cda07fc has been approved by Mark-Simulacrum

It is now in the queue for this repository.

[inkreasing]

I was a bit confused if this is actually sound, since what it looked like to me was creating a mut ref to the vec (hole) then using the old ref (self.pop()) and then using the mut ref again (*hole = swap).

From what i can see this depends on Vec::pop not creating a reference to the allocation of the Vec. Since the created ref points to the allocation and the pop only uses the reference to the Vec itself hole isn't invalidated.

I think an implementation of Vec::pop that invalidates all pointers to the allocation would be allowed from the documentation.

Since this is also std code maybe this is fine and if the implementation of pop changes this would probably be caught by miri in CI, but i was very confused by it, so maybe a comment would be nice.

(same in the code above)

[Mark-Simulacrum]

I think there's no intentional problem here, and I'd expect miri to (eventually, at least) catch things here if it's an issue. I don't think comments are super helpful here personally, they seem more distracting than useful as long as we have the right safeguards (via tests etc.)