nyanoblog/src/util/slice.rs

use std::ops::{Add, AddAssign, Not, Sub, SubAssign};

/// Helper utility for parsers operating on slices
pub struct SliceCursor<'a, T> {
    data: &'a [T],
    pos: Position,
    chop: usize,
}

/// Helper for the [`SliceCursor`] helper.
#[derive(Copy, Clone)]
struct Position {
    /// Always within -1 and `end` (both inclusive).
    pos: isize,
    /// Length of the slice (i.e. first out-of-bound index).
    end: usize,
}

impl<'a, T> Clone for SliceCursor<'a, T> {
    fn clone(&self) -> Self {
        Self {
            data: self.data,
            pos: self.pos,
            chop: self.chop,
        }
    }
}

impl<'a, T> SliceCursor<'a, T> {
    pub fn new(data: &'a [T]) -> Self {
        assert!(data.len() <= isize::MAX as usize);

        Self {
            data,
            pos: Position::new(data.len()),
            chop: 0,
        }
    }

    /// Return the item at the current position, if any.
    pub fn current(&self) -> Option<&'a T> {
        self.pos.index().map(|index| &self.data[index])
    }

    /// Advance the cursor to the next item and return that item.
    pub fn next(&mut self) -> Option<&'a T> {
        self.pos.advance().map(|index| &self.data[index])
    }

    /// Return the next item without advancing the cursor.
    pub fn peek(&self) -> Option<&'a T> {
        self.pos.next_index().map(|index| &self.data[index])
    }

    /// Reverse the cursor to the previous item and return that item.
    pub fn prev(&mut self) -> Option<&'a T> {
        self.pos.reverse().map(|index| &self.data[index])
    }

    /// Peek for the next item and advance the cursor if `predicate` is true.
    pub fn next_if<F>(&mut self, predicate: F) -> Option<&'a T>
    where
        F: FnOnce(&'a T) -> bool,
    {
        let next = self.peek()?;
        if predicate(next) {
            self.pos.advance();
            Some(next)
        } else {
            None
        }
    }

    /// Advance to the last item for which `predicate` is true and return a
    /// slice from the current position up to and including that last item.
    ///
    /// Besides the fact that it will not modify the chop position,
    /// this operation is functionally equivalent to:
    ///
    /// ```
    /// cursor.chop();
    /// while let Some(c) = cursor.peek() {
    ///     if predicate(c) {
    ///         cursor.next();
    ///     } else {
    ///         break;
    ///     }
    /// }
    /// let result = cursor.chop();
    /// ```
    pub fn next_while<F>(&mut self, mut predicate: F) -> &'a [T]
    where
        F: FnMut(&'a T) -> bool,
    {
        let start = self.pos.next_index_or_end();
        let len = self.data[start..]
            .iter()
            .enumerate()
            .find_map(|(i, v)| predicate(v).not().then_some(i))
            .unwrap_or(self.data.len() - start);
        let end = start + len;
        self.pos += len;
        &self.data[start..end]
    }

    /// Save the cursor's state and perform an arbitrary operation on it.
    /// If the operation failed (i.e. yielded `None`), restore the cursor's
    /// state.  Passes on the return value of `op`.
    ///
    /// `op` SHOULD NOT redefine the cursor unless you want buggy code.
    /// This is because the original slice is only restored if `op` succeeded.
    pub fn attempt<F, U>(&mut self, op: F) -> Option<U>
    where
        F: FnOnce(&mut Self) -> Option<U>,
    {
        let backup = self.save();
        let result = op(self);
        if result.is_none() {
            self.restore(backup);
        }
        result
    }

    /// Return a slice over all elements since the last time this method was called.
    /// If the cursor went backwards, the slice is empty.
    pub fn chop(&mut self) -> &'a [T] {
        let start = self.chop;
        let end = self.pos.next_index_or_end();

        let slice = if start < end {
            &self.data[start..end]
        } else {
            &[]
        };

        self.chop = end;
        slice
    }

    /// Return how many items remain in the cursor.
    /// This indicates how many times it is possible
    /// to call [`Self::next()`] before it returns `None`.
    pub fn remaining(&self) -> usize {
        self.data.len() - self.pos.next_index_or_end()
    }

    fn save(&self) -> Self {
        self.clone()
    }

    fn restore(&mut self, backup: Self) {
        self.data = backup.data;
        self.chop = backup.chop;
        self.pos = backup.pos;
    }
}

impl Position {
    pub fn new(end: usize) -> Position {
        assert!(end <= isize::MAX as usize);
        Position { pos: -1, end }
    }

    pub fn advance(&mut self) -> Option<usize> {
        *self += 1;
        self.index()
    }

    pub fn reverse(&mut self) -> Option<usize> {
        *self -= 1;
        self.index()
    }

    pub fn jump_to(&mut self, index: usize) {
        assert!(index <= self.end);
        self.pos = index as isize;
    }

    pub fn index(&self) -> Option<usize> {
        (0..self.end as isize)
            .contains(&self.pos)
            .then_some(self.pos as usize)
    }

