#!amber

DEVICE_EXTENSION VK_KHR_acceleration_structure
DEVICE_EXTENSION VK_KHR_ray_tracing_pipeline
DEVICE_EXTENSION VK_KHR_buffer_device_address
DEVICE_EXTENSION VK_KHR_pipeline_library
DEVICE_FEATURE   RayTracingPipelineFeaturesKHR.rayTracingPipeline
DEVICE_FEATURE   AccelerationStructureFeaturesKHR.accelerationStructure
DEVICE_FEATURE   BufferDeviceAddressFeatures.bufferDeviceAddress

SHADER ray_generation raygen1 GLSL
  #version 460 core
  #extension GL_EXT_ray_tracing : require
  layout(location = 0) rayPayloadEXT uint payload;
  layout(set = 0, binding = 0) uniform accelerationStructureEXT topLevelAS;
  layout(set = 0, binding = 1, r32ui) uniform uimage2D result;
  void main()
  {
    for (uint i = 0; i < 3; i++)
    {
      payload = 0;
      // Set origin to trigger miss
      traceRayEXT(
        topLevelAS,           // accelerationStructureEXT topLevel,
        0u,                   // uint                     rayFlags
        0xFFu,                // uint                     cullMask,
        0,                    // uint                     sbtRecordOffset,
        0,                    // uint                     sbtRecordStride,
        i,                    // uint                     missIndex,
        vec3(0.0, 0.0, 0.0),  // vec3                     origin,
        0.0,                  // float                    Tmin,
        vec3(0.0, 0.0, -1.0), // vec3                     direction,
        0.25,                 // float                    Tmax,
        0                     // int                      payload
      );
      imageStore(result, ivec2(i, 0), uvec4(payload, 0u, 0u, 0u));
    }
  }
END

SHADER miss miss1 GLSL
  #version 460 core
  #extension GL_EXT_ray_tracing : require
  layout(location = 0) rayPayloadInEXT uint payload;
  void main()
  {
    payload = 1;
  }
END

SHADER miss miss2 GLSL
  #version 460 core
  #extension GL_EXT_ray_tracing : require
  layout(location = 0) rayPayloadInEXT uint payload;
  void main()
  {
    payload = 2;
  }
END

SHADER miss miss3 GLSL
  #version 460 core
  #extension GL_EXT_ray_tracing : require
  layout(location = 0) rayPayloadInEXT uint payload;
  void main()
  {
    payload = 3;
  }
END

SHADER closest_hit closesthit1 GLSL
  #version 460 core
  #extension GL_EXT_ray_tracing : require
  hitAttributeEXT vec3 attribs;
  layout(location = 0) rayPayloadInEXT uint payload;
  void main()
  {
  }
END

SHADER any_hit anyhit1 GLSL
  #version 460 core
  #extension GL_EXT_ray_tracing : require
  hitAttributeEXT vec3 attribs;
  layout(location = 0) rayPayloadInEXT uint payload;
  void main()
  {
  }
END

ACCELERATION_STRUCTURE BOTTOM_LEVEL triangles
  GEOMETRY TRIANGLES
    -1 -1  1
    -1  1  1
     1  1  1
  END
  GEOMETRY TRIANGLES
     1  1  1
     1 -1  1
    -1 -1  1
  END
END

ACCELERATION_STRUCTURE TOP_LEVEL tlas1
  BOTTOM_LEVEL_INSTANCE triangles
  END
END

IMAGE img1 FORMAT R32_UINT DIM_2D WIDTH 4 HEIGHT 1 FILL 0

PIPELINE raytracing my_rtlib_bottom
  MAX_RAY_PAYLOAD_SIZE 16
  MAX_RAY_HIT_ATTRIBUTE_SIZE 16
  BIND ACCELERATION_STRUCTURE tlas1 DESCRIPTOR_SET 0 BINDING 0
  BIND BUFFER img1 AS storage_image DESCRIPTOR_SET 0 BINDING 1

  FLAGS LIBRARY
  SHADER_GROUP g_miss3 miss3
END

PIPELINE raytracing my_rtlib
  USE_LIBRARY my_rtlib_bottom
  MAX_RAY_PAYLOAD_SIZE 16
  MAX_RAY_HIT_ATTRIBUTE_SIZE 16
  BIND ACCELERATION_STRUCTURE tlas1 DESCRIPTOR_SET 0 BINDING 0
  BIND BUFFER img1 AS storage_image DESCRIPTOR_SET 0 BINDING 1

  FLAGS LIBRARY
  SHADER_GROUP g_miss2 miss2
END

PIPELINE raytracing my_rtpipeline
  USE_LIBRARY my_rtlib
  MAX_RAY_PAYLOAD_SIZE 16
  MAX_RAY_HIT_ATTRIBUTE_SIZE 16
  BIND ACCELERATION_STRUCTURE tlas1 DESCRIPTOR_SET 0 BINDING 0
  BIND BUFFER img1 AS storage_image DESCRIPTOR_SET 0 BINDING 1
  SHADER_GROUP g_rgen raygen1
  SHADER_GROUP g_miss1 miss1
  SHADER_GROUP g_hits anyhit1 closesthit1
  SHADER_BINDING_TABLE sbt_raygen
    g_rgen
  END
  SHADER_BINDING_TABLE sbt_miss
    g_miss1
    g_miss2
    g_miss3
  END
  SHADER_BINDING_TABLE sbt_hit
    g_hits
  END
END

RUN my_rtpipeline RAYGEN sbt_raygen MISS sbt_miss HIT sbt_hit 1 1 1

EXPECT img1 IDX 0 EQ 1 2 3