greaseweazle 0.2.0

use crate::Ticks;
use num_enum::{IntoPrimitive, TryFromPrimitive};
use std::{error::Error, fmt, iter, mem};

/// A wrapper around a Greaseweazle flux stream, encoded in the format that the Greaseweazle sends
/// or receives.
///
/// # Encoding format
///
/// The encoding is a variable-length multibyte encoding, vaguely in the style of UTF-8. Shorter
/// flux values are encoded in fewer bytes than longer values. The stream is zero-terminated,
/// meaning that the value `0` is disallowed in the middle of the stream, and must be present as
/// the last byte of the stream.
///
/// - Flux values less than 250 are encoded with a single byte as-is.
/// - Flux values greater than or equal to 250, and less than 1525, are encoded with two bytes,
///   neither of which can be 0 or 255:
///   - first byte: `250 + (flux - 250) / 255`
///   - second byte: `1 + (flux - 250) % 255`
/// - Flux values greater than or equal to 1525 are encoded by a "space" flux-op, followed by a
///   regular flux value, typically with the value 249.
///
/// A flux-op is a special escape sequence that starts with the value 255, followed by a byte
/// indicating the type, and then a 28-bit operand value encoded into the 4 bytes after that.
/// - **1**: index pulse. The operand specifies the time delta since the most recent flux
///   transition or "space" flux-op, where the index pulse occurs.
///   Index pulses are only used when reading flux, and not when writing.
/// - **2**: empty space, with no flux activity for a prolonged period. The operand specifies the
///   time, which is then added to the first following regular flux value (and index pulse, if one
///   occurs first).
/// - **3**: astable flux pattern. This is used when writing to create a "no flux area" of
///   alternating flux. It is not used when reading.
///
/// The 28-bit operand is encoded as follows:
/// - byte 1: `1 | (value << 1) & 255`
/// - byte 2: `1 | (value >> 6) & 255`
/// - byte 3: `1 | (value >> 13) & 255`
/// - byte 4: `1 | (value >> 20) & 255`
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct EncodedFlux(Vec<u8>);

impl EncodedFlux {
    /// Creates an encoded flux stream from the provided iterator of flux timings.
    ///
    /// If a flux value in `flux_timings` is greater than the provided `nfa_threshold`, then
    /// it is replaced with a "no flux area": the Greaseweazle generates an alternating flux
    /// pattern, with a period of `nfa_period`. The values used by the official Greaseweazle host
    /// software are:
    /// - `nfa_threshold`: 150 microseconds.
    /// - `nfa_period`: 1.25 microseconds.
    ///
    /// In some cases, the Greaseweazle can have difficulty accurately writing the very end of the
    /// flux stream. The official Greaseweazle host software tries to compensate for this by adding
    /// a dummy flux value to the end of the stream, of about 100 microseconds. If you experience
    /// difficulties getting a correct write, you could try adding this dummy value as well.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use std::time::Duration;
    /// #
    /// # let mut gw = greaseweazle::Greaseweazle::new("foo")?;
    /// # let sample_freq = 0;
    /// # let flux_timings: Vec<Ticks> = Vec::new();
    /// use greaseweazle::{EncodedFlux, Ticks};
    ///
    /// let dummy = Ticks::from_duration(Duration::from_nanos(100_000), sample_freq);
    /// let nfa_threshold = Ticks::from_duration(Duration::from_nanos(150_000), sample_freq);
    /// let nfa_period = Ticks::from_duration(Duration::from_nanos(1_250), sample_freq);
    /// let encoded = EncodedFlux::new(
    ///     flux_timings.into_iter().chain(std::iter::once(dummy)),
    ///     nfa_threshold,
    ///     nfa_period,
    /// );
    /// gw.write_flux(&encoded, true, true, Ticks(0))?;
    /// # Ok::<(), greaseweazle::CommandError>(())
    /// ```
    pub fn new(
        flux_timings: impl IntoIterator<Item = Ticks>,
        nfa_threshold: Ticks,
        nfa_period: Ticks,
    ) -> Self {
        fn write_28bit(encoded: &mut Vec<u8>, value: u32) {
            assert!(
                value & 0xf0000000 == 0,
                "flux-op operand overflowed 28 bits"
            );
            encoded.extend(
                [value << 1, value >> 6, value >> 13, value >> 20]
                    .into_iter()
                    .map(|value| (1 | value as u8)),
            );
        }

        let mut bytes: Vec<u8> = Vec::new();

        for Ticks(flux) in flux_timings {
            if flux == 0 {
                continue;
            } else if flux > nfa_threshold.0 {
                bytes.push(255);
                bytes.push(FluxOp::Space.into());
                write_28bit(&mut bytes, flux);
                bytes.push(255);
                bytes.push(FluxOp::Astable.into());
                write_28bit(&mut bytes, nfa_period.into());
            } else if flux < 250 {
                bytes.push(flux as u8);
            } else if flux < 1525 {
                let low = ((flux - 250) % 255) as u8;
                let high = ((flux - 250) / 255) as u8;
                bytes.push(250 + high);
                bytes.push(1 + low);
            } else {
                bytes.push(255);
                bytes.push(FluxOp::Space.into());
                write_28bit(&mut bytes, flux - 249);
                bytes.push(249);
            }
        }

        bytes.push(0);
        Self(bytes)
    }

