From 7bbba155f10ce1344724ea00ca70c4d3bb469272 Mon Sep 17 00:00:00 2001
From: Daniel Schadt <kingdread@gmx.de>
Date: Thu, 12 Jan 2023 23:59:46 +0100
Subject: parallelize PNG encoding

This gives a massive speedup
---
 Cargo.lock      |  1 +
 Cargo.toml      |  1 +
 src/layer.rs    | 13 +++++++++++++
 src/main.rs     |  2 +-
 src/renderer.rs | 40 ++++++++++++++++++++++------------------
 5 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7e07963..d705490 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -509,6 +509,7 @@ dependencies = [
  "indicatif",
  "nalgebra 0.31.4",
  "num-traits",
+ "rayon",
  "roxmltree",
 ]
 
diff --git a/Cargo.toml b/Cargo.toml
index 3b48cca..468fafc 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,7 @@ imageproc = "0.23.0"
 indicatif = "0.17.2"
 nalgebra = "0.31.4"
 num-traits = "0.2.15"
+rayon = "1.6.1"
 roxmltree = "0.17.0"
 
 [dev-dependencies]
diff --git a/src/layer.rs b/src/layer.rs
index a8057a0..ff69f7d 100644
--- a/src/layer.rs
+++ b/src/layer.rs
@@ -12,6 +12,7 @@ use image::{
     ColorType, ImageBuffer, ImageEncoder, Pixel, RgbaImage,
 };
 use num_traits::Zero;
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
 
 pub const TILE_HEIGHT: u64 = 256;
 pub const TILE_WIDTH: u64 = 256;
@@ -104,6 +105,18 @@ impl<P: Pixel> TileLayer<P> {
     }
 }
 
+impl<P> TileLayer<P>
+where
+    P: Pixel + Send,
+    P::Subpixel: Send,
+{
+    pub fn into_parallel_tiles(
+        self,
+    ) -> impl ParallelIterator<Item = (u64, u64, ImageBuffer<P, Vec<P::Subpixel>>)> {
+        IntoParallelIterator::into_par_iter(self.tiles).map(|((x, y), t)| (x, y, t))
+    }
+}
+
 pub fn compress_png<P: AsRef<Path>>(image: &RgbaImage, path: P) -> Result<()> {
     let outstream = BufWriter::new(File::create(path)?);
     let encoder =
diff --git a/src/main.rs b/src/main.rs
index 1859243..269d95c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -39,7 +39,7 @@ fn main() -> Result<()> {
         let bar = ProgressBar::new(counter.tile_count().try_into().unwrap())
             .with_style(progress_style.clone());
         bar.set_prefix("Saving heat tiles");
-        renderer::lazy_colorization(&counter, &target, |x| bar.inc(x.try_into().unwrap()))?;
+        renderer::lazy_colorization(counter, &target, |x| bar.inc(x.try_into().unwrap()))?;
         bar.finish();
     }
 
diff --git a/src/renderer.rs b/src/renderer.rs
index cece8f3..2f29828 100644
--- a/src/renderer.rs
+++ b/src/renderer.rs
@@ -3,6 +3,7 @@ use std::{fs, path::Path};
 use color_eyre::eyre::{bail, Result};
 use image::{ImageBuffer, Luma, Pixel, RgbaImage};
 use nalgebra::{vector, Vector2};
+use rayon::iter::ParallelIterator;
 
 use super::{
     gpx::Coordinates,
@@ -131,10 +132,10 @@ 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: FnMut(usize)>(
-    layer: &HeatCounter,
+pub fn lazy_colorization<P: AsRef<Path>, F: Fn(usize) + Send + Sync>(
+    layer: HeatCounter,
     base_dir: P,
-    mut progress_callback: F,
+    progress_callback: F,
 ) -> Result<()> {
     let base_dir = base_dir.as_ref();
     let max = layer.pixels().map(|l| l.0[0]).max().unwrap_or_default();
@@ -142,27 +143,30 @@ pub fn lazy_colorization<P: AsRef<Path>, F: FnMut(usize)>(
         return Ok(());
     }
 
-    for (tile_x, tile_y, tile) in layer.enumerate_tiles() {
-        let colorized = colorize_tile(tile, max.into());
-        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"));
-        layer::compress_png(&colorized, file)?;
-        progress_callback(1);
-    }
+    layer
+        .into_parallel_tiles()
+        .try_for_each(|(tile_x, tile_y, tile)| {
+            let colorized = colorize_tile(&tile, max.into());
+            let folder = base_dir.join(tile_x.to_string());
+            let metadata = folder.metadata();
+            match metadata {
+                Err(_) => fs::create_dir_all(&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(())
+        })?;
     Ok(())
 }
 
 /// Renders the heat counter for the given zoom level and track points.
-pub fn render_heatcounter<F: FnMut(usize)>(
+pub fn render_heatcounter<F: Fn(usize) + Send + Sync>(
     zoom: u32,
     tracks: &[Vec<Coordinates>],
-    mut progress_callback: F,
+    progress_callback: F,
 ) -> HeatCounter {
     let mut heatcounter = TileLayer::from_pixel([0].into());
 
-- 
cgit v1.2.3