From 04f1c09f90abebd0c6a7658105dec57099a63caa Mon Sep 17 00:00:00 2001
From: Alan Egerton <eggyal@gmail.com>
Date: Sat, 27 Nov 2021 14:19:24 +0000
Subject: [PATCH] Avoid UB when short-circuiting try_map_id for Vec

---
 compiler/rustc_data_structures/src/functor.rs | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/compiler/rustc_data_structures/src/functor.rs b/compiler/rustc_data_structures/src/functor.rs
index 1307c68ba0b..d3715a998ce 100644
--- a/compiler/rustc_data_structures/src/functor.rs
+++ b/compiler/rustc_data_structures/src/functor.rs
@@ -87,7 +87,6 @@ impl<T> IdFunctor for Vec<T> {
         // FIXME: We don't really care about panics here and leak
         // far more than we should, but that should be fine for now.
         let len = self.len();
-        let mut error = Ok(());
         unsafe {
             self.set_len(0);
             let start = self.as_mut_ptr();
@@ -96,8 +95,16 @@ impl<T> IdFunctor for Vec<T> {
                 match f(ptr::read(p)) {
                     Ok(value) => ptr::write(p, value),
                     Err(err) => {
-                        error = Err(err);
-                        break;
+                        // drop all other elements in self
+                        // (current element was "moved" into the call to f)
+                        for j in (0..i).chain(i + 1..len) {
+                            let p = start.add(j);
+                            ptr::drop_in_place(p);
+                        }
+
+                        // returning will drop self, releasing the allocation
+                        // (len is 0 so elements will not be re-dropped)
+                        return Err(err);
                     }
                 }
             }
@@ -105,7 +112,7 @@ impl<T> IdFunctor for Vec<T> {
             // so we don't leak memory.
             self.set_len(len);
         }
-        error.map(|()| self)
+        Ok(self)
     }
 }