tinyraytracer.cpp

// tinyraytracer rewritten for the sake of performance

#include <limits>
#include <iostream>
#include <fstream>
#include <vector>
#include <algorithm>

#include <cstdint>
#include <cmath>

#include <pspkernel.h>
#include <pspdebug.h>
#include <psprtc.h>
#include <psppower.h>

#include "geometry.h"


PSP_MODULE_INFO("TinyRayTracer", 0, 1, 0);
PSP_MAIN_THREAD_ATTR(THREAD_ATTR_VFPU | THREAD_ATTR_USER);
PSP_HEAP_SIZE_KB(-1024);


struct Light {
    Light(const Vec3f &p, const float i) : position(p), intensity(i) {}
    Vec3f position;
    float intensity;
};

struct Material {
    Material(const float r, const Vec4f &a, const Vec3f &color, const float spec) : refractive_index(r), albedo(a), diffuse_color(color), specular_exponent(spec) {}
    Material() : refractive_index(1), albedo(1, 0, 0, 0), diffuse_color(), specular_exponent() {}
    float refractive_index;
    Vec4f albedo;
    Vec3f diffuse_color;
    float specular_exponent;
};

struct Sphere {
    Vec3f center;
    const float radius, radius_pow2;
    Material material;

    Sphere(const Vec3f &c, const float r, const Material &m) : center(c), radius(r), radius_pow2(r * r), material(m) {}

    bool ray_intersect(const Vec3f &orig, const Vec3f &dir, float &t0) const {
        Vec3f L = center - orig;
        const float tca = L * dir;
        const float d2 = L * L - tca * tca;
        if (d2 > radius_pow2) return false;
        const float thc = sqrtf(radius_pow2 - d2);
        t0       = tca - thc;
        const float t1 = tca + thc;
        if (t0 < 0) t0 = t1;
        if (t1 < 0) return false;
        return true;
    }
};


Vec3f reflect(const Vec3f &I, const Vec3f &N) {
    return I - N * 2.f * (I * N);
}

Vec3f refract(const Vec3f &I, const Vec3f &N, const float eta_t, const float eta_i=1.f) { // Snell's law
    const float cosi = -std::max(-1.f, std::min(1.f, I * N));
    if (cosi < 0) return refract(I, -N, eta_i, eta_t); // if the ray comes from the inside the object, swap the air and the media
    float eta = eta_i / eta_t;
    float k = 1 - eta * eta * (1 - cosi * cosi);
    // As i checked, k is never less than zero, so just removing this branch
    //return k < 0 ? Vec3f(1, 0, 0) : I * eta + N * (eta * cosi - sqrtf(k)); // k<0 = total reflection, no ray to refract. I refract it anyways, this has no physical meaning
    return I * eta + N * (eta * cosi - sqrtf(k));
}

bool scene_intersect(const Vec3f &orig, const Vec3f &dir, const std::vector<Sphere> &spheres, Vec3f &hit, Vec3f &N, Material &material) {
    float spheres_dist = std::numeric_limits<float>::max();
    for (size_t i=0; i < spheres.size(); ++i) {
        float dist_i;
        if (spheres[i].ray_intersect(orig, dir, dist_i) && dist_i < spheres_dist) {
            spheres_dist = dist_i;
            hit = orig + dir * dist_i;
            N = (hit - spheres[i].center).normalize();
            material = spheres[i].material;
        }
    }

    float checkerboard_dist = std::numeric_limits<float>::max();
    if (fabs(dir.y) > 1e-3f)  {
        float d = -(orig.y + 4) / dir.y; // the checkerboard plane has equation y = -4
        Vec3f pt = orig + dir * d;
        if (d > 0 && fabs(pt.x) < 10 && pt.z < -10 && pt.z > -30 && d < spheres_dist) {
            checkerboard_dist = d;
            hit = pt;
            N = Vec3f(0.f, 1.f, 0.f);
            material.diffuse_color = (int(.5f * hit.x + 1000) + int(.5f * hit.z)) & 1 ? Vec3f(.3f, .3f, .3f) : Vec3f(.3f, .2f, .1f);
        }
    }
    return std::min(spheres_dist, checkerboard_dist) < 1000;
}

