diff --git a/Cargo.toml b/Cargo.toml index 69b3aac..5630031 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ harness = false [features] default = ["std"] std = [] +unstable = [] test = ["std", "arbitrary", "arbitrary/derive"] [dependencies] diff --git a/src/boxed.rs b/src/boxed.rs index 081ffee..25ecaa5 100644 --- a/src/boxed.rs +++ b/src/boxed.rs @@ -32,11 +32,16 @@ pub(crate) struct BoxedString { /// /// Returns `true` if aligned to an odd address, `false` if even. The sense of /// the boolean is "does this look like an InlineString? true/false" +#[inline] fn check_alignment(ptr: *const u8) -> bool { ptr.align_offset(2) > 0 } impl GenericString for BoxedString { + fn cap(&self) -> usize { + self.cap + } + fn set_size(&mut self, size: usize) { self.len = size; debug_assert!(self.len <= self.cap); @@ -53,6 +58,7 @@ impl GenericString for BoxedString { impl BoxedString { const MINIMAL_CAPACITY: usize = MAX_INLINE * 2; + #[inline] pub(crate) fn check_alignment(this: &Self) -> bool { check_alignment(this.ptr.as_ptr()) } diff --git a/src/inline.rs b/src/inline.rs index 2815137..3c2eb32 100644 --- a/src/inline.rs +++ b/src/inline.rs @@ -56,6 +56,10 @@ impl DerefMut for InlineString { } impl GenericString for InlineString { + fn cap(&self) -> usize { + MAX_INLINE + } + fn set_size(&mut self, size: usize) { self.marker.set_data(size as u8); } diff --git a/src/lib.rs b/src/lib.rs index 33e0023..0d03fca 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -88,6 +88,7 @@ //! | Feature | Description | //! | ------- | ----------- | //! | [`arbitrary`](https://crates.io/crates/arbitrary) | [`Arbitrary`][Arbitrary] implementation for [`SmartString`]. | +//! | `unstable` | Features only available with a nightly Rust compiler. Currently this is only the [`extend_reserve`][core::iter::traits::collect::Extend::extend_reserve] implementation for [`SmartString`]. | //! | [`proptest`](https://crates.io/crates/proptest) | A strategy for generating [`SmartString`]s from a regular expression. | //! | [`serde`](https://crates.io/crates/serde) | [`Serialize`][Serialize] and [`Deserialize`][Deserialize] implementations for [`SmartString`]. | //! @@ -102,6 +103,7 @@ #![warn(unreachable_pub, missing_debug_implementations, missing_docs)] #![cfg_attr(not(feature = "std"), no_std)] #![cfg_attr(needs_allocator_feature, feature(allocator_api))] +#![cfg_attr(feature = "unstable", feature(extend_one))] extern crate alloc; @@ -421,6 +423,17 @@ impl SmartString { } } + /// Ensures that the string has a capacity of at least the given number of bytes. + /// This function will reallocate the string (and therefore unbox a boxed string) + /// in order to fit the new capacity. + /// + /// Note that if the string's capacity is already large enough, this function does nothing. + /// This also applies to inline strings, when `capacity` is less than or equal to + /// [`MAX_INLINE`]. + pub fn ensure_capacity(&mut self, target_cap: usize) { + string_op_grow!(ops::EnsureCapacity, self, target_cap) + } + /// Push a character to the end of the string. pub fn push(&mut self, ch: char) { string_op_grow!(ops::Push, self, ch) @@ -710,6 +723,11 @@ impl<'a, Mode: SmartStringMode> Extend<&'a str> for SmartString { self.push_str(item); } } + + #[cfg(feature = "unstable")] + fn extend_one(&mut self, item: &'a str) { + self.push_str(item) + } } impl<'a, Mode: SmartStringMode> Extend<&'a char> for SmartString { @@ -718,6 +736,16 @@ impl<'a, Mode: SmartStringMode> Extend<&'a char> for SmartString { self.push(*item); } } + + #[cfg(feature = "unstable")] + fn extend_one(&mut self, item: &'a char) { + self.push(*item) + } + + #[cfg(feature = "unstable")] + fn extend_reserve(&mut self, additional: usize) { + self.ensure_capacity(self.capacity() + additional) + } } impl Extend for SmartString { @@ -726,6 +754,16 @@ impl Extend for SmartString { self.push(item); } } + + #[cfg(feature = "unstable")] + fn extend_one(&mut self, item: char) { + self.push(item) + } + + #[cfg(feature = "unstable")] + fn extend_reserve(&mut self, additional: usize) { + self.ensure_capacity(self.capacity() + additional) + } } impl Extend> for SmartString { @@ -742,6 +780,11 @@ impl Extend for SmartString { self.push_str(&item); } } + + #[cfg(feature = "unstable")] + fn extend_one(&mut self, item: String) { + self.push_str(&item) + } } impl<'a, Mode: SmartStringMode + 'a> Extend<&'a SmartString> for SmartString { @@ -750,6 +793,11 @@ impl<'a, Mode: SmartStringMode + 'a> Extend<&'a SmartString> for SmartStri self.push_str(item); } } + + #[cfg(feature = "unstable")] + fn extend_one(&mut self, item: &'a SmartString) { + self.push_str(item) + } } impl<'a, Mode: SmartStringMode> Extend<&'a String> for SmartString { @@ -758,6 +806,11 @@ impl<'a, Mode: SmartStringMode> Extend<&'a String> for SmartString { self.push_str(item); } } + + #[cfg(feature = "unstable")] + fn extend_one(&mut self, item: &'a String) { + self.push_str(item) + } } impl Add for SmartString { diff --git a/src/ops.rs b/src/ops.rs index f662264..9fb53f0 100644 --- a/src/ops.rs +++ b/src/ops.rs @@ -18,6 +18,7 @@ use core::{ }; pub(crate) trait GenericString: Deref + DerefMut { + fn cap(&self) -> usize; fn set_size(&mut self, size: usize); fn as_mut_capacity_slice(&mut self) -> &mut [u8]; } @@ -135,6 +136,15 @@ impl Truncate { } } +pub(crate) struct EnsureCapacity; +impl EnsureCapacity { + pub(crate) fn cap(this: &S, target_cap: usize) -> usize { + this.cap().max(target_cap) + } + + pub(crate) fn op(_this: &mut S, _target_cap: usize) {} +} + pub(crate) struct Pop; impl Pop { pub(crate) fn op(this: &mut S) -> Option { diff --git a/src/test.rs b/src/test.rs index 8abba06..9b6c075 100644 --- a/src/test.rs +++ b/src/test.rs @@ -552,6 +552,42 @@ mod tests { assert_eq!(Discriminant::Inline, s.discriminant()); } + #[test] + fn capacity_modification() { + use crate::inline::InlineString; + + const SHORT_TEXT: &'static str = "short enough"; + static_assertions::const_assert!(SHORT_TEXT.len() < MAX_INLINE - 1); + let mut string = + SmartString::::from_inline(InlineString::try_from(SHORT_TEXT).unwrap()); + + string.ensure_capacity(string.len() + 1); + assert!(string.capacity() >= SHORT_TEXT.len()); + assert!(string.is_inline()); + + string.ensure_capacity(MAX_INLINE + 1); + assert!(!string.is_inline()); + + const LARGE_CAPACITY: usize = MAX_INLINE * 40; + string.ensure_capacity(LARGE_CAPACITY); + assert!(string.capacity() >= LARGE_CAPACITY); + + string.ensure_capacity(LARGE_CAPACITY / 2); + assert!(string.capacity() >= LARGE_CAPACITY); + + // Check memory corruptions or size mishandling + assert!(string.len() == SHORT_TEXT.len()); + assert!(string.as_str() == SHORT_TEXT); + + for _ in 0..20 { + string = string + "1234567890"; + } + // Read the entire string to check memory corruptions + format!("{}", string); + assert!(string.capacity() >= LARGE_CAPACITY); + assert!(string.len() == SHORT_TEXT.len() + 10 * 20); + } + #[test] fn from_string() { let std_s =