factor out saving logic

Since we want to support SQLite at some point, it makes sense to have
the exact storage method abstracted away.

4 files changed, 99 insertions, 51 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 06f8f31..670678f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,3 +2,4 @@
 pub mod gpx;
 pub mod layer;
 pub mod renderer;
+pub mod storage;
diff --git a/src/main.rs b/src/main.rs
index 53a1301..7230dce 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,17 +1,14 @@
-use std::{
-    fs,
-    io::{self, ErrorKind},
-    path::{Path, PathBuf},
-};
+use std::{io, path::PathBuf};
 
 use clap::Parser;
 use color_eyre::{
-    eyre::{bail, eyre, Context, Result},
+    eyre::{bail, eyre, Result},
     Report,
 };
 use hittekaart::{
     gpx::{self, Compression},
     renderer,
+    storage::{Folder, Storage},
 };
 use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle};
 use is_terminal::IsTerminal;
@@ -88,7 +85,8 @@ fn main() -> Result<()> {
     let tracks = tracks.into_iter().collect::<Result<Vec<_>>>()?;
     bar.finish();
 
-    ensure_output_directory(&args.output_directory)?;
+    let mut storage = Folder::new(args.output_directory.clone());
+    storage.prepare()?;
 
     let multibar = MultiProgress::new();
     if !use_progress_bars {
@@ -107,37 +105,23 @@ fn main() -> Result<()> {
         bar.finish();
         multibar.remove(&bar);
 
-        let target = [&args.output_directory, &zoom.to_string().into()]
-            .iter()
-            .collect::<PathBuf>();
-        fs::create_dir(&target)?;
+        storage.prepare_zoom(zoom)?;
 
         let bar =
             make_bar(counter.tile_count().try_into().unwrap()).with_style(progress_style.clone());
         multibar.insert_from_back(1, bar.clone());
         bar.set_prefix("Saving heat tiles");
-        renderer::lazy_colorization(counter, &target, |x| bar.inc(x.try_into().unwrap()))?;
+        renderer::lazy_colorization(counter, |rendered_tile| {
+            storage.store(zoom, rendered_tile.x, rendered_tile.y, &rendered_tile.data)?;
+            bar.inc(1);
+            Ok(())
+        })?;
         bar.finish();
         multibar.remove(&bar);
         zoom_bar.inc(1);
     }
+    storage.finish()?;
     zoom_bar.finish();
 
     Ok(())
 }
-
-fn ensure_output_directory<P: AsRef<Path>>(path: P) -> Result<()> {
-    let path = path.as_ref();
-    let metadata = fs::metadata(path);
-    match metadata {
-        Err(e) if e.kind() == ErrorKind::NotFound => {
-            let parent = path.parent().unwrap_or_else(|| Path::new("/"));
-            fs::create_dir(path)
-                .context(format!("Could not create output directory at {parent:?}"))?
-        }
-        Err(e) => Err(e).context("Error while checking output directory")?,
-        Ok(m) if m.is_dir() => (),
-        Ok(_) => bail!("Output directory is not a directory"),
-    }
-    Ok(())
-}
diff --git a/src/renderer.rs b/src/renderer.rs
index bb33ddc..4840e73 100644
--- a/src/renderer.rs
+++ b/src/renderer.rs
@@ -1,9 +1,6 @@
-use std::{fs, path::Path, thread};
+use std::thread;
 
-use color_eyre::{
-    eyre::{bail, Result},
-    Report,
-};
+use color_eyre::{eyre::Result, Report};
 use image::{ImageBuffer, Luma, Pixel, RgbaImage};
 use nalgebra::{vector, Vector2};
 use rayon::iter::ParallelIterator;
@@ -13,6 +10,13 @@ use super::{
     layer::{self, TileLayer},
 };
 
+#[derive(Debug, Clone)]
+pub struct RenderedTile {
+    pub x: u64,
+    pub y: u64,
+    pub data: Vec<u8>,
+}
+
 pub type HeatCounter = TileLayer<Luma<u8>>;
 
 fn render_circle<P: Pixel>(layer: &mut TileLayer<P>, center: (u64, u64), radius: u64, pixel: P) {
@@ -135,32 +139,21 @@ fn colorize_tile(tile: &ImageBuffer<Luma<u8>, Vec<u8>>, max: u32) -> RgbaImage {
 /// rendering the next one.
 ///
 /// This has a way lower memory usage than [`colorize_heatcounter`].
-pub fn lazy_colorization<P: AsRef<Path>, F: Fn(usize) + Send + Sync>(
+pub fn lazy_colorization<F: FnMut(RenderedTile) -> Result<()> + Send + Sync>(
     layer: HeatCounter,
-    base_dir: P,
-    progress_callback: F,
+    mut save_callback: F,
 ) -> Result<()> {
-    let base_dir = base_dir.as_ref();
     let max = layer.pixels().map(|l| l.0[0]).max().unwrap_or_default();
     if max == 0 {
         return Ok(());
     }
 
-    type Job = (u64, u64, Vec<u8>);
-    let (tx, rx) = crossbeam_channel::bounded::<Job>(30);
+    let (tx, rx) = crossbeam_channel::bounded::<RenderedTile>(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(&folder)?,
-                Ok(m) if !m.is_dir() => bail!("Output path is not a directory"),
-                _ => {}
-            }
-            let file = folder.join(format!("{tile_y}.png"));
-            fs::write(file, data)?;
+            let Ok(tile) = rx.recv() else { return Ok::<_, Report>(()) };
+            save_callback(tile)?;
         });
 
         layer
@@ -168,8 +161,11 @@ pub fn lazy_colorization<P: AsRef<Path>, F: Fn(usize) + Send + Sync>(
             .try_for_each_with(tx, |tx, (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);
+                tx.send(RenderedTile {
+                    x: tile_x,
+                    y: tile_y,
+                    data,
+                })?;
                 Ok::<(), Report>(())
             })?;
 
diff --git a/src/storage.rs b/src/storage.rs
new file mode 100644
index 0000000..fef8194
--- /dev/null
+++ b/src/storage.rs
@@ -0,0 +1,67 @@
+use color_eyre::{eyre::{bail, WrapErr}, Result};
+use std::{
+    fs,
+    io::ErrorKind,
+    path::{Path, PathBuf},
+};
+
+pub trait Storage {
+    fn prepare(&mut self) -> Result<()>;
+    fn prepare_zoom(&mut self, zoom: u32) -> Result<()>;
+    fn store(&mut self, zoom: u32, x: u64, y: u64, data: &[u8]) -> Result<()>;
+    fn finish(&mut self) -> Result<()>;
+}
+
+#[derive(Debug)]
+pub struct Folder {
+    base_dir: PathBuf,
+}
+
+impl Folder {
+    pub fn new(base_dir: PathBuf) -> Self {
+        Folder { base_dir }
+    }
+}
+
+impl Storage for Folder {
+    fn prepare(&mut self) -> Result<()> {
+        let path = &self.base_dir;
+        let metadata = fs::metadata(path);
+        match metadata {
+            Err(e) if e.kind() == ErrorKind::NotFound => {
+                let parent = path.parent().unwrap_or_else(|| Path::new("/"));
+                fs::create_dir(path)
+                    .context(format!("Could not create output directory at {parent:?}"))?
+            }
+            Err(e) => Err(e).context("Error while checking output directory")?,
+            Ok(m) if m.is_dir() => (),
+            Ok(_) => bail!("Output directory is not a directory"),
+        }
+        Ok(())
+    }
+
+    fn prepare_zoom(&mut self, zoom: u32) -> Result<()> {
+        let target = [&self.base_dir, &zoom.to_string().into()]
+            .iter()
+            .collect::<PathBuf>();
+        fs::create_dir(target)?;
+        Ok(())
+    }
+
+    fn store(&mut self, zoom: u32, x: u64, y: u64, data: &[u8]) -> Result<()> {
+        let folder = self.base_dir.join(zoom.to_string()).join(x.to_string());
+        let metadata = folder.metadata();
+        match metadata {
+            Err(_) => fs::create_dir(&folder)?,
+            Ok(m) if !m.is_dir() => bail!("Output path is not a directory"),
+            _ => {}
+        }
+        let file = folder.join(format!("{y}.png"));
+        fs::write(file, data)?;
+        Ok(())
+    }
+
+    fn finish(&mut self) -> Result<()> {
+        Ok(())
+    }
+}