rust-stdf 1.1.0

//
// stdf_file.rs
// Author: noonchen - chennoon233@foxmail.com
// Created Date: October 3rd 2022
// -----
// Last Modified: Sun Aug 23 2026
// Modified By: noonchen
// -----
// Copyright (c) 2022 noonchen
//

use crate::records::*;
use crate::stdf_codec::*;
use crate::stdf_error::{StdfError, StdfErrorKind};
#[cfg(feature = "bzip")]
use bzip2::bufread::BzDecoder;
#[cfg(feature = "gzip")]
use flate2::bufread::GzDecoder;
use rust_stdf_derive::stdf_match_expr;
use std::io::{self, BufReader, SeekFrom}; // struct or enum
use std::io::{BufRead, Read, Seek, Write};
use std::{fs, path::Path}; // trait
#[cfg(feature = "zipfile")]
use zip::{read::ZipFile, ZipArchive};

/// `Unsafe` struct for coupling
/// file and `ZipArchive`
/// in order to get access to
/// the same `ZipFile`
#[cfg(feature = "zipfile")]
pub(crate) struct ZipBundle<R> {
    // put `ZipFile` on top
    // to ensure it is dropped
    // before `ZipArchive`
    file: Option<ZipFile<'static>>,
    archive: Box<ZipArchive<R>>,
}

#[allow(clippy::large_enum_variant)]
pub(crate) enum StdfStream<R> {
    Binary(R),
    #[cfg(feature = "gzip")]
    Gz(GzDecoder<R>),
    #[cfg(feature = "bzip")]
    Bz(BzDecoder<R>),
    #[cfg(feature = "zipfile")]
    Zip(ZipBundle<R>),
}

/// STDF Reader
///
/// This reader can process STDF datalogs of Version V4 and V4-2007
///
/// Supported compression:
///  - Uncompressed
///  - Gzip (.gz)
///  - Bzip (.bz2)
///
/// # Example
///
/// ```
/// use rust_stdf::{stdf_file::*, stdf_record_type::*, StdfRecord};
///
/// let stdf_path = "demo_file.stdf";
/// let mut reader = match StdfReader::new(&stdf_path) {
///     Ok(r) => r,
///     Err(e) => {
///         println!("{}", e);
///         return;
///     }
/// };
///
/// // we will count total DUT# in the file
/// // and put test result of PTR named
/// // "continuity test" in a vector.
/// let mut dut_count: u64 = 0;
/// let mut continuity_rlt = vec![];
///
/// // use type filter to work on certain types,
/// // use `|` to combine multiple typs
/// let rec_types = REC_PIR | REC_PTR;
/// // iterator starts from current file position,
/// // if file hits EOF, it will NOT redirect to 0.
/// for rec in reader
///     .get_record_iter()
///     .map(|x| x.unwrap())
///     .filter(|x| x.is_type(rec_types))
/// {
///     match rec {
///         StdfRecord::PIR(_) => {dut_count += 1;}
///         StdfRecord::PTR(ref ptr_rec) => {
///             if ptr_rec.test_txt == "continuity test" {
///                 continuity_rlt.push(ptr_rec.result);
///             }
///         }
///         _ => {}
///     }
/// }
/// println!("Total duts {} \n continuity result {:?}",
///         dut_count,
///         continuity_rlt);
/// ```
pub struct StdfReader<R> {
    endianness: ByteOrder,
    stream: StdfStream<R>,
}

/// STDF Writer
///
/// This writer accepts streams that implement [`Write`] trait,
/// and writes STDF records in given byte order.
///
/// Multiple APIs are provided to write STDF records, including:
/// - `write_record` for writing parsed STDF V4-2007 records, not including reserved/unknown records.
/// - `write_stdf_record` for writing any parsed STDF records.
/// - `write_stdf_record_view` for writing any STDF record views.
/// - `write_raw` for writing any unprocessed STDF records.
/// - `write_raw_view` for writing any unprocessed STDF record views.
///
/// ```
/// use rust_stdf::{stdf_file::StdfWriter, ByteOrder, FAR};
///
/// let mut far = FAR::new();
/// far.cpu_type = 2;
/// far.stdf_ver = 4;
///
/// let mut out = Vec::new();
/// let mut writer = StdfWriter::new(&mut out, ByteOrder::LittleEndian);
/// writer.write_record(&far).unwrap();
///
/// assert_eq!(out, [2, 0, 0, 10, 2, 4]);
/// ```
pub struct StdfWriter<W> {
    writer: W,
    byte_order: ByteOrder,
}

