faf_replay_parser/iter/

mod.rs

//! Iterators for composing flexible replay parsing pipelines
//!
//! Traits from this module are automatically implemented and are made available by star importing
//! from [`prelude`].
//!
//! ```rust
//! use faf_replay_parser::iter::prelude::*;
//! ```
//!
//! This module defines three traits for creating composable replay command parsing pipelines. Note
//! that you can't parse replay headers this way, only the command stream. Each pipeline will start
//! with one of the entrypoint methods defined by [`CommandIter`] depending on the level of parsing
//! that the pipeline needs as a starting point.
//!
//! The pipeline can then use the standard library iterator features such as `.filter` and `.map`,
//! or even third party iterator adapters, to define whatever processing logic is desired. This
//! module also contains iterator adapters for parsing the command frames in case the pipeline
//! needs to start from the raw bytes in order to inject custom logic before the parsing step.
//! These adapter methods are defined by [`CommandFrameAdapters`].
//!
//! There are also some iterator adapters to simplify the handling of errors for common cases,
//! since the parsing functions return result types. These are implemented by [`ResultAdapters`].

use std::fmt::Debug;

use crate::replay::ReplayCommandFrameSpan;
use crate::version::Version;

use byteorder::{ByteOrder, LittleEndian};

mod parse;
pub mod prelude;
mod verify;
mod while_;

pub use parse::*;
pub use verify::*;
pub use while_::*;

/// Allow objects to return iterators over contained command frames.
pub trait CommandIter<'a> {
    type Iter: Iterator<Item = ReplayCommandFrameSpan<'a>>;

    /// Return an iterator over the raw command frame data. The iterator should not validate the
    /// contents of the command identifier or the command data. It is up to the implementor to
    /// decide how to handle frames with sizes less than the frame size.
    ///
    /// # Examples
    /// ```rust
    /// use faf_replay_parser::iter::prelude::*;
    /// use faf_replay_parser::iter::{parse_command, verify_frame_header};
    /// use faf_replay_parser::SCFA;
    /// use faf_replay_parser::scfa::ReplayCommand::*;
    ///
    /// let data: &[u8] = &[
    ///     1, 4, 0, 1,          // SetCommandSource
    ///     0, 7, 0, 1, 0, 0, 0, // Advance
    ///     5, 3, 0              // RequestPause
    /// ];
    ///
    /// let mut iter = data.iter_command_frames_raw()
    ///     .filter(|frame| frame.cmd < 3)
    ///     .filter_map(|frame| verify_frame_header::<SCFA>(frame).ok())
    ///     .filter_map(|frame| parse_command::<SCFA>(frame).ok());
    ///
    /// match iter.next() {
    ///     Some(SetCommandSource { id: 1 }) => {},
    ///     e => panic!("wrong result {:?}", e)
    /// }
    /// match iter.next() {
    ///     Some(Advance { ticks: 1 }) => {},
    ///     e => panic!("wrong result {:?}", e)
    /// }
    /// assert!(iter.next().is_none());
    ///
    /// ```
    fn iter_command_frames_raw(&self) -> Self::Iter;

    /// Return an iterator over the command frames. The command frame header will be checked for
    /// validity, and any errors returned.
    ///
    /// **Note**: If only a subset of commands are needed, it can be more efficient to compose
    /// `iter_command_frames_raw` with a filter and map adapter to call [`verify_frame_header`]
    /// only for frames that have the desired command id.
    fn iter_command_frames<V: Version>(&self) -> VerifyFrame<'a, V, Self::Iter> {
        VerifyFrame::new(self.iter_command_frames_raw())
    }

    /// Return an iterator over the parsed `Command`s, verifying the command headers and
    /// returning any errors.
    ///
    /// # Examples
    /// ```rust
    /// use faf_replay_parser::iter::prelude::*;
    /// use faf_replay_parser::SCFA;
    /// use faf_replay_parser::{ReplayResult, scfa::ReplayCommand};
    ///
    /// let data: &[u8] = &[0, 7, 0, 1, 0, 0, 0, 1, 0, 0, 23, 3, 0];
    /// let commands = data.iter_commands::<SCFA>()
    ///     .collect::<Vec<ReplayResult<ReplayCommand>>>();
    ///
    /// assert_eq!(
    ///     format!("{:?}", commands),
    ///     "[Ok(Advance { ticks: 1 }), Err(Malformed(\"invalid command size\")), Ok(EndGame)]"
    /// )
    /// ```
    ///
    /// **Note**: If only a subset of commands are needed, it can be more efficient to compose
    /// `iter_command_frames` with a filter and map adapter to call [`parse_command`]
    /// only for frames that have the desired command id.
    fn iter_commands<V: Version>(&self) -> ParseCommand<'a, V, Self::Iter> {
        ParseCommand::new(self.iter_command_frames_raw())
    }
}

