Rust code of a tree fractal

Question

I just picked up Rust recently and this is my first program (longer than 10 lines at least) so I'm looking for constructs that are more native and natural to Rust. I come from a c++/ python background. If there are some more compile-time optimizations that would also be nice since I am interested in zero overhead abstractions.

use std::fs::File;
use std::io::Write;
use std::cmp;

const SCALING_FACTOR: f64 = 1.4 as f64;

#[derive(Clone)]
struct RGB {
    red: u8,
    green: u8,
    blue: u8
}

fn serialize_rgb(pixels: &Vec<RGB>, size: usize) -> Vec<u8> {
    // for saving to a file. Is there any way we could do this
    // without constructing a new array? Would be much faster
    let mut output: Vec<u8> = Vec::with_capacity(size * 3);
    for pix in pixels {
        output.push(pix.red);
        output.push(pix.green);
        output.push(pix.blue);
    }
    output
}

struct Canvas {
    // Using 1D array so the bytes are together in memory, should be more efficient than Vec<Vec>
    // since that would store pointers to vectors?
    pixels: Vec<RGB>,
    width: i32,
    height: i32
}

impl Canvas {
    fn set_colour(&mut self, x: i32, y: i32, colour: &RGB) {
        // make this more natural? In C++ you can overload () to get a functor
        if x > 0 && y > 0 &&  x < self.width && y <  self.height {
            self.pixels[(self.width * y + x) as usize] = colour.clone();
        }
    }

    fn write_to_file(&mut self, filename: &str) {
        let mut file = init_ppm(filename, self.width, self.height);
        file.write_all(&serialize_rgb(&self.pixels, (self.width * self.height) as usize)).expect("error");
        /* slow
        for pixel in &self.pixels {
            file.write_all(&[pixel.red, pixel.green, pixel.blue]).expect("error writing to a file");
        }*/
    }

    fn new(width: i32, height: i32) -> Canvas {
        Canvas {
            width,
            height,
            pixels: vec![RGB{red:0, green:0, blue:0}; (width * height) as usize]
        }
    }

    fn draw_square(&mut self, center: &Point, width: i32, colour: &RGB) {
        for y in cmp::max(0, center.y - width) .. cmp::min(self.height, center.y + width) {
            for x in cmp::max(0, center.x - width) .. cmp::min(self.width, center.x + width) {
                self.set_colour(x ,y, &colour);
            }
        }
    }

    fn draw_line(&mut self, from: &Point, to: &Point, width: i32, colour: &RGB) {
        // function that connects two points on the grid with a line
        if from.x == to.x { // vertical lines
            let startx = cmp::max(from.x - width, 0);
            let endx = cmp::min(from.x + width, self.width);
            let endy = cmp::max(from.y, to.y) + 1;
            let starty = cmp::min(from.y, to.y);
            for y in starty .. endy {
                for x in  startx .. endx {
                    self.set_colour(x, y, colour);
                }

            }
        }
        else {
            let k = (to.y - from.y) as f64 / (to.x - from.x) as f64;
            let n = to.y as f64 - k * to.x as f64;
            let lower = cmp::min(from.x, to.x);
            let upper = cmp::max(from.x, to.x) + 1;
            for x in lower .. upper {
                // We colour y's as a function of x's
                self.draw_square(
                    &Point {x: x, y: (k * x as f64 + n) as i32},
                    width,
                    colour
                );
            }
            if k.abs() > 1.0 {
                // for steep lines, we also have to consider x as a function of y to get good results
                let lower = cmp::min(from.y, to.y);
                let upper = cmp::max(from.y, to.y) + 1;
                for y in lower .. upper {
                    self.draw_square(
                        &Point {x: ((y as f64 - n) / k) as i32, y: y},
                        width,
                        colour
                    );
                }
            }
        }
    }

}

fn rotate_point(center: &Point, point: &Point, angle: f64) -> Point {
    // also scales down a bit
    let (sin, cos) = angle.sin_cos();
    let translated = Point {x: ((point.x - center.x) as f64 / SCALING_FACTOR) as i32,
                            y: ((point.y - center.y) as f64 / SCALING_FACTOR) as i32};
    let rotated = Point {x: (translated.x as f64 * cos - translated.y as f64 * sin) as i32,
                         y: (translated.x as f64 * sin + translated.y as f64 * cos) as i32
    };
    Point {x: rotated.x + center.x, y: rotated.y + center.y}
}

fn init_ppm(filename: &str, width: i32, height: i32) -> File {
    let mut file = File::create(format!("{}.ppm",filename)).expect("couldn't create");
    file.write_all(format!("P6 {} {} 255 ", width, height).as_bytes()).expect("error writing to a file");
    file
}