Vec3f cast_ray(const Vec3f &orig, const Vec3f &dir, const std::vector<Sphere> &spheres, const std::vector<Light> &lights, const size_t depth = 0) {
    Vec3f point, N;
    Material material;

    if (depth > 4 || !scene_intersect(orig, dir, spheres, point, N, material))
        return Vec3f(0.2f, 0.7f, 0.8f); // background color

    Vec3f reflect_dir = reflect(dir, N).normalize();
    Vec3f refract_dir = refract(dir, N, material.refractive_index).normalize();
    Vec3f reflect_orig = reflect_dir * N < 0 ? point - N * 1e-3f : point + N * 1e-3f; // offset the original point to avoid occlusion by the object itself
    Vec3f refract_orig = refract_dir * N < 0 ? point - N * 1e-3f : point + N * 1e-3f;
    Vec3f reflect_color = cast_ray(reflect_orig, reflect_dir, spheres, lights, depth + 1);
    Vec3f refract_color = cast_ray(refract_orig, refract_dir, spheres, lights, depth + 1);

    float diffuse_light_intensity = 0, specular_light_intensity = 0;
    for (size_t i = 0; i < lights.size(); ++i) {
        Vec3f light_dir      = (lights[i].position - point).normalize();
        float light_distance = (lights[i].position - point).norm();

        Vec3f shadow_orig = light_dir * N < 0 ? point - N * 1e-3f : point + N * 1e-3f; // checking if the point lies in the shadow of the lights[i]
        Vec3f shadow_pt, shadow_N;
        Material tmpmaterial;
        if (scene_intersect(shadow_orig, light_dir, spheres, shadow_pt, shadow_N, tmpmaterial) && (shadow_pt-shadow_orig).norm() < light_distance)
            continue;

        diffuse_light_intensity  += lights[i].intensity * std::max(0.f, light_dir * N);
        specular_light_intensity += powf(std::max(0.f, -reflect(-light_dir, N) * dir), material.specular_exponent) * lights[i].intensity;
    }
    return material.diffuse_color * diffuse_light_intensity * material.albedo[0] + Vec3f(1.f, 1.f, 1.f) * specular_light_intensity * material.albedo[1] + reflect_color * material.albedo[2] + refract_color * material.albedo[3];
}


u64 finalTick;
u32 tickRes;
u64 lastTick;

