Array texture example (#5077)

# Objective

- Make the reusable PBR shading functionality a little more reusable
  - Add constructor functions for `StandardMaterial` and `PbrInput` structs to populate them with default values
  - Document unclear `PbrInput` members
- Demonstrate how to reuse the bevy PBR shading functionality
- The final important piece from #3969 as the initial shot at making the PBR shader code reusable in custom materials

## Solution

- Add back and rework the 'old' `array_texture` example from pre-0.6.
- Create a custom shader material
  - Use a single array texture binding and sampler for the material bind group
  - Use a shader that calls `pbr()` from the `bevy_pbr::pbr_functions` import
- Spawn a row of cubes using the custom material
- In the shader, select the array texture layer to sample by using the world position x coordinate modulo the number of array texture layers

<img width="1392" alt="Screenshot 2022-06-23 at 12 28 05" src="https://user-images.githubusercontent.com/302146/175278593-2296f519-f577-4ece-81c0-d842283784a1.png">

Co-authored-by: Carter Anderson <mcanders1@gmail.com>

Author: superdump

Cargo.toml

@@ -1183,6 +1183,16 @@ description = "A compute shader that simulates Conway's Game of Life"
 category = "Shaders"
 wasm = false
 
+[[example]]
+name = "array_texture"
+path = "examples/shader/array_texture.rs"
+
+[package.metadata.example.array_texture]
+name = "Array Texture"
+description = "A shader that shows how to reuse the core bevy PBR shading functionality in a custom material that obtains the base color from an array texture."
+category = "Shaders"
+wasm = true
+
 # Stress tests
 [[package.metadata.category]]
 name = "Stress Tests"

assets/shaders/array_texture.wgsl

@@ -0,0 +1,63 @@
+#import bevy_pbr::mesh_view_bindings
+#import bevy_pbr::mesh_bindings
+
+#import bevy_pbr::pbr_types
+#import bevy_pbr::utils
+#import bevy_pbr::clustered_forward
+#import bevy_pbr::lighting
+#import bevy_pbr::shadows
+#import bevy_pbr::pbr_functions
+
+[[group(1), binding(0)]]
+var my_array_texture: texture_2d_array<f32>;
+[[group(1), binding(1)]]
+var my_array_texture_sampler: sampler;
+
+struct FragmentInput {
+    [[builtin(front_facing)]] is_front: bool;
+    [[builtin(position)]] frag_coord: vec4<f32>;
+    [[location(0)]] world_position: vec4<f32>;
+    [[location(1)]] world_normal: vec3<f32>;
+    [[location(2)]] uv: vec2<f32>;
+#ifdef VERTEX_TANGENTS
+    [[location(3)]] world_tangent: vec4<f32>;
+#endif
+#ifdef VERTEX_COLORS
+    [[location(4)]] color: vec4<f32>;
+#endif
+};
+
+[[stage(fragment)]]
+fn fragment(in: FragmentInput) -> [[location(0)]] vec4<f32> {
+    let layer = i32(in.world_position.x) & 0x3;
+
+    // Prepare a 'processed' StandardMaterial by sampling all textures to resolve
+    // the material members
+    var pbr_input: PbrInput = pbr_input_new();
+
+    pbr_input.material.base_color = textureSample(my_array_texture, my_array_texture_sampler, in.uv, layer);
+#ifdef VERTEX_COLORS
+    pbr_input.material.base_color = pbr_input.material.base_color * in.color;
+#endif
+
+    pbr_input.frag_coord = in.frag_coord;
+    pbr_input.world_position = in.world_position;
+    pbr_input.world_normal = in.world_normal;
+
+    pbr_input.is_orthographic = view.projection[3].w == 1.0;
+
+    pbr_input.N = prepare_normal(
+        pbr_input.material.flags,
+        in.world_normal,
+#ifdef VERTEX_TANGENTS
+#ifdef STANDARDMATERIAL_NORMAL_MAP
+        in.world_tangent,
+#endif
+#endif
+        in.uv,
+        in.is_front,
+    );
+    pbr_input.V = calculate_view(in.world_position, pbr_input.is_orthographic);
+
+    return pbr(pbr_input);
+}

crates/bevy_pbr/src/render/pbr.wgsl

@@ -74,6 +74,7 @@ fn fragment(in: FragmentInput) -> [[location(0)]] vec4<f32> {
         pbr_input.is_orthographic = view.projection[3].w == 1.0;
 
         pbr_input.N = prepare_normal(
+            material.flags,
             in.world_normal,
 #ifdef VERTEX_TANGENTS
 #ifdef STANDARDMATERIAL_NORMAL_MAP

crates/bevy_pbr/src/render/pbr_functions.wgsl

@@ -3,6 +3,7 @@
 // NOTE: This ensures that the world_normal is normalized and if
 // vertex tangents and normal maps then normal mapping may be applied.
 fn prepare_normal(
+    standard_material_flags: u32,
     world_normal: vec3<f32>,
 #ifdef VERTEX_TANGENTS
 #ifdef STANDARDMATERIAL_NORMAL_MAP
@@ -25,7 +26,7 @@ fn prepare_normal(
 #endif
 #endif
 
-    if ((material.flags & STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT) != 0u) {
+    if ((standard_material_flags & STANDARD_MATERIAL_FLAGS_DOUBLE_SIDED_BIT) != 0u) {
         if (!is_front) {
             N = -N;
 #ifdef VERTEX_TANGENTS
@@ -41,15 +42,15 @@ fn prepare_normal(
 #ifdef STANDARDMATERIAL_NORMAL_MAP
     // Nt is the tangent-space normal.
     var Nt: vec3<f32>;
-    if ((material.flags & STANDARD_MATERIAL_FLAGS_TWO_COMPONENT_NORMAL_MAP) != 0u) {
+    if ((standard_material_flags & STANDARD_MATERIAL_FLAGS_TWO_COMPONENT_NORMAL_MAP) != 0u) {
         // Only use the xy components and derive z for 2-component normal maps.
         Nt = vec3<f32>(textureSample(normal_map_texture, normal_map_sampler, uv).rg * 2.0 - 1.0, 0.0);
         Nt.z = sqrt(1.0 - Nt.x * Nt.x - Nt.y * Nt.y);
     } else {
         Nt = textureSample(normal_map_texture, normal_map_sampler, uv).rgb * 2.0 - 1.0;
     }
     // Normal maps authored for DirectX require flipping the y component
-    if ((material.flags & STANDARD_MATERIAL_FLAGS_FLIP_NORMAL_MAP_Y) != 0u) {
+    if ((standard_material_flags & STANDARD_MATERIAL_FLAGS_FLIP_NORMAL_MAP_Y) != 0u) {
         Nt.y = -Nt.y;
     }
     // NOTE: The mikktspace method of normal mapping applies maps the tangent-space normal from
@@ -86,12 +87,36 @@ struct PbrInput {
     occlusion: f32;
     frag_coord: vec4<f32>;
     world_position: vec4<f32>;
+    // Normalized world normal used for shadow mapping as normal-mapping is not used for shadow
+    // mapping
     world_normal: vec3<f32>;
+    // Normalized normal-mapped world normal used for lighting
     N: vec3<f32>;
+    // Normalized view vector in world space, pointing from the fragment world position toward the
+    // view world position
     V: vec3<f32>;
     is_orthographic: bool;
 };
 
+// Creates a PbrInput with default values
+fn pbr_input_new() -> PbrInput {
+    var pbr_input: PbrInput;
+
+    pbr_input.material = standard_material_new();
+    pbr_input.occlusion = 1.0;
+
+    pbr_input.frag_coord = vec4<f32>(0.0, 0.0, 0.0, 1.0);
+    pbr_input.world_position = vec4<f32>(0.0, 0.0, 0.0, 1.0);
+    pbr_input.world_normal = vec3<f32>(0.0, 0.0, 1.0);
+
+    pbr_input.is_orthographic = false;
+
+    pbr_input.N = vec3<f32>(0.0, 0.0, 1.0);
+    pbr_input.V = vec3<f32>(1.0, 0.0, 0.0);
+
+    return pbr_input;
+}
+
 fn pbr(
     in: PbrInput,
 ) -> vec4<f32> {

crates/bevy_pbr/src/render/pbr_types.wgsl

@@ -22,3 +22,19 @@ let STANDARD_MATERIAL_FLAGS_ALPHA_MODE_MASK: u32                = 128u;
 let STANDARD_MATERIAL_FLAGS_ALPHA_MODE_BLEND: u32               = 256u;
 let STANDARD_MATERIAL_FLAGS_TWO_COMPONENT_NORMAL_MAP: u32       = 512u;
 let STANDARD_MATERIAL_FLAGS_FLIP_NORMAL_MAP_Y: u32              = 1024u;
+
+// Creates a StandardMaterial with default values
+fn standard_material_new() -> StandardMaterial {
+    var material: StandardMaterial;
+
+    // NOTE: Keep in-sync with src/pbr_material.rs!
+    material.base_color = vec4<f32>(1.0, 1.0, 1.0, 1.0);
+    material.emissive = vec4<f32>(0.0, 0.0, 0.0, 1.0);
+    material.perceptual_roughness = 0.089;
+    material.metallic = 0.01;
+    material.reflectance = 0.5;
+    material.flags = STANDARD_MATERIAL_FLAGS_ALPHA_MODE_OPAQUE;
+    material.alpha_cutoff = 0.5;
+
+    return material;
+}

examples/README.md

@@ -253,6 +253,7 @@ There are also compute shaders which are used for more general processing levera
 Example | Description
 --- | ---
 [Animated](../examples/shader/animate_shader.rs) | A shader that uses dynamic data like the time since startup
+[Array Texture](../examples/shader/array_texture.rs) | A shader that shows how to reuse the core bevy PBR shading functionality in a custom material that obtains the base color from an array texture.
 [Compute - Game of Life](../examples/shader/compute_shader_game_of_life.rs) | A compute shader that simulates Conway's Game of Life
 [Custom Vertex Attribute](../examples/shader/custom_vertex_attribute.rs) | A shader that reads a mesh's custom vertex attribute
 [Instancing](../examples/shader/shader_instancing.rs) | A shader that renders a mesh multiple times in one draw call

examples/shader/array_texture.rs

@@ -0,0 +1,190 @@
+use bevy::{
+    asset::LoadState,
+    ecs::system::{lifetimeless::SRes, SystemParamItem},
+    pbr::MaterialPipeline,
+    prelude::*,
+    reflect::TypeUuid,
+    render::{
+        render_asset::{PrepareAssetError, RenderAsset, RenderAssets},
+        render_resource::{
+            BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout,
+            BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingResource, BindingType,
+            SamplerBindingType, ShaderStages, TextureSampleType, TextureViewDimension,
+        },
+        renderer::RenderDevice,
+    },
+};
+
+/// This example illustrates how to create a texture for use with a `texture_2d_array<f32>` shader
+/// uniform variable.
+fn main() {
+    App::new()
+        .add_plugins(DefaultPlugins)
+        .add_plugin(MaterialPlugin::<ArrayTextureMaterial>::default())
+        .add_startup_system(setup)
+        .add_system(create_array_texture)
+        .run();
+}
+
+struct LoadingTexture {
+    is_loaded: bool,
+    handle: Handle<Image>,
+}
+
+fn setup(mut commands: Commands, asset_server: Res<AssetServer>) {
+    // Start loading the texture.
+    commands.insert_resource(LoadingTexture {
+        is_loaded: false,
+        handle: asset_server.load("textures/array_texture.png"),
+    });
+
+    // light
+    commands.spawn_bundle(PointLightBundle {
+        point_light: PointLight {
+            intensity: 3000.0,
+            ..Default::default()
+        },
+        transform: Transform::from_xyz(-3.0, 2.0, -1.0),
+        ..Default::default()
+    });
+    commands.spawn_bundle(PointLightBundle {
+        point_light: PointLight {
+            intensity: 3000.0,
+            ..Default::default()
+        },
+        transform: Transform::from_xyz(3.0, 2.0, 1.0),
+        ..Default::default()
+    });
+
+    // camera
+    commands.spawn_bundle(Camera3dBundle {
+        transform: Transform::from_xyz(5.0, 5.0, 5.0).looking_at(Vec3::new(1.5, 0.0, 0.0), Vec3::Y),
+        ..Default::default()
+    });
+}
+
+fn create_array_texture(
+    mut commands: Commands,
+    asset_server: Res<AssetServer>,
+    mut loading_texture: ResMut<LoadingTexture>,
+    mut images: ResMut<Assets<Image>>,
+    mut meshes: ResMut<Assets<Mesh>>,
+    mut materials: ResMut<Assets<ArrayTextureMaterial>>,
+) {
+    if loading_texture.is_loaded
+        || asset_server.get_load_state(loading_texture.handle.clone()) != LoadState::Loaded
+    {
+        return;
+    }
+    loading_texture.is_loaded = true;
+    let image = images.get_mut(&loading_texture.handle).unwrap();
+
+    // Create a new array texture asset from the loaded texture.
+    let array_layers = 4;
+    image.reinterpret_stacked_2d_as_array(array_layers);
+
+    // Spawn some cubes using the array texture
+    let mesh_handle = meshes.add(Mesh::from(shape::Cube { size: 1.0 }));
+    let material_handle = materials.add(ArrayTextureMaterial {
+        array_texture: loading_texture.handle.clone(),
+    });
+    for x in -5..=5 {
+        commands.spawn_bundle(MaterialMeshBundle {
+            mesh: mesh_handle.clone(),
+            material: material_handle.clone(),
+            transform: Transform::from_xyz(x as f32 + 0.5, 0.0, 0.0),
+            ..Default::default()
+        });
+    }
+}
+
+#[derive(Debug, Clone, TypeUuid)]
+#[uuid = "9c5a0ddf-1eaf-41b4-9832-ed736fd26af3"]
+struct ArrayTextureMaterial {
+    array_texture: Handle<Image>,
+}
+
+#[derive(Clone)]
+pub struct GpuArrayTextureMaterial {
+    bind_group: BindGroup,
+}
+
+impl RenderAsset for ArrayTextureMaterial {
+    type ExtractedAsset = ArrayTextureMaterial;
+    type PreparedAsset = GpuArrayTextureMaterial;
+    type Param = (
+        SRes<RenderDevice>,
+        SRes<MaterialPipeline<Self>>,
+        SRes<RenderAssets<Image>>,
+    );
+    fn extract_asset(&self) -> Self::ExtractedAsset {
+        self.clone()
+    }
+
+    fn prepare_asset(
+        extracted_asset: Self::ExtractedAsset,
+        (render_device, material_pipeline, gpu_images): &mut SystemParamItem<Self::Param>,
+    ) -> Result<Self::PreparedAsset, PrepareAssetError<Self::ExtractedAsset>> {
+        let (array_texture_texture_view, array_texture_sampler) = if let Some(result) =
+            material_pipeline
+                .mesh_pipeline
+                .get_image_texture(gpu_images, &Some(extracted_asset.array_texture.clone()))
+        {
+            result
+        } else {
+            return Err(PrepareAssetError::RetryNextUpdate(extracted_asset));
+        };
+        let bind_group = render_device.create_bind_group(&BindGroupDescriptor {
+            entries: &[
+                BindGroupEntry {
+                    binding: 0,
+                    resource: BindingResource::TextureView(array_texture_texture_view),
+                },
+                BindGroupEntry {
+                    binding: 1,
+                    resource: BindingResource::Sampler(array_texture_sampler),
+                },
+            ],
+            label: Some("array_texture_material_bind_group"),
+            layout: &material_pipeline.material_layout,
+        });
+
+        Ok(GpuArrayTextureMaterial { bind_group })
+    }
+}
+
+impl Material for ArrayTextureMaterial {
+    fn fragment_shader(asset_server: &AssetServer) -> Option<Handle<Shader>> {
+        Some(asset_server.load("shaders/array_texture.wgsl"))
+    }
+
+    fn bind_group(render_asset: &<Self as RenderAsset>::PreparedAsset) -> &BindGroup {
+        &render_asset.bind_group
+    }
+
+    fn bind_group_layout(render_device: &RenderDevice) -> BindGroupLayout {
+        render_device.create_bind_group_layout(&BindGroupLayoutDescriptor {
+            entries: &[
+                // Array Texture
+                BindGroupLayoutEntry {
+                    binding: 0,
+                    visibility: ShaderStages::FRAGMENT,
+                    ty: BindingType::Texture {
+                        multisampled: false,
+                        sample_type: TextureSampleType::Float { filterable: true },
+                        view_dimension: TextureViewDimension::D2Array,
+                    },
+                    count: None,
+                },
+                // Array Texture Sampler
+                BindGroupLayoutEntry {
+                    binding: 1,
+                    visibility: ShaderStages::FRAGMENT,
+                    ty: BindingType::Sampler(SamplerBindingType::Filtering),
+                    count: None,
+                },
+            ],
+            label: None,
+        })
+    }
+}