bevy/crates/bevy_render/src/shader/uniform.rs

use crate::{
    color::ColorSource,
    pipeline::{BindType, VertexBufferDescriptor},
    texture::Texture,
    Renderable,
};

use bevy_asset::{AssetStorage, Handle};
use bevy_core::bytes::GetBytes;
use legion::prelude::*;

pub trait AsUniforms: Send + Sync + 'static {
    fn get_field_infos() -> &'static [FieldInfo];
    fn get_uniform_bytes(&self, name: &str) -> Option<Vec<u8>>;
    fn get_uniform_texture(&self, name: &str) -> Option<Handle<Texture>>;
    fn get_shader_defs(&self) -> Option<Vec<String>>;
    fn get_field_bind_type(&self, name: &str) -> Option<FieldBindType>;
    fn get_uniform_bytes_ref(&self, name: &str) -> Option<&[u8]>;
    fn get_vertex_buffer_descriptor() -> Option<&'static VertexBufferDescriptor>;
}

pub fn shader_def_system<T>(uniforms: Ref<T>, mut renderable: RefMut<Renderable>)
where
    T: AsUniforms + Send + Sync + 'static,
{
    if let Some(shader_defs) = uniforms.get_shader_defs() {
        renderable
            .render_resource_assignments
            .pipeline_specialization
            .shader_specialization
            .shader_defs
            .extend(shader_defs)
    }
}

pub fn asset_handle_shader_def_system<T>(
    asset_storage: Resource<AssetStorage<T>>,
    asset_handle: Ref<Handle<T>>,
    mut renderable: RefMut<Renderable>,
) where
    T: AsUniforms + Send + Sync + 'static,
{
    if !renderable.is_visible || renderable.is_instanced {
        return;
    }

    let uniforms = asset_storage.get(&asset_handle).unwrap();
    if let Some(shader_defs) = uniforms.get_shader_defs() {
        renderable
            .render_resource_assignments
            .pipeline_specialization
            .shader_specialization
            .shader_defs
            .extend(shader_defs)
    }
}

pub trait ShaderDefSuffixProvider {
    fn get_shader_def(&self) -> Option<&'static str>;
}

impl ShaderDefSuffixProvider for bool {
    fn get_shader_def(&self) -> Option<&'static str> {
        match *self {
            true => Some(""),
            false => None,
        }
    }
}

pub enum FieldBindType {
    Uniform { size: usize },
    Texture,
}

pub struct FieldInfo {
    pub name: &'static str,
    pub uniform_name: &'static str,
    pub texture_name: &'static str,
    pub sampler_name: &'static str,
    pub is_instanceable: bool,
}

pub trait AsFieldBindType {
    fn get_bind_type(&self) -> Option<FieldBindType>;
}

impl AsFieldBindType for ColorSource {
    fn get_bind_type(&self) -> Option<FieldBindType> {
        match *self {
            ColorSource::Texture(_) => Some(FieldBindType::Texture),
            ColorSource::Color(color) => color.get_bind_type(),
        }
    }
}

impl AsFieldBindType for Option<Handle<Texture>> {
    fn get_bind_type(&self) -> Option<FieldBindType> {
        match *self {
            Some(_) => Some(FieldBindType::Texture),
            None => None,
        }
    }
}

impl AsFieldBindType for Handle<Texture> {
    fn get_bind_type(&self) -> Option<FieldBindType> {
        Some(FieldBindType::Texture)
    }
}

impl<T> AsFieldBindType for T
where
    T: GetBytes,
{
    // TODO: this breaks if get_bytes_ref() isn't supported for a datatype
    default fn get_bind_type(&self) -> Option<FieldBindType> {
        Some(FieldBindType::Uniform {
            size: self.get_bytes_ref().unwrap().len(),
        })
    }
}

pub trait GetTexture {
    fn get_texture(&self) -> Option<Handle<Texture>> {
        None
    }
}

impl<T> GetTexture for T
where
    T: GetBytes,
{
    default fn get_texture(&self) -> Option<Handle<Texture>> {
        None
    }
}

impl GetTexture for Handle<Texture> {
    fn get_texture(&self) -> Option<Handle<Texture>> {
        Some(self.clone())
    }
}

impl GetTexture for Option<Handle<Texture>> {
    fn get_texture(&self) -> Option<Handle<Texture>> {
        *self
    }
}

impl GetTexture for ColorSource {
    fn get_texture(&self) -> Option<Handle<Texture>> {
        match self {
            ColorSource::Color(_) => None,
            ColorSource::Texture(texture) => Some(texture.clone()),
        }
    }
}

pub struct UniformInfo<'a> {
    pub name: &'a str,
    pub bind_type: BindType,
}