Fix UB in ThinVec::flat_map_in_place

thin_vec.as_ptr() goes through the Deref impl of ThinVec, which will
not allow access to any memory as we did call set_len(0) first.

Author: bjorn3

compiler/rustc_data_structures/src/flat_map_in_place.rs

@@ -1,4 +1,4 @@
-use std::ptr;
+use std::{mem, ptr};
 
 use smallvec::{Array, SmallVec};
 use thin_vec::ThinVec;
@@ -69,5 +69,55 @@ impl<T, A: Array<Item = T>> FlatMapInPlace<T> for SmallVec<A> {
 }
 
 impl<T> FlatMapInPlace<T> for ThinVec<T> {
-    flat_map_in_place!();
+    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>(&'a mut ThinVec<T>);
+
+        impl<'a, T> 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
+                }
+            }
+        }
+
+        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);
+        }
+    }
 }