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

use flate2::write::GzEncoder;
use flate2::Compression;

use crate::chopper::error::ChopperResult;
use crate::util::buf::writer::ChopperBufWriter;

const GZ: &str = ".gz";
const ZST: &str = ".zst";

pub enum CompressionFormat {
    GZ,
    ZST,
}

pub fn compress_if_possible(
    writer: ChopperBufWriter,
    format: &str,
) -> ChopperResult<(ChopperBufWriter, String)> {
    Ok(
        if let Some((format, remaining_format)) = resolve_format(format) {
            (compress(format, writer)?, remaining_format)
        } else {
            (writer, format.to_string())
        },
    )
}

fn resolve_format(format: &str) -> Option<(CompressionFormat, String)> {
    if format.ends_with(GZ) {
        return Some((
            CompressionFormat::GZ,
            format[..(format.len() - GZ.len())].to_owned(),
        ));
    }
    if format.ends_with(ZST) {
        return Some((
            CompressionFormat::ZST,
            format[..(format.len() - ZST.len())].to_owned(),
        ));
    }
    None
}

fn compress(
    format: CompressionFormat,
    writer: ChopperBufWriter,
) -> ChopperResult<ChopperBufWriter> {
    match format {
        CompressionFormat::GZ => compress_gz(writer),
        CompressionFormat::ZST => compress_zst(writer),
    }
}

fn compress_gz(writer: ChopperBufWriter) -> ChopperResult<ChopperBufWriter> {
    let encoder = GzEncoder::new(writer, Compression::default());
    Ok(ChopperBufWriter::new(encoder))
}

fn compress_zst(writer: ChopperBufWriter) -> ChopperResult<ChopperBufWriter> {
    let encoder = zstd::stream::Encoder::new(writer, 1)?.auto_finish();
    Ok(ChopperBufWriter::new(encoder))
}