/// Adapters for iterators over command frames.
pub trait CommandFrameAdapters<'a>: Iterator<Item = ReplayCommandFrameSpan<'a>> + Sized {
    /// Verify the validity of the frame headers. See [`verify_frame_header`].
    fn verify_frame<V: Version>(self) -> VerifyFrame<'a, V, Self> {
        VerifyFrame::new(self)
    }

    /// Parse the command data into `Command`s without verifying the frame headers. It is
    /// recommended to always include
    /// [`verify_frame`](trait.CommandFrameAdapters.html#method.verify_frame) earlier in the
    /// pipeline.
    fn parse_command_data<V: Version>(self) -> ParseCommandData<'a, V, Self> {
        ParseCommandData::new(self)
    }

    /// Parse the command data into `Command`s. See [`parse_command`].
    fn parse_command<V: Version>(self) -> ParseCommand<'a, V, Self> {
        ParseCommand::new(self)
    }
}

impl<'a, I: Iterator<Item = ReplayCommandFrameSpan<'a>> + Sized> CommandFrameAdapters<'a> for I {}

/// Adapters for iterators over results.
pub trait ResultAdapters<T, U>: Iterator<Item = Result<T, U>> + Sized {
    /// Take items while they are `Ok` and unwrap them. This is equivalent to:
    ///
    /// ```rust
    /// let list = [Ok(0), Err("foo")];
    /// list.iter()
    ///     .take_while(|result| result.is_ok())
    ///     .map(|result| result.unwrap());
    /// ```
    fn while_ok(self) -> WhileOk<Self> {
        WhileOk::new(self)
    }

    /// Take items while they are `Err`s and unwrap them. This is equivalent to:
    ///
    /// ```rust
    /// let list = [Ok(0), Err("foo")];
    /// list.iter()
    ///     .take_while(|result| result.is_err())
    ///     .map(|result| result.unwrap_err());
    /// ```
    fn while_err(self) -> WhileErr<Self> {
        WhileErr::new(self)
    }
}

impl<T, U: Debug, I: Iterator<Item = Result<T, U>>> ResultAdapters<T, U> for I {}

// ------- Implementation for byte slices ------- //

impl<'a> CommandIter<'a> for &'a [u8] {
    type Iter = BytesCommandFrameRawIterator<'a>;

    fn iter_command_frames_raw(&self) -> BytesCommandFrameRawIterator<'a> {
        BytesCommandFrameRawIterator::new(self)
    }
}

/// Implements [`CommandIter`] for byte slices
pub struct BytesCommandFrameRawIterator<'a> {
    data: &'a [u8],
}
impl<'a> BytesCommandFrameRawIterator<'a> {
    pub fn new(data: &'a [u8]) -> BytesCommandFrameRawIterator<'a> {
        BytesCommandFrameRawIterator { data }
    }
}
impl<'a> Iterator for BytesCommandFrameRawIterator<'a> {
    type Item = ReplayCommandFrameSpan<'a>;