    pub fn next_index(&self) -> Option<usize> {
        let next_pos = assert_usize(self.pos + 1);
        (next_pos < self.end).then_some(next_pos)
    }

    pub fn next_index_or_end(&self) -> usize {
        self.next_index().unwrap_or(self.end)
    }
}

impl Add<usize> for Position {
    type Output = Position;
    fn add(self, rhs: usize) -> Position {
        let rhs = assert_isize(rhs);
        let result = self.pos.saturating_add(rhs);
        Position {
            pos: result.min(self.end as isize),
            end: self.end,
        }
    }
}

impl AddAssign<usize> for Position {
    fn add_assign(&mut self, rhs: usize) {
        let rhs = assert_isize(rhs);
        self.pos = self.pos.saturating_add(rhs).min(self.end as isize);
    }
}

impl Sub<usize> for Position {
    type Output = Position;
    fn sub(self, rhs: usize) -> Position {
        let rhs = assert_isize(rhs);
        let result = self.pos.saturating_sub(rhs);
        Position {
            pos: result.max(-1),
            end: self.end,
        }
    }
}

impl SubAssign<usize> for Position {
    fn sub_assign(&mut self, rhs: usize) {
        let rhs = assert_isize(rhs);
        self.pos = self.pos.saturating_sub(rhs).max(-1);
    }
}

fn assert_usize(i: isize) -> usize {
    assert!(i >= 0);
    i as usize
}

fn assert_isize(u: usize) -> isize {
    assert!(u <= isize::MAX as usize);
    u as isize
}

#[cfg(test)]
mod tests {
    use crate::util::slice::SliceCursor;

    #[test]
    fn next_and_prev() {
        let data: Vec<u8> = (0..10).collect();
        let mut cursor = SliceCursor::new(&data);
        assert_eq!(cursor.remaining(), 10);

        for (i, v) in data.iter().enumerate() {
            assert_eq!(cursor.remaining(), data.len() - i);
            assert_eq!(cursor.next(), Some(v));
        }
        assert_eq!(cursor.remaining(), 0);
        assert!(cursor.next().is_none());
        assert_eq!(cursor.remaining(), 0);

        for (i, v) in data.iter().rev().enumerate() {
            assert_eq!(cursor.prev(), Some(v));
            assert_eq!(cursor.remaining(), i);
        }
        assert_eq!(cursor.prev(), None);
        assert_eq!(cursor.remaining(), 10);
    }

    #[test]
    fn chop() {
        let data: Vec<u8> = (0..10).collect();
        let mut cursor = SliceCursor::new(&data);

        assert_eq!(cursor.chop(), &data[0..0]);

        cursor.next();
        assert_eq!(cursor.chop(), &data[0..1]);

        for _ in 0..3 {
            cursor.next();
        }
        assert_eq!(cursor.chop(), &data[1..4]);

        cursor.prev();
        assert_eq!(cursor.chop(), &data[3..3]);

        while cursor.next().is_some() {}
        assert_eq!(cursor.chop(), &data[3..10]);

        for i in (0u8..10).rev() {
            assert_eq!(cursor.prev(), Some(&i));
        }
        assert!(cursor.prev().is_none());
        assert_eq!(cursor.chop(), &data[0..0]);

        while cursor.next().is_some() {}
        assert_eq!(cursor.chop(), &data[..]);
        assert_eq!(cursor.chop(), &data[10..10]);
    }

    #[test]
    fn predicates() {
        let data: Vec<u8> = (0..10).collect();
        let mut cursor = SliceCursor::new(&data);

        assert_eq!(cursor.next_if(|c| *c == 0), Some(&data[0]));
        assert_eq!(cursor.next_if(|c| *c == 0), None);

        assert_eq!(cursor.next_while(|c| *c < 5), &data[1..5]);
        assert_eq!(cursor.current(), Some(&4));
        assert_eq!(cursor.chop(), &data[0..5]);
    }

    #[test]
    fn attempt() {
        let data: Vec<u8> = (0..10).collect();
        let mut cursor = SliceCursor::new(&data);

        let result = cursor.attempt(|cursor| cursor.next().copied().filter(|c| *c == 0));
        assert_eq!(result, Some(0));
        assert_eq!(cursor.remaining(), 9);
        assert_eq!(cursor.current(), Some(&0));

        let result = cursor.attempt(|cursor| cursor.next().copied().filter(|c| *c == 0));
        assert_eq!(result, None);
        assert_eq!(cursor.remaining(), 9);
        assert_eq!(cursor.current(), Some(&0));

        let data2: Vec<u8> = (10..20).collect();

        let _: Option<()> = cursor.attempt(|cursor| {
            *cursor = SliceCursor::new(&data2);
            cursor.next();
            None
        });
        assert_eq!(cursor.current(), Some(&0));

        cursor.attempt(|cursor| {
            *cursor = SliceCursor::new(&data2);
            cursor.next();
            Some(())
        });
        assert_eq!(cursor.current(), Some(&10));
    }
}