Added support for custom shader widget for iced_wgpu backend.

6 files changed, 182 insertions, 3 deletions

diff --git a/wgpu/src/backend.rs b/wgpu/src/backend.rs
index 2bd29f42..907611d9 100644
--- a/wgpu/src/backend.rs
+++ b/wgpu/src/backend.rs
@@ -3,9 +3,7 @@ use crate::graphics::backend;
 use crate::graphics::color;
 use crate::graphics::{Transformation, Viewport};
 use crate::primitive::{self, Primitive};
-use crate::quad;
-use crate::text;
-use crate::triangle;
+use crate::{custom, quad, text, triangle};
 use crate::{Layer, Settings};
 
 #[cfg(feature = "tracing")]
@@ -25,6 +23,7 @@ pub struct Backend {
     quad_pipeline: quad::Pipeline,
     text_pipeline: text::Pipeline,
     triangle_pipeline: triangle::Pipeline,
+    pipeline_storage: custom::Storage,
 
     #[cfg(any(feature = "image", feature = "svg"))]
     image_pipeline: image::Pipeline,
@@ -50,6 +49,7 @@ impl Backend {
             quad_pipeline,
             text_pipeline,
             triangle_pipeline,
+            pipeline_storage: custom::Storage::default(),
 
             #[cfg(any(feature = "image", feature = "svg"))]
             image_pipeline,
@@ -66,6 +66,7 @@ impl Backend {
         queue: &wgpu::Queue,
         encoder: &mut wgpu::CommandEncoder,
         clear_color: Option<Color>,
+        format: wgpu::TextureFormat,
         frame: &wgpu::TextureView,
         primitives: &[Primitive],
         viewport: &Viewport,
@@ -88,6 +89,7 @@ impl Backend {
         self.prepare(
             device,
             queue,
+            format,
             encoder,
             scale_factor,
             target_size,
@@ -117,6 +119,7 @@ impl Backend {
         &mut self,
         device: &wgpu::Device,
         queue: &wgpu::Queue,
+        format: wgpu::TextureFormat,
         _encoder: &mut wgpu::CommandEncoder,
         scale_factor: f32,
         target_size: Size<u32>,
@@ -179,6 +182,20 @@ impl Backend {
                     target_size,
                 );
             }
+
+            if !layer.shaders.is_empty() {
+                for shader in &layer.shaders {
+                    shader.primitive.prepare(
+                        format,
+                        device,
+                        queue,
+                        target_size,
+                        scale_factor,
+                        transformation,
+                        &mut self.pipeline_storage,
+                    );
+                }
+            }
         }
     }
 
@@ -302,6 +319,47 @@ impl Backend {
 
                 text_layer += 1;
             }
+
+            // kill render pass to let custom shaders get mut access to encoder
+            let _ = ManuallyDrop::into_inner(render_pass);
+
+            if !layer.shaders.is_empty() {
+                for shader in &layer.shaders {
+                    //This extra check is needed since each custom pipeline must set it's own
+                    //scissor rect, which will panic if bounds.w/h < 1
+                    let bounds = shader.bounds * scale_factor;
+
+                    if bounds.width < 1.0 || bounds.height < 1.0 {
+                        continue;
+                    }
+
+                    shader.primitive.render(
+                        &self.pipeline_storage,
+                        bounds.into(),
+                        target,
+                        target_size,
+                        encoder,
+                    );
+                }
+            }
+
+            // recreate and continue processing layers
+            render_pass = ManuallyDrop::new(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,
+                },
+            ));
         }
 
         let _ = ManuallyDrop::into_inner(render_pass);
diff --git a/wgpu/src/custom.rs b/wgpu/src/custom.rs
new file mode 100644
index 00000000..65dd0496
--- /dev/null
+++ b/wgpu/src/custom.rs
@@ -0,0 +1,66 @@
+use crate::core::{Rectangle, Size};
+use crate::graphics::Transformation;
+use std::any::{Any, TypeId};
+use std::collections::HashMap;
+use std::fmt::Debug;
+
+/// Stores custom, user-provided pipelines.
+#[derive(Default, Debug)]
+pub struct Storage {
+    pipelines: HashMap<TypeId, Box<dyn Any>>,
+}
+
+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>(&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.")
+        })
+    }
+}
+
+/// 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,
+        target_size: Size<u32>,
+        scale_factor: f32,
+        transform: Transformation,
+        storage: &mut Storage,
+    );
+
+    /// Renders the [`Primitive`].
+    fn render(
+        &self,
+        storage: &Storage,
+        bounds: Rectangle<u32>,
+        target: &wgpu::TextureView,
+        target_size: Size<u32>,
+        encoder: &mut wgpu::CommandEncoder,
+    );
+}
diff --git a/wgpu/src/layer.rs b/wgpu/src/layer.rs
index 286801e6..d451cbfd 100644
--- a/wgpu/src/layer.rs
+++ b/wgpu/src/layer.rs
@@ -34,6 +34,9 @@ pub struct Layer<'a> {
 
     /// The images of the [`Layer`].
     pub images: Vec<Image>,