impl<W: Write> StdfWriter<W> {
    /// Create a new writer with the given output stream and byte order.
    #[inline(always)]
    pub fn new(writer: W, byte_order: ByteOrder) -> Self {
        StdfWriter { writer, byte_order }
    }

    /// Write a STDF V4-2007 record to the output stream.
    ///
    /// Validation will be performed on the record before writing, such as
    /// length check of Cn field, count check for Kx* fields, returns `StdfError`
    /// if validation fails.
    ///
    /// Note: `ReservedRec` and `UnknownRec` cannot be written with this method,
    /// use `write_stdf_record` instead.
    #[inline]
    #[allow(private_bounds)]
    pub fn write_record<R: StdfRecordWrite>(&mut self, r: &R) -> Result<(), StdfError> {
        r.validate()?;
        let n = r.payload_len();
        self.write_header(R::REC_TYP, R::REC_SUB, n)?;
        r.write_payload(&mut self.writer, &self.byte_order)
    }

    /// Write a STDF V4-2007 record to the output stream, also supports
    /// `ReservedRec` and `UnknownRec`.
    ///
    /// Validation will be performed on the record before writing.
    ///
    /// Note: reserved/unknown record must match the writer's byte order,
    /// otherwise, error [`StdfErrorKind::ByteOrderMismatch`] is returned.
    #[inline]
    pub fn write_stdf_record(&mut self, r: &StdfRecord) -> Result<(), StdfError> {
        stdf_match_expr!(record_write)
    }

    /// Similar to `write_stdf_record`, but takes a borrowed [`StdfRecordView`].
    ///
    /// If known record types have different byte order than the writer, they are re-encoded
    /// via implicit `to_owned()`. Reserved/unknown views cannot be re-encoded,
    /// so a byte order mismatch returns error
    /// [`StdfErrorKind::ByteOrderMismatch`].
    #[inline]
    pub fn write_stdf_record_view(&mut self, v: &StdfRecordView) -> Result<(), StdfError> {
        stdf_match_expr!(view_write)
    }

    /// Write an unprocessed record [`RawDataElement`] to the output stream.
    ///
    /// Byte order of the raw record must match the writer's.
    pub fn write_raw(&mut self, raw: &RawDataElement) -> Result<(), StdfError> {
        self.write_guarded_record(
            raw.header.typ,
            raw.header.sub,
            &raw.raw_data,
            raw.byte_order,
        )
    }

    /// Write an unprocessed record view [`RawDataElementView`] to the output stream.
    ///
    /// Byte order of the raw record must match the writer's.
    pub fn write_raw_view(&mut self, v: &RawDataElementView) -> Result<(), StdfError> {
        self.write_guarded_record(v.header.typ, v.header.sub, v.raw_data, v.byte_order)
    }

    /// Manually flush any buffered data to the output stream.
    ///
    /// `StdfWriter` never buffers internally, but a wrapping `BufWriter` does;
    /// use this method to force wrapped writers to flush.
    #[inline]
    pub fn flush(&mut self) -> Result<(), StdfError> {
        self.writer.flush()?;
        Ok(())
    }

    /// Check byte order before writing raw payload.
    #[inline]
    fn write_guarded_record(
        &mut self,
        typ: u8,
        sub: u8,
        raw_data: &[u8],
        record_order: ByteOrder,
    ) -> Result<(), StdfError> {
        if record_order != self.byte_order {
            return Err(StdfError::new(
                StdfErrorKind::ByteOrderMismatch,
                format!(
                    "record byte order is {:?}, writer uses {:?}",
                    record_order, self.byte_order
                ),
            ));
        }
        self.write_header_and_payload(typ, sub, raw_data)
    }

