feat: add truncate() methods

Author: ryota2357

src/lib.rs

@@ -769,6 +769,49 @@ impl LeanString {
         self.0.insert_str(idx, string)
     }
 
+    /// Shortens a [`LeanString`] to the specified length.
+    ///
+    /// If `new_len` is greater than or equal to the string's current length, this has no effect.
+    ///
+    /// # Panics
+    ///
+    /// Panics if **any** of the following conditions is met:
+    ///
+    /// 1. `new_len` does not lie on a [`char`] boundary.
+    /// 2. The system is out-of-memory when cloning the [`LeanString`].
+    ///
+    /// For 2, If you want to handle such a problem manually, use [`LeanString::try_truncate()`].
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use lean_string::LeanString;
+    /// let mut s = LeanString::from("hello");
+    /// s.truncate(2);
+    /// assert_eq!(s, "he");
+    ///
+    /// // Truncating to a larger length does nothing:
+    /// s.truncate(10);
+    /// assert_eq!(s, "he");
+    /// ```
+    #[inline]
+    pub fn truncate(&mut self, new_len: usize) {
+        self.try_truncate(new_len).unwrap_with_msg()
+    }
+
+    /// Fallible version of [`LeanString::truncate()`].
+    ///
+    /// This method won't panic if the system is out-of-memory, but return an [`ReserveError`].
+    /// Otherwise it behaves the same as [`LeanString::truncate()`].
+    ///
+    /// # Panics
+    ///
+    /// This method still panics if `new_len` does not lie on a [`char`] boundary.
+    #[inline]
+    pub fn try_truncate(&mut self, new_len: usize) -> Result<(), ReserveError> {
+        self.0.truncate(new_len)
+    }
+
     /// Reduces the length of the [`LeanString`] to zero.
     ///
     /// If the [`LeanString`] is unique, this method will not change the capacity.

src/repr.rs

@@ -353,46 +353,10 @@ impl Repr {
         // SAFETY: We know this is a valid length which falls on a char boundary
         let new_len = self.len() - ch.len_utf8();
 
-        if self.is_heap_buffer() {
-            // SAFETY: We just checked that `self` is HeapBuffer
-            let heap = unsafe { self.as_heap_buffer_mut() };
-
-            if !heap.is_len_on_heap() {
-                // Since len is inlined and we don't modify the buffer by popping a char, it is ok
-                // to just set the new length.
-                // SAFETY: `new_len <= len <= capacity`
-                unsafe { heap.set_len(new_len) };
-            } else {
-                // See `reverse` method for the explanation of the ordering.
-                if heap.reference_count().fetch_sub(1, Release) == 1 {
-                    // `heap` is unique, we can set the new length in place.
-
-                    // See `reverse` method for the explanation of the ordering.
-                    heap.reference_count().fetch_add(1, Acquire);
-
-                    // SAFETY: `heap` is unique, we can reallocate in place.
-                    unsafe { heap.set_len(new_len) };
-                } else {
-                    // SAFETY: `ptr` is valid for `len` bytes, and `HeapBuffer` contains valid UTF-8.
-                    let str = unsafe {
-                        let ptr = self.0 as *mut u8;
-                        let slice = slice::from_raw_parts_mut(ptr, new_len);
-                        str::from_utf8_unchecked_mut(slice)
-                    };
-                    *self = Repr::from_str(str)?;
-                }
-            }
-        } else if self.is_static_buffer() {
-            // SAFETY:
-            // - We just checked that `self` is StaticBuffer
-            // - `new_len <= len <= capacity`
-            unsafe { self.as_static_buffer_mut().set_len(new_len) };
-        } else {
-            // SAFETY:
-            // - The number of types of buffer is 3, and the remaining is InlineBuffer.
-            // - From `#Safety`, `new_len <= MAX_INLINE_SIZE` is true.
-            unsafe { self.as_inline_buffer_mut().set_len(new_len) };
-        }
+        // SAFETY:
+        // - `new_len` is less than `len()` because we calculated it from `len() - ch.len_utf8()`.
+        // - `new_len` is a valid char boundary because `ch` is a valid char.
+        unsafe { self.truncate_unchecked(new_len) }?;
 
         Ok(Some(ch))
     }
