Voxel Face Render with Bevy PBR

material.rs

use bevy::mesh::{MeshVertexAttribute, MeshVertexBufferLayoutRef};
use bevy::pbr::wireframe::WireframeConfig;
use bevy::pbr::{
    Material, MaterialPipeline, MaterialPipelineKey, MeshPipelineKey, TONEMAPPING_LUT_SAMPLER_BINDING_INDEX,
    TONEMAPPING_LUT_TEXTURE_BINDING_INDEX,
};
use bevy::prelude::*;
use bevy::render::render_resource::{
    AsBindGroup, Face, RenderPipelineDescriptor, SpecializedMeshPipelineError, VertexFormat,
};
use bevy::shader::{ShaderDefVal, ShaderRef};
use bytemuck::{Pod, Zeroable};

/// Custom vertex attribute for packed voxel face data.
pub const ATTRIBUTE_VOXEL_FACE: MeshVertexAttribute =
    MeshVertexAttribute::new("VoxelFace", 0x7E57_FA00, VertexFormat::Uint32x2);

bitflags::bitflags! {
    #[repr(transparent)]
    #[derive(Clone, Copy, Debug, Default, Pod, Zeroable)]
    pub struct PBRFlags : u32 {
        /// Receives shadows.
        const SHADOW_RECEIVER_BIT = 1u32 << 0u32;
        /// Receives transmitted shadows (e.g. colored light through leaves). Only has effect if SHADOW_RECEIVER_BIT is also set.
        const TRANSMITTED_SHADOW_RECEIVER_BIT = 1u32 << 1u32;
        /// Skips all lighting calculations, including shadows. Use for fully emissive or unlit materials.
        const UNLIT_BIT = 1u32 << 2u32;
        /// Render wireframe borders for debugging.
        const SHOW_WIREFRAME_BIT = 1u32 << 3u32;
    }
}

impl Into<u32> for PBRFlags {
    fn into(self) -> u32 {
        self.bits()
    }
}

/// Simple voxel material with custom vertex layout.
#[derive(Asset, AsBindGroup, Debug, Clone, TypePath)]
pub struct VoxelMaterial {
    /// Chunk origin in world space (xyz = position, w = unused)
    #[uniform(0)]
    pub chunk_origin: Vec4,

    /// Surface atlas texture containing block face materials (6 faces × 8×8 per block)
    #[texture(1)]
    #[sampler(2)]
    pub surface_atlas_color: Handle<Image>,

    /// Surface atlas texture containing block face materials (6 faces × 8×8 per block)
    #[texture(3)]
    #[sampler(4)]
    pub surface_atlas_emissive: Handle<Image>,

    /// Surface atlas texture containing block face materials (6 faces × 8×8 per block)
    #[texture(5)]
    #[sampler(6)]
    pub surface_atlas_metallic_roughness: Handle<Image>,

    /// Flags defined by PBRFlags bitfield (e.g. shadow receiver, unlit)
    #[uniform(7)]
    pub shader_pbr_flags: u32,
}

impl VoxelMaterial {
    pub fn set_wireframe(&mut self, enabled: bool) {
        if enabled {
            self.shader_pbr_flags |= PBRFlags::SHOW_WIREFRAME_BIT.bits();
        } else {
            self.shader_pbr_flags &= !PBRFlags::SHOW_WIREFRAME_BIT.bits();
        }
    }

    pub fn set_receives_shadows(&mut self, enabled: bool) {
        if enabled {
            self.shader_pbr_flags |= PBRFlags::SHADOW_RECEIVER_BIT.bits();
        } else {
            self.shader_pbr_flags &= !PBRFlags::SHADOW_RECEIVER_BIT.bits();
        }
    }
}

impl Material for VoxelMaterial {
    fn vertex_shader() -> ShaderRef {
        "shaders/voxel_face_instancing.wgsl".into()
    }

    fn fragment_shader() -> ShaderRef {
        "shaders/voxel_face_instancing.wgsl".into()
    }

    fn prepass_vertex_shader() -> ShaderRef {
        "shaders/voxel_face_instancing.wgsl".into()
    }

    //fn prepass_fragment_shader() -> ShaderRef {
    //    "shaders/voxel_face_instancing.wgsl".into()
    //}

