Reintroduce support for custom primitives in `iced_wgpu`

6 files changed, 176 insertions, 140 deletions

diff --git a/wgpu/src/engine.rs b/wgpu/src/engine.rs
index e45b62b2..96cd6db8 100644
--- a/wgpu/src/engine.rs
+++ b/wgpu/src/engine.rs
@@ -1,19 +1,21 @@
 use crate::buffer;
 use crate::graphics::Antialiasing;
-use crate::primitive::pipeline;
+use crate::primitive;
 use crate::quad;
 use crate::text;
 use crate::triangle;
 
 #[allow(missing_debug_implementations)]
 pub struct Engine {
+    pub(crate) staging_belt: wgpu::util::StagingBelt,
+    pub(crate) format: wgpu::TextureFormat,
+
     pub(crate) quad_pipeline: quad::Pipeline,
     pub(crate) text_pipeline: text::Pipeline,
     pub(crate) triangle_pipeline: triangle::Pipeline,
-    pub(crate) _pipeline_storage: pipeline::Storage,
     #[cfg(any(feature = "image", feature = "svg"))]
     pub(crate) image_pipeline: crate::image::Pipeline,
-    pub(crate) staging_belt: wgpu::util::StagingBelt,
+    pub(crate) primitive_storage: primitive::Storage,
 }
 
 impl Engine {
@@ -43,13 +45,16 @@ impl Engine {
             staging_belt: wgpu::util::StagingBelt::new(
                 buffer::MAX_WRITE_SIZE as u64,
             ),
+            format,
+
             quad_pipeline,
             text_pipeline,
             triangle_pipeline,
-            _pipeline_storage: pipeline::Storage::default(),
 
             #[cfg(any(feature = "image", feature = "svg"))]
             image_pipeline,
+
+            primitive_storage: primitive::Storage::default(),
         }
     }
 
diff --git a/wgpu/src/layer.rs b/wgpu/src/layer.rs
index c8c27c61..7a18e322 100644
--- a/wgpu/src/layer.rs
+++ b/wgpu/src/layer.rs
@@ -4,6 +4,7 @@ use crate::graphics::color;
 use crate::graphics::text::{Editor, Paragraph};
 use crate::graphics::Mesh;
 use crate::image::{self, Image};
+use crate::primitive::{self, Primitive};
 use crate::quad::{self, Quad};
 use crate::text::{self, Text};
 use crate::triangle;
@@ -13,6 +14,7 @@ pub struct Layer {
     pub bounds: Rectangle,
     pub quads: quad::Batch,
     pub triangles: triangle::Batch,
+    pub primitives: primitive::Batch,
     pub text: text::Batch,
     pub images: image::Batch,
 }
@@ -23,6 +25,7 @@ impl Default for Layer {
             bounds: Rectangle::INFINITE,
             quads: quad::Batch::default(),
             triangles: triangle::Batch::default(),
+            primitives: primitive::Batch::default(),
             text: text::Batch::default(),
             images: image::Batch::default(),
         }
@@ -222,6 +225,18 @@ impl Stack {
         });
     }
 
