1#import bevy_render::view::View
2#import bevy_render::maths
3#import bevy_pbr::mesh_types::{SkinnedMesh, MorphAttributes, MorphDescriptor, MorphWeights}
4#import bevy_pbr::skinning::joint_matrices
5#import bevy_mod_outline::common::{OutlineViewUniform, VertexOutput}
6
7struct Instance {
8    world_from_local: mat3x4<f32>,
9    world_plane_origin: vec3<f32>,
10    world_plane_offset: vec3<f32>,
11    volume_colour: vec4<f32>,
12    volume_offset: f32,
13    stencil_offset: f32,
14    alpha_mask_threshold: f32,
15    first_vertex_index: u32,
16    current_skin_index: u32,
17    current_morph_index: u32,
18};
19
20struct Vertex {
21    @location(0) position: vec3<f32>,
22    @builtin(instance_index) instance_index: u32,
23#ifndef VERTEX_OFFSET_ZERO
24    @location(1) outline_normal: vec3<f32>,
25#endif
26#ifdef ALPHA_MASK_TEXTURE
27    @location(2) uv: vec2<f32>,
28#endif
29#ifdef SKINNED
30    @location(5) joint_indices: vec4<u32>,
31    @location(6) joint_weights: vec4<f32>,
32#endif
33#ifdef MORPH_TARGETS
34    @builtin(vertex_index) index: u32,
35#endif
36};
37
38@group(0) @binding(0)
39var<uniform> view_uniform: OutlineViewUniform;
40
41#import bevy_pbr::skinning
42
43#ifdef INSTANCE_BATCH_SIZE
44@group(1) @binding(0) var<uniform> mesh: array<Instance, #{INSTANCE_BATCH_SIZE}u>;
45#else
46@group(1) @binding(0) var<storage> mesh: array<Instance>;
47#endif
48
49#ifdef MORPH_TARGETS
50#ifdef SKINS_USE_UNIFORM_BUFFERS
51@group(2) @binding(2) var<uniform> morph_weights: MorphWeights;
52@group(2) @binding(3) var morph_targets: texture_3d<f32>;
53#else
54@group(2) @binding(2) var<storage> morph_weights: array<f32>;
55@group(2) @binding(3) var<storage> morph_targets: array<MorphAttributes>;
56@group(2) @binding(8) var<storage> morph_descriptors: array<MorphDescriptor>;
57#endif
58
59fn morph_layer_count(descriptor_index: u32) -> u32 {
60#ifdef SKINS_USE_UNIFORM_BUFFERS
61    let dimensions = textureDimensions(morph_targets);
62    return u32(dimensions.z);
63#else
64    return morph_descriptors[descriptor_index].weight_count;
65#endif
66}
67
68fn morph_weight_at(weight_index: u32, descriptor_index: u32) -> f32 {
69#ifdef SKINS_USE_UNIFORM_BUFFERS
70    let i = weight_index;
71    return morph_weights.weights[i / 4u][i % 4u];
72#else
73    let weights_offset = morph_descriptors[descriptor_index].current_weights_offset;
74    return morph_weights[weights_offset + weight_index];
75#endif
76}
77
78#ifdef SKINS_USE_UNIFORM_BUFFERS
79const morph_position_offset: u32 = 0u;
80const morph_total_component_count: u32 = 9u;
81
82fn morph_component_texture_coord(vertex_index: u32, component_offset: u32) -> vec2<u32> {
83    let width = u32(textureDimensions(morph_targets).x);
84    let component_index = morph_total_component_count * vertex_index + component_offset;
85    return vec2<u32>(component_index % width, component_index / width);
86}
87
88fn morph_pixel(vertex: u32, component: u32, weight: u32) -> f32 {
89    let coord = morph_component_texture_coord(vertex, component);
90    return textureLoad(morph_targets, vec3(coord, weight), 0).r;
91}
92
93fn morph(vertex_index: u32, component_offset: u32, weight_index: u32) -> vec3<f32> {
94    return vec3<f32>(
95        morph_pixel(vertex_index, component_offset, weight_index),
96        morph_pixel(vertex_index, component_offset + 1u, weight_index),
97        morph_pixel(vertex_index, component_offset + 2u, weight_index),
98    );
99}
100
101fn morph_position(vertex_index: u32, weight_index: u32, instance_index: u32) -> vec3<f32> {
102    return morph(vertex_index, morph_position_offset, weight_index);
103}
104#else
105fn get_morph_target(vertex_index: u32, weight_index: u32, descriptor_index: u32) -> MorphAttributes {
106    let targets_offset = morph_descriptors[descriptor_index].targets_offset;
107    let vertex_count = morph_descriptors[descriptor_index].vertex_count;
108    return morph_targets[targets_offset + weight_index * vertex_count + vertex_index];
109}
110
111fn morph_position(vertex_index: u32, weight_index: u32, descriptor_index: u32) -> vec3<f32> {
112    return get_morph_target(vertex_index, weight_index, descriptor_index).position;
113}
114#endif
115
116fn morph_vertex(vertex_in: Vertex, instance_index: u32) -> Vertex {
117    var vertex = vertex_in;
118    let first_vertex = mesh[instance_index].first_vertex_index;
119    var morph_index = mesh[instance_index].current_morph_index;
120    let vertex_index = vertex.index - first_vertex;
121
122    let weight_count = morph_layer_count(morph_index);
123    for (var i: u32 = 0u; i < weight_count; i ++) {
124        let weight = morph_weight_at(i, morph_index);
125        if weight == 0.0 {
126            continue;
127        }
128        vertex.position += weight * morph_position(vertex_index, i, morph_index);
129    }
130    return vertex;
131}
132#endif
133
134#ifdef SKINNED
135fn skin_model(
136    indexes: vec4<u32>,
137    weights: vec4<f32>,
138    instance_index: u32,
139) -> mat4x4<f32> {
140#ifdef SKINS_USE_UNIFORM_BUFFERS
141    return weights.x * joint_matrices.data[indexes.x]
142        + weights.y * joint_matrices.data[indexes.y]
143        + weights.z * joint_matrices.data[indexes.z]
144        + weights.w * joint_matrices.data[indexes.w];
145#else
146    var skin_index = mesh[instance_index].current_skin_index;
147    return weights.x * joint_matrices[skin_index + indexes.x]
148        + weights.y * joint_matrices[skin_index + indexes.y]
149        + weights.z * joint_matrices[skin_index + indexes.z]
150        + weights.w * joint_matrices[skin_index + indexes.w];
151#endif
152}
153#endif
154
155fn mat4to3(m: mat4x4<f32>) -> mat3x3<f32> {
156    return mat3x3<f32>(
157        m[0].xyz, m[1].xyz, m[2].xyz
158    );
159}
160
161fn model_origin_z(plane: vec3<f32>, view_proj: mat4x4<f32>) -> f32 {
162    var proj_zw = mat4x2<f32>(
163        view_proj[0].zw, view_proj[1].zw,
164        view_proj[2].zw, view_proj[3].zw);
165    var zw = proj_zw * vec4<f32>(plane, 1.0);
166    return zw.x / zw.y;
167}
168
169@vertex
170fn vertex(vertex_no_morph: Vertex) -> VertexOutput {
171    let iid = vertex_no_morph.instance_index;
172#ifdef MORPH_TARGETS
173    var vertex = morph_vertex(vertex_no_morph, iid);
174#else
175    var vertex = vertex_no_morph;
176#endif
177#ifdef SKINNED
178    let model = skin_model(vertex.joint_indices, vertex.joint_weights, iid);
179#else
180    let model = bevy_render::maths::affine3_to_square(mesh[iid].world_from_local);
181#endif
182    let clip_pos = view_uniform.clip_from_world * (model * vec4<f32>(vertex.position, 1.0));
183#ifdef VOLUME
184    let offset = mesh[iid].volume_offset;
185#else
186    let offset = mesh[iid].stencil_offset;
187#endif
188#ifdef VERTEX_OFFSET_ZERO
189    let out_xy = clip_pos.xy;
190#else
191    let clip_norm = mat4to3(view_uniform.clip_from_world) * (mat4to3(model) * vertex.outline_normal);
192    let corrected_norm = normalize(clip_norm.xy * vec2<f32>(view_uniform.aspect, 1.0));
193    let ndc_delta = offset * corrected_norm * view_uniform.scale_clip_from_logical * clip_pos.w;
194    let out_xy = clip_pos.xy + ndc_delta;
195#endif
196    var out: VertexOutput;
197    out.position = vec4<f32>(out_xy, clip_pos.zw);
198#ifdef FLAT_DEPTH
199#ifdef PLANE_OFFSET_ZERO
200    let world_plane = mesh[iid].world_plane_origin;
201#else
202    let world_from_view = bevy_render::maths::mat2x4_f32_to_mat3x3_unpack(view_uniform.world_from_view_a, view_uniform.world_from_view_b);
203    let model_eye = normalize(world_from_view * vec3<f32>(0.0, 0.0, -1.0));
204    let world_plane = mesh[iid].world_plane_origin + model_eye * mesh[iid].world_plane_offset;
205#endif
206    out.flat_depth = model_origin_z(world_plane, view_uniform.clip_from_world);
207#endif
208#ifdef VOLUME
209    out.volume_colour = mesh[iid].volume_colour;
210#endif
211#ifdef ALPHA_MASK_TEXTURE
212    out.alpha_mask_threshold = mesh[iid].alpha_mask_threshold;
213    out.uv = vertex.uv;
214#endif
215    return out;
216}