Refactor `image::Pipeline` into `prepare` and `render` architecture

4 files changed, 190 insertions, 147 deletions

diff --git a/wgpu/src/backend.rs b/wgpu/src/backend.rs
index 395d28d5..90a511ef 100644
--- a/wgpu/src/backend.rs
+++ b/wgpu/src/backend.rs
@@ -106,7 +106,7 @@ impl Backend {
         self.text_pipeline.end_frame();
 
         #[cfg(any(feature = "image", feature = "svg"))]
-        self.image_pipeline.trim_cache(device, encoder);
+        self.image_pipeline.end_frame(device, queue, encoder);
     }
 
     fn flush(
@@ -177,16 +177,32 @@ impl Backend {
                 let scaled = transformation
                     * Transformation::scale(scale_factor, scale_factor);
 
-                self.image_pipeline.draw(
+                self.image_pipeline.prepare(
                     device,
-                    staging_belt,
+                    queue,
                     encoder,
                     &layer.images,
                     scaled,
-                    bounds,
-                    target,
                     scale_factor,
                 );
+
+                let mut render_pass =
+                    encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+                        label: Some("iced_wgpu::image 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.image_pipeline.render(bounds, &mut render_pass);
             }
         }
 
diff --git a/wgpu/src/image.rs b/wgpu/src/image.rs
index a5e63b17..f6c39f1a 100644
--- a/wgpu/src/image.rs
+++ b/wgpu/src/image.rs
@@ -6,7 +6,7 @@ use iced_graphics::image::raster;
 #[cfg(feature = "svg")]
 use iced_graphics::image::vector;
 
-use crate::Transformation;
+use crate::{Buffer, Transformation};
 use atlas::Atlas;
 
 use iced_graphics::layer;
@@ -34,15 +34,108 @@ pub struct Pipeline {
     vector_cache: RefCell<vector::Cache<Atlas>>,
 
     pipeline: wgpu::RenderPipeline,
-    uniforms: wgpu::Buffer,
     vertices: wgpu::Buffer,
     indices: wgpu::Buffer,
-    instances: wgpu::Buffer,
-    constants: wgpu::BindGroup,
+    sampler: wgpu::Sampler,
     texture: wgpu::BindGroup,
     texture_version: usize,
-    texture_layout: wgpu::BindGroupLayout,
     texture_atlas: Atlas,
+    texture_layout: wgpu::BindGroupLayout,
+    constant_layout: wgpu::BindGroupLayout,
+
+    layers: Vec<Layer>,
+    prepare_layer: usize,
+    render_layer: usize,
+}
+
+#[derive(Debug)]
+struct Layer {
+    uniforms: wgpu::Buffer,
+    constants: wgpu::BindGroup,
+    instances: Buffer<Instance>,
+    instance_count: usize,
+}
+
+impl Layer {
+    fn new(
+        device: &wgpu::Device,
+        constant_layout: &wgpu::BindGroupLayout,
+        sampler: &wgpu::Sampler,
+    ) -> Self {
+        let uniforms = device.create_buffer(&wgpu::BufferDescriptor {
+            label: Some("iced_wgpu::image uniforms buffer"),
+            size: mem::size_of::<Uniforms>() as u64,
+            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+
+        let constants = device.create_bind_group(&wgpu::BindGroupDescriptor {
+            label: Some("iced_wgpu::image constants bind group"),
+            layout: &constant_layout,
+            entries: &[
+                wgpu::BindGroupEntry {
+                    binding: 0,
+                    resource: wgpu::BindingResource::Buffer(
+                        wgpu::BufferBinding {
+                            buffer: &uniforms,
+                            offset: 0,
+                            size: None,
+                        },
+                    ),
+                },
+                wgpu::BindGroupEntry {
+                    binding: 1,
+                    resource: wgpu::BindingResource::Sampler(&sampler),
+                },
+            ],
+        });
+
+        let instances = Buffer::new(
+            device,
+            "iced_wgpu::image instance buffer",
+            Instance::INITIAL,
+            wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
+        );
+
+        Self {
+            uniforms,
+            constants,
+            instances,
+            instance_count: 0,
+        }
+    }
+
+    fn prepare(
+        &mut self,
+        device: &wgpu::Device,
+        queue: &wgpu::Queue,
+        instances: &[Instance],
+        transformation: Transformation,
+    ) {
+        queue.write_buffer(
+            &self.uniforms,
+            0,
+            bytemuck::bytes_of(&Uniforms {
+                transform: transformation.into(),
+            }),
+        );
+
+        let _ = self.instances.resize(device, instances.len());
+        self.instances.write(queue, 0, instances);
+
+        self.instance_count = instances.len();
+    }
+
+    fn render<'a>(&'a self, render_pass: &mut wgpu::RenderPass<'a>) {
+        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,
+        );
+    }
 }
 
 impl Pipeline {
@@ -86,35 +179,6 @@ impl Pipeline {
                 ],
             });
 
-        let uniforms_buffer = device.create_buffer(&wgpu::BufferDescriptor {
-            label: Some("iced_wgpu::image uniforms buffer"),
-            size: mem::size_of::<Uniforms>() as u64,
-            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
-            mapped_at_creation: false,
-        });
-
-        let constant_bind_group =
-            device.create_bind_group(&wgpu::BindGroupDescriptor {
-                label: Some("iced_wgpu::image constants bind group"),
-                layout: &constant_layout,
-                entries: &[
-                    wgpu::BindGroupEntry {
-                        binding: 0,
-                        resource: wgpu::BindingResource::Buffer(
-                            wgpu::BufferBinding {
-                                buffer: &uniforms_buffer,
-                                offset: 0,
-                                size: None,
-                            },
-                        ),
-                    },
-                    wgpu::BindGroupEntry {
-                        binding: 1,
-                        resource: wgpu::BindingResource::Sampler(&sampler),
-                    },
-                ],
-            });
-
         let texture_layout =
             device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                 label: Some("iced_wgpu::image texture atlas layout"),
@@ -225,13 +289,6 @@ impl Pipeline {
                 usage: wgpu::BufferUsages::INDEX,
             });
 
