Rollup merge of rust-lang#155582 - nnethercote:overhaul-flat_map_in_place, r=chenyukang

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.

r? @chenyukang

Author: JonathanBrouwer

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<T>: Sized {
+pub trait FlatMapInPlace<T> {
+    /// `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<F, I>(&mut self, f: F)
     where
         F: FnMut(T) -> I,
         I: IntoIterator<Item = T>;
 }
 
-// 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<F, I>(&mut self, mut f: F)
-        where
-            F: FnMut(T) -> I,
-            I: IntoIterator<Item = T>,
-        {
-            struct LeakGuard<'a, T $(: $bound)?>(&'a mut $vec<T>);
-
-            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
-                    }
+// Blanket impl for all vec-like types that impl `FlatMapInPlaceVec`.
+impl<V: FlatMapInPlaceVec> FlatMapInPlace<V::Elem> for V {
+    fn flat_map_in_place<F, I>(&mut self, mut f: F)
+    where
+        F: FnMut(V::Elem) -> I,
+        I: IntoIterator<Item = V::Elem>,
+    {
+        struct LeakGuard<'a, V: FlatMapInPlaceVec>(&'a mut V);
+
+        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;
-                        }
+        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.
-                this.0.set_len(write_i);
+            // `write_i` tracks the number of actually written new items.
+            guard.0.set_len(write_i);
 
-                // The ThinVec is in a sane state again. Prevent the LeakGuard from leaking the data.
-                mem::forget(this);
-            }
+            // `vec` is in a sane state again. Prevent the LeakGuard from leaking the data.
+            mem::forget(guard);
         }
-    };
+    }
+}
+
+// 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<T> FlatMapInPlace<T> for Vec<T> {
-    flat_map_in_place!(Vec);
+impl<T> FlatMapInPlaceVec for Vec<T> {
+    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<T, A: Array<Item = T>> FlatMapInPlace<T> for SmallVec<A> {
-    flat_map_in_place!(SmallVec where T: Array);
+impl<T> FlatMapInPlaceVec for ThinVec<T> {
+    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<T> FlatMapInPlace<T> for ThinVec<T> {
-    flat_map_in_place!(ThinVec);
+impl<T, const N: usize> 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);
+    }
 }