    /// Creates a flux stream directly from already-encoded bytes.
    ///
    /// The contents is not validated, so you could end up sending bad data to the Greaseweazle
    /// this way. Use with caution.
    #[inline]
    pub fn from_encoded_bytes(bytes: Vec<u8>) -> Self {
        Self(bytes)
    }

    /// Returns a reference to the raw encoded bytes, in Greaseweazle format.
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        self.0.as_slice()
    }

    /// Iteratively decodes the flux stream.
    ///
    /// The iterator yields `Ok` if a flux or index value was successfully decoded,
    /// `Err` if an error was encountered.
    /// The iterator can continue to be used after `Err` was returned; iteration will continue
    /// after the bytes that failed to decode.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # let mut gw = greaseweazle::Greaseweazle::new("foo")?;
    /// # let sample_freq = 0;
    /// use greaseweazle::Ticks;
    ///
    /// let encoded = gw.read_flux(Ticks(0), 2)?;
    /// let (flux, index) = encoded
    ///     .decode()
    ///     .map(|result| {
    ///          result.map(|decode_value| {
    ///              decode_value.map(|ticks| ticks.to_duration(sample_freq))
    ///          })
    ///      })
    ///      .collect::<Result<(Vec<_>, Vec<_>), _>>()?;
    /// # Ok::<(), Box<dyn std::error::Error>>(())
    /// ```
    pub fn decode(&self) -> impl Iterator<Item = Result<FluxDecodeValue<Ticks>, FluxDecodeError>> {
        let mut iter = self.0.iter().copied().enumerate();
        let mut flux_ticks: u32 = 0;
        let mut index_ticks: u32 = 0;
        let mut last_index_value: u32 = 0;

        iter::from_fn(
            move || -> Option<Result<FluxDecodeValue<Ticks>, FluxDecodeError>> {
                while let Some((i, byte)) = iter.next() {
                    macro_rules! next {
                        () => {
                            match iter.next() {
                                None => {
                                    return Some(Err(FluxDecodeError::UnexpectedEndOfStream {
                                        index: i,
                                    }))
                                }
                                Some((i, 0)) => {
                                    return Some(Err(FluxDecodeError::ZeroByte { index: i }))
                                }
                                Some((_, byte)) => byte,
                            }
                        };
                    }

                    if byte == 0 {
                        if iter.len() == 0 {
                            return None;
                        } else {
                            return Some(Err(FluxDecodeError::NotZeroTerminated));
                        }
                    } else if byte == 255 {
                        let byte = next!();
                        let bytes = [next!(), next!(), next!(), next!()];

                        let Ok(flux_op) = FluxOp::try_from(byte) else {
                            return Some(Err(FluxDecodeError::InvalidFluxOp {
                                index: i,
                                flux_op: byte,
                            }));
                        };

                        let bytes = bytes.map(|byte| u32::from(byte) & 254);
                        let value =
                            (bytes[0] >> 1) + (bytes[1] << 6) + (bytes[2] << 13) + (bytes[3] << 20);

                        match flux_op {
                            FluxOp::Index => {
                                index_ticks = index_ticks
                                    .checked_add(value)
                                    .expect("index value overflowed a u32")
                                    .checked_sub(last_index_value)
                                    .expect("last_index_value was too big!");
                                last_index_value = value;
                                return Some(Ok(FluxDecodeValue::Index(Ticks(mem::take(
                                    &mut index_ticks,
                                )))));
                            }
                            FluxOp::Space => {
                                flux_ticks = flux_ticks
                                    .checked_add(value)
                                    .expect("flux value overflowed a u32");
                                index_ticks = index_ticks
                                    .checked_add(value)
                                    .expect("index value overflowed a u32");
                            }
                            FluxOp::Astable => {
                                // Ignore the value and convert back to regular flux.
                                return Some(Ok(FluxDecodeValue::Flux(Ticks(mem::take(
                                    &mut flux_ticks,
                                )))));
                            }
                        }
                    } else {
                        let byte = u32::from(byte);
                        let value = if let Some(high) = byte.checked_sub(250) {
                            let byte = next!();
                            let low = u32::from(byte) - 1;
                            250 + high * 255 + low
                        } else {
                            byte
                        };

                        flux_ticks = flux_ticks
                            .checked_add(value)
                            .expect("flux value overflowed a u32");
                        index_ticks = index_ticks
                            .checked_add(value)
                            .expect("index value overflowed a u32");
                        return Some(Ok(FluxDecodeValue::Flux(Ticks(mem::take(&mut flux_ticks)))));
                    }
                }

                None
            },
        )
    }
}

