✨ Add image save support

Once the render is done it saves an PNG.

6 files changed, 89 insertions, 14 deletions

diff --git a/racer-tracer/Cargo.toml b/racer-tracer/Cargo.toml
index 0e80de3..e5ba59b 100644
--- a/racer-tracer/Cargo.toml
+++ b/racer-tracer/Cargo.toml
@@ -14,4 +14,6 @@ structopt = "0.3"
 num_cpus = "1.15.0"
 synchronoise = "1.0.1"
 config = "0.13.3"
+sha2 = "0.10.6"
+image = "0.24.5"
 serde = { version = "1", features = ["derive"] }
diff --git a/racer-tracer/config.yml b/racer-tracer/config.yml
index 66af4a6..26aea31 100644
--- a/racer-tracer/config.yml
+++ b/racer-tracer/config.yml
@@ -17,3 +17,5 @@ screen:
   height: 720
 
 scene: ./scene.yml
+
+image_output_dir: "../"
diff --git a/racer-tracer/src/config.rs b/racer-tracer/src/config.rs
index cf747ec..7bd7887 100644
--- a/racer-tracer/src/config.rs
+++ b/racer-tracer/src/config.rs
@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use config::File;
 use serde::Deserialize;
 use structopt::StructOpt;
@@ -59,6 +61,9 @@ pub struct Config {
 
     #[serde(default)]
     pub scene: Option<String>,
+
+    #[serde(default)]
+    pub image_output_dir: Option<PathBuf>,
 }
 
 impl Config {
diff --git a/racer-tracer/src/error.rs b/racer-tracer/src/error.rs
index 6e1b305..1725d1a 100644
--- a/racer-tracer/src/error.rs
+++ b/racer-tracer/src/error.rs
@@ -5,6 +5,9 @@ pub enum TracerError {
     #[error("Unknown error: {message}")]
     Unknown { message: String, exit_code: i32 },
 
+    #[error("Error: {0}")]
+    Generic(String),
+
     #[error("Failed to create window: {0}")]
     FailedToCreateWindow(String),
 
@@ -31,6 +34,9 @@ pub enum TracerError {
 
     #[error("Cancel event")]
     CancelEvent,
+
+    #[error("Image save error: {0}")]
+    ImageSave(String),
 }
 
 impl From<TracerError> for i32 {
@@ -49,6 +55,8 @@ impl From<TracerError> for i32 {
             TracerError::FailedToAcquireLock(_) => 8,
             TracerError::ExitEvent => 9,
             TracerError::CancelEvent => 10,
+            TracerError::Generic(_) => 11,
+            TracerError::ImageSave(_) => 12,
         }
     }
 }
diff --git a/racer-tracer/src/main.rs b/racer-tracer/src/main.rs
index 7d7510f..aa1d974 100644
--- a/racer-tracer/src/main.rs
+++ b/racer-tracer/src/main.rs
@@ -11,14 +11,18 @@ mod scene;
 mod util;
 mod vec3;
 
+extern crate image as img;
+
 use std::{
     convert::TryFrom,
+    path::PathBuf,
     sync::RwLock,
     time::{Duration, Instant},
     vec::Vec,
 };
 
 use minifb::{Key, Window, WindowOptions};
+use sha2::{Digest, Sha256};
 use synchronoise::SignalEvent;
 
 use crate::{
@@ -70,7 +74,7 @@ fn run(config: Config) -> Result<(), TracerError> {
                                 },
                                 |_| {
                                     let render_time = Instant::now();
-                                    let res = render(
+                                    render(
                                         &screen_buffer,
                                         &camera,
                                         &image,
@@ -78,21 +82,74 @@ fn run(config: Config) -> Result<(), TracerError> {
                                         &config.render,
                                         Some(&cancel_render),
                                         None,
-                                    );
-                                    render_image.reset();
-
-                                    println!(
-                                        "It took {} seconds to render the image.",
-                                        Instant::now().duration_since(render_time).as_secs()
-                                    );
-
-                                    // TODO: Output the image
-
-                                    res
+                                    )
+                                    .and_then(|_| {
+                                        render_image.reset();
+
+                                        println!(
+                                            "It took {} seconds to render the image.",
+                                            Instant::now().duration_since(render_time).as_secs()
+                                        );
+                                        screen_buffer
+                                            .read()
+                                            .map_err(|e| {
+                                                TracerError::FailedToAcquireLock(e.to_string())
+                                            })
+                                            .map(|buf| {
+                                                // Convert ARGB8 to RGBA8
+                                                buf.iter()
+                                                    .map(|v| {
+                                                        let a: u32 = (v >> 24) & 0xff;
+                                                        let r: u32 = (v >> 16) & 0xff;
+                                                        let g: u32 = (v >> 8) & 0xff;
+                                                        let b: u32 = v & 0xff;
+
+                                                        (r << 24) | (g << 16) | (b << 8) | a
+                                                    })
+                                                    .flat_map(|val| val.to_be_bytes())
+                                                    .collect::<Vec<u8>>()
+                                            })
+                                    })
+                                    .and_then(|buf| {
+                                        match &config.image_output_dir {
+                                            Some(image_dir) => {
+                                                println!("Saving image...");
+                                                let mut sha = Sha256::new();
+
+                                                sha.update(buf.as_slice());
+
+                                                let mut file_path = PathBuf::from(image_dir);
+                                                file_path.push(format!("{:X}.png", sha.finalize()));
+
+                                                img::save_buffer(
+                                                    file_path.as_path(),
+                                                    buf.as_slice(),
+                                                    config.screen.width as u32,
+                                                    config.screen.height as u32,
+                                                    img::ColorType::Rgba8,
+                                                )
+                                                .map_err(|e| {
+                                                    let error = e.to_string();
+                                                    TracerError::ImageSave(error)
+                                                })
+                                                .map(|_| {
+                                                    println!(
+                                                        "Saved image to: {}",
+                                                        file_path.to_string_lossy()
+                                                    );
+                                                })
+                                            }
+                                            None => Ok(()),
+                                        }
+                                    })
                                 },
                             )
                     });
             }
+
+            if render_res.is_err() {
+                exit.signal();
+            }
         });
 
         s.spawn(|_| {
@@ -109,7 +166,7 @@ fn run(config: Config) -> Result<(), TracerError> {
             })
             .and_then(|mut window| {
                 let mut t = Instant::now();
-                while window.is_open() && !window.is_key_down(Key::Escape) {
+                while window.is_open() && !window.is_key_down(Key::Escape) && !exit.status() {
                     let dt = t.elapsed().as_micros() as f64 / 1000000.0;
                     t = Instant::now();
                     // Sleep a bit to not hog the lock on the buffer all the time.
diff --git a/racer-tracer/src/vec3.rs b/racer-tracer/src/vec3.rs
index c5edf32..8530bfe 100644
--- a/racer-tracer/src/vec3.rs
+++ b/racer-tracer/src/vec3.rs
@@ -87,7 +87,8 @@ impl Vec3 {
         let red: u32 = (self.data[0] * 255.0) as u32;
         let green: u32 = (self.data[1] * 255.0) as u32;
         let blue: u32 = (self.data[2] * 255.0) as u32;
-        (red << 16) | green << 8 | blue
+        // ARGB
+        (255 << 24) | (red << 16) | green << 8 | blue
     }
 
     pub fn random() -> Self {