split_buf.rs

use core::{
    cmp::Ordering,
    iter::{Chain, Copied, Enumerate},
    ops::{Index, RangeFrom, RangeTo},
};

use nom::CompareResult;

#[cfg_attr(feature = "defmt", derive(defmt::Format))]
#[derive(Clone, Copy, Debug)]
pub struct SplitBuf<'a>(pub &'a [u8], pub &'a [u8]);

impl PartialEq for SplitBuf<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.len() == other.len() && self.iter().zip(other.iter()).all(|(a, b)| a == b)
    }
}

impl PartialEq<&[u8]> for SplitBuf<'_> {
    fn eq(&self, other: &&[u8]) -> bool {
        *self == SplitBuf::from(*other)
    }
}

impl PartialEq<&str> for SplitBuf<'_> {
    fn eq(&self, other: &&str) -> bool {
        *self == SplitBuf::from(other.as_bytes())
    }
}

impl<'a> SplitBuf<'a> {
    pub fn len(&self) -> usize {
        self.0.len() + self.1.len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn iter(&self) -> Chain<core::slice::Iter<'a, u8>, core::slice::Iter<'a, u8>> {
        self.0.iter().chain(self.1.iter())
    }

    pub fn index_from(&self, index: RangeFrom<usize>) -> Self {
        assert!(index.start <= self.len());

        if index.start < self.0.len() {
            SplitBuf(&self.0[index.start..], self.1)
        } else {
            SplitBuf(&self.1[index.start - self.0.len()..], &[])
        }
    }

    pub fn index_to(&self, index: RangeTo<usize>) -> Self {
        assert!(index.end <= self.len());
        assert!(self.1.is_empty() || !self.0.is_empty());

        if index.end < self.0.len() {
            SplitBuf(&self.0[..index.end], &[])
        } else {
            SplitBuf(self.0, &self.1[..index.end - self.0.len()])
        }
    }

    pub fn split_at(&self, index: usize) -> (Self, Self) {
        (self.index_to(..index), self.index_from(index..))
    }

    pub fn split_once(&self, substr: &[u8]) -> Option<(Self, Self)> {
        let offset = self.find_substring(substr)?;

        Some((self.index_to(..offset), self.index_from(offset + substr.len()..)))
    }

    pub fn copy_to_slice(&self, buf: &mut [u8]) {
        assert_eq!(buf.len(), self.len());
        assert!(self.1.is_empty() || !self.0.is_empty());

        if self.1.is_empty() {
            buf.copy_from_slice(self.0);
        } else {
            let (left, right) = buf.split_at_mut(self.0.len());
            left.copy_from_slice(self.0);
            right.copy_from_slice(self.1);
        }
    }

    pub fn starts_with(&self, v: &[u8]) -> bool {
        if v.len() > self.len() {
            return false;
        }

        self.iter().zip(v.iter()).all(|(a, b)| a == b)
    }

    pub fn strip_prefix(&self, prefix: &[u8]) -> Option<Self> {
        if self.starts_with(prefix) {
            Some(self.index_from(prefix.len()..))
        } else {
            None
        }
    }

    pub fn find_substring(&self, substr: &[u8]) -> Option<usize> {
        (0..=self.len().saturating_sub(substr.len()))
            .find(|&offset| self.index_from(offset..).starts_with(substr))
    }

    /// Like `find_substring`, but checks if `substr` actually references the same memory as `self`
    #[cfg(test)]
    pub(crate) fn find_substring_by_ref(&self, substr: SplitBuf<'_>) -> Option<usize> {
        extern crate std;
        use std::dbg;
        if dbg!(self.0.as_ptr_range()).contains(&substr.0.as_ptr()) {
            Some((substr.0.as_ptr() as isize - self.0.as_ptr() as isize) as usize)
        } else if dbg!(self.1.as_ptr_range()).contains(dbg!(&substr.0.as_ptr())) {
            Some((substr.1.as_ptr() as isize - self.0.as_ptr() as isize) as usize + self.0.len())
        } else {
            None
        }
    }
}

impl<'a> From<&'a [u8]> for SplitBuf<'a> {
    fn from(buf: &'a [u8]) -> Self {
        SplitBuf(buf, &[])
    }
}

