twgpu 0.4.1

use crate::sprites::AtlasSegments;

use super::blit::Blit;

const LABEL: Option<&str> = Some("BLIT");

impl Blit {
    /// Assumes that the mip level 0 is initialized.
    /// Fills in all smaller mip levels.
    /// Compatible with all texture dimensions except for 3D.
    ///
    /// Panics
    ///
    /// - On 3D textures
    pub fn generate_mipmaps(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        texture: &wgpu::Texture,
        device: &wgpu::Device,
    ) {
        assert_ne!(texture.dimension(), wgpu::TextureDimension::D3);
        // Array textures bound as non-array textures is not yet supported on webgl.
        // The webgpu backend will be caught in the cross-fire for the time being,
        // but this workaround is not too significant.
        // TODO: Remove this workaround when webgl is ready
        #[cfg(not(target_family = "wasm"))]
        for layer in 0..texture.size().depth_or_array_layers {
            self.generate_2d_mipmap(encoder, texture, device, layer);
        }
        #[cfg(target_family = "wasm")]
        self.webgl_generate_2d_array_mipmap(encoder, texture, device);
    }

    #[allow(unused)]
    fn webgl_generate_2d_array_mipmap(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        texture: &wgpu::Texture,
        device: &wgpu::Device,
    ) {
        let tmp_size = wgpu::Extent3d {
            depth_or_array_layers: 1,
            ..texture.size()
        };
        let tmp_texture = device.create_texture(&wgpu::TextureDescriptor {
            label: LABEL,
            size: tmp_size,
            mip_level_count: texture.mip_level_count(),
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: texture.format(),
            usage: wgpu::TextureUsages::COPY_SRC
                | wgpu::TextureUsages::COPY_DST
                | wgpu::TextureUsages::RENDER_ATTACHMENT
                | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });
        for layer in 0..texture.size().depth_or_array_layers {
            encoder.copy_texture_to_texture(
                wgpu::TexelCopyTextureInfo {
                    mip_level: 0,
                    origin: wgpu::Origin3d {
                        x: 0,
                        y: 0,
                        z: layer,
                    },
                    ..texture.as_image_copy()
                },
                tmp_texture.as_image_copy(),
                tmp_size,
            );
            self.generate_2d_mipmap(encoder, &tmp_texture, device, 0);
            for mip_level in 1..texture.mip_level_count() {
                encoder.copy_texture_to_texture(
                    wgpu::TexelCopyTextureInfo {
                        mip_level,
                        ..tmp_texture.as_image_copy()
                    },
                    wgpu::TexelCopyTextureInfo {
                        mip_level,
                        origin: wgpu::Origin3d {
                            x: 0,
                            y: 0,
                            z: layer,
                        },
                        ..texture.as_image_copy()
                    },
                    tmp_size.mip_level_size(mip_level, wgpu::TextureDimension::D2),
                );
            }
        }
    }

    fn generate_2d_mipmap(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        texture: &wgpu::Texture,
        device: &wgpu::Device,
        layer: u32,
    ) {
        for to_mip_level in 1..texture.mip_level_count() {
            let from = wgpu::TexelCopyTextureInfo {
                texture,
                mip_level: to_mip_level - 1,
                origin: wgpu::Origin3d {
                    z: layer,
                    ..wgpu::Origin3d::default()
                },
                aspect: wgpu::TextureAspect::All,
            };
            let to = wgpu::TexelCopyTextureInfo {
                texture,
                mip_level: to_mip_level,
                origin: wgpu::Origin3d {
                    z: layer,
                    ..wgpu::Origin3d::default()
                },
                aspect: wgpu::TextureAspect::All,
            };
            self.blit(from, to, encoder, device);
        }
    }

    /// Transform an RGBA8 image into a fully mipmapped texture.
    pub fn upload_mipmapped(
        &self,
        image: &image::RgbaImage,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
    ) -> wgpu::Texture {
        let width = image.width();
        let height = image.height();
        let size = wgpu::Extent3d {
            width,
            height,
            depth_or_array_layers: 1,
        };
        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: LABEL,
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
            mip_level_count: size.max_mips(wgpu::TextureDimension::D2),
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: wgpu::TextureUsages::TEXTURE_BINDING
                | wgpu::TextureUsages::COPY_SRC
                | wgpu::TextureUsages::COPY_DST
                | wgpu::TextureUsages::RENDER_ATTACHMENT,
            view_formats: &[],
        });
        queue.write_texture(
            texture.as_image_copy(),
            image.as_raw(),
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(4 * width),
                rows_per_image: Some(height),
            },
            wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: 1,
            },
        );
        let mut encoder =
            device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: LABEL });
        self.generate_mipmaps(&mut encoder, &texture, device);
        queue.submit([encoder.finish()]);
        texture
    }

    pub fn upload_mipmapped_atlas<T: AtlasSegments>(
        &self,
        atlas: &image::RgbaImage,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
    ) -> Vec<wgpu::Texture> {
        T::ALL_ORDERED
            .iter()
            .map(|segment| T::sub_view(*segment, atlas).to_image())
            .map(|image| self.upload_mipmapped(&image, device, queue))
            .collect()
    }
}