Jai adaptation of Casey Muratori's Handemade Ray example, day 0 (with modifications)

raytracer.jai

// Casey Muratori's Handmade Ray tutorial, translated to Jai
// https://guide.handmadehero.org/ray/ray00/

main :: () {
  image := allocate_image(1920, 1080);

  materials := Material.[
    .{emit_color = srgb_to_linear(.{0.3, 0.4, 0.5})},
    .{ref_color  = srgb_to_linear(.{0.5, 0.5, 0.5})},
    .{ref_color  = srgb_to_linear(.{0.7, 0.5, 0.3})},
    .{ref_color  = srgb_to_linear(.{0.9, 0.0, 0.0}), emit_color  = srgb_to_linear(.{0.9, 0.0, 0.0})},
    .{ref_color  = srgb_to_linear(.{0.2, 0.8, 0.2}), scatter = 1},
    .{ref_color  = srgb_to_linear(.{0.8, 0.9, 0.8}), scatter = 0.99},
    .{ref_color  = srgb_to_linear(.{0.7, 0.7, 0.9}), scatter = 0.99, solidity = 0.3},
  ];

  planes := Plane.[
    .{n = .{0,  0, 1}, d =  0, mat_index = 1},
    // .{n = normalize(Vector3.{-.4, -1, 0}), d =  8, mat_index = 5},
  ];

  spheres := Sphere.[
    .{p = .{ 0,  0, 0}, r = 1, mat_index = 2},
    .{p = .{ 3, -2, 0}, r = 1, mat_index = 3},
    .{p = .{-2, -1, 2}, r = 1, mat_index = 4},
    .{p = .{1, -5, 1}, r = 0.8, mat_index = 6},
  ];

  world := World.{
    materials = materials,
    planes    = planes,
    spheres   = spheres,
  };

  // Construct camera coordinate system.
  // Note that our world is defined in a right-handed, z-up coordinate system.
  // Our camera will be defined such that backward in its coordinate system's z-direction is what it's pointing at.
  camera_p := Vector3.{0, -10, 1};  // Currently, it's just pointing at the origin from wherever it's located.
  camera_z := normalize(camera_p);
  camera_x := normalize(cross_product(.{0, 0, 1}, camera_z));
  camera_y := normalize(cross_product(camera_z, camera_x));

  film_dist   := 1.0;
  film_center := camera_p - film_dist * camera_z;

  // Correct the film's aspect ratio to our output image's.
  film_w, film_h := 1.0, 1.0;
  if image.width > image.height {
    film_h = 1.0*image.height/image.width;
  }
  else if image.height > image.width {
    film_w = 1.0*image.width/image.height;
  }
  half_film_w, half_film_h := film_w/2, film_h/2;

  rays_per_pixel   := 16;
  contribution     := 1.0 / (rays_per_pixel);
  pixel_w, pixel_h := film_w / image.width, film_h / image.height;

  frame_draw_start := current_time_monotonic();
  for y: 0..image.height-1 {
    film_y := -1.0 + 2.0 * (cast(float) y + 0.5) / (cast(float) image.height);
    for x: 0..image.width-1 {
      // Project each filter onto an imaginary film in front of the camera.
      film_x := -1.0 + 2.0 * (cast(float) x + 0.5) / (cast(float) image.width);

      film_p := film_center + film_x*half_film_w*camera_x + film_y*half_film_h*camera_y;

      ray_origin := camera_p;

      // Allow diffuse materials to work better by casting multiple rays per pixel and averaging them.
      // This simulates an area being partially lit from bounced-off light rays from multiple locations.
      color: Vector3;
      for 0..rays_per_pixel-1 {
        // Randomly jitter the samples around within a pixel's distance - this will give a kind of anti-aliasing effect.
        offset        := Vector3.{random_bilateral() * 0.5 * pixel_w, random_bilateral() * 0.5 * pixel_h, 0};
        ray_direction := normalize(film_p + offset - camera_p);
        color         += contribution * ray_cast(world, ray_origin, ray_direction);
      }

      bmp_value := rgbaf_to_rgba8(Vector4.{xyz = linear_to_srgb(color), w = 1});
      image.pixels[x + y*image.width] = bmp_value;
    }
  }
  frame_draw_time := current_time_monotonic() - frame_draw_start;
  log("Draw took: % ms (% FPS)", to_float64_seconds(frame_draw_time) * 1000, 1.0 / to_float64_seconds(frame_draw_time));

  bmp_header := BMP_Header.{
    file_type           = 0x4D42,
    file_size           = xx (size_of(BMP_Header) + image.byte_size),
    reserved_1          = 0,
    reserved_2          = 0,
    bitmap_offset       = size_of(BMP_Header),
    size                = size_of(BMP_Header) - 14,
    width               = xx image.width,
    height              = xx image.height,
    planes              = 1,
    bits_per_pixel      = xx (size_of(type_of(image.pixels[0])) * 8),
    compression         = 0,
    size_of_bitmap      = xx (image.byte_size),
    horz_resolution     = 0,
    vert_resolution     = 0,
    colors_used         = 0,
    colors_important    = 0,
  };

  {
    file, success := file_open("raytracer.bmp", for_writing = true, keep_existing_content = false);
    defer file_close(*file);

    // @Incomplete: Need to pad each row to 4 bytes to be spec-correct.
    // row_size := (bmp_header.bits_per_pixel * image.width + 31) / 32 * 4;
    // for y: 0..image.height-1 {

    // }

    file_write(*file, *bmp_header, size_of(type_of(bmp_header)));
    file_write(*file, image.pixels.data, image.byte_size);
  }
}

