raytracing weekend thing implementation in plain C

raytracing-weekend.c

// compile with (linux) gcc -std=c11 -Ofast -lcglm -lm -fopenmp raytracing-weekend.c
// dependencies: cglm (available on AUR), openmp (probably already installed on ur system)

#include <cglm/vec3.h>
#include <inttypes.h>
#include <math.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb/stb_image_write.h>

typedef uint8_t u8;
typedef int32_t i32;
typedef uint32_t u32;
typedef size_t usize;
typedef ptrdiff_t size;

typedef float f32;
typedef double f64;

inline void dump_image_header(FILE *output, i32 width, i32 height) {
  fprintf(output, "P3\n");
  fprintf(output, "%" PRIi32 " %" PRIi32 "\n", width, height);
  fprintf(output, "255\n");
}

inline void dump_color(FILE *output, vec3 color) {
  i32 r = (i32)round(color[0] * 255);
  i32 g = (i32)round(color[1] * 255);
  i32 b = (i32)round(color[2] * 255);
  fprintf(output, "%" PRIi32 " %" PRIi32 " %" PRIi32 "\n", r, g, b);
}

typedef struct {
  vec3 origin;
  vec3 direction;
} ray;

void point_on_ray(ray *r, f32 t, vec3 result) {
  glm_vec3_copy(r->origin, result);
  glm_vec3_muladds(r->direction, t, result);
}

typedef struct {
  f32 min;
  f32 max;
} range;

bool range_test(const range *r, f32 value) {
  return value >= r->min && value <= r->max;
}

f32 clamp(const range *r, f32 value) {
  if (value <= r->min)
    value = r->min;
  if (value >= r->max)
    value = r->max;
  return value;
}

f32 random_float() { return (f32)((f64)rand() / (RAND_MAX + 1.0)); }

void random_vec3(vec3 v) {
  v[0] = random_float();
  v[1] = random_float();
  v[2] = random_float();
}

void random_vec3_in_unit_sphere(vec3 v) {
  do {
    random_vec3(v);
  } while (glm_vec3_norm2(v) >= 1);
}

void sphere_normal(vec3 center, f32 radius, vec3 hit_point, vec3 normal) {
  glm_vec3_sub(hit_point, center, normal);
  glm_vec3_divs(normal, radius, normal);
}

typedef struct hit_record hit_record;

typedef enum {
  material_tag_lambertian,
  material_tag_metal,
  material_tag_dielectric
} material_tag;

typedef struct {
  vec3 albedo;
} lambertian_data;

typedef struct {
  vec3 albedo;
  f32 fuzz;
} metal_data;

typedef struct {
  f32 index_of_refraction;
} dielectric_data;

typedef union {
  lambertian_data lambertian;
  metal_data metal;
  dielectric_data dielectric;
} material_data;

typedef struct material {
  material_tag tag;
  material_data data;
  bool (*scatter)(struct material *mat, ray *r, hit_record *hit,
                  vec3 attenuation, ray *scattered);
} material;

struct hit_record {
  vec3 hit_point;
  vec3 normal;
  f32 t;
  bool front;
  material *mat;
};

bool lambertian_scatter(material *mat, ray *r, hit_record *hit,
                        vec3 attenuation, ray *scattered) {
  vec3 direction;
  random_vec3_in_unit_sphere(direction);
  glm_vec3_normalize(direction);
  glm_vec3_add(direction, hit->normal, direction);
  if (glm_vec3_norm2(direction) < 1e-4) {
    glm_vec3_copy(hit->normal, direction);
  }

  glm_vec3_copy(hit->hit_point, scattered->origin);
  glm_vec3_copy(direction, scattered->direction);
  glm_vec3_copy(mat->data.lambertian.albedo, attenuation);
  return true;
}

material make_lambertian(vec3 albedo) {
  lambertian_data data;
  glm_vec3_copy(albedo, data.albedo);
  return (material){
      .tag = material_tag_lambertian,
      .data = {.lambertian = data},
      .scatter = lambertian_scatter,
  };
}

void reflect(vec3 dir, vec3 normal, vec3 out) {
  glm_vec3_normalize_to(dir, out);
  glm_vec3_muladds(normal, -2 * glm_vec3_dot(out, normal), out);
}

