raytracer/src/main.cpp

#include <SDL.h>
#include <array>
#include <future>
#include <glm/glm.hpp>
#include <vector>

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb_image_write.h>

#include "camera.h"
#include "glad/glad.h"
#include "image.h"
#include "imgui.h"
#include "imgui_impl_opengl3.h"
#include "imgui_impl_sdl.h"
#include "intersections.h"
#include "lighting.h"
#include "scene.h"

#define TINYOBJLOADER_IMPLEMENTATION
#include <tiny_obj_loader.h>

// scene information
constexpr int32_t width = 256, height = 256;
bool use_bvh = true;

const Camera camera = [] {
    Camera camera;
    camera.look_at(glm::vec3(0, 0, 3), glm::vec3(0));

    return camera;
}();

// internal variables
constexpr int32_t tile_size = 32;
constexpr int32_t num_tiles_x = width / tile_size;
constexpr int32_t num_tiles_y = height / tile_size;

// globals
Scene scene = {};
Image<glm::vec4, width, height> colors = {};
bool image_dirty = false;

enum class DisplayMode {
    Combined,
    Direct,
    Indirect,
    Reflect
};

const std::array diplay_mode_strings = {"Combined", "Direct", "Indirect", "Reflect"};

inline DisplayMode display_mode;

bool calculate_tile(const int32_t from_x, const int32_t to_width, const int32_t from_y, const int32_t to_height) {
    for (int32_t y = from_y; y < (from_y + to_height); y++) {
        for (int32_t x = from_x; x < (from_x + to_width); x++) {
            Ray ray_camera = camera.get_ray(x, y, width, height);

            if (auto result = cast_scene(ray_camera, scene, use_bvh)) {
                glm::vec3 chosen_display;
                switch (display_mode) {
                    case DisplayMode::Combined:
                        chosen_display = result->combined;
                        break;
                    case DisplayMode::Direct:
                        chosen_display = result->direct;
                        break;
                    case DisplayMode::Indirect:
                        chosen_display = result->indirect;
                        break;
                    case DisplayMode::Reflect:
                        chosen_display = result->reflect;
                        break;
                }

                colors.get(x, y) = glm::vec4(chosen_display, 1.0f);

                image_dirty = true;
            }
        }
    }

    return true;
}

GLuint quad_vao = 0;
GLuint pixel_program = 0;
GLuint pixels_texture = 0;

void setup_gfx() {
    // create quad for pixel rendering
    constexpr std::array vertices = {
        -0.5f, 0.5f, 0.0f, 0.0f,  0.0f,  0.5f,  0.5f, 0.0f, 1.0f, 0.0f, 0.5f,
        -0.5f, 0.0f, 1.0f, -1.0f, -0.5f, -0.5f, 0.0f, 0.0f, -1.0f // we render it upside down (to opengl) so we flip our
                                                                  // tex coord
    };

    constexpr std::array elements = {0, 1, 2, 2, 3, 0};

    glGenVertexArrays(1, &quad_vao);
    glBindVertexArray(quad_vao);

    GLuint vbo = 0;
    glGenBuffers(1, &vbo);
    glBindBuffer(GL_ARRAY_BUFFER, vbo);
    glBufferData(GL_ARRAY_BUFFER, vertices.size() * sizeof(float), vertices.data(), GL_STATIC_DRAW);

    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 5 * sizeof(float), nullptr);
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, 5 * sizeof(float), reinterpret_cast<void*>(3 * sizeof(float)));

    glEnableVertexAttribArray(0);
    glEnableVertexAttribArray(1);

    GLuint ebo = 0;
    glGenBuffers(1, &ebo);
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ebo);
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, elements.size() * sizeof(uint32_t), elements.data(), GL_STATIC_DRAW);

    constexpr std::string_view vertex_glsl = "#version 330 core\n"
                                             "layout (location = 0) in vec3 in_position;\n"
                                             "layout (location = 1) in vec2 in_uv;\n"
                                             "out vec2 uv;\n"
                                             "void main()\n"
                                             "{\n"
                                             "   gl_Position = vec4(in_position.xyz, 1.0);\n"
                                             "   uv = in_uv;\n"
                                             "}\n";
    const char* vertex_src = vertex_glsl.data();

    constexpr std::string_view fragment_glsl = "#version 330 core\n"
                                               "in vec2 uv;\n"
                                               "out vec4 out_color;\n"
                                               "uniform sampler2D pixel_texture;\n"
                                               "void main()\n"
                                               "{\n"
                                               "    out_color = texture(pixel_texture, uv);\n"
                                               "}\n";
    const char* fragment_src = fragment_glsl.data();

    GLuint vertex_shader = glCreateShader(GL_VERTEX_SHADER);
    glShaderSource(vertex_shader, 1, &vertex_src, nullptr);
    glCompileShader(vertex_shader);

    GLuint fragment_shader = glCreateShader(GL_FRAGMENT_SHADER);
    glShaderSource(fragment_shader, 1, &fragment_src, nullptr);
    glCompileShader(fragment_shader);

    pixel_program = glCreateProgram();

    glAttachShader(pixel_program, vertex_shader);
    glAttachShader(pixel_program, fragment_shader);
    glLinkProgram(pixel_program);

    glDeleteShader(vertex_shader);
    glDeleteShader(fragment_shader);

    glGenTextures(1, &pixels_texture);
    glBindTexture(GL_TEXTURE_2D, pixels_texture);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, width, height, 0, GL_RGBA, GL_FLOAT, nullptr);
    glBindTexture(GL_TEXTURE_2D, 0);
}