    #[inline]
    fn write_header_and_payload(
        &mut self,
        typ: u8,
        sub: u8,
        raw_data: &[u8],
    ) -> Result<(), StdfError> {
        self.write_header(typ, sub, raw_data.len())?;
        self.writer.write_all(raw_data)?;
        Ok(())
    }

    #[inline]
    fn write_header(&mut self, typ: u8, sub: u8, payload_len: usize) -> Result<(), StdfError> {
        if payload_len > u16::MAX as usize {
            return Err(StdfError::new(
                StdfErrorKind::RecordTooLarge,
                format!(
                    "record ({typ}, {sub}) payload length {payload_len} exceeds {}",
                    u16::MAX
                ),
            ));
        }
        let len = payload_len as u16;
        match self.byte_order {
            ByteOrder::LittleEndian => self.writer.write_all(&len.to_le_bytes())?,
            ByteOrder::BigEndian => self.writer.write_all(&len.to_be_bytes())?,
        }
        self.writer.write_all(&[typ, sub])?;
        Ok(())
    }
}

/// Iterator over [`StdfRecord`]s, created by [`StdfReader::get_record_iter`].
///
/// Each item is a fully parsed [`StdfRecord`].
pub struct RecordIter<'a, R> {
    inner: &'a mut StdfReader<R>,
    // reused across records, bytes are only borrowed during parsing
    buffer: Vec<u8>,
}

/// Iterator over unprocessed STDF records, created by [`StdfReader::get_rawdata_iter`].
///
/// Each item is an unprocessed [`RawDataElement`] that holds an owned record data buffer.
pub struct RawDataIter<'a, R> {
    offset: u64,
    inner: &'a mut StdfReader<R>,
}

/// Iterator over unprocessed STDF records that borrows the record bytes,
/// created by [`StdfReader::get_rawdata_view_iter`].
///
/// Each item is an unprocessed [`RawDataElementView`] that borrows record data buffer.
///
/// This iterator outperforms [`RawDataIter`] because there is no allocation, but
/// it cannot implement [`Iterator`], so it can only be driven by `next`.
///
/// The yielded item is only valid until the following `next` call.
///
/// # Example
///
/// ```no_run
/// use rust_stdf::{stdf_file::*, stdf_record_type::*, StdfRecordView};
///
/// let mut reader = StdfReader::new("demo_file.stdf").unwrap();
/// let mut iter = reader.get_rawdata_view_iter();
/// while let Some(item) = iter.next() {
///     let raw = item.unwrap();
///     if !raw.is_type(REC_PTR) {
///         continue;
///     }
///     let rec_view: StdfRecordView = (&raw).into();
///     if let StdfRecordView::PTR(ptr) = rec_view {
///         println!("{}", ptr.result());
///     }
/// }
/// ```
pub struct RawDataViewIter<'a, R> {
    offset: u64,
    inner: &'a mut StdfReader<R>,
    // reused across records, bytes are borrowed by the yielded view
    buffer: Vec<u8>,
}

// implementations

impl StdfReader<BufReader<fs::File>> {
    /// Open the given file and return a StdfReader, if successful
    #[inline(always)]
    pub fn new<P>(path: P) -> Result<Self, StdfError>
    where
        P: AsRef<Path>,
    {
        // determine the compress type by file extension
        let path_string = path.as_ref().display().to_string();
        let file_ext = path_string.rsplit('.').next();
        let compress_type = match file_ext {
            Some(ext) => match ext {
                #[cfg(feature = "gzip")]
                "gz" => CompressType::GzipCompressed,
                #[cfg(feature = "bzip")]
                "bz2" => CompressType::BzipCompressed,
                #[cfg(feature = "zipfile")]
                "zip" => CompressType::ZipCompressed,
                _ => CompressType::Uncompressed,
            },
            None => CompressType::Uncompressed,
        };
        let fp = fs::OpenOptions::new().read(true).open(path)?;
        let br = BufReader::with_capacity(2 << 20, fp);
        StdfReader::from(br, &compress_type)
    }
}