struct Point {
    x: i32,
    y: i32
}

fn main() {
    const WIDTH: i32 = 1500;
    const HEIGHT: i32 = 1500;
    let mut picture = Canvas::new(WIDTH, HEIGHT);
    draw_tree(&mut picture,
              &Point {x: WIDTH/2, y: HEIGHT},
              &Point {x: WIDTH/2, y: 3*HEIGHT/4},
              15,
              &RGB {red: 255, blue: 255, green: 255},
              0.6,
              2);
    picture.write_to_file("test");
}


fn draw_tree(mut canvas: &mut Canvas, prev: &Point, next: &Point, iter: i32, colour: &RGB, angle: f64, branches: i32) {
    // recursively generates branches.
    if iter == 0 {
        return;
    }
    canvas.draw_line(prev, next, 1, colour);
    let prev = Point {x: 2 * next.x - prev.x, y: 2 * next.y - prev.y};

    if branches % 2 == 1 {
        draw_tree(&mut canvas, &next, &rotate_point(next, &prev, 0.0), iter - 1, colour, angle, branches);
    }
    for i in 1 .. branches / 2 + 1 {
        // colours are hardcoded currently but that's not really important
        let rot_left = rotate_point(next, &prev, i as f64 * angle);
        draw_tree(&mut canvas, &next, &rot_left, iter - 1, &RGB{red: 247, green: 97, blue: 74}, angle, branches);

        let rot_right = rotate_point(next, &prev, - i as f64 * angle);
        draw_tree(&mut canvas, &next, &rot_right, iter - 1, &RGB{red: 26, green: 121, blue: 244}, angle, branches);
    }
}

End result of the code is a .ppm image file test.ppm that looks like this (depending on how you set the settings in main and SCALING_FACTOR)

Answer 1

This was great work, especially for a first project. I made a bunch of changes and tried to comment my reasoning where things stood out. Here's a few things to note:

• Your code was at the size where you could go either way, but I decided to split it up into modules.
• I made Point and Rgb implement Copy since they are just very small collections of numbers. Now you can have functions take those by value, rather than reference.
• I store the color data as bytes rather than Rgbs, which lets you write that data directly to the ppm. Another possible way to do it is to use unsafe to change a &[Rgb] into a &[u8], but it's trickier and you have to be very careful.
• You should generally try to use usize when describing the size of something.
• Make it a goal to have doc comments on at least everything public. You can add #![warn(missing_docs)] to the top of main.rs to help you with this.
• Check out the tools rustfmt and clippy.

main.rs

mod canvas;
mod point;
mod rgb;

pub use canvas::Canvas;
pub use point::Point;
pub use rgb::Rgb;

fn main() {
    const WIDTH: usize = 1500;
    const HEIGHT: usize = 1500;

    let mut picture = Canvas::new(WIDTH, HEIGHT);

    draw_tree(
        &mut picture,
        Point {
            x: WIDTH as i32 / 2,
            y: HEIGHT as i32,
        },
        Point {
            x: WIDTH as i32 / 2,
            y: 3 * HEIGHT as i32 / 4,
        },
        15,
        Rgb::WHITE,
        0.6,
        2,
    );

    if let Err(e) = picture.save("test".as_ref()) {
        eprintln!("Error saving image: {}", e);
        // Lets the caller know that our program failed.
        std::process::exit(1);
    }
}

fn draw_tree(
    canvas: &mut Canvas,
    prev: Point,
    next: Point,
    iter: i32,
    colour: Rgb,
    angle: f64,
    branches: usize,
) {
    if iter == 0 {
        return;
    }
    canvas.draw_line(prev, next, 1, colour);
    let prev = Point {
        x: 2 * next.x - prev.x,
        y: 2 * next.y - prev.y,
    };

    if branches % 2 == 1 {
        let straight = next.rotate_and_scale(prev, 0.0);
        draw_tree(canvas, next, straight, iter - 1, colour, angle, branches);
    }

    for i in 1..=branches / 2 {
        let left = next.rotate_and_scale(prev, i as f64 * angle);
        draw_tree(canvas, next, left, iter - 1, Rgb::LEFT, angle, branches);

        let right = next.rotate_and_scale(prev, -(i as f64) * angle);
        draw_tree(canvas, next, right, iter - 1, Rgb::RIGHT, angle, branches);
    }
}

canvas.rs

use std::{cmp, fs::File, io::Write, path::Path};

use crate::{Point, Rgb};

