str_stack/
lib.rs

1//! A string allocation library. This is primarily useful when you want to allocate a bunch of
2//! small strings, use them, and then destroy them all together.
3//!
4//! ## Example
5//!
6//! ```
7//! use str_stack::StrStack;
8//!
9//! let mut stack = StrStack::new();
10//! let first = stack.push("one");
11//! let second = stack.push("two");
12//! let third = stack.push("three");
13//!
14//! assert_eq!(&stack[first], "one");
15//! assert_eq!(&stack[second], "two");
16//! assert_eq!(&stack[third], "three");
17//! ```
18//!
19use std::ops::Index;
20use std::fmt::{self, Write};
21use std::io::{self, Read};
22use std::iter::FromIterator;
23use std::slice;
24
25#[derive(Clone, Default)]
26pub struct StrStack {
27    data: String,
28    ends: Vec<usize>,
29}
30
31impl Index<usize> for StrStack {
32    type Output = str;
33    #[inline]
34    fn index(&self, index: usize) -> &str {
35        unsafe {
36            assert!(index < self.len(), "index out of bounds");
37            self.get_unchecked(index)
38        }
39    }
40}
41
42#[derive(Clone)]
43pub struct Iter<'a> {
44    data: &'a str,
45    ends: &'a [usize],
46}
47
48impl fmt::Debug for StrStack {
49    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
50        f.debug_list().entries(self.iter()).finish()
51    }
52}
53
54impl<'a> Iterator for Iter<'a> {
55    type Item = &'a str;
56    #[inline]
57    fn next(&mut self) -> Option<&'a str> {
58        unsafe {
59            let len = self.ends.len();
60            if len == 1 {
61                None
62            } else {
63                let start = *self.ends.get_unchecked(0);
64                let end = *self.ends.get_unchecked(1);
65                self.ends = slice::from_raw_parts(self.ends.as_ptr().offset(1), len - 1);
66                Some(self.data.slice_unchecked(start, end))
67            }
68        }
69    }
70
71    fn count(self) -> usize {
72        self.size_hint().0
73    }
74
75    fn last(mut self) -> Option<&'a str> {
76        self.next_back()
77    }
78
79    #[inline]
80    fn size_hint(&self) -> (usize, Option<usize>) {
81        let len = self.ends.len() - 1;
82        (len, Some(len))
83    }
84}
85
86impl<'a> ExactSizeIterator for Iter<'a> {}
87
88impl<'a> DoubleEndedIterator for Iter<'a> {
89    #[inline]
90    fn next_back(&mut self) -> Option<&'a str> {
91        unsafe {
92            let len = self.ends.len();
93            if len == 1 {
94                None
95            } else {
96                let start = *self.ends.get_unchecked(len-2);
97                let end = *self.ends.get_unchecked(len-1);
98                self.ends = slice::from_raw_parts(self.ends.as_ptr(), len - 1);
99                Some(self.data.slice_unchecked(start, end))
100            }
101        }
102    }
103}
104
105impl<'a> IntoIterator for &'a StrStack {
106    type IntoIter = Iter<'a>;
107    type Item = &'a str;
108    #[inline]
109    fn into_iter(self) -> Iter<'a> {
110        self.iter()
111    }
112}
113
114impl StrStack {
115    /// Create a new StrStack.
116    #[inline]
117    pub fn new() -> StrStack {
118        StrStack::with_capacity(0, 0)
119    }
120
121    /// Create a new StrStack with the given capacity.
122    ///
123    /// You will be able to push `bytes` bytes and create `strings` strings before reallocating.
124    #[inline]
125    pub fn with_capacity(bytes: usize, strings: usize) -> StrStack {
126        let mut stack = StrStack {
127            data: String::with_capacity(bytes),
128            ends: Vec::with_capacity(strings+1)
129        };
130        // Yes, I know I don't need this. However, putting this here avoids checks later which
131        // makes this much faster.
132        stack.ends.push(0);
133        stack
134    }
135
136    /// Push a string onto the string stack.
137    ///
138    /// This returns the index of the string on the stack.
139    #[inline]
140    pub fn push(&mut self, s: &str) -> usize {
141        self.data.push_str(s);
142        self.ends.push(self.data.len());
143        self.len() - 1
144    }
145
146    /// Iterate over the strings on the stack.
147    #[inline]
148    pub fn iter(&self) -> Iter {
149        Iter {
150            data: &self.data,
151            ends: &self.ends,
152        }
153    }
154
155    /// Remove the top string from the stack.
156    ///
157    /// Returns true iff a string was removed.
158    #[inline]
159    pub fn pop(&mut self) -> bool {
160        if self.ends.len() <= 1 {
161            false
162        } else {
163            self.ends.pop();
164            self.data.truncate(*self.ends.last().unwrap());
165            true
166        }
167    }
168
169    /// Clear the stack.
170    #[inline]
171    pub fn clear(&mut self) {
172        self.ends.truncate(1);
173        self.data.clear();
174    }
175
176    /// Returns the number of strings on the stack.
177    #[inline]
178    pub fn len(&self) -> usize {
179        self.ends.len() - 1
180    }
181
182    /// Truncate the stack to `len` strings.
183    #[inline]
184    pub fn truncate(&mut self, len: usize) {
185        self.ends.truncate(len.saturating_add(1));
186        self.data.truncate(*self.ends.last().unwrap());
187    }
188
189    /// Read from `source` into the string stack.
190    ///
191    /// Returns the index of the new string or an IO Error.
192    pub fn consume<R: io::Read>(&mut self, mut source: R) -> io::Result<usize> {
193        match source.read_to_string(&mut self.data) {
194            Ok(_) => {
195                self.ends.push(self.data.len());
196                Ok(self.len() - 1)
197            },
198            Err(e) => Err(e),
199        }
200    }
201
202    /// Returns a writer helper for this string stack.
203    ///
204    /// This is useful for building a string in-place on the string-stack.
205    ///
206    /// Example:
207    ///
208    /// ```
209    /// use std::fmt::Write;
210    /// use str_stack::StrStack;
211    ///
212    /// let mut s = StrStack::new();
213    /// let index = {
214    ///     let mut writer = s.writer();
215    ///     writer.write_str("Hello");
216    ///     writer.write_char(' ');
217    ///     writer.write_str("World");
218    ///     writer.write_char('!');
219    ///     writer.finish()
220    /// };
221    /// assert_eq!(&s[index], "Hello World!");
222    /// ```
223    #[inline]
224    pub fn writer(&mut self) -> Writer {
225        Writer(self)
226    }
227
228    /// Allows calling the write! macro directly on the string stack:
229    ///
230    /// Example:
231    ///
232    /// ```
233    /// use std::fmt::Write;
234    /// use str_stack::StrStack;
235    ///
236    /// let mut s = StrStack::new();
237    /// let index = write!(&mut s, "Hello {}!", "World");
238    /// assert_eq!(&s[index], "Hello World!");
239    /// ```
240    #[inline]
241    pub fn write_fmt(&mut self, args: fmt::Arguments) -> usize {
242        let mut writer = self.writer();
243        let _ = writer.write_fmt(args);
244        writer.finish()
245    }
246
247    #[inline]
248    pub unsafe fn get_unchecked(&self, index: usize) -> &str {
249        let start = *self.ends.get_unchecked(index);
250        let end = *self.ends.get_unchecked(index+1);
251        self.data.slice_unchecked(start, end)
252    }
253}
254
255impl<S> Extend<S> for StrStack where S: AsRef<str> {
256    fn extend<T>(&mut self, iterator: T) where T: IntoIterator<Item=S> {
257        let iterator = iterator.into_iter();
258        let (min, _) = iterator.size_hint();
259        self.ends.reserve(min);
260        for v in iterator {
261            self.push(v.as_ref());
262        }
263    }
264}
265
266impl<S> FromIterator<S> for StrStack where S: AsRef<str> {
267    fn from_iter<T>(iterator: T) -> Self where T: IntoIterator<Item=S> {
268        let mut stack = StrStack::new();
269        stack.extend(iterator);
270        stack
271    }
272}
273
274pub struct Writer<'a>(&'a mut StrStack);
275
276impl<'a> Writer<'a> {
277    /// Finish pushing the string onto the stack and return its index.
278    #[inline]
279    pub fn finish(self) -> usize {
280        // We push on drop.
281        self.0.len()
282    }
283}
284
285impl<'a> fmt::Write for Writer<'a> {
286    #[inline]
287    fn write_str(&mut self, s: &str) -> fmt::Result {
288        self.0.data.push_str(s);
289        Ok(())
290    }
291    #[inline]
292    fn write_char(&mut self, c: char) -> fmt::Result {
293        self.0.data.push(c);
294        Ok(())
295    }
296}
297
298impl<'a> Drop for Writer<'a> {
299    fn drop(&mut self) {
300        self.0.ends.push(self.0.data.len());
301    }
302}
303
304#[test]
305fn test_basic() {
306    let mut stack = StrStack::new();
307    let first = stack.push("one");
308    let second = stack.push("two");
309    let third = stack.push("three");
310
311    assert_eq!(&stack[first], "one");
312    assert_eq!(&stack[second], "two");
313    assert_eq!(&stack[third], "three");
314
315    assert_eq!(stack.len(), 3);
316
317    assert!(stack.pop());
318
319    assert_eq!(stack.len(), 2);
320
321    assert_eq!(&stack[first], "one");
322    assert_eq!(&stack[second], "two");
323
324    assert!(stack.pop());
325    assert!(stack.pop());
326
327    assert_eq!(stack.len(), 0);
328    assert!(!stack.pop());
329}
330
331#[test]
332fn test_consume() {
333    let mut stack = StrStack::new();
334    let idx = stack.consume("testing".as_bytes()).unwrap();
335    assert_eq!(&stack[idx], "testing");
336}
337
338#[test]
339fn test_writer() {
340    let mut stack = StrStack::new();
341    let first = {
342        let mut w = stack.writer();
343        write!(w, "{}", "first ").unwrap();
344        write!(w, "{}", "second").unwrap();
345        w.finish()
346    };
347
348    let second = {
349        let mut w = stack.writer();
350        write!(w, "{}", "third ").unwrap();
351        write!(w, "{}", "fourth").unwrap();
352        w.finish()
353    };
354    assert_eq!(&stack[first], "first second");
355    assert_eq!(&stack[second], "third fourth");
356}
357
358#[test]
359fn test_iter() {
360    let mut stack = StrStack::new();
361    stack.push("one");
362    stack.push("two");
363    stack.push("three");
364
365    let v1: Vec<_> = stack.iter().collect();
366    let v2: Vec<_> = stack.iter().rev().collect();
367
368    assert_eq!(&v1[..], &["one", "two", "three"]);
369    assert_eq!(&v2[..], &["three", "two", "one"]);
370}