impl<R: BufRead + Seek> StdfReader<R> {
    /// Consume a input stream and generate a StdfReader, if successful
    #[inline(always)]
    pub fn from(in_stream: R, compress_type: &CompressType) -> Result<Self, StdfError> {
        let mut stream = match compress_type {
            #[cfg(feature = "gzip")]
            CompressType::GzipCompressed => StdfStream::Gz(GzDecoder::new(in_stream)),
            #[cfg(feature = "bzip")]
            CompressType::BzipCompressed => StdfStream::Bz(BzDecoder::new(in_stream)),
            #[cfg(feature = "zipfile")]
            CompressType::ZipCompressed => StdfStream::Zip(ZipBundle::new(in_stream, 0)?),
            _ => StdfStream::Binary(in_stream),
        };

        // read FAR header from file
        let mut buf = [0u8; 4];
        stream.read_exact(&mut buf)?;
        // parse header assuming little endian
        let far_header = RecordHeader::new().read_from_bytes(&buf, &ByteOrder::LittleEndian)?;
        let endianness = match far_header.len {
            2 => Ok(ByteOrder::LittleEndian),
            512 => Ok(ByteOrder::BigEndian),
            _ => Err(StdfError::new(
                StdfErrorKind::InvalidStdfFile,
                String::from("Cannot determine endianness"),
            )),
        }?;
        // check if it's FAR
        if (far_header.typ, far_header.sub) != (0, 10) {
            return Err(StdfError::new(
                StdfErrorKind::InvalidStdfFile,
                format!(
                    "FAR header (0, 10) expected, but {:?} is found",
                    (far_header.typ, far_header.sub)
                ),
            ));
        }
        // restore file position
        // current flate2 does not support fseek, we need to consume
        // old stream and create a new one.
        // If seek is supported, this function can be replaced by:
        //
        // stream.seek(SeekFrom::Start(0))?;
        //
        stream = rewind_stream_position(stream)?;

        Ok(StdfReader { endianness, stream })
    }

    /// Return the endianness of the STDF file.
    #[inline(always)]
    pub fn get_byte_order(&self) -> ByteOrder {
        self.endianness
    }

    #[inline(always)]
    fn read_header(&mut self) -> Result<RecordHeader, StdfError> {
        let mut buf = [0u8; 4];
        self.stream.read_exact(&mut buf)?;
        RecordHeader::new().read_from_bytes(&buf, &self.endianness)
    }

    /// Return an iterator for [`StdfRecord`].
    ///
    /// The iterator yields records starting from the current
    /// file position onward.
    #[inline(always)]
    pub fn get_record_iter(&mut self) -> RecordIter<'_, R> {
        RecordIter {
            inner: self,
            buffer: Vec::with_capacity(1024),
        }
    }

    /// Return an iterator for unprocessed STDF bytes [`RawDataElement`].
    ///
    /// The iterator yields records starting from the current file position
    /// onward, and the `offset` field is counted relative to that position.
    #[inline(always)]
    pub fn get_rawdata_iter(&mut self) -> RawDataIter<'_, R> {
        RawDataIter {
            offset: 0,
            inner: self,
        }
    }

    /// Return a buffer-reusing iterator for unprocessed STDF bytes [`RawDataElementView`].
    ///
    /// Unlike [`get_rawdata_iter`](Self::get_rawdata_iter), the yielded
    /// [`RawDataElementView`] borrows the iterator's reused buffer,
    /// so it performs no per-record allocation. See [`RawDataViewIter`] for usage.
    ///
    /// The iterator yields records starting from the current file position
    /// onward, and the `offset` field is counted relative to that position.
    #[inline(always)]
    pub fn get_rawdata_view_iter(&mut self) -> RawDataViewIter<'_, R> {
        RawDataViewIter {
            offset: 0,
            inner: self,
            buffer: Vec::with_capacity(1024),
        }
    }
}

#[cfg(feature = "zipfile")]
impl<R: BufRead + Seek> ZipBundle<R> {
    /// the following code is modified from this SO post:
    /// https://stackoverflow.com/questions/67823680/open-a-single-file-from-a-zip-archive-and-pass-on-as-read-instance/
    pub(crate) fn new(stream: R, file_index: usize) -> Result<ZipBundle<R>, StdfError> {
        let archive = ZipArchive::new(stream)?;
        let mut archive = Box::new(archive);

        // SAFETY: fakes a 'static borrow of `archive`. Sound because `archive` is
        // boxed (address stable across the move below), `file` drops before
        // `archive` (whose Drop reads it), and `archive` is never aliased while
        // `file` is alive (reopen_file clears `file` before re-borrowing).
        let file = unsafe {
            std::mem::transmute::<ZipFile<'_>, ZipFile<'static>>(archive.by_index(file_index)?)
        };
        Ok(ZipBundle {
            archive,
            file: Some(file),
        })
    }