/// Value that results from decoding a Greaseweazle-encoded flux byte stream.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FluxDecodeValue<T> {
    /// A flux timing value was decoded.
    Flux(T),

    /// An index timing value was decoded.
    Index(T),
}

impl<T> FluxDecodeValue<T> {
    /// Maps a `FluxDecodeValue<T>` to `FluxDecodeValue<U>` by applying a function to the
    /// contained value.
    #[inline]
    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> FluxDecodeValue<U> {
        match self {
            Self::Flux(value) => FluxDecodeValue::Flux(f(value)),
            Self::Index(value) => FluxDecodeValue::Index(f(value)),
        }
    }
}

impl<T> FromIterator<FluxDecodeValue<T>> for (Vec<T>, Vec<T>) {
    fn from_iter<I: IntoIterator<Item = FluxDecodeValue<T>>>(iter: I) -> Self {
        let mut flux_timings: Vec<T> = Vec::new();
        let mut index_timings: Vec<T> = Vec::new();

        for val in iter {
            match val {
                FluxDecodeValue::Flux(value) => flux_timings.push(value),
                FluxDecodeValue::Index(value) => index_timings.push(value),
            }
        }

        (flux_timings, index_timings)
    }
}

/// Error that can occur when decoding a Greaseweazle-encoded flux byte stream.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum FluxDecodeError {
    /// The last byte of the stream was not zero.
    NotZeroTerminated,

    /// A zero byte was encountered.
    ZeroByte {
        /// The index in the stream at which it was encountered.
        index: usize,
    },

    /// The end of the stream was reached in the middle of a multi-byte sequence.
    UnexpectedEndOfStream {
        /// The index in the stream at which the multi-byte sequence starts.
        index: usize,
    },

    /// An invalid flux-op type was encountered.
    InvalidFluxOp {
        /// The index in the stream at which the multi-byte sequence starts.
        index: usize,

        /// The invalid flux-op code.
        flux_op: u8,
    },
}

impl fmt::Display for FluxDecodeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            FluxDecodeError::NotZeroTerminated => {
                write!(f, "the last byte of the stream was not zero")
            }
            FluxDecodeError::ZeroByte { index } => {
                write!(f, "a zero byte was encountered at index {index}")
            }
            FluxDecodeError::UnexpectedEndOfStream { index } => {
                write!(
                    f,
                    "the end of the stream was reached in the middle of a multi-byte sequence \
                    starting at index {index}"
                )
            }
            FluxDecodeError::InvalidFluxOp { index, flux_op } => {
                write!(
                    f,
                    "an invalid flux-op type {flux_op} was encountered at index {index}"
                )
            }
        }
    }
}

