diff --git a/src/lib.rs b/src/lib.rs index fd3e14a..0f07a95 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,6 +5,7 @@ use core::ops::{Deref, DerefMut}; use core::mem::replace; +use core::mem::needs_drop; extern crate alloc; use alloc::vec::Vec; @@ -113,4 +114,22 @@ impl TinyVec { } } } + + pub fn truncate(&mut self, new_len: usize) { + match &mut self.0 { + Payload::Inline { len, data } => { + if needs_drop::() { + while *len > new_len { + replace(&mut data[*len - 1], T::default()); + *len -= 1; + } + } else { + *len = (*len).min(new_len); + } + } + Payload::Heap(ref mut vec) => { + vec.truncate(new_len) + } + } + } } diff --git a/tests/basic.rs b/tests/basic.rs index e36b7e5..4818699 100644 --- a/tests/basic.rs +++ b/tests/basic.rs @@ -1,10 +1,11 @@ - -extern crate std; +#![allow(bad_style)] use tinyvec::*; +use core::ops::Deref; + #[test] -fn test_push_pop() { +fn TinyVec_push_pop() { let mut tv = TinyVec::new(); assert_eq!(tv.pop(), None); @@ -31,3 +32,32 @@ fn test_push_pop() { assert_eq!(tv.pop(), None); assert_eq!(tv.pop(), None); } + +#[test] +fn TinyVec_truncate() { + // Truncating a five element vector to two elements + let mut tv = TinyVec::new(); + tv.push(1_i32); + tv.push(2); + tv.push(3); + tv.push(4); + tv.push(5); + tv.truncate(2); + assert_eq!(tv.deref(), [1, 2]); + + // No truncation occurs when len is greater than the vector's current length + let mut tv = TinyVec::new(); + tv.push(1_i32); + tv.push(2); + tv.push(3); + tv.truncate(8); + assert_eq!(tv.deref(), [1, 2, 3]); + + // Truncating when len == 0 is equivalent to calling the clear method. + let mut tv = TinyVec::new(); + tv.push(1_i32); + tv.push(2); + tv.push(3); + tv.truncate(0); + assert_eq!(tv.deref(), []); +}