    pub(crate) fn reopen_file(&mut self, file_index: usize) -> Result<(), StdfError> {
        // Drop the outstanding borrow before re-borrowing `archive`.
        self.file = None;
        // SAFETY: same invariants as `new`. `file` was just cleared above.
        let file = unsafe {
            std::mem::transmute::<ZipFile<'_>, ZipFile<'static>>(self.archive.by_index(file_index)?)
        };
        self.file = Some(file);
        Ok(())
    }
}

#[cfg(feature = "zipfile")]
impl<R: BufRead + Seek> Read for ZipBundle<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.file.as_mut().unwrap().read(buf)
    }
}

impl<R: BufRead + Seek> StdfStream<R> {
    #[cfg(feature = "atdf")]
    #[inline(always)]
    pub(crate) fn read_until(&mut self, delim: u8, buf: &mut Vec<u8>) -> io::Result<usize> {
        match self {
            StdfStream::Binary(bstream) => bstream.read_until(delim, buf),
            #[cfg(feature = "gzip")]
            StdfStream::Gz(gzstream) => general_read_until(gzstream, delim, buf),
            #[cfg(feature = "bzip")]
            StdfStream::Bz(bzstream) => general_read_until(bzstream, delim, buf),
            #[cfg(feature = "zipfile")]
            StdfStream::Zip(zipstream) => general_read_until(zipstream, delim, buf),
        }
    }
}

impl<R: BufRead + Seek> Read for StdfStream<R> {
    #[inline(always)]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self {
            StdfStream::Binary(bstream) => bstream.read(buf),
            #[cfg(feature = "gzip")]
            StdfStream::Gz(gzstream) => gzstream.read(buf),
            #[cfg(feature = "bzip")]
            StdfStream::Bz(bzstream) => bzstream.read(buf),
            #[cfg(feature = "zipfile")]
            StdfStream::Zip(zipstream) => zipstream.read(buf),
        }
    }
}

// impl<R: Seek> Seek for StdfStream<R> {
//     fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
//         match self {
//             StdfStream::Binary(bstream) => bstream.seek(pos),
//             // arm that does not support seek
//             _ => Ok(0),
//         }
//     }
// }

impl<R: BufRead + Seek> Iterator for RecordIter<'_, R> {
    type Item = Result<StdfRecord, StdfError>;

    #[inline(always)]
    fn next(&mut self) -> Option<Self::Item> {
        let header = match self.inner.read_header() {
            Ok(h) => h,
            Err(e) => {
                return match e.kind() {
                    // only 2 error will be returned by `read_header`
                    // Eof indicates normal EOF, UnexpectedEof indicates unexpected EOF
                    StdfErrorKind::Eof => None,
                    _ => Some(Err(e)),
                };
            }
        };
        // grow only when this record is larger than any seen so far, then
        // fill and parse just the [..len] prefix
        let len = header.len as usize;
        if self.buffer.len() < len {
            self.buffer.resize(len, 0);
        }
        if let Err(io_e) = self.inner.stream.read_exact(&mut self.buffer[..len]) {
            return Some(Err(StdfError::new(StdfErrorKind::Io, io_e.to_string())));
        }

        let mut rec = StdfRecord::new_from_header(&header);
        rec.read_from_bytes(&self.buffer[..len], &self.inner.endianness);
        Some(Ok(rec))
    }
}