    fn specialize(
        _pipeline: &MaterialPipeline,
        descriptor: &mut RenderPipelineDescriptor,
        layout: &MeshVertexBufferLayoutRef,
        key: MaterialPipelineKey<Self>,
    ) -> Result<(), SpecializedMeshPipelineError> {
        let vertex_layout = layout.0.get_layout(&[ATTRIBUTE_VOXEL_FACE.at_shader_location(0)])?;
        descriptor.vertex.buffers = vec![vertex_layout];
        descriptor.primitive.cull_mode = Some(Face::Back);

        // print backtrace
        info!("{:?}", key.mesh_key);

        let mut shader_defs_to_add = Vec::new();

        // TODO: get these from bevy pbr someway
        if key.mesh_key.intersects(
            MeshPipelineKey::NORMAL_PREPASS
                | MeshPipelineKey::MOTION_VECTOR_PREPASS
                | MeshPipelineKey::DEFERRED_PREPASS,
        ) {
            info!("Adding PREPASS_FRAGMENT shader def for VoxelMaterial");
            shader_defs_to_add.push("PREPASS_FRAGMENT".into());
        }

        // shader_defs_to_add.push("TONEMAP_IN_SHADER".into()); // BROKEN
        shader_defs_to_add.push(ShaderDefVal::UInt(
            "TONEMAPPING_LUT_TEXTURE_BINDING_INDEX".into(),
            TONEMAPPING_LUT_TEXTURE_BINDING_INDEX,
        ));
        shader_defs_to_add.push(ShaderDefVal::UInt(
            "TONEMAPPING_LUT_SAMPLER_BINDING_INDEX".into(),
            TONEMAPPING_LUT_SAMPLER_BINDING_INDEX,
        ));

        shader_defs_to_add.push("TONEMAP_METHOD_TONY_MC_MAPFACE".into());

        if key.mesh_key.msaa_samples() > 1 {
            shader_defs_to_add.push("MULTISAMPLED".into());
        };

        // shader_defs_to_add.push("SCREEN_SPACE_AMBIENT_OCCLUSION".into());

        descriptor.vertex.shader_defs.extend(shader_defs_to_add.clone());
        if let Some(fragment) = &mut descriptor.fragment {
            fragment.shader_defs.extend(shader_defs_to_add);
        }

        Ok(())
    }
}

/// Create a VoxelMaterial with given chunk origin and surface atlas.
pub fn create_voxel_material(
    chunk_origin: Vec3,
    color_map: Handle<Image>,
    emissive_map: Handle<Image>,
    metallic_roughness_map: Handle<Image>,
) -> VoxelMaterial {
    VoxelMaterial {
        chunk_origin: chunk_origin.extend(0.0),
        surface_atlas_color: color_map.clone(),
        surface_atlas_emissive: emissive_map.clone(),
        surface_atlas_metallic_roughness: metallic_roughness_map.clone(),
        shader_pbr_flags: (PBRFlags::SHADOW_RECEIVER_BIT | PBRFlags::TRANSMITTED_SHADOW_RECEIVER_BIT).into(),
    }
}

pub(crate) fn update_wireframe(
    wireframe_config: Res<WireframeConfig>,
    mut voxel_materials: ResMut<Assets<VoxelMaterial>>,
) {
    let material_ids: Vec<_> = voxel_materials.iter().map(|(id, _)| id).collect();
    for material_id in material_ids {
        voxel_materials.get_mut(material_id).unwrap().set_wireframe(wireframe_config.global);
    }
}

shader.wgsl

// Simple voxel shader with custom vertex attribute
#import consts

#import bevy_pbr::{
    mesh_functions::{get_world_from_local, mesh_position_local_to_clip, mesh_position_local_to_world},
    mesh_view_bindings::view,
    pbr_types::{STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT, PbrInput, pbr_input_new},
    pbr_bindings,
    pbr_functions as fns,
    mesh_types::{MESH_FLAGS_SHADOW_RECEIVER_BIT, MESH_FLAGS_TRANSMITTED_SHADOW_RECEIVER_BIT},
}

#ifdef TONEMAP_IN_SHADER
#import bevy_core_pipeline::tonemapping
#endif



// Our custom material (AsBindGroup generates the binding)
struct VoxelMaterial {
    chunk_origin: vec4<f32>,
}

