stomp.c

#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdint.h>
#include <pthread.h>
#include <stdatomic.h>

/*
  Some terminology:

  - Tile: A sub region of the image which can be rendered separately from any other tiles
  - Fragment: A single pixel of the output image
  - Sample: A single raytrace within a single fragment
  - Object: A geometric shape which reacts to light
  - Bounces: The number of times a ray will bounce before halting
*/

#define WIDTH 640
#define HEIGHT 480
#define THREADS 16
#define SAMPLES 512
#define BOUNCES 13
#define BOUNCE_SAMPLES 1
#define TILE_WIDTH 16
#define TILES ((WIDTH + (TILE_WIDTH - 1)) / TILE_WIDTH)

#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define MAX(a, b) ((a) > (b) ? (a) : (b))

#define LENGTHOF(x) (sizeof(x) / sizeof((x)[0]))

// ========================================
//
// Vector.

typedef float Scalar;
typedef __attribute__((aligned(16))) struct Vector {
  Scalar x, y, z, w;
} Vector;

Vector vectorAdd(Vector a, Vector b) {
  return (Vector) { a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w };
}

Vector vectorAdds(Vector a, Scalar b) {
  return (Vector) { a.x + b, a.y + b, a.z + b, a.w + b };
}

Vector vectorSub(Vector a, Vector b) {
  return (Vector) { a.x - b.x, a.y - b.y, a.z - b.z, a.w - b.w };
}

Vector vectorSubs(Vector a, Scalar b) {
  return (Vector) { a.x - b, a.y - b, a.z - b, a.w - b };
}

Vector vectorMul(Vector a, Vector b) {
  return (Vector) { a.x * b.x, a.y * b.y, a.z * b.z, a.w * b.w };
}

Vector vectorMuls(Vector a, Scalar b) {
  return (Vector) { a.x * b, a.y * b, a.z * b, a.w * b };
}

Vector vectorDiv(Vector a, Vector b) {
  return (Vector) { a.x / b.x, a.y / b.y, a.z / b.z, a.w / b.w };
}

Vector vectorDivs(Vector a, Scalar b) {
  return (Vector) { a.x / b, a.y / b, a.z / b, a.w / b };
}

Vector vectorMulsAdd(Vector a, Vector b, Scalar c) {
  return (Vector) { a.x + b.x * c, a.y + b.y * c, a.z + b.z * c, a.w + b.w * c };
}

Vector vectorLerp(Vector a, Vector b, Scalar r) {
  return vectorAdd(vectorMuls(a, 1.0 - r), vectorMuls(b, r));
}

Scalar vectorDot(Vector a, Vector b) {
  return a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
}

Scalar vectorDot3(Vector a, Vector b) {
  return a.x * b.x + a.y * b.y + a.z * b.z;
}

Scalar vectorLength(Vector a) {
  return sqrtf(vectorDot(a, a));
}

Scalar vectorLength3(Vector a) {
  return sqrtf(vectorDot3(a, a));
}

Vector vectorNorm(Vector a) {
  return vectorDivs(a, vectorLength(a));
}

Vector vectorNorm3(Vector a) {
  return vectorDivs(a, vectorLength3(a));
}

Vector vectorReflect3(Vector i, Vector n) {
  return vectorMulsAdd(i, n, -2.0 * vectorDot3(i, n));
}

// ========================================
//
// Random.

__thread uint32_t random_seed;

uint32_t randInt(void) {
  random_seed = (random_seed * 8761381) + 1;
  return random_seed;
}

Scalar randUnorm(void) {
  return ((Scalar) randInt()) / (Scalar) UINT32_MAX;
}

Scalar randNorm(void) {
  return randUnorm() * 2.0 - 1.0;
}

Vector randVectorSphere3(void) {
  Vector v = { 2.0 };
  while(vectorDot3(v, v) > 1.0) {
    v = (Vector) { randNorm(), randNorm(), randNorm() };
  }
  return v;
}

void randSeed(uint32_t seed) {
  random_seed = seed;
  randInt();
}

// ========================================
//
// Raytrace.

typedef struct Hit {
  bool hit;
  Scalar distance;
  Vector position;
  Vector normal;
} Hit;

