From a9b8635c4d5ff08c5745787ef63fcf167c52aef0 Mon Sep 17 00:00:00 2001
From: Matt Brubeck <mbrubeck@limpet.net>
Date: Fri, 4 Jan 2019 11:31:52 -0800
Subject: [PATCH] Don't leak on panic in extend

This ensures that the length of the SmallVec is updated even if the
iterator panics in `next`.

This uses `SetLenOnDrop` rather than setting the length inside the loop,
because otherwise this suffers from the same optimization issue as
rust-lang/rust#36355.

Fixes #136.
---
 lib.rs | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/lib.rs b/lib.rs
index c6642a6..09f925f 100644
--- a/lib.rs
+++ b/lib.rs
@@ -1340,18 +1340,16 @@ impl<A: Array> Extend<A::Item> for SmallVec<A> {
         self.reserve(lower_size_bound);
 
         unsafe {
-            let len = self.len();
-            let ptr = self.as_mut_ptr().offset(len as isize);
-            let mut count = 0;
-            while count < lower_size_bound {
+            let (ptr, len_ptr, cap) = self.triple_mut();
+            let mut len = SetLenOnDrop::new(len_ptr);
+            while len.get() < cap {
                 if let Some(out) = iter.next() {
-                    ptr::write(ptr.offset(count as isize), out);
-                    count += 1;
+                    ptr::write(ptr.offset(len.get() as isize), out);
+                    len.increment_len(1);
                 } else {
                     break;
                 }
             }
-            self.set_len(len + count);
         }
 
         for elem in iter {
@@ -1560,6 +1558,11 @@ impl<'a> SetLenOnDrop<'a> {
         SetLenOnDrop { local_len: *len, len: len }
     }
 
+    #[inline]
+    fn get(&self) -> usize {
+        self.local_len
+    }
+
     #[inline]
     fn increment_len(&mut self, increment: usize) {
         self.local_len += increment;