raytracer.mank

# Based on: https://www.scratchapixel.com/code.php?id=3&origin=/lessons/3d-basic-rendering/introduction-to-ray-tracing
# Made nice and manky!

const MAX_RAY_DEPTH := 10;
const INF := 10'000'000'000.0;
const M_PI := 3.141592653589793;

fun trace: Vec3f (
  ray_origin: Vec3f,
  ray_dir: Vec3f,
  spheres: Sphere[],
  depth: i32
) {
  t_near := INF;
  hit_sphere_idx := -1;
  for s_idx in 0 .. spheres.length {
    bind (hit, t0, t1) = sphere_intersect(spheres[s_idx], ray_origin, ray_dir);
    if hit {
      if t0 < 0. { t0 = t1; }
      if t0 < t_near {
        t_near = t0;
        hit_sphere_idx = s_idx;
      }
    }
  }

  if hit_sphere_idx < 0 {
    return vec3f_from_f64(2.); # sky needs to be a very bright white
  }
  hit_colour := new_vec3f();

  hit_sphere := ref spheres[hit_sphere_idx];
  hit_point := vec3_add(ray_origin, vec3_scalar_mult(ray_dir, t_near)); # walk down ray
  hit_normal := vec3_sub(hit_point, hit_sphere.centre);
  vec3f_normalize(hit_normal);

  if hit_sphere_idx == 2 {
    # Hack to add a texture!
    hit_sphere.colour = get_earth_texel(hit_normal);
    hit_sphere.emission_colour = hit_sphere.colour;
  }

  # used to help with normals (just move them a tiny bit so their not on the surface)
  bias := 1./10000.;

  inside_sphere := false;
  if vec3_dot(ray_dir, hit_normal) > 0. {
    # view dir and normal not in opposite directions (inside)
    hit_normal = vec3_minus(hit_normal);
    inside_sphere = true;
  }

  if (hit_sphere.transparency > 0. || hit_sphere.reflection > 0.) && depth < MAX_RAY_DEPTH {
    facing_ratio := -vec3_dot(ray_dir, hit_normal);
    fresnele := mix(pow(1. - facing_ratio, 3.), 1., 0.1);

    reflect_dir := vec3_sub(ray_dir, vec3_scalar_mult(hit_normal, 2. * vec3_dot(ray_dir, hit_normal)));
    vec3f_normalize(reflect_dir);
    reflection := trace(vec3_add(hit_point, vec3_scalar_mult(hit_normal, bias)), reflect_dir, spheres, depth + 1);

    # Now refraction/transmission
    refraction := new_vec3f();
    if hit_sphere.transparency > 0. {
      # Spooky maths...
      ior := 1.1;
      eta := if inside_sphere { ior } else { 1. / ior };
      cosi := -vec3_dot(hit_normal, ray_dir);
      k := 1. - eta * eta * (1. - cosi * cosi);
      refract_dir := vec3_add(
        vec3_scalar_mult(ray_dir, eta),
        vec3_scalar_mult(hit_normal, eta * cosi - sqrt(k)));
      vec3f_normalize(reflect_dir);
      refraction = trace(vec3_sub(hit_point, vec3_scalar_mult(hit_normal, bias)), refract_dir, spheres, depth + 1);
    }

    hit_colour = vec3_mult(
      vec3_add(
        vec3_scalar_mult(reflection, fresnele),
        vec3_scalar_mult(refraction, (1. - fresnele) * hit_sphere.transparency)),
      hit_sphere.colour);
  } else {
    # Boring faked diffuse light
    path_blocked := \light_idx, light_dir: Vec3f -> {
      for s_idx in 0 .. spheres.length {
        if s_idx == light_idx { continue; }
        bind (hit, _, __) = sphere_intersect(spheres[s_idx],
          vec3_add(hit_point, vec3_scalar_mult(hit_normal, bias)), light_dir);
        if hit { return true; }
      }
      false
    }

    max := \a,b -> { if a > b { a } else { b } }

    for s_idx in 0 .. spheres.length {
      emission := ref spheres[s_idx].emission_colour;
      if vec3_length(emission) > 0. {
        # Light source
        light_dir := vec3_sub(spheres[s_idx].centre, hit_point);
        vec3f_normalize(light_dir);

        if path_blocked(s_idx, light_dir) {
          continue;
        }

        # attenuation & emission
        hit_colour = vec3_add(hit_colour, vec3_mult(hit_sphere.colour,
          vec3_scalar_mult(emission, max(0., vec3_dot(hit_normal, light_dir)))));
      }
    }
  }

  vec3_add(hit_colour, hit_sphere.emission_colour)
}

