use crabslab::{Id, Slab};
use glam::{Vec2, Vec4, Vec4Swizzles};
use spirv_std::{image::Image2d, spirv, Sampler};
#[cfg(not(target_arch = "spirv"))]
mod cpu;
#[cfg(not(target_arch = "spirv"))]
pub use cpu::*;
#[cfg(feature = "bloom_vertex")]
#[spirv(vertex)]
pub fn bloom_vertex(
#[spirv(vertex_index)] vertex_index: u32,
#[spirv(instance_index)] in_id: u32,
out_uv: &mut Vec2,
#[spirv(flat)] out_id: &mut u32,
#[spirv(position)] out_clip_pos: &mut Vec4,
) {
let i = (vertex_index % 6) as usize;
*out_uv = crate::math::UV_COORD_QUAD_CCW[i];
*out_clip_pos = crate::math::CLIP_SPACE_COORD_QUAD_CCW[i];
*out_id = in_id;
}
#[cfg(feature = "bloom_downsample_fragment")]
#[spirv(fragment)]
pub fn bloom_downsample_fragment(
#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] slab: &[u32],
#[spirv(descriptor_set = 0, binding = 1)] texture: &Image2d,
#[spirv(descriptor_set = 0, binding = 2)] sampler: &Sampler,
in_uv: Vec2,
#[spirv(flat)] in_pixel_size_id: Id<Vec2>,
downsample: &mut Vec4,
) {
use glam::Vec3;
let Vec2 { x, y } = slab.read(in_pixel_size_id);
let a = texture.sample(*sampler, Vec2::new(in_uv.x - 2.0 * x, in_uv.y + 2.0 * y));
let b = texture.sample(*sampler, Vec2::new(in_uv.x, in_uv.y + 2.0 * y));
let c = texture.sample(*sampler, Vec2::new(in_uv.x + 2.0 * x, in_uv.y + 2.0 * y));
let d = texture.sample(*sampler, Vec2::new(in_uv.x - 2.0 * x, in_uv.y));
let e = texture.sample(*sampler, Vec2::new(in_uv.x, in_uv.y));
let f = texture.sample(*sampler, Vec2::new(in_uv.x + 2.0 * x, in_uv.y));
let g = texture.sample(*sampler, Vec2::new(in_uv.x - 2.0 * x, in_uv.y - 2.0 * y));
let h = texture.sample(*sampler, Vec2::new(in_uv.x, in_uv.y - 2.0 * y));
let i = texture.sample(*sampler, Vec2::new(in_uv.x + 2.0 * x, in_uv.y - 2.0 * y));
let j = texture.sample(*sampler, Vec2::new(in_uv.x - x, in_uv.y + y));
let k = texture.sample(*sampler, Vec2::new(in_uv.x + x, in_uv.y + y));
let l = texture.sample(*sampler, Vec2::new(in_uv.x - x, in_uv.y - y));
let m = texture.sample(*sampler, Vec2::new(in_uv.x + x, in_uv.y - y));
let f1 = 0.125;
let f2 = 0.0625;
let f3 = 0.03125;
let center = e * f1;
let inner = (j + k + l + m) * f1;
let outer = (b + d + h + f) * f2;
let furthest = (a + c + g + i) * f3;
let min = Vec3::splat(f32::EPSILON).extend(1.0);
*downsample = (center + inner + outer + furthest).max(min);
}
#[cfg(feature = "bloom_upsample_fragment")]
#[spirv(fragment)]
pub fn bloom_upsample_fragment(
#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] slab: &[u32],
#[spirv(descriptor_set = 0, binding = 1)] texture: &Image2d,
#[spirv(descriptor_set = 0, binding = 2)] sampler: &Sampler,
in_uv: Vec2,
#[spirv(flat)] filter_radius_id: Id<Vec2>,
upsample: &mut Vec4,
) {
let Vec2 { x, y } = slab.read(filter_radius_id);
let a = texture
.sample(*sampler, Vec2::new(in_uv.x - x, in_uv.y + y))
.xyz();
let b = texture
.sample(*sampler, Vec2::new(in_uv.x, in_uv.y + y))
.xyz();
let c = texture
.sample(*sampler, Vec2::new(in_uv.x + x, in_uv.y + y))
.xyz();
let d = texture
.sample(*sampler, Vec2::new(in_uv.x - x, in_uv.y))
.xyz();
let e = texture.sample(*sampler, Vec2::new(in_uv.x, in_uv.y)).xyz();
let f = texture
.sample(*sampler, Vec2::new(in_uv.x + x, in_uv.y))
.xyz();
let g = texture
.sample(*sampler, Vec2::new(in_uv.x - x, in_uv.y - y))
.xyz();
let h = texture
.sample(*sampler, Vec2::new(in_uv.x, in_uv.y - y))
.xyz();
let i = texture
.sample(*sampler, Vec2::new(in_uv.x + x, in_uv.y - y))
.xyz();
let mut sample = e * 4.0;
sample += (b + d + f + h) * 2.0;
sample += a + c + g + i;
sample *= 1.0 / 16.0;
*upsample = sample.extend(0.5);
}
#[cfg(feature = "bloom_mix_fragment")]
#[spirv(fragment)]
#[allow(clippy::too_many_arguments)]
pub fn bloom_mix_fragment(
#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] slab: &[u32],
#[spirv(descriptor_set = 0, binding = 1)] hdr_texture: &Image2d,
#[spirv(descriptor_set = 0, binding = 2)] hdr_sampler: &Sampler,
#[spirv(descriptor_set = 0, binding = 3)] bloom_texture: &Image2d,
#[spirv(descriptor_set = 0, binding = 4)] bloom_sampler: &Sampler,
in_uv: Vec2,
#[spirv(flat)] in_bloom_strength_id: Id<f32>,
frag_color: &mut Vec4,
) {
let bloom_strength = slab.read(in_bloom_strength_id);
let hdr = hdr_texture.sample(*hdr_sampler, in_uv).xyz();
let bloom = bloom_texture.sample(*bloom_sampler, in_uv).xyz();
let color = hdr.lerp(bloom, bloom_strength);
*frag_color = color.extend(1.0)
}