Desktop: Fix UI and viewport blending

desktop/src/render/composite_shader.wgsl

@@ -44,38 +44,60 @@ var s_diffuse: sampler;
 
 @fragment
 fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
-	let ui = textureSample(t_ui, s_diffuse, in.tex_coords);
-	if (ui.a >= 0.999) {
-		return ui;
+	let ui_linear = textureSample(t_ui, s_diffuse, in.tex_coords);
+	if (ui_linear.a >= 0.999) {
+		return ui_linear;
 	}
 
 	let viewport_coordinate = (in.tex_coords - constants.viewport_offset) * constants.viewport_scale;
 
-	// Vello renders its values to an `RgbaUnorm` texture, but if we try to use this in the main rendering pipeline
-	// which renders to an `Srgb` surface, gamma mapping is applied twice. This converts back to linear to compensate.
-	let overlay_raw = textureSample(t_overlays, s_diffuse, viewport_coordinate);
-	let overlay = vec4<f32>(srgb_to_linear(overlay_raw.rgb), overlay_raw.a);
-	let viewport_raw = textureSample(t_viewport, s_diffuse, viewport_coordinate);
-	let viewport = vec4<f32>(srgb_to_linear(viewport_raw.rgb), viewport_raw.a);
+	let overlay_srgb = textureSample(t_overlays, s_diffuse, viewport_coordinate);
+	let viewport_srgb = textureSample(t_viewport, s_diffuse, viewport_coordinate);
 
-	if (overlay.a < 0.001) {
-		return blend(ui, viewport);
+	// UI texture is premultiplied, we need to unpremultiply before blending
+	let ui_srgb = linear_to_srgb(unpremultiply(ui_linear));
+
+	if (overlay_srgb.a < 0.001) {
+		if (ui_srgb.a < 0.001) {
+			return srgb_to_linear(viewport_srgb);
+		} else {
+			return srgb_to_linear(blend(ui_srgb, viewport_srgb));
+		}
 	}
 
-	let composite = blend(overlay, viewport);
-	return blend(ui, composite);
-}
+	let composite_linear = blend(srgb_to_linear(overlay_srgb), srgb_to_linear(viewport_srgb));
 
-fn srgb_to_linear(srgb: vec3<f32>) -> vec3<f32> {
-	return select(
-		pow((srgb + 0.055) / 1.055, vec3<f32>(2.4)),
-		srgb / 12.92,
-		srgb <= vec3<f32>(0.04045)
-	);
+	if (ui_srgb.a < 0.001) {
+		return composite_linear;
+	}
+
+	return srgb_to_linear(blend(ui_srgb, linear_to_srgb(composite_linear)));
 }
 
 fn blend(fg: vec4<f32>, bg: vec4<f32>) -> vec4<f32> {
 	let a = fg.a + bg.a * (1.0 - fg.a);
 	let rgb = fg.rgb * fg.a + bg.rgb * bg.a * (1.0 - fg.a);
 	return vec4<f32>(rgb, a);
 }
+
+fn linear_to_srgb(in: vec4<f32>) -> vec4<f32> {
+	let cutoff = vec3<f32>(0.0031308);
+	let lo = in.rgb * 12.92;
+	let hi = 1.055 * pow(max(in.rgb, vec3<f32>(0.0)), vec3<f32>(1.0/2.4)) - 0.055;
+	return vec4<f32>(select(lo, hi, in.rgb > cutoff), in.a);
+}
+
+fn srgb_to_linear(in: vec4<f32>) -> vec4<f32> {
+	let cutoff = vec3<f32>(0.04045);
+	let lo = in.rgb / 12.92;
+	let hi = pow((in.rgb + 0.055) / 1.055, vec3<f32>(2.4));
+	return vec4<f32>(select(lo, hi, in.rgb > cutoff), in.a);
+}
+
+fn unpremultiply(in: vec4<f32>) -> vec4<f32> {
+	if (in.a > 0.0) {
+		return vec4<f32>((in.rgb / in.a), in.a);
+	} else {
+		return vec4<f32>(0.0);
+	}
+}