Refactor: Add add/sub to TaggedLen (#570) (#572)

* Refactor TaggedLen to use add and sub directly (#570)

* refactor: remove increment, decrement methods in TaggedLen

---------

Signed-off-by: Alejandro Vaz <alejandro.vaz.myt@gmail.com>
Co-authored-by: Alejandro Vaz <alejandro.vaz.myt@gmail.com>

Author: iamjaeholee

src/lib.rs

@@ -289,7 +289,7 @@ impl<'a, T: 'a, const N: usize> Drop for Drain<'a, T, N> {
                             let dst = ptr.add(start);
                             core::ptr::copy(src, dst, self.0.tail_len);
                         }
-                        source_vec.set_len(start + self.0.tail_len);
+                        source_vec.len.add(self.0.tail_len);
                     }
                 }
             }
@@ -308,7 +308,7 @@ impl<'a, T: 'a, const N: usize> Drop for Drain<'a, T, N> {
             unsafe {
                 let vec = vec.as_mut();
                 let old_len = vec.len();
-                vec.set_len(old_len + drop_len + self.tail_len);
+                vec.len.add(drop_len + self.tail_len);
                 vec.truncate(old_len + self.tail_len);
             }
 
@@ -374,7 +374,7 @@ impl<T, const N: usize> Drain<'_, T, N> {
             if let Some(new_item) = replace_with.next() {
                 unsafe {
                     core::ptr::write(place, new_item);
-                    vec.set_len(vec.len() + 1);
+                    vec.len.add(1);
                 }
             } else {
                 return false;
@@ -1137,7 +1137,7 @@ impl<T, const N: usize> SmallVec<T, N> {
             debug_assert!(len < self.capacity());
             // SAFETY: we have wrote the value to the address already
             unsafe {
-                self.len.increment();
+                self.len.add(1);
             }
         }
 
@@ -1157,7 +1157,7 @@ impl<T, const N: usize> SmallVec<T, N> {
         // SAFETY: new_len < len since len is non-zero and
         // we are returning ownership of the current value.
         unsafe {
-            self.len.decrement();
+            self.len.sub(1);
         }
         // SAFETY: this element was initialized and we just gave up ownership of
         // it, so we can give it away
@@ -1188,7 +1188,7 @@ impl<T, const N: usize> SmallVec<T, N> {
         // SAFETY: we have a mutable reference to each vector and each uniquely
         // owns its memory. so the ranges can't overlap
         unsafe { copy_nonoverlapping(other.as_ptr(), ptr, other_len) };
-        unsafe { self.set_len(total_len) }
+        unsafe { self.len.add(other_len) }
     }
 
     #[inline]
@@ -1388,7 +1388,7 @@ impl<T, const N: usize> SmallVec<T, N> {
             let value = core::ptr::read(self.as_ptr().add(index));
             let base_ptr = self.as_mut_ptr();
             core::ptr::copy(base_ptr.add(new_len), base_ptr.add(index), 1);
-            self.set_len(new_len);
+            self.len.sub(1);
             value
         }
     }
@@ -1423,7 +1423,7 @@ impl<T, const N: usize> SmallVec<T, N> {
         let new_len = len - 1;
         unsafe {
             // SAFETY: new_len < len
-            self.set_len(new_len);
+            self.len.sub(1);
             let ptr = self.as_mut_ptr();
             let ith = ptr.add(index);
             // This item is initialized since index < len
@@ -1478,7 +1478,7 @@ impl<T, const N: usize> SmallVec<T, N> {
             debug_assert!(len < self.capacity());
             // SAFETY: we have wrote the value to the address already
             unsafe {
-                self.len.increment();
+                self.len.add(1);
             }
         }
 
@@ -1908,7 +1908,7 @@ impl<T: Clone, const N: usize> SmallVec<T, N> {
         unsafe {
             let dst = self.as_mut_ptr().add(l);
             copy_nonoverlapping(src, dst, len);
-            self.set_len(l + len);
+            self.len.add(len);
         }
     }
 
@@ -1932,7 +1932,7 @@ impl<T: Clone, const N: usize> SmallVec<T, N> {
             let l = self.len();
             let ptr = self.as_mut_ptr();
             copy_nonoverlapping(ptr.add(start), ptr.add(l), len);
-            self.set_len(l + len);
+            self.len.add(len);
         }
     }
 
@@ -1953,7 +1953,7 @@ impl<T: Clone, const N: usize> SmallVec<T, N> {
             copy_nonoverlapping(other.as_ptr(), ith_ptr, len);
 
             // SAFETY: all the elements are initialized
-            self.set_len(l + len);
+            self.len.add(len);
         }
     }
 
@@ -2238,7 +2238,7 @@ impl<T, const N: usize> SmallVec<T, N> {
 
         // SAFETY: The elements were initialized in the loop above.
         unsafe {
-            self.set_len(old_len + len);
+            self.len.add(len);
         }
     }
 
@@ -2496,7 +2496,7 @@ unsafe impl<const N: usize> BufMut for SmallVec<u8, N> {
         }
 
         // Addition will not overflow since the sum is at most the capacity.
-        unsafe { self.set_len(len + cnt) };
+        unsafe { self.len.add(cnt) };
     }
 
     #[inline]

src/taggedlen.rs

@@ -52,36 +52,20 @@ impl<T> TaggedLen<T> {
         if Self::IS_ZST { self.0 } else { self.0 >> 1 }
     }
 
-    /// Returns the same tag with the length increased by one.
-    ///
-    /// This increases the length without rereading the `on heap` flag.
-    ///
-    /// # Safety
-    ///
-    /// The caller must ensure that after incrementing, the length would still
-    /// be less than [`isize::MAX`] in bytes. For non-ZSTs this means the
-    /// length must be less than `isize::MAX - 1` before the call.
     #[inline]
-    pub const unsafe fn increment(&mut self) {
+    pub const unsafe fn add(&mut self, n: usize) {
         self.0 += if Self::IS_ZST {
-            1
+            n
         } else {
-            debug_assert!(self.value() + 1 < isize::MAX as usize);
-            0b10
+            debug_assert!(self.value() + n < isize::MAX as usize);
+            n << 1
         }
     }
 
-    /// Returns the same tag with the length decreased by one.
-    ///
-    /// This decreases the length without rereading the `on heap` flag.
-    ///
-    /// # Safety
-    ///
-    /// The caller must ensure that the length is greater than zero before the
-    /// call.
     #[inline]
-    pub const unsafe fn decrement(&mut self) {
-        debug_assert!(self.value() > 0);
-        self.0 -= if Self::IS_ZST { 1 } else { 0b10 };
+    pub const unsafe fn sub(&mut self, n: usize) {
+        debug_assert!(self.value() >= n);
+
+        self.0 -= if Self::IS_ZST { n } else { n << 1 };
     }
 }