webgpu-infinite-world-grid.cc

struct WorldGrid {
  static const size_t vertexDataStride = 6;

  wgpu::Buffer            _indexBuffer;
  wgpu::BindGroup         _bindGroup;
  wgpu::DepthStencilState _depthStencil;
  wgpu::RenderPipeline    _pipeline;

  WorldGrid() {
    _depthStencil.format = wgpu::TextureFormat::Depth24PlusStencil8;
    _depthStencil.depthCompare = wgpu::CompareFunction::Less;
    _depthStencil.depthWriteEnabled = true;
  }

  void init(
    const wgpu::Device& device,
    const wgpu::Buffer& cameraBuffer,
    const FramebufferInfo& framebuffer)
  {
    // based on http://asliceofrendering.com/scene%20helper/2020/01/05/InfiniteGrid/

    // shaders
    wgpu::ShaderModule vsModule = pbwgpu::createShaderMod(device, R"(
      type mat4 = mat4x4<f32>;

      fn inverse(m :mat4) -> mat4 {
        // Note: wgsl does not have an inverse() (matrix inverse) function built in.
        // Source adapted from https://github.com/glslify/glsl-inverse/blob/master/index.glsl
        let a00 = m[0][0];
        let a01 = m[0][1];
        let a02 = m[0][2];
        let a03 = m[0][3];

        let a10 = m[1][0];
        let a11 = m[1][1];
        let a12 = m[1][2];
        let a13 = m[1][3];

        let a20 = m[2][0];
        let a21 = m[2][1];
        let a22 = m[2][2];
        let a23 = m[2][3];

        let a30 = m[3][0];
        let a31 = m[3][1];
        let a32 = m[3][2];
        let a33 = m[3][3];

        let b00 = a00 * a11 - a01 * a10;
        let b01 = a00 * a12 - a02 * a10;
        let b02 = a00 * a13 - a03 * a10;
        let b03 = a01 * a12 - a02 * a11;
        let b04 = a01 * a13 - a03 * a11;
        let b05 = a02 * a13 - a03 * a12;
        let b06 = a20 * a31 - a21 * a30;
        let b07 = a20 * a32 - a22 * a30;
        let b08 = a20 * a33 - a23 * a30;
        let b09 = a21 * a32 - a22 * a31;
        let b10 = a21 * a33 - a23 * a31;
        let b11 = a22 * a33 - a23 * a32;

        let det = b00 * b11 - b01 * b10 + b02 * b09 + b03 * b08 - b04 * b07 + b05 * b06;

        return mat4(
          vec4<f32>(
            a11 * b11 - a12 * b10 + a13 * b09,
            a02 * b10 - a01 * b11 - a03 * b09,
            a31 * b05 - a32 * b04 + a33 * b03,
            a22 * b04 - a21 * b05 - a23 * b03),
          vec4<f32>(
            a12 * b08 - a10 * b11 - a13 * b07,
            a00 * b11 - a02 * b08 + a03 * b07,
            a32 * b02 - a30 * b05 - a33 * b01,
            a20 * b05 - a22 * b02 + a23 * b01),
          vec4<f32>(
            a10 * b10 - a11 * b08 + a13 * b06,
            a01 * b08 - a00 * b10 - a03 * b06,
            a30 * b04 - a31 * b02 + a33 * b00,
            a21 * b02 - a20 * b04 - a23 * b00),
          vec4<f32>(
            a11 * b07 - a10 * b09 - a12 * b06,
            a00 * b09 - a01 * b07 + a02 * b06,
            a31 * b01 - a30 * b03 - a32 * b00,
            a20 * b03 - a21 * b01 + a22 * b00)
        ) * (1.0 / det);
      }


      [[block]] struct Camera {
        view : mat4x4<f32>;
        proj : mat4x4<f32>;
      };
      [[group(0), binding(0)]] var<uniform> camera : Camera;

      struct VertexOut {
        [[builtin(position)]] position : vec4<f32>;
        [[location(1)]] nearPoint : vec3<f32>;
        [[location(2)]] farPoint : vec3<f32>;
      };

      // TODO: figure out how to make this a module constant.
      // See https://www.w3.org/TR/WGSL/#module-constants
      // When doing the below, and then indexing with an input, we get this error:
      //   error: index must be signed or unsigned integer literal
      //           let p = gridPlane[vertexIndex];
      //                             ^^^^^^^^^^^
      // // Grid position are in xy clipped space
      // let gridPlane = array<vec3<f32>,6>(
      //   vec3<f32>(1.0, 1.0, 0.0), vec3<f32>(-1.0, -1.0, 0.0), vec3<f32>(-1.0, 1.0, 0.0),
      //   vec3<f32>(-1.0, -1.0, 0.0), vec3<f32>(1.0, 1.0, 0.0), vec3<f32>(1.0, -1.0, 0.0)
      // );

      fn UnprojectPoint(x :f32, y :f32, z :f32, view :mat4, proj :mat4) -> vec3<f32> {
        let viewInv = inverse(view); // expensive -- could do this on cpu and pass as uniform
        let projInv = inverse(proj);
        let unprojectedPoint :vec4<f32> = viewInv * projInv * vec4<f32>(x, y, z, 1.0);
        return unprojectedPoint.xyz / unprojectedPoint.w;
      }

      [[stage(vertex)]] fn main(
        [[builtin(vertex_index)]] vertexIndex : u32) -> VertexOut
      {
        // Grid position
        var gridPlane : array<vec3<f32>,6> = array<vec3<f32>,6>(
          // in xy clipped space
          vec3<f32>(1.0, 1.0, 0.0), vec3<f32>(-1.0, -1.0, 0.0), vec3<f32>(-1.0, 1.0, 0.0),
          vec3<f32>(-1.0, -1.0, 0.0), vec3<f32>(1.0, 1.0, 0.0), vec3<f32>(1.0, -1.0, 0.0)
        );
        var output : VertexOut;

        let p = gridPlane[vertexIndex].xyz;

        // position directly at the clipped coordinates
        output.position = vec4<f32>(p, 1.0);

        // unprojecting on the near plane and far plane
        output.nearPoint = UnprojectPoint(p.x, p.y, 0.0, camera.view, camera.proj);
        output.farPoint = UnprojectPoint(p.x, p.y, 1.0, camera.view, camera.proj);

        return output;
      })");

    wgpu::ShaderModule fsModule = pbwgpu::createShaderMod(device, R"(
      [[block]] struct Camera {
        view : mat4x4<f32>;
        proj : mat4x4<f32>;
      };
      [[group(0), binding(0)]] var<uniform> camera : Camera;

      fn grid(fragPos3D :vec3<f32>, scale :f32) -> vec4<f32> {
        let thickness = 1.0;
        let opacity = 0.3;

        // use the scale variable to set the distance between the lines
        let coord :vec2<f32> = fragPos3D.xy * scale;

        let derivative :vec2<f32> = fwidth(coord);
        let grid :vec2<f32> = abs(fract(coord - 0.5) - 0.5) / derivative;
        let line :f32 = min(grid.x, grid.y);
        let minimumz :f32 = min(derivative.y, 1.0);
        let minimumx :f32 = min(derivative.x, 1.0);
        let axisLineThreshold = thickness / scale;

        var color :vec4<f32> = vec4<f32>(1.0, 1.0, 1.0, thickness - min(line, thickness));

        // x axis
        if (fragPos3D.y > -axisLineThreshold * minimumz &&
            fragPos3D.y < axisLineThreshold * minimumz)
        {
          color.r = 1.0;
          color.g = 0.35;
          color.b = 0.3;
        } elseif (fragPos3D.x > -axisLineThreshold * minimumx &&
                  fragPos3D.x < axisLineThreshold * minimumx)
        {
          // y axis
          color.r = 0.1;
          color.g = 1.0;
          color.b = 0.3;
        } else {
          color.a = color.a * opacity;
        }
        return color;
      }

      fn computeDepth(pos :vec3<f32>) -> f32 {
        let clip_space_pos :vec4<f32> = camera.proj * camera.view * vec4<f32>(pos.xyz, 1.0);
        return (clip_space_pos.z / clip_space_pos.w);
      }

      let near :f32 = 1.0; // FIXME pass as uniform
      let far  :f32 = 1000.0; // FIXME pass as uniform

      fn computeLinearDepth(pos :vec3<f32>) -> f32 {
        let clip_space_pos :vec4<f32> = camera.proj * camera.view * vec4<f32>(pos.xyz, 1.0);
        let clip_space_depth :f32 = (clip_space_pos.z / clip_space_pos.w) * 2.0 - 1.0; // put back between -1 and 1
        let linearDepth :f32 = (2.0 * near * far) / (far + near - clip_space_depth * (far - near)); // get linear value between 0.01 and 100
        return linearDepth / far;  // normalize
      }

      struct FragmentOut {
        [[location(0)]]         color     : vec4<f32>;
        [[builtin(frag_depth)]] fragDepth : f32;
      };

      [[stage(fragment)]] fn main(
        // [[builtin(position)]] fragCoord : vec4<f32>,
        [[location(1)]] nearPoint : vec3<f32>,
        [[location(2)]] farPoint : vec3<f32>
      ) -> FragmentOut
      {
        var output : FragmentOut;

        let t :f32 = -nearPoint.z / (farPoint.z - nearPoint.z);
        var tf = 0.3; // opacity = 1 when t > 0, opacity = 0 otherwise
        if (t > 0.0) {
          tf = 0.7;
        }

        let fragPos3D :vec3<f32> = nearPoint + t * (farPoint - nearPoint);

        // avoid drawing over other things
        output.fragDepth = computeDepth(fragPos3D);

        // fade out when the grid gets close to "near" and "far" limits
        let linearDepth :f32 = computeLinearDepth(fragPos3D);
        let fading :f32 = max(0.0, (0.5 - linearDepth));

        // grid lines
        output.color = (
          (
            grid(fragPos3D, 0.05)  // major lines
          + grid(fragPos3D, 0.5) // minor lines
          )
          * tf * fading
        );

        // DEBUG
        // output.color = vec4<f32>(1.0, 0.5, 0.1, tf);
        // output.color = vec4<f32>(1.0, 0.5, 0.1, 0.5);
        // output.color = vec4<f32>(mix(vec3<f32>(1.0, 1.0, 1.0), fragPos3D, vec3<f32>(0.8)), 0.5);

        return output;
      })");

    static const u32 indexData[6 * 6] = {0, 1, 2, 3, 4, 5};
    _indexBuffer = pbwgpu::createBuffer(
      device, indexData, sizeof(indexData), wgpu::BufferUsage::Index);

    // shader bind group layout (uniforms)
    wgpu::BindGroupLayout bgl = pbwgpu::createBindGroupLayout(
      device,
      std::vector<wgpu::BindGroupLayoutEntry>{
        wgpu::BindGroupLayoutEntry{ // camera
          .binding = 0,
          .visibility = wgpu::ShaderStage::Vertex | wgpu::ShaderStage::Fragment,
          .buffer = {
            .type = wgpu::BufferBindingType::Uniform,
            .hasDynamicOffset = false,
            .minBindingSize = 0,
          },
        },
      }
    );

    // bind group
    _bindGroup = pbwgpu::createBindGroup(device, bgl, std::vector<wgpu::BindGroupEntry>{
      // [[group(0), binding(0)]]
      { .binding = 0, .buffer = cameraBuffer, .size = sizeof(CameraData) },
    });

    // pipeline
    {
      wgpu::PipelineLayout pipelineLayout = pbwgpu::createBasicPipelineLayout(device, &bgl);
      wgpu::BlendState blend{
        // enable alpha blending:
        .color = {
          .operation = wgpu::BlendOperation::Add,
          .srcFactor = wgpu::BlendFactor::SrcAlpha,
          .dstFactor = wgpu::BlendFactor::OneMinusSrcAlpha,
        },
        // .alpha = {
        //   .operation = wgpu::BlendOperation::Add,
        //   .srcFactor = wgpu::BlendFactor::SrcAlpha,
        //   .dstFactor = wgpu::BlendFactor::OneMinusSrcAlpha,
        // },
      };
      wgpu::ColorTargetState fragTarget{
        .format = framebuffer.textureFormat,
        .blend = &blend,
      };
      wgpu::FragmentState fragState{
        .targetCount = 1,
        .targets = &fragTarget,
        .entryPoint = "main",
        .module = fsModule,
      };
      wgpu::RenderPipelineDescriptor pipelineDesc{
        .layout = pipelineLayout,
        .depthStencil = &_depthStencil,
        // .primitive = { .topology = wgpu::PrimitiveTopology::LineList, },
        .vertex = {
          .entryPoint = "main",
          .bufferCount = 0,
          .module = vsModule,
        },
        .fragment = &fragState,
      };
      _pipeline = device.CreateRenderPipeline(&pipelineDesc);
    }
  }

  void render(wgpu::Queue& queue, wgpu::RenderPassEncoder& pass) {
    pass.SetPipeline(_pipeline);
    pass.SetBindGroup(0, _bindGroup);
    pass.SetIndexBuffer(_indexBuffer, wgpu::IndexFormat::Uint32);
    pass.DrawIndexed(6);
  }
}; // WorldGrid