1#import bevy_sprite::{
2    mesh2d_vertex_output::VertexOutput,
3    mesh2d_view_bindings::view,
4}
5
6#ifdef TONEMAP_IN_SHADER
7#import bevy_core_pipeline::tonemapping
8#endif
9#ifdef SRGB_OUTPUT
10#import bevy_render::color_operations::linear_to_srgb
11#endif
12#ifdef OKLAB_OUTPUT
13#import bevy_render::color_operations::linear_rgb_to_oklab
14#endif
15
16struct ColorMaterial {
17    color: vec4<f32>,
18    uv_transform: mat3x3<f32>,
19    // 'flags' is a bit field indicating various options. u32 is 32 bits so we have up to 32 options.
20    flags: u32,
21    alpha_cutoff: f32,
22};
23
24const COLOR_MATERIAL_FLAGS_TEXTURE_BIT: u32              = 1u;
25const COLOR_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS: u32 = 3221225472u; // (0b11u32 << 30)
26const COLOR_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE: u32        = 0u;          // (0u32 << 30)
27const COLOR_MATERIAL_FLAGS_ALPHA_MODE_MASK: u32          = 1073741824u; // (1u32 << 30)
28const COLOR_MATERIAL_FLAGS_ALPHA_MODE_BLEND: u32         = 2147483648u; // (2u32 << 30)
29
30@group(#{MATERIAL_BIND_GROUP}) @binding(0) var<uniform> material: ColorMaterial;
31@group(#{MATERIAL_BIND_GROUP}) @binding(1) var texture: texture_2d<f32>;
32@group(#{MATERIAL_BIND_GROUP}) @binding(2) var texture_sampler: sampler;
33
34@fragment
35fn fragment(
36    mesh: VertexOutput,
37) -> @location(0) vec4<f32> {
38    var output_color: vec4<f32> = material.color;
39
40#ifdef VERTEX_COLORS
41    output_color = output_color * mesh.color;
42#endif
43
44    let uv = (material.uv_transform * vec3(mesh.uv, 1.0)).xy;
45
46    if ((material.flags & COLOR_MATERIAL_FLAGS_TEXTURE_BIT) != 0u) {
47        output_color = output_color * textureSample(texture, texture_sampler, uv);
48    }
49
50    output_color = alpha_discard(material, output_color);
51
52#ifdef TONEMAP_IN_SHADER
53    output_color = tonemapping::tone_mapping(output_color, view.color_grading);
54#endif
55#ifdef SRGB_OUTPUT
56    output_color = vec4(linear_to_srgb(output_color.rgb), output_color.a);
57#endif
58#ifdef OKLAB_OUTPUT
59    output_color = vec4(linear_rgb_to_oklab(output_color.rgb), output_color.a);
60#endif
61    return output_color;
62}
63
64fn alpha_discard(material: ColorMaterial, output_color: vec4<f32>) -> vec4<f32> {
65    var color = output_color;
66    let alpha_mode = material.flags & COLOR_MATERIAL_FLAGS_ALPHA_MODE_RESERVED_BITS;
67    if alpha_mode == COLOR_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE {
68        // NOTE: If rendering as opaque, alpha should be ignored so set to 1.0
69        color.a = 1.0;
70    }
71#ifdef MAY_DISCARD
72    else if alpha_mode == COLOR_MATERIAL_FLAGS_ALPHA_MODE_MASK {
73       if color.a >= material.alpha_cutoff {
74            // NOTE: If rendering as masked alpha and >= the cutoff, render as fully opaque
75            color.a = 1.0;
76        } else {
77            // NOTE: output_color.a < in.material.alpha_cutoff should not be rendered
78            discard;
79        }
80    }
81#endif // MAY_DISCARD
82
83    return color;
84}