twgpu 0.4.1

Render Teeworlds and DDNet maps
Documentation
use std::ops::Range;

use crate::buffer::{GpuArrayBuffer, GpuDeviceExt, GpuVecBuffer};
use crate::Camera;
use wgpu::{
    include_wgsl, BindGroup, BlendState, BufferBinding, BufferUsages, ColorTargetState,
    ColorWrites, Device, FragmentState, IndexFormat, MultisampleState, PipelineCompilationOptions,
    PrimitiveState, RenderPipeline, RenderPipelineDescriptor, ShaderModule, TextureFormat,
    VertexState,
};

use crate::sprites::texture_storage::SpriteTextures;
use crate::sprites::{SpriteVertex, TextureToken};

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

pub struct BatchSpriteRender {
    index_buffer: GpuArrayBuffer<[u32; 6]>,
    shader: ShaderModule,
    format: TextureFormat,
    pipeline: Option<RenderPipeline>,
    bind_groups: Option<(BindGroup, BindGroup)>,
    current_textures: usize,
    cache: Vec<(Range<usize>, Option<wgpu::RenderBundle>)>,
}

fn adjust_index_buffer(idx: &mut GpuArrayBuffer<[u32; 6]>, amount: usize, device: &Device) -> bool {
    if idx.len() >= amount {
        return false;
    }
    let required_size = u32::try_from(amount.next_power_of_two()).expect("Too big index buffer");
    let indices: Vec<_> = (0..required_size)
        .map(|i| [0, 1, 3, 0, 2, 3].map(|x| x + i * 4))
        .collect();
    *idx = device.array_buffer(indices.as_slice(), Some("Index "), BufferUsages::INDEX);
    true
}

impl BatchSpriteRender {
    pub fn render<'pass>(
        &'pass mut self,
        camera_buffer: BufferBinding,
        camera_index: usize,
        sprite_buffer: &'pass GpuVecBuffer<[SpriteVertex; 4]>,
        texture_buffer: &'pass GpuVecBuffer<TextureToken>,
        textures: &'pass SpriteTextures,
        render_pass: &mut wgpu::RenderPass,
        device: &Device,
    ) {
        self.render_range(
            0..sprite_buffer.len(),
            camera_buffer,
            camera_index,
            sprite_buffer,
            texture_buffer,
            textures,
            render_pass,
            device,
        )
    }

    pub fn render_range<'pass>(
        &'pass mut self,
        range: Range<usize>,
        camera_buffer: BufferBinding,
        camera_index: usize,
        sprite_buffer: &'pass GpuVecBuffer<[SpriteVertex; 4]>,
        texture_buffer: &'pass GpuVecBuffer<TextureToken>,
        textures: &'pass SpriteTextures,
        render_pass: &mut wgpu::RenderPass,
        device: &Device,
    ) {
        assert_eq!(sprite_buffer.len(), texture_buffer.len());
        if range.is_empty() {
            return;
        }
        if camera_index >= self.cache.len() {
            self.cache.resize(camera_index + 1, (0..0, None));
        }
        let (largest_range, cached_bundle) = &mut self.cache[camera_index];
        if adjust_index_buffer(&mut self.index_buffer, sprite_buffer.len(), device)
            || textures.texture_count() != self.current_textures
        {
            // Reset cached bundle
            *largest_range = 0..0;
            let chunked_group_layouts = textures.chunked_group_layouts(device);
            let bind_groups = textures.chunked_bind_groups(
                &chunked_group_layouts,
                camera_buffer,
                texture_buffer,
                device,
            );
            self.current_textures = textures.texture_count();
            let vertex = VertexState {
                module: &self.shader,
                entry_point: Some("vs_main"),
                compilation_options: PipelineCompilationOptions::default(),
                buffers: &[SpriteVertex::vertex_buffer_layout()],
            };
            let fragment = FragmentState {
                module: &self.shader,
                entry_point: Some("fs_main"),
                compilation_options: PipelineCompilationOptions::default(),
                targets: &[Some(ColorTargetState {
                    format: self.format,
                    blend: Some(BlendState::ALPHA_BLENDING),
                    write_mask: ColorWrites::all(),
                })],
            };
            let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: LABEL,
                bind_group_layouts: &[&chunked_group_layouts.0, &chunked_group_layouts.1],
                push_constant_ranges: &[],
            });
            let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
                label: LABEL,
                layout: Some(&pipeline_layout),
                vertex,
                primitive: PrimitiveState::default(),
                depth_stencil: None,
                multisample: MultisampleState::default(),
                fragment: Some(fragment),
                multiview: None,
                cache: None,
            });
            self.pipeline = Some(pipeline);
            self.bind_groups = Some(bind_groups);
        }
        if *largest_range == range {
            // Happiest path!
            // Our render bundle matches the thing to render
            render_pass.execute_bundles([cached_bundle.as_ref().unwrap()]);
        } else if largest_range.len() < range.len() {
            // Semi-happy path: Amount of sprite is still growing.
            // We rebuild the render bundle for the new size.
            let mut bundle_encoder =
                device.create_render_bundle_encoder(&wgpu::RenderBundleEncoderDescriptor {
                    label: LABEL,
                    color_formats: &[Some(self.format)],
                    depth_stencil: None,
                    sample_count: 1,
                    multiview: None,
                });
            bundle_encoder.set_pipeline(self.pipeline.as_ref().unwrap());
            bundle_encoder.set_vertex_buffer(0, sprite_buffer.full_slice());
            bundle_encoder.set_index_buffer(self.index_buffer.full_slice(), IndexFormat::Uint32);
            let camera_offset = camera_index * std::mem::size_of::<Camera>();
            let bind_groups = self.bind_groups.as_ref().unwrap();
            bundle_encoder.set_bind_group(0, &bind_groups.0, &[camera_offset as u32]);
            bundle_encoder.set_bind_group(1, &bind_groups.1, &[]);
            let count = range.end - range.start;
            bundle_encoder.draw_indexed(0..6 * count as u32, range.start as i32 * 4, 0..1);

            let bundle = bundle_encoder.finish(&wgpu::RenderBundleDescriptor { label: LABEL });
            render_pass.execute_bundles([&bundle]);
            *largest_range = range;
            *cached_bundle = Some(bundle);
        } else {
            // Sad path,
            // too few sprites or different offset, will do everything manually.
            render_pass.set_pipeline(self.pipeline.as_ref().unwrap());
            render_pass.set_vertex_buffer(0, sprite_buffer.full_slice());
            render_pass.set_index_buffer(self.index_buffer.full_slice(), IndexFormat::Uint32);
            let camera_offset = camera_index * std::mem::size_of::<Camera>();
            let bind_groups = self.bind_groups.as_ref().unwrap();
            render_pass.set_bind_group(0, &bind_groups.0, &[camera_offset as u32]);
            render_pass.set_bind_group(1, &bind_groups.1, &[]);
            let count = range.end - range.start;
            render_pass.draw_indexed(0..6 * count as u32, range.start as i32 * 4, 0..1);
        }
    }

    pub fn new(format: TextureFormat, device: &Device) -> Self {
        let shader = device.create_shader_module(include_wgsl!("batched.wgsl"));
        let index_buffer = device.array_buffer(&[], Some("Zero-sized Init"), BufferUsages::INDEX);
        Self {
            index_buffer,
            shader,
            format,
            pipeline: None,
            bind_groups: None,
            current_textures: 0,
            cache: Vec::new(),
        }
    }
}