proc ray_trace(image_width: i32, image_height: i32, spheres: Sphere[]) {
  # image pixel array not needed... if we just print pixels on the fly
  # (vectors still very needed!)
  # TODO: Vectors! Can only handle tiny images with stack memory :(
  # image := [=[=new_vec3f();60];60];

  bind (fov, aspect) = (30., image_width as f64 / image_height as f64);

  inverse_width := 1. / image_width as f64;
  inverse_height := 1. / image_height as f64;
  angle := tan(M_PI * 0.5 * fov / 180.);

  float_to_char := \f -> {
    (if f >= 1. {
      1.
    } else {
      f
    } * 255.) as char
  };

  output_rgb := \rgb: Vec3f -> {
    bind {.x/r, .y/g, .z/b} = rgb;
    putchar(float_to_char(r));
    putchar(float_to_char(g));
    putchar(float_to_char(b));
  }

  # Make .ppm image
  println!("P6\n{width} {height}\n255",
    int_to_string(image_width), int_to_string(image_height));
  for y in 0 .. image_height {
    for x in 0 .. image_width {
      xx := (2. * ((x as f64 + 0.5) * inverse_width) - 1.) * angle * aspect;
      yy := (1. - 2. * ((y as f64 + 0.5) * inverse_height)) * angle;
      ray_dir := vec3f_from_xyz(xx, yy, -1.);
      vec3f_normalize(ray_dir);
      output_rgb(trace(new_vec3f(), ray_dir, spheres, 0));
    }
  }
}

proc main {
  spheres := vec!([
    # Middle sphere
    make_sphere(
      vec3f_from_xyz(0., 0., -20.), 4., vec3f_from_xyz(1., 0.32, 0.36), 1., 0.5, new_vec3f()),
    # Right sphere
    make_sphere(
      vec3f_from_xyz(5., -1., -15.), 2., vec3f_from_xyz(0.90, 0.76, 0.46), 1., 0., new_vec3f()),
    # Earth
    make_sphere(
      vec3f_from_xyz(6., 3., -20.), 2.4, new_vec3f(), 0., 0., new_vec3f()),
    # Left sphere
    make_sphere(
      vec3f_from_xyz(-5.5, 0., -15.), 3., vec3f_from_xyz(0.67, 0.84, 0.90), 1., 0., new_vec3f()),
    # Ground
    make_sphere(
      vec3f_from_xyz(0., -10004., -20.), 10000., vec3f_from_xyz(0.5,0.25,0.155), 0., 0., new_vec3f()),
    # Light (behind view plane)
    make_sphere(
      vec3f_from_xyz(0., 20., -30.), 3., new_vec3f(), 0.0, 0.0, vec3f_from_f64(3.))
  ]);

  bind (WIDTH, HEIGHT) = (1920, 1080); # fine for my PC -- adjust as needed
  eprintln("Tracing...");
  ray_trace(WIDTH, HEIGHT, spheres);
  eprintln("Done!");
}

# Alpha mixing

fun mix: f64 (a: f64, b: f64, mix: f64) {
  b * mix + a * (1. - mix)
}

# Sphere

pod Sphere {
  # World
  centre: Vec3f,
  radius: f64,
  radius_squared: f64,
  # Props
  transparency: f64,
  reflection: f64,
  # Colours
  colour: Vec3f,
  emission_colour: Vec3f
}

fun make_sphere: Sphere (
  centre: Vec3f, radius: f64, colour: Vec3f,
  reflection: f64, transparency: f64, emission_colour: Vec3f
) {
  Sphere {
    .centre = centre,
    .radius = radius,
    .radius_squared = radius * radius,
    .colour = colour,
    .emission_colour = emission_colour,
    .transparency = transparency,
    .reflection = reflection
  }
}

