GPU-accelerated UI toolkit (Vulkan)
git clone https://git.lucas.co/cce-ui.git
src/draw/scene3d_lit.wgsl (5.8K)
1 // A lit, textured mesh in the 3D scene pass: `LitDraw` (draw/lit.rs). The
2 // uniform block's head (mvp, window size, corner radius and shape) is
3 // scene3d.wgsl's, so the window-corner cut below is the same curve; the
4 // rest is the draw's material and the scene's light. Group 1 is a user
5 // image's descriptor set, as `SceneImage` binds it: the base-colour texture,
6 // or a 1x1 white one for an untextured draw (a set must be bound either way).
7 //
8 // Shading is the metallic-roughness model at its plainest: Lambert diffuse
9 // and a GGX / Smith / Schlick specular for each of two directional lights,
10 // plus a sky/ground hemisphere that is both the diffuse ambient and what a
11 // metal sees reflected. The diffuse term is albedo x light x N.L with no
12 // 1/pi, the scale the host's baked lighting has always used, and the
13 // specular is scaled to match, so a matte lit draw looks like a baked one.
14 struct Uniforms {
15 mvp: mat4x4<f32>,
16 window_size: vec2<f32>,
17 window_radius: f32,
18 corner_shape: f32,
19 eye: vec4<f32>,
20 // rgb the base colour, a the draw's opacity.
21 base: vec4<f32>,
22 // metallic, roughness, textured (0/1), unused.
23 surface: vec4<f32>,
24 key_toward: vec4<f32>,
25 key_color: vec4<f32>,
26 fill_toward: vec4<f32>,
27 fill_color: vec4<f32>,
28 sky: vec4<f32>,
29 ground: vec4<f32>,
30 }
31
32 @group(0) @binding(0) var<uniform> u: Uniforms;
33 @group(1) @binding(0) var t_base: texture_2d<f32>;
34 @group(1) @binding(1) var s_base: sampler;
35
36 const PI: f32 = 3.14159265;
37
38 // scene3d.wgsl's window_corner_distance, kept in lockstep with it.
39 fn window_corner_distance(pos: vec2<f32>) -> f32 {
40 let w = u.window_size.x;
41 let h = u.window_size.y;
42 let r = u.window_radius;
43
44 if (pos.x < 0.0 || pos.x > w || pos.y < 0.0 || pos.y > h) {
45 return 1e5;
46 }
47 if (r <= 0.0) {
48 return -1e5;
49 }
50 let q = abs(pos - vec2f(w * 0.5, h * 0.5)) - vec2f(w * 0.5 - r, h * 0.5 - r);
51 if (q.x > 0.0 && q.y > 0.0) {
52 let shape = u.corner_shape;
53 if (shape > 2.001) {
54 let lp = max(pow(pow(q.x, shape) + pow(q.y, shape), 1.0 / shape), 1e-4);
55 let g = vec2f(pow(q.x / lp, shape - 1.0), pow(q.y / lp, shape - 1.0));
56 return (lp - r) / max(length(g), 1e-4);
57 }
58 return length(q) - r;
59 }
60 return -1e5;
61 }
62
63 struct VertexOutput {
64 @builtin(position) position: vec4f,
65 @location(0) world: vec3f,
66 @location(1) normal: vec3f,
67 @location(2) uv: vec2f,
68 @location(3) color: vec3f,
69 };
70
71 @vertex
72 fn vs_main(
73 @location(0) position: vec3f,
74 @location(1) normal: vec3f,
75 @location(2) uv: vec2f,
76 @location(3) color: vec3f,
77 ) -> VertexOutput {
78 var out: VertexOutput;
79 out.position = u.mvp * vec4f(position, 1.0);
80 out.world = position;
81 out.normal = normal;
82 out.uv = uv;
83 out.color = color;
84 return out;
85 }
86
87 // The GGX distribution, for alpha = roughness^2.
88 fn ggx(n_h: f32, a: f32) -> f32 {
89 let a2 = a * a;
90 let d = n_h * n_h * (a2 - 1.0) + 1.0;
91 return a2 / (PI * d * d);
92 }
93
94 // Smith-Schlick geometry term for both directions.
95 fn smith(n_v: f32, n_l: f32, roughness: f32) -> f32 {
96 let k = (roughness + 1.0) * (roughness + 1.0) / 8.0;
97 return (n_v / (n_v * (1.0 - k) + k)) * (n_l / (n_l * (1.0 - k) + k));
98 }
99
100 fn schlick(f0: vec3f, cos_theta: f32) -> vec3f {
101 return f0 + (vec3f(1.0) - f0) * pow(1.0 - cos_theta, 5.0);
102 }
103
104 // One directional light's contribution.
105 fn shade(n: vec3f, v: vec3f, l: vec3f, color: vec3f, diffuse: vec3f, f0: vec3f, roughness: f32) -> vec3f {
106 let n_l = dot(n, l);
107 if (n_l <= 0.0) {
108 return vec3f(0.0);
109 }
110 let h = normalize(l + v);
111 let n_v = max(dot(n, v), 1e-4);
112 let n_h = max(dot(n, h), 0.0);
113 let v_h = max(dot(v, h), 0.0);
114 let a = roughness * roughness;
115 let spec = ggx(n_h, a) * smith(n_v, n_l, roughness) * schlick(f0, v_h) / (4.0 * n_v * n_l + 1e-4);
116 // x pi: the diffuse carries no 1/pi, so the specular is scaled alike.
117 return color * n_l * (diffuse + spec * PI);
118 }
119
120 @fragment
121 fn fs_main(in: VertexOutput, @builtin(front_facing) front: bool) -> @location(0) vec4f {
122 // The sampler clamps, so the coordinates are wrapped here (glTF repeats
123 // by default), with the UNWRAPPED coordinates' gradients: fract jumps
124 // where the wrap falls, and its own gradients there would pick the
125 // smallest mip and draw a seam. Sampled ahead of any discard: implicit
126 // derivatives are undefined after non-uniform control flow.
127 let gx = dpdx(in.uv);
128 let gy = dpdy(in.uv);
129 let texel = textureSampleGrad(t_base, s_base, fract(in.uv), gx, gy);
130 let textured = u.surface.z > 0.5;
131
132 var albedo = u.base.rgb * in.color;
133 var alpha = u.base.a;
134 if (textured) {
135 albedo = albedo * texel.rgb;
136 alpha = alpha * texel.a;
137 }
138 let cov = 1.0 - smoothstep(-0.5, 0.5, window_corner_distance(in.position.xy));
139 alpha = alpha * cov;
140 if (alpha <= 0.002) {
141 discard;
142 }
143
144 // Both sides lit: a face seen from behind is shaded by its back.
145 var n = normalize(in.normal);
146 if (!front) {
147 n = -n;
148 }
149 let v = normalize(u.eye.xyz - in.world);
150 let metallic = u.surface.x;
151 let roughness = u.surface.y;
152 let diffuse = albedo * (1.0 - metallic);
153 let f0 = mix(vec3f(0.04), albedo, metallic);
154
155 let ambient = mix(u.ground.rgb, u.sky.rgb, 0.5 + 0.5 * n.y);
156 let r = reflect(-v, n);
157 let seen = mix(u.ground.rgb, u.sky.rgb, 0.5 + 0.5 * r.y);
158 let n_v = max(dot(n, v), 1e-4);
159 // Schlick with roughness: a rough surface reflects less at grazing.
160 let fr = f0 + (max(vec3f(1.0 - roughness), f0) - f0) * pow(1.0 - n_v, 5.0);
161
162 var color = diffuse * ambient + seen * fr;
163 color += shade(n, v, u.key_toward.xyz, u.key_color.rgb, diffuse, f0, roughness);
164 color += shade(n, v, u.fill_toward.xyz, u.fill_color.rgb, diffuse, f0, roughness);
165 return vec4f(color, alpha);
166 }