impl<'a, const N: usize> TryFrom<SplitBuf<'a>> for heapless::String<N> {
    type Error = Option<core::str::Utf8Error>;

    fn try_from(value: SplitBuf<'a>) -> Result<Self, Self::Error> {
        let mut vec = heapless::Vec::from_slice(value.0).map_err(|()| None)?;
        vec.extend_from_slice(value.1).map_err(|()| None)?;
        heapless::String::from_utf8(vec).map_err(Some)
    }
}

impl Index<usize> for SplitBuf<'_> {
    type Output = u8;

    fn index(&self, index: usize) -> &Self::Output {
        if index < self.0.len() {
            &self.0[index]
        } else {
            &self.1[index - self.0.len()]
        }
    }
}

impl<'a> nom::Input for SplitBuf<'a> {
    type Item = u8;
    type Iter = Copied<Chain<core::slice::Iter<'a, u8>, core::slice::Iter<'a, u8>>>;
    type IterIndices = Enumerate<Self::Iter>;

    fn input_len(&self) -> usize {
        self.len()
    }

    fn take(&self, count: usize) -> Self {
        if count < self.0.len() {
            SplitBuf(&self.0[..count], &[])
        } else {
            SplitBuf(self.0, &self.1[..count - self.0.len()])
        }
    }

    fn take_from(&self, index: usize) -> Self {
        self.index_from(index..)
    }

    fn take_split(&self, count: usize) -> (Self, Self) {
        match count.cmp(&self.0.len()) {
            Ordering::Less => (SplitBuf(&self.0[count..], self.1), SplitBuf(&self.0[..count], &[])),
            Ordering::Equal => (SplitBuf(self.1, &[]), SplitBuf(self.0, &[])),
            Ordering::Greater => (
                SplitBuf(&self.1[count - self.0.len()..], &[]),
                SplitBuf(self.0, &self.1[..count - self.0.len()]),
            ),
        }
    }

    fn position<P: Fn(Self::Item) -> bool>(&self, predicate: P) -> Option<usize> {
        self.iter().position(|&c| predicate(c))
    }

    fn iter_elements(&self) -> Self::Iter {
        self.iter().copied()
    }

    fn iter_indices(&self) -> Self::IterIndices {
        self.iter().copied().enumerate()
    }

    fn slice_index(&self, count: usize) -> Result<usize, nom::Needed> {
        if self.len() >= count {
            Ok(count)
        } else {
            Err(nom::Needed::new(count - self.len()))
        }
    }
}

impl<'a> nom::Compare<&'a str> for SplitBuf<'_> {
    fn compare(&self, t: &'a str) -> CompareResult {
        let pos = self
            .iter()
            .zip(t.as_bytes().iter())
            .position(|(a, b)| a != b);

        match pos {
            Some(_) => CompareResult::Error,
            None => {
                if self.len() >= t.len() {
                    CompareResult::Ok
                } else {
                    CompareResult::Incomplete
                }
            }
        }
    }

    fn compare_no_case(&self, _: &'a str) -> CompareResult {
        unimplemented!()
    }
}

impl<'a> nom::Compare<&'a [u8]> for SplitBuf<'_> {
    fn compare(&self, t: &'a [u8]) -> CompareResult {
        let pos = self.iter().zip(t.iter()).position(|(a, b)| a != b);

        match pos {
            Some(_) => CompareResult::Error,
            None => {
                if self.len() >= t.len() {
                    CompareResult::Ok
                } else {
                    CompareResult::Incomplete
                }
            }
        }
    }

    fn compare_no_case(&self, _: &'a [u8]) -> CompareResult {
        unimplemented!()
    }
}

// impl nom::InputTakeAtPosition for SplitBuf<'_> {
//     fn split_at_position<P, E: ParseError<Self>>(&self, predicate: P) -> IResult<Self, Self, E>
//     where
//         P: Fn(Self::Item) -> bool,
//     {
//         match self.iter().position(|c| predicate(*c)) {
//             Some(i) => Ok(self.take_split(i)),
//             None => Err(nom::Err::Incomplete(nom::Needed::new(1))),
//         }
//     }