// Thanks iq! https://iquilezles.org/www/articles/intersectors/intersectors.htm
Hit rayVsSphere(Vector ray_origin, Vector ray_dir, Vector position, Scalar radius) {
  Vector rel = vectorSub(ray_origin, position);
  Scalar b = vectorDot3(rel, ray_dir);
  Scalar c = vectorDot3(rel, rel) - (radius * radius);
  Scalar h = (b * b) - c;

  Hit hit = { .hit = false };

  if(h >= 0.0 && -b > 0.0) {
    hit.hit = true;
    hit.distance = -b - sqrtf(h);
    hit.position = vectorAdd(ray_origin, vectorMuls(ray_dir, hit.distance));
    hit.normal = vectorNorm3(vectorSub(hit.position, position));
  }

  return hit;
}

// Thanks iq! https://iquilezles.org/www/articles/intersectors/intersectors.htm
Hit rayVsPlane(Vector ray_origin, Vector ray_dir, Vector plane) {
  Hit hit = { .hit = false };

  hit.distance = -(vectorDot3(ray_origin, plane) - plane.w) / vectorDot3(ray_dir, plane);

  if(hit.distance > 0.0) {
    hit.hit = true;
    hit.position = vectorAdd(ray_origin, vectorMuls(ray_dir, hit.distance));
    hit.normal = plane;
  }

  return hit;
}

// ========================================
//
// Scene definition.

typedef enum ObjectType {
  OBJECT_TYPE_SKY = 0,
  OBJECT_TYPE_PLANE,
  OBJECT_TYPE_SPHERE,
} ObjectType;

typedef struct Object {
  ObjectType type;

  // Transformation.
  Vector position;
  Vector normal;
  Scalar radius;

  // Material.
  Vector albedo;
  Scalar roughness;
  bool emissive;
} Object;

Object objects[] = {
  // Sky.
  {
    .type = OBJECT_TYPE_SKY,
    .albedo = { 0.0, 0.0, 0.0 },
    .emissive = true,
  },

  // P's "Cornell Box"
#if 1
  // Floor.
  {
    .type = OBJECT_TYPE_PLANE,
    .normal = { 0.0, 0.0, 1.0, -1.0 },
    .albedo = { 1.0, 1.0, 1.0 },
    .roughness = 1.0,
  },

  // Ceiling.
  {
    .type = OBJECT_TYPE_PLANE,
    .normal = { 0.0, 0.0, -1.0, -2.0 },
    .albedo = { 1.0, 1.0, 1.0 },
    .emissive = true,
  },

  // Left wall.
  {
    .type = OBJECT_TYPE_PLANE,
    .normal = { 1.0, 0.0, 0.0, -2.0 },
    .albedo = { 1.0, 0.5, 0.3 },
    .roughness = 1.0,
  },

  // Back wall.
  {
    .type = OBJECT_TYPE_PLANE,
    .normal = { 0.0, -1.0, 0.0, -3.0 },
    .albedo = { 0.3, 0.6, 1.0 },
    .roughness = 1.0,
  },

  // Right wall.
  {
    .type = OBJECT_TYPE_PLANE,
    .normal = { -1.0, 0.0, 0.0, -2.0 },
    .albedo = { 0.6, 0.1, 0.6 },
    .roughness = 1.0,
  },

  // Left sphere.
  {
    .type = OBJECT_TYPE_SPHERE,
    .position = { -1.0, 0.5, -0.5 },
    .radius = 0.5,
    .albedo = { 1.0, 1.0, 1.0 },
    .roughness = 1.0,
  },

  // Center sphere.
  {
    .type = OBJECT_TYPE_SPHERE,
    .position = { 0.0, -1.5, -0.75 },
    .radius = 0.25,
    .albedo = { 0.95, 0.98, 0.97 },
    .roughness = 0.0,
  },

  // Right sphere.
  {
    .type = OBJECT_TYPE_SPHERE,
    .position = { 1.0, 0.5, -0.5 },
    .radius = 0.5,
    .albedo = { 1.0, 1.0, 1.0 },
    .roughness = 0.1,
  },
#endif
};

Vector view_origin = { 0.0, -4.0, 0.0 };

// ========================================
//
// Path tracing.

