Improved wgpu callbacks (#3253)

* Improved wgpu callbacks

* update documentation on egui_wgpu callbacks

* make shared callback resource map pub

* make it nicer to create epaint::PaintCallback from egui_wgpu callback

* constrain ClippedPrimitive lifetime to outlive wgpu::RenderPass

* Revert callback resources to TypeMap, put finish_prepare on callback trait

* doc string fixes

Cargo.toml

@@ -36,4 +36,4 @@ opt-level = 2
 
 [workspace.dependencies]
 thiserror = "1.0.37"
-wgpu = { version = "0.17.0", features = ["fragile-send-sync-non-atomic-wasm"] }
+wgpu = "0.17.0"

crates/egui-wgpu/src/lib.rs

@@ -10,8 +10,8 @@ pub use wgpu;
 
 /// Low-level painting of [`egui`](https://github.com/emilk/egui) on [`wgpu`].
 pub mod renderer;
-pub use renderer::CallbackFn;
 pub use renderer::Renderer;
+pub use renderer::{Callback, CallbackResources, CallbackTrait};
 
 /// Module for painting [`egui`](https://github.com/emilk/egui) with [`wgpu`] on [`winit`].
 #[cfg(feature = "winit")]

crates/egui-wgpu/src/renderer.rs

@@ -1,72 +1,60 @@
 #![allow(unsafe_code)]
 
-use std::num::NonZeroU64;
+use std::{borrow::Cow, num::NonZeroU64, ops::Range};
-use std::ops::Range;
+
-use std::{borrow::Cow, collections::HashMap};
+use epaint::{ahash::HashMap, emath::NumExt, PaintCallbackInfo, Primitive, Vertex};
 
-use type_map::concurrent::TypeMap;
 use wgpu;
 use wgpu::util::DeviceExt as _;
 
-use epaint::{emath::NumExt, PaintCallbackInfo, Primitive, Vertex};
+// Only implements Send + Sync on wasm32 in order to allow storing wgpu resources on the type map.
+#[cfg(not(target_arch = "wasm32"))]
+pub type CallbackResources = type_map::concurrent::TypeMap;
+#[cfg(target_arch = "wasm32")]
+pub type CallbackResources = type_map::TypeMap;
 
-/// A callback function that can be used to compose an [`epaint::PaintCallback`] for custom WGPU
+pub struct Callback(Box<dyn CallbackTrait>);
-/// rendering.
+
+impl Callback {
+    /// Creates a new [`epaint::PaintCallback`] from a callback trait instance.
+    pub fn new_paint_callback(
+        rect: epaint::emath::Rect,
+        callback: impl CallbackTrait + 'static,
+    ) -> epaint::PaintCallback {
+        epaint::PaintCallback {
+            rect,
+            callback: std::sync::Arc::new(Self(Box::new(callback))),
+        }
+    }
+}
+
+/// A callback trait that can be used to compose an [`epaint::PaintCallback`] via [`Callback`]
+/// for custom WGPU rendering.
 ///
-/// The callback is composed of two functions: `prepare` and `paint`:
+/// Callbacks in [`Renderer`] are done in three steps:
-/// - `prepare` is called every frame before `paint`, and can use the passed-in
+/// * [`CallbackTrait::prepare`]: called for all registered callbacks before the main egui render pass.
-///   [`wgpu::Device`] and [`wgpu::Buffer`] to allocate or modify GPU resources such as buffers.
+/// * [`CallbackTrait::finish_prepare`]: called for all registered callbacks after all callbacks finished calling prepare.
-/// - `paint` is called after `prepare` and is given access to the [`wgpu::RenderPass`] so
+/// * [`CallbackTrait::paint`]: called for all registered callbacks during the main egui render pass.
-///   that it can issue draw commands into the same [`wgpu::RenderPass`] that is used for
-///   all other egui elements.
 ///
-/// The final argument of both the `prepare` and `paint` callbacks is a the
+/// Each callback has access to an instance of [`CallbackResources`] that is stored in the [`Renderer`].
-/// [`paint_callback_resources`][crate::renderer::Renderer::paint_callback_resources].
+/// This can be used to store wgpu resources that need to be accessed during the [`CallbackTrait::paint`] step.
-/// `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
+/// The callbacks implementing [`CallbackTrait`] itself must always be Send + Sync, but resources stored in
+/// [`Renderer::callback_resources`] are not required to implement Send + Sync when building for wasm.
+/// (this is because wgpu stores references to the JS heap in most of its resources which can not be shared with other threads).
 ///
-/// See the [`custom3d_wgpu`](https://github.com/emilk/egui/blob/master/crates/egui_demo_app/src/apps/custom3d_wgpu.rs) demo source for a detailed usage example.
+///
-pub struct CallbackFn {
+/// # Command submission
-    prepare: Box<PrepareCallback>,
+///
-    paint: Box<PaintCallback>,
+/// ## Command Encoder
-}
-
-type PrepareCallback = dyn Fn(
-        &wgpu::Device,
-        &wgpu::Queue,
-        &mut wgpu::CommandEncoder,
-        &mut TypeMap,
-    ) -> Vec<wgpu::CommandBuffer>
-    + 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(|_, _, _, _| Vec::new()),
-            paint: Box::new(|_, _, _| ()),
-        }
-    }
-}
-
-impl CallbackFn {
-    pub fn new() -> Self {
-        Self::default()
-    }
-
-    /// Set the prepare callback.
 ///
 /// The passed-in `CommandEncoder` is egui's and can be used directly to register
 /// wgpu commands for simple use cases.
 /// This allows reusing the same [`wgpu::CommandEncoder`] for all callbacks and egui
 /// rendering itself.
 ///
+/// ## Command Buffers
+///
 /// For more complicated use cases, one can also return a list of arbitrary
 /// `CommandBuffer`s and have complete control over how they get created and fed.
 /// In particular, this gives an opportunity to parallelize command registration and
@@ -75,33 +63,47 @@ impl CallbackFn {
 /// When using eframe, the main egui command buffer, as well as all user-defined
 /// command buffers returned by this function, are guaranteed to all be submitted
 /// at once in a single call.
-    pub fn prepare<F>(mut self, prepare: F) -> Self
+///
-    where
+/// Command Buffers returned by [`CallbackTrait::finish_prepare`] will always be issued *after*
-        F: Fn(
+/// those returned by [`CallbackTrait::prepare`].
-                &wgpu::Device,
+/// Order within command buffers returned by [`CallbackTrait::prepare`] is dependent
-                &wgpu::Queue,
+/// on the order the respective [`epaint::Shape::Callback`]s were submitted in.
-                &mut wgpu::CommandEncoder,
+///
-                &mut TypeMap,
+/// # Example
-            ) -> Vec<wgpu::CommandBuffer>
+///
-            + Sync
+/// See the [`custom3d_wgpu`](https://github.com/emilk/egui/blob/master/crates/egui_demo_app/src/apps/custom3d_wgpu.rs) demo source for a detailed usage example.
-            + Send
+pub trait CallbackTrait: Send + Sync {
-            + 'static,
+    fn prepare(
-    {
+        &self,
-        self.prepare = Box::new(prepare) as _;
+        _device: &wgpu::Device,
-        self
+        _queue: &wgpu::Queue,
+        _egui_encoder: &mut wgpu::CommandEncoder,
+        _callback_resources: &mut CallbackResources,
+    ) -> Vec<wgpu::CommandBuffer> {
+        Vec::new()
     }
 
-    /// Set the paint callback
+    /// Called after all [`CallbackTrait::prepare`] calls are done.
-    pub fn paint<F>(mut self, paint: F) -> Self
+    fn finish_prepare(
-    where
+        &self,
-        F: for<'a, 'b> Fn(PaintCallbackInfo, &'a mut wgpu::RenderPass<'b>, &'b TypeMap)
+        _device: &wgpu::Device,
-            + Sync
+        _queue: &wgpu::Queue,
-            + Send
+        _egui_encoder: &mut wgpu::CommandEncoder,
-            + 'static,
+        _callback_resources: &mut CallbackResources,
-    {
+    ) -> Vec<wgpu::CommandBuffer> {
-        self.paint = Box::new(paint) as _;
+        Vec::new()
-        self
     }
+
+    /// Called after all [`CallbackTrait::finish_prepare`] calls are done.
+    ///
+    /// It is given access to the [`wgpu::RenderPass`] so that it can issue draw commands
+    /// into the same [`wgpu::RenderPass`] that is used for all other egui elements.
+    fn paint<'a>(
+        &'a self,
+        info: PaintCallbackInfo,
+        render_pass: &mut wgpu::RenderPass<'a>,
+        callback_resources: &'a CallbackResources,
+    );
 }
 
 /// Information about the screen used for rendering.
@@ -164,9 +166,10 @@ pub struct Renderer {
     next_user_texture_id: u64,
     samplers: HashMap<epaint::textures::TextureOptions, wgpu::Sampler>,
 
-    /// Storage for use by [`epaint::PaintCallback`]'s that need to store resources such as render
+    /// Storage for resources shared with all invocations of [`CallbackTrait`]'s methods.
-    /// pipelines that must have the lifetime of the renderpass.
+    ///
-    pub paint_callback_resources: TypeMap,
+    /// See also [`CallbackTrait`].
+    pub callback_resources: CallbackResources,
 }
 
 impl Renderer {
@@ -346,10 +349,10 @@ impl Renderer {
             },
             uniform_bind_group,
             texture_bind_group_layout,
-            textures: HashMap::new(),
+            textures: HashMap::default(),
             next_user_texture_id: 0,
-            samplers: HashMap::new(),
+            samplers: HashMap::default(),
-            paint_callback_resources: TypeMap::default(),
+            callback_resources: CallbackResources::default(),
         }
     }
 
@@ -357,7 +360,7 @@ impl Renderer {
     pub fn render<'rp>(
         &'rp self,
         render_pass: &mut wgpu::RenderPass<'rp>,
-        paint_jobs: &[epaint::ClippedPrimitive],
+        paint_jobs: &'rp [epaint::ClippedPrimitive],
         screen_descriptor: &ScreenDescriptor,
     ) {
         crate::profile_function!();
@@ -432,7 +435,7 @@ impl Renderer {
                     }
                 }
                 Primitive::Callback(callback) => {
-                    let cbfn = if let Some(c) = callback.callback.downcast_ref::<CallbackFn>() {
+                    let cbfn = if let Some(c) = callback.callback.downcast_ref::<Callback>() {
                         c
                     } else {
                         // We already warned in the `prepare` callback
@@ -467,7 +470,7 @@ impl Renderer {
                             );
                         }
 
-                        (cbfn.paint)(
+                        cbfn.0.paint(
                             PaintCallbackInfo {
                                 viewport: callback.rect,
                                 clip_rect: *clip_rect,
@@ -475,7 +478,7 @@ impl Renderer {
                                 screen_size_px: size_in_pixels,
                             },
                             render_pass,
-                            &self.paint_callback_resources,
+                            &self.callback_resources,
                         );
                     }
                 }
@@ -751,7 +754,7 @@ impl Renderer {
     /// Uploads the uniform, vertex and index data used by the renderer.
     /// Should be called before `render()`.
     ///
-    /// Returns all user-defined command buffers gathered from prepare callbacks.
+    /// Returns all user-defined command buffers gathered from [`CallbackTrait::prepare`] & [`CallbackTrait::finish_prepare`] callbacks.
     pub fn update_buffers(
         &mut self,
         device: &wgpu::Device,
@@ -778,7 +781,8 @@ impl Renderer {
             self.previous_uniform_buffer_content = uniform_buffer_content;
         }
 
-        // Determine how many vertices & indices need to be rendered.
+        // Determine how many vertices & indices need to be rendered, and gather prepare callbacks
+        let mut callbacks = Vec::new();
         let (vertex_count, index_count) = {
             crate::profile_scope!("count_vertices_indices");
             paint_jobs.iter().fold((0, 0), |acc, clipped_primitive| {
@@ -786,7 +790,14 @@ impl Renderer {
                     Primitive::Mesh(mesh) => {
                         (acc.0 + mesh.vertices.len(), acc.1 + mesh.indices.len())
                     }
-                    Primitive::Callback(_) => acc,
+                    Primitive::Callback(callback) => {
+                        if let Some(c) = callback.callback.downcast_ref::<Callback>() {
+                            callbacks.push(c.0.as_ref());
+                        } else {
+                            log::warn!("Unknown paint callback: expected `egui_wgpu::Callback`");
+                        };
+                        acc
+                    }
                 }
             })
         };
@@ -861,34 +872,33 @@ impl Renderer {
             }
         }
 
+        let mut user_cmd_bufs = Vec::new();
         {
-            crate::profile_scope!("user command buffers");
+            crate::profile_scope!("prepare callbacks");
-            let mut user_cmd_bufs = Vec::new(); // collect user command buffers
+            for callback in &callbacks {
-            for epaint::ClippedPrimitive { primitive, .. } in paint_jobs.iter() {
+                user_cmd_bufs.extend(callback.prepare(
-                match primitive {
-                    Primitive::Mesh(_) => {}
-                    Primitive::Callback(callback) => {
-                        let cbfn = if let Some(c) = callback.callback.downcast_ref::<CallbackFn>() {
-                            c
-                        } else {
-                            log::warn!("Unknown paint callback: expected `egui_wgpu::CallbackFn`");
-                            continue;
-                        };
-
-                        crate::profile_scope!("callback");
-                        user_cmd_bufs.extend((cbfn.prepare)(
                     device,
                     queue,
                     encoder,
-                            &mut self.paint_callback_resources,
+                    &mut self.callback_resources,
                 ));
             }
         }
+        {
+            crate::profile_scope!("finish prepare callbacks");
+            for callback in &callbacks {
+                user_cmd_bufs.extend(callback.finish_prepare(
+                    device,
+                    queue,
+                    encoder,
+                    &mut self.callback_resources,
+                ));
             }
+        }
+
         user_cmd_bufs
     }
 }
-}
 
 fn create_sampler(
     options: epaint::textures::TextureOptions,
@@ -969,6 +979,9 @@ impl ScissorRect {
     }
 }
 
+// Wgpu objects contain references to the JS heap on the web, therefore they are not Send/Sync.
+// It follows that egui_wgpu::Renderer can not be Send/Sync either when building with wasm.
+#[cfg(not(target_arch = "wasm32"))]
 #[test]
 fn renderer_impl_send_sync() {
     fn assert_send_sync<T: Send + Sync>() {}

crates/egui_demo_app/src/apps/custom3d_wgpu.rs

@@ -1,4 +1,4 @@
-use std::{num::NonZeroU64, sync::Arc};
+use std::num::NonZeroU64;
 
 use eframe::{
     egui_wgpu::wgpu::util::DeviceExt,
@@ -84,7 +84,7 @@ impl Custom3d {
         wgpu_render_state
             .renderer
             .write()
-            .paint_callback_resources
+            .callback_resources
             .insert(TriangleRenderResources {
                 pipeline,
                 bind_group,
@@ -119,46 +119,64 @@ impl eframe::App for Custom3d {
     }
 }
 
+// Callbacks in egui_wgpu have 3 stages:
+// * prepare (per callback impl)
+// * finish_prepare (once)
+// * paint (per callback impl)
+//
+// The prepare callback is called every frame before paint and is given access to the wgpu
+// Device and Queue, which can be used, for instance, to update buffers and uniforms before
+// rendering.
+// If [`egui_wgpu::Renderer`] has [`egui_wgpu::FinishPrepareCallback`] registered,
+// it will be called after all `prepare` callbacks have been called.
+// You can use this to update any shared resources that need to be updated once per frame
+// after all callbacks have been processed.
+//
+// On both prepare methods you can use the main `CommandEncoder` that is passed-in,
+// return an arbitrary number of user-defined `CommandBuffer`s, or both.
+// The main command buffer, as well as all user-defined ones, will be submitted together
+// to the GPU in a single call.
+//
+// The paint callback is called after finish prepare and is given access to egui's main render pass,
+// which can be used to issue draw commands.
+struct CustomTriangleCallback {
+    angle: f32,
+}
+
+impl egui_wgpu::CallbackTrait for CustomTriangleCallback {
+    fn prepare(
+        &self,
+        device: &wgpu::Device,
+        queue: &wgpu::Queue,
+        _egui_encoder: &mut wgpu::CommandEncoder,
+        resources: &mut egui_wgpu::CallbackResources,
+    ) -> Vec<wgpu::CommandBuffer> {
+        let resources: &TriangleRenderResources = resources.get().unwrap();
+        resources.prepare(device, queue, self.angle);
+        Vec::new()
+    }
+
+    fn paint<'a>(
+        &self,
+        _info: egui::PaintCallbackInfo,
+        render_pass: &mut wgpu::RenderPass<'a>,
+        resources: &'a egui_wgpu::CallbackResources,
+    ) {
+        let resources: &TriangleRenderResources = resources.get().unwrap();
+        resources.paint(render_pass);
+    }
+}
+
 impl Custom3d {
     fn custom_painting(&mut self, ui: &mut egui::Ui) {
         let (rect, response) =
             ui.allocate_exact_size(egui::Vec2::splat(300.0), egui::Sense::drag());
 
         self.angle += response.drag_delta().x * 0.01;
-
+        ui.painter().add(egui_wgpu::Callback::new_paint_callback(
-        // Clone locals so we can move them into the paint callback:
-        let angle = self.angle;
-
-        // The callback function for WGPU is in two stages: prepare, and paint.
-        //
-        // The prepare callback is called every frame before paint and is given access to the wgpu
-        // Device and Queue, which can be used, for instance, to update buffers and uniforms before
-        // rendering.
-        //
-        // You can use the main `CommandEncoder` that is passed-in, return an arbitrary number
-        // of user-defined `CommandBuffer`s, or both.
-        // The main command buffer, as well as all user-defined ones, will be submitted together
-        // to the GPU in a single call.
-        //
-        // The paint callback is called after prepare and is given access to the render pass, which
-        // can be used to issue draw commands.
-        let cb = egui_wgpu::CallbackFn::new()
-            .prepare(move |device, queue, _encoder, paint_callback_resources| {
-                let resources: &TriangleRenderResources = paint_callback_resources.get().unwrap();
-                resources.prepare(device, queue, angle);
-                Vec::new()
-            })
-            .paint(move |_info, render_pass, paint_callback_resources| {
-                let resources: &TriangleRenderResources = paint_callback_resources.get().unwrap();
-                resources.paint(render_pass);
-            });
-
-        let callback = egui::PaintCallback {
             rect,
-            callback: Arc::new(cb),
+            CustomTriangleCallback { angle: self.angle },
-        };
+        ));
-
-        ui.painter().add(callback);
     }
 }
 

crates/epaint/src/shape.rs

@@ -852,7 +852,7 @@ pub struct PaintCallback {
     ///
     /// The concrete value of `callback` depends on the rendering backend used. For instance, the
     /// `glow` backend requires that callback be an `egui_glow::CallbackFn` while the `wgpu`
-    /// backend requires a `egui_wgpu::CallbackFn`.
+    /// backend requires a `egui_wgpu::Callback`.
     ///
     /// If the type cannot be downcast to the type expected by the current backend the callback
     /// will not be drawn.
@@ -862,7 +862,9 @@ pub struct PaintCallback {
     ///
     /// The rendering backend is also responsible for restoring any state, such as the bound shader
     /// program, vertex array, etc.
-    pub callback: Arc<dyn Any + Sync + Send>,
+    ///
+    /// Shape has to be clone, therefore this has to be an `Arc` instead of a `Box`.
+    pub callback: Arc<dyn Any + Send + Sync>,
 }
 
 impl std::fmt::Debug for PaintCallback {