fractal.rs

fractal.rs

// Based on the code in O'Reilly's Programming Rust, Chapter 2

extern crate num;
extern crate image;
extern crate crossbeam;

use std::str::FromStr;
use std::io::Write;
use std::fs::File;
use num::Complex;
use image::ColorType;
use image::png::PNGEncoder;    

#[derive(Copy, Clone)]
enum FractalKind {
    Mandelbrot,
    Mandelbrot3,
    BurningShip,
}

impl FromStr for FractalKind {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "mandelbrot" => Ok(FractalKind::Mandelbrot),
            "mandelbrot3" => Ok(FractalKind::Mandelbrot3),
            "burning_ship" => Ok(FractalKind::BurningShip),
            _ => Err(())
        }
    }
}

/// Return new Complex with abs() on both real and imaginary parts
#[inline]
fn abs_re_im(c: Complex<f64>) -> Complex<f64> {
    Complex {
        re: c.re.abs(),
        im: c.im.abs()
    }
}

/// Try to determine if `c` is in the Mandelbrot set, using at most `limit`
/// iterations to decide.
///
/// If `c` is not a member, return `Some(i)`, where `i` is the number of
/// iterations it took for `c` to leave the circle of radius two centered on the
/// origin. If `c` seems to be a member (more precisely, if we reached the
/// iteration limit without being able to prove that `c` is not a member),
/// return `None`.
#[inline]
fn escape_time(c: Complex<f64>, limit: u32, kind: FractalKind) -> Option<u32> {
    let mut z = Complex { re: 0.0, im: 0.0 };
    for i in 0..limit {
        z = match kind {
            FractalKind::Mandelbrot => z * z + c,
            FractalKind::Mandelbrot3 => z * z * z + c,
            FractalKind::BurningShip => {
                let z_abs = abs_re_im(z);
                z_abs * z_abs + c
            }
        };
        if z.norm_sqr() > 4.0 {
            return Some(i);
        }
    }
    None
}

/// Parse the string `s` as a coordinate pair, like `"400x600"` or `"1.0,0.5"`.
///
/// Specifically, `s` should have the form <left><sep><right>, where <sep> is
/// the character given by the `separator` argument, and <left> and <right> are both
/// strings that can be parsed by `T::from_str`.
///
/// If `s` has the proper form, return `Some<(x, y)>`. If it doesn't parse
/// correctly, return `None`.
fn parse_pair<T: FromStr>(s: &str, separator: char) -> Option<(T, T)> {
    match s.find(separator) {
        None => None,
        Some(index) => {
            match (T::from_str(&s[..index]), T::from_str(&s[index + 1..])) {
                (Ok(l), Ok(r)) => Some((l, r)),
                _ => None
            }
        }
    }
}

/// Parse a pair of floating-point numbers separated by a comma as a complex
/// number.
fn parse_complex(s: &str) -> Option<Complex<f64>> {
    match parse_pair(s, ',') {
        Some((re, im)) => Some(Complex { re, im }),
        None => None
    }
}

/// Given the row and column of a pixel in the output image, return the
/// corresponding point on the complex plane.
///
/// `bounds` is a pair giving the width and height of the image in pixels.
/// `pixel` is a (column, row) pair indicating a particular pixel in that image.
/// The `upper_left` and `lower_right` parameters are points on the complex
/// plane designating the area our image covers.
#[inline]
fn pixel_to_point(
    bounds: (usize, usize),
    pixel: (usize, usize),
    upper_left: Complex<f64>,
    lower_right: Complex<f64>) -> Complex<f64>
{
    let plane_width = lower_right.re - upper_left.re;
    let plane_height = upper_left.im - lower_right.im;

    let ratio_x = pixel.0 as f64 / bounds.0 as f64;
    let ratio_y = pixel.1 as f64 / bounds.1 as f64;

    Complex {
        re: upper_left.re + (plane_width * ratio_x),
        im: upper_left.im - (plane_height * ratio_y)
    }
}

/// Render a rectangle of the Mandelbrot set into a buffer of pixels.
///
/// The `bounds` argument gives the width and height of the buffer `pixels`,
/// which holds one grayscale pixel per byte. The `upper_left` and `lower_right`
/// arguments specify points on the complex plane corresponding to the upper-
/// left and lower-right corners of the pixel buffer.
fn render(
    pixels: &mut [u8],
    bounds: (usize, usize),
    upper_left: Complex<f64>,
    lower_right: Complex<f64>,
    kind: FractalKind)
{
    assert!(pixels.len() == bounds.0 * bounds.1);

    for y in 0 .. bounds.1 {
        for x in 0 .. bounds.0 {
            let point = pixel_to_point(bounds, (x, y), upper_left, lower_right);
            pixels[y * bounds.0 + x] =
                match escape_time(point, 255, kind) {
                    None => 0,
                    Some(count) => 255 - count as u8
                };
        }
    }
}

