ray tracing in one weekend implemented in ReScript

README.md

Ray Tracing in One Weekend

周末光追的 ReScript 实现

直接编译运行即可

Rt.res

@inline
let v3create = (a, b, c) => {
    [a, b, c]
}
@inline
let v3x = (a) => Js.Array2.unsafe_get(a, 0)
@inline
let v3y = (a) => Js.Array2.unsafe_get(a, 1)
@inline
let v3z = (a) => Js.Array2.unsafe_get(a, 2)
@inline
let v3vneg = (a) => [-.v3x(a), -.v3y(a), -.v3z(a)]
@inline
let v3vadd = (a, b) => [v3x(a) +. v3x(b), v3y(a) +. v3y(b), v3z(a) +. v3z(b)]
@inline
let v3vsub = (a, b) => [v3x(a) -. v3x(b), v3y(a) -. v3y(b), v3z(a) -. v3z(b)]
@inline
let v3vmul = (a, b) => [v3x(a) *. v3x(b), v3y(a) *. v3y(b), v3z(a) *. v3z(b)]
@inline
let v3vdiv = (a, b) => [v3x(a) /. v3x(b), v3y(a) /. v3y(b), v3z(a) /. v3z(b)]
@inline
let v3mul = (a, t) => [v3x(a) *. t, v3y(a) *. t, v3z(a) *. t]
@inline
let v3div = (a, t) => [v3x(a) /. t, v3y(a) /. t, v3z(a) /. t]
@inline
let v3setx = (a, v) => Js.Array2.unsafe_set(a, 0, v)
@inline
let v3sety = (a, v) => Js.Array2.unsafe_set(a, 1, v)
@inline
let v3setz = (a, v) => Js.Array2.unsafe_set(a, 2, v)
@inline
let v3set = (a, b) => { Js.Array2.unsafe_set(a, 0, Js.Array2.unsafe_get(b, 0));
			Js.Array2.unsafe_set(a, 1, Js.Array2.unsafe_get(b, 1));
			Js.Array2.unsafe_set(a, 2, Js.Array2.unsafe_get(b, 2));}
@inline
let v3vnegInPlace = (a) => { v3setx(a, -.v3x(a));
			     v3sety(a, -.v3y(a));
			     v3setz(a, -.v3z(a)) }
@inline
let v3vaddInPlace = (a, b) => { v3setx(a, v3x(a) +. v3x(b));
				v3sety(a, v3y(a) +. v3y(b));
				v3setz(a, v3z(a) +. v3z(b))}
@inline
let v3vsubInPlace = (a, b) => { v3setx(a, v3x(a) -. v3x(b));
				v3sety(a, v3y(a) -. v3y(b));
				v3setz(a, v3z(a) -. v3z(b))}
@inline
let v3vmulInPlace = (a, b) => { v3setx(a, v3x(a) *. v3x(b));
				v3sety(a, v3y(a) *. v3y(b));
				v3setz(a, v3z(a) *. v3z(b))}
@inline
let v3vdivInPlace = (a, b) => { v3setx(a, v3x(a) /. v3x(b));
				v3sety(a, v3y(a) /. v3y(b));
				v3setz(a, v3z(a) /. v3z(b))}
@inline
let v3mulInPlace = (a, b) => { v3setx(a, v3x(a) *. b);
			       v3sety(a, v3y(a) *. b);
			       v3setz(a, v3z(a) *. b);}
@inline
let v3divInPlace = (a, b) => { v3setx(a, v3x(a) /. b);
			       v3sety(a, v3y(a) /. b);
			       v3setz(a, v3z(a) /. b);}
