Add minimal infrastructure for overriding an Image

vello/src/lib.rs

@@ -89,7 +89,7 @@ mod shaders;
 mod wgpu_engine;
 
 #[cfg(feature = "wgpu")]
-use std::num::NonZeroUsize;
+use std::{num::NonZeroUsize, sync::Arc};
 
 /// Styling and composition primitives.
 pub use peniko;
@@ -329,6 +329,20 @@ impl Renderer {
         })
     }
 
+    /// Overwrite the `Image` with the `Texture` texture.
+    ///
+    /// If texture is `None`, removes the override.
+    pub fn override_image(
+        &mut self,
+        image: &peniko::Image,
+        texture: Option<Arc<wgpu::ImageCopyTextureBase<wgpu::Texture>>>,
+    ) -> Option<Arc<wgpu::ImageCopyTextureBase<wgpu::Texture>>> {
+        match texture {
+            Some(texture) => self.engine.image_overrides.insert(image.data.id(), texture),
+            None => self.engine.image_overrides.remove(&image.data.id()),
+        }
+    }
+
     /// Renders a scene to the target texture.
     ///
     /// The texture is assumed to be of the specified dimensions and have been created with

vello/src/recording.rs

@@ -4,6 +4,8 @@
 use std::num::NonZeroU64;
 use std::sync::atomic::{AtomicU64, Ordering};
 
+use peniko::Image;
+
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Default)]
 pub struct ShaderId(pub usize);
 
@@ -61,7 +63,7 @@ pub enum Command {
     UploadUniform(BufferProxy, Vec<u8>),
     /// Commands the data to be uploaded to the given image.
     UploadImage(ImageProxy, Vec<u8>),
-    WriteImage(ImageProxy, [u32; 4], Vec<u8>),
+    WriteImage(ImageProxy, [u32; 2], Image),
     // Discussion question: third argument is vec of resources?
     // Maybe use tricks to make more ergonomic?
     // Alternative: provide bufs & images as separate sequences
@@ -133,17 +135,8 @@ impl Recording {
         image_proxy
     }
 
-    pub fn write_image(
-        &mut self,
-        image: ImageProxy,
-        x: u32,
-        y: u32,
-        width: u32,
-        height: u32,
-        data: impl Into<Vec<u8>>,
-    ) {
-        let data = data.into();
-        self.push(Command::WriteImage(image, [x, y, width, height], data));
+    pub fn write_image(&mut self, proxy: ImageProxy, x: u32, y: u32, image: Image) {
+        self.push(Command::WriteImage(proxy, [x, y], image));
     }
 
     pub fn dispatch<R>(&mut self, shader: ShaderId, wg_size: (u32, u32, u32), resources: R)

vello/src/render.rs

@@ -113,14 +113,7 @@ impl Render {
             ImageProxy::new(images.width, images.height, ImageFormat::Rgba8)
         };
         for image in images.images {
-            recording.write_image(
-                image_atlas,
-                image.1,
-                image.2,
-                image.0.width,
-                image.0.height,
-                image.0.data.data(),
-            );
+            recording.write_image(image_atlas, image.1, image.2, image.0.clone());
         }
         let cpu_config =
             RenderConfig::new(&layout, params.width, params.height, &params.base_color);

vello/src/wgpu_engine.rs

@@ -5,6 +5,7 @@ use std::borrow::Cow;
 use std::cell::RefCell;
 use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet};
+use std::sync::Arc;
 
 use vello_shaders::cpu::CpuBinding;
 
@@ -36,6 +37,10 @@ pub struct WgpuEngine {
     #[cfg(not(target_arch = "wasm32"))]
     shaders_to_initialise: Option<Vec<UninitialisedShader>>,
     pub(crate) use_cpu: bool,
+    /// Overrides from a specific `Image`'s [`id`](peniko::Image::id) to a wgpu `Texture`.
+    ///
+    /// The `Texture` should have the same size as the `Image`.
+    pub(crate) image_overrides: HashMap<u64, Arc<wgpu::ImageCopyTextureBase<Texture>>>,
 }
 
 struct WgpuShader {
@@ -433,31 +438,53 @@ impl WgpuEngine {
                     self.bind_map
                         .insert_image(image_proxy.id, texture, texture_view);
                 }
-                Command::WriteImage(proxy, [x, y, width, height], data) => {
+                Command::WriteImage(proxy, [x, y], image) => {
                     let (texture, _) = self.bind_map.get_or_create_image(*proxy, device);
                     let format = proxy.format.to_wgpu();
                     let block_size = format
                         .block_copy_size(None)
                         .expect("ImageFormat must have a valid block size");
-                    queue.write_texture(
-                        wgpu::ImageCopyTexture {
-                            texture,
-                            mip_level: 0,
-                            origin: wgpu::Origin3d { x: *x, y: *y, z: 0 },
-                            aspect: TextureAspect::All,
-                        },
-                        &data[..],
-                        wgpu::ImageDataLayout {
-                            offset: 0,
-                            bytes_per_row: Some(*width * block_size),
-                            rows_per_image: None,
-                        },
-                        wgpu::Extent3d {
-                            width: *width,
-                            height: *height,
-                            depth_or_array_layers: 1,
-                        },
-                    );
+                    if let Some(overrider) = self.image_overrides.get(&image.data.id()) {
+                        encoder.copy_texture_to_texture(
+                            wgpu::ImageCopyTexture {
+                                texture: &overrider.texture,
+                                mip_level: overrider.mip_level,
+                                origin: overrider.origin,
+                                aspect: overrider.aspect,
+                            },
+                            wgpu::ImageCopyTexture {
+                                texture,
+                                mip_level: 0,
+                                origin: wgpu::Origin3d { x: *x, y: *y, z: 0 },
+                                aspect: TextureAspect::All,
+                            },
+                            wgpu::Extent3d {
+                                width: image.width,
+                                height: image.height,
+                                depth_or_array_layers: 1,
+                            },
+                        );
+                    } else {
+                        queue.write_texture(
+                            wgpu::ImageCopyTexture {
+                                texture,
+                                mip_level: 0,
+                                origin: wgpu::Origin3d { x: *x, y: *y, z: 0 },
+                                aspect: TextureAspect::All,
+                            },
+                            image.data.data(),
+                            wgpu::ImageDataLayout {
+                                offset: 0,
+                                bytes_per_row: Some(image.width * block_size),
+                                rows_per_image: None,
+                            },
+                            wgpu::Extent3d {
+                                width: image.width,
+                                height: image.height,
+                                depth_or_array_layers: 1,
+                            },
+                        );
+                    }
                 }
                 Command::Dispatch(shader_id, wg_size, bindings) => {
                     // println!("dispatching {:?} with {} bindings", wg_size, bindings.len());