@@ -516,6 +480,74 @@ impl Repr {
         Ok(())
     }
 
+    #[inline]
+    pub(crate) fn truncate(&mut self, new_len: usize) -> Result<(), ReserveError> {
+        if new_len >= self.len() {
+            return Ok(());
+        }
+
+        let str = self.as_str();
+        assert!(
+            str.is_char_boundary(new_len),
+            "index is not a char boundary or out of bounds (index: {new_len})",
+        );
+
+        // SAFETY: We just checked that `new_len < len()` and `new_len` is a valid char
+        unsafe { self.truncate_unchecked(new_len) }
+    }
+
+    /// # Safety
+    ///
+    /// - `new_len` must be less than or equal to `len()`
+    /// - `new_len` must be a valid char boundary.
+    unsafe fn truncate_unchecked(&mut self, new_len: usize) -> Result<(), ReserveError> {
+        debug_assert!(new_len <= self.len());
+        debug_assert!(self.as_str().is_char_boundary(new_len));
+
+        if self.is_heap_buffer() {
+            // SAFETY: We just checked that `self` is HeapBuffer
+            let heap = unsafe { self.as_heap_buffer_mut() };
+
+            if !heap.is_len_on_heap() {
+                // Since len is inlined and we don't modify the buffer by popping a char, it is ok
+                // to just set the new length.
+                // SAFETY: `new_len <= len <= capacity`
+                unsafe { heap.set_len(new_len) };
+            } else {
+                // See `reverse` method for the explanation of the ordering.
+                if heap.reference_count().fetch_sub(1, Release) == 1 {
+                    // `heap` is unique, we can set the new length in place.
+
+                    // See `reverse` method for the explanation of the ordering.
+                    heap.reference_count().fetch_add(1, Acquire);
+
+                    // SAFETY: `heap` is unique, we can reallocate in place.
+                    unsafe { heap.set_len(new_len) };
+                } else {
+                    // SAFETY: `ptr` is valid for `len` bytes, and `HeapBuffer` contains valid UTF-8.
+                    let str = unsafe {
+                        let ptr = self.0 as *mut u8;
+                        let slice = slice::from_raw_parts_mut(ptr, new_len);
+                        str::from_utf8_unchecked_mut(slice)
+                    };
+                    *self = Repr::from_str(str)?;
+                }
+            }
+        } else if self.is_static_buffer() {
+            // SAFETY:
+            // - We just checked that `self` is StaticBuffer
+            // - `new_len <= len <= capacity`
+            unsafe { self.as_static_buffer_mut().set_len(new_len) };
+        } else {
+            // SAFETY:
+            // - The number of types of buffer is 3, and the remaining is InlineBuffer.
+            // - From `#Safety`, `new_len <= MAX_INLINE_SIZE` is true.
+            unsafe { self.as_inline_buffer_mut().set_len(new_len) };
+        }
+
+        Ok(())
+    }
+
     #[inline]
     pub(crate) fn is_unique(&self) -> bool {
         if self.is_heap_buffer() {

tests/alloc_string.rs

@@ -341,37 +341,37 @@ fn test_pop() {
 //     assert_eq!(nihon.capacity(), orig_capacity);
 // }
 
-// #[test]
-// fn test_str_truncate() {
-//     let mut s = String::from("12345");
-//     s.truncate(5);
-//     assert_eq!(s, "12345");
-//     s.truncate(3);
-//     assert_eq!(s, "123");
-//     s.truncate(0);
-//     assert_eq!(s, "");
-//
-//     let mut s = String::from("12345");
-//     let p = s.as_ptr();
-//     s.truncate(3);
-//     s.push_str("6");
-//     let p_ = s.as_ptr();
-//     assert_eq!(p_, p);
-// }
+#[test]
+fn test_str_truncate() {
+    let mut s = LeanString::from("12345");
+    s.truncate(5);
+    assert_eq!(s, "12345");
+    s.truncate(3);
+    assert_eq!(s, "123");
+    s.truncate(0);
+    assert_eq!(s, "");
 
-// #[test]
-// fn test_str_truncate_invalid_len() {
-//     let mut s = String::from("12345");
-//     s.truncate(6);
-//     assert_eq!(s, "12345");
-// }
+    let mut s = LeanString::from("12345");
+    let p = s.as_ptr();
+    s.truncate(3);
+    s.push_str("6");
+    let p_ = s.as_ptr();
+    assert_eq!(p_, p);
+}
 
-// #[test]
-// #[should_panic]
-// fn test_str_truncate_split_codepoint() {
-//     let mut s = String::from("\u{FC}"); // ü
-//     s.truncate(1);
-// }
+#[test]
+fn test_str_truncate_invalid_len() {
+    let mut s = LeanString::from("12345");
+    s.truncate(6);
+    assert_eq!(s, "12345");
+}
+
+#[test]
+#[should_panic(expected = "index is not a char boundary or out of bounds (index: 1)")]
+fn test_str_truncate_split_codepoint() {
+    let mut s = LeanString::from("\u{FC}"); // ü
+    s.truncate(1);
+}
 
 #[test]
 fn test_str_clear() {

tests/handmade.rs

@@ -202,6 +202,10 @@ fn pop_share_buffer() {
 
     // buffer is shared
     assert_eq!(s.as_ptr(), s2.as_ptr());
+
+    // modify makes a new buffer
+    s2.push('0');
+    assert_ne!(s.as_ptr(), s2.as_ptr());
 }
 
 #[test]
@@ -368,6 +372,76 @@ fn insert_fail() {
     s.insert(7, 'a');
 }
 
+#[test]
+fn truncate_keep_capacity() {
+    let mut inline = LeanString::from("abcdef");
+
+    assert!(!inline.is_heap_allocated());
+    inline.truncate(3);
+    assert_eq!(inline, "abc");
+    assert_eq!(inline.len(), 3);
+    assert!(!inline.is_heap_allocated());
+    assert_eq!(inline.capacity(), INLINE_LIMIT);
+
+    let mut heap = LeanString::from("a".repeat(INLINE_LIMIT + 10).as_str());
+    let original_capacity = heap.capacity();
+    assert!(heap.is_heap_allocated());
+
+    heap.truncate(INLINE_LIMIT + 1);
+    assert_eq!(heap, "a".repeat(INLINE_LIMIT + 1));
+    assert_eq!(heap.len(), INLINE_LIMIT + 1);
+    assert!(heap.is_heap_allocated());
+    assert_eq!(heap.capacity(), original_capacity);
+
+    heap.truncate(1);
+    assert_eq!(heap, "a");
+    assert_eq!(heap.len(), 1);
+    assert!(heap.is_heap_allocated());
+    assert_eq!(heap.capacity(), original_capacity);
+}
+
+#[test]
+fn truncate_from_static() {
+    let mut static_ = LeanString::from_static_str("abcdefghijklmnopqrstuvwxyz");
+    assert_eq!(static_.len(), 26);
+    assert!(!static_.is_heap_allocated());
+
+    static_.truncate(20);
+    assert_eq!(static_, "abcdefghijklmnopqrst");
+    assert_eq!(static_.len(), 20);
+    assert_eq!(static_.capacity(), 20);
+    assert!(!static_.is_heap_allocated());
+}
+
+#[test]
+fn truncate_share_buffer() {
+    // s is inlined
+    let mut s = LeanString::from("abcdefgh");
+    assert_eq!(s.len(), 8);
+
+    let mut s1 = s.clone();
+    s1.truncate(4);
+    assert_eq!(s1, "abcd");
+    assert_eq!(s1.len(), 4);
+
+    // s is not changed
+    assert_eq!(s, "abcdefgh");
+
+    // s into heap
+    s.push_str("ijklmnopqrstuvwxyz");
+    assert_eq!(s.len(), 26);
+    assert!(s.is_heap_allocated());
+
+    // buffer is shared
+    let mut s2 = s.clone();
+    s2.truncate(20);
+    assert_eq!(s.as_ptr(), s2.as_ptr());
+
+    // modify makes a new buffer
+    s2.push('0');
+    assert_ne!(s.as_ptr(), s2.as_ptr());
+}
+
 #[test]
 fn convert_static_to_inline_with_reserve() {
     let s: &'static str = "1234567890ABCDEFGHIJ";