@inline
let v3len = (a) => {
    Js.Math.sqrt(v3x(a) *. v3x(a) +. v3y(a) *. v3y(a) +. v3z(a) *. v3z(a))
}
@inline
let v3lenS = (a) => v3x(a) *. v3x(a) +. v3y(a) *. v3y(a) +. v3z(a) *. v3z(a)
@inline
let v3unit = (a) => {
    let le = v3len(a)
    [v3x(a) /. le, v3y(a) /. le, v3z(a) /. le]
}
@inline
let v3unitInPlace = (a) => {
    let le = v3len(a)
    v3setx(a, v3x(a) /. le)
    v3sety(a, v3y(a) /. le)
    v3setz(a, v3z(a) /. le)
}
@inline
let v3dot = (a, b) => {
    v3x(a) *. v3x(b) +. v3y(a) *. v3y(b) +. v3z(a) *. v3z(b)
}
@inline
let v3cross = (a, b) => {
    let r = v3y(a) *. v3z(b) -. v3z(a) *. v3y(b)
    let g = v3z(a) *. v3x(b) -. v3x(a) *. v3z(b)
    let b = v3x(a) *. v3y(b) -. v3y(a) *. v3x(b)
    [r, g, b]
}
@inline
let v3nearZero = (a) => {
    let epsilon = 1e-8
    let abs = Js.Math.abs_float
    abs(v3x(a)) < epsilon && abs(v3y(a)) < epsilon && abs(v3z(a)) < epsilon
}

type v3 = array<float>

type fd

@module("fs")
external openSync: (string, string) => fd = "openSync"

@module("fs")
external closeSync: (fd) => () = "closeSync"

@module("fs")
external writeSync: (fd, string) => () = "writeSync"

let newline = '\n'->String.make(1, _)

let writeV3 = (file: fd, v: v3) => {
    let r = Belt.Float.toInt(v3x(v) *. 255.999)->Belt.Int.toString
    let g = Belt.Float.toInt(v3y(v) *. 255.999)->Belt.Int.toString
    let b = Belt.Float.toInt(v3z(v) *. 255.999)->Belt.Int.toString
    writeSync(file, r ++ " " ++ g ++ " " ++ b ++ newline)
}

let gen = (~width: int, ~height: int, file: string) => {
    let f = openSync(file, "w")
    let rest = ref(width * height)
    writeSync(f, "P3" ++ newline)
    writeSync(f, width->Belt.Int.toString ++ " "
	      ++ height->Belt.Int.toString ++ newline)
    writeSync(f, "256" ++ newline)
    (. v: v3) => {
	if rest.contents == -1 {
	    Js.log("already finish")
	} else {
	    writeV3(f, v)
	    rest := rest.contents - 1
	    if (rest.contents == 0) {
		closeSync(f)
		rest := -1
	    }
	}
    }
}

let writemat = (mat: array<array<v3>>, file: string) => {
    let out = gen(~width = mat[0]->Js.Array2.length,
		  ~height = mat->Js.Array2.length,
		  file)
    mat->Belt.Array.forEachU((. a) => {
	a->Belt.Array.forEachU((. v) => {
	    out(. v)
	})
    })
}

type ray = {
    origin: v3,
    direction: v3
}

@inline
let rayorigin = (a: ray) => a.origin
@inline
let raydirection = (a: ray) => a.direction
@inline
let rayat = (a: ray, t: float) => v3vadd(rayorigin(a), raydirection(a)->v3mul(t))

@inline
let pi = 3.1415926535897932385
@val external inf: float = "Infinity"
@inline
let d2r = (deg) => deg *. pi /. 180.0
let rand = Js.Math.random
let abs = Js.Math.abs_float
let min = Js.Math.min_float

type rec hit_record = {
    p: v3,
    normal: v3,
    t: float,
    front: bool,
    mater: material
}
and material = (. ray, hit_record) => option<(v3, ray)>

@inline
let hit_set_normal = (ra, nor) => v3dot(ra->raydirection, nor) < 0.0 ? (true, nor) : (false, nor->v3vneg)

type hittable = (. ray, float, float) => option<hit_record>

let hitByList = (. ray, tmin, tmax, hit_arr: array<hittable>) => {
    hit_arr->Belt.Array.reduceU((tmax, None), (. c, a) => {
	let (curr_max, _) = c
	let res = a(. ray, tmin, curr_max)
	switch res {
		| None => c
		| Some(hi) => {
		    (hi.t, res)
		}
	}
    })->((_, ret))=>ret
}

