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