Skip to main content

nt_string/unicode_string/
string.rs

1// Copyright 2023-2026 Colin Finck <colin@reactos.org>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use ::alloc::alloc::{self, Layout};
5use ::alloc::string::String;
6use widestring::{U16CStr, U16Str};
7
8use core::cmp::Ordering;
9use core::iter::once;
10use core::ops::{Add, AddAssign, Deref, DerefMut};
11use core::{fmt, mem, ptr};
12
13use crate::error::{NtStringError, Result};
14use crate::helpers::RawNtString;
15use crate::traits::TryExtend;
16
17use super::{impl_eq, impl_partial_cmp, NtUnicodeStr, NtUnicodeStrMut};
18
19/// An allocated, owned, and growable variant of `UNICODE_STRING` (equivalent of [`String`]).
20///
21/// See the [module-level documentation](super) for more details.
22#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
23#[derive(Debug)]
24#[repr(transparent)]
25pub struct NtUnicodeString {
26    raw: RawNtString<*mut u16>,
27}
28
29impl NtUnicodeString {
30    /// Creates an empty [`NtUnicodeString`].
31    ///
32    /// This operation won't allocate any buffer.
33    /// Therefore, length and capacity will both be zero.
34    pub fn new() -> Self {
35        Self {
36            raw: RawNtString {
37                length: 0,
38                maximum_length: 0,
39                buffer: ptr::null_mut(),
40            },
41        }
42    }
43
44    /// Returns a mutable [`NtUnicodeStrMut`] reference for this string.
45    pub fn as_unicode_str_mut(&mut self) -> &mut NtUnicodeStrMut<'static> {
46        self.deref_mut()
47    }
48
49    fn layout(&self) -> Layout {
50        Layout::array::<u16>(self.capacity_in_elements()).unwrap()
51    }
52
53    /// Creates an [`NtUnicodeString`] from an existing [`u16`] string buffer without a terminating NUL character.
54    ///
55    /// The string is expected to consist of valid UTF-16 characters.
56    ///
57    /// The given buffer becomes the internal buffer of the [`NtUnicodeString`] and therefore won't be NUL-terminated.
58    /// See the [module-level documentation](super) for the implications of that.
59    ///
60    /// This function has *O*(1) complexity.
61    ///
62    /// If you have a NUL-terminated buffer, either use [`try_from_u16_until_nul`] or convert from a [`U16CStr`]
63    /// using the corresponding [`TryFrom`] implementation.
64    ///
65    /// [`try_from_u16_until_nul`]: Self::try_from_u16_until_nul
66    pub fn try_from_u16(buffer: &[u16]) -> Result<Self> {
67        let unicode_str = NtUnicodeStr::try_from_u16(buffer)?;
68        Ok(Self::from(&unicode_str))
69    }
70
71    /// Creates an [`NtUnicodeString`] from an existing [`u16`] string buffer that contains at least one NUL character.
72    ///
73    /// The string is expected to consist of valid UTF-16 characters.
74    ///
75    /// The string will be terminated at the NUL character.
76    /// An [`NtStringError::NulNotFound`] error is returned if no NUL character could be found.
77    /// As a consequence, this function has *O*(*n*) complexity.
78    ///
79    /// The resulting internal `buffer` of [`NtUnicodeString`] will be NUL-terminated.
80    /// See the [module-level documentation](super) for the implications of that.
81    ///
82    /// Use [`try_from_u16`] if you have a buffer that is not NUL-terminated.
83    /// You can also convert from a NUL-terminated [`U16CStr`] in *O*(1) via the corresponding [`TryFrom`] implementation.
84    ///
85    /// [`try_from_u16`]: Self::try_from_u16
86    pub fn try_from_u16_until_nul(buffer: &[u16]) -> Result<Self> {
87        let unicode_str = NtUnicodeStr::try_from_u16_until_nul(buffer)?;
88        Ok(Self::from(&unicode_str))
89    }
90
91    /// Appends the given [`char`] to the end of this string.
92    ///
93    /// Returns an [`NtStringError::BufferSizeExceedsU16`] error if the resulting string would exceed
94    /// 65532 bytes.
95    /// This is due to the fact that a `UNICODE_STRING` internally uses a 16-bit field to store the length.
96    /// Additionally, this function allocates one more character for NUL termination of the internal
97    /// buffer.
98    /// See the [module-level documentation](super) for the implications of that.
99    ///
100    /// Note that every UTF-16 character consumes 2 or 4 bytes.
101    pub fn try_push(&mut self, c: char) -> Result<()> {
102        // Determine the required additional capacity.
103        //
104        // Add one element for the terminating NUL character
105        // (for applications that mistakenly treat the buffer of a Unicode String as a NUL-terminated string).
106        let encoded_length = c.len_utf16();
107        let additional_elements = encoded_length + 1;
108
109        // Reserve the additional capacity.
110        let additional = (additional_elements * mem::size_of::<u16>()) as u16;
111        self.try_reserve(additional)?;
112
113        // Encode the character as UTF-16 at the end of the buffer.
114        let end_index = self.len_in_elements();
115        self.raw.length += additional;
116
117        let dest_slice = &mut self.as_mut_slice()[end_index..];
118        c.encode_utf16(dest_slice);
119
120        // NUL-terminate it.
121        dest_slice[encoded_length] = 0;
122
123        // Set the final length (without including the terminating NUL character).
124        self.raw.length -= mem::size_of::<u16>() as u16;
125
126        Ok(())
127    }
128
129    /// Appends the given string slice to the end of this string.
130    ///
131    /// Returns an [`NtStringError::BufferSizeExceedsU16`] error if the resulting string would exceed
132    /// 65532 bytes.
133    /// This is due to the fact that a `UNICODE_STRING` internally uses a 16-bit field to store the length.
134    /// Additionally, this function allocates one more character for NUL termination of the internal
135    /// buffer.
136    /// See the [module-level documentation](super) for the implications of that.
137    ///
138    /// Note that every UTF-16 character consumes 2 or 4 bytes.
139    pub fn try_push_str(&mut self, s: &str) -> Result<()> {
140        // Determine the required additional capacity.
141        //
142        // Add one element for the terminating NUL character
143        // (for applications that mistakenly treat the buffer of a Unicode String as a NUL-terminated string).
144        let additional_elements = s
145            .encode_utf16()
146            .count()
147            .checked_add(1)
148            .ok_or(NtStringError::BufferSizeExceedsU16)?;
149
150        // Reserve the additional capacity.
151        let additional_bytes = additional_elements
152            .checked_mul(mem::size_of::<u16>())
153            .ok_or(NtStringError::BufferSizeExceedsU16)?;
154        let additional =
155            u16::try_from(additional_bytes).map_err(|_| NtStringError::BufferSizeExceedsU16)?;
156        self.try_reserve(additional)?;
157
158        // Copy over the string and NUL-terminate it.
159        let end_index = self.len_in_elements();
160        self.raw.length += additional;
161
162        for (string_item, utf16_item) in self.as_mut_slice()[end_index..]
163            .iter_mut()
164            .zip(s.encode_utf16().chain(once(0)))
165        {
166            *string_item = utf16_item;
167        }
168
169        // Set the final length (without the terminating NUL character).
170        self.raw.length -= mem::size_of::<u16>() as u16;
171
172        Ok(())
173    }
174
175    /// Appends the given [`u16`] string buffer (without a terminating NUL character) to the end of this string.
176    ///
177    /// The string is expected to consist of valid UTF-16 characters.
178    ///
179    /// Returns an [`NtStringError::BufferSizeExceedsU16`] error if the resulting string would exceed
180    /// 65532 bytes.
181    /// This is due to the fact that a `UNICODE_STRING` internally uses a 16-bit field to store the length.
182    /// Additionally, this function allocates one more character for NUL termination of the internal
183    /// buffer.
184    /// See the [module-level documentation](super) for the implications of that.
185    ///
186    /// Note that every UTF-16 character consumes 2 or 4 bytes.
187    ///
188    /// This function has *O*(1) complexity.
189    ///
190    /// See [`try_push_u16_until_nul`] or [`try_push_u16cstr`] if you have a NUL-terminated buffer.
191    ///
192    /// [`try_push_u16_until_nul`]: Self::try_push_u16_until_nul
193    /// [`try_push_u16cstr`]: Self::try_push_u16cstr
194    pub fn try_push_u16(&mut self, buffer: &[u16]) -> Result<()> {
195        // Determine the required additional capacity.
196        //
197        // Add one element for the terminating NUL character
198        // (for applications that mistakenly treat the buffer of a Unicode String as a NUL-terminated string).
199        let additional_elements = buffer
200            .len()
201            .checked_add(1)
202            .ok_or(NtStringError::BufferSizeExceedsU16)?;
203
204        // Reserve the additional capacity.
205        let additional_bytes = additional_elements
206            .checked_mul(mem::size_of::<u16>())
207            .ok_or(NtStringError::BufferSizeExceedsU16)?;
208        let additional =
209            u16::try_from(additional_bytes).map_err(|_| NtStringError::BufferSizeExceedsU16)?;
210        self.try_reserve(additional)?;
211
212        // Copy over the string.
213        let end_index = self.len_in_elements();
214        self.raw.length += additional;
215
216        let dest_slice = &mut self.as_mut_slice()[end_index..];
217        dest_slice[..buffer.len()].copy_from_slice(buffer);
218
219        // NUL-terminate it.
220        dest_slice[buffer.len()] = 0;
221
222        // Set the final length (without the terminating NUL character).
223        self.raw.length -= mem::size_of::<u16>() as u16;
224
225        Ok(())
226    }
227
228    /// Appends the given [`u16`] string buffer, which contains at least one NUL character,
229    /// to the end of this string.
230    ///
231    /// The string is expected to consist of valid UTF-16 characters.
232    ///
233    /// The string will be terminated at the NUL character.
234    /// An [`NtStringError::NulNotFound`] error is returned if no NUL character could be found.
235    /// As a consequence, this function has *O*(*n*) complexity.
236    ///
237    /// Returns an [`NtStringError::BufferSizeExceedsU16`] error if the resulting string would exceed
238    /// 65532 bytes.
239    /// This is due to the fact that a `UNICODE_STRING` internally uses a 16-bit field to store the length.
240    /// Additionally, this function allocates one more character for NUL termination of the internal
241    /// buffer.
242    /// See the [module-level documentation](super) for the implications of that.
243    ///
244    /// Note that every UTF-16 character consumes 2 or 4 bytes.
245    ///
246    /// Use [`try_push_u16`] if you have a buffer that is not NUL-terminated.
247    /// You can also push a NUL-terminated [`U16CStr`] in *O*(1) via [`try_push_u16cstr`].
248    ///
249    /// [`try_push_u16`]: Self::try_push_u16
250    /// [`try_push_u16cstr`]: Self::try_push_u16cstr
251    pub fn try_push_u16_until_nul(&mut self, buffer: &[u16]) -> Result<()> {
252        match buffer.iter().position(|x| *x == 0) {
253            Some(nul_pos) => self.try_push_u16(&buffer[..nul_pos]),
254            None => Err(NtStringError::NulNotFound),
255        }
256    }
257
258    /// Appends the given [`U16CStr`] to the end of this string.
259    ///
260    /// Returns an [`NtStringError::BufferSizeExceedsU16`] error if the resulting string would exceed
261    /// 65532 bytes.
262    /// This is due to the fact that a `UNICODE_STRING` internally uses a 16-bit field to store the length.
263    /// Additionally, this function allocates one more character for NUL termination of the internal
264    /// buffer.
265    /// See the [module-level documentation](super) for the implications of that.
266    ///
267    /// Note that every UTF-16 character consumes 2 or 4 bytes.
268    ///
269    /// This function has *O*(1) complexity.
270    pub fn try_push_u16cstr(&mut self, u16cstr: &U16CStr) -> Result<()> {
271        self.try_push_u16(u16cstr.as_slice())
272    }
273
274    /// Appends the given [`U16Str`] to the end of this string.
275    ///
276    /// Returns an [`NtStringError::BufferSizeExceedsU16`] error if the resulting string would exceed
277    /// 65532 bytes.
278    /// This is due to the fact that a `UNICODE_STRING` internally uses a 16-bit field to store the length.
279    /// Additionally, this function allocates one more character for NUL termination of the internal
280    /// buffer.
281    /// See the [module-level documentation](super) for the implications of that.
282    ///
283    /// Note that every UTF-16 character consumes 2 or 4 bytes.
284    ///
285    /// This function has *O*(1) complexity.
286    pub fn try_push_u16str(&mut self, u16str: &U16Str) -> Result<()> {
287        self.try_push_u16(u16str.as_slice())
288    }
289
290    /// Reserves capacity for `additional` bytes more than the current length.
291    ///
292    /// Returns an [`NtStringError::BufferSizeExceedsU16`] error if the resulting capacity would exceed
293    /// 65535 bytes.
294    ///
295    /// Note that every UTF-16 character consumes 2 or 4 bytes.
296    pub fn try_reserve(&mut self, additional: u16) -> Result<()> {
297        if self.remaining_capacity() >= additional {
298            return Ok(());
299        }
300
301        let new_capacity = self
302            .len()
303            .checked_add(additional)
304            .ok_or(NtStringError::BufferSizeExceedsU16)?;
305
306        if self.raw.buffer.is_null() {
307            self.raw.maximum_length = new_capacity;
308            let new_layout = self.layout();
309
310            self.raw.buffer = unsafe { alloc::alloc(new_layout) } as *mut u16;
311        } else {
312            let old_layout = self.layout();
313
314            self.raw.buffer = unsafe {
315                alloc::realloc(
316                    self.raw.buffer as *mut u8,
317                    old_layout,
318                    usize::from(new_capacity),
319                )
320            } as *mut u16;
321
322            self.raw.maximum_length = new_capacity;
323        }
324
325        Ok(())
326    }
327
328    /// Creates an empty [`NtUnicodeString`] with at least the specified capacity.
329    ///
330    /// This will preallocate a buffer with the given capacity.
331    /// If the given capacity is `0`, no allocation will occur, and this method is identical to the [`new`] method.
332    ///
333    /// [`new`]: Self::new
334    pub fn with_capacity(capacity: u16) -> Self {
335        let mut string = Self::new();
336        string.try_reserve(capacity).unwrap();
337        string
338    }
339}
340
341impl Add<&str> for NtUnicodeString {
342    type Output = NtUnicodeString;
343
344    fn add(mut self, rhs: &str) -> Self::Output {
345        if let Err(e) = self.try_push_str(rhs) {
346            panic!("{e}");
347        }
348
349        self
350    }
351}
352
353impl Add<&U16CStr> for NtUnicodeString {
354    type Output = NtUnicodeString;
355
356    fn add(mut self, rhs: &U16CStr) -> Self::Output {
357        if let Err(e) = self.try_push_u16cstr(rhs) {
358            panic!("{e}");
359        }
360
361        self
362    }
363}
364
365impl Add<&U16Str> for NtUnicodeString {
366    type Output = NtUnicodeString;
367
368    fn add(mut self, rhs: &U16Str) -> Self::Output {
369        if let Err(e) = self.try_push_u16str(rhs) {
370            panic!("{e}");
371        }
372
373        self
374    }
375}
376
377impl AddAssign<&str> for NtUnicodeString {
378    fn add_assign(&mut self, rhs: &str) {
379        if let Err(e) = self.try_push_str(rhs) {
380            panic!("{e}");
381        }
382    }
383}
384
385impl AddAssign<&U16CStr> for NtUnicodeString {
386    fn add_assign(&mut self, rhs: &U16CStr) {
387        if let Err(e) = self.try_push_u16cstr(rhs) {
388            panic!("{e}");
389        }
390    }
391}
392
393impl AddAssign<&U16Str> for NtUnicodeString {
394    fn add_assign(&mut self, rhs: &U16Str) {
395        if let Err(e) = self.try_push_u16str(rhs) {
396            panic!("{e}");
397        }
398    }
399}
400
401impl Clone for NtUnicodeString {
402    /// Creates a copy of this [`NtUnicodeString`].
403    ///
404    /// This implementation keeps the original capacity.
405    fn clone(&self) -> Self {
406        NtUnicodeString::from(self.as_unicode_str())
407    }
408}
409
410impl Default for NtUnicodeString {
411    fn default() -> Self {
412        Self::new()
413    }
414}
415
416impl Deref for NtUnicodeString {
417    // The NtUnicodeString can't be dropped while someone still holds a reference to it,
418    // so the `buffer` lifetime is effectively 'static for the duration of the reference.
419    type Target = NtUnicodeStrMut<'static>;
420
421    fn deref(&self) -> &Self::Target {
422        // SAFETY: `NtUnicodeStrMut` and `NtUnicodeString` have the same memory layout,
423        // so we can safely transmute `NtUnicodeString` to `NtUnicodeStrMut`.
424        unsafe { mem::transmute(self) }
425    }
426}
427
428impl DerefMut for NtUnicodeString {
429    fn deref_mut(&mut self) -> &mut Self::Target {
430        // SAFETY: `NtUnicodeStrMut` and `NtUnicodeString` have the same memory layout,
431        // so we can safely transmute `NtUnicodeString` to `NtUnicodeStrMut`.
432        unsafe { mem::transmute(self) }
433    }
434}
435
436impl fmt::Display for NtUnicodeString {
437    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
438        fmt::Display::fmt(self.deref(), f)
439    }
440}
441
442impl Drop for NtUnicodeString {
443    fn drop(&mut self) {
444        if !self.raw.buffer.is_null() {
445            let layout = self.layout();
446            unsafe { alloc::dealloc(self.raw.buffer as *mut u8, layout) }
447        }
448    }
449}
450
451impl Eq for NtUnicodeString {}
452
453impl From<char> for NtUnicodeString {
454    /// Creates an [`NtUnicodeString`] from a single [`char`].
455    fn from(c: char) -> Self {
456        let mut string = Self::new();
457        string.try_push(c).unwrap();
458        string
459    }
460}
461
462impl<'a> From<&NtUnicodeStr<'a>> for NtUnicodeString {
463    /// Creates an [`NtUnicodeString`] from an existing [`NtUnicodeStr`].
464    ///
465    /// This implementation keeps the original capacity.
466    fn from(unicode_str: &NtUnicodeStr) -> Self {
467        let mut new_string = Self::with_capacity(unicode_str.capacity());
468
469        if !unicode_str.is_empty() {
470            new_string.raw.length = unicode_str.len();
471            new_string
472                .as_mut_slice()
473                .copy_from_slice(unicode_str.as_slice());
474        }
475
476        new_string
477    }
478}
479
480impl Ord for NtUnicodeString {
481    fn cmp(&self, other: &Self) -> Ordering {
482        Ord::cmp(self.deref(), other.deref())
483    }
484}
485
486impl PartialOrd<NtUnicodeString> for NtUnicodeString {
487    fn partial_cmp(&self, other: &NtUnicodeString) -> Option<Ordering> {
488        Some(self.cmp(other))
489    }
490}
491
492impl_eq! { NtUnicodeString, NtUnicodeString }
493impl_eq! { NtUnicodeStr<'a>, NtUnicodeString }
494impl_eq! { NtUnicodeString, NtUnicodeStr<'a> }
495impl_eq! { NtUnicodeStrMut<'a>, NtUnicodeString }
496impl_eq! { NtUnicodeString, NtUnicodeStrMut<'a> }
497impl_eq! { NtUnicodeString, str }
498impl_eq! { str, NtUnicodeString }
499impl_eq! { NtUnicodeString, &str }
500impl_eq! { &str, NtUnicodeString }
501
502impl_partial_cmp! { NtUnicodeStr<'a>, NtUnicodeString }
503impl_partial_cmp! { NtUnicodeString, NtUnicodeStr<'a> }
504impl_partial_cmp! { NtUnicodeStrMut<'a>, NtUnicodeString }
505impl_partial_cmp! { NtUnicodeString, NtUnicodeStrMut<'a> }
506impl_partial_cmp! { NtUnicodeString, str }
507impl_partial_cmp! { str, NtUnicodeString }
508impl_partial_cmp! { NtUnicodeString, &str }
509impl_partial_cmp! { &str, NtUnicodeString }
510
511impl TryExtend<char> for NtUnicodeString {
512    type Error = NtStringError;
513
514    fn try_extend<I: IntoIterator<Item = char>>(&mut self, iter: I) -> Result<()> {
515        let iterator = iter.into_iter();
516        let (lower_bound, _) = iterator.size_hint();
517
518        // Add one element for the terminating NUL character
519        // (for applications that mistakenly treat the buffer of a Unicode String as a NUL-terminated string).
520        let additional_elements = lower_bound + 1;
521
522        // Reserve the additional capacity once to save on allocations.
523        let additional_bytes = u16::try_from(additional_elements * mem::size_of::<u16>())
524            .map_err(|_| NtStringError::BufferSizeExceedsU16)?;
525        self.try_reserve(additional_bytes)?;
526
527        for ch in iterator {
528            self.try_push(ch)?;
529        }
530
531        Ok(())
532    }
533
534    fn try_extend_one(&mut self, item: char) -> Result<()> {
535        self.try_push(item)
536    }
537}
538
539impl<'a> TryExtend<&'a str> for NtUnicodeString {
540    type Error = NtStringError;
541
542    fn try_extend<I: IntoIterator<Item = &'a str>>(&mut self, iter: I) -> Result<()> {
543        for s in iter.into_iter() {
544            self.try_push_str(s)?;
545        }
546
547        Ok(())
548    }
549}
550
551impl<'a> TryExtend<&'a U16CStr> for NtUnicodeString {
552    type Error = NtStringError;
553
554    fn try_extend<I: IntoIterator<Item = &'a U16CStr>>(&mut self, iter: I) -> Result<()> {
555        for s in iter.into_iter() {
556            self.try_push_u16cstr(s)?;
557        }
558
559        Ok(())
560    }
561}
562
563impl<'a> TryExtend<&'a U16Str> for NtUnicodeString {
564    type Error = NtStringError;
565
566    fn try_extend<I: IntoIterator<Item = &'a U16Str>>(&mut self, iter: I) -> Result<()> {
567        for s in iter.into_iter() {
568            self.try_push_u16str(s)?;
569        }
570
571        Ok(())
572    }
573}
574
575impl TryFrom<&str> for NtUnicodeString {
576    type Error = NtStringError;
577
578    /// Converts a string slice into an owned [`NtUnicodeString`].
579    ///
580    /// This allocates a buffer of matching size on the heap and NUL-terminates it internally.
581    /// See the [module-level documentation](super) for the implications of that.
582    fn try_from(s: &str) -> Result<Self> {
583        let mut string = Self::new();
584        string.try_push_str(s)?;
585        Ok(string)
586    }
587}
588
589impl TryFrom<String> for NtUnicodeString {
590    type Error = NtStringError;
591
592    /// Converts a [`String`] into an owned [`NtUnicodeString`].
593    ///
594    /// This allocates a buffer of matching size on the heap and NUL-terminates it internally.
595    /// See the [module-level documentation](super) for the implications of that.
596    fn try_from(s: String) -> Result<Self> {
597        NtUnicodeString::try_from(s.as_str())
598    }
599}
600
601impl TryFrom<&String> for NtUnicodeString {
602    type Error = NtStringError;
603
604    /// Converts a [`String`] reference into an owned [`NtUnicodeString`].
605    ///
606    /// This allocates a buffer of matching size on the heap and NUL-terminates it internally.
607    /// See the [module-level documentation](super) for the implications of that.
608    fn try_from(s: &String) -> Result<Self> {
609        NtUnicodeString::try_from(s.as_str())
610    }
611}
612
613impl TryFrom<&U16CStr> for NtUnicodeString {
614    type Error = NtStringError;
615
616    /// Converts a [`U16CStr`] reference into an owned [`NtUnicodeString`].
617    ///
618    /// The internal buffer will be NUL-terminated.
619    /// See the [module-level documentation](super) for the implications of that.
620    fn try_from(value: &U16CStr) -> Result<Self> {
621        let unicode_str = NtUnicodeStr::try_from(value)?;
622        Ok(Self::from(&unicode_str))
623    }
624}
625
626impl TryFrom<&U16Str> for NtUnicodeString {
627    type Error = NtStringError;
628
629    /// Converts a [`U16Str`] reference into an owned [`NtUnicodeString`].
630    ///
631    /// The internal buffer will NOT be NUL-terminated.
632    /// See the [module-level documentation](super) for the implications of that.
633    fn try_from(value: &U16Str) -> Result<Self> {
634        let unicode_str = NtUnicodeStr::try_from(value)?;
635        Ok(Self::from(&unicode_str))
636    }
637}
638
639#[cfg(test)]
640mod tests {
641    use alloc::vec::Vec;
642
643    use crate::error::NtStringError;
644    use crate::traits::TryExtend;
645    use crate::unicode_string::NtUnicodeString;
646
647    #[test]
648    fn test_add() {
649        let mut string = NtUnicodeString::new();
650        string += "๐Ÿ‘";
651        assert_eq!(string, "๐Ÿ‘");
652
653        let string2 = string + "๐Ÿ‘Ž";
654        assert_eq!(string2, "๐Ÿ‘๐Ÿ‘Ž");
655    }
656
657    #[test]
658    fn test_chars() {
659        // Verify that ASCII characters work and allocate 2 bytes per character.
660        // Verify that one more element (2 bytes) is allocated for the terminating NUL character.
661        let moin = NtUnicodeString::try_from("Moin").unwrap();
662        assert_eq!(moin.capacity(), 10);
663        assert_eq!(moin.len(), 8);
664        let vec = moin.chars_lossy().collect::<Vec<char>>();
665        assert_eq!(vec, ['M', 'o', 'i', 'n']);
666
667        // Verify that Unicode characters inside the Basic Multilingual Plane work and allocate 2 bytes per character.
668        // Verify that one more element (2 bytes) is allocated for the terminating NUL character.
669        let ไปŠๆ—ฅใฏ = NtUnicodeString::try_from("ไปŠๆ—ฅใฏ").unwrap();
670        assert_eq!(ไปŠๆ—ฅใฏ.capacity(), 8);
671        assert_eq!(ไปŠๆ—ฅใฏ.len(), 6);
672        let vec = ไปŠๆ—ฅใฏ.chars_lossy().collect::<Vec<char>>();
673        assert_eq!(vec, ['ไปŠ', 'ๆ—ฅ', 'ใฏ']);
674
675        // Verify that Unicode characters outside the Basic Multilingual Plane (e.g. emojis) work and allocate 4 bytes per character.
676        // Verify that one more element (2 bytes) is allocated for the terminating NUL character.
677        let smile = NtUnicodeString::try_from("๐Ÿ˜€").unwrap();
678        assert_eq!(smile.capacity(), 6);
679        assert_eq!(smile.len(), 4);
680        let vec = smile.chars_lossy().collect::<Vec<char>>();
681        assert_eq!(vec, ['๐Ÿ˜€']);
682    }
683
684    #[test]
685    fn test_cmp() {
686        let a = NtUnicodeString::try_from("a").unwrap();
687        let b = NtUnicodeString::try_from("b").unwrap();
688        assert!(a < b);
689    }
690
691    #[test]
692    fn test_eq() {
693        let hello = NtUnicodeString::try_from("Hello").unwrap();
694        let hello_again = NtUnicodeString::try_from("Hello again").unwrap();
695        assert_ne!(hello, hello_again);
696
697        let mut hello_clone = hello.clone();
698        assert_eq!(hello, hello_clone);
699
700        hello_clone.try_reserve(42).unwrap();
701        assert_eq!(hello, hello_clone);
702    }
703
704    #[test]
705    fn test_extend_and_pop() {
706        // Verify that 32766 characters still work.
707        // Verify that one more element (2 bytes) is allocated for the terminating NUL character.
708        let a_string = "a".repeat(32766);
709        let mut string = NtUnicodeString::try_from(a_string).unwrap();
710        assert_eq!(string.capacity(), 65534);
711        assert_eq!(string.len(), 65532);
712
713        // Verify that extending by a single character fails.
714        assert_eq!(
715            string.try_extend(Some('b')),
716            Err(NtStringError::BufferSizeExceedsU16)
717        );
718
719        // Pop a character to append a new one.
720        assert_eq!(string.pop(), Some(Ok('a')));
721        assert_eq!(string.capacity(), 65534);
722        assert_eq!(string.len(), 65530);
723        string.try_extend_one('c').unwrap();
724        assert_eq!(string.capacity(), 65534);
725        assert_eq!(string.len(), 65532);
726
727        // Pop two characters to append an emoji.
728        assert_eq!(string.pop(), Some(Ok('c')));
729        assert_eq!(string.pop(), Some(Ok('a')));
730        assert_eq!(string.capacity(), 65534);
731        assert_eq!(string.len(), 65528);
732        string.try_extend_one('๐Ÿ˜€').unwrap();
733        assert_eq!(string.capacity(), 65534);
734        assert_eq!(string.len(), 65532);
735
736        // Pop the emoji and another character to append 3 ASCII characters.
737        assert_eq!(string.pop(), Some(Ok('๐Ÿ˜€')));
738        assert_eq!(string.pop(), Some(Ok('a')));
739        assert_eq!(string.capacity(), 65534);
740        assert_eq!(string.len(), 65526);
741        string.try_extend("def".chars()).unwrap();
742        assert_eq!(string.capacity(), 65534);
743        assert_eq!(string.len(), 65532);
744    }
745
746    #[test]
747    fn test_from_u16() {
748        // Verify that `try_from_u16` fails if we exceed a `u16`.
749        let mut a_vec = "a".repeat(32768).encode_utf16().collect::<Vec<u16>>();
750        assert_eq!(
751            NtUnicodeString::try_from_u16(&a_vec),
752            Err(NtStringError::BufferSizeExceedsU16)
753        );
754
755        // Verify that `try_from_u16` does not reserve any space for a terminating NUL character.
756        a_vec.pop();
757        let string = NtUnicodeString::try_from_u16(&a_vec).unwrap();
758        assert_eq!(string.capacity(), 65534);
759        assert_eq!(string.len(), 65534);
760
761        // Verify that `try_from_u16_until_nul` does reserve space for a terminating NUL character.
762        a_vec[4] = 0;
763        let string = NtUnicodeString::try_from_u16_until_nul(&a_vec).unwrap();
764        assert_eq!(string.capacity(), 10);
765        assert_eq!(string.len(), 8);
766        assert_eq!(string, "aaaa");
767    }
768
769    #[test]
770    fn test_push_str() {
771        let mut string = NtUnicodeString::new();
772        string.try_push_str("Hey").unwrap();
773        assert_eq!(string, "Hey");
774        assert_eq!(string.capacity(), 8);
775        assert_eq!(string.len(), 6);
776
777        string.try_push_str("Ho").unwrap();
778        assert_eq!(string, "HeyHo");
779        assert_eq!(string.capacity(), 12);
780        assert_eq!(string.len(), 10);
781    }
782
783    #[test]
784    fn test_reserve() {
785        let mut string = NtUnicodeString::new();
786        assert_eq!(string.capacity(), 0);
787
788        string.try_reserve(5).unwrap();
789        assert_eq!(string.capacity(), 5);
790
791        string.try_reserve(3).unwrap();
792        assert_eq!(string.capacity(), 5);
793
794        string.try_push_str("a").unwrap();
795        assert_eq!(string, "a");
796        assert_eq!(string.capacity(), 5);
797
798        string.try_push_str("b").unwrap();
799        assert_eq!(string, "ab");
800        assert_eq!(string.capacity(), 6);
801    }
802}