full color!

src/lib.rs

@@ -234,30 +234,45 @@ impl Core {
 		const BASE: Sphere = Sphere {
 			position: vec3(0., 0., -BASE_R),
 			radius: BASE_R,
+			color: vec3(0.1, 0.5, 0.0),
 		};
-		fn sphere(pos: Vec3) -> Sphere {
+		fn sphere(pos: Vec3, color: Vec3) -> Sphere {
 			Sphere {
 				position: pos,
 				radius: BASE_POS.distance(pos) - BASE_R,
+				color,
 			}
 		}
 		let scene = Scene {
 			objects: vec![
 				BASE,
-				sphere(vec3(0., 0., 0.1)),
-				sphere(vec3(0.3, 0., 0.1)),
-				sphere(vec3(0.1, 0.3, 0.1)),
+				sphere(vec3(0., 0., 0.1), vec3(1., 0.3, 0.0)),
+				sphere(vec3(0.3, 0., 0.1), vec3(0.3, 1., 0.0)),
+				sphere(vec3(0.1, 0.3, 0.1), vec3(0.3, 0.0, 1.)),
 			],
 		};
 
+		#[derive(Debug, Clone, Copy)]
+		struct LightHit<'a> {
+			incident: Ray,
+			normal: Vec3,
+			light: Vec3,
+			object: &'a Sphere,
+		}
 		let mut prng = rand_pcg::Pcg64::new(42, 0);
 		let source_rays: Vec<Ray> = (0..10240).map(|_| source.make_ray(&mut prng)).collect();
 		let camera_rays: Vec<Ray> = (0..10240).map(|_| camera.make_ray(&mut prng)).collect();
 		let mut source_ray_display: Vec<Vertex> = Vec::with_capacity(source_rays.len());
-		let mut hits: Vec<Hit> = Vec::with_capacity(source_rays.len());
+		let mut hits: Vec<LightHit> = Vec::with_capacity(source_rays.len());
 		for ray in source_rays {
+			let light = Vec3::splat(1.);
 			if let Some(hit) = scene.trace_ray(ray) {
-				hits.push(hit);
+				hits.push(LightHit {
+					incident: hit.incident,
+					normal: hit.normal,
+					light,
+					object: hit.object,
+				});
 				if args.show_hit_emission {
 					source_ray_display.extend([
 						Vertex {
@@ -274,7 +289,7 @@ impl Core {
 					source_ray_display.extend([
 						Vertex {
 							pos: hit.incident.base - 0.02 * hit.incident.dir,
-							color: vec3(0., 0., 1.),
+							color: light,
 						},
 						Vertex {
 							pos: hit.incident.base,
@@ -300,20 +315,26 @@ impl Core {
 		if args.reflections > 0 {
 			let mut hits1 = hits.clone();
 			for _ in 0..args.reflections {
-				let mut hits2: Vec<Hit> = Vec::with_capacity(hits1.len());
+				let mut hits2: Vec<LightHit> = Vec::with_capacity(hits1.len());
 				for hit in &hits1 {
 					let reflector = Lambertian;
 					let reflected = reflector.reflect(&mut prng, hit.normal, hit.incident.dir);
+					let light = hit.light * hit.object.color;
 					let ray = Ray::new(hit.incident.base, reflected);
 					let Some(hit2) = scene.trace_ray(ray) else {
 						continue;
 					};
-					hits2.push(hit2);
+					hits2.push(LightHit {
+						incident: hit2.incident,
+						normal: hit2.normal,
+						light: light,
+						object: hit2.object,
+					});
 					if args.show_indirect_hit {
 						source_ray_display.extend([
 							Vertex {
 								pos: hit2.incident.base - 0.02 * hit2.incident.dir,
-								color: vec3(1., 0., 1.),
+								color: light,
 							},
 							Vertex {
 								pos: hit2.incident.base,
@@ -330,8 +351,8 @@ impl Core {
 			let sigma2 = args.accum_sigma.powi(2);
 			let accum_normalizator = (2. * PI * sigma2).recip();
 			let hits = &hits;
-			move |hit: Hit| -> f32 {
-				let mut total_cd = 0.0f32;
+			move |hit: Hit| -> Vec3 {
+				let mut total_cd = Vec3::splat(0.);
 				for light_hit in hits {
 					let d2 = hit.incident.base.distance_squared(light_hit.incident.base);
 					if d2 > 9. * sigma2 {
@@ -344,16 +365,18 @@ impl Core {
 					assert!(normal.is_normalized());
 					assert!(hit.incident.dir.is_normalized());
 					let reflector = Lambertian;
-					let in_lm = 1.0;
-					let out_cd =
-						in_lm * reflector.brdf(normal, light_hit.incident.dir, -hit.incident.dir);
+					let in_lm = light_hit.light;
+					let out_cd = in_lm
+						* reflector.brdf(normal, light_hit.incident.dir, -hit.incident.dir)
+						* hit.object.color;
 					let weight = accum_normalizator * (-0.5 * d2 / sigma2).exp();
 					total_cd += weight * out_cd;
 				}
 				total_cd * args.accum_scale
 			}
 		};
-		let colormap = |light: f32| {
+		let colormap = |light: Vec3| {
+			let light = light.dot(Vec3::splat(1. / 3.));
 			let brightness = 3. * (1. - (1. + light).recip());
 			vec3(brightness, brightness - 1., brightness - 2.)
 				.clamp(Vec3::splat(0.), Vec3::splat(1.))
@@ -397,11 +420,11 @@ impl Core {
 							if !dir.is_finite() {
 								dir = camera.direction();
 							}
-							let light = light_at(Hit {
+							let color = light_at(Hit {
 								incident: Ray::new(pos, dir),
 								normal,
+								object: obj,
 							});
-							let color = Vec3::splat(light);
 							render::faces::Vertex { pos, color }
 						})
 						.collect::<Vec<_>>(),

src/render/colormap.wgsl

@@ -14,9 +14,10 @@ struct Varying {
 
 @group(0) @binding(0) var<uniform> look: LookParams;
 
-fn colormap(light: f32) -> vec3f {
-	let brightness = 3. * (1. - 1. / (1. + light));
-	return vec3(brightness, brightness - 1., brightness - 2.);
+fn colormap(light: vec3f) -> vec3f {
+	let avg = dot(light, vec3f(1. / 3.));
+	let scale = 1. / (1. + avg);
+	return scale * light;
 }
 
 @vertex
@@ -28,5 +29,5 @@ fn on_vertex(in: Vertex) -> Varying {
 
 @fragment
 fn on_fragment(in: Varying) -> @location(0) vec4f {
-	return vec4f(colormap(in.color.x), 1.0);
+	return vec4f(colormap(in.color.rgb), in.color.a);
 }

src/trace.rs

@@ -76,6 +76,7 @@ impl Source {
 pub struct Sphere {
 	pub position: Vec3,
 	pub radius: f32,
+	pub color: Vec3,
 }
 
 struct Hit1 {
@@ -113,26 +114,28 @@ pub struct Scene {
 }
 
 #[derive(Debug, Clone, Copy)]
-pub struct Hit {
+pub struct Hit<'a> {
 	pub incident: Ray,
 	pub normal: Vec3,
+	pub object: &'a Sphere,
 }
 
 impl Scene {
-	pub fn trace_ray(&self, ray: Ray) -> Option<Hit> {
+	pub fn trace_ray(&self, ray: Ray) -> Option<Hit<'_>> {
 		const EPS: f32 = -1e-3;
 		let hit = self
 			.objects
 			.iter()
-			.filter_map(|obj| obj.trace_ray(ray))
-			.filter(|h| h.dist >= EPS)
-			.min_by(|a, b| f32::total_cmp(&a.dist, &b.dist))?;
+			.filter_map(|obj| obj.trace_ray(ray).map(|hit| (obj, hit)))
+			.filter(|(_obj, h)| h.dist >= EPS)
+			.min_by(|a, b| f32::total_cmp(&a.1.dist, &b.1.dist))?;
 		Some(Hit {
 			incident: Ray {
-				base: hit.pos,
+				base: hit.1.pos,
 				dir: ray.dir,
 			},
-			normal: hit.normal,
+			normal: hit.1.normal,
+			object: hit.0,
 		})
 	}
 }