mandelbrot/src/mandelbrot.rs

use num_complex::Complex;
use rayon::prelude::*;

const MAX_ITERATIONS: u16 = 512;

#[derive(Debug, Clone)]
pub struct Mandelbrot {
    texture_buffer: Vec<u8>,
}

impl Mandelbrot {
    const IMG_WIDTH: usize = 1280;
    const IMG_HEIGHT: usize = 720;

    pub fn new() -> Self {
        Mandelbrot {
            texture_buffer: vec![0; (Self::IMG_WIDTH * Self::IMG_HEIGHT) * 4],
        }
    }

    pub fn width() -> usize {
        Self::IMG_WIDTH
    }

    pub fn height() -> usize {
        Self::IMG_HEIGHT
    }

    pub fn generate_image(&mut self) -> Vec<u8> {
        self.texture_buffer
            .par_chunks_mut(4)
            .enumerate()
            .for_each(|(i, buf)| {
                let iters = Self::new_escape_time(i, (-2.5, 1.0), (-1.0, 1.0), 64);
                let normalized_iters = iters / MAX_ITERATIONS as f64;
                let h = 0.5 + (10.0 * normalized_iters);
                buf.copy_from_slice(&Self::hsv_to_rgb(h, 0.6, 1.0));
            });

        self.texture_buffer.clone()
    }

    pub fn generate_scaled_image(
        &mut self,
        x_bounds: (f64, f64),
        y_bounds: (f64, f64),
        max_iterations: u32,
    ) -> &[u8] {
        self.texture_buffer
            .par_chunks_mut(4)
            .enumerate()
            .for_each(|(i, buf)| {
                let max_iters = max_iterations as f64;

                let iters = Self::new_escape_time(i, x_bounds, y_bounds, max_iterations);
                let normalized_iters = iters / max_iters;

                let h = normalized_iters * 350.0;
                let v = if iters == max_iters { 0.0 } else { 1.0 };
                buf.copy_from_slice(&Self::hsv_to_rgb(h, 1.0, v))
            });

        &self.texture_buffer
    }

    fn new_escape_time(
        i: usize,
        x_bounds: (f64, f64),
        y_bounds: (f64, f64),
        max_iterations: u32,
    ) -> f64 {
        let px = i % Self::IMG_WIDTH;
        let py = i / Self::IMG_WIDTH;
        let c = Self::coords_to_complex(px, py, x_bounds, y_bounds);

        Self::new_calc_num_iters(c, max_iterations)
    }

    #[inline]
    fn hsv_to_rgb(h: f64, s: f64, v: f64) -> [u8; 4] {
        let c = v * s;
        let x = c * (1.0 - (((h / 60.0) % 2.0) - 1.0).abs());
        let m = v - c;

        let (r_prime, g_prime, b_prime) = {
            if h < 60.0 {
                (c, x, 0.0)
            } else if h < 120.0 {
                (x, c, 0.0)
            } else if h < 180.0 {
                (0.0, c, x)
            } else if h < 240.0 {
                (0.0, x, c)
            } else if h < 300.0 {
                (x, 0.0, c)
            } else if h < 360.0 {
                (c, 0.0, x)
            } else {
                unreachable!()
            }
        };

        let r = ((r_prime + m) * 255.0) as u8;
        let g = ((g_prime + m) * 255.0) as u8;
        let b = ((b_prime + m) * 255.0) as u8;
        let a = 0xFF;

        [r, g, b, a]
    }

    fn new_calc_num_iters(c: Complex<f64>, max_iterations: u32) -> f64 {
        let mut z: Complex<f64> = Complex::new(0.0, 0.0);
        let mut num_iters: u32 = 0;

        loop {
            if z.norm_sqr() > 4.0 {
                break;
            }

            if num_iters >= max_iterations {
                break;
            }

            z = (z * z) + c;
            num_iters += 1;
        }

        if num_iters < max_iterations {
            (num_iters as f64 + 1.0) - z.norm().ln().ln() / 2f64.ln()
        } else {
            num_iters as f64
        }
    }

    fn coords_to_complex(
        px: usize,
        py: usize,
        x_bounds: (f64, f64),
        y_bounds: (f64, f64),
    ) -> Complex<f64> {
        Complex::new(
            Self::scale_width(px, x_bounds),
            Self::scale_height(py, y_bounds),
        )
    }

    fn scale_width(px: usize, x_bounds: (f64, f64)) -> f64 {
        // const X_MIN: f64 = -2.5;
        // const X_MAX: f64 = 1.0;

        x_bounds.0 + ((x_bounds.1 - x_bounds.0) * (px as f64 - 0.0)) / Self::IMG_WIDTH as f64 - 0.0
    }

    fn scale_height(py: usize, y_bounds: (f64, f64)) -> f64 {
        // const Y_MIN: f64 = -1.0;
        // const Y_MAX: f64 = 1.0;

        y_bounds.0 + ((y_bounds.1 - y_bounds.0) * (py as f64 - 0.0)) / Self::IMG_HEIGHT as f64 - 0.0
    }
}