//     fn split_at_position1<P, E: ParseError<Self>>(
//         &self,
//         predicate: P,
//         e: nom::error::ErrorKind,
//     ) -> IResult<Self, Self, E>
//     where
//         P: Fn(Self::Item) -> bool,
//     {
//         match self.iter().position(|c| predicate(*c)) {
//             Some(0) => Err(nom::Err::Error(E::from_error_kind(*self, e))),
//             Some(i) => Ok(self.take_split(i)),
//             None => Err(nom::Err::Incomplete(nom::Needed::new(1))),
//         }
//     }

//     fn split_at_position_complete<P, E: ParseError<Self>>(
//         &self,
//         predicate: P,
//     ) -> IResult<Self, Self, E>
//     where
//         P: Fn(Self::Item) -> bool,
//     {
//         match self.iter().position(|c| predicate(*c)) {
//             Some(i) => Ok(self.take_split(i)),
//             None => Ok(self.take_split(self.input_len())),
//         }
//     }

//     fn split_at_position1_complete<P, E: ParseError<Self>>(
//         &self,
//         predicate: P,
//         e: nom::error::ErrorKind,
//     ) -> IResult<Self, Self, E>
//     where
//         P: Fn(Self::Item) -> bool,
//     {
//         match self.iter().position(|c| predicate(*c)) {
//             Some(0) => Err(nom::Err::Error(E::from_error_kind(*self, e))),
//             Some(i) => Ok(self.take_split(i)),
//             None => {
//                 if self.is_empty() {
//                     Err(nom::Err::Error(E::from_error_kind(*self, e)))
//                 } else {
//                     Ok(self.take_split(self.input_len()))
//                 }
//             }
//         }
//     }
// }

impl<'a> nom::FindSubstring<&'a str> for SplitBuf<'_> {
    fn find_substring(&self, substr: &'a str) -> Option<usize> {
        self.find_substring(substr.as_bytes())
    }
}

#[cfg(test)]
mod tests {
    use nom::{Compare as _, IResult, Input as _, Parser as _, bytes::streaming::tag};

    use super::*;

