raytracing-rs

https://raytracing.github.io in Rust

git clone https://code.pdelong.com/raytracing-rs.git

  1use std::collections::BinaryHeap;
  2use std::collections::binary_heap::PeekMut;
  3use std::io::Write;
  4use std::sync::Arc;
  5use std::sync::mpsc;
  6
  7use rand::Rng;
  8use rand::SeedableRng;
  9use rayon::prelude::*;
 10
 11use crate::hittable::Hittable;
 12use crate::hittable::HittableList;
 13use crate::math::Color;
 14use crate::math::Interval;
 15use crate::math::Point;
 16use crate::math::Ray;
 17use crate::math::Vec3;
 18
 19#[derive(Default)]
 20pub struct CameraBuilder {
 21    aspect_ratio: Option<f64>,
 22    image_width: Option<u32>,
 23    samples_per_pixel: Option<u32>,
 24    max_depth: Option<i32>,
 25}
 26
 27impl CameraBuilder {
 28    pub fn new() -> Self {
 29        Self::default()
 30    }
 31
 32    pub fn build(self) -> Camera {
 33        let aspect_ratio = self.aspect_ratio.unwrap_or(16. / 9.);
 34        let image_width = self.image_width.unwrap_or(1024);
 35        let image_height = ((image_width as f64) / aspect_ratio) as u32;
 36        let samples_per_pixel = self.samples_per_pixel.unwrap_or(10);
 37        let max_depth = self.max_depth.unwrap_or(10);
 38
 39        Camera::new(CameraArgs {
 40            image_width,
 41            image_height,
 42            samples_per_pixel,
 43            max_depth,
 44        })
 45    }
 46
 47    pub fn image_width(mut self, image_width: u32) -> Self {
 48        self.image_width = Some(image_width);
 49
 50        self
 51    }
 52
 53    pub fn aspect_ratio(mut self, aspect_ratio: f64) -> Self {
 54        self.aspect_ratio = Some(aspect_ratio);
 55
 56        self
 57    }
 58
 59    pub fn samples_per_pixel(mut self, samples_per_pixel: u32) -> Self {
 60        self.samples_per_pixel = Some(samples_per_pixel);
 61
 62        self
 63    }
 64
 65    pub fn max_depth(mut self, max_depth: i32) -> Self {
 66        self.max_depth = Some(max_depth);
 67
 68        self
 69    }
 70}
 71
 72pub struct CameraArgs {
 73    image_width: u32,
 74    image_height: u32,
 75    samples_per_pixel: u32,
 76    max_depth: i32,
 77}
 78
 79pub struct Camera {
 80    image_width: u32,
 81    image_height: u32,
 82    center: Point,
 83    pixel_delta_u: Vec3,
 84    pixel_delta_v: Vec3,
 85    pixel00_loc: Point,
 86    samples_per_pixel: u32,
 87    max_depth: i32,
 88}
 89impl Camera {
 90    pub fn new(args: CameraArgs) -> Self {
 91        let image_width = args.image_width;
 92        let image_height = args.image_height;
 93        let viewport_height = 2.;
 94        let viewport_width = viewport_height * ((image_width as f64) / (image_height as f64));
 95        let focal_length = 1.;
 96        let center = Vec3 {
 97            x: 0.,
 98            y: 0.,
 99            z: 0.,
100        };
101
102        let viewport_u = Vec3 {
103            x: viewport_width,
104            y: 0.,
105            z: 0.,
106        };
107        let viewport_v = Vec3 {
108            x: 0.,
109            y: -viewport_height,
110            z: 0.,
111        };
112        let pixel_delta_u = viewport_u / (image_width as f64);
113        let pixel_delta_v = viewport_v / (image_height as f64);
114
115        let viewport_upper_left = center
116            - Vec3 {
117                x: 0.,
118                y: 0.,
119                z: focal_length,
120            }
121            - viewport_v / 2.
122            - viewport_u / 2.;
123        let pixel00_loc = viewport_upper_left + (pixel_delta_u + pixel_delta_v) / 2.;
124
125        Self {
126            image_width,
127            image_height,
128            center,
129            pixel_delta_u,
130            pixel_delta_v,
131            pixel00_loc,
132            samples_per_pixel: args.samples_per_pixel,
133            max_depth: args.max_depth,
134        }
135    }
136
137    pub fn render<W: Write + Send>(&mut self, mut w: W, world: HittableList) {
138        let world = Arc::new(world);
139        rayon::scope(|s| {
140            let (tx, rx) = mpsc::channel();
141            s.spawn(|_| {
142                (0..self.image_height)
143                    .into_par_iter()
144                    .map(|y| {
145                        let mut rng = rand::rngs::SmallRng::from_os_rng();
146                        let world = world.clone();
147                        let mut row = Vec::new();
148                        for x in 0..self.image_width {
149                            let mut color = Color {
150                                x: 0.,
151                                y: 0.,
152                                z: 0.,
153                            };
154
155                            for _ in 0..self.samples_per_pixel {
156                                let sample_loc = self.pixel00_loc
157                                    + self.pixel_delta_v
158                                        * (y as f64 + rng.random_range(-0.5..=0.5))
159                                    + self.pixel_delta_u
160                                        * (x as f64 + rng.random_range(-0.5..=0.5));
161
162                                let dir = (sample_loc - self.center).normalize();
163                                let r = Ray {
164                                    orig: self.center,
165                                    dir,
166                                };
167
168                                let ray_color = ray_color(&mut rng, self.max_depth, &r, &world);
169                                color += ray_color / self.samples_per_pixel as f64;
170                            }
171                            row.push(color);
172                        }
173
174                        (row, y)
175                    })
176                    .for_each(|(row, y)| tx.send((row, y as usize)).unwrap());
177
178                drop(tx);
179            });
180
181            writeln!(w, "P6\n{} {}\n65535", self.image_width, self.image_height).unwrap();
182            let mut remaining = self.image_height;
183            let iter = ReorderIterator::new(
184                rx.iter().inspect(|_| {
185                    remaining -= 1;
186                    eprint!("\rScanlines remaining: {remaining} ");
187                }),
188                0,
189            );
190            for row in iter {
191                for color in row {
192                    w.write_all(&color.bytes()).unwrap();
193                }
194            }
195        });
196
197        eprintln!("\rDone!                  ");
198    }
199}
200
201fn ray_color<R: rand::Rng>(rng: &mut R, depth: i32, r: &Ray, world: &HittableList) -> Color {
202    if depth <= 0 {
203        return Color {
204            x: 0.,
205            y: 0.,
206            z: 0.,
207        };
208    }
209
210    if let Some(hit) = world.hit(r, &Interval::new(0.0001, f64::INFINITY).unwrap()) {
211        let new_dir = Vec3::random_on_unit_hemisphere(rng, &hit.normal);
212        let new_ray = Ray {
213            orig: hit.point,
214            dir: new_dir,
215        };
216        ray_color(rng, depth - 1, &new_ray, world) / 2.
217    } else {
218        let a = (r.dir.y() + 1.) / 2.;
219        let white = Color {
220            x: 1.,
221            y: 1.,
222            z: 1.,
223        };
224        let blue = Color {
225            x: 0.5,
226            y: 0.7,
227            z: 1.,
228        };
229
230        white * (1. - a) + blue * a
231    }
232}
233
234struct State<T> {
235    row: T,
236    y: usize,
237}
238
239impl<T> PartialEq for State<T> {
240    fn eq(&self, other: &Self) -> bool {
241        self.y == other.y
242    }
243}
244
245impl<T> Eq for State<T> {}
246
247impl<T> PartialOrd for State<T> {
248    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
249        Some(self.cmp(other))
250    }
251}
252
253impl<T> Ord for State<T> {
254    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
255        // note that this will put things in the opposite order, which is what we want here.
256        other.y.cmp(&self.y)
257    }
258}
259
260struct ReorderIterator<I, T> {
261    inner: I,
262    heap: BinaryHeap<State<T>>,
263    next: usize,
264}
265
266impl<I, T> ReorderIterator<I, T> {
267    fn new(inner: I, first: usize) -> Self {
268        ReorderIterator {
269            inner,
270            heap: BinaryHeap::new(),
271            next: first,
272        }
273    }
274}
275
276impl<I: Iterator<Item = (T, usize)>, T> Iterator for ReorderIterator<I, T> {
277    type Item = T;
278
279    fn next(&mut self) -> Option<Self::Item> {
280        if let Some(s) = self.heap.peek_mut() {
281            if s.y == self.next {
282                self.next += 1;
283                return Some(PeekMut::pop(s).row);
284            } else if s.y < self.next {
285                panic!("there were duplicates!")
286            }
287        }
288
289        for (v, i) in self.inner.by_ref() {
290            if i == self.next {
291                self.next += 1;
292                return Some(v);
293            }
294
295            self.heap.push(State { row: v, y: i });
296        }
297
298        if !self.heap.is_empty() {
299            panic!("gaps!")
300        }
301
302        // In either case, we need to read out of the underlying iterator until
303        None
304    }
305}