@group(#{MATERIAL_BIND_GROUP}) @binding(0)
var<uniform> voxel_material: VoxelMaterial;

@group(#{MATERIAL_BIND_GROUP}) @binding(1)
var color_map: texture_2d<f32>;

@group(#{MATERIAL_BIND_GROUP}) @binding(2)
var color_map_sampler: sampler;


@group(#{MATERIAL_BIND_GROUP}) @binding(3)
var emissive: texture_2d<f32>;

@group(#{MATERIAL_BIND_GROUP}) @binding(4)
var emissive_sampler: sampler;


@group(#{MATERIAL_BIND_GROUP}) @binding(5)
var metallic_roughness: texture_2d<f32>;

@group(#{MATERIAL_BIND_GROUP}) @binding(6)
var metallic_roughness_sampler: sampler;


@group(#{MATERIAL_BIND_GROUP}) @binding(7)
var<uniform> shader_pbr_flags: u32;

const MASK3: u32 = 7u;
const MASK8: u32 = 255u;

struct FragmentOutput {
    @location(0) color: vec4<f32>,
    @builtin(frag_depth) depth: f32,
}

struct VertexInput {
    @builtin(instance_index) instance_index: u32,
    @builtin(vertex_index) vertex_index: u32,
    @location(0) voxel_face: vec2<u32>,
}

struct CustomVertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) world_position: vec4<f32>,
    @location(1) world_normal: vec3<f32>,
    @location(2) color: vec4<f32>,
    @location(3) uv: vec2<f32>,
    @location(4) face_dims: vec2<f32>,      // quad dimensions in sub-voxels
    @location(5) voxel_id: u32,             // Block atlas ID
    @location(6) face: u32,                 // Face direction (0-5)
    @location(7) instance_index: u32,       // Pass through instance index for fragment shader
    @location(8) flags: u32,                 // Custom flags (unused for now)
}

fn get_face_normal(face: u32) -> vec3<f32> {
    switch (face) {
        case 0u: { return vec3<f32>(1.0, 0.0, 0.0); }
        case 1u: { return vec3<f32>(-1.0, 0.0, 0.0); }
        case 2u: { return vec3<f32>(0.0, 1.0, 0.0); }
        case 3u: { return vec3<f32>(0.0, -1.0, 0.0); }
        case 4u: { return vec3<f32>(0.0, 0.0, 1.0); }
        case 5u: { return vec3<f32>(0.0, 0.0, -1.0); }
        default: { return vec3<f32>(0.0, 1.0, 0.0); }
    }
}

fn get_face_color(face: u32) -> vec3<f32> {
    switch (face) {
        case 0u: { return vec3<f32>(1.0, 0.0, 0.0); }  // +X red
        case 1u: { return vec3<f32>(1.0, 0.5, 0.0); }  // -X orange
        case 2u: { return vec3<f32>(0.0, 1.0, 0.0); }  // +Y green
        case 3u: { return vec3<f32>(1.0, 1.0, 0.0); }  // -Y yellow
        case 4u: { return vec3<f32>(0.0, 0.0, 1.0); }  // +Z blue
        case 5u: { return vec3<f32>(1.0, 0.0, 1.0); }  // -Z magenta
        default: { return vec3<f32>(1.0, 1.0, 1.0); }
    }
}

fn get_corner_uv(corner: u32) -> vec2<f32> {
    switch (corner) {
        case 0u: { return vec2<f32>(0.0, 0.0); }
        case 1u: { return vec2<f32>(1.0, 0.0); }
        case 2u: { return vec2<f32>(1.0, 1.0); }
        case 3u: { return vec2<f32>(0.0, 1.0); }
        default: { return vec2<f32>(0.0, 0.0); }
    }
}

fn compute_sub_uv(world_pos: vec3<f32>, face: u32) -> vec2<u32> {

    // Get sub-voxel coordinates (0-7) based on face direction
    var sub_u: u32;
    var sub_v: u32;

    // Extract sub-voxel coordinates based on face direction
    // Each face maps different world axes to U/V texture coordinates
    switch (face) {
        case 0u: {  // +X: U=Y, V=Z
            sub_u = u32(floor(fract(world_pos.y / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
            sub_v = u32(floor(fract(world_pos.z / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
        }
        case 1u: {  // -X: U=Y, V=Z (flipped)
            sub_u = u32(floor(fract(world_pos.y / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
            sub_v = u32(floor(fract(world_pos.z / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
        }
        case 2u: {  // +Y: U=X, V=Z (flipped)
            sub_u = u32(floor(fract(world_pos.x / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
            sub_v = u32(floor(fract(world_pos.z / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
        }
        case 3u: {  // -Y: U=X, V=Z
            sub_u = u32(floor(fract(world_pos.x / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
            sub_v = u32(floor(fract(world_pos.z / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
        }
        case 4u: {  // +Z: U=X, V=Y
            sub_u = u32(floor(fract(world_pos.x / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
            sub_v = u32(floor(fract(world_pos.y / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
        }
        case 5u: {  // -Z: U=X, V=Y (flipped)
            sub_u = u32(floor(fract(world_pos.x / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
            sub_v = u32(floor(fract(world_pos.y / consts::UNIT_VOXEL_SIZE / consts::BLOCK_RESOLUTION) * consts::BLOCK_RESOLUTION));
        }
        default: {
            sub_u = 0u;
            sub_v = 0u;
        }
    }

    // Clamp to valid range (0-7)
    sub_u = clamp(sub_u, 0u, 7u);
    sub_v = clamp(sub_v, 0u, 7u);

    return vec2<u32>(sub_u, sub_v);
}

@vertex
fn vertex(in: VertexInput) -> CustomVertexOutput {
    var out: CustomVertexOutput;

    // Unpack face data
    // lo layout: x(9 bits) | y(9 bits) | z(9 bits) | face(3 bits) | unused(2 bits)
    let lo = in.voxel_face.x;
    let hi = in.voxel_face.y;

    let block_x = lo & 0x1FFu;  // 9 bits
    let block_y = (lo >> 9u) & 0x1FFu;  // 9 bits
    let block_z = (lo >> 18u) & 0x1FFu;  // 9 bits
    let face = (lo >> 27u) & MASK3;  // 3 bits
    // du and dv now use 9 bits each: du (0-8), dv (9-17)
    let du = hi & 0x1FFu;  // 9 bits
    let dv = (hi >> 9u) & 0x1FFu;  // 9 bits
    let voxel_id = (hi >> 18u) & 0x3FFFu;  // 14 bits (bits 18-31)

    // Map vertex index to corner
    let local_vertex = in.vertex_index % 6u;
    var corner: u32;
    switch (local_vertex) {
        case 0u: { corner = 0u; }
        case 1u: { corner = 1u; }
        case 2u: { corner = 2u; }
        case 3u: { corner = 0u; }
        case 4u: { corner = 2u; }
        case 5u: { corner = 3u; }
        default: { corner = 0u; }
    }

    let corner_uv = get_corner_uv(corner);

    let origin_meters = vec3<f32>(
        f32(block_x) * consts::UNIT_VOXEL_SIZE,
        f32(block_y) * consts::UNIT_VOXEL_SIZE,
        f32(block_z) * consts::UNIT_VOXEL_SIZE,
    );

    let quad_width = f32(du) * consts::UNIT_VOXEL_SIZE;
    let quad_height = f32(dv) * consts::UNIT_VOXEL_SIZE;

    let u_pos = corner_uv.x * quad_width;
    let v_pos = corner_uv.y * quad_height;

    // Map to 3D position based on face
    var local_pos: vec3<f32>;
    switch (face) {
        case 0u: { local_pos = origin_meters + vec3<f32>(0.0, u_pos, v_pos); }         // +X
        case 1u: { local_pos = origin_meters + vec3<f32>(0.0, quad_width - u_pos, v_pos); }    // -X
        case 2u: { local_pos = origin_meters + vec3<f32>(quad_width - u_pos, 0.0, v_pos); } // +Y
        case 3u: { local_pos = origin_meters + vec3<f32>(u_pos, 0.0, v_pos); }                 // -Y
        case 4u: { local_pos = origin_meters + vec3<f32>(u_pos, v_pos, 0.0); }         // +Z
        case 5u: { local_pos = origin_meters + vec3<f32>(quad_width - u_pos, v_pos, 0.0); }    // -Z
        default: { local_pos = origin_meters; }
    }

    let position = vec4<f32>(local_pos, 1.0);

    out.position = mesh_position_local_to_clip(
        get_world_from_local(in.instance_index),
        position,
    );
    out.world_position = mesh_position_local_to_world(
        get_world_from_local(in.instance_index),
        position,
    );
    out.world_normal = get_face_normal(face);

    let face_color = get_face_color(face);
    out.color = vec4<f32>(face_color, 1.0);
    out.uv = corner_uv;
    out.face_dims = vec2<f32>(f32(du), f32(dv));
    out.voxel_id = voxel_id;
    out.face = face;
    out.instance_index = in.instance_index;

    return out;
}

fn get_tex_coord(in: CustomVertexOutput) -> vec2<i32> {
    let world_pos = in.world_position.xyz;
    let sub_uv = compute_sub_uv(world_pos, in.face);
    let sub_u = sub_uv.x;
    let sub_v = sub_uv.y;
    let atlas_dims = textureDimensions(color_map);
    let face_index = in.voxel_id * 6u + in.face;
    let faces_per_row = atlas_dims.x / 8u;
    let face_col = face_index % faces_per_row;
    let face_row = face_index / faces_per_row;

    // Final texture coordinates (pixel position in atlas)
    let tex_x = face_col * 8u + sub_u;
    let tex_y = face_row * 8u + sub_v;

    return vec2<i32>(i32(tex_x), i32(tex_y));
}


@fragment
fn fragment(
    in: CustomVertexOutput,
    @builtin(front_facing) is_front: bool,
) -> FragmentOutput {
    var out: FragmentOutput;
    // Compute depth
    out.depth = in.position.z;


#ifndef PREPASS_PIPELINE
    if ((shader_pbr_flags & consts::SHOW_WIREFRAME_BIT) != 0u) {
        // Wireframe rendering - scale line thickness based on face dimensions
        // We want constant thickness in sub-voxel units, not UV space
        let line_thickness_subvoxels = 0.1;  // Half a sub-voxel width

        // Scale threshold by face dimensions (larger faces need smaller UV threshold)
        let edge_threshold_u = line_thickness_subvoxels / in.face_dims.x;
        let edge_threshold_v = line_thickness_subvoxels / in.face_dims.y;
        let dist_to_edge_u = min(in.uv.x, 1.0 - in.uv.x);
        let dist_to_edge_v = min(in.uv.y, 1.0 - in.uv.y);

        // Check if we're near an edge (use dimension-specific thresholds)
        let near_u_edge = dist_to_edge_u < edge_threshold_u;
        let near_v_edge = dist_to_edge_v < edge_threshold_v;

        if near_u_edge || near_v_edge {
            out.color = vec4<f32>(1.0, 1.0, 1.0, 1.0);
            return out;
        }
    }
#endif

    let tex_coords = get_tex_coord(in);
    let tex_x = tex_coords.x;
    let tex_y = tex_coords.y;

#ifdef DEPTH_PREPASS
    // In depth prepass, we only output depth
    return out;
#else

    // Sample the texture (use nearest neighbor)
    let material_color = textureLoad(color_map, vec2<i32>(i32(tex_x), i32(tex_y)), 0);
    let material_emissive = textureLoad(emissive, vec2<i32>(i32(tex_x), i32(tex_y)), 0);
    let material_metallic_roughness = textureLoad(metallic_roughness, vec2<i32>(i32(tex_x), i32(tex_y)), 0);

    let is_orthographic = false;

    var pbr_in: PbrInput = pbr_input_new();
    pbr_in.flags = MESH_FLAGS_SHADOW_RECEIVER_BIT | MESH_FLAGS_TRANSMITTED_SHADOW_RECEIVER_BIT;
    pbr_in.material.base_color = material_color;
    pbr_in.material.base_color = fns::alpha_discard(pbr_in.material, pbr_in.material.base_color);
    pbr_in.material.emissive = vec4<f32>(material_emissive.rgb, 1.0);
    pbr_in.material.metallic = material_metallic_roughness.b;
    pbr_in.material.perceptual_roughness = material_metallic_roughness.g;

    pbr_in.frag_coord = in.position;
    pbr_in.world_position = in.world_position;
    let double_sided = (pbr_in.material.flags & STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT) != 0u;
    pbr_in.world_normal = fns::prepare_world_normal(
        in.world_normal,
        double_sided,
        is_front,
    );

    pbr_in.N = normalize(pbr_in.world_normal);
    pbr_in.V = fns::calculate_view(in.world_position, is_orthographic);

    // var out: FragmentOutput;
    out.color = fns::apply_pbr_lighting(pbr_in) * pbr_in.material.base_color;
    out.color = fns::main_pass_post_lighting_processing(pbr_in, out.color);

#endif
    return out;
}