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