fun sphere_intersect: (bool, f64, f64) (
  sphere: ref Sphere,
  ray_origin: Vec3f,
  ray_dir: Vec3f
) {
  # Some maths (copied -- including meaningless var names :()
  MISS := (false, -1., -1.);
  l := vec3_sub(sphere.centre, ray_origin);
  tca := vec3_dot(l, ray_dir);
  if (tca < 0.) { return MISS; }
  d2 := vec3_dot(l,l) - tca * tca;
  if (d2 > sphere.radius_squared) { return MISS; }
  thc := sqrt(sphere.radius_squared - d2);
  (true, tca - thc, tca + thc)
}

# Vector

pod Vec3f {
  x: f64, y: f64, z: f64
}

fun new_vec3f: Vec3f {
  Vec3f { .x = 0., .y = 0., .z = 0. }
}

fun vec3f_from_f64: Vec3f (f: f64) {
  Vec3f { .x = f, .y = f, .z = f}
}

fun vec3f_from_xyz: Vec3f (x: f64, y: f64, z: f64) {
  Vec3f { .x = x, .y = y, .z = z}
}

fun vec3_length_squared: f64 (vec: ref Vec3f) {
  bind {.x, .y, .z} = vec;
  x * x + y * y + z * z
}

fun vec3_length: f64 (vec: ref Vec3f) {
  sqrt(vec3_length_squared(vec))
}

proc vec3f_normalize(vec: ref Vec3f) {
  l2 := vec3_length_squared(vec);
  if l2 > 0. {
    scale := 1. / sqrt(l2);
    vec.x *= scale;
    vec.y *= scale;
    vec.z *= scale;
  }
}

fun vec3_scalar_mult: Vec3f (vec: Vec3f, scalar: f64) {
  vec.x *= scalar;
  vec.y *= scalar;
  vec.z *= scalar;
  vec
}

fun vec3_mult: Vec3f (a: Vec3f, b: Vec3f) {
  vec3f_from_xyz(a.x * b.x, a.y * b.y, a.z * b.z)
}

fun vec3_dot: f64 (a: Vec3f, b: Vec3f) {
  a.x * b.x + a.y * b.y + a.z * b.z
}

fun vec3_sub: Vec3f (a: Vec3f, b: Vec3f) {
  vec3f_from_xyz(a.x - b.x, a.y - b.y, a.z - b.z)
}

fun vec3_add: Vec3f (a: Vec3f, b: Vec3f) {
  vec3f_from_xyz(a.x + b.x, a.y + b.y, a.z + b.z)
}

fun vec3_minus: Vec3f (vec: Vec3f) {
  vec3f_from_xyz(-vec.x, -vec.y, -vec.z)
}

# Mank logo

fun sample_mank_logo: i32 (x: i32, y: i32) {
  # TODO: make better
  logo := [
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xfb, 0xbb,
    0x9a, 0xdf, 0xf9, 0x35, 0xaa, 0xbf, 0xfa, 0xae, 0xb2,
    0x7f, 0xfb, 0xa0, 0xba, 0x7f, 0xfb, 0xae, 0xba, 0xbf,
    0xfb, 0xae, 0xba, 0xdf, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
    0xff, 0xff ];

  px := y * 32 + x;
  cell := px / 8;
  (logo[cell] >> (7 - px % 8)) & 1
}

fun get_uv: (f64, f64) (normal: Vec3f) {
  PI := 3.14159;
  (
    0.5 + atan2(normal.z, normal.x) / (2. * PI),
    0.5 - asin(normal.y) / PI
  )
}

fun get_mank_texel: Vec3f (normal: Vec3f) {
  bind (u, v) = get_uv(normal);
  s := if sample_mank_logo((u * 32.) as i32, (v * 32.) as i32) == 1 { 1. } else { 0.4 }
  vec3f_from_f64(s)
}

# Earth