void refract(vec3 dir, vec3 normal, float refraction_ratio, vec3 out) {
  vec3 unit_dir;
  glm_vec3_normalize_to(dir, unit_dir);
  f32 cos_theta = fmin(1.0, -glm_vec3_dot(unit_dir, normal));

  vec3 perp, parallel;
  glm_vec3_copy(unit_dir, perp);
  glm_vec3_muladds(normal, cos_theta, perp);
  glm_vec3_scale(perp, refraction_ratio, perp);

  glm_vec3_scale(normal, -sqrt(fabs(1.0 - glm_vec3_norm2(perp))), parallel);

  glm_vec3_add(perp, parallel, out);
}

bool metal_scatter(material *mat, ray *r, hit_record *hit, vec3 attenuation,
                   ray *scattered) {
  vec3 direction;
  reflect(r->direction, hit->normal, direction);

  vec3 fuzz_vector;
  random_vec3_in_unit_sphere(fuzz_vector);
  glm_vec3_normalize(fuzz_vector);
  glm_vec3_muladds(fuzz_vector, mat->data.metal.fuzz, direction);

  glm_vec3_copy(hit->hit_point, scattered->origin);
  glm_vec3_copy(direction, scattered->direction);
  glm_vec3_copy(mat->data.metal.albedo, attenuation);
  return true;
}

material make_metal(vec3 albedo, f32 fuzz) {
  metal_data data;
  glm_vec3_copy(albedo, data.albedo);
  data.fuzz = fuzz;
  return (material){
      .tag = material_tag_lambertian,
      .data = {.metal = data},
      .scatter = metal_scatter,
  };
}

f32 reflectance(f32 cos, f32 ref_index) {
  f32 r0 = (1 - ref_index) / (1 + ref_index);
  r0 = r0 * r0;
  return r0 + (1 - r0) * pow(1 - cos, 5.0);
}

bool dielectric_scatter(material *mat, ray *r, hit_record *hit,
                        vec3 attenuation, ray *scattered) {
  f32 refraction_ratio = mat->data.dielectric.index_of_refraction;
  if (hit->front) {
    refraction_ratio = 1.0 / refraction_ratio;
  }

  glm_vec3_copy(hit->hit_point, scattered->origin);

  vec3 unit_dir;
  glm_vec3_normalize_to(r->direction, unit_dir);

  f32 cos_theta = fmin(1.0, -glm_vec3_dot(unit_dir, hit->normal));
  f32 sin_theta = sqrt(1.0 - cos_theta * cos_theta);

  bool cannot_refract = refraction_ratio * sin_theta > 1.0;

  if (cannot_refract ||
      reflectance(cos_theta, refraction_ratio) > random_float()) {
    reflect(r->direction, hit->normal, scattered->direction);
  } else {
    refract(r->direction, hit->normal, refraction_ratio, scattered->direction);
  }

  attenuation[0] = 1.0;
  attenuation[1] = 1.0;
  attenuation[2] = 1.0;

  return true;
}

material make_dielectric(f32 index_of_refraction) {
  return (material){
      .tag = material_tag_dielectric,
      .data = {.dielectric = {.index_of_refraction = index_of_refraction}},
      .scatter = dielectric_scatter,
  };
}

typedef enum {
  hittable_tag_sphere,
} hittable_tag;

typedef struct {
  vec3 center;
  f32 radius;
} sphere_data;

typedef union {
  sphere_data sphere;
} hittable_data;

typedef struct hittable {
  hittable_tag tag;
  hittable_data data;
  material *mat;

  bool (*hit)(struct hittable *hit, ray *r, const range *t_range,
              hit_record *rec);
} hittable;

bool sphere_hit(hittable *hit, ray *r, const range *t_range, hit_record *rec) {
  vec3 to_origin_vec;
  glm_vec3_sub(r->origin, hit->data.sphere.center, to_origin_vec);
  f32 radius = hit->data.sphere.radius;

  f32 a, half_b, c;
  a = glm_vec3_norm2(r->direction);
  half_b = glm_vec3_dot(r->direction, to_origin_vec);
  c = glm_vec3_norm2(to_origin_vec) - radius * radius;

  f32 det_prime = half_b * half_b - a * c;
  if (det_prime < 0) {
    return false;
  }

  f32 sqrt_det = sqrt(det_prime);
  f32 t = (-half_b - sqrt_det) / a;
  if (!range_test(t_range, t)) {
    t = (-half_b + sqrt_det) / a;
    if (!range_test(t_range, t)) {
      return false;
    }
  }

  if (rec) {
    rec->t = t;
    rec->mat = hit->mat;
    point_on_ray(r, t, rec->hit_point);
    glm_vec3_sub(rec->hit_point, hit->data.sphere.center, rec->normal);
    glm_vec3_divs(rec->normal, radius, rec->normal);
    rec->front = glm_vec3_dot(r->direction, rec->normal) < 0;
    if (!rec->front) {
      glm_vec3_negate(rec->normal);
    }
  }

  return true;
}