impl<R: BufRead + Seek> Iterator for RawDataIter<'_, R> {
    type Item = Result<RawDataElement, StdfError>;

    #[inline(always)]
    fn next(&mut self) -> Option<Self::Item> {
        let header = match self.inner.read_header() {
            Ok(h) => h,
            Err(e) => {
                return match e.kind() {
                    // Eof indicates normal EOF, UnexpectedEof indicates unexpected EOF
                    StdfErrorKind::Eof => None,
                    _ => Some(Err(e)),
                };
            }
        };
        // advance position by 4 after reading a header successfully
        self.offset += 4;
        let data_offset = self.offset;
        // create a buffer to store record raw data
        let mut buffer = vec![0u8; header.len as usize];
        if let Err(io_e) = self.inner.stream.read_exact(&mut buffer) {
            return Some(Err(StdfError::new(StdfErrorKind::Io, io_e.to_string())));
        }
        self.offset += header.len as u64;

        Some(Ok(RawDataElement {
            offset: data_offset,
            header,
            raw_data: buffer,
            byte_order: self.inner.endianness,
        }))
    }
}

impl<R: BufRead + Seek> RawDataViewIter<'_, R> {
    /// Advance the iterator, returning the next [`RawDataElementView`], or `None` at EOF.
    ///
    /// The returned item will be invalidated by the following call.
    ///
    /// It mirrors [`Iterator::next`] but cannot be part of the [`Iterator`] trait
    /// because the item borrows from `self`.
    #[inline(always)]
    #[allow(clippy::should_implement_trait)] // lending iterator: item borrows `self`
    pub fn next(&mut self) -> Option<Result<RawDataElementView<'_>, StdfError>> {
        let header = match self.inner.read_header() {
            Ok(h) => h,
            Err(e) => {
                return match e.kind() {
                    // Eof indicates normal EOF, UnexpectedEof indicates unexpected EOF
                    StdfErrorKind::Eof => None,
                    _ => Some(Err(e)),
                };
            }
        };
        // advance position by 4 after reading a header successfully
        self.offset += 4;
        let data_offset = self.offset;
        // grow only when this record is larger than any seen so far, then
        // fill just the [..len] prefix and borrow it in the view
        let len = header.len as usize;
        if self.buffer.len() < len {
            self.buffer.resize(len, 0);
        }
        if let Err(io_e) = self.inner.stream.read_exact(&mut self.buffer[..len]) {
            return Some(Err(StdfError::new(StdfErrorKind::Io, io_e.to_string())));
        }
        self.offset += header.len as u64;

        Some(Ok(RawDataElementView {
            offset: data_offset,
            header,
            raw_data: &self.buffer[..len],
            byte_order: self.inner.endianness,
        }))
    }
}

// help functions

#[inline(always)]
pub(crate) fn rewind_stream_position<R: BufRead + Seek>(
    old_stream: StdfStream<R>,
) -> Result<StdfStream<R>, StdfError> {
    let new_stream = match old_stream {
        StdfStream::Binary(mut br) => {
            br.seek(SeekFrom::Start(0))?;
            StdfStream::Binary(br)
        }
        #[cfg(feature = "gzip")]
        StdfStream::Gz(gzr) => {
            // get the inner handle and create a new stream after seek
            let mut fp = gzr.into_inner();
            fp.seek(SeekFrom::Start(0))?;
            StdfStream::Gz(GzDecoder::new(fp))
        }
        #[cfg(feature = "bzip")]
        StdfStream::Bz(bzr) => {
            // get the inner handle and create a new stream after seek
            let mut fp = bzr.into_inner();
            fp.seek(SeekFrom::Start(0))?;
            StdfStream::Bz(BzDecoder::new(fp))
        }
        #[cfg(feature = "zipfile")]
        StdfStream::Zip(mut zipr) => {
            zipr.reopen_file(0)?;
            StdfStream::Zip(zipr)
        }
    };
    Ok(new_stream)
}

#[cfg(all(feature = "atdf", any(feature = "gzip", feature = "bzip",)))]
#[inline(always)]
fn general_read_until<T: Read>(r: &mut T, delim: u8, buf: &mut Vec<u8>) -> io::Result<usize> {
    let mut one_byte = [0u8; 1];
    let mut n: usize = 0;
    loop {
        // read one byte at a time
        match r.read(&mut one_byte) {
            Ok(num) => {
                if num == 0 {
                    // EOF reached
                    break;
                }
            }
            Err(e) => return Err(e),
        };
        buf.extend_from_slice(&one_byte);
        n += 1;
        // break at delimiter
        if delim == one_byte[0] {
            break;
        }
    }
    Ok(n)
}