fn parallel_render(
    pixels: &mut [u8],
    bounds: (usize, usize),
    upper_left: Complex<f64>,
    lower_right: Complex<f64>,
    kind: FractalKind)
{
    let threads = 8;
    let rows_per_band = bounds.1 / threads + 1;

    {
        let bands: Vec<&mut [u8]> = pixels.chunks_mut(rows_per_band * bounds.0).collect();
        crossbeam::scope(|spawner| {
            for (i, band) in bands.into_iter().enumerate() {
                let top = rows_per_band * i;
                let height = band.len() / bounds.0;
                let band_bounds = (bounds.0, height);
                let band_upper_left = pixel_to_point(bounds, (0, top), upper_left, lower_right);
                let band_lower_right = pixel_to_point(bounds, (bounds.0, top + height), upper_left, lower_right);

                spawner.spawn(move || {
                    render(band, band_bounds, band_upper_left, band_lower_right, kind);
                });
            }
        });
    }
}

fn write_image(filename: &str, pixels: &[u8], bounds: (usize, usize)) -> Result<(), std::io::Error> {
    let output = File::create(filename)?;
    let encoder = PNGEncoder::new(output);
    encoder.encode(&pixels, bounds.0 as u32, bounds.1 as u32, ColorType::Gray(8))?;
    Ok(())
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() != 6 {
        writeln!(std::io::stderr(), "Usage: {} FILE PIXELS UPPERLEFT LOWERRIGHT KIND", args[0]).unwrap();
        writeln!(std::io::stderr(), "Example: {} mandel.png 1000x750 -1.20,0.35 -1,0.20 mandelbrot", args[0]).unwrap();
        writeln!(std::io::stderr(), "Example: {} ship.png 4000x4000 -2,2 2,-2 burning_ship", args[0]).unwrap();
        std::process::exit(1);
    }

    let bounds = parse_pair(&args[2], 'x').expect("error parsing image dimensions");
    let upper_left = parse_complex(&args[3]).expect("error parsing upper left corner point");
    let lower_right = parse_complex(&args[4]).expect("error parsing lower right corner point");
    let kind = FractalKind::from_str(&args[5]).expect("error parsing fractal kind");
    let mut pixels = vec![0; bounds.0 * bounds.1];
    parallel_render(&mut pixels, bounds, upper_left, lower_right, kind);
    write_image(&args[1], &pixels, bounds).expect("error writing PNG file");
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_parse_pair() {
        assert_eq!(parse_pair::<i32>("", ','), None);
        assert_eq!(parse_pair::<i32>("10,", ','), None);
        assert_eq!(parse_pair::<i32>(",10", ','), None);
        assert_eq!(parse_pair::<i32>("10,20", ','), Some((10, 20)));
        assert_eq!(parse_pair::<i32>("10,20xy", ','), None);
        assert_eq!(parse_pair::<f64>("0.5x", 'x'), None);
        assert_eq!(parse_pair::<f64>("0.5x1.5", 'x'), Some((0.5, 1.5)));
    }

    #[test]
    fn test_parse_complex() {
        assert_eq!(parse_complex("1.25,-0.0625"), Some(Complex { re: 1.25, im: -0.0625 }));
        assert_eq!(parse_complex(",-0.0625"), None);
    }

    #[test]
    fn test_pixel_to_point() {
        assert_eq!(
            pixel_to_point(
                (100, 100),
                (25, 75),
                Complex { re: -1.0, im: 1.0 },
                Complex { re: 1.0, im: -1.0 }
            ),
            Complex { re: -0.5, im: -0.5 }
        );

        assert_eq!(
            pixel_to_point(
                (400, 100),
                (0, 0),
                Complex { re: -2.0, im: 1.0 },
                Complex { re: 2.0, im: -1.0 }
            ),
            Complex { re: -2.0, im: 1.0 }
        );

        assert_eq!(
            pixel_to_point(
                (400, 100),
                (400, 100),
                Complex { re: -2.0, im: 1.0 },
                Complex { re: 2.0, im: -1.0 }
            ),
            Complex { re: 2.0, im: -1.0 }
        );

        assert_eq!(
            pixel_to_point(
                (400, 100),
                (200, 100),
                Complex { re: -2.0, im: 1.0 },
                Complex { re: 2.0, im: -1.0 }
            ),
            Complex { re: 0.0, im: -1.0 }
        );

        assert_eq!(
            pixel_to_point(
                (400, 100),
                (400, 50),
                Complex { re: -2.0, im: 1.0 },
                Complex { re: 2.0, im: -1.0 }
            ),
            Complex { re: 2.0, im: 0.0 }
        );        
    }
}