path: root/src/3d/visual_debug.cc

// This file is part of the 64k demo project.
// Implementation of visual debugging tools.

#include "3d/visual_debug.h"

#if !defined(STRIP_ALL)

#include <cstdio>
#include <cstring>

// Simple shader for drawing colored lines
static const char* kDebugShaderCode = R"(
struct Uniforms {
    viewProj : mat4x4<f32>,
}
@group(0) @binding(0) var<uniform> uniforms : Uniforms;

struct VertexInput {
    @location(0) position : vec3<f32>,
    @location(1) color : vec3<f32>,
}

struct VertexOutput {
    @builtin(position) position : vec4<f32>,
    @location(0) color : vec3<f32>,
}

@vertex
fn vs_main(in : VertexInput) -> VertexOutput {
    var out : VertexOutput;
    out.position = uniforms.viewProj * vec4<f32>(in.position, 1.0);
    out.color = in.color;
    return out;
}

@fragment
fn fs_main(in : VertexOutput) -> @location(0) vec4<f32> {
    return vec4<f32>(in.color, 1.0);
}
)";

void VisualDebug::init(WGPUDevice device, WGPUTextureFormat format) {
  device_ = device;
  create_pipeline(format);

  // Initial capacity for vertex buffer (e.g., 1024 lines)
  vertex_buffer_capacity_ =
      1024 * 2 * sizeof(float) * 6; // 2 verts per line, 6 floats per vert

  WGPUBufferDescriptor vb_desc = {};
  vb_desc.usage = WGPUBufferUsage_Vertex | WGPUBufferUsage_CopyDst;
  vb_desc.size = vertex_buffer_capacity_;
  vertex_buffer_ = wgpuDeviceCreateBuffer(device_, &vb_desc);

  WGPUBufferDescriptor ub_desc = {};
  ub_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
  ub_desc.size = sizeof(mat4);
  uniform_buffer_ = wgpuDeviceCreateBuffer(device_, &ub_desc);
}

void VisualDebug::shutdown() {
  if (pipeline_)
    wgpuRenderPipelineRelease(pipeline_);
  if (bind_group_layout_)
    wgpuBindGroupLayoutRelease(bind_group_layout_);
  if (vertex_buffer_)
    wgpuBufferRelease(vertex_buffer_);
  if (uniform_buffer_)
    wgpuBufferRelease(uniform_buffer_);
  if (bind_group_)
    wgpuBindGroupRelease(bind_group_);
}

void VisualDebug::create_pipeline(WGPUTextureFormat format) {
  // Bind Group Layout
  WGPUBindGroupLayoutEntry bgl_entry = {};
  bgl_entry.binding = 0;
  bgl_entry.visibility = WGPUShaderStage_Vertex;
  bgl_entry.buffer.type = WGPUBufferBindingType_Uniform;
  bgl_entry.buffer.minBindingSize = sizeof(mat4);

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 1;
  bgl_desc.entries = &bgl_entry;
  bind_group_layout_ = wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc);

  // Pipeline Layout
  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bind_group_layout_;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(device_, &pl_desc);

  // Shader
#if defined(DEMO_CROSS_COMPILE_WIN32)
  WGPUShaderModuleWGSLDescriptor wgsl_desc = {};
  wgsl_desc.chain.sType = WGPUSType_ShaderModuleWGSLDescriptor;
  wgsl_desc.code = kDebugShaderCode;
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = (const WGPUChainedStruct*)&wgsl_desc.chain;
#else
  WGPUShaderSourceWGSL wgsl_desc = {};
  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_desc.code = {kDebugShaderCode, strlen(kDebugShaderCode)};
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = (const WGPUChainedStruct*)&wgsl_desc.chain;
#endif
  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(device_, &shader_desc);

  // Vertex State
  WGPUVertexAttribute attributes[2];
  // Position
  attributes[0].format = WGPUVertexFormat_Float32x3;
  attributes[0].offset = 0;
  attributes[0].shaderLocation = 0;
  // Color
  attributes[1].format = WGPUVertexFormat_Float32x3;
  attributes[1].offset = sizeof(float) * 3;
  attributes[1].shaderLocation = 1;

  WGPUVertexBufferLayout vertex_layout = {};
  vertex_layout.arrayStride = sizeof(float) * 6;
  vertex_layout.stepMode = WGPUVertexStepMode_Vertex;
  vertex_layout.attributeCount = 2;
  vertex_layout.attributes = attributes;

  // Pipeline
  WGPURenderPipelineDescriptor pipeline_desc = {};
  pipeline_desc.layout = pipeline_layout;
  pipeline_desc.vertex.module = shader_module;
#if defined(DEMO_CROSS_COMPILE_WIN32)
  pipeline_desc.vertex.entryPoint = "vs_main";
#else
  pipeline_desc.vertex.entryPoint = {"vs_main", 7};
#endif
  pipeline_desc.vertex.bufferCount = 1;
  pipeline_desc.vertex.buffers = &vertex_layout;

  WGPUFragmentState fragment_state = {};
  fragment_state.module = shader_module;
#if defined(DEMO_CROSS_COMPILE_WIN32)
  fragment_state.entryPoint = "fs_main";
#else
  fragment_state.entryPoint = {"fs_main", 7};