let sphere2hittable = (center: v3, rad: float, mat: material) => {
    let hit: hittable = (. r, tmin, tmax) => {
	let oc = r->rayorigin->v3vsub(center)
	let a = r->raydirection->v3lenS
	let hb = v3dot(oc, r->raydirection)
	let c = oc->v3lenS -. rad *. rad
	let de = hb *. hb -. a *. c
	de < 0.0 ? None : {
	    let sqrtde = de->Js.Math.sqrt
	    let root = (-.hb -. sqrtde) /. a
	    if (root < tmin || tmax < root) {
		let root = (-.hb +. sqrtde) /. a
		if (root < tmin || tmax < root) {
		    None
		} else {
		    let p = r->rayat(root)
		    let nor = p->v3vsub(center)->v3div(rad)
		    let (front, normal) = hit_set_normal(r, nor)
		    Some({t: root,
			  p: p,
			  normal: normal,
			  front: front,
			  mater: mat})
		}
	    } else {
		let p = r->rayat(root)
		let nor = p->v3vsub(center)->v3div(rad)
		let (front, normal) = hit_set_normal(r, nor)
		Some({t: root,
		      p: p,
		      normal: normal,
		      front: front,
		      mater: mat})
	    }
	}
    }
    hit
}

let random_in_unit_disk = () => {
    let res = ref(v3create(0.0, 0.0, 0.0))
    let flag = ref(true)
    while flag.contents {
	let p = v3create(rand() *. 2.0 -. 1.0, rand() *. 2.0 -. 1.0, 0.0)
	if (p->v3lenS < 1.0) {
	    res := p
	    flag := false
	}
    }
    res.contents
}

type cam = {
    origin: v3,
    lower_left_corner: v3,
    horizontal: v3,
    vertical: v3,
    u: v3,
    v: v3,
    w: v3,
    lens_radius: float
}

let camcreate = (~lookfrom,
		 ~lookat,
		 ~vup,
		 ~aspect_ratio,
		 ~viewport_height as vh,
		 ~vfov,
		 ~aperture,
		 ~focus_dist) =>
    {
	let theta = d2r(vfov)
	let h = Js.Math.tan(theta /. 2.0)
	let viewport_height = vh *. h
	let viewport_width = aspect_ratio *. viewport_height

	let w = lookfrom->v3vsub(lookat)->v3unit
	let u = v3cross(vup, w)->v3unit
	let v = v3cross(w, u)

	let origin = lookfrom
	let horizontal = v3mul(u, viewport_width *. focus_dist)
	let vertical = v3mul(v, viewport_height *. focus_dist)
	let lower_left_corner = origin
	    ->v3vsub(v3div(horizontal, 2.0))
	    ->v3vsub(v3div(vertical, 2.0))
	    ->v3vsub(v3mul(w, focus_dist))
	{
	    origin: origin,
	    lower_left_corner: lower_left_corner,
	    horizontal: horizontal,
	    vertical: vertical,
	    w: w,
	    u: u,
	    v: v,
	    lens_radius: aperture /. 2.0
	}
    }

let camgetray = (. came: cam, s: float, t: float) => {
    let rd = (random_in_unit_disk())->v3mul(came.lens_radius)
    let offset = v3vadd(came.u->v3mul(v3x(rd)), came.v->v3mul(v3y(rd)))
    {origin: came.origin->v3vadd(offset),
     direction: came.lower_left_corner
     ->v3vadd(v3mul(came.horizontal, s))
     ->v3vadd(v3mul(came.vertical, t))
     ->v3vsub(came.origin)
     ->v3vsub(offset)}
}

@inline
let clamp = (x: float, min: float, max: float) => {
    x < min ? min : x > max ? max : x
}

let random_in_unit_sphere = () => {
    let flag = ref(true)
    let r = ref(v3create(0.0, 0.0, 0.0))
    while (flag.contents) {
	let r1 = v3create(rand() *. 2.0 -. 1.0,
			  rand() *. 2.0 -. 1.0,
			  rand() *. 2.0 -. 1.0)
	if (v3lenS(r1) <= 1.0) {
	    flag := false
	    r := r1
	}
    }
    r.contents
}