#[derive(Clone, Debug)]
pub struct Canvas {
    // Since we don't need to resize data, we can store a boxed slice, which will prevent us from
    // accidently pushing, poping, resizing, etc since those are Vec methods.
    // We can store the bytes rather than the colors, which makes writing it easier.
    data: Box<[u8]>,
    // Use usize or at least an unsigned type whereever it makes sense to do so, such as a size or
    // count. I would've changed more things to unsigned, but some of your logic depends on signed.
    width: usize,
    height: usize,
}

impl Canvas {
    pub fn set_colour(&mut self, x: usize, y: usize, colour: Rgb) {
        // Using usize for parameters, we don't have to check against negative.
        if x < self.width && y < self.height {
            // Make sure to multiply by 3 here to account for directly storing the bytes...
            let offset = (self.width * y + x) * std::mem::size_of::<Rgb>();
            self.data[offset] = colour.red;
            self.data[offset + 1] = colour.green;
            self.data[offset + 2] = colour.blue;
        }
    }

    // Rather than panicing, we should return a Result, especially since it is easy because all
    // operations return io::Result.
    // Using a Path for the filename is more precise.
    pub fn save(&mut self, filename: &Path) -> std::io::Result<()> {
        let mut file = File::create(filename.with_extension("ppm"))?;

        write!(file, "P6 {} {} 255 ", self.width, self.height)?;
        file.write_all(&self.data)?;

        Ok(())
    }

    pub fn new(width: usize, height: usize) -> Canvas {
        Canvas {
            width,
            height,
            data: vec![0; width * height * std::mem::size_of::<Rgb>()].into_boxed_slice(),
        }
    }

    // I didn't look much at draw_square or draw_line...
    pub fn draw_square(&mut self, center: Point, width: i32, colour: Rgb) {
        for y in cmp::max(0, center.y - width)..cmp::min(self.height as i32, center.y + width) {
            for x in cmp::max(0, center.x - width)..cmp::min(self.width as i32, center.x + width) {
                self.set_colour(x as usize, y as usize, colour);
            }
        }
    }

    pub fn draw_line(&mut self, from: Point, to: Point, width: i32, colour: Rgb) {
        if from.x == to.x {
            let startx = cmp::max(from.x - width, 0);
            let endx = cmp::min(from.x + width, self.width as i32);
            let endy = cmp::max(from.y, to.y) + 1;
            let starty = cmp::min(from.y, to.y);
            for y in starty..endy {
                for x in startx..endx {
                    self.set_colour(x as usize, y as usize, colour);
                }
            }
        } else {
            let k = f64::from(to.y - from.y) / f64::from(to.x - from.x);
            let n = f64::from(to.y) - k * f64::from(to.x);
            let lower = cmp::min(from.x, to.x);
            let upper = cmp::max(from.x, to.x) + 1;
            for x in lower..upper {
                self.draw_square(
                    Point {
                        x,
                        y: (k * f64::from(x) + n) as i32,
                    },
                    width,
                    colour,
                );
            }
            if k.abs() > 1.0 {
                let lower = cmp::min(from.y, to.y);
                let upper = cmp::max(from.y, to.y) + 1;
                for y in lower..upper {
                    self.draw_square(
                        Point {
                            x: ((f64::from(y) - n) / k) as i32,
                            y,
                        },
                        width,
                        colour,
                    );
                }
            }
        }
    }
}

point.rs

const SCALING_FACTOR: f64 = 1.4 as f64;

#[derive(Copy, Clone, Debug)]
pub struct Point {
    pub x: i32,
    pub y: i32,
}

impl Point {
    /// Rotates a point around self by the given angle, and scales it down.
    pub fn rotate_and_scale(self, point: Point, angle: f64) -> Point {
        // also scales down a bit
        let (sin, cos) = angle.sin_cos();
        let translated = Point {
            x: ((point.x - self.x) as f64 / SCALING_FACTOR) as i32,
            y: ((point.y - self.y) as f64 / SCALING_FACTOR) as i32,
        };
        let rotated = Point {
            x: (translated.x as f64 * cos - translated.y as f64 * sin) as i32,
            y: (translated.x as f64 * sin + translated.y as f64 * cos) as i32,
        };
        Point {
            x: rotated.x + self.x,
            y: rotated.y + self.y,
        }
    }
}

rgb.rs

#[derive(Copy, Clone, Debug)]
pub struct Rgb {
    pub red: u8,
    pub green: u8,
    pub blue: u8,
}

impl Rgb {
    // Makes it a little nicer to create colors, IMO.
    pub const fn new(red: u8, green: u8, blue: u8) -> Self {
        Self { red, green, blue }
    }

    pub const WHITE: Self = Self::new(255, 255, 255);
    pub const LEFT: Self = Self::new(247, 97, 74);
    pub const RIGHT: Self = Self::new(26, 121, 244);
}