impl Error for FluxDecodeError {}

#[derive(TryFromPrimitive, IntoPrimitive)]
#[repr(u8)]
enum FluxOp {
    /// Index pulse.
    /// Only used when reading; not generated when encoding and thus not sent to the Greaseweazle.
    Index = 1,

    /// Increments the tick count without generating any flux.
    Space = 2,

    /// Generates regular flux transitions at specified astable period.
    /// Only used when writing; never sent from the Greaseweazle, and converted to regular flux
    /// when decoding.
    Astable = 3,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{fs::File, io::Read, path::Path, time::Duration};

    #[test]
    fn decode_flux() {
        let bytes = read_u8s("test_files/decode_bytes.bin");
        let expected_flux = read_u32s("test_files/decode_flux.bin");
        let expected_index = read_u32s("test_files/decode_index.bin");

        let (decoded_flux, decoded_index) = EncodedFlux::from_encoded_bytes(bytes)
            .decode()
            .collect::<Result<_, _>>()
            .unwrap();

        for (i, (Ticks(decoded), expected)) in
            decoded_flux.iter().zip(expected_flux.iter()).enumerate()
        {
            if decoded != expected {
                panic!(
                    "decoded flux mismatch at index {i}\n\
                    decoded:  {:?}\n\
                    expected: {:?}",
                    decoded_flux[i..].chunks(8).next().unwrap(),
                    expected_flux[i..].chunks(8).next().unwrap(),
                );
            }
        }

        for (i, (Ticks(decoded), expected)) in
            decoded_index.iter().zip(expected_index.iter()).enumerate()
        {
            if decoded != expected {
                panic!(
                    "decoded index mismatch at index {i}\n\
                    decoded:  {:?}\n\
                    expected: {:?}",
                    decoded_index[i..].chunks(8).next().unwrap(),
                    expected_index[i..].chunks(8).next().unwrap(),
                );
            }
        }

        assert_eq!(
            decoded_flux.len(),
            expected_flux.len(),
            "decoded flux length mismatch"
        );
        assert_eq!(
            decoded_index.len(),
            expected_index.len(),
            "decoded index length mismatch"
        );
    }

    #[test]
    fn encode_flux() {
        let flux = read_u32s("test_files/encode_flux.bin");
        let expected_bytes = read_u8s("test_files/encode_bytes.bin");
        const SAMPLE_FREQ: u32 = 72_000_000;

        let encoded = EncodedFlux::new(
            flux.into_iter().map(Ticks),
            Ticks::from_duration(Duration::from_nanos(150_000), SAMPLE_FREQ),
            Ticks::from_duration(Duration::from_nanos(1_250), SAMPLE_FREQ),
        );
        let encoded_bytes = encoded.as_bytes();

        for (i, (encoded, expected)) in encoded_bytes.iter().zip(expected_bytes.iter()).enumerate()
        {
            if encoded != expected {
                panic!(
                    "encoded bytes mismatch at index {i}\n\
                    encoded:  {:?}\n\
                    expected: {:?}",
                    encoded_bytes[i..].chunks(8).next().unwrap(),
                    expected_bytes[i..].chunks(8).next().unwrap(),
                );
            }
        }

        assert_eq!(
            encoded.as_bytes().len(),
            expected_bytes.len(),
            "encoded bytes length mismatch"
        );
    }

    fn read_u8s(path: impl AsRef<Path>) -> Vec<u8> {
        let mut buf: Vec<u8> = Vec::new();
        File::open(path).unwrap().read_to_end(&mut buf).unwrap();
        buf
    }

    fn read_u32s(path: impl AsRef<Path>) -> Vec<u32> {
        read_u8s(path)
            .chunks_exact(size_of::<u32>())
            .map(|bytes| u32::from_le_bytes(bytes.try_into().unwrap()))
            .collect()
    }
}