rogue/source/rogue.cpp

#include "rogue.hpp"

#define STB_IMAGE_IMPLEMENTATION
#include <stb/stb_image.h>

#include "vuk/ImageAttachment.hpp"
#include "vuk/RenderGraph.hpp"
#include "vuk/Types.hpp"
#include "vuk/Value.hpp"
#include "vuk/runtime/CommandBuffer.hpp"
#include "vuk/runtime/vk/Allocator.hpp"
#include "vuk/runtime/vk/DeviceFrameResource.hpp"
#include "vuk/runtime/vk/Image.hpp"
#include "vuk/runtime/vk/VkRuntime.hpp"
#include "vuk/vsl/Core.hpp"
#include <backends/imgui_impl_glfw.h>
#include <glm/gtc/matrix_transform.hpp>

#include <clip.h>
#include <imgui.h>
#include <vuk/extra/ImGuiIntegration.hpp>
#include <vuk/runtime/vk/Pipeline.hpp>
#include <vuk/runtime/vk/PipelineTypes.hpp>
#include <vuk/vsl/BindlessArray.hpp>

[[nodiscard]] auto raytrace() {
  return vuk::make_pass("composite", [](vuk::CommandBuffer &command_buffer,
                                        VUK_BA(vuk::eComputeRead) camera_buffer,
                                        VUK_IA(vuk::eComputeRW) target) {
    command_buffer.bind_compute_pipeline("rt_compute")
        .bind_buffer(0, 0, camera_buffer)
        .bind_image(0, 1, target)
        .dispatch_invocations_per_pixel(target);

    return std::make_tuple(std::move(camera_buffer), std::move(target));
  });
}

[[nodiscard]] auto composite() {
  return vuk::make_pass("composite", [](vuk::CommandBuffer &command_buffer,
                                        VUK_IA(vuk::eFragmentSampled) rt_image,
                                        VUK_IA(vuk::eColorWrite) target) {
    command_buffer.bind_graphics_pipeline("display_rt")
        .set_dynamic_state(vuk::DynamicStateFlagBits::eScissor |
                           vuk::DynamicStateFlagBits::eViewport)
        .set_viewport(0, vuk::Rect2D::framebuffer())
        .set_scissor(0, vuk::Rect2D::framebuffer())
        .set_rasterization({})
        .set_depth_stencil({
            .depthTestEnable = false,
            .depthWriteEnable = false,
            .depthCompareOp = vuk::CompareOp::eLess,
        })
        .set_color_blend(target, {})
        .bind_image(0, 0, rt_image)
        .bind_sampler(0, 0, {})
        .draw(3, 1, 0, 0);

    return std::make_tuple(std::move(target));
  });
}

