1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
//
// atdf_file.rs
// Author: noonchen - chennoon233@foxmail.com
// Created Date: October 6th 2022
// -----
// Last Modified: Wed Nov 02 2022
// Modified By: noonchen
// -----
// Copyright (c) 2022 noonchen
//

use crate::atdf_types::AtdfRecord;
use crate::stdf_error::StdfError;
use crate::stdf_file::{rewind_stream_position, StdfStream};
use crate::stdf_types::{bytes_to_string, CompressType};
use bzip2::bufread::BzDecoder;
use flate2::bufread::GzDecoder;
use std::io::{BufRead, BufReader, Seek};
use std::{fs, mem, str};

pub struct AtdfReader<R> {
    delimiter: char,
    scale_flag: bool,
    stream: StdfStream<R>,
}

pub struct AtdfRecordIter<'a, R> {
    inner: &'a mut AtdfReader<R>,
    // ATDF record might be divided
    // into multiple lines.
    incomplete_rec: String,
}

// impl

impl AtdfReader<BufReader<fs::File>> {
    #[inline(always)]
    pub fn new(path: &str) -> Result<Self, StdfError> {
        // determine the compress type by file extension
        let compress_type = if path.ends_with(".gz") {
            CompressType::GzipCompressed
        } else if path.ends_with(".bz2") {
            CompressType::BzipCompressed
        } else if path.ends_with(".zip") {
            CompressType::ZipCompressed
        } else {
            CompressType::Uncompressed
        };

        let fp = fs::OpenOptions::new().read(true).open(path)?;
        let br = BufReader::with_capacity(2 << 20, fp);
        AtdfReader::from(br, &compress_type)
    }
}

impl<R: BufRead + Seek> AtdfReader<R> {
    #[inline(always)]
    pub fn from(in_stream: R, compress_type: &CompressType) -> Result<Self, StdfError> {
        let mut stream = match compress_type {
            CompressType::GzipCompressed => StdfStream::Gz(GzDecoder::new(in_stream)),
            CompressType::BzipCompressed => StdfStream::Bz(BzDecoder::new(in_stream)),
            _ => StdfStream::Binary(in_stream),
        };

        let mut far_bytes = vec![];
        stream.read_until(b'\n', &mut far_bytes)?;
        let far_str = bytes_to_string(&far_bytes);
        if !far_str.starts_with("FAR:A") || far_bytes.len() < 9 {
            return Err(StdfError {
                code: 6,
                msg: format!(
                    "FAR record pattern 'FAR:A' not detected or required fields missing, found {}",
                    far_str
                ),
            });
        }
        // according to atdf spec, delimiter is the byte after 'A'
        let delimiter = far_bytes[5] as char;
        // parametric scale flag, default is false
        let scale_flag = {
            let far_str_vec: Vec<_> = far_str.split(delimiter).collect();
            if far_str_vec.len() > 3 {
                far_str_vec[3] == "S"
            } else {
                false
            }
        };
        // reset file position
        stream = rewind_stream_position(stream)?;

        Ok(AtdfReader {
            delimiter,
            scale_flag,
            stream,
        })
    }

    #[inline(always)]
    pub fn get_record_iter(&mut self) -> AtdfRecordIter<R> {
        AtdfRecordIter {
            inner: self,
            incomplete_rec: String::new(),
        }
    }
}

// implement of ATDF iterator

impl<R: BufRead + Seek> Iterator for AtdfRecordIter<'_, R> {
    type Item = AtdfRecord;

    #[inline(always)]
    fn next(&mut self) -> Option<Self::Item> {
        // if next_rec is empty, means
        // the previous rec is not completed yet
        loop {
            // read a line
            let mut tmp_buf = Vec::with_capacity(127);
            let eof = match self.inner.stream.read_until(b'\n', &mut tmp_buf) {
                Ok(n) => n == 0,
                Err(e) => {
                    println!("Error when reading ATDF file => {}", e);
                    return None;
                }
            };

            let tmp_line = match str::from_utf8(&tmp_buf) {
                Ok(s) => s,
                Err(_) => {
                    println!("String error: ATDF should only contains ascii symbols, ");
                    return None;
                }
            };

            if !tmp_line.is_empty() && tmp_line.starts_with(' ') {
                // starts with space, means it belongs to incomplete_rec
                // remove prefix space and suffix \n
                self.incomplete_rec.push_str(str_trim(tmp_line));
                // directly goes to the next loop iteration
                continue;
            }

            // not starts with space, trim \r\n first
            let clean_line = str_trim(tmp_line);
            // if current line is empty, but eof is not reach
            // skip this empty line...
            if !eof && clean_line.is_empty() {
                continue;
            }

            // a possible new rec found! or EOF reached
            // store clean_line to the completed_rec then swap with incomplete_rec
            let mut complete_rec = String::from(clean_line);
            mem::swap(&mut self.incomplete_rec, &mut complete_rec);
            // if previous incomplete_rec is empty && EOF, we should stop
            if eof && complete_rec.is_empty() {
                return None;
            } else if complete_rec.is_empty() {
                // not eof, but not content in complete_rec
                // happens in the beginning
                continue;
            }

            // send...
            return match AtdfRecord::from_atdf_string(
                &complete_rec,
                self.inner.delimiter,
                self.inner.scale_flag,
            ) {
                Ok(atdf_rec) => Some(atdf_rec),
                Err(e) => {
                    println!("{}", e);
                    None
                }
            };
        }
    }
}

#[inline(always)]
pub(crate) fn str_trim(input: &str) -> &str {
    let no_pre_space = input.strip_prefix(' ').unwrap_or(input);
    no_pre_space
        .strip_suffix("\r\n")
        .or_else(|| input.strip_suffix('\n'))
        .unwrap_or(no_pre_space)
}