Vector accumulateSample(Vector ray_origin, Vector ray_dir, int bounce) {
  Scalar distance = 100000.0;
  Object *object = &objects[0];
  Hit hit = { .hit = false };

  for(int i = 1; i < LENGTHOF(objects); i++) {
    Object *object_ = &objects[i];
    Hit hit_;

    if(object_->type == OBJECT_TYPE_PLANE) {
      hit_ = rayVsPlane(ray_origin, ray_dir, object_->normal);
    } else if(object_->type == OBJECT_TYPE_SPHERE) {
      hit_ = rayVsSphere(ray_origin, ray_dir, object_->position, object_->radius);
    }

    if(hit_.hit && hit_.distance < distance) {
      distance = hit_.distance;
      object = object_;
      hit = hit_;
    }
  }

  if(object->emissive) {
    return object->albedo;
  } else {
    Vector lighting = { 0 };

    if(hit.hit && bounce > 0) {
      Vector reflect_dir = vectorReflect3(ray_dir, hit.normal);

      for(int i = 0; i < BOUNCE_SAMPLES; i++) {
        Vector sample_dir = randVectorSphere3();
        if(vectorDot3(sample_dir, hit.normal) < 0.0) {
          vectorMuls(sample_dir, -1.0);
        }

        sample_dir = vectorLerp(reflect_dir, sample_dir, object->roughness);
        sample_dir = vectorNorm3(sample_dir);

        Vector sample = accumulateSample(vectorMulsAdd(hit.position, hit.normal, 0.00001), sample_dir, bounce - 1);

        lighting = vectorAdd(lighting, sample);
      }
    }

    return vectorMul(lighting, object->albedo);
  }
}

// ========================================
//
// Initial sample dispatch.

uint8_t fragments[WIDTH][HEIGHT][3];

Vector renderSample(Scalar x, Scalar y) {
  Vector ray_dir = {
    ((x / WIDTH) - 0.5) * (WIDTH / (Scalar) HEIGHT),
    1.0,
    0.5 - (y / HEIGHT)
  };
  ray_dir = vectorNorm3(ray_dir);

  return accumulateSample(view_origin, ray_dir, BOUNCES);
}

void renderTile(int tile_index) {
  fprintf(stderr, "Rendering tile %d/%d...\n", tile_index + 1, TILES);

  randSeed(tile_index);

  int tile_x = tile_index * TILE_WIDTH;
  int tile_width = TILE_WIDTH;
  if((tile_x + tile_width) > WIDTH) {
    tile_width = WIDTH - tile_x;
  }

  for(int x = tile_x; x < (tile_x + tile_width); x++) {
    for(int y = 0; y < HEIGHT; y++) {
      Vector lighting = { 0 };
      for(int i = 0; i < SAMPLES; i++) {
        Scalar sample_x = (Scalar) x;
        Scalar sample_y = (Scalar) y;

        if(i > 0) {
          sample_x += randNorm() * 0.5;
          sample_y += randNorm() * 0.5;
        }

        lighting = vectorMulsAdd(lighting, renderSample(sample_x, sample_y), (1.0 / SAMPLES) / BOUNCE_SAMPLES);
      }

      //lighting = vectorMuls(lighting, 0.4);

      fragments[x][y][0] = (uint8_t) MAX(0, MIN(lighting.x * 255, 255));
      fragments[x][y][1] = (uint8_t) MAX(0, MIN(lighting.y * 255, 255));
      fragments[x][y][2] = (uint8_t) MAX(0, MIN(lighting.z * 255, 255));
    }
  }
}

// ========================================
//
// Threaded tile dispatch.

_Atomic int tiles_complete;

void *renderThread(void *user_pointer) {
  int tile_index;
  while((tile_index = atomic_fetch_add(&tiles_complete, 1)) < TILES) {
    renderTile(tile_index);
  }

  return NULL;
}

// ========================================
//
// Entrypoint.

int main(int argc, char **argv) {
  atomic_init(&tiles_complete, 0);

  // Thread dispatch.
  pthread_t threads[THREADS];
  for(int i = 0; i < THREADS; i++) {
    pthread_create(&threads[i], NULL, renderThread, NULL);
  }
  for(int i = 0; i < THREADS; i++) {
    pthread_join(threads[i], NULL);
  }

  // Add simple watermark.
  for(int x = 0; x < 10; x++) {
    for(int y = 0; y < 10; y++) {
      fragments[x][y][0] = 127;
      fragments[x][y][1] = 255;
      fragments[x][y][2] = 0;
    }
  }

  // Write to standard out.
  printf("P3 %d %d 255", WIDTH, HEIGHT);
  for(int y = 0; y < HEIGHT; y++) {
    for(int x = 0; x < WIDTH; x++) {
      printf(" %d %d %d", fragments[x][y][0], fragments[x][y][1], fragments[x][y][2]);
    }
  }
  printf("\n");

  return 0;
}