allocate_image :: (width: u32, height: u32) -> Image {
  result: Image;
  result.width     = width;
  result.height    = height;
  result.pixels    = NewArray(width * height, u32);
  result.byte_size = result.pixels.count * size_of(type_of(result.pixels[0]));
  return result;
}

ray_cast :: (world: World, origin: Vector3, direction: Vector3) -> Vector3 {
  TOLERANCE        :: 0.0001;
  MIN_HIT_DISTANCE :: 0.001;

  result: Vector3;
  attenuation := Vector3.{1, 1, 1};

  ray_origin    := origin;
  ray_direction := direction;

  for bounce_count: 0..8-1 {
    hit_distance      := FLOAT32_MAX;
    hit_mat_index: u32 = 0;
    next_origin:   Vector3;
    next_normal:   Vector3;

    for plane: world.planes {
      // Equation of a plane is:
      // n.p + d = 0  (n is the normal vector, d is distance from origin)
      //
      // Ray equation is:
      // p = o + t*r   (o and r are vectors)
      //
      // Plugging those together and solving for t gets you:
      // t = (-d - n.o) / (n.r)

      denom := dot_product(plane.n, ray_direction);

      if abs(denom) > TOLERANCE {
        this_distance := (-plane.d - dot_product(plane.n, ray_origin)) / denom;
        if this_distance > MIN_HIT_DISTANCE && this_distance < hit_distance {
          hit_distance  = this_distance;
          hit_mat_index = plane.mat_index;

          next_origin = this_distance*ray_direction + ray_origin;
          next_normal = plane.n;
        }
      }
    }

    for sphere: world.spheres {
      // Equation of a sphere at origin is:
      // p_x^2 * p_y^2 * p_z^2 - r^2 = 0  (p is a vector)
      //
      // Ray equation is:
      // p = o + t*d   (o and d are vectors)
      //
      // Plugging those together and simplifying using the definition of a dot product gets you:
      // (d.d) * t^2 + (2*d.o) * t + (o.o - r*r) = 0
      //
      // You can then solve for t using the quadratic equation to get two potential values for t,
      // then sneakily account for the sphere's center by substituting o with (o - center).

      // Use quadratic eq
      sphere_relative_origin := ray_origin - sphere.p;
      a := dot_product(ray_direction, ray_direction);
      b := 2 * dot_product(ray_direction, sphere_relative_origin);
      c := dot_product(sphere_relative_origin, sphere_relative_origin) - sphere.r*sphere.r;

      inside_root := b*b - 4*a*c;
      denom       := 2*a;

      if inside_root > 0 && abs(denom) > TOLERANCE {
        rooted       := sqrt(inside_root);
        sol_1, sol_2 := (-b + rooted) / denom, (-b - rooted) / denom;

        first_hit := sol_1;
        if sol_2 > MIN_HIT_DISTANCE && sol_2 < first_hit  first_hit = sol_2;

        if first_hit > MIN_HIT_DISTANCE && first_hit < hit_distance {
          hit_distance  = first_hit;
          hit_mat_index = sphere.mat_index;

          next_origin   = first_hit*ray_direction + ray_origin;
          next_normal   = normalize(next_origin - sphere.p);
        }
      }
    }

    mat    := world.materials[hit_mat_index];
    result += attenuation * mat.emit_color;
    if hit_mat_index {
      pass_through_factor := cast(float) (random_get_zero_to_one() < mat.solidity);
      if pass_through_factor  attenuation = attenuation * mat.ref_color;

      ray_origin = next_origin;

      pure_bounce   := ray_direction - 2 * dot_product(next_normal, ray_direction) * next_normal;
      random_bounce := normalize(next_normal + .{random_bilateral(), random_bilateral(), random_bilateral()});
      bounced       := lerp(random_bounce, pure_bounce, mat.scatter);
      ray_direction  = lerp(ray_direction, bounced,     pass_through_factor);
    }
    else {
      break;
    }
  }


  return result;
}