+    pub fn draw_primitive(
+        &mut self,
+        bounds: Rectangle,
+        primitive: Box<dyn Primitive>,
+    ) {
+        let bounds = bounds * self.transformation();
+
+        self.layers[self.current]
+            .primitives
+            .push(primitive::Instance { bounds, primitive });
+    }
+
     pub fn push_clip(&mut self, bounds: Rectangle) {
         self.previous.push(self.current);
 
@@ -282,6 +297,7 @@ impl Stack {
 
             live.quads.clear();
             live.triangles.clear();
+            live.primitives.clear();
             live.text.clear();
             live.images.clear();
             pending_meshes.clear();
diff --git a/wgpu/src/lib.rs b/wgpu/src/lib.rs
index 030bcade..0580399d 100644
--- a/wgpu/src/lib.rs
+++ b/wgpu/src/lib.rs
@@ -101,8 +101,6 @@ impl Renderer {
         }
     }
 
-    pub fn draw_primitive(&mut self, _primitive: Primitive) {}
-
     pub fn present<T: AsRef<str>>(
         &mut self,
         engine: &mut Engine,
@@ -158,6 +156,19 @@ impl Renderer {
                 );
             }
 
+            if !layer.primitives.is_empty() {
+                for instance in &layer.primitives {
+                    instance.primitive.prepare(
+                        device,
+                        queue,
+                        engine.format,
+                        &mut engine.primitive_storage,
+                        &instance.bounds,
+                        viewport,
+                    );
+                }
+            }
+
             if !layer.text.is_empty() {
                 engine.text_pipeline.prepare(
                     device,
@@ -247,7 +258,9 @@ impl Renderer {
                 continue;
             };
 
-            let scissor_rect = physical_bounds.snap();
+            let Some(scissor_rect) = physical_bounds.snap() else {
+                continue;
+            };
 
             if !layer.quads.is_empty() {
                 engine.quad_pipeline.render(
@@ -293,6 +306,43 @@ impl Renderer {
                 ));
             }
 
+            if !layer.primitives.is_empty() {
+                let _ = ManuallyDrop::into_inner(render_pass);
+
+                for instance in &layer.primitives {
+                    if let Some(clip_bounds) = (instance.bounds * scale)
+                        .intersection(&physical_bounds)
+                        .and_then(Rectangle::snap)
+                    {
+                        instance.primitive.render(
+                            encoder,
+                            &engine.primitive_storage,
+                            frame,
+                            &clip_bounds,
+                        );
+                    }
+                }
+
+                render_pass = ManuallyDrop::new(encoder.begin_render_pass(
+                    &wgpu::RenderPassDescriptor {
+                        label: Some("iced_wgpu render pass"),
+                        color_attachments: &[Some(
+                            wgpu::RenderPassColorAttachment {
+                                view: frame,
+                                resolve_target: None,
+                                ops: wgpu::Operations {
+                                    load: wgpu::LoadOp::Load,
+                                    store: wgpu::StoreOp::Store,
+                                },
+                            },
+                        )],
+                        depth_stencil_attachment: None,
+                        timestamp_writes: None,
+                        occlusion_query_set: None,
+                    },
+                ));
+            }
+
             if !layer.text.is_empty() {
                 text_layer += engine.text_pipeline.render(
                     &self.text_storage,
@@ -520,6 +570,12 @@ impl graphics::geometry::Renderer for Renderer {
     }
 }
 
+impl primitive::Renderer for Renderer {
+    fn draw_primitive(&mut self, bounds: Rectangle, primitive: impl Primitive) {
+        self.layers.draw_primitive(bounds, Box::new(primitive));
+    }
+}
+
 impl graphics::compositor::Default for crate::Renderer {
     type Compositor = window::Compositor;
 }
diff --git a/wgpu/src/primitive.rs b/wgpu/src/primitive.rs
index 8e311d2b..4ba1ed9a 100644
--- a/wgpu/src/primitive.rs
+++ b/wgpu/src/primitive.rs
@@ -1,20 +1,95 @@
-//! Draw using different graphical primitives.
-pub mod pipeline;
+//! Draw custom primitives.
+use crate::core::{self, Rectangle};
+use crate::graphics::Viewport;
 
-pub use pipeline::Pipeline;
+use rustc_hash::FxHashMap;
+use std::any::{Any, TypeId};
+use std::fmt::Debug;
 
-use crate::graphics::Mesh;
+/// A batch of primitives.
+pub type Batch = Vec<Instance>;
 
-use std::fmt::Debug;
+/// A set of methods which allows a [`Primitive`] to be rendered.
+pub trait Primitive: Debug + Send + Sync + 'static {
+    /// Processes the [`Primitive`], allowing for GPU buffer allocation.
+    fn prepare(
+        &self,
+        device: &wgpu::Device,
+        queue: &wgpu::Queue,
+        format: wgpu::TextureFormat,
+        storage: &mut Storage,
+        bounds: &Rectangle,
+        viewport: &Viewport,
+    );
+
+    /// Renders the [`Primitive`].
+    fn render(
+        &self,
+        encoder: &mut wgpu::CommandEncoder,
+        storage: &Storage,
+        target: &wgpu::TextureView,
+        clip_bounds: &Rectangle<u32>,
+    );
+}
+
+#[derive(Debug)]
+/// An instance of a specific [`Primitive`].
+pub struct Instance {
+    /// The bounds of the [`Instance`].
+    pub bounds: Rectangle,
+
+    /// The [`Primitive`] to render.
+    pub primitive: Box<dyn Primitive>,
+}
+
+impl Instance {
+    /// Creates a new [`Pipeline`] with the given [`Primitive`].
+    pub fn new(bounds: Rectangle, primitive: impl Primitive) -> Self {
+        Instance {
+            bounds,
+            primitive: Box::new(primitive),
+        }
+    }
+}
+
+/// A renderer than can draw custom primitives.
+pub trait Renderer: core::Renderer {
+    /// Draws a custom primitive.
+    fn draw_primitive(&mut self, bounds: Rectangle, primitive: impl Primitive);
+}
+
+/// Stores custom, user-provided types.
+#[derive(Default, Debug)]
+pub struct Storage {
+    pipelines: FxHashMap<TypeId, Box<dyn Any + Send>>,
+}
+
+impl Storage {
+    /// Returns `true` if `Storage` contains a type `T`.
+    pub fn has<T: 'static>(&self) -> bool {
+        self.pipelines.get(&TypeId::of::<T>()).is_some()
+    }
+
+    /// Inserts the data `T` in to [`Storage`].
+    pub fn store<T: 'static + Send>(&mut self, data: T) {
+        let _ = self.pipelines.insert(TypeId::of::<T>(), Box::new(data));
+    }
 
-/// The graphical primitives supported by `iced_wgpu`.
-pub type Primitive = crate::graphics::Primitive<Custom>;
+    /// Returns a reference to the data with type `T` if it exists in [`Storage`].
+    pub fn get<T: 'static>(&self) -> Option<&T> {
+        self.pipelines.get(&TypeId::of::<T>()).map(|pipeline| {
+            pipeline
+                .downcast_ref::<T>()
+                .expect("Pipeline with this type does not exist in Storage.")
+        })
+    }
 
-/// The custom primitives supported by `iced_wgpu`.
-#[derive(Debug, Clone, PartialEq)]
-pub enum Custom {
-    /// A mesh primitive.
-    Mesh(Mesh),
-    /// A custom pipeline primitive.
-    Pipeline(Pipeline),
+    /// Returns a mutable reference to the data with type `T` if it exists in [`Storage`].
+    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
+        self.pipelines.get_mut(&TypeId::of::<T>()).map(|pipeline| {
+            pipeline
+                .downcast_mut::<T>()
+                .expect("Pipeline with this type does not exist in Storage.")
+        })
+    }
 }