namespace rog {

rogue::rogue(logger *log)
    : log(log), display(&user_input, {1920, 1080}, "Rogue PT"),
      context(&display), shader_compiler("../assets/shaders") {
  create_pipelines();
  if (!context.rtx_supported) {
    log->warn("rtx is not supported on {}. inline ray tracing and ray tracing "
              "pipeline will not be available",
              context.gpu_name);
  }
  IMGUI_CHECKVERSION();
  ImGui::CreateContext();

  ImGui_ImplGlfw_InitForOther(display.raw(), true);
  imgui_data = vuk::extra::ImGui_ImplVuk_Init(context.superframe_allocator);

  texture_of_raytrace_accumulation = vuk::ImageAttachment::from_preset(
      vuk::ImageAttachment::Preset::eMap2D, vuk::Format::eR32G32B32A32Sfloat,
      vuk::Extent3D{1920u, 1080u, 1u}, vuk::Samples::e1);
  auto bad_data = std::vector<float>(1920 * 1080 * 4);
  for (int i = 0; i < bad_data.size(); i++) {
    bad_data[i] = (i + 1) % 4 == 0 ? 1.0f : 0.0f;
  }
  texture_of_raytrace_accumulation.level_count = 1;
  auto [image, view, future] = vuk::create_image_and_view_with_data(
      context.superframe_allocator, vuk::DomainFlagBits::eTransferOnTransfer,
      texture_of_raytrace_accumulation, std::span(bad_data));
  raytrace_accumulation_image = std::move(image);
  raytrace_accumulation_view = std::move(view);

  auto futures = std::vector<vuk::UntypedValue>();
  futures.push_back(std::move(future.as_released(
      vuk::Access::eFragmentSampled, vuk::DomainFlagBits::eGraphicsQueue)));

  vuk::wait_for_values_explicit(context.superframe_allocator, context.compiler,
                                futures);
}

void rogue::stop() { close_requested = true; }

void rogue::do_tick() {
  auto [frame_allocator, frame_target] = get_next_frame_resources();

  auto rt_val =
      vuk::declare_ia("rt_accumulation", texture_of_raytrace_accumulation);

  auto wants_final_image = false;

  glm::vec3 camPos = glm::vec3(0.0f, 0.0f, 50.0f);
  glm::vec3 target = glm::vec3(0.0f, 0.0f, 0.0f);
  glm::vec3 up = glm::vec3(0.0f, 1.0f, 0.0f);

  float fov = 60.0f;
  float aspect = 1920.0f / 1080.0f;
  float near = 0.1f, far = 100.0f;

  auto proj = glm::perspective(glm::radians(fov), aspect, near, far);
  struct {
    glm::mat4 camera_to_world;
    glm::mat4 inverse_projection;
    glm::uvec4 resolution;
  } camera_data{
      .camera_to_world =
          glm::transpose(glm::inverse(glm::lookAt(camPos, target, up))),
      .inverse_projection = glm::transpose(glm::inverse(proj)),
      .resolution = glm::uvec4(1920, 1080, 0, 0),
  };
  auto [camera_buffer, camera_future] = vuk::create_buffer(
      frame_allocator, vuk::MemoryUsage::eCPUtoGPU,
      vuk::DomainFlagBits::eTransferOperation, std::span(&camera_data, 1));

  std::tie(camera_future, rt_val) =
      raytrace()(std::move(camera_future), std::move(rt_val));

  std::tie(frame_target) =
      composite()(std::move(rt_val), std::move(frame_target));

  ImGui_ImplGlfw_NewFrame();
  ImGui::NewFrame();

  ImGui::Begin("Console");
  if (ImGui::Button("Copy image")) {
    wants_final_image = true;
  }
  ImGui::End();

  ImGui::Render();
  frame_target = vuk::extra::ImGui_ImplVuk_Render(
      frame_allocator, std::move(frame_target), imgui_data);

  if (wants_final_image) {
    log->warn("image clipboard is not implemented yet");
  }
  auto entire_thing = vuk::enqueue_presentation(frame_target);
  entire_thing.submit(frame_allocator, context.compiler);
}

void rogue::run() {
  while (!close_requested.load() && !display.should_close()) {
    do_tick();

    display.poll_events();
  }
}

vuk::PipelineBaseCreateInfo
rogue::create_compute_pipeline(std::string_view comp_module,
                               std::string_view comp_entry) {
  const auto comp_bytes = shader_compiler.compile(comp_module, comp_entry);

  auto pipeline_info = vuk::PipelineBaseCreateInfo();
  pipeline_info.add_spirv(comp_bytes, comp_module.data(), comp_entry.data());

  return pipeline_info;
}

vuk::PipelineBaseCreateInfo rogue::create_graphics_pipeline(
    std::string_view vs_module, std::string_view vs_entry,
    std::string_view fs_module, std::string_view fs_entry) {
  const auto vs_bytes = shader_compiler.compile(vs_module, vs_entry);
  const auto fs_bytes = shader_compiler.compile(fs_module, fs_entry);

  auto pipeline_info = vuk::PipelineBaseCreateInfo();
  pipeline_info.add_spirv(vs_bytes, vs_module.data(), vs_entry.data());
  pipeline_info.add_spirv(fs_bytes, fs_module.data(), fs_entry.data());
  return pipeline_info;
}

vuk::PipelineBaseCreateInfo
rogue::create_graphics_pipeline(std::string_view module_name) {
  return create_graphics_pipeline(module_name, "vert_main", module_name,
                                  "frag_main");
}

void rogue::create_pipelines() {
  log->debug("creating pipelines");
  context.runtime.create_named_pipeline("display_rt",
                                        create_graphics_pipeline("test.slang"));
  context.runtime.create_named_pipeline(
      "rt_compute", create_compute_pipeline("raytrace/compute.slang", "main"));
  log->debug("created pipelines");
}

vuk::Value<vuk::ImageAttachment> rogue::get_swap_target() {
  auto imported_swapchain = vuk::acquire_swapchain(context.swapchain);
  auto swapchain_image =
      vuk::acquire_next_image("swp_img", std::move(imported_swapchain));
  return vuk::clear_image(std::move(swapchain_image),
                          vuk::ClearColor{0.0f, 0.0f, 1.0f, 1.0f});
}

rogue::vuk_frame_resources rogue::get_next_frame_resources() {
  auto &frame_resource = context.superframe_resource.get_next_frame();
  context.runtime.next_frame();
  return {
      .allocator = vuk::Allocator(frame_resource),
      .swap_target = get_swap_target(),
  };
}

} // namespace rog