-        let instances = device.create_buffer(&wgpu::BufferDescriptor {
-            label: Some("iced_wgpu::image instance buffer"),
-            size: mem::size_of::<Instance>() as u64 * Instance::MAX as u64,
-            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
-            mapped_at_creation: false,
-        });
-
         let texture_atlas = Atlas::new(device);
 
         let texture = device.create_bind_group(&wgpu::BindGroupDescriptor {
@@ -253,15 +310,18 @@ impl Pipeline {
             vector_cache: RefCell::new(vector::Cache::default()),
 
             pipeline,
-            uniforms: uniforms_buffer,
             vertices,
             indices,
-            instances,
-            constants: constant_bind_group,
+            sampler,
             texture,
             texture_version: texture_atlas.layer_count(),
-            texture_layout,
             texture_atlas,
+            texture_layout,
+            constant_layout,
+
+            layers: Vec::new(),
+            prepare_layer: 0,
+            render_layer: 0,
         }
     }
 
@@ -281,18 +341,19 @@ impl Pipeline {
         svg.viewport_dimensions()
     }
 
-    pub fn draw(
+    pub fn prepare(
         &mut self,
         device: &wgpu::Device,
-        staging_belt: &mut wgpu::util::StagingBelt,
+        queue: &wgpu::Queue,
         encoder: &mut wgpu::CommandEncoder,
         images: &[layer::Image],
         transformation: Transformation,
-        bounds: Rectangle<u32>,
-        target: &wgpu::TextureView,
         _scale: f32,
     ) {
         #[cfg(feature = "tracing")]
+        let _ = info_span!("Wgpu::Image", "PREPARE").entered();
+
+        #[cfg(feature = "tracing")]
         let _ = info_span!("Wgpu::Image", "DRAW").entered();
 
         let instances: &mut Vec<Instance> = &mut Vec::new();
@@ -309,7 +370,7 @@ impl Pipeline {
                 layer::Image::Raster { handle, bounds } => {
                     if let Some(atlas_entry) = raster_cache.upload(
                         handle,
-                        &mut (device, encoder),
+                        &mut (device, queue, encoder),
                         &mut self.texture_atlas,
                     ) {
                         add_instances(
@@ -336,7 +397,7 @@ impl Pipeline {
                         *color,
                         size,
                         _scale,
-                        &mut (device, encoder),
+                        &mut (device, queue, encoder),
                         &mut self.texture_atlas,
                     ) {
                         add_instances(
@@ -376,68 +437,27 @@ impl Pipeline {
             self.texture_version = texture_version;
         }
 
-        {
-            let mut uniforms_buffer = staging_belt.write_buffer(
-                encoder,
-                &self.uniforms,
-                0,
-                wgpu::BufferSize::new(mem::size_of::<Uniforms>() as u64)
-                    .unwrap(),
+        if self.layers.len() <= self.prepare_layer {
+            self.layers.push(Layer::new(
                 device,
-            );
-
-            uniforms_buffer.copy_from_slice(bytemuck::bytes_of(&Uniforms {
-                transform: transformation.into(),
-            }));
+                &self.constant_layout,
+                &self.sampler,
+            ));
         }
 
-        let mut i = 0;
-        let total = instances.len();
-
-        while i < total {
-            let end = (i + Instance::MAX).min(total);
-            let amount = end - i;
-
-            let mut instances_buffer = staging_belt.write_buffer(
-                encoder,
-                &self.instances,
-                0,
-                wgpu::BufferSize::new(
-                    (amount * std::mem::size_of::<Instance>()) as u64,
-                )
-                .unwrap(),
-                device,
-            );
-
-            instances_buffer.copy_from_slice(bytemuck::cast_slice(
-                &instances[i..i + amount],
-            ));
+        let layer = &mut self.layers[self.prepare_layer];
+        layer.prepare(device, queue, instances, transformation);
 
-            let mut render_pass =
-                encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
-                    label: Some("iced_wgpu::image 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.prepare_layer += 1;
+    }
 
+    pub fn render<'a>(
+        &'a mut self,
+        bounds: Rectangle<u32>,
+        render_pass: &mut wgpu::RenderPass<'a>,
+    ) {
+        if let Some(layer) = self.layers.get(self.render_layer) {
             render_pass.set_pipeline(&self.pipeline);
-            render_pass.set_bind_group(0, &self.constants, &[]);
-            render_pass.set_bind_group(1, &self.texture, &[]);
-            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,
@@ -446,30 +466,37 @@ impl Pipeline {
                 bounds.height,
             );
 
-            render_pass.draw_indexed(
-                0..QUAD_INDICES.len() as u32,
-                0,
-                0..amount as u32,
+            render_pass.set_bind_group(1, &self.texture, &[]);
+            render_pass.set_index_buffer(
+                self.indices.slice(..),
+                wgpu::IndexFormat::Uint16,
             );
+            render_pass.set_vertex_buffer(0, self.vertices.slice(..));
 
-            i += Instance::MAX;
+            layer.render(render_pass);
+
+            self.render_layer += 1;
         }
     }
 
-    pub fn trim_cache(
+    pub fn end_frame(
         &mut self,
         device: &wgpu::Device,
+        queue: &wgpu::Queue,
         encoder: &mut wgpu::CommandEncoder,
     ) {
         #[cfg(feature = "image")]
         self.raster_cache
             .borrow_mut()
-            .trim(&mut self.texture_atlas, &mut (device, encoder));
+            .trim(&mut self.texture_atlas, &mut (device, queue, encoder));
 
         #[cfg(feature = "svg")]
         self.vector_cache
             .borrow_mut()
-            .trim(&mut self.texture_atlas, &mut (device, encoder));
+            .trim(&mut self.texture_atlas, &mut (device, queue, encoder));
+
+        self.prepare_layer = 0;
+        self.render_layer = 0;
     }
 }
 
@@ -507,7 +534,7 @@ struct Instance {
 }
 
 impl Instance {
-    pub const MAX: usize = 1_000;
+    pub const INITIAL: usize = 1_000;
 }
 
 #[repr(C)]
diff --git a/wgpu/src/image/atlas.rs b/wgpu/src/image/atlas.rs
index eafe2f96..7df67abd 100644
--- a/wgpu/src/image/atlas.rs
+++ b/wgpu/src/image/atlas.rs
@@ -185,13 +185,13 @@ impl Atlas {
 
     fn upload_allocation(
         &mut self,
-        buffer: &wgpu::Buffer,
+        data: &[u8],
         image_width: u32,
         image_height: u32,
         padding: u32,
         offset: usize,
         allocation: &Allocation,
-        encoder: &mut wgpu::CommandEncoder,
+        queue: &wgpu::Queue,
     ) {
         let (x, y) = allocation.position();
         let Size { width, height } = allocation.size();
@@ -203,15 +203,7 @@ impl Atlas {
             depth_or_array_layers: 1,
         };
 
-        encoder.copy_buffer_to_texture(
-            wgpu::ImageCopyBuffer {
-                buffer,
-                layout: wgpu::ImageDataLayout {
-                    offset: offset as u64,
-                    bytes_per_row: NonZeroU32::new(4 * image_width + padding),
-                    rows_per_image: NonZeroU32::new(image_height),
-                },
-            },
+        queue.write_texture(
             wgpu::ImageCopyTexture {
                 texture: &self.texture,
                 mip_level: 0,
@@ -222,6 +214,12 @@ impl Atlas {
                 },
                 aspect: wgpu::TextureAspect::default(),
             },
+            data,
+            wgpu::ImageDataLayout {
+                offset: offset as u64,
+                bytes_per_row: NonZeroU32::new(4 * image_width + padding),
+                rows_per_image: NonZeroU32::new(image_height),
+            },
             extent,
         );
     }
@@ -301,17 +299,19 @@ impl Atlas {
 
 impl image::Storage for Atlas {
     type Entry = Entry;
-    type State<'a> = (&'a wgpu::Device, &'a mut wgpu::CommandEncoder);
+    type State<'a> = (
+        &'a wgpu::Device,
+        &'a wgpu::Queue,
+        &'a mut wgpu::CommandEncoder,
+    );
 
     fn upload(
         &mut self,
         width: u32,
         height: u32,
         data: &[u8],
-        (device, encoder): &mut Self::State<'_>,
+        (device, queue, encoder): &mut Self::State<'_>,
     ) -> Option<Self::Entry> {
-        use wgpu::util::DeviceExt;
-
         let entry = {
             let current_size = self.layers.len();
             let entry = self.allocate(width, height)?;
@@ -344,17 +344,16 @@ impl image::Storage for Atlas {
             )
         }
 
-        let buffer =
-            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
-                label: Some("iced_wgpu::image staging buffer"),
-                contents: &padded_data,
-                usage: wgpu::BufferUsages::COPY_SRC,
-            });
-
         match &entry {
             Entry::Contiguous(allocation) => {
                 self.upload_allocation(
-                    &buffer, width, height, padding, 0, allocation, encoder,
+                    &padded_data,
+                    width,
+                    height,
+                    padding,
+                    0,
+                    allocation,
+                    queue,
                 );
             }
             Entry::Fragmented { fragments, .. } => {
@@ -363,13 +362,13 @@ impl image::Storage for Atlas {
                     let offset = (y * padded_width as u32 + 4 * x) as usize;
 
                     self.upload_allocation(
-                        &buffer,
+                        &padded_data,
                         width,
                         height,
                         padding,
                         offset,
                         &fragment.allocation,
-                        encoder,
+                        queue,
                     );
                 }
             }
diff --git a/wgpu/src/lib.rs b/wgpu/src/lib.rs
index 6d6d3fd6..9da40572 100644
--- a/wgpu/src/lib.rs
+++ b/wgpu/src/lib.rs
@@ -56,7 +56,8 @@ pub use wgpu;
 pub use backend::Backend;
 pub use settings::Settings;
 
-pub(crate) use iced_graphics::Transformation;
+use crate::buffer::Buffer;
+use iced_graphics::Transformation;
 
 #[cfg(any(feature = "image", feature = "svg"))]
 mod image;