let random_in_hemisphere = (normal) => {
    let ve = random_in_unit_sphere()
    v3dot(ve, normal) > 0.0 ? ve : v3vneg(ve)
}

let lambertian = (color) => {
    let ret: material = (. _, hitre) => {
	let scatter_direction = hitre.normal->v3vadd(v3unit(random_in_unit_sphere()))
	let dir = v3nearZero(scatter_direction) ? hitre.normal : scatter_direction
	let ra = {origin: hitre.p, direction: dir}
	Some((color, ra))
    }
    ret
}

@inline
let reflect = (v: v3, n: v3) => {
    v->v3vsub(n->v3mul(v3dot(v, n) *. 2.0))
}

let metal = (color, fuzz) => {
    let f = fuzz < 0.0 ? 0.0 : fuzz < 1.0 ? fuzz : 1.0
    let ret: material = (. ray, hitre) => {
	let reflected = reflect(ray->raydirection->v3unit, hitre.normal)
	let scattered = {origin: hitre.p,
			 direction: reflected
			 ->v3vadd(v3mul(random_in_unit_sphere(), f))}
	if (v3dot(scattered->raydirection, hitre.normal) > 0.0) {
	    Some((color, scattered))
	} else {
	    None
	}
    }
    ret
}

@inline
let refract = (uv, normal, etai_div_etat) => {
    let cos_theta = min(-.v3dot(uv, normal), 1.0)
    let r_out_perp = uv->v3vadd(v3mul(normal, cos_theta))->v3mul(etai_div_etat)
    let r_out_parallel = -.Js.Math.sqrt(abs(1.0 -. r_out_perp->v3lenS))->v3mul(normal, _)
    v3vadd(r_out_parallel, r_out_perp)
}

@inline
let reflectance = (cosine, ref_idx) => {
    let r0 = (1.0 -. ref_idx) /. (1.0 +. ref_idx)
    let r1 = r0 *. r0
    let temp = (1.0 -. cosine)
    let t1 = temp *. temp
    let t2 = t1 *. t1
    let tf = t2 *. temp
    r1 +. (1.0 -. r1) *. tf
}

let dielectric = (irate) => {
    let ret: material = (. ray, hitre) => {
	let color = v3create(1.0, 1.0, 1.0)
	let ratio = hitre.front ? 1.0 /. irate : irate
	let unit_dir = ray->raydirection->v3unit
	let cos_theta = min(-1.0 *. v3dot(unit_dir, hitre.normal), 1.0)
	let sin_theta = Js.Math.sqrt(1.0 -. cos_theta *. cos_theta)
	let cannot_refract = ratio *. sin_theta > 1.0
	let direction = cannot_refract || reflectance(cos_theta, ratio) > rand() ?
	    reflect(unit_dir, hitre.normal) :
	    refract(unit_dir, hitre.normal, ratio)
	let scattered = {origin: hitre.p,
			 direction: direction}
	Some(color, scattered)
    }
    ret
}

