Add --sparsify

src/main.rs

@@ -1,11 +1,14 @@
 mod qcow2;
+mod sparsify;
+mod utils;
 
 use std::ops::Range;
 use std::path::Path;
 
 use qcow2::StreamingQcow2Writer;
+use sparsify::sparsify_layout;
 
-const USAGE: &'static str = "Usage: streaming-qcow2-writer input.img [layout.json] > output.qcow2";
+const USAGE: &'static str = "Usage: streaming-qcow2-writer [--sparsify] input.img [layout.json] > output.qcow2";
 
 #[cfg(unix)]
 const BLKGETSIZE64_CODE: u8 = 0x12; // Defined in linux/fs.h
@@ -48,12 +51,19 @@ fn get_file_size(file: &std::fs::File) -> std::io::Result<u64> {
 
 fn main() {
     // Read command-line arguments
-    let mut args = std::env::args_os();
+    let mut args = std::env::args_os().peekable();
+    let mut sparsify = false;
     if let None = args.next() {
         eprintln!("Not enough arguments");
         eprintln!("{}", USAGE);
         std::process::exit(2);
     }
+    if let Some(arg) = args.peek() {
+        if arg == "--sparsify" {
+            sparsify = true;
+            args.next().unwrap();
+        }
+    }
     let Some(input) = args.next() else {
         eprintln!("Not enough arguments");
         eprintln!("{}", USAGE);
@@ -66,7 +76,7 @@ fn main() {
     }
 
     // Open input
-    let (input, input_size) = match std::fs::File::open(input)
+    let (mut input, input_size) = match std::fs::File::open(input)
         .and_then(|f| get_file_size(&f).map(|s| (f, s)))
     {
         Ok(o) => o,
@@ -78,7 +88,7 @@ fn main() {
     eprintln!("Input is {} bytes", input_size);
 
     // Read layout
-    let layout = match layout {
+    let mut layout = match layout {
         Some(arg) => match load_layout_file(Path::new(&arg)) {
             Ok(l) => l,
             Err(e) => {
@@ -89,6 +99,17 @@ fn main() {
         None => vec![0..input_size],
     };
 
+    // Optional first pass: find holes
+    if sparsify {
+        layout = match sparsify_layout(&mut input, &layout) {
+            Ok(l) => l,
+            Err(e) => {
+                eprintln!("Error reading file: {}", e);
+                std::process::exit(1);
+            }
+        };
+    }
+
     // Initialize writer
     let qcow2_writer = StreamingQcow2Writer::new(input_size, layout.iter().cloned());
 

src/qcow2.rs

@@ -3,6 +3,8 @@ use std::collections::HashMap;
 use std::io::{Read, Seek, SeekFrom, Write};
 use std::ops::Range;
 
+use crate::utils::divide_and_round_up;
+
 const CLUSTER_SIZE: u64 = 65536;
 
 const REPORT_INTERVAL_BYTES: u64 = 500_000_000; // 500 MB
@@ -16,10 +18,6 @@ pub struct StreamingQcow2Writer {
     data_clusters: Vec<u64>,
 }
 
-fn divide_and_round_up(a: u64, b: u64) -> u64 {
-    (a + b - 1) / b
-}
-
 impl StreamingQcow2Writer {
     pub fn new<I: Iterator<Item=Range<u64>>>(input_size: u64, ranges: I) -> StreamingQcow2Writer {
         // Build a list of clusters
@@ -32,7 +30,7 @@ impl StreamingQcow2Writer {
 
             if let Some(last_cluster) = last_cluster {
                 if from_cluster < last_cluster {
-                    panic!("Data clusters are not sorted");
+                    panic!("Layout is not sorted");
                 } else if from_cluster == last_cluster {
                     // It is possible for the start of this range to fall in
                     // the same cluster where the last range ended