hittable make_sphere(material *mat, vec3 center, f32 radius) {
  sphere_data data;
  glm_vec3_copy(center, data.center);
  data.radius = radius;

  return (hittable){
      .tag = hittable_tag_sphere,
      .data =
          {
              .sphere = data,
          },
      .hit = sphere_hit,
      .mat = mat,
  };
}

void get_ray_color(ray *r, hittable *objects, i32 num_objects, vec3 color,
                   i32 num_rays) {
  if (num_rays > 50) {
    color[0] = 0;
    color[1] = 0;
    color[2] = 0;
    return;
  }

  bool hit = false;
  hit_record nearest_hit;
  nearest_hit.t = 1000.0;
  for (i32 i = 0; i < num_objects; ++i) {
    hittable *object = &objects[i];
    hit_record current;
    if (object->hit(object, r, &(range){.min = 0.001, .max = nearest_hit.t},
                    &current)) {
      hit = true;
      nearest_hit = current;
    }
  }

  if (hit) {
    ray scattered_ray;
    vec3 attenuation;
    if (nearest_hit.mat->scatter(nearest_hit.mat, r, &nearest_hit, attenuation,
                                 &scattered_ray)) {
      get_ray_color(&scattered_ray, objects, num_objects, color, num_rays + 1);
      glm_vec3_mul(color, attenuation, color);
    } else {
      glm_vec3_zero(color);
    }

    return;
  }

  vec3 blend_from = {1.0, 1.0, 1.0}, blend_to = {0.5, 0.7, 1.0};
  vec3 normalized_direction;
  glm_vec3_normalize_to(r->direction, normalized_direction);
  f32 blend_factor = 0.5 * (normalized_direction[1] + 1.0);
  glm_vec3_lerpc(blend_from, blend_to, blend_factor, color);
}

typedef struct {
  i32 width, height;
  f32 viewport_width, viewport_height;
  vec3 center, up, right;
  vec3 viewport_u, viewport_v;
  vec3 pixel_delta_u, pixel_delta_v;
  vec3 top_left;
} camera;

void init_camera(camera *cam, i32 width, i32 height, i32 samples, f32 fov,
                 vec3 center, vec3 up, vec3 right) {
  cam->width = width;
  cam->height = height;
  cam->viewport_height = 2.0 * tan(fov * GLM_PI / 360);
  cam->viewport_width = cam->viewport_height * width / height;
  glm_vec3_copy(center, cam->center);
  glm_vec3_copy(up, cam->up);
  glm_vec3_copy(right, cam->right);
  glm_vec3_scale(right, cam->viewport_width, cam->viewport_u);
  glm_vec3_scale(up, -cam->viewport_height, cam->viewport_v);
  glm_vec3_divs(cam->viewport_u, width, cam->pixel_delta_u);
  glm_vec3_divs(cam->viewport_v, height, cam->pixel_delta_v);

  vec3 direction;
  glm_vec3_cross(up, right, cam->top_left);
  glm_vec3_add(cam->top_left, center, cam->top_left);
  glm_vec3_muladds(cam->viewport_u, -0.5, cam->top_left);
  glm_vec3_muladds(cam->viewport_v, -0.5, cam->top_left);
  glm_vec3_muladds(cam->pixel_delta_u, 0.5, cam->top_left);
  glm_vec3_muladds(cam->pixel_delta_v, 0.5, cam->top_left);
}

void get_trace_ray(camera *cam, f32 u, f32 v, ray *trace_ray) {
  glm_vec3_copy(cam->center, trace_ray->origin);
  glm_vec3_sub(cam->top_left, cam->center, trace_ray->direction);
  glm_vec3_muladds(cam->pixel_delta_u, u, trace_ray->direction);
  glm_vec3_muladds(cam->pixel_delta_v, v, trace_ray->direction);
}

void gamma_correct(vec3 color) {
  static const f32 gamma = 2.0;
  for (i32 i = 0; i < 3; ++i) {
    color[i] = pow(color[i], 1.0 / gamma);
  }
}

