3 files changed, 46 insertions, 13 deletions

diff --git a/src/layer.rs b/src/layer.rs
index ff69f7d..2726c81 100644
--- a/src/layer.rs
+++ b/src/layer.rs
@@ -3,7 +3,11 @@
 //! This supports OSM-style "tiled" images, but not all of the tiles have to be present. If a tile
 //! is not present, a default pixel is returned. The tile is allocated with the first call to a
 //! mutating operation.
-use std::{fs::File, io::BufWriter, path::Path};
+use std::{
+    fs::File,
+    io::{BufWriter, Write},
+    path::Path,
+};
 
 use color_eyre::eyre::Result;
 use fnv::FnvHashMap;
@@ -119,6 +123,10 @@ where
 
 pub fn compress_png<P: AsRef<Path>>(image: &RgbaImage, path: P) -> Result<()> {
     let outstream = BufWriter::new(File::create(path)?);
+    compress_png_stream(image, outstream)
+}
+
+pub fn compress_png_stream<W: Write>(image: &RgbaImage, outstream: W) -> Result<()> {
     let encoder =
         PngEncoder::new_with_quality(outstream, CompressionType::Best, FilterType::Adaptive);
 
@@ -127,6 +135,12 @@ pub fn compress_png<P: AsRef<Path>>(image: &RgbaImage, path: P) -> Result<()> {
     Ok(())
 }
 
+pub fn compress_png_as_bytes(image: &RgbaImage) -> Result<Vec<u8>> {
+    let mut buffer = Vec::new();
+    compress_png_stream(image, &mut buffer)?;
+    Ok(buffer)
+}
+
 fn zero_pixel<P: Pixel>() -> P {
     let zeroes = vec![Zero::zero(); P::CHANNEL_COUNT as usize];
     *P::from_slice(&zeroes)
diff --git a/src/main.rs b/src/main.rs
index e2dacb7..a508fdc 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -123,7 +123,7 @@ fn ensure_output_directory<P: AsRef<Path>>(path: P) -> Result<()> {
     let metadata = fs::metadata(path);
     match metadata {
         Err(e) if e.kind() == ErrorKind::NotFound => {
-            let parent = path.parent().unwrap_or(Path::new("/"));
+            let parent = path.parent().unwrap_or_else(|| Path::new("/"));
             fs::create_dir(path)
                 .context(format!("Could not create output directory at {parent:?}"))?
         }
diff --git a/src/renderer.rs b/src/renderer.rs
index 2f29828..de65da6 100644
--- a/src/renderer.rs
+++ b/src/renderer.rs
@@ -1,6 +1,9 @@
-use std::{fs, path::Path};
+use std::{fs, mem, path::Path, sync::mpsc, thread};
 
-use color_eyre::eyre::{bail, Result};
+use color_eyre::{
+    eyre::{bail, Result},
+    Report,
+};
 use image::{ImageBuffer, Luma, Pixel, RgbaImage};
 use nalgebra::{vector, Vector2};
 use rayon::iter::ParallelIterator;
@@ -143,22 +146,38 @@ pub fn lazy_colorization<P: AsRef<Path>, F: Fn(usize) + Send + Sync>(
         return Ok(());
     }
 
-    layer
-        .into_parallel_tiles()
-        .try_for_each(|(tile_x, tile_y, tile)| {
-            let colorized = colorize_tile(&tile, max.into());
+    type Job = (u64, u64, Vec<u8>);
+    let (tx, rx) = mpsc::sync_channel::<Job>(30);
+
+    thread::scope(|s| {
+        let saver = s.spawn(move || loop {
+            let Ok((tile_x, tile_y, data)) = rx.recv() else { return Ok(()) };
             let folder = base_dir.join(tile_x.to_string());
             let metadata = folder.metadata();
             match metadata {
-                Err(_) => fs::create_dir_all(&folder)?,
+                Err(_) => fs::create_dir(&folder)?,
                 Ok(m) if !m.is_dir() => bail!("Output path is not a directory"),
                 _ => {}
             }
             let file = folder.join(format!("{tile_y}.png"));
-            layer::compress_png(&colorized, file)?;
-            progress_callback(1);
-            Ok(())
-        })?;
+            fs::write(file, data)?;
+        });
+
+        layer
+            .into_parallel_tiles()
+            .try_for_each(|(tile_x, tile_y, tile)| {
+                let colorized = colorize_tile(&tile, max.into());
+                let data = layer::compress_png_as_bytes(&colorized)?;
+                tx.send((tile_x, tile_y, data))?;
+                progress_callback(1);
+                Ok::<(), Report>(())
+            })?;
+
+        mem::drop(tx);
+        saver.join().unwrap()?;
+        Ok::<_, Report>(())
+    })?;
+
     Ok(())
 }