1#import bevy_render::view::View
2#import bevy_ui::ui_node::{
3    draw_uinode_background,
4    draw_uinode_border,
5}
6
7#import bevy_render::color_operations::{
8    convert_to_linear_rgba,
9    mix_oklch,
10    mix_oklch_long,
11    mix_hsv,
12    mix_hsv_long,
13    mix_hsl,
14    mix_hsl_long,
15    oklch_to_linear_rgb,
16    hsv_to_linear_rgb,
17    hsl_to_linear_rgb,
18    oklab_to_linear_rgb,
19}
20
21#import bevy_render::maths::PI
22
23const TAU: f32 = 2. * PI;
24
25const RIGHT_VERTEX = 2u;
26const BOTTOM_VERTEX = 4u;
27// must align with BORDER_* shader_flags from bevy_ui/render/mod.rs
28const RADIAL: u32 = 16u;
29const FILL_START: u32 = 32u;
30const FILL_END: u32 = 64u;
31const CONIC: u32 = 128u;
32const BORDER_LEFT: u32 = 256u;
33const BORDER_TOP: u32 = 512u;
34const BORDER_RIGHT: u32 = 1024u;
35const BORDER_BOTTOM: u32 = 2048u;
36const BORDER_ANY: u32 = BORDER_LEFT + BORDER_TOP + BORDER_RIGHT + BORDER_BOTTOM;
37
38fn enabled(flags: u32, mask: u32) -> bool {
39    return (flags & mask) != 0u;
40}
41
42@group(0) @binding(0) var<uniform> view: View;
43
44struct GradientVertexOutput {
45    @location(0) uv: vec2<f32>,
46    @location(1) @interpolate(flat) size: vec2<f32>,
47    @location(2) @interpolate(flat) flags: u32,
48    @location(3) @interpolate(flat) radius: vec4<f32>,
49    @location(4) @interpolate(flat) border: vec4<f32>,    
50
51    // Position relative to the center of the rectangle.
52    @location(5) point: vec2<f32>,
53    @location(6) @interpolate(flat) g_start: vec2<f32>,
54    @location(7) @interpolate(flat) dir: vec2<f32>,
55    @location(8) @interpolate(flat) start_color: vec4<f32>,
56    @location(9) @interpolate(flat) start_len: f32,
57    @location(10) @interpolate(flat) end_len: f32,
58    @location(11) @interpolate(flat) end_color: vec4<f32>,
59    @location(12) @interpolate(flat) hint: f32,
60    @builtin(position) position: vec4<f32>,
61};
62
63@vertex
64fn vertex(
65    @location(0) vertex_position: vec3<f32>,
66    @location(1) vertex_uv: vec2<f32>,
67    @location(2) flags: u32,
68
69    // x: top left, y: top right, z: bottom right, w: bottom left.
70    @location(3) radius: vec4<f32>,
71
72    // x: left, y: top, z: right, w: bottom.
73    @location(4) border: vec4<f32>,
74    @location(5) size: vec2<f32>,
75    @location(6) point: vec2<f32>,
76    @location(7) @interpolate(flat) g_start: vec2<f32>,
77    @location(8) @interpolate(flat) dir: vec2<f32>,
78    @location(9) @interpolate(flat) start_color: vec4<f32>,
79    @location(10) @interpolate(flat) start_len: f32,
80    @location(11) @interpolate(flat) end_len: f32,
81    @location(12) @interpolate(flat) end_color: vec4<f32>,
82    @location(13) @interpolate(flat) hint: f32
83) -> GradientVertexOutput {
84    var out: GradientVertexOutput;
85    out.position = view.clip_from_world * vec4(vertex_position, 1.0);
86    out.uv = vertex_uv;
87    out.size = size;
88    out.flags = flags;
89    out.radius = radius;
90    out.border = border;
91    out.point = point;
92    out.dir = dir;
93    out.start_color = start_color;
94    out.start_len = start_len;
95    out.end_len = end_len;
96    out.end_color = end_color;
97    out.g_start = g_start;
98    out.hint = hint;
99
100    return out;
101}
102
103@fragment
104fn fragment(in: GradientVertexOutput) -> @location(0) vec4<f32> {
105    var g_distance: f32;
106    if enabled(in.flags, RADIAL) {
107        g_distance = radial_distance(in.point, in.g_start, in.dir.x);
108    } else if enabled(in.flags, CONIC) {
109        g_distance = conic_distance(in.dir.x, in.point, in.g_start);
110    } else {
111        g_distance = linear_distance(in.point, in.g_start, in.dir);
112    }
113
114    let gradient_color = interpolate_gradient(
115        g_distance,
116        in.start_color,
117        in.start_len,
118        in.end_color,
119        in.end_len,
120        in.hint,
121        in.flags
122    );
123
124    if enabled(in.flags, BORDER_ANY) {
125        return draw_uinode_border(gradient_color, in.point, in.size, in.radius, in.border, in.flags);
126    } else {
127        return draw_uinode_background(gradient_color, in.point, in.size, in.radius, in.border, in.flags);
128    }
129}
130
131fn rem_euclid(a: f32, b: f32) -> f32 {
132    return ((a % b) + b) % b;
133}
134
135// These functions are used to calculate the distance in gradient space from the start of the gradient to the point.
136// The distance in gradient space is then used to interpolate between the start and end colors.
137
138fn linear_distance(
139    point: vec2<f32>,
140    g_start: vec2<f32>,
141    g_dir: vec2<f32>,
142) -> f32 {
143    return dot(point - g_start, g_dir);
144}
145
146fn radial_distance(
147    point: vec2<f32>,
148    center: vec2<f32>,
149    ratio: f32,
150) -> f32 {
151    let d = point - center;
152    return length(vec2(d.x, d.y * ratio));
153}
154
155fn conic_distance(
156    start: f32,
157    point: vec2<f32>,
158    center: vec2<f32>,
159) -> f32 {
160    let d = point - center;
161    let angle = atan2(-d.x, d.y) + PI;
162    return (((angle - start) % TAU) + TAU) % TAU;
163}
164
165// Mix the colors, choosing the appropriate interpolation method for the given color space
166fn mix_colors(
167    start_color: vec3<f32>,
168    end_color: vec3<f32>,
169    t: f32,
170) -> vec3<f32> {
171#ifdef IN_OKLCH
172    return mix_oklch(start_color, end_color, t);
173#else ifdef IN_OKLCH_LONG
174    return mix_oklch_long(start_color, end_color, t);
175#else ifdef IN_HSV
176    return mix_hsv(start_color, end_color, t);
177#else ifdef IN_HSV_LONG
178    return mix_hsv_long(start_color, end_color, t);
179#else ifdef IN_HSL
180    return mix_hsl(start_color, end_color, t);
181#else ifdef IN_HSL_LONG
182    return mix_hsl_long(start_color, end_color, t);
183#else
184    // Just lerp in linear RGBA, OkLab and SRGBA spaces
185    return mix(start_color, end_color, t);
186#endif
187}
188
189// Convert a color from the interpolation color space to linear rgba
190fn convert_to_linear_rgba(
191    color: vec4<f32>,
192) -> vec4<f32> {
193#ifdef IN_OKLCH
194    let rgb = oklch_to_linear_rgb(color.xyz);
195#else ifdef IN_OKLCH_LONG
196    let rgb = oklch_to_linear_rgb(color.xyz);
197#else ifdef IN_HSV
198    let rgb = hsv_to_linear_rgb(color.xyz);
199#else ifdef IN_HSV_LONG
200    let rgb = hsv_to_linear_rgb(color.xyz);
201#else ifdef IN_HSL
202    let rgb = hsl_to_linear_rgb(color.xyz);
203#else ifdef IN_HSL_LONG
204    let rgb = hsl_to_linear_rgb(color.xyz);
205#else ifdef IN_OKLAB
206    let rgb = oklab_to_linear_rgb(color.xyz);
207#else ifdef IN_SRGB
208    let rgb = pow(color.xyz, vec3(2.2));
209#else
210    // Color is already in linear rgba space
211    let rgb = color.rgb;
212#endif
213    return vec4(rgb, color.a);
214}
215
216fn interpolate_gradient(
217    distance: f32,
218    start_color: vec4<f32>,
219    start_distance: f32,
220    end_color: vec4<f32>,
221    end_distance: f32,
222    hint: f32,
223    flags: u32,
224) -> vec4<f32> {
225    if start_distance == end_distance {
226        if distance <= start_distance && enabled(flags, FILL_START) {
227            return convert_to_linear_rgba(start_color);
228        }
229        if start_distance <= distance && enabled(flags, FILL_END) {
230            return convert_to_linear_rgba(end_color);
231        }
232        return vec4(0.);
233    }
234
235    var t = (distance - start_distance) / (end_distance - start_distance);
236
237    if t < 0.0 {
238        if enabled(flags, FILL_START) {
239            return convert_to_linear_rgba(start_color);
240        }
241        return vec4(0.0);
242    }
243
244    if 1. < t {
245        if enabled(flags, FILL_END) {
246            return convert_to_linear_rgba(end_color);
247        }
248        return vec4(0.0);
249    }
250
251    if t < hint {
252        t = 0.5 * t / hint;
253    } else {
254        t = 0.5 * (1 + (t - hint) / (1.0 - hint));
255    }
256
257    return convert_to_linear_rgba(vec4(mix_colors(start_color.xyz, end_color.xyz, t), mix(start_color.a, end_color.a, t)));
258}