int main() {
  srand(time(NULL));
  const i32 width = 1280, height = 720, samples = 500;
  const f32 fov = 20;

  camera cam;
  {
    vec3 pos = {13, 2, 3};
    vec3 lookat = {0, 0, 0};
    vec3 up = {0, 1, 0};
    vec3 right, look_dir;
    glm_vec3_sub(lookat, pos, look_dir);
    glm_vec3_normalize(look_dir);
    glm_vec3_crossn(look_dir, up, right);
    glm_vec3_cross(right, look_dir, up);
    init_camera(&cam, width, height, samples, fov, pos, up, right);
  }

  material ground = make_lambertian((vec3){0.8, 0.8, 0.0});
  material center = make_lambertian((vec3){0.1, 0.2, 0.5});
  material left = make_dielectric(1.5);
  material right = make_metal((vec3){0.8, 0.6, 0.2}, 0.0);
  material ground2 = make_lambertian((vec3){0.5, 0.5, 0.5});
  material m1 = make_dielectric(1.5);
  material m2 = make_lambertian((vec3){0.4, 0.2, 0.1});
  material m3 = make_metal((vec3){0.7, 0.6, 0.5}, 0.0);

#define MAX_SPHERES 600
  material random_materials[MAX_SPHERES];
  f32 R = cos(GLM_PI / 4);
  hittable objects[MAX_SPHERES] = {
#if 0
      make_sphere(&center, (vec3){0, 0, -1}, 0.5),
      make_sphere(&ground, (vec3){0, -100.5, -1}, 100.0),
      make_sphere(&left, (vec3){-1, 0, -1}, 0.5),
      make_sphere(&left, (vec3){-1, 0, -1}, -0.4),
      make_sphere(&right, (vec3){1, 0, -1}, 0.5),
#endif
  };

  i32 num_objects;
  {
    i32 i = 0;
    for (i32 a = -11; a < 11; ++a) {
      for (i32 b = -11; b < 11; ++b) {
        f32 choose_mat = random_float();
        vec3 center = {a + 0.9 * random_float(), 0.2, b + 0.9 * random_float()};

        vec3 dist;
        glm_vec3_sub(center, (vec3){4, 0.2, 0}, dist);
        if (glm_vec3_norm2(dist) > 0.9 * 0.9) {
          if (i > MAX_SPHERES) {
            fprintf(
                stderr,
                "too many spheres! please increase the MAX_SPHERES variable.");
            return 1;
          }

          if (choose_mat < 0.8) {
            vec3 albedo, v1, v2;
            random_vec3(v1);
            random_vec3(v2);
            glm_vec3_mul(v1, v2, albedo);
            random_materials[i] = make_lambertian(albedo);
          } else if (choose_mat < 0.95) {
            vec3 albedo;
            random_vec3(albedo);
            glm_vec3_scale(albedo, 0.5, albedo);
            glm_vec3_adds(albedo, 0.5, albedo);
            f32 fuzz = random_float();
            random_materials[i] = make_metal(albedo, fuzz);
          } else {
            random_materials[i] = make_dielectric(1.5);
          }

          objects[i] = make_sphere(&random_materials[i], center, 0.2);
          ++i;
        }
      }
    }

    num_objects = i;
  }

  if (MAX_SPHERES >= num_objects + 4) {
    objects[num_objects] = make_sphere(&ground2, (vec3){0, -1000, 0}, 1000);
    objects[num_objects + 1] = make_sphere(&m1, (vec3){0, 1, 0}, 1);
    objects[num_objects + 2] = make_sphere(&m2, (vec3){-4, 1, 0}, 1);
    objects[num_objects + 3] = make_sphere(&m3, (vec3){4, 1, 0}, 1);
    num_objects += 4;
  }

  // dump_image_header(stdout, width, height);
  u8 *pixels = malloc(width * height * 4);
#pragma omp parallel for
  for (i32 y = 0; y < height; ++y) {
    fprintf(stderr, "Rendering y = %" PRIi32 "\n", y);
    for (i32 x = 0; x < width; ++x) {
      vec3 color;
      glm_vec3_zero(color);
      for (i32 sample = 0; sample < samples; ++sample) {
        f32 msaa_x = x + random_float();
        f32 msaa_y = y + random_float();
        ray trace_ray;
        get_trace_ray(&cam, msaa_x, msaa_y, &trace_ray);
        vec3 msaa_color;
        get_ray_color(&trace_ray, objects, num_objects, msaa_color, 0);
        glm_vec3_add(color, msaa_color, color);
      }

      glm_vec3_divs(color, samples, color);
      gamma_correct(color);
      i32 off = (y * width + x) * 4;
      pixels[off] = (i32)round(255 * color[0]);
      pixels[off + 1] = (i32)round(255 * color[1]);
      pixels[off + 2] = (i32)round(255 * color[2]);
      pixels[off + 3] = 255;
    }
  }

  stbi_write_png("output.png", width, height, 4, pixels, 0);

  return 0;
}