use core::iter::FusedIterator;
use core::marker::PhantomData;
use core::mem::ManuallyDrop;
use core::ops::{Range, RangeBounds};
use core::ptr::{self, NonNull};
use core::{slice, str};
use crate::TinyString;
pub struct Drain<'a, const N: usize> {
string: NonNull<TinyString<N>>,
remaining_start: usize,
remaining_len: usize,
tail_start: usize,
tail_len: usize,
_marker: PhantomData<&'a mut TinyString<N>>,
}
impl<const N: usize> Drain<'_, N> {
pub const fn as_str(&mut self) -> &str {
if self.remaining_len == 0 {
return ""
}
unsafe {
let str = self.string.as_ref().as_ptr().add(self.remaining_start);
let slice = slice::from_raw_parts(str, self.remaining_len);
str::from_utf8_unchecked(slice)
}
}
pub fn keep_rest(self) {
let mut slf = ManuallyDrop::new(self);
unsafe {
let string = slf.string.as_mut();
let start = string.len();
let buf = string.as_mut_ptr();
let dst = buf.add(start);
let src = buf.add(slf.remaining_start);
ptr::copy(src, dst, slf.remaining_len);
let src = buf.add(slf.tail_start);
let dst = dst.add(slf.remaining_len);
ptr::copy(src, dst, slf.tail_len);
string.buf.set_len(start + slf.remaining_len + slf.tail_len);
}
}
}
impl<const N: usize> Iterator for Drain<'_, N> {
type Item = char;
fn next(&mut self) -> Option<Self::Item> {
let c = self.as_str().chars().next()?;
self.remaining_start += c.len_utf8();
self.remaining_len -= c.len_utf8();
Some(c)
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.remaining_len, Some(self.remaining_len))
}
}
impl<const N: usize> DoubleEndedIterator for Drain<'_, N> {
fn next_back(&mut self) -> Option<Self::Item> {
let c = self.as_str().chars().next_back()?;
self.remaining_len -= c.len_utf8();
Some(c)
}
}
impl<const N: usize> ExactSizeIterator for Drain<'_, N> {}
impl<const N: usize> FusedIterator for Drain<'_, N> {}
impl<const N: usize> Drop for Drain<'_, N> {
fn drop(&mut self) {
unsafe {
let string = self.string.as_mut();
let len = string.len();
let ptr = string.as_mut_ptr();
let src = ptr.add(self.tail_start);
let dst = ptr.add(len);
ptr::copy(src, dst, self.tail_len);
string.buf.set_len(len + self.tail_len);
}
}
}
impl<const N: usize> TinyString<N> {
pub fn drain<R>(&mut self, range: R) -> Drain<'_, N>
where
R: RangeBounds<usize>
{
let len = self.len();
let Range { start, end } = self.slice_range(range, len);
unsafe {
self.buf.set_len(start);
let string = NonNull::new_unchecked(self as *mut _);
Drain {
tail_len: len - end,
tail_start: end,
remaining_start: start,
remaining_len: end - start,
string,
_marker: PhantomData,
}
}
}
}