    #[test]
    fn split_buf_index() {
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6])[0], 1);
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6])[1], 2);
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6])[3], 4);
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6])[4], 5);
    }

    #[test]
    fn split_buf_index_from() {
        assert_eq!(
            SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_from(0..),
            SplitBuf(&[1, 2, 3], &[4, 5, 6])
        );
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_from(1..), SplitBuf(&[2, 3], &[4, 5, 6]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_from(3..), SplitBuf(&[4, 5, 6], &[]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_from(4..), SplitBuf(&[5, 6], &[]));
    }

    #[test]
    fn split_buf_index_to() {
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_to(..0), SplitBuf(&[], &[]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_to(..1), SplitBuf(&[1], &[]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_to(..2), SplitBuf(&[1, 2], &[]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_to(..3), SplitBuf(&[1, 2, 3], &[]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_to(..4), SplitBuf(&[1, 2, 3], &[4]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_to(..5), SplitBuf(&[1, 2, 3], &[4, 5]));
        assert_eq!(
            SplitBuf(&[1, 2, 3], &[4, 5, 6]).index_to(..6),
            SplitBuf(&[1, 2, 3], &[4, 5, 6])
        );
    }

    #[test]
    fn split_buf_copy_to_slice() {
        let mut buf = [0; 5];
        SplitBuf(&[1, 2, 3], &[4, 5]).copy_to_slice(&mut buf);
        assert_eq!(buf, [1, 2, 3, 4, 5,]);

        let mut buf = [0; 3];
        SplitBuf(&[1, 2, 3], &[]).copy_to_slice(&mut buf);
        assert_eq!(buf, [1, 2, 3]);
    }

    #[test]
    fn split_buf_starts_with() {
        assert!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[]));
        assert!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[1]));
        assert!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[1, 2]));
        assert!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[1, 2, 3]));
        assert!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[1, 2, 3, 4]));
        assert!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[1, 2, 3, 4, 5]));
        assert!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[1, 2, 3, 4, 5, 6]));

        assert!(!SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[1, 2, 3, 4, 5, 6, 7]));
        assert!(!SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[1, 2, 3, 4, 5, 6, 7, 8]));
        assert!(!SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[2]));
        assert!(!SplitBuf(&[1, 2, 3], &[4, 5, 6]).starts_with(&[4]));

        assert!(!SplitBuf(&[], &[]).starts_with(&[1]));
    }

    #[test]
    fn split_buf_try_into_heapless_string() {
        assert_eq!(
            heapless::String::<6>::try_from(SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36])),
            Ok(heapless::String::try_from("123456").unwrap())
        );

        assert_eq!(
            heapless::String::<4>::try_from(SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36])),
            Err(None)
        );
    }

    #[test]
    fn split_buf_input_take() {
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).take(0), SplitBuf(&[], &[]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).take(1), SplitBuf(&[1], &[]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).take(3), SplitBuf(&[1, 2, 3], &[]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).take(4), SplitBuf(&[1, 2, 3], &[4]));
        assert_eq!(SplitBuf(&[1, 2, 3], &[4, 5, 6]).take(6), SplitBuf(&[1, 2, 3], &[4, 5, 6]));
    }

    #[test]
    fn split_buf_input_take_split() {
        assert_eq!(
            SplitBuf(&[1, 2, 3], &[4, 5, 6]).take_split(0),
            (SplitBuf(&[1, 2, 3], &[4, 5, 6]), SplitBuf(&[], &[]))
        );
        assert_eq!(
            SplitBuf(&[1, 2, 3], &[4, 5, 6]).take_split(1),
            (SplitBuf(&[2, 3], &[4, 5, 6]), SplitBuf(&[1], &[]))
        );
        assert_eq!(
            SplitBuf(&[1, 2, 3], &[4, 5, 6]).take_split(3),
            (SplitBuf(&[4, 5, 6], &[]), SplitBuf(&[1, 2, 3], &[]))
        );
        assert_eq!(
            SplitBuf(&[1, 2, 3], &[4, 5, 6]).take_split(4),
            (SplitBuf(&[5, 6], &[]), SplitBuf(&[1, 2, 3], &[4]))
        );
        assert_eq!(
            SplitBuf(&[1, 2, 3], &[4, 5, 6]).take_split(6),
            (SplitBuf(&[], &[]), SplitBuf(&[1, 2, 3], &[4, 5, 6]))
        );
    }

    #[test]
    fn split_buf_compare_str() {
        assert_eq!(
            SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36]).compare("34"),
            CompareResult::Error
        );
        assert_eq!(
            SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36]).compare("12"),
            CompareResult::Ok
        );
        assert_eq!(
            SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36]).compare("123"),
            CompareResult::Ok
        );
        assert_eq!(
            SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36]).compare("1234"),
            CompareResult::Ok
        );
        assert_eq!(
            SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36]).compare("12345"),
            CompareResult::Ok
        );
        assert_eq!(
            SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36]).compare("123456"),
            CompareResult::Ok
        );
        assert_eq!(
            SplitBuf(&[0x31, 0x32, 0x33], &[0x34, 0x35, 0x36]).compare("1234567"),
            CompareResult::Incomplete
        );
    }

    #[test]
    fn split_buf_tag() {
        assert_eq!(
            tag("def")(SplitBuf::from(&b"abc"[..])),
            IResult::Err(nom::Err::Error(((&b"abc"[..]).into(), nom::error::ErrorKind::Tag)))
        );

        assert_eq!(
            tag("abc")(SplitBuf::from(&b"abcdef"[..])),
            IResult::<_, _, ()>::Ok(((&b"def"[..]).into(), (&b"abc"[..]).into()))
        );

        assert_eq!(
            (tag("abc"), tag("def")).parse(SplitBuf::from(&b"abcdef"[..])),
            IResult::<_, _, ()>::Ok((
                (&b""[..]).into(),
                ((&b"abc"[..]).into(), (&b"def"[..]).into())
            ))
        );
    }
}