1// these functions cause additional bindings to be brought in, which
2// fails the shader compilations. We should improve the ability to import these
3// in the future and then use the imports here.
4// import bevy_pbr::render::view_transformations::position_world_to_clip;
5// import bevy_pbr::render::mesh_functions::mesh_position_local_to_world;
6import super::{
7    bindings::{globals, view, VertexOutput, TaskPayload},
8};
9
10enable wgpu_mesh_shader;
11
12fn position_world_to_clip(world_pos: vec3<f32>) -> vec4<f32> {
13    let clip_pos = view.clip_from_world * vec4(world_pos, 1.0);
14    return clip_pos;
15}
16
17fn mesh_position_local_to_world(world_from_local: mat4x4<f32>, vertex_position: vec4<f32>) -> vec4<f32> {
18    return world_from_local * vertex_position;
19}
20
21// This data represents a cube, for simplicity reasons it is hardcoded.
22// half_size is half the size of the cube
23const half_size = vec3(0.1);
24const min = -half_size;
25const max = half_size;
26
27// Suppose Y-up right hand, and camera look from +Z to -Z
28// 192 items
29const vertices = array(
30    // xyz, normal.xyz, uv.xy
31    // Front
32    min.x, min.y, max.z, 0.0, 0.0, 1.0, 0.0, 0.0,
33    max.x, min.y, max.z, 0.0, 0.0, 1.0, 1.0, 0.0,
34    max.x, max.y, max.z, 0.0, 0.0, 1.0, 1.0, 1.0,
35    min.x, max.y, max.z, 0.0, 0.0, 1.0, 0.0, 1.0,
36    // Back
37    min.x, max.y, min.z, 0.0, 0.0, -1.0, 1.0, 0.0,
38    max.x, max.y, min.z, 0.0, 0.0, -1.0, 0.0, 0.0,
39    max.x, min.y, min.z, 0.0, 0.0, -1.0, 0.0, 1.0,
40    min.x, min.y, min.z, 0.0, 0.0, -1.0, 1.0, 1.0,
41    // Right
42    max.x, min.y, min.z, 1.0, 0.0, 0.0, 0.0, 0.0,
43    max.x, max.y, min.z, 1.0, 0.0, 0.0, 1.0, 0.0,
44    max.x, max.y, max.z, 1.0, 0.0, 0.0, 1.0, 1.0,
45    max.x, min.y, max.z, 1.0, 0.0, 0.0, 0.0, 1.0,
46    // Left
47    min.x, min.y, max.z, -1.0, 0.0, 0.0, 1.0, 0.0,
48    min.x, max.y, max.z, -1.0, 0.0, 0.0, 0.0, 0.0,
49    min.x, max.y, min.z, -1.0, 0.0, 0.0, 0.0, 1.0,
50    min.x, min.y, min.z, -1.0, 0.0, 0.0, 1.0, 1.0,
51    // Top
52    max.x, max.y, min.z, 0.0, 1.0, 0.0, 1.0, 0.0,
53    min.x, max.y, min.z, 0.0, 1.0, 0.0, 0.0, 0.0,
54    min.x, max.y, max.z, 0.0, 1.0, 0.0, 0.0, 1.0,
55    max.x, max.y, max.z, 0.0, 1.0, 0.0, 1.0, 1.0,
56    // Bottom
57    max.x, min.y, max.z, 0.0, -1.0, 0.0, 0.0, 0.0,
58    min.x, min.y, max.z, 0.0, -1.0, 0.0, 1.0, 0.0,
59    min.x, min.y, min.z, 0.0, -1.0, 0.0, 1.0, 1.0,
60    max.x, min.y, min.z, 0.0, -1.0, 0.0, 0.0, 1.0
61);
62
63// 36 items
64const indices = array(
65    0, 1, 2, 2, 3, 0, // front
66    4, 5, 6, 6, 7, 4, // back
67    8, 9, 10, 10, 11, 8, // right
68    12, 13, 14, 14, 15, 12, // left
69    16, 17, 18, 18, 19, 16, // top
70    20, 21, 22, 22, 23, 20, // bottom
71);
72
73
74struct PrimitiveOutput {
75    @builtin(triangle_indices) indices: vec3<u32>,
76    @builtin(cull_primitive) cull: bool,
77    @per_primitive @location(3) colorMask: vec4<f32>,
78}
79
80var<task_payload> taskPayload: TaskPayload;
81
82struct MeshOutput {
83    @builtin(vertices) vertices: array<VertexOutput, 24>,
84    @builtin(primitives) primitives: array<PrimitiveOutput, 12>,
85    @builtin(vertex_count) vertex_count: u32,
86    @builtin(primitive_count) primitive_count: u32,
87}
88
89var<workgroup> mesh_output: MeshOutput;
90
91struct MeshInput {
92    @builtin(global_invocation_id) global_invocation_id: vec3u,
93    @builtin(num_workgroups) num_workgroups: vec3u,
94    @builtin(workgroup_id) workgroup_id: vec3u
95}
96
97@mesh(mesh_output)
98@payload(taskPayload)
99@workgroup_size(1)
100fn mesh(mesh_input: MeshInput) {
101    mesh_output.vertex_count = 24;
102    mesh_output.primitive_count = 12;
103    
104    // we "fake" the `world_from_local` matrix that normally exists for
105    // meshes because we're generating and positioning cubes in arbitrary
106    // positions. The global_invocation_id will give us an x/y/z for every
107    // point in the 3d space we're filling.
108    let world_from_local = mat4x4(
109        vec4(1., 0., 0., 0,),
110        vec4(0., 1., 0., 0.),
111        vec4(0., 0., 1., 0.),
112        vec4(
113            f32(mesh_input.global_invocation_id.x),
114            f32(mesh_input.global_invocation_id.y),
115            f32(mesh_input.global_invocation_id.z),
116            1.,
117        ),
118    );
119
120    // Set the data for each vertex, including position, normal, and uv.
121    for (var i: i32 = 0; i < 24; i++) {
122        let offset = i * 8;
123        mesh_output.vertices[i].position =  position_world_to_clip(
124           mesh_position_local_to_world(world_from_local, 
125                vec4(vertices[offset] , vertices[offset + 1], vertices[offset + 2], 1.)
126           ).xyz
127       );
128        mesh_output.vertices[i].normal = vec3(vertices[offset + 3], vertices[offset + 4], vertices[offset + 5]);
129        mesh_output.vertices[i].uv = vec2(vertices[offset + 6], vertices[offset + 7]);
130    };
131
132    // Set primitive data
133
134    // Define a color based on the position in the total set of workgroups
135    let colorMask = vec4<f32>(
136        vec3f(mesh_input.global_invocation_id.xyz) / vec3f(mesh_input.num_workgroups.xyz),
137        1.0
138    );
139
140    for (var i: i32 = 0; i < 12; i++) {
141        let offset = i * 3;
142        // Set the indices for this primitive
143        mesh_output.primitives[i].indices = vec3<u32>(u32(indices[offset]), u32(indices[offset + 1]), u32(indices[offset + 2]));
144        // cull primitives based on the current time and offset by their position
145        mesh_output.primitives[i].cull = sin(globals.time * 10. + f32(mesh_input.global_invocation_id.x) + f32(mesh_input.global_invocation_id.y) + f32(mesh_input.global_invocation_id.z) ) > 0.8;
146        // Apply a color to this primitive, multiplied by the task payload's color
147        mesh_output.primitives[i].colorMask = colorMask * taskPayload.colorMask;
148    };
149}