+
+    /// The custom shader primitives of this [`Layer`].
+    pub shaders: Vec<primitive::Shader>,
 }
 
 impl<'a> Layer<'a> {
@@ -45,6 +48,7 @@ impl<'a> Layer<'a> {
             meshes: Vec::new(),
             text: Vec::new(),
             images: Vec::new(),
+            shaders: Vec::new(),
         }
     }
 
@@ -308,6 +312,18 @@ impl<'a> Layer<'a> {
                         }
                     }
                 },
+                primitive::Custom::Shader(shader) => {
+                    let layer = &mut layers[current_layer];
+
+                    let bounds = Rectangle::new(
+                        Point::new(translation.x, translation.y),
+                        shader.bounds.size(),
+                    );
+
+                    if layer.bounds.intersection(&bounds).is_some() {
+                        layer.shaders.push(shader.clone());
+                    }
+                }
             },
         }
     }
diff --git a/wgpu/src/lib.rs b/wgpu/src/lib.rs
index 424dfeb3..13d8e886 100644
--- a/wgpu/src/lib.rs
+++ b/wgpu/src/lib.rs
@@ -29,6 +29,7 @@
     rustdoc::broken_intra_doc_links
 )]
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
+pub mod custom;
 pub mod layer;
 pub mod primitive;
 pub mod settings;
diff --git a/wgpu/src/primitive.rs b/wgpu/src/primitive.rs
index 8dbf3008..4347dcda 100644
--- a/wgpu/src/primitive.rs
+++ b/wgpu/src/primitive.rs
@@ -1,6 +1,10 @@
 //! Draw using different graphical primitives.
 use crate::core::Rectangle;
+use crate::custom;
 use crate::graphics::{Damage, Mesh};
+use std::any::Any;
+use std::fmt::Debug;
+use std::sync::Arc;
 
 /// The graphical primitives supported by `iced_wgpu`.
 pub type Primitive = crate::graphics::Primitive<Custom>;
@@ -10,12 +14,44 @@ pub type Primitive = crate::graphics::Primitive<Custom>;
 pub enum Custom {
     /// A mesh primitive.
     Mesh(Mesh),
+    /// A custom shader primitive
+    Shader(Shader),
+}
+
+impl Custom {
+    /// Create a custom [`Shader`] primitive.
+    pub fn shader<P: custom::Primitive>(
+        bounds: Rectangle,
+        primitive: P,
+    ) -> Self {
+        Self::Shader(Shader {
+            bounds,
+            primitive: Arc::new(primitive),
+        })
+    }
 }
 
 impl Damage for Custom {
     fn bounds(&self) -> Rectangle {
         match self {
             Self::Mesh(mesh) => mesh.bounds(),
+            Self::Shader(shader) => shader.bounds,
         }
     }
 }
+
+#[derive(Clone, Debug)]
+/// A custom primitive which can be used to render primitives associated with a custom pipeline.
+pub struct Shader {
+    /// The bounds of the [`Shader`].
+    pub bounds: Rectangle,
+
+    /// The [`custom::Primitive`] to render.
+    pub primitive: Arc<dyn custom::Primitive>,
+}
+
+impl PartialEq for Shader {
+    fn eq(&self, other: &Self) -> bool {
+        self.primitive.type_id() == other.primitive.type_id()
+    }
+}
diff --git a/wgpu/src/window/compositor.rs b/wgpu/src/window/compositor.rs
index 1ddbe5fe..90d64e17 100644
--- a/wgpu/src/window/compositor.rs
+++ b/wgpu/src/window/compositor.rs
@@ -178,6 +178,7 @@ pub fn present<Theme, T: AsRef<str>>(
                 &compositor.queue,
                 &mut encoder,
                 Some(background_color),
+                frame.texture.format(),
                 view,
                 primitives,
                 viewport,
@@ -357,6 +358,7 @@ pub fn screenshot<Theme, T: AsRef<str>>(
         &compositor.queue,
         &mut encoder,
         Some(background_color),
+        texture.format(),
         &view,
         primitives,
         viewport,