void render(const std::vector<Sphere> &spheres, const std::vector<Light> &lights) {
    constexpr int width = 480, height = 272, half_w = width / 2, half_h  = height / 2;
    constexpr float fov = M_PI / 3.f;
    const float tan_half_fov = tanf(fov / 2.f);
    const float dir_z = -height / (2.f * tan_half_fov);
    std::vector<Vec3f> framebuffer(width*height);
    pspDebugScreenPrintf("INIT RENDER!\n");

#ifdef TRT_REALTIME_OUTPUT
    struct Pixel { uint8_t r, g, b, _pad; };
    volatile Pixel *realFB = (Pixel*)0x44000000;
    auto clamp = [](auto val, auto from, auto to){ return std::max(from, std::min(to, val)); };
#endif
    for (size_t j = 0; j < height; ++j) { // actual rendering loop
        for (size_t i = 0; i < width; ++i) {
            float dir_x =  (i + 0.5f) - half_w;
            float dir_y = -(j + 0.5f) + half_h;    // this flips the image at the same time
            auto res = cast_ray(Vec3f(0, 0, 0), Vec3f(dir_x, dir_y, dir_z).normalize(), spheres, lights);
            framebuffer[i + j * width] = res;
#ifdef TRT_REALTIME_OUTPUT
            /*
             * PSP's framebuffer has kinda weird layout, it is 512x272x4 bytes memory chunk actually (not 480x272x4)
             * But last 32 pixels (or 32x4=128 bytes) of every row are unused
             */
            realFB[i + j * (width + 32)].r = (uint8_t)(clamp(res[0], 0.f, 1.f) * 255);
            realFB[i + j * (width + 32)].g = (uint8_t)(clamp(res[1], 0.f, 1.f) * 255);
            realFB[i + j * (width + 32)].b = (uint8_t)(clamp(res[2], 0.f, 1.f) * 255);
#endif
        }
    }

    sceRtcGetCurrentTick(&finalTick);

    pspDebugScreenPrintf("OUTPUT!\n");
    std::ofstream ofs; // save the framebuffer to file
    ofs.open("./out.ppm",std::ios::binary);

    pspDebugScreenPrintf("OPEN FILE!\n");
    ofs << "P6\n" << width << " " << height << "\n255\n";
    pspDebugScreenPrintf("HEADER!\n");
    pspDebugScreenPrintf("WRITING!\n");
    for (size_t i = 0; i < height*width; ++i) {
        Vec3f &c = framebuffer[i];
        float max = std::max(c[0], std::max(c[1], c[2]));
        if (max>1) c = c*(1.f/max);
        for (size_t j = 0; j<3; j++) {
            ofs << (char)(255 * std::max(0.f, std::min(1.f, framebuffer[i][j])));
        }
    }
    ofs.close();
    pspDebugScreenPrintf("CLOSE!\n");

    
    double dt = ((double)finalTick - (double)lastTick) / (double)tickRes;
    pspDebugScreenPrintf("BENCHMARK TIME: %.3f seconds", dt);
    // also write benchmark time to file
    std::ofstream ofs2;
    ofs2.open("./out.ppm.txt");
    ofs2 << dt << "s";
    ofs2.flush();
    ofs2.close();
    sceKernelDelayThread(1000 * 1000);

#ifndef TRT_REALTIME_OUTPUT
    for (size_t i = 0; i < height * width; ++i) {
        Vec3f& c = framebuffer[i];
        float max = std::max(c[0], std::max(c[1], c[2]));
        if (max > 1) c = c * (1.f / max);

        char* ccc = (char*)0x44000000;

        int rows = i / 480;
        int columnOff = i % 480;

        for (size_t j = 0; j < 4; j++) {
            if (j == 3) {
                ccc[rows * 4 + j] = 255;
                break;
            }
            ccc[rows * 512 * 4 + columnOff * 4 + j] = (char)(255 * std::max(0.f, std::min(1.f, framebuffer[i][j])));
        }
    }
#endif
}

int main() {

    scePowerSetClockFrequency(333, 333, 166);
    pspDebugScreenInit();

    tickRes = sceRtcGetTickResolution();


    Material      ivory(1.0f, Vec4f(.6f,  .3f, .1f, 0.f), Vec3f(.4f, .4f, .3f),   50.f);
    Material      glass(1.5f, Vec4f(0.f,  .5f, .1f, .8f), Vec3f(.6f, .7f, .8f),  125.f);
    Material red_rubber(1.0f, Vec4f(.9f,  .1f, 0.f, 0.f), Vec3f(.3f, .1f, .1f),   10.f);
    Material     mirror(1.0f, Vec4f(0.f, 10.f, .8f, 0.f), Vec3f(1.f, 1.f, 1.f), 1425.f);

    std::vector<Sphere> spheres;
    spheres.push_back(Sphere(Vec3f(-3,    0,   -16), 2,      ivory));
    spheres.push_back(Sphere(Vec3f(-1.0f, -1.5f, -12), 2,      glass));
    spheres.push_back(Sphere(Vec3f( 1.5f, -0.5f, -18), 3, red_rubber));
    spheres.push_back(Sphere(Vec3f( 7,    5,   -18), 4,     mirror));

    std::vector<Light>  lights;
    lights.push_back(Light(Vec3f(-20, 20,  20), 1.5f));
    lights.push_back(Light(Vec3f( 30, 50, -25), 1.8f));
    lights.push_back(Light(Vec3f( 30, 20,  30), 1.7f));
		
    pspDebugScreenSetXY(0, 0);
	pspDebugScreenPrintf("RENDER!\n");

    sceRtcGetCurrentTick(&lastTick);
    render(spheres, lights);


    sceKernelDelayThread(5000 * 1000);

    sceKernelExitGame();
    return 0;
}