let random_scene = () => {
    let ground_material = lambertian(v3create(0.5, 0.5, 0.5))
    let world = []
    let add = (element)=>ignore(world->Js.Array2.push(element))

    add(sphere2hittable(v3create(0.0, -1000.0, 0.0), 1000.0, ground_material))

    for a in -11 to 10 {
	for b in -11 to 10 {
	    let fa = a->Belt.Int.toFloat
	    let fb = b->Belt.Int.toFloat
	    let choose_mat = rand()
	    let center = v3create(fa +. 0.9 *. rand(),
				  0.2,
				  fb +. 0.9 *. rand())
	    if v3vsub(center, v3create(4.0, 0.2, 0.0))->v3len > 0.9 {
		if choose_mat < 0.8 {
		    let a1 = v3create(rand(), rand(), rand())
		    let a2 = v3create(rand(), rand(), rand())
		    let albedo = v3vmul(a1, a2)
		    let m = lambertian(albedo)
		    add(sphere2hittable(center, 0.2, m))
		} else if choose_mat < 0.95 {
		    let albedo = v3create(rand() /. 2.0 +. 0.5,
					  rand() /. 2.0 +. 0.5,
					  rand() /. 2.0 +. 0.5)
		    let fuzz = rand() /. 2.0
		    let m = metal(albedo, fuzz)
		    add(sphere2hittable(center, 0.2, m))
		} else {
		    let m =dielectric(1.5)
		    add(sphere2hittable(center, 0.2, m))
		}
	    }
	}
    }
    let m1 = dielectric(1.5)
    add(sphere2hittable(v3create(0.0, 1.0, 0.0), 1.0, m1))
    let m2 = lambertian(v3create(0.4, 0.2, 0.1))
    add(sphere2hittable(v3create(-4.0, 1.0, 0.0), 1.0, m2))
    let m3 = metal(v3create(0.7, 0.6, 0.5), 0.0)
    add(sphere2hittable(v3create(4.0, 1.0, 0.0), 1.0, m3))
    world
}

let rec ray_color = (r: ray, world: array<hittable>, depth: int, cv) => {
    if (depth <= 0) {
	v3create(0.0, 0.0, 0.0)
    } else {
	let somehit = hitByList(. r, 0.001, inf, world)
	switch somehit {
		| Some(a) => {
		    let b = a.mater(. r, a)
		    switch b {
			    | Some((color, scattered)) => {
				ray_color(scattered, world, depth - 1, v3vmul(color, cv))
			    }
			    | None => {
				v3create(0.0, 0.0, 0.0)
			    }
		    }
		}
		| None => {
		    let un = r->raydirection->v3unit
		    let ti = (v3y(un) +. 1.0) *. 0.5
		    let r1 = v3create(1.0, 1.0, 1.0)->v3mul(1.0 -. ti)
		    let r2 = v3create(0.5, 0.7, 1.0)->v3mul(ti)
		    let r = v3vadd(r1, r2)
		    v3vmul(r, cv)
		}
	}
    }
}

let main = () => {
    let aspect_ratio = 3.0 /. 2.0
    let image_width = 400
    let image_height = (Belt.Int.toFloat(image_width) /. aspect_ratio)
	->Belt.Float.toInt
    let samples_per_pixel = 50
    let fper = samples_per_pixel->Belt.Int.toFloat
    let max_depth = 50

    let aperture = 0.1
    let dist_to_focus = 10.0

    let world = random_scene()
    let came = camcreate(~lookfrom=v3create(13.0, 2.0, 3.0),
			 ~lookat=v3create(0.0, 0.0, 0.0),
			 ~vup=v3create(0.0, 1.0, 0.0),
			 ~aspect_ratio=aspect_ratio,
			 ~viewport_height=2.0,
			 ~vfov=20.0,
			 ~focus_dist=dist_to_focus,
			 ~aperture=aperture)

    Js.Console.timeStart("test")
    Belt.Array.makeBy(image_height, (j) => {
	Belt.Array.makeBy(image_width, (i) => {
	    let j = image_height - 1 - j
	    let color = v3create(0.0, 0.0, 0.0)
	    for _ in 0 to samples_per_pixel - 1 {
		let u = (Belt.Int.toFloat(i) +. rand()) /. Belt.Int.toFloat(image_width - 1)
		let v = (Belt.Int.toFloat(j) +. rand()) /. Belt.Int.toFloat(image_height - 1)
		let r = came->camgetray(. _, u, v)
		v3vaddInPlace(color, ray_color(r, world, max_depth, v3create(1.0, 1.0, 1.0)))
	    }
	    v3create((v3x(color) /. fper)->clamp(0.0, 0.999)->Js.Math.sqrt,
		     (v3y(color) /. fper)->clamp(0.0, 0.999)->Js.Math.sqrt,
		     (v3z(color) /. fper)->clamp(0.0, 0.999)->Js.Math.sqrt)

	})
    })->writemat("fb.ppm")
    Js.Console.timeEnd("test")
}

main()