zvault/src/cli/algotest.rs

use prelude::*;

use std::io::{self, Cursor, Read, Write};
use std::fs::File;
use std::collections::HashSet;

use chrono::Duration;


struct ChunkSink {
    chunks: Vec<(usize, usize)>,
    pos: usize,
    written: usize
}

impl ChunkSink {
    fn end_chunk(&mut self) {
        self.chunks.push((self.pos, self.written));
        self.pos += self.written;
        self.written = 0;
    }
}

impl Write for ChunkSink {
    fn write(&mut self, data: &[u8]) -> Result<usize, io::Error> {
        self.written += data.len();
        Ok(data.len())
    }

    fn flush(&mut self) -> Result<(), io::Error> {
        Ok(())
    }
}

fn chunk(data: &[u8], mut chunker: Box<Chunker>, sink: &mut ChunkSink) {
    let mut cursor = Cursor::new(data);
    while chunker.chunk(&mut cursor, sink).unwrap() == ChunkerStatus::Continue {
        sink.end_chunk();
    }
    sink.end_chunk();
}

#[allow(dead_code)]
pub fn run(
    path: &str,
    bundle_size: usize,
    chunker: ChunkerType,
    compression: Option<Compression>,
    encrypt: bool,
    hash: HashMethod,
) {
    let mut total_write_time = 0.0;
    let mut total_read_time = 0.0;

    println!("Reading input file ...");
    let mut file = File::open(path).unwrap();
    let total_size = file.metadata().unwrap().len();
    let mut size = total_size;
    let mut data = Vec::with_capacity(size as usize);
    let read_time = Duration::span(|| { file.read_to_end(&mut data).unwrap(); })
        .num_milliseconds() as f32 / 1_000.0;
    println!(
        "- {}, {}",
        to_duration(read_time),
        to_speed(size, read_time)
    );

    println!();

    println!(
        "Chunking data with {}, avg chunk size {} ...",
        chunker.name(),
        to_file_size(chunker.avg_size() as u64)
    );
    let mut chunk_sink = ChunkSink {
        chunks: Vec::with_capacity(2 * size as usize / chunker.avg_size()),
        written: 0,
        pos: 0
    };
    let chunker = chunker.create();
    let chunk_time = Duration::span(|| chunk(&data, chunker, &mut chunk_sink))
        .num_milliseconds() as f32 / 1_000.0;
    total_write_time += chunk_time;
    println!(
        "- {}, {}",
        to_duration(chunk_time),
        to_speed(size, chunk_time)
    );
    let mut chunks = chunk_sink.chunks;
    assert_eq!(chunks.iter().map(|c| c.1).sum::<usize>(), size as usize);
    let chunk_size_avg = size as f32 / chunks.len() as f32;
    let chunk_size_stddev = (chunks
                                 .iter()
                                 .map(|c| (c.1 as f32 - chunk_size_avg).powi(2))
                                 .sum::<f32>() /
                                 (chunks.len() as f32 - 1.0))
        .sqrt();
    println!(
        "- {} chunks, avg size: {} ±{}",
        chunks.len(),
        to_file_size(chunk_size_avg as u64),
        to_file_size(chunk_size_stddev as u64)
    );

    println!();

    println!("Hashing chunks with {} ...", hash.name());
    let mut hashes = Vec::with_capacity(chunks.len());
    let hash_time = Duration::span(|| for &(pos, len) in &chunks {
        hashes.push(hash.hash(&data[pos..pos + len]))
    }).num_milliseconds() as f32 / 1_000.0;
    total_write_time += hash_time;
    println!(
        "- {}, {}",
        to_duration(hash_time),
        to_speed(size, hash_time)
    );
    let mut seen_hashes = HashSet::with_capacity(hashes.len());
    let mut dups = Vec::new();
    for (i, hash) in hashes.into_iter().enumerate() {
        if !seen_hashes.insert(hash) {
            dups.push(i);
        }
    }
    let mut dup_size = 0;
    dups.reverse();
    for i in &dups {
        let (_, len) = chunks.remove(*i);
        dup_size += len;
    }
    println!(
        "- {} duplicate chunks, {}, {:.1}% saved",
        dups.len(),
        to_file_size(dup_size as u64),
        dup_size as f32 / size as f32 * 100.0
    );
    size -= dup_size as u64;

    let mut bundles = Vec::new();

    if let Some(compression) = compression.clone() {
        println!();

        println!("Compressing chunks with {} ...", compression.to_string());
        let compress_time = Duration::span(|| {
            let mut bundle = Vec::with_capacity(bundle_size + 2 * chunk_size_avg as usize);
            let mut c = compression.compress_stream().unwrap();
            for &(pos, len) in &chunks {
                c.process(&data[pos..pos + len], &mut bundle).unwrap();
                if bundle.len() >= bundle_size {
                    c.finish(&mut bundle).unwrap();
                    bundles.push(bundle);
                    bundle = Vec::with_capacity(bundle_size + 2 * chunk_size_avg as usize);
                    c = compression.compress_stream().unwrap();
                }
            }
            c.finish(&mut bundle).unwrap();
            bundles.push(bundle);
        }).num_milliseconds() as f32 / 1_000.0;
        total_write_time += compress_time;
        println!(
            "- {}, {}",
            to_duration(compress_time),
            to_speed(size, compress_time)
        );
        let compressed_size = bundles.iter().map(|b| b.len()).sum::<usize>();
        println!(
            "- {} bundles, {}, {:.1}% saved",
            bundles.len(),
            to_file_size(compressed_size as u64),
            (size as f32 - compressed_size as f32) / size as f32 * 100.0
        );
        size = compressed_size as u64;
    } else {
        let mut bundle = Vec::with_capacity(bundle_size + 2 * chunk_size_avg as usize);
        for &(pos, len) in &chunks {
            bundle.extend_from_slice(&data[pos..pos + len]);
            if bundle.len() >= bundle_size {
                bundles.push(bundle);
                bundle = Vec::with_capacity(bundle_size + 2 * chunk_size_avg as usize);
            }
        }
        bundles.push(bundle);
    }

    if encrypt {
        println!();

        let (public, secret) = Crypto::gen_keypair();
        let mut crypto = Crypto::dummy();
        crypto.add_secret_key(public, secret);
        let encryption = (EncryptionMethod::Sodium, public[..].to_vec().into());

        println!("Encrypting bundles...");
        let mut encrypted_bundles = Vec::with_capacity(bundles.len());

        let encrypt_time = Duration::span(|| for bundle in bundles {
            encrypted_bundles.push(crypto.encrypt(&encryption, &bundle).unwrap());
        }).num_milliseconds() as f32 / 1_000.0;
        println!(
            "- {}, {}",
            to_duration(encrypt_time),
            to_speed(size, encrypt_time)
        );
        total_write_time += encrypt_time;

        println!();

        println!("Decrypting bundles...");
        bundles = Vec::with_capacity(encrypted_bundles.len());
        let decrypt_time = Duration::span(|| for bundle in encrypted_bundles {
            bundles.push(crypto.decrypt(&encryption, &bundle).unwrap());
        }).num_milliseconds() as f32 / 1_000.0;
        println!(
            "- {}, {}",
            to_duration(decrypt_time),
            to_speed(size, decrypt_time)
        );
        total_read_time += decrypt_time;
    }

    if let Some(compression) = compression {
        println!();

        println!("Decompressing bundles with {} ...", compression.to_string());
        let mut dummy = ChunkSink {
            chunks: vec![],
            written: 0,
            pos: 0
        };
        let decompress_time = Duration::span(|| for bundle in &bundles {
            let mut c = compression.decompress_stream().unwrap();
            c.process(bundle, &mut dummy).unwrap();
            c.finish(&mut dummy).unwrap();
        }).num_milliseconds() as f32 / 1_000.0;
        println!(
            "- {}, {}",
            to_duration(decompress_time),
            to_speed(total_size - dup_size as u64, decompress_time)
        );
        total_read_time += decompress_time;
    }

    println!();

    println!(
        "Total storage size: {} / {}, ratio: {:.1}%",
        to_file_size(size as u64),
        to_file_size(total_size as u64),
        size as f32 / total_size as f32 * 100.0
    );
    println!(
        "Total processing speed: {}",
        to_speed(total_size, total_write_time)
    );
    println!(
        "Total read speed: {}",
        to_speed(total_size, total_read_time)
    );
}