Skip to main content

wimbd/
io.rs

1//! IO helpers.
2
3use std::{
4    fs::File,
5    io::{self, prelude::*},
6    rc::Rc,
7};
8
9use anyhow::Result;
10use flate2::read::MultiGzDecoder;
11
12/// A buffered reader for gzip files.
13pub struct GzBufReader {
14    reader: io::BufReader<MultiGzDecoder<File>>,
15    buf: Rc<String>,
16}
17
18fn new_buf() -> Rc<String> {
19    Rc::new(String::with_capacity(2048))
20}
21
22impl GzBufReader {
23    pub fn open(path: impl AsRef<std::path::Path>) -> Result<Self> {
24        let reader = io::BufReader::new(MultiGzDecoder::new(File::open(path)?));
25        let buf = new_buf();
26
27        Ok(Self { reader, buf })
28    }
29}
30
31type DataIteratorItem = io::Result<Rc<String>>;
32
33impl Iterator for GzBufReader {
34    type Item = DataIteratorItem;
35
36    fn next(&mut self) -> Option<Self::Item> {
37        let buf = match Rc::get_mut(&mut self.buf) {
38            Some(buf) => {
39                buf.clear();
40                buf
41            }
42            None => {
43                self.buf = new_buf();
44                Rc::make_mut(&mut self.buf)
45            }
46        };
47
48        self.reader
49            .read_line(buf)
50            .map(|u| {
51                if u == 0 {
52                    None
53                } else {
54                    Some(Rc::clone(&self.buf))
55                }
56            })
57            .transpose()
58    }
59}