diff --git a/wgpu/src/primitive/pipeline.rs b/wgpu/src/primitive/pipeline.rs
index 59c54db9..8b137891 100644
--- a/wgpu/src/primitive/pipeline.rs
+++ b/wgpu/src/primitive/pipeline.rs
@@ -1,116 +1 @@
-//! Draw primitives using custom pipelines.
-use crate::core::{self, Rectangle, Size};
 
-use rustc_hash::FxHashMap;
-use std::any::{Any, TypeId};
-use std::fmt::Debug;
-use std::sync::Arc;
-
-#[derive(Clone, Debug)]
-/// A custom primitive which can be used to render primitives associated with a custom pipeline.
-pub struct Pipeline {
-    /// The bounds of the [`Pipeline`].
-    pub bounds: Rectangle,
-
-    /// The [`Primitive`] to render.
-    pub primitive: Arc<dyn Primitive>,
-}
-
-impl Pipeline {
-    /// Creates a new [`Pipeline`] with the given [`Primitive`].
-    pub fn new(bounds: Rectangle, primitive: impl Primitive) -> Self {
-        Pipeline {
-            bounds,
-            primitive: Arc::new(primitive),
-        }
-    }
-}
-
-impl PartialEq for Pipeline {
-    fn eq(&self, other: &Self) -> bool {
-        self.primitive.type_id() == other.primitive.type_id()
-    }
-}
-
-/// A set of methods which allows a [`Primitive`] to be rendered.
-pub trait Primitive: Debug + Send + Sync + 'static {
-    /// Processes the [`Primitive`], allowing for GPU buffer allocation.
-    fn prepare(
-        &self,
-        format: wgpu::TextureFormat,
-        device: &wgpu::Device,
-        queue: &wgpu::Queue,
-        bounds: Rectangle,
-        target_size: Size<u32>,
-        scale_factor: f32,
-        storage: &mut Storage,
-    );
-
-    /// Renders the [`Primitive`].
-    fn render(
-        &self,
-        storage: &Storage,
-        target: &wgpu::TextureView,
-        target_size: Size<u32>,
-        viewport: Rectangle<u32>,
-        encoder: &mut wgpu::CommandEncoder,
-    );
-}
-
-/// A renderer than can draw custom pipeline primitives.
-pub trait Renderer: core::Renderer {
-    /// Draws a custom pipeline primitive.
-    fn draw_pipeline_primitive(
-        &mut self,
-        bounds: Rectangle,
-        primitive: impl Primitive,
-    );
-}
-
-impl Renderer for crate::Renderer {
-    fn draw_pipeline_primitive(
-        &mut self,
-        bounds: Rectangle,
-        primitive: impl Primitive,
-    ) {
-        self.draw_primitive(super::Primitive::Custom(super::Custom::Pipeline(
-            Pipeline::new(bounds, primitive),
-        )));
-    }
-}
-
-/// Stores custom, user-provided pipelines.
-#[derive(Default, Debug)]
-pub struct Storage {
-    pipelines: FxHashMap<TypeId, Box<dyn Any + Send>>,
-}
-
-impl Storage {
-    /// Returns `true` if `Storage` contains a pipeline with type `T`.
-    pub fn has<T: 'static>(&self) -> bool {
-        self.pipelines.get(&TypeId::of::<T>()).is_some()
-    }
-
-    /// Inserts the pipeline `T` in to [`Storage`].
-    pub fn store<T: 'static + Send>(&mut self, pipeline: T) {
-        let _ = self.pipelines.insert(TypeId::of::<T>(), Box::new(pipeline));
-    }
-
-    /// Returns a reference to pipeline with type `T` if it exists in [`Storage`].
-    pub fn get<T: 'static>(&self) -> Option<&T> {
-        self.pipelines.get(&TypeId::of::<T>()).map(|pipeline| {
-            pipeline
-                .downcast_ref::<T>()
-                .expect("Pipeline with this type does not exist in Storage.")
-        })
-    }
-
-    /// Returns a mutable reference to pipeline `T` if it exists in [`Storage`].
-    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
-        self.pipelines.get_mut(&TypeId::of::<T>()).map(|pipeline| {
-            pipeline
-                .downcast_mut::<T>()
-                .expect("Pipeline with this type does not exist in Storage.")
-        })
-    }
-}
diff --git a/wgpu/src/triangle.rs b/wgpu/src/triangle.rs
index 98ba41b3..8470ea39 100644
--- a/wgpu/src/triangle.rs
+++ b/wgpu/src/triangle.rs
@@ -501,14 +501,13 @@ impl Layer {
         let mut last_is_solid = None;
 
         for (index, mesh) in meshes.iter().enumerate() {
-            let Some(clip_bounds) =
-                bounds.intersection(&(mesh.clip_bounds() * transformation))
+            let Some(clip_bounds) = bounds
+                .intersection(&(mesh.clip_bounds() * transformation))
+                .and_then(Rectangle::snap)
             else {
                 continue;
             };
 
-            let clip_bounds = clip_bounds.snap();
-
             render_pass.set_scissor_rect(
                 clip_bounds.x,
                 clip_bounds.y,