    // Inlining here seems to be BAD for performance?
    fn next(&mut self) -> Option<Self::Item> {
        let frame = get_frame_raw(self.data)?;
        self.data = &self.data[frame.data.len()..];
        Some(frame)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        // Minimum command frame size is 3 bytes
        (0, Some(self.data.len() / 3))
    }
}

/// Get a command frame from the start of a buffer if enough data is available.
fn get_frame_raw(data: &[u8]) -> Option<ReplayCommandFrameSpan> {
    if data.len() < 3 {
        return None;
    }
    let cmd = unsafe { *data.get_unchecked(0) };
    let size = unsafe { LittleEndian::read_u16(data.get_unchecked(1..3)) as usize };
    if data.len() < size {
        return None;
    }
    // In case sizes are invalid (less than 3)
    let amount = std::cmp::max(size, 3);

    unsafe {
        Some(ReplayCommandFrameSpan {
            cmd,
            size: size as u16,
            data: &data.get_unchecked(..amount),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    use crate::reader::ReplayReadError;
    use crate::scfa::SCFA;
    use crate::scfa::{replay_command, ReplayCommand};

    #[test]
    fn test_basic() {
        let data: &[u8] = &[
            0, 3, 0, // Frame with no data
            100, 0, 0, // Size field is zero
            3, 5, 0, 0, 0, // Frame with 2 bytes of data
        ];

        let mut raw_iter = data.iter_command_frames_raw();
        assert_eq!(
            raw_iter.next(),
            Some(ReplayCommandFrameSpan {
                cmd: 0,
                size: 3,
                data: &data[0..3]
            })
        );
        assert_eq!(
            raw_iter.next(),
            Some(ReplayCommandFrameSpan {
                cmd: 100,
                size: 0,
                data: &data[3..6]
            })
        );
        assert_eq!(
            raw_iter.next(),
            Some(ReplayCommandFrameSpan {
                cmd: 3,
                size: 5,
                data: &data[6..11]
            })
        );
        assert_eq!(raw_iter.next(), None);

        let mut iter = data.iter_command_frames::<SCFA>();
        assert_eq!(
            iter.next().unwrap().unwrap(),
            ReplayCommandFrameSpan {
                cmd: 0,
                size: 3,
                data: &data[0..3]
            }
        );
        assert!(iter.next().unwrap().is_err());
        assert_eq!(
            iter.next().unwrap().unwrap(),
            ReplayCommandFrameSpan {
                cmd: 3,
                size: 5,
                data: &data[6..11]
            }
        );
        assert!(iter.next().is_none());
    }

    #[test]
    fn test_basic_commands() {
        let data: &[u8] = &[
            0, 3, 0, // Frame with no data
            23, 0, 0, // Size field is zero
            0, 7, 0, 1, 0, 0, 0, // Advance command
        ];

        let mut iter = data.iter_commands::<SCFA>();

        match iter.next().unwrap().unwrap_err() {
            ReplayReadError::IO(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
                assert_eq!(format!("{}", e), "failed to fill whole buffer");
            }
            _ => panic!("wrong error type"),
        };
        match iter.next().unwrap().unwrap_err() {
            ReplayReadError::Malformed(msg) => assert_eq!(msg, "invalid command size"),
            _ => panic!("wrong error type"),
        };
        match iter.next().unwrap().unwrap() {
            ReplayCommand::Advance { ticks: 1 } => {}
            _ => panic!("wrong command type or data"),
        }
        assert!(iter.next().is_none());
    }

    /// An example of composing a custom data parser out of iterator adapters from `std`.
    #[test]
    fn test_std_adapters() {
        use ReplayCommand::Advance;

        let data: &[u8] = &[
            0, 7, 0, 1, 0, 0, 0, // Ok
            0, 3, 0, // Malformed, not enough data
            0, 7, 0, 1, 0, 0, 0, // Ok
        ];

        // Parse all `Advance` commands stopping on the first malformed command
        let mut iter = data
            .iter_command_frames_raw()
            .filter(|frame| frame.cmd == replay_command::ADVANCE)
            .map(verify_frame_header::<SCFA>)
            .take_while(|r| r.is_ok())
            .map(|r| parse_command::<SCFA>(r.unwrap()))
            .take_while(|r| r.is_ok())
            .map(|r| r.unwrap());

        match iter.next() {
            Some(Advance { ticks: 1 }) => {}
            e => panic!("wrong result {:?}", e),
        }
        assert_eq!(iter.next(), None);
    }

    /// An example of composing a custom data parser using the provided iterator adapters.
    #[test]
    fn test_adapters() {
        use ReplayCommand::Advance;

        let data: &[u8] = &[
            0, 7, 0, 1, 0, 0, 0, // Ok
            0, 3, 0, // Malformed, not enough data
            0, 7, 0, 1, 0, 0, 0, // Ok
        ];

        // Parse all `Advance` commands stopping on the first malformed command
        let mut iter = data
            .iter_command_frames_raw()
            .filter(|frame| frame.cmd == replay_command::ADVANCE)
            .verify_frame::<SCFA>()
            .while_ok()
            .parse_command_data::<SCFA>()
            .while_ok();

        match iter.next() {
            Some(Advance { ticks: 1 }) => {}
            e => panic!("wrong result {:?}", e),
        }
        assert_eq!(iter.next(), None);
    }
}