Refactor `quad::Pipeline` to `prepare` and `render` architecture

3 files changed, 230 insertions, 112 deletions

diff --git a/wgpu/src/backend.rs b/wgpu/src/backend.rs
index f19959bd..2c35936c 100644
--- a/wgpu/src/backend.rs
+++ b/wgpu/src/backend.rs
@@ -88,13 +88,14 @@ impl Backend {
         let mut layers = Layer::generate(primitives, viewport);
         layers.push(Layer::overlay(overlay_text, viewport));
 
-        for layer in layers {
+        for (i, layer) in layers.iter().enumerate() {
             self.flush(
                 device,
                 queue,
                 scale_factor,
                 transformation,
                 &layer,
+                i,
                 staging_belt,
                 encoder,
                 frame,
@@ -102,6 +103,7 @@ impl Backend {
             );
         }
 
+        self.quad_pipeline.end_frame();
         self.text_pipeline.end_frame();
 
         #[cfg(any(feature = "image", feature = "svg"))]
@@ -115,6 +117,7 @@ impl Backend {
         scale_factor: f32,
         transformation: Transformation,
         layer: &Layer<'_>,
+        layer_index: usize,
         staging_belt: &mut wgpu::util::StagingBelt,
         encoder: &mut wgpu::CommandEncoder,
         target: &wgpu::TextureView,
@@ -127,16 +130,32 @@ impl Backend {
         }
 
         if !layer.quads.is_empty() {
-            self.quad_pipeline.draw(
+            self.quad_pipeline.prepare(
                 device,
-                staging_belt,
-                encoder,
+                queue,
                 &layer.quads,
                 transformation,
                 scale_factor,
-                bounds,
-                target,
             );
+
+            let mut render_pass =
+                encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+                    label: Some("iced_wgpu::quad render pass"),
+                    color_attachments: &[Some(
+                        wgpu::RenderPassColorAttachment {
+                            view: target,
+                            resolve_target: None,
+                            ops: wgpu::Operations {
+                                load: wgpu::LoadOp::Load,
+                                store: true,
+                            },
+                        },
+                    )],
+                    depth_stencil_attachment: None,
+                });
+
+            self.quad_pipeline
+                .render(layer_index, bounds, &mut render_pass);
         }
 
         if !layer.meshes.is_empty() {
diff --git a/wgpu/src/buffer.rs b/wgpu/src/buffer.rs
index 7c092d0b..c210dd4e 100644
--- a/wgpu/src/buffer.rs
+++ b/wgpu/src/buffer.rs
@@ -1,3 +1,89 @@
 //! Utilities for buffer operations.
 pub mod dynamic;
 pub mod r#static;
+
+use std::marker::PhantomData;
+use std::ops::RangeBounds;
+
+#[derive(Debug)]
+pub struct Buffer<T> {
+    label: &'static str,
+    size: u64,
+    usage: wgpu::BufferUsages,
+    raw: wgpu::Buffer,
+    type_: PhantomData<T>,
+}
+
+impl<T: bytemuck::Pod> Buffer<T> {
+    pub fn new(
+        device: &wgpu::Device,
+        label: &'static str,
+        amount: usize,
+        usage: wgpu::BufferUsages,
+    ) -> Self {
+        let size = next_copy_size::<T>(amount);
+
+        let raw = device.create_buffer(&wgpu::BufferDescriptor {
+            label: Some(label),
+            size,
+            usage,
+            mapped_at_creation: false,
+        });
+
+        Self {
+            label,
+            size,
+            usage,
+            raw,
+            type_: PhantomData,
+        }
+    }
+
+    pub fn resize(&mut self, device: &wgpu::Device, new_count: usize) -> bool {
+        let new_size = (std::mem::size_of::<T>() * new_count) as u64;
+
+        if self.size < new_size {
+            self.raw = device.create_buffer(&wgpu::BufferDescriptor {
+                label: Some(self.label),
+                size: new_size,
+                usage: self.usage,
+                mapped_at_creation: false,
+            });
+
+            self.size = new_size;
+
+            true
+        } else {
+            false
+        }
+    }
+
+    pub fn write(
+        &self,
+        queue: &wgpu::Queue,
+        offset_count: usize,
+        contents: &[T],
+    ) {
+        queue.write_buffer(
+            &self.raw,
+            (std::mem::size_of::<T>() * offset_count) as u64,
+            bytemuck::cast_slice(contents),
+        );
+    }
+
+    pub fn slice(
+        &self,
+        bounds: impl RangeBounds<wgpu::BufferAddress>,
+    ) -> wgpu::BufferSlice<'_> {
+        self.raw.slice(bounds)
+    }
+}
+
+fn next_copy_size<T>(amount: usize) -> u64 {
+    let align_mask = wgpu::COPY_BUFFER_ALIGNMENT - 1;
+
+    (((std::mem::size_of::<T>() * amount).next_power_of_two() as u64
+        + align_mask)
+        & !align_mask)
+        .max(wgpu::COPY_BUFFER_ALIGNMENT)
+}
diff --git a/wgpu/src/quad.rs b/wgpu/src/quad.rs
index 2f5fcc6b..ebbe7a9d 100644
--- a/wgpu/src/quad.rs
+++ b/wgpu/src/quad.rs
@@ -1,3 +1,4 @@
+use crate::buffer::Buffer;
 use crate::Transformation;
 use iced_graphics::layer;
 use iced_native::Rectangle;
@@ -12,11 +13,11 @@ use tracing::info_span;
 #[derive(Debug)]
 pub struct Pipeline {
     pipeline: wgpu::RenderPipeline,
-    constants: wgpu::BindGroup,
-    constants_buffer: wgpu::Buffer,
+    constant_layout: wgpu::BindGroupLayout,
     vertices: wgpu::Buffer,
     indices: wgpu::Buffer,
-    instances: wgpu::Buffer,
+    layers: Vec<Layer>,
+    current_layer: usize,
 }
 
 impl Pipeline {
@@ -38,22 +39,6 @@ impl Pipeline {
                 }],
             });
 
-        let constants_buffer = device.create_buffer(&wgpu::BufferDescriptor {
-            label: Some("iced_wgpu::quad uniforms buffer"),
-            size: mem::size_of::<Uniforms>() as wgpu::BufferAddress,
-            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
-            mapped_at_creation: false,
-        });
-
-        let constants = device.create_bind_group(&wgpu::BindGroupDescriptor {
-            label: Some("iced_wgpu::quad uniforms bind group"),
-            layout: &constant_layout,
-            entries: &[wgpu::BindGroupEntry {
-                binding: 0,
-                resource: constants_buffer.as_entire_binding(),
-            }],
-        });
-
         let layout =
             device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                 label: Some("iced_wgpu::quad pipeline layout"),
@@ -148,118 +133,146 @@ impl Pipeline {
                 usage: wgpu::BufferUsages::INDEX,
             });
 
