julia_set/
cpu.rs

1//! CPU backend for Julia set rendering.
2
3use std::convert::Infallible;
4
5use num_complex::Complex32;
6use rayon::prelude::*;
7
8use crate::{Backend, ImageBuffer, Params, Render};
9
10/// Backend that uses CPU for computations.
11///
12/// The current implementation is based on the [`rayon`] crate.
13///
14/// [`rayon`]: https://crates.io/crates/rayon
15#[cfg_attr(docsrs, doc(cfg(feature = "cpu_backend")))]
16#[derive(Debug, Clone, Copy, Default)]
17pub struct Cpu;
18
19impl<F: ComputePoint> Backend<F> for Cpu {
20    type Error = Infallible;
21    type Program = CpuProgram<F>;
22
23    fn create_program(&self, function: F) -> Result<Self::Program, Self::Error> {
24        Ok(CpuProgram::new(function))
25    }
26}
27
28#[derive(Debug, Clone, Copy)]
29struct CpuParams {
30    image_size: [u32; 2],
31    image_size_f32: [f32; 2],
32    view_size: [f32; 2],
33    view_center: Complex32,
34    inf_distance_sq: f32,
35    max_iterations: u8,
36}
37
38impl CpuParams {
39    #[allow(clippy::cast_precision_loss)] // The loss of precision is acceptable
40    fn new(params: &Params) -> Self {
41        Self {
42            image_size: params.image_size,
43            image_size_f32: [params.image_size[0] as f32, params.image_size[1] as f32],
44            view_size: [params.view_width(), params.view_height],
45            view_center: Complex32::new(params.view_center[0], params.view_center[1]),
46            inf_distance_sq: params.inf_distance * params.inf_distance,
47            max_iterations: params.max_iterations,
48        }
49    }
50
51    #[allow(clippy::cast_precision_loss)] // The loss of precision is acceptable
52    fn map_pixel(self, pixel_row: u32, pixel_col: u32) -> Complex32 {
53        let [width, height] = self.image_size_f32;
54        let [view_width, view_height] = self.view_size;
55
56        let re = (pixel_col as f32 + 0.5) / width;
57        let re = (re - 0.5) * view_width;
58        let im = (pixel_row as f32 + 0.5) / height;
59        let im = (0.5 - im) * view_height;
60        Complex32::new(re, im) + self.view_center
61    }
62}
63
64/// Complex-valued function of a single variable.
65#[cfg_attr(docsrs, doc(cfg(feature = "cpu_backend")))]
66pub trait ComputePoint: Sync {
67    /// Computes the function value at the specified point.
68    fn compute_point(&self, z: Complex32) -> Complex32;
69}
70
71/// Programs output by the [`Cpu`] backend. Come in two varieties depending on the type param:
72/// native closures, or interpreted [`Function`](crate::Function)s.
73#[cfg_attr(docsrs, doc(cfg(feature = "cpu_backend")))]
74#[derive(Debug)]
75pub struct CpuProgram<F> {
76    function: F,
77}
78
79impl<F: Fn(Complex32) -> Complex32 + Sync> ComputePoint for F {
80    fn compute_point(&self, z: Complex32) -> Complex32 {
81        self(z)
82    }
83}
84
85impl<F: ComputePoint> CpuProgram<F> {
86    fn new(function: F) -> Self {
87        Self { function }
88    }
89
90    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
91    fn compute_row(&self, params: CpuParams, pixel_row: u32) -> Vec<u8> {
92        let [image_width, _] = params.image_size;
93
94        let pixels = (0..image_width).map(|pixel_col| {
95            let mut z = params.map_pixel(pixel_row, pixel_col);
96            let mut iter = params.max_iterations;
97
98            for i in 0..params.max_iterations {
99                z = self.function.compute_point(z);
100                if z.is_nan() || z.is_infinite() || z.norm_sqr() > params.inf_distance_sq {
101                    iter = i;
102                    break;
103                }
104            }
105
106            let color = f32::from(iter) / f32::from(params.max_iterations);
107            (color * 255.0).round() as u8 // cannot truncate or lose sign by construction
108        });
109        pixels.collect()
110    }
111}
112
113impl<F: ComputePoint> Render for CpuProgram<F> {
114    type Error = Infallible;
115
116    fn render(&self, params: &Params) -> Result<ImageBuffer, Self::Error> {
117        let [width, height] = params.image_size;
118        let pixel_size = (width * height) as usize;
119        let params = CpuParams::new(params);
120
121        let buffer: Vec<u8> = (0..height)
122            .into_par_iter()
123            .fold(
124                || Vec::with_capacity(pixel_size),
125                |mut buffer, pixel_row| {
126                    let line = self.compute_row(params, pixel_row);
127                    buffer.extend_from_slice(&line);
128                    buffer
129                },
130            )
131            .flatten()
132            .collect();
133        Ok(ImageBuffer::from_raw(width, height, buffer).unwrap())
134    }
135}
136
137#[cfg(feature = "dyn_cpu_backend")]
138mod dynamic {
139    use std::{collections::HashMap, convert::Infallible};
140
141    use arithmetic_parser::BinaryOp;
142    use num_complex::Complex32;
143
144    use super::{ComputePoint, Cpu, CpuProgram};
145    use crate::{Backend, Function, function::Evaluated};
146
147    impl Backend<&Function> for Cpu {
148        type Error = Infallible;
149        type Program = CpuProgram<Function>;
150
151        fn create_program(&self, function: &Function) -> Result<Self::Program, Self::Error> {
152            Ok(CpuProgram::new(function.clone()))
153        }
154    }
155
156    fn eval(expr: &Evaluated, variables: &HashMap<&str, Complex32>) -> Complex32 {
157        match expr {
158            Evaluated::Variable(s) => variables[s.as_str()],
159            Evaluated::Value(val) => *val,
160            Evaluated::Negation(inner) => -eval(inner, variables),
161            Evaluated::Binary { op, lhs, rhs } => {
162                let lhs_value = eval(lhs, variables);
163                let rhs_value = eval(rhs, variables);
164                match op {
165                    BinaryOp::Add => lhs_value + rhs_value,
166                    BinaryOp::Sub => lhs_value - rhs_value,
167                    BinaryOp::Mul => lhs_value * rhs_value,
168                    BinaryOp::Div => lhs_value / rhs_value,
169                    BinaryOp::Power => lhs_value.powc(rhs_value),
170                    _ => unreachable!(),
171                }
172            }
173            Evaluated::FunctionCall { function, arg } => {
174                let evaluated_arg = eval(arg, variables);
175                function.eval(evaluated_arg)
176            }
177        }
178    }
179
180    impl ComputePoint for Function {
181        fn compute_point(&self, z: Complex32) -> Complex32 {
182            let mut variables = HashMap::new();
183            variables.insert("z", z);
184
185            for (var_name, expr) in self.assignments() {
186                let expr = eval(expr, &variables);
187                variables.insert(var_name, expr);
188            }
189            eval(self.return_value(), &variables)
190        }
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    fn assert_close(x: Complex32, y: Complex32) {
199        assert!((x.re - y.re).abs() <= f32::EPSILON, "{x:?}, {y:?}");
200        assert!((x.im - y.im).abs() <= f32::EPSILON, "{x:?}, {y:?}");
201    }
202
203    #[test]
204    fn mapping_pixels() {
205        let params = Params::new([100, 100], 1.0);
206        let params = CpuParams::new(&params);
207        assert_close(params.map_pixel(0, 0), Complex32::new(-0.495, 0.495));
208        assert_close(params.map_pixel(0, 50), Complex32::new(0.005, 0.495));
209        assert_close(params.map_pixel(0, 100), Complex32::new(0.505, 0.495));
210        assert_close(params.map_pixel(50, 0), Complex32::new(-0.495, -0.005));
211        assert_close(params.map_pixel(50, 50), Complex32::new(0.005, -0.005));
212        assert_close(params.map_pixel(50, 100), Complex32::new(0.505, -0.005));
213        assert_close(params.map_pixel(100, 0), Complex32::new(-0.495, -0.505));
214        assert_close(params.map_pixel(100, 50), Complex32::new(0.005, -0.505));
215        assert_close(params.map_pixel(100, 100), Complex32::new(0.505, -0.505));
216    }
217
218    #[test]
219    #[cfg(feature = "dyn_cpu_backend")]
220    fn compute() {
221        use crate::Function;
222
223        let program: Function = "z * z + 0.5i".parse().unwrap();
224        assert_eq!(
225            program.compute_point(Complex32::new(0.0, 0.0)),
226            Complex32::new(0.0, 0.5)
227        );
228        assert_eq!(
229            program.compute_point(Complex32::new(1.0, 0.0)),
230            Complex32::new(1.0, 0.5)
231        );
232        assert_eq!(
233            program.compute_point(Complex32::new(-1.0, 0.0)),
234            Complex32::new(1.0, 0.5)
235        );
236        assert_eq!(
237            program.compute_point(Complex32::new(0.0, 1.0)),
238            Complex32::new(-1.0, 0.5)
239        );
240    }
241
242    #[test]
243    #[cfg(feature = "dyn_cpu_backend")]
244    fn compute_does_not_panic() {
245        use crate::Function;
246
247        let program: Function = "1.0 / z + 0.5i".parse().unwrap();
248        let z = program.compute_point(Complex32::new(0.0, 0.0));
249        assert!(z.is_nan());
250    }
251}