Re-implement PaintCallbacks With Support for WGPU (#1684)

* Re-implement PaintCallbacks With Support for WGPU

This makes breaking changes to the PaintCallback system, but makes it
flexible enough to support both the WGPU and glow backends with custom
rendering.

Also adds a WGPU equivalent to the glow demo for custom painting.

egui-wgpu/src/lib.rs

@@ -6,7 +6,10 @@ pub use wgpu;
 
 /// Low-level painting of [`egui`] on [`wgpu`].
 pub mod renderer;
+pub use renderer::CallbackFn;
 
 /// Module for painting [`egui`] with [`wgpu`] on [`winit`].
 #[cfg(feature = "winit")]
 pub mod winit;
+#[cfg(feature = "winit")]
+pub use crate::winit::RenderState;

egui-wgpu/src/renderer.rs

@@ -2,10 +2,79 @@
 
 use std::{borrow::Cow, collections::HashMap, num::NonZeroU32};
 
-use egui::epaint::Primitive;
+use egui::{epaint::Primitive, PaintCallbackInfo};
+use type_map::TypeMap;
 use wgpu;
 use wgpu::util::DeviceExt as _;
 
+/// A callback function that can be used to compose an [`egui::PaintCallback`] for custom WGPU
+/// rendering.
+///
+/// The callback is composed of two functions: `prepare` and `paint`.
+///
+/// `prepare` is called every frame before `paint`, and can use the passed-in [`wgpu::Device`] and
+/// [`wgpu::Buffer`] to allocate or modify GPU resources such as buffers.
+///
+/// `paint` is called after `prepare` and is given access to the the [`wgpu::RenderPass`] so that it
+/// can issue draw commands.
+///
+/// The final argument of both the `prepare` and `paint` callbacks is a the
+/// [`paint_callback_resources`][crate::renderer::RenderPass::paint_callback_resources].
+/// `paint_callback_resources` has the same lifetime as the Egui render pass, so it can be used to
+/// store buffers, pipelines, and other information that needs to be accessed during the render
+/// pass.
+///
+/// # Example
+///
+/// See the [custom3d_wgpu] demo source for a detailed usage example.
+///
+/// [custom3d_wgpu]:
+///     https://github.com/emilk/egui/blob/master/egui_demo_app/src/apps/custom3d_wgpu.rs
+pub struct CallbackFn {
+    prepare: Box<PrepareCallback>,
+    paint: Box<PaintCallback>,
+}
+
+type PrepareCallback = dyn Fn(&wgpu::Device, &wgpu::Queue, &mut TypeMap) + Sync + Send;
+type PaintCallback =
+    dyn for<'a, 'b> Fn(PaintCallbackInfo, &'a mut wgpu::RenderPass<'b>, &'b TypeMap) + Sync + Send;
+
+impl Default for CallbackFn {
+    fn default() -> Self {
+        CallbackFn {
+            prepare: Box::new(|_, _, _| ()),
+            paint: Box::new(|_, _, _| ()),
+        }
+    }
+}
+
+impl CallbackFn {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Set the prepare callback
+    pub fn prepare<F>(mut self, prepare: F) -> Self
+    where
+        F: Fn(&wgpu::Device, &wgpu::Queue, &mut TypeMap) + Sync + Send + 'static,
+    {
+        self.prepare = Box::new(prepare) as _;
+        self
+    }
+
+    /// Set the paint callback
+    pub fn paint<F>(mut self, paint: F) -> Self
+    where
+        F: for<'a, 'b> Fn(PaintCallbackInfo, &'a mut wgpu::RenderPass<'b>, &'b TypeMap)
+            + Sync
+            + Send
+            + 'static,
+    {
+        self.paint = Box::new(paint) as _;
+        self
+    }
+}
+
 /// Enum for selecting the right buffer type.
 #[derive(Debug)]
 enum BufferType {
@@ -61,6 +130,9 @@ pub struct RenderPass {
     /// sampler.
     textures: HashMap<egui::TextureId, (Option<wgpu::Texture>, wgpu::BindGroup)>,
     next_user_texture_id: u64,
+    /// Storage for use by [`egui::PaintCallback`]'s that need to store resources such as render
+    /// pipelines that must have the lifetime of the renderpass.
+    pub paint_callback_resources: type_map::TypeMap,
 }
 
 impl RenderPass {
@@ -214,6 +286,7 @@ impl RenderPass {
             texture_bind_group_layout,
             textures: HashMap::new(),
             next_user_texture_id: 0,
+            paint_callback_resources: TypeMap::default(),
         }
     }
 
@@ -258,13 +331,13 @@ impl RenderPass {
         paint_jobs: &[egui::epaint::ClippedPrimitive],
         screen_descriptor: &ScreenDescriptor,
     ) {
-        rpass.set_pipeline(&self.render_pipeline);
-
-        rpass.set_bind_group(0, &self.uniform_bind_group, &[]);
-
         let pixels_per_point = screen_descriptor.pixels_per_point;
         let size_in_pixels = screen_descriptor.size_in_pixels;
 
+        // Whether or not we need to reset the renderpass state because a paint callback has just
+        // run.
+        let mut needs_reset = true;
+
         for (
             (
                 egui::ClippedPrimitive {
@@ -279,41 +352,34 @@ impl RenderPass {
             .zip(&self.vertex_buffers)
             .zip(&self.index_buffers)
         {
-            // Transform clip rect to physical pixels.
-            let clip_min_x = pixels_per_point * clip_rect.min.x;
-            let clip_min_y = pixels_per_point * clip_rect.min.y;
-            let clip_max_x = pixels_per_point * clip_rect.max.x;
-            let clip_max_y = pixels_per_point * clip_rect.max.y;
-
-            // Make sure clip rect can fit within an `u32`.
-            let clip_min_x = clip_min_x.clamp(0.0, size_in_pixels[0] as f32);
-            let clip_min_y = clip_min_y.clamp(0.0, size_in_pixels[1] as f32);
-            let clip_max_x = clip_max_x.clamp(clip_min_x, size_in_pixels[0] as f32);
-            let clip_max_y = clip_max_y.clamp(clip_min_y, size_in_pixels[1] as f32);
-
-            let clip_min_x = clip_min_x.round() as u32;
-            let clip_min_y = clip_min_y.round() as u32;
-            let clip_max_x = clip_max_x.round() as u32;
-            let clip_max_y = clip_max_y.round() as u32;
-
-            let width = (clip_max_x - clip_min_x).max(1);
-            let height = (clip_max_y - clip_min_y).max(1);
-
-            {
-                // Clip scissor rectangle to target size.
-                let x = clip_min_x.min(size_in_pixels[0]);
-                let y = clip_min_y.min(size_in_pixels[1]);
-                let width = width.min(size_in_pixels[0] - x);
-                let height = height.min(size_in_pixels[1] - y);
-
-                // Skip rendering with zero-sized clip areas.
-                if width == 0 || height == 0 {
-                    continue;
-                }
-
-                rpass.set_scissor_rect(x, y, width, height);
+            if needs_reset {
+                rpass.set_viewport(
+                    0.0,
+                    0.0,
+                    size_in_pixels[0] as f32,
+                    size_in_pixels[1] as f32,
+                    0.0,
+                    1.0,
+                );
+                rpass.set_pipeline(&self.render_pipeline);
+                rpass.set_bind_group(0, &self.uniform_bind_group, &[]);
+                needs_reset = false;
             }
 
+            let PixelRect {
+                x,
+                y,
+                width,
+                height,
+            } = calculate_pixel_rect(clip_rect, pixels_per_point, size_in_pixels);
+
+            // Skip rendering with zero-sized clip areas.
+            if width == 0 || height == 0 {
+                continue;
+            }
+
+            rpass.set_scissor_rect(x, y, width, height);
+
             match primitive {
                 Primitive::Mesh(mesh) => {
                     if let Some((_texture, bind_group)) = self.textures.get(&mesh.texture_id) {
@@ -328,8 +394,57 @@ impl RenderPass {
                         tracing::warn!("Missing texture: {:?}", mesh.texture_id);
                     }
                 }
-                Primitive::Callback(_) => {
-                    // already warned about earlier
+                Primitive::Callback(callback) => {
+                    let cbfn = if let Some(c) = callback.callback.downcast_ref::<CallbackFn>() {
+                        c
+                    } else {
+                        // We already warned in the `prepare` callback
+                        continue;
+                    };
+
+                    if callback.rect.is_positive() {
+                        needs_reset = true;
+
+                        // Set the viewport rect
+                        let PixelRect {
+                            x,
+                            y,
+                            width,
+                            height,
+                        } = calculate_pixel_rect(&callback.rect, pixels_per_point, size_in_pixels);
+                        rpass.set_viewport(
+                            x as f32,
+                            y as f32,
+                            width as f32,
+                            height as f32,
+                            0.0,
+                            1.0,
+                        );
+
+                        // Set the scissor rect
+                        let PixelRect {
+                            x,
+                            y,
+                            width,
+                            height,
+                        } = calculate_pixel_rect(clip_rect, pixels_per_point, size_in_pixels);
+                        // Skip rendering with zero-sized clip areas.
+                        if width == 0 || height == 0 {
+                            continue;
+                        }
+                        rpass.set_scissor_rect(x, y, width, height);
+
+                        (cbfn.paint)(
+                            PaintCallbackInfo {
+                                viewport: callback.rect,
+                                clip_rect: *clip_rect,
+                                pixels_per_point,
+                                screen_size_px: size_in_pixels,
+                            },
+                            rpass,
+                            &self.paint_callback_resources,
+                        );
+                    }
                 }
             }
         }
@@ -448,7 +563,6 @@ impl RenderPass {
         };
     }
 
-    /// Should be called before `execute()`.
     pub fn free_texture(&mut self, id: &egui::TextureId) {
         self.textures.remove(id);
     }
@@ -587,8 +701,15 @@ impl RenderPass {
                         });
                     }
                 }
-                Primitive::Callback(_) => {
-                    tracing::warn!("Painting callbacks not supported by egui-wgpu (yet)");
+                Primitive::Callback(callback) => {
+                    let cbfn = if let Some(c) = callback.callback.downcast_ref::<CallbackFn>() {
+                        c
+                    } else {
+                        tracing::warn!("Unknown paint callback: expected `egui_gpu::CallbackFn`");
+                        continue;
+                    };
+
+                    (cbfn.prepare)(device, queue, &mut self.paint_callback_resources);
                 }
             }
         }
@@ -633,3 +754,51 @@ impl RenderPass {
         }
     }
 }
+
+/// A Rect in physical pixel space, used for setting viewport and cliipping rectangles.
+struct PixelRect {
+    x: u32,
+    y: u32,
+    width: u32,
+    height: u32,
+}
+
+/// Convert the Egui clip rect to a physical pixel rect we can use for the GPU viewport/scissor
+fn calculate_pixel_rect(
+    clip_rect: &egui::Rect,
+    pixels_per_point: f32,
+    target_size: [u32; 2],
+) -> PixelRect {
+    // Transform clip rect to physical pixels.
+    let clip_min_x = pixels_per_point * clip_rect.min.x;
+    let clip_min_y = pixels_per_point * clip_rect.min.y;
+    let clip_max_x = pixels_per_point * clip_rect.max.x;
+    let clip_max_y = pixels_per_point * clip_rect.max.y;
+
+    // Make sure clip rect can fit within an `u32`.
+    let clip_min_x = clip_min_x.clamp(0.0, target_size[0] as f32);
+    let clip_min_y = clip_min_y.clamp(0.0, target_size[1] as f32);
+    let clip_max_x = clip_max_x.clamp(clip_min_x, target_size[0] as f32);
+    let clip_max_y = clip_max_y.clamp(clip_min_y, target_size[1] as f32);
+
+    let clip_min_x = clip_min_x.round() as u32;
+    let clip_min_y = clip_min_y.round() as u32;
+    let clip_max_x = clip_max_x.round() as u32;
+    let clip_max_y = clip_max_y.round() as u32;
+
+    let width = (clip_max_x - clip_min_x).max(1);
+    let height = (clip_max_y - clip_min_y).max(1);
+
+    // Clip scissor rectangle to target size.
+    let x = clip_min_x.min(target_size[0]);
+    let y = clip_min_y.min(target_size[1]);
+    let width = width.min(target_size[0] - x);
+    let height = height.min(target_size[1] - y);
+
+    PixelRect {
+        x,
+        y,
+        width,
+        height,
+    }
+}

egui-wgpu/src/winit.rs

@@ -1,13 +1,19 @@
+use std::sync::Arc;
+
+use egui::mutex::RwLock;
 use tracing::error;
 use wgpu::{Adapter, Instance, Surface, TextureFormat};
 
 use crate::renderer;
 
-struct RenderState {
-    device: wgpu::Device,
-    queue: wgpu::Queue,
-    target_format: TextureFormat,
-    egui_rpass: renderer::RenderPass,
+/// Access to the render state for egui, which can be useful in combination with
+/// [`egui::PaintCallback`]s for custom rendering using WGPU.
+#[derive(Clone)]
+pub struct RenderState {
+    pub device: Arc<wgpu::Device>,
+    pub queue: Arc<wgpu::Queue>,
+    pub target_format: TextureFormat,
+    pub egui_rpass: Arc<RwLock<renderer::RenderPass>>,
 }
 
 struct SurfaceState {
@@ -66,6 +72,13 @@ impl<'a> Painter<'a> {
         }
     }
 
+    /// Get the [`RenderState`].
+    ///
+    /// Will return [`None`] if the render state has not been initialized yet.
+    pub fn get_render_state(&self) -> Option<RenderState> {
+        self.render_state.as_ref().cloned()
+    }
+
     async fn init_render_state(
         &self,
         adapter: &Adapter,
@@ -74,13 +87,13 @@ impl<'a> Painter<'a> {
         let (device, queue) =
             pollster::block_on(adapter.request_device(&self.device_descriptor, None)).unwrap();
 
-        let egui_rpass = renderer::RenderPass::new(&device, target_format, self.msaa_samples);
+        let rpass = renderer::RenderPass::new(&device, target_format, self.msaa_samples);
 
         RenderState {
-            device,
-            queue,
+            device: Arc::new(device),
+            queue: Arc::new(queue),
             target_format,
-            egui_rpass,
+            egui_rpass: Arc::new(RwLock::new(rpass)),
         }
     }
 
@@ -246,27 +259,22 @@ impl<'a> Painter<'a> {
             pixels_per_point,
         };
 
-        for (id, image_delta) in &textures_delta.set {
-            render_state.egui_rpass.update_texture(
+        {
+            let mut rpass = render_state.egui_rpass.write();
+            for (id, image_delta) in &textures_delta.set {
+                rpass.update_texture(&render_state.device, &render_state.queue, *id, image_delta);
+            }
+
+            rpass.update_buffers(
                 &render_state.device,
                 &render_state.queue,
-                *id,
-                image_delta,
+                clipped_primitives,
+                &screen_descriptor,
             );
         }
-        for id in &textures_delta.free {
-            render_state.egui_rpass.free_texture(id);
-        }
-
-        render_state.egui_rpass.update_buffers(
-            &render_state.device,
-            &render_state.queue,
-            clipped_primitives,
-            &screen_descriptor,
-        );
 
         // Record all render passes.
-        render_state.egui_rpass.execute(
+        render_state.egui_rpass.read().execute(
             &mut encoder,
             &output_view,
             clipped_primitives,
@@ -279,6 +287,13 @@ impl<'a> Painter<'a> {
             }),
         );
 
+        {
+            let mut rpass = render_state.egui_rpass.write();
+            for id in &textures_delta.free {
+                rpass.free_texture(id);
+            }
+        }
+
         // Submit the commands.
         render_state.queue.submit(std::iter::once(encoder.finish()));