void update_texture() {
    glBindTexture(GL_TEXTURE_2D, pixels_texture);
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, width, height, 0, GL_RGBA, GL_FLOAT, colors.array.data());
    glBindTexture(GL_TEXTURE_2D, 0);
}

std::vector<std::future<bool>> futures;

void render() {
    futures.clear();
    colors.reset();

    for (int32_t y = 0; y < num_tiles_y; y++) {
        for (int32_t x = 0; x < num_tiles_x; x++) {
            auto f = std::async(std::launch::async, calculate_tile, x * tile_size, tile_size, y * tile_size, tile_size);
            futures.push_back(std::move(f));
        }
    }
}

void dump_to_file() {
    uint8_t pixels[width * height * 3] = {};

    int i = 0;
    for (int32_t y = height - 1; y >= 0; y--) {
        for (int32_t x = 0; x < width; x++) {
            const glm::ivec4 c = colors.get(x, y);
            pixels[i++] = c.r;
            pixels[i++] = c.g;
            pixels[i++] = c.b;
        }
    }

    stbi_write_png("output.png", width, height, 3, pixels, width * 3);
}

template<typename UnderlyingType> void walk_node(Node<UnderlyingType>& node) {
    if (ImGui::TreeNode(
            &node,
            "min: (%f %f %f)\n max: (%f %f %f)",
            node.extent.min.x,
            node.extent.min.y,
            node.extent.min.z,
            node.extent.max.x,
            node.extent.max.y,
            node.extent.max.z)) {
        ImGui::Text("Is split: %i", node.is_split);
        ImGui::Text("Contained triangles: %lu", node.contained_objects.size());

        if (node.is_split) {
            for (auto& child : node.children)
                walk_node(*child);
        }

        ImGui::TreePop();
    }
}

void walk_object(Object& object) {
    walk_node(object.octree->root);
}

int main(int, char*[]) {
    SDL_Init(SDL_INIT_VIDEO | SDL_INIT_TIMER);

    SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 3);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_FLAGS, SDL_GL_CONTEXT_FORWARD_COMPATIBLE_FLAG);
    SDL_GL_SetAttribute(SDL_GL_CONTEXT_PROFILE_MASK, SDL_GL_CONTEXT_PROFILE_CORE);

    SDL_Window* window = SDL_CreateWindow(
        "raytracer",
        SDL_WINDOWPOS_CENTERED,
        SDL_WINDOWPOS_CENTERED,
        800,
        600,
        SDL_WINDOW_OPENGL | SDL_WINDOW_RESIZABLE);

    SDL_GLContext gl_context = SDL_GL_CreateContext(window);
    SDL_GL_MakeCurrent(window, gl_context);
    SDL_GL_SetSwapInterval(1);

    gladLoadGL();
    setup_gfx();

    ImGui::CreateContext();
    ImGui::StyleColorsDark();

    ImGui_ImplSDL2_InitForOpenGL(window, gl_context);
    ImGui_ImplOpenGL3_Init("#version 330 core");

    bool running = true;
    while (running) {
        SDL_Event event = {};
        while (SDL_PollEvent(&event)) {
            ImGui_ImplSDL2_ProcessEvent(&event);

            if (event.type == SDL_QUIT)
                running = false;
        }

        ImGui_ImplOpenGL3_NewFrame();
        ImGui_ImplSDL2_NewFrame(window);

        ImGui::NewFrame();

        if (ImGui::BeginMainMenuBar()) {
            if (ImGui::BeginMenu("File")) {
                if (ImGui::Button("Load Example Models")) {
                    auto& sphere = scene.load_from_file("misc/sphere.obj");
                    sphere.color = {1, 1, 1};

                    auto& plane = scene.load_from_file("misc/plane.obj");
                    plane.position.y = -1;
                    plane.color = {1, 0, 0};

                    scene.generate_acceleration();
                }

                ImGui::EndMenu();
            }

            ImGui::EndMainMenuBar();
        }

        ImGui::Begin("Render Options");

        ImGui::Checkbox("Use BVH", &use_bvh);
        ImGui::InputInt("Indirect Samples", &num_indirect_samples);

        if (ImGui::BeginCombo("Channel Selection", diplay_mode_strings[static_cast<int>(display_mode)])) {
            if (ImGui::Selectable("Combined"))
                display_mode = DisplayMode::Combined;

            if (ImGui::Selectable("Direct"))
                display_mode = DisplayMode::Direct;

            if (ImGui::Selectable("Indirect"))
                display_mode = DisplayMode::Indirect;

            if (ImGui::Selectable("Reflect"))
                display_mode = DisplayMode::Reflect;

            ImGui::EndCombo();
        }

        if (ImGui::Button("Render"))
            render();

        ImGui::SameLine();

        if (ImGui::Button("Save Output"))
            dump_to_file();

        if (image_dirty) {
            update_texture();
            image_dirty = false;
        }

        for (auto& object : scene.objects) {
            if (ImGui::TreeNode("Object")) {
                walk_object(*object);

                ImGui::TreePop();
            }
        }

        ImGui::End();

        ImGui::Render();

        auto& io = ImGui::GetIO();
        glViewport(0, 0, (int)io.DisplaySize.x, (int)io.DisplaySize.y);
        glClear(GL_COLOR_BUFFER_BIT);

        glUseProgram(pixel_program);
        glBindVertexArray(quad_vao);
        glBindTexture(GL_TEXTURE_2D, pixels_texture);
        glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, 0);

        ImGui_ImplOpenGL3_RenderDrawData(ImGui::GetDrawData());

        SDL_GL_SwapWindow(window);
    }

    return 0;
}