From 21cd762cd16ff9806209c605a750491e223a75b7 Mon Sep 17 00:00:00 2001 From: Nicholas Nethercote Date: Tue, 21 Apr 2026 15:06:58 +1000 Subject: [PATCH] Rewrite `FlatMapInPlace`. Replace the hacky macro with a generic function and a new `FlatMapInPlaceVec` trait. More verbose but more readable and typical. LLM disclosure: I asked Claude Code to critique this file and it suggested the generic function + trait idea. I implemented the idea entirely by hand. --- .../src/flat_map_in_place.rs | 196 ++++++++++++------ 1 file changed, 136 insertions(+), 60 deletions(-) diff --git a/compiler/rustc_data_structures/src/flat_map_in_place.rs b/compiler/rustc_data_structures/src/flat_map_in_place.rs index 6d718059f9c8..d1fdd999d36a 100644 --- a/compiler/rustc_data_structures/src/flat_map_in_place.rs +++ b/compiler/rustc_data_structures/src/flat_map_in_place.rs @@ -1,81 +1,157 @@ use std::{mem, ptr}; -use smallvec::{Array, SmallVec}; +use smallvec::SmallVec; use thin_vec::ThinVec; -pub trait FlatMapInPlace: Sized { +pub trait FlatMapInPlace { + /// `f` turns each element into 0..many elements. This function will consume the existing + /// elements in a vec-like structure and replace them with any number of new elements — fewer, + /// more, or the same number — as efficiently as possible. fn flat_map_in_place(&mut self, f: F) where F: FnMut(T) -> I, I: IntoIterator; } -// The implementation of this method is syntactically identical for all the -// different vector types. -macro_rules! flat_map_in_place { - ($vec:ident $( where T: $bound:path)?) => { - fn flat_map_in_place(&mut self, mut f: F) - where - F: FnMut(T) -> I, - I: IntoIterator, - { - struct LeakGuard<'a, T $(: $bound)?>(&'a mut $vec); +// Blanket impl for all vec-like types that impl `FlatMapInPlaceVec`. +impl FlatMapInPlace for V { + fn flat_map_in_place(&mut self, mut f: F) + where + F: FnMut(V::Elem) -> I, + I: IntoIterator, + { + struct LeakGuard<'a, V: FlatMapInPlaceVec>(&'a mut V); - impl<'a, T $(: $bound)?> Drop for LeakGuard<'a, T> { - fn drop(&mut self) { - unsafe { - self.0.set_len(0); // make sure we just leak elements in case of panic - } + impl<'a, V: FlatMapInPlaceVec> Drop for LeakGuard<'a, V> { + fn drop(&mut self) { + unsafe { + // Leak all elements in case of panic. + self.0.set_len(0); } } - - let this = LeakGuard(self); - - let mut read_i = 0; - let mut write_i = 0; - unsafe { - while read_i < this.0.len() { - // move the read_i'th item out of the vector and map it - // to an iterator - let e = ptr::read(this.0.as_ptr().add(read_i)); - let iter = f(e).into_iter(); - read_i += 1; - - for e in iter { - if write_i < read_i { - ptr::write(this.0.as_mut_ptr().add(write_i), e); - write_i += 1; - } else { - // If this is reached we ran out of space - // in the middle of the vector. - // However, the vector is in a valid state here, - // so we just do a somewhat inefficient insert. - this.0.insert(write_i, e); - - read_i += 1; - write_i += 1; - } - } - } - - // write_i tracks the number of actually written new items. - this.0.set_len(write_i); - - // The ThinVec is in a sane state again. Prevent the LeakGuard from leaking the data. - mem::forget(this); - } } - }; + + let guard = LeakGuard(self); + + let mut read_i = 0; + let mut write_i = 0; + unsafe { + while read_i < guard.0.len() { + // Move the read_i'th item out of the vector and map it to an iterator. + let e = ptr::read(guard.0.as_ptr().add(read_i)); + let iter = f(e).into_iter(); + read_i += 1; + + for e in iter { + if write_i < read_i { + ptr::write(guard.0.as_mut_ptr().add(write_i), e); + write_i += 1; + } else { + // If this is reached we ran out of space in the middle of the vector. + // However, the vector is in a valid state here, so we just do a somewhat + // inefficient insert. + guard.0.insert(write_i, e); + + read_i += 1; + write_i += 1; + } + } + } + + // `write_i` tracks the number of actually written new items. + guard.0.set_len(write_i); + + // `vec` is in a sane state again. Prevent the LeakGuard from leaking the data. + mem::forget(guard); + } + } } -impl FlatMapInPlace for Vec { - flat_map_in_place!(Vec); +// A vec-like type must implement these operations to support `flat_map_in_place`. +pub trait FlatMapInPlaceVec { + type Elem; + + fn len(&self) -> usize; + unsafe fn set_len(&mut self, len: usize); + fn as_ptr(&self) -> *const Self::Elem; + fn as_mut_ptr(&mut self) -> *mut Self::Elem; + fn insert(&mut self, idx: usize, elem: Self::Elem); } -impl> FlatMapInPlace for SmallVec { - flat_map_in_place!(SmallVec where T: Array); +impl FlatMapInPlaceVec for Vec { + type Elem = T; + + fn len(&self) -> usize { + self.len() + } + + unsafe fn set_len(&mut self, len: usize) { + unsafe { + self.set_len(len); + } + } + + fn as_ptr(&self) -> *const Self::Elem { + self.as_ptr() + } + + fn as_mut_ptr(&mut self) -> *mut Self::Elem { + self.as_mut_ptr() + } + + fn insert(&mut self, idx: usize, elem: Self::Elem) { + self.insert(idx, elem); + } } -impl FlatMapInPlace for ThinVec { - flat_map_in_place!(ThinVec); +impl FlatMapInPlaceVec for ThinVec { + type Elem = T; + + fn len(&self) -> usize { + self.len() + } + + unsafe fn set_len(&mut self, len: usize) { + unsafe { + self.set_len(len); + } + } + + fn as_ptr(&self) -> *const Self::Elem { + self.as_slice().as_ptr() + } + + fn as_mut_ptr(&mut self) -> *mut Self::Elem { + self.as_mut_slice().as_mut_ptr() + } + + fn insert(&mut self, idx: usize, elem: Self::Elem) { + self.insert(idx, elem); + } +} + +impl FlatMapInPlaceVec for SmallVec<[T; N]> { + type Elem = T; + + fn len(&self) -> usize { + self.len() + } + + unsafe fn set_len(&mut self, len: usize) { + unsafe { + self.set_len(len); + } + } + + fn as_ptr(&self) -> *const Self::Elem { + self.as_ptr() + } + + fn as_mut_ptr(&mut self) -> *mut Self::Elem { + self.as_mut_ptr() + } + + fn insert(&mut self, idx: usize, elem: Self::Elem) { + self.insert(idx, elem); + } }