rust_stdf/
stdf_file.rs

1//
2// stdf_file.rs
3// Author: noonchen - chennoon233@foxmail.com
4// Created Date: October 3rd 2022
5// -----
6// Last Modified: Mon Nov 14 2022
7// Modified By: noonchen
8// -----
9// Copyright (c) 2022 noonchen
10//
11
12use crate::stdf_error::StdfError;
13use crate::stdf_types::*;
14#[cfg(feature = "bzip")]
15use bzip2::bufread::BzDecoder;
16#[cfg(feature = "gzip")]
17use flate2::bufread::GzDecoder;
18use std::io::{self, BufReader, SeekFrom}; // struct or enum
19use std::io::{BufRead, Read, Seek};
20use std::{fs, path::Path}; // trait
21#[cfg(feature = "zipfile")]
22use zip::{read::ZipFile, ZipArchive};
23
24/// `Unsafe` struct for coupling
25/// file and `ZipArchive`
26/// in order to get access to
27/// the same `ZipFile`
28#[cfg(feature = "zipfile")]
29pub(crate) struct ZipBundle<R> {
30    // put `ZipFile` on top
31    // to ensure it is dropped
32    // before `ZipArchive`
33    file: Option<ZipFile<'static>>,
34    archive: Box<ZipArchive<R>>,
35}
36
37#[allow(clippy::large_enum_variant)]
38pub(crate) enum StdfStream<R> {
39    Binary(R),
40    #[cfg(feature = "gzip")]
41    Gz(GzDecoder<R>),
42    #[cfg(feature = "bzip")]
43    Bz(BzDecoder<R>),
44    #[cfg(feature = "zipfile")]
45    Zip(ZipBundle<R>),
46}
47
48/// STDF Reader
49///
50/// This reader can process STDF datalogs of Version V4 and V4-2007
51///
52/// Supported compression:
53///  - Uncompressed
54///  - Gzip (.gz)
55///  - Bzip (.bz2)
56///
57/// # Example
58///
59/// ```
60/// use rust_stdf::{stdf_file::*, stdf_record_type::*, StdfRecord};
61///
62/// let stdf_path = "demo_file.stdf";
63/// let mut reader = match StdfReader::new(&stdf_path) {
64///     Ok(r) => r,
65///     Err(e) => {
66///         println!("{}", e);
67///         return;
68///     }
69/// };
70///
71/// // we will count total DUT# in the file
72/// // and put test result of PTR named
73/// // "continuity test" in a vector.
74/// let mut dut_count: u64 = 0;
75/// let mut continuity_rlt = vec![];
76///
77/// // use type filter to work on certain types,
78/// // use `|` to combine multiple typs
79/// let rec_types = REC_PIR | REC_PTR;
80/// // iterator starts from current file position,
81/// // if file hits EOF, it will NOT redirect to 0.
82/// for rec in reader
83///     .get_record_iter()
84///     .map(|x| x.unwrap())
85///     .filter(|x| x.is_type(rec_types))
86/// {
87///     match rec {
88///         StdfRecord::PIR(_) => {dut_count += 1;}
89///         StdfRecord::PTR(ref ptr_rec) => {
90///             if ptr_rec.test_txt == "continuity test" {
91///                 continuity_rlt.push(ptr_rec.result);
92///             }
93///         }
94///         _ => {}
95///     }
96/// }
97/// println!("Total duts {} \n continuity result {:?}",
98///         dut_count,
99///         continuity_rlt);
100/// ```
101pub struct StdfReader<R> {
102    endianness: ByteOrder,
103    stream: StdfStream<R>,
104}
105
106pub struct RecordIter<'a, R> {
107    inner: &'a mut StdfReader<R>,
108}
109
110pub struct RawDataIter<'a, R> {
111    offset: u64,
112    inner: &'a mut StdfReader<R>,
113}
114
115// implementations
116
117impl StdfReader<BufReader<fs::File>> {
118    /// Open the given file and return a StdfReader, if successful
119    #[inline(always)]
120    pub fn new<P>(path: P) -> Result<Self, StdfError>
121    where
122        P: AsRef<Path>,
123    {
124        // determine the compress type by file extension
125        let path_string = path.as_ref().display().to_string();
126        let file_ext = path_string.rsplit('.').next();
127        let compress_type = match file_ext {
128            Some(ext) => match ext {
129                #[cfg(feature = "gzip")]
130                "gz" => CompressType::GzipCompressed,
131                #[cfg(feature = "bzip")]
132                "bz2" => CompressType::BzipCompressed,
133                #[cfg(feature = "zipfile")]
134                "zip" => CompressType::ZipCompressed,
135                _ => CompressType::Uncompressed,
136            },
137            None => CompressType::Uncompressed,
138        };
139        let fp = fs::OpenOptions::new().read(true).open(path)?;
140        let br = BufReader::with_capacity(2 << 20, fp);
141        StdfReader::from(br, &compress_type)
142    }
143}
144
145impl<R: BufRead + Seek> StdfReader<R> {
146    /// Consume a input stream and generate a StdfReader, if successful
147    #[inline(always)]
148    pub fn from(in_stream: R, compress_type: &CompressType) -> Result<Self, StdfError> {
149        let mut stream = match compress_type {
150            #[cfg(feature = "gzip")]
151            CompressType::GzipCompressed => StdfStream::Gz(GzDecoder::new(in_stream)),
152            #[cfg(feature = "bzip")]
153            CompressType::BzipCompressed => StdfStream::Bz(BzDecoder::new(in_stream)),
154            #[cfg(feature = "zipfile")]
155            CompressType::ZipCompressed => StdfStream::Zip(ZipBundle::new(in_stream, 0)?),
156            _ => StdfStream::Binary(in_stream),
157        };
158
159        // read FAR header from file
160        let mut buf = [0u8; 4];
161        stream.read_exact(&mut buf)?;
162        // parse header assuming little endian
163        let far_header = RecordHeader::new().read_from_bytes(&buf, &ByteOrder::LittleEndian)?;
164        let endianness = match far_header.len {
165            2 => Ok(ByteOrder::LittleEndian),
166            512 => Ok(ByteOrder::BigEndian),
167            _ => Err(StdfError {
168                code: 1,
169                msg: String::from("Cannot determine endianness"),
170            }),
171        }?;
172        // check if it's FAR
173        if (far_header.typ, far_header.sub) != (0, 10) {
174            return Err(StdfError {
175                code: 1,
176                msg: format!(
177                    "FAR header (0, 10) expected, but {:?} is found",
178                    (far_header.typ, far_header.sub)
179                ),
180            });
181        }
182        // restore file position
183        // current flate2 does not support fseek, we need to consume
184        // old stream and create a new one.
185        // If seek is supported, this function can be replaced by:
186        //
187        // stream.seek(SeekFrom::Start(0))?;
188        //
189        stream = rewind_stream_position(stream)?;
190
191        Ok(StdfReader { endianness, stream })
192    }
193
194    #[inline(always)]
195    fn read_header(&mut self) -> Result<RecordHeader, StdfError> {
196        let mut buf = [0u8; 4];
197        self.stream.read_exact(&mut buf)?;
198        RecordHeader::new().read_from_bytes(&buf, &self.endianness)
199    }
200
201    /// return an iterator for StdfRecord
202    ///
203    /// Only the records after the current file position
204    /// can be read.
205    #[inline(always)]
206    pub fn get_record_iter(&mut self) -> RecordIter<R> {
207        RecordIter { inner: self }
208    }
209
210    /// return an iterator for unprocessed STDF bytes
211    ///
212    /// beware that internal `offset` counter is starting
213    /// from the current position.
214    #[inline(always)]
215    pub fn get_rawdata_iter(&mut self) -> RawDataIter<R> {
216        RawDataIter {
217            offset: 0,
218            inner: self,
219        }
220    }
221}
222
223#[cfg(feature = "zipfile")]
224impl<R: BufRead + Seek> ZipBundle<R> {
225    /// the following code is modified from this SO post:
226    /// https://stackoverflow.com/questions/67823680/open-a-single-file-from-a-zip-archive-and-pass-on-as-read-instance/
227    pub(crate) fn new(stream: R, file_index: usize) -> Result<ZipBundle<R>, StdfError> {
228        let archive = ZipArchive::new(stream)?;
229        let mut archive = Box::new(archive);
230
231        let file =
232            unsafe { std::mem::transmute::<_, ZipFile<'static>>(archive.by_index(file_index)?) };
233        Ok(ZipBundle {
234            archive,
235            file: Some(file),
236        })
237    }
238
239    pub(crate) fn reopen_file(&mut self, file_index: usize) -> Result<(), StdfError> {
240        self.file = None;
241        let file = unsafe {
242            std::mem::transmute::<_, ZipFile<'static>>(self.archive.by_index(file_index)?)
243        };
244        self.file = Some(file);
245        Ok(())
246    }
247}
248
249#[cfg(feature = "zipfile")]
250impl<R: BufRead + Seek> Read for ZipBundle<R> {
251    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
252        self.file.as_mut().unwrap().read(buf)
253    }
254}
255
256impl<R: BufRead + Seek> StdfStream<R> {
257    #[cfg(feature = "atdf")]
258    #[inline(always)]
259    pub(crate) fn read_until(&mut self, delim: u8, buf: &mut Vec<u8>) -> io::Result<usize> {
260        match self {
261            StdfStream::Binary(bstream) => bstream.read_until(delim, buf),
262            #[cfg(feature = "gzip")]
263            StdfStream::Gz(gzstream) => general_read_until(gzstream, delim, buf),
264            #[cfg(feature = "bzip")]
265            StdfStream::Bz(bzstream) => general_read_until(bzstream, delim, buf),
266            #[cfg(feature = "zipfile")]
267            StdfStream::Zip(zipstream) => general_read_until(zipstream, delim, buf),
268        }
269    }
270}
271
272impl<R: BufRead + Seek> Read for StdfStream<R> {
273    #[inline(always)]
274    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
275        match self {
276            StdfStream::Binary(bstream) => bstream.read(buf),
277            #[cfg(feature = "gzip")]
278            StdfStream::Gz(gzstream) => gzstream.read(buf),
279            #[cfg(feature = "bzip")]
280            StdfStream::Bz(bzstream) => bzstream.read(buf),
281            #[cfg(feature = "zipfile")]
282            StdfStream::Zip(zipstream) => zipstream.read(buf),
283        }
284    }
285}
286
287// impl<R: Seek> Seek for StdfStream<R> {
288//     fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
289//         match self {
290//             StdfStream::Binary(bstream) => bstream.seek(pos),
291//             // arm that does not support seek
292//             _ => Ok(0),
293//         }
294//     }
295// }
296
297impl<R: BufRead + Seek> Iterator for RecordIter<'_, R> {
298    type Item = Result<StdfRecord, StdfError>;
299
300    #[inline(always)]
301    fn next(&mut self) -> Option<Self::Item> {
302        let header = match self.inner.read_header() {
303            Ok(h) => h,
304            Err(e) => {
305                return match e.code {
306                    // only 2 error will be returned by `read_header`
307                    // code = 4, indicates normal EOF
308                    4 => None,
309                    // code = 5, indicates unexpected EOF
310                    _ => Some(Err(e)),
311                };
312            }
313        };
314        // create a buffer to store record raw data
315        let mut buffer = vec![0u8; header.len as usize];
316        if let Err(io_e) = self.inner.stream.read_exact(&mut buffer) {
317            return Some(Err(StdfError {
318                code: 3,
319                msg: io_e.to_string(),
320            }));
321        }
322
323        let mut rec = StdfRecord::new_from_header(header);
324        rec.read_from_bytes(&buffer, &self.inner.endianness);
325        Some(Ok(rec))
326    }
327}
328
329impl<R: BufRead + Seek> Iterator for RawDataIter<'_, R> {
330    type Item = Result<RawDataElement, StdfError>;
331
332    #[inline(always)]
333    fn next(&mut self) -> Option<Self::Item> {
334        let header = match self.inner.read_header() {
335            Ok(h) => h,
336            Err(e) => {
337                return match e.code {
338                    // code = 4, indicates normal EOF
339                    4 => None,
340                    // code = 5, indicates unexpected EOF
341                    _ => Some(Err(e)),
342                };
343            }
344        };
345        // advance position by 4 after reading a header successfully
346        self.offset += 4;
347        let data_offset = self.offset;
348        // create a buffer to store record raw data
349        let mut buffer = vec![0u8; header.len as usize];
350        if let Err(io_e) = self.inner.stream.read_exact(&mut buffer) {
351            return Some(Err(StdfError {
352                code: 3,
353                msg: io_e.to_string(),
354            }));
355        }
356        self.offset += header.len as u64;
357
358        Some(Ok(RawDataElement {
359            offset: data_offset,
360            header,
361            raw_data: buffer,
362            byte_order: self.inner.endianness,
363        }))
364    }
365}
366
367// help functions
368
369#[inline(always)]
370pub(crate) fn rewind_stream_position<R: BufRead + Seek>(
371    old_stream: StdfStream<R>,
372) -> Result<StdfStream<R>, StdfError> {
373    let new_stream = match old_stream {
374        StdfStream::Binary(mut br) => {
375            br.seek(SeekFrom::Start(0))?;
376            StdfStream::Binary(br)
377        }
378        #[cfg(feature = "gzip")]
379        StdfStream::Gz(gzr) => {
380            // get the inner handle and create a new stream after seek
381            let mut fp = gzr.into_inner();
382            fp.seek(SeekFrom::Start(0))?;
383            StdfStream::Gz(GzDecoder::new(fp))
384        }
385        #[cfg(feature = "bzip")]
386        StdfStream::Bz(bzr) => {
387            // get the inner handle and create a new stream after seek
388            let mut fp = bzr.into_inner();
389            fp.seek(SeekFrom::Start(0))?;
390            StdfStream::Bz(BzDecoder::new(fp))
391        }
392        #[cfg(feature = "zipfile")]
393        StdfStream::Zip(mut zipr) => {
394            zipr.reopen_file(0)?;
395            StdfStream::Zip(zipr)
396        }
397    };
398    Ok(new_stream)
399}
400
401#[cfg(all(feature = "atdf", any(feature = "gzip", feature = "bzip",)))]
402#[inline(always)]
403fn general_read_until<T: Read>(r: &mut T, delim: u8, buf: &mut Vec<u8>) -> io::Result<usize> {
404    let mut one_byte = [0u8; 1];
405    let mut n: usize = 0;
406    loop {
407        // read one byte at a time
408        match r.read(&mut one_byte) {
409            Ok(num) => {
410                if num == 0 {
411                    // EOF reached
412                    break;
413                }
414            }
415            Err(e) => return Err(e),
416        };
417        buf.extend_from_slice(&one_byte);
418        n += 1;
419        // break at delimiter
420        if delim == one_byte[0] {
421            break;
422        }
423    }
424    Ok(n)
425}