-        let instances = device.create_buffer(&wgpu::BufferDescriptor {
-            label: Some("iced_wgpu::quad instance buffer"),
-            size: mem::size_of::<layer::Quad>() as u64 * MAX_INSTANCES as u64,
-            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
-            mapped_at_creation: false,
-        });
-
         Pipeline {
             pipeline,
-            constants,
-            constants_buffer,
+            constant_layout,
             vertices,
             indices,
-            instances,
+            layers: Vec::new(),
+            current_layer: 0,
         }
     }
 
-    pub fn draw(
+    pub fn prepare(
         &mut self,
         device: &wgpu::Device,
-        staging_belt: &mut wgpu::util::StagingBelt,
-        encoder: &mut wgpu::CommandEncoder,
+        queue: &wgpu::Queue,
         instances: &[layer::Quad],
         transformation: Transformation,
         scale: f32,
-        bounds: Rectangle<u32>,
-        target: &wgpu::TextureView,
     ) {
-        #[cfg(feature = "tracing")]
-        let _ = info_span!("Wgpu::Quad", "DRAW").entered();
+        if self.layers.len() <= self.current_layer {
+            self.layers.push(Layer::new(device, &self.constant_layout));
+        }
 
-        let uniforms = Uniforms::new(transformation, scale);
+        let layer = &mut self.layers[self.current_layer];
+        layer.prepare(device, queue, instances, transformation, scale);
 
-        {
-            let mut constants_buffer = staging_belt.write_buffer(
-                encoder,
-                &self.constants_buffer,
-                0,
-                wgpu::BufferSize::new(mem::size_of::<Uniforms>() as u64)
-                    .unwrap(),
-                device,
+        self.current_layer += 1;
+    }
+
+    pub fn render<'a>(
+        &'a self,
+        layer: usize,
+        bounds: Rectangle<u32>,
+        render_pass: &mut wgpu::RenderPass<'a>,
+    ) {
+        if let Some(layer) = self.layers.get(layer) {
+            render_pass.set_pipeline(&self.pipeline);
+
+            render_pass.set_scissor_rect(
+                bounds.x,
+                bounds.y,
+                bounds.width,
+                bounds.height,
+            );
+
+            render_pass.set_index_buffer(
+                self.indices.slice(..),
+                wgpu::IndexFormat::Uint16,
             );
+            render_pass.set_vertex_buffer(0, self.vertices.slice(..));
 
-            constants_buffer.copy_from_slice(bytemuck::bytes_of(&uniforms));
+            layer.draw(render_pass);
         }
+    }
 
-        let mut i = 0;
-        let total = instances.len();
+    pub fn end_frame(&mut self) {
+        self.current_layer = 0;
+    }
+}
 
-        while i < total {
-            let end = (i + MAX_INSTANCES).min(total);
-            let amount = end - i;
+#[derive(Debug)]
+struct Layer {
+    constants: wgpu::BindGroup,
+    constants_buffer: wgpu::Buffer,
+    instances: Buffer<layer::Quad>,
+    instance_count: usize,
+}
 
-            let instance_bytes = bytemuck::cast_slice(&instances[i..end]);
+impl Layer {
+    pub fn new(
+        device: &wgpu::Device,
+        constant_layout: &wgpu::BindGroupLayout,
+    ) -> Self {
+        let constants_buffer = device.create_buffer(&wgpu::BufferDescriptor {
+            label: Some("iced_wgpu::quad uniforms buffer"),
+            size: mem::size_of::<Uniforms>() as wgpu::BufferAddress,
+            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
 
-            let mut instance_buffer = staging_belt.write_buffer(
-                encoder,
-                &self.instances,
-                0,
-                wgpu::BufferSize::new(instance_bytes.len() as u64).unwrap(),
-                device,
-            );
+        let constants = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("iced_wgpu::quad uniforms bind group"),
+            layout: &constant_layout,
+            entries: &[wgpu::BindGroupEntry {
+                binding: 0,
+                resource: constants_buffer.as_entire_binding(),
+            }],
+        });
 
-            instance_buffer.copy_from_slice(instance_bytes);
-
-            #[cfg(feature = "tracing")]
-            let _ = info_span!("Wgpu::Quad", "BEGIN_RENDER_PASS").enter();
-
-            {
-                let mut render_pass =
-                    encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
-                        label: Some("iced_wgpu::quad render pass"),
-                        color_attachments: &[Some(
-                            wgpu::RenderPassColorAttachment {
-                                view: target,
-                                resolve_target: None,
-                                ops: wgpu::Operations {
-                                    load: wgpu::LoadOp::Load,
-                                    store: true,
-                                },
-                            },
-                        )],
-                        depth_stencil_attachment: None,
-                    });
-
-                render_pass.set_pipeline(&self.pipeline);
-                render_pass.set_bind_group(0, &self.constants, &[]);
-                render_pass.set_index_buffer(
-                    self.indices.slice(..),
-                    wgpu::IndexFormat::Uint16,
-                );
-                render_pass.set_vertex_buffer(0, self.vertices.slice(..));
-                render_pass.set_vertex_buffer(1, self.instances.slice(..));
-
-                render_pass.set_scissor_rect(
-                    bounds.x,
-                    bounds.y,
-                    bounds.width,
-                    // TODO: Address anti-aliasing adjustments properly
-                    bounds.height,
-                );
-
-                render_pass.draw_indexed(
-                    0..QUAD_INDICES.len() as u32,
-                    0,
-                    0..amount as u32,
-                );
-            }
-
-            i += MAX_INSTANCES;
+        let instances = Buffer::new(
+            device,
+            "iced_wgpu::quad instance buffer",
+            MAX_INSTANCES,
+            wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
+        );
+
+        Self {
+            constants,
+            constants_buffer,
+            instances,
+            instance_count: 0,
         }
     }
+
+    pub fn prepare(
+        &mut self,
+        device: &wgpu::Device,
+        queue: &wgpu::Queue,
+        instances: &[layer::Quad],
+        transformation: Transformation,
+        scale: f32,
+    ) {
+        #[cfg(feature = "tracing")]
+        let _ = info_span!("Wgpu::Quad", "PREPARE").entered();
+
+        let uniforms = Uniforms::new(transformation, scale);
+
+        queue.write_buffer(
+            &self.constants_buffer,
+            0,
+            bytemuck::bytes_of(&uniforms),
+        );
+
+        let _ = self.instances.resize(device, instances.len());
+        self.instances.write(queue, 0, instances);
+        self.instance_count = instances.len();
+    }
+
+    pub fn draw<'a>(&'a self, render_pass: &mut wgpu::RenderPass<'a>) {
+        #[cfg(feature = "tracing")]
+        let _ = info_span!("Wgpu::Quad", "DRAW").entered();
+
+        render_pass.set_bind_group(0, &self.constants, &[]);
+        render_pass.set_vertex_buffer(1, self.instances.slice(..));
+
+        render_pass.draw_indexed(
+            0..QUAD_INDICES.len() as u32,
+            0,
+            0..self.instance_count as u32,
+        );
+    }
 }
 
 #[repr(C)]