3 files changed, 79 insertions(+), 45 deletions(-)

M src/main.rs
M src/piet_image.rs
M test.png

@@ 1,3 1,4 @@
+use std::fmt::format;
 use std::fs;
 
 use clap::Parser;


@@ 10,8 11,6 @@ mod utils;
 use crate::instructions::PietEncoder;
 use anyhow::{Context, Result};
 
-use spiral::*;
-
 /// A text to piet encoder
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about=None)]


@@ 38,15 37,19 @@ fn main() -> Result<()> {
 
     let instrs = encoder.encode_text(&input);
 
-    let mut img = piet_image::PietImager::with_instructions(instrs);
+    let mut imager = piet_image::PietImager::with_instructions(instrs);
 
     if let Some(img_size) = args.img_size {
-        // TODO: add image_size
-        img.to_image()? // something
+        imager
+            .set_img_size(img_size)
+            .to_image()
+            .context(format!("Failed to convert to image with size {}", img_size))?
             .save(args.output.unwrap_or("test.png".to_string()))
             .context("Failed to write ouput image")
     } else {
-        img.to_image()?
+        imager
+            .to_image()
+            .context("Failed to convert to image with minimum size")?
             .save(args.output.unwrap_or("test.png".to_string()))
             .context("Failed to write ouput image")
     }

@@ 8,45 8,60 @@ use anyhow::{bail, Context, Result};
 use crate::colors::{self, Color};
 use crate::instructions;
 use crate::instructions::{Instructions::*, *};
-use crate::spiral::Spiral2D;
+use crate::spiral::{Direction, Spiral2D};
 use crate::utils;
-use crate::Direction;
 
+// TODO: Make img an Option<RgbImage>
 pub struct PietImager {
-    img: RgbImage,
+    img_size: Option<u32>,
+    img: Option<RgbImage>,
     instructions: Vec<Instructions>,
 }
 
 impl PietImager {
     pub fn new() -> Self {
         Self {
-            img: RgbImage::new(0, 0),
+            img_size: None,
+            img: None,
             instructions: vec![],
         }
     }
 
+    pub fn set_img_size(&mut self, img_size: u32) -> &mut Self {
+        self.img_size = Some(img_size);
+        self
+    }
+
     pub fn with_instructions(instructions: Vec<Instructions>) -> Self {
         Self {
+            img_size: None,
+            img: None,
             instructions,
-            img: RgbImage::new(0, 0),
         }
     }
 
     pub fn to_image(&mut self) -> Result<RgbImage> {
-        let (min_img_size, max_img_size) = instructions::expected_size(&self.instructions)?;
-        eprintln!(
-            "(min_img_size, max_img_size: {:?}",
-            (min_img_size, max_img_size)
-        );
-
-        let sat_img_size: u32 = utils::search_for_minimum_satisfier(
-            |size| self.try_image_with_size(*size).is_ok(),
-            min_img_size..=max_img_size,
-        )
-        .unwrap();
-
-        self.try_image_with_size(sat_img_size)
-            .context("Failed to create a piet image!")
+        if let Some(img_size) = self.img_size {
+            self.try_image_with_size(img_size)
+                .context("Failed to create a piet image!")
+        } else {
+            let (min_img_size, max_img_size) = instructions::expected_size(&self.instructions)?;
+            eprintln!(
+                "(min_img_size, max_img_size: {:?}",
+                (min_img_size, max_img_size)
+            );
+
+            let sat_img_size: u32 = utils::search_for_minimum_satisfier(
+                |size| self.try_image_with_size(*size).is_ok(),
+                min_img_size..=max_img_size,
+            )
+            .unwrap();
+
+            self.img_size = Some(sat_img_size);
+
+            self.try_image_with_size(sat_img_size)
+                .context("Failed to create a piet image!")
+        }
     }
 
     fn instructions_to_blocks(&self, img_size: u32) -> Result<Vec<Color>> {


@@ 71,8 86,14 @@ impl PietImager {
                         //    "SPLIT PUSH - remaining_space: {:?}, n: {:?}, prev: {:?}",
                         //    remaining_space, n, prev
                         //);
-                        let first_push = remaining_space as usize - corner_turn.len();
-                        let second_push = n - first_push;
+                        let first_push = (remaining_space as usize)
+                            .checked_sub(corner_turn.len())
+                            .context("Not enough space on the spiral to perform a turn!")?;
+
+                        let second_push = n
+                            .checked_sub(first_push)
+                            .context("Not enough space on the spiral to perform a turn!")?;
+
                         //eprintln!("first: {}, second: {}", first_push, second_push);
                         for _ in 0..first_push {
                             colors.push(prev);


@@ 154,7 175,7 @@ impl PietImager {
     }
 
     fn colors_to_img(&mut self, colors: Vec<Color>, img_size: u32) -> Result<RgbImage> {
-        self.img = RgbImage::new(img_size, img_size);
+        let mut img = RgbImage::new(img_size, img_size);
 
         let mut final_pos = (0, 0);
         let mut final_dir = Direction::Right;


@@ 162,7 183,7 @@ impl PietImager {
         let sp = Spiral2D::new(0, 0, Direction::Right, img_size, colors.len() as u32);
 
         for ((x, y, dir), color) in sp.zip(colors.iter()) {
-            self.img.put_pixel(x, y, Rgb(color.to_bytes()));
+            img.put_pixel(x, y, Rgb(color.to_bytes()));
             final_pos = (x, y);
             final_dir = dir;
         }


@@ 203,12 224,10 @@ impl PietImager {
 
         for (x, y, _) in filler_spiral {
             let neighbors: HashSet<[u8; 3]> = vec![
-                self.img
-                    .get_pixel_checked(if x == 0 { 0 } else { x - 1 }, y),
-                self.img.get_pixel_checked(x + 1, y),
-                self.img
-                    .get_pixel_checked(x, if y == 0 { 0 } else { y - 1 }),
-                self.img.get_pixel_checked(x, y + 1),
+                img.get_pixel_checked(if x == 0 { 0 } else { x - 1 }, y),
+                img.get_pixel_checked(x + 1, y),
+                img.get_pixel_checked(x, if y == 0 { 0 } else { y - 1 }),
+                img.get_pixel_checked(x, y + 1),
             ]
             .into_iter()
             .map(|c| if let Some(rgb) = c { rgb.0 } else { [0, 0, 0] })


@@ 217,7 236,7 @@ impl PietImager {
             let legal_colors: Vec<&[u8; 3]> = all_colors.difference(&neighbors).collect();
             let picked_color = *legal_colors.choose(&mut rand::thread_rng()).unwrap();
 
-            self.img.put_pixel(x, y, Rgb(*picked_color));
+            img.put_pixel(x, y, Rgb(*picked_color));
         }
 
         let (x, y) = final_pos;


@@ 237,7 256,7 @@ impl PietImager {
 
             let all_pixels_in_region_are_black = terminator_region
                 .iter()
-                .all(|&(x, y)| self.img.get_pixel(x, y) == &Rgb(Color::Black.to_bytes()));
+                .all(|&(x, y)| img.get_pixel(x, y) == &Rgb(Color::Black.to_bytes()));
 
             if !all_pixels_in_region_are_black {
                 bail!("Unable to add terminator to piet image due to insuficient space");


@@ 245,8 264,8 @@ impl PietImager {
         }
         // Add terminator
         self.add_terminator(final_pos, color, &final_dir);
-
-        Ok(self.img.clone())
+        self.img = Some(img.clone());
+        Ok(img.clone())
     }
 
     /// Creates a square that starts at the given position, clockwise


@@ 264,25 283,37 @@ impl PietImager {
 
     fn add_terminator(&mut self, pos: (u32, u32), color: &Color, dir: &Direction) {
         let (mut x, mut y) = pos;
-        self.img.put_pixel(x, y, Rgb(color.to_bytes()));
+        self.img
+            .as_mut()
+            .map(|s| s.put_pixel(x, y, Rgb(color.to_bytes())));
 
         let color = colors::instruction_to_color(PUSH(1), *color);
         (x, y) = self.update_x_y(x, y, 1, &dir.clockwise());
-        self.img.put_pixel(x, y, Rgb(color.to_bytes()));
+        self.img
+            .as_mut()
+            .map(|s| s.put_pixel(x, y, Rgb(color.to_bytes())));
 
         let color = colors::instruction_to_color(PUSH(1), color);
         (x, y) = self.update_x_y(x, y, 1, &dir.clockwise());
-        self.img.put_pixel(x, y, Rgb(color.to_bytes()));
+        self.img
+            .as_mut()
+            .map(|s| s.put_pixel(x, y, Rgb(color.to_bytes())));
 
         (x, y) = self.update_x_y(x, y, 1, &dir);
-        self.img.put_pixel(x, y, Rgb(color.to_bytes()));
+        self.img
+            .as_mut()
+            .map(|s| s.put_pixel(x, y, Rgb(color.to_bytes())));
 
         (x, y) = self.update_x_y(x, y, 2, &dir.clockwise().clockwise());
-        self.img.put_pixel(x, y, Rgb(color.to_bytes()));
+        self.img
+            .as_mut()
+            .map(|s| s.put_pixel(x, y, Rgb(color.to_bytes())));
 
         (x, y) = self.update_x_y(x, y, 1, &dir);
         (x, y) = self.update_x_y(x, y, 1, &dir.clockwise());
-        self.img.put_pixel(x, y, Rgb(color.to_bytes()));
+        self.img
+            .as_mut()
+            .map(|s| s.put_pixel(x, y, Rgb(color.to_bytes())));
     }
 
     fn update_x_y(&self, x: u32, y: u32, step: u32, dir: &Direction) -> (u32, u32) {