#endif
  fragment_state.targetCount = 1;

  WGPUColorTargetState color_target = {};
  color_target.format = format;
  color_target.writeMask = WGPUColorWriteMask_All;
  // Enable simple alpha blending if needed, but opaque lines are fine for now
  fragment_state.targets = &color_target;
  pipeline_desc.fragment = &fragment_state;

  pipeline_desc.primitive.topology = WGPUPrimitiveTopology_LineList;
  pipeline_desc.primitive.cullMode = WGPUCullMode_None;
  pipeline_desc.primitive.frontFace = WGPUFrontFace_CCW;

  WGPUDepthStencilState depth_stencil = {};
  depth_stencil.format = WGPUTextureFormat_Depth24Plus;
  depth_stencil.depthWriteEnabled = WGPUOptionalBool_False; // Don't write depth
  depth_stencil.depthCompare =
      WGPUCompareFunction_Less; // But do test against it
  pipeline_desc.depthStencil = &depth_stencil;

  pipeline_desc.multisample.count = 1;
  pipeline_desc.multisample.mask = 0xFFFFFFFF;

  pipeline_ = wgpuDeviceCreateRenderPipeline(device_, &pipeline_desc);

  wgpuPipelineLayoutRelease(pipeline_layout);
  wgpuShaderModuleRelease(shader_module);
}

void VisualDebug::add_box(const mat4& transform, const vec3& local_extent,
                          const vec3& color) {
  float lx = local_extent.x;
  float ly = local_extent.y;
  float lz = local_extent.z;

  // 8 corners of transformed box
  vec4 p[] = {
      transform * vec4(-lx, -ly, -lz, 1), transform * vec4(lx, -ly, -lz, 1),
      transform * vec4(lx, ly, -lz, 1),   transform * vec4(-lx, ly, -lz, 1),
      transform * vec4(-lx, -ly, lz, 1),  transform * vec4(lx, -ly, lz, 1),
      transform * vec4(lx, ly, lz, 1),    transform * vec4(-lx, ly, lz, 1)};

  // 12 edges (each 2 vertices)
  DebugLine edges[] = {
      {p[0].xyz(), p[1].xyz(), color},
      {p[1].xyz(), p[2].xyz(), color},
      {p[2].xyz(), p[3].xyz(), color},
      {p[3].xyz(), p[0].xyz(), color}, // Front face
      {p[4].xyz(), p[5].xyz(), color},
      {p[5].xyz(), p[6].xyz(), color},
      {p[6].xyz(), p[7].xyz(), color},
      {p[7].xyz(), p[4].xyz(), color}, // Back face
      {p[0].xyz(), p[4].xyz(), color},
      {p[1].xyz(), p[5].xyz(), color},
      {p[2].xyz(), p[6].xyz(), color},
      {p[3].xyz(), p[7].xyz(), color} // Connecting edges
  };

  for (const auto& l : edges) {
    lines_.push_back(l);
  }
}

void VisualDebug::update_buffers(const mat4& view_proj) {
  // Update Uniforms
  wgpuQueueWriteBuffer(wgpuDeviceGetQueue(device_), uniform_buffer_, 0,
                       &view_proj, sizeof(mat4));

  // Update Vertices
  size_t required_size = lines_.size() * 2 * sizeof(float) * 6;
  if (required_size > vertex_buffer_capacity_) {
    // Resize buffer
    wgpuBufferRelease(vertex_buffer_);
    vertex_buffer_capacity_ = required_size * 2; // Double capacity
    WGPUBufferDescriptor vb_desc = {};
    vb_desc.usage = WGPUBufferUsage_Vertex | WGPUBufferUsage_CopyDst;
    vb_desc.size = vertex_buffer_capacity_;
    vertex_buffer_ = wgpuDeviceCreateBuffer(device_, &vb_desc);
  }

  if (required_size > 0) {
    std::vector<float> vertex_data;
    vertex_data.reserve(lines_.size() * 12); // 2 verts * 6 floats
    for (const auto& line : lines_) {
      vertex_data.push_back(line.start.x);
      vertex_data.push_back(line.start.y);
      vertex_data.push_back(line.start.z);
      vertex_data.push_back(line.color.x);
      vertex_data.push_back(line.color.y);
      vertex_data.push_back(line.color.z);

      vertex_data.push_back(line.end.x);
      vertex_data.push_back(line.end.y);
      vertex_data.push_back(line.end.z);
      vertex_data.push_back(line.color.x);
      vertex_data.push_back(line.color.y);
      vertex_data.push_back(line.color.z);
    }
    wgpuQueueWriteBuffer(wgpuDeviceGetQueue(device_), vertex_buffer_, 0,
                         vertex_data.data(),
                         vertex_data.size() * sizeof(float));
  }

  // Re-create bind group if needed (e.g. if uniform buffer changed, though here
  // it's static)
  if (!bind_group_) {
    WGPUBindGroupEntry bg_entry = {};
    bg_entry.binding = 0;
    bg_entry.buffer = uniform_buffer_;
    bg_entry.size = sizeof(mat4);

    WGPUBindGroupDescriptor bg_desc = {};
    bg_desc.layout = bind_group_layout_;
    bg_desc.entryCount = 1;
    bg_desc.entries = &bg_entry;
    bind_group_ = wgpuDeviceCreateBindGroup(device_, &bg_desc);
  }
}

void VisualDebug::render(WGPURenderPassEncoder pass, const mat4& view_proj) {
  if (lines_.empty())
    return;

  update_buffers(view_proj);

  wgpuRenderPassEncoderSetPipeline(pass, pipeline_);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);
  wgpuRenderPassEncoderSetVertexBuffer(pass, 0, vertex_buffer_, 0,
                                       lines_.size() * 2 * sizeof(float) * 6);
  wgpuRenderPassEncoderDraw(pass, (uint32_t)lines_.size() * 2, 1, 0, 0);

  lines_.clear(); // Clear for next frame
}

#endif // !defined(STRIP_ALL)