fun sample_earth_texture: Vec3f (x: i32, y: i32) {
  earth := [
    0x0b2551, 0x082250, 0x051d48, 0x051d49, 0x0e2958, 0x23406f, 0x224377,
    0x3d5982, 0x6d7d99, 0x5f7598, 0x798dac, 0x8c99b1, 0xa0a9ba, 0x9fa7b5,
    0x33486c, 0x071e44, 0x0c2149, 0x233960, 0x193360, 0x0d2c5c, 0x123261,
    0x1b3867, 0x133364, 0x0c2d5f, 0x20375f, 0x2c3d5f, 0x10254d, 0x061c44,
    0x0d2753, 0x13386e, 0x0d2f62, 0x09285a, 0x1c3f6d, 0x1f3041, 0x1d242d,
    0x19273a, 0x283b51, 0x64707b, 0x68798c, 0x808fa1, 0x798896, 0x4f637d,
    0x445e87, 0xd8dee1, 0xffffff, 0xadb8c3, 0x27446c, 0x0a2553, 0x0d2653,
    0x254065, 0x213b51, 0x193c67, 0x1f4068, 0x2f475f, 0x2a3a48, 0x283334,
    0x3f4538, 0x4e4f3a, 0x363e34, 0x283337, 0x313c3f, 0x323b2f, 0x243844,
    0x273f57, 0x1e4271, 0x253434, 0x293234, 0x343c3c, 0x2b3017, 0x1a2108,
    0x27280f, 0x313a3d, 0x23436d, 0x3b4342, 0x10294a, 0x3a4b6c, 0x3f5175,
    0x001e50, 0x173562, 0x183b69, 0x203c61, 0x1f3030, 0x172615, 0x1a2507,
    0x1c2709, 0x192203, 0x232808, 0x1f2404, 0x172306, 0x152103, 0x1d2906,
    0x222c08, 0x2a3215, 0x324044, 0x293934, 0x242c2c, 0x051129, 0x040e25,
    0x000922, 0x051942, 0x1b2b3a, 0x363719, 0x3d4021, 0x18250b, 0x162212,
    0x19270f, 0x1e3544, 0x082855, 0x000b2a, 0x0c1d41, 0x06183b, 0x1e364c,
    0x2a3931, 0x2c3621, 0x2e3a1d, 0x27331d, 0x41432c, 0x605a3e, 0x5d5335,
    0x464528, 0x4b482d, 0x4a472b, 0x4b482d, 0x2a330f, 0x192b1e, 0x12376a,
    0x091b37, 0x000920, 0x00010e, 0x01020e, 0x030513, 0x01061b, 0x101a32,
    0x6d5e43, 0x6a5b3b, 0x333b16, 0x22300e, 0x122535, 0x04122d, 0x060e22,
    0x091634, 0x0d2249, 0x00031c, 0x373739, 0x555350, 0x26384e, 0x263646,
    0x5f5a4c, 0x636051, 0x958364, 0x7e6e4b, 0x85785e, 0x8f7d5e, 0x766747,
    0x53502d, 0x283d3f, 0x152f4a, 0x070f1f, 0x000111, 0x030617, 0x030a1e,
    0x03091b, 0x020515, 0x030818, 0x000a23, 0x283148, 0x5b573b, 0x27363b,
    0x183350, 0x050f26, 0x000006, 0x040a1d, 0x071330, 0x020516, 0x272d43,
    0xab8f75, 0xc59e73, 0xac9074, 0xa08d79, 0x98836e, 0xac9472, 0x79726a,
    0x767159, 0x636039, 0x444827, 0x1f2b0b, 0x27381c, 0x183254, 0x071634,
    0x03091c, 0x030617, 0x020516, 0x03091c, 0x03081a, 0x020413, 0x010414,
    0x030a1e, 0x011336, 0x0f2139, 0x182a31, 0x132c48, 0x122440, 0x08152e,
    0x030b25, 0x060a1a, 0x000215, 0x42474d, 0xa18f59, 0x9d885b, 0xac9467,
    0xb19360, 0x6a6753, 0x8b7e71, 0x222d4a, 0x1a283f, 0x384039, 0x152d48,
    0x2a3d29, 0x163054, 0x061327, 0x010516, 0x051129, 0x020413, 0x03081a,
    0x020414, 0x030a1e, 0x020718, 0x03091e, 0x030c22, 0x06122d, 0x07193e,
    0x091d44, 0x132948, 0x212e19, 0x1d2d20, 0x0e1f32, 0x021031, 0x05102b,
    0x021024, 0x0e1e16, 0x1e2e27, 0x223312, 0x1d300d, 0x524d29, 0x41383b,
    0x000421, 0x0b1c3d, 0x0d1c3b, 0x091e43, 0x214160, 0x18324e, 0x0e1e30,
    0x061027, 0x081634, 0x06122f, 0x040d24, 0x03091c, 0x03091c, 0x02081b,
    0x040e27, 0x030c23, 0x061331, 0x091a3b, 0x05183c, 0x162237, 0x1c2a0a,
    0x142403, 0x2f3619, 0x2f3032, 0x02091d, 0x06102c, 0x030d2a, 0x020d28,
    0x212c21, 0x323214, 0x473d34, 0x010f2c, 0x091838, 0x071736, 0x020617,
    0x030819, 0x102031, 0x193145, 0x09254a, 0x11304e, 0x0d233b, 0x0c244e,
    0x071736, 0x040d23, 0x03081a, 0x05112b, 0x071534, 0x071533, 0x091c42,
    0x071534, 0x050f2a, 0x040d27, 0x32342e, 0x39391d, 0x544523, 0x2d2f2e,
    0x000112, 0x07132f, 0x061128, 0x01071b, 0x534942, 0x634d2a, 0x343639,
    0x192235, 0x08193b, 0x071533, 0x030617, 0x04091b, 0x000111, 0x0e1a3a,
    0x47444c, 0x5b5044, 0x303d52, 0x08224b, 0x0b1f41, 0x020413, 0x020414,
    0x030a1f, 0x05112c, 0x071738, 0x0a1e45, 0x081a3e, 0x0a1939, 0x000e2f,
    0x322f39, 0x544a24, 0x2b3535, 0x071c42, 0x030a1f, 0x06122e, 0x07112a,
    0x000823, 0x4e494d, 0x715a3e, 0x11213d, 0x10182e, 0x02081c, 0x071432,
    0x061331, 0x07132c, 0x00061e, 0x302f35, 0x6d4a2d, 0x875d3a, 0x6b5332,
    0x0e264c, 0x0c2755, 0x040b1f, 0x020413, 0x010310, 0x020414, 0x030b21,
    0x081a3e, 0x081a3e, 0x091838, 0x011439, 0x343533, 0x3e4440, 0x041129,
    0x000213, 0x050f29, 0x091c42, 0x071635, 0x010415, 0x060e22, 0x0d172f,
    0x04112e, 0x081d44, 0x040a1d, 0x040e27, 0x0a1e45, 0x0a1c40, 0x040f28,
    0x071021, 0x081122, 0x141d2d, 0x232b2e, 0x040c22, 0x142f5a, 0x030717,
    0x020718, 0x040e24, 0x050d23, 0x07132f, 0x0b2048, 0x091736, 0x050d22,
    0x061330, 0x3c4754, 0x163768, 0x051433, 0x07132c, 0x040e25, 0x091838,
    0x0d2149, 0x0b1b3b, 0x07112b, 0x03091d, 0x091735, 0x07193c, 0x020c23,
    0x112b58, 0x091d43, 0x041334, 0x031539, 0x03102f, 0x000921, 0x000d2b,
    0x03163a, 0x040c21, 0x152d51, 0x05102b, 0x07183a, 0x081a40, 0x02102e,
    0x00061e, 0x000318, 0x000113, 0x00010d, 0x000113, 0x04112a, 0x1c273e,
    0x071026, 0x07132e, 0x030d23, 0x000111, 0x000219, 0x000214, 0x00010f,
    0x00010c, 0x00010b, 0x060a1a, 0x070b1b, 0x000b24, 0x06132c, 0x0f162a,
    0x192644, 0x0f1d3a, 0x0e1a35, 0x12203e, 0x091c41, 0x00163d, 0x021232,
    0x091d43, 0x061535, 0x0a1634, 0x202b47, 0x384460, 0x334261, 0x3c4a66,
    0x59647d, 0x4f6180, 0x99a0ac, 0x9097a2, 0x041737, 0x091835, 0x0c1730,
    0x494d59, 0x808086, 0x8b8d93, 0x8e8f96, 0x8f9198, 0xa2a4a8, 0xcaced1,
    0xd0d3d6, 0xaaacaf, 0xbcbec0, 0xdde0e4, 0xe2e7ef, 0xd9dee8, 0xd8dee8,
    0xdde2ea, 0xc6cdd9, 0x9fabc0, 0x43597f ];
  px := y * 32 + x;
  s := earth[px];
  b := s & 255;
  g := (s >> 8) & 255;
  r := (s >> 16) & 255;

  vec3f_from_xyz(r as f64 / 255.,  g as f64 / 255., b as f64 / 255.)
}

fun get_earth_texel: Vec3f (normal: Vec3f) {
  bind (u, v) = get_uv(normal);
  sample_earth_texture(31 - (u * 32.) as i32, (v * 15.) as i32)
}