linear_to_srgb :: (color: Vector3) -> Vector3 {
  result := color;
  for 0..3-1 {
    comp := 1.055 * pow(result.component[it], 1.0/2.4) - 0.055;
    if result.component[it] <= 0.0031308 {
      comp = result.component[it] * 12.92;
    }
    result.component[it] = comp;
  }
  return result;
}

srgb_to_linear :: (color: Vector3) -> Vector3 {
  result := color;
  for 0..3-1 {
    comp := pow((result.component[it]+0.055)/1.055, 2.4);
    if result.component[it] <= 0.04045 {
      comp = result.component[it] / 12.92;
    }
    result.component[it] = comp;
  }
  return result;
}

rgbaf_to_rgba8 :: (color: Vector4) -> u32 {
  r: u32 = xx (255.0 * clamp(color.x, 0, 1) + 0.5);
  g: u32 = xx (255.0 * clamp(color.y, 0, 1) + 0.5);
  b: u32 = xx (255.0 * clamp(color.z, 0, 1) + 0.5);
  a: u32 = xx (255.0 * clamp(color.w, 0, 1) + 0.5);

  return b | (g << 8) | (r << 16) | (a << 24);
}

random_bilateral :: () -> float {
  // return 2.0 * (cast(float) random_get_zero_to_one_new()) - 1.0;
  return 2.0 * random_get_zero_to_one() - 1.0;
}

BMP_Header :: struct {
  file_type:        u16;
  file_size:        u32 #align 2;
  reserved_1:       u16;
  reserved_2:       u16;
  bitmap_offset:    u32 #align 2;
  size:             u32 #align 2;
  width:            s32 #align 2;
  height:           s32 #align 2;
  planes:           u16;
  bits_per_pixel:   u16;
  compression:      u32 #align 2;
  size_of_bitmap:   u32 #align 2;
  horz_resolution:  s32 #align 2;
  vert_resolution:  s32 #align 2;
  colors_used:      u32 #align 2;
  colors_important: u32 #align 2;
} #no_padding

Image :: struct {
  width:     u32;
  height:    u32;
  byte_size: int;
  pixels:    [] u32;
}

Material :: struct {
  scatter:    float;          // 0 means complete scatter, 1 means perfectly reflective
  solidity:   float    =  1;  // 0 means complete scatter, 1 means perfectly reflective
  emit_color: Vector3;
  ref_color:  Vector3;
}

Plane :: struct {
  n:         Vector3;   // The vector normal to the plane.
  d:         float;     // The distance plane lines from the origin, in the opposite direction of n.
  mat_index: u32;
}

Sphere :: struct {
  p:         Vector3;   // The sphere's center.
  r:         float;     // The sphere's radius.
  mat_index: u32;
}

World :: struct {
  materials: [] Material;
  planes:    [] Plane;
  spheres:   [] Sphere;
}

#import "Basic";
#import "File";
#import "Math";
#import "Random";