Add `Image` rotation support

Co-authored-by: DKolter <68352124+DKolter@users.noreply.github.com>

4 files changed, 116 insertions, 21 deletions

diff --git a/wgpu/src/image/mod.rs b/wgpu/src/image/mod.rs
index 8b831a3c..69f8a8ca 100644
--- a/wgpu/src/image/mod.rs
+++ b/wgpu/src/image/mod.rs
@@ -135,14 +135,20 @@ impl Pipeline {
                         attributes: &wgpu::vertex_attr_array!(
                             // Position
                             0 => Float32x2,
-                            // Scale
+                            // Center
                             1 => Float32x2,
-                            // Atlas position
+                            // Image size
                             2 => Float32x2,
+                            // Rotation
+                            3 => Float32,
+                            // Scale
+                            4 => Float32x2,
+                            // Atlas position
+                            5 => Float32x2,
                             // Atlas scale
-                            3 => Float32x2,
+                            6 => Float32x2,
                             // Layer
-                            4 => Sint32,
+                            7 => Sint32,
                         ),
                     }],
                 },
@@ -208,9 +214,10 @@ impl Pipeline {
         belt: &mut wgpu::util::StagingBelt,
         images: &Batch,
         transformation: Transformation,
-        scale: f32,
+        global_scale: f32,
     ) {
-        let transformation = transformation * Transformation::scale(scale);
+        let transformation =
+            transformation * Transformation::scale(global_scale);
 
         let nearest_instances: &mut Vec<Instance> = &mut Vec::new();
         let linear_instances: &mut Vec<Instance> = &mut Vec::new();
@@ -224,6 +231,8 @@ impl Pipeline {
                     handle,
                     filter_method,
                     bounds,
+                    rotation,
+                    scale,
                 } => {
                     if let Some(atlas_entry) =
                         cache.upload_raster(device, encoder, handle)
@@ -231,6 +240,8 @@ impl Pipeline {
                         add_instances(
                             [bounds.x, bounds.y],
                             [bounds.width, bounds.height],
+                            *rotation,
+                            [scale.width, scale.height],
                             atlas_entry,
                             match filter_method {
                                 crate::core::image::FilterMethod::Nearest => {
@@ -251,15 +262,24 @@ impl Pipeline {
                     handle,
                     color,
                     bounds,
+                    rotation,
+                    scale,
                 } => {
                     let size = [bounds.width, bounds.height];
 
                     if let Some(atlas_entry) = cache.upload_vector(
-                        device, encoder, handle, *color, size, scale,
+                        device,
+                        encoder,
+                        handle,
+                        *color,
+                        size,
+                        global_scale,
                     ) {
                         add_instances(
                             [bounds.x, bounds.y],
                             size,
+                            *rotation,
+                            [scale.width, scale.height],
                             atlas_entry,
                             nearest_instances,
                         );
@@ -487,7 +507,10 @@ impl Data {
 #[derive(Debug, Clone, Copy, Zeroable, Pod)]
 struct Instance {
     _position: [f32; 2],
+    _center: [f32; 2],
     _size: [f32; 2],
+    _rotation: f32,
+    _scale: [f32; 2],
     _position_in_atlas: [f32; 2],
     _size_in_atlas: [f32; 2],
     _layer: u32,
@@ -506,12 +529,27 @@ struct Uniforms {
 fn add_instances(
     image_position: [f32; 2],
     image_size: [f32; 2],
+    rotation: f32,
+    scale: [f32; 2],
     entry: &atlas::Entry,
     instances: &mut Vec<Instance>,
 ) {
+    let center = [
+        image_position[0] + image_size[0] / 2.0,
+        image_position[1] + image_size[1] / 2.0,
+    ];
+
     match entry {
         atlas::Entry::Contiguous(allocation) => {
-            add_instance(image_position, image_size, allocation, instances);
+            add_instance(
+                image_position,
+                center,
+                image_size,
+                rotation,
+                scale,
+                allocation,
+                instances,
+            );
         }
         atlas::Entry::Fragmented { fragments, size } => {
             let scaling_x = image_size[0] / size.width as f32;
@@ -537,7 +575,10 @@ fn add_instances(
                     fragment_height as f32 * scaling_y,
                 ];
 
-                add_instance(position, size, allocation, instances);
+                add_instance(
+                    position, center, size, rotation, scale, allocation,
+                    instances,
+                );
             }
         }
     }
@@ -546,7 +587,10 @@ fn add_instances(
 #[inline]
 fn add_instance(
     position: [f32; 2],
+    center: [f32; 2],
     size: [f32; 2],
+    rotation: f32,
+    scale: [f32; 2],
     allocation: &atlas::Allocation,
     instances: &mut Vec<Instance>,
 ) {
@@ -556,7 +600,10 @@ fn add_instance(
 
     let instance = Instance {
         _position: position,
+        _center: center,
         _size: size,
+        _rotation: rotation,
+        _scale: scale,
         _position_in_atlas: [
             (x as f32 + 0.5) / atlas::SIZE as f32,
             (y as f32 + 0.5) / atlas::SIZE as f32,
diff --git a/wgpu/src/layer.rs b/wgpu/src/layer.rs
index 9526c5a8..648ec476 100644
--- a/wgpu/src/layer.rs
+++ b/wgpu/src/layer.rs
@@ -1,5 +1,6 @@
-use crate::core::renderer;
-use crate::core::{Background, Color, Point, Rectangle, Transformation};
+use crate::core::{
+    renderer, Background, Color, Point, Rectangle, Size, Transformation,
+};
 use crate::graphics;
 use crate::graphics::color;
 use crate::graphics::layer;
@@ -117,11 +118,15 @@ impl Layer {
         filter_method: crate::core::image::FilterMethod,
         bounds: Rectangle,
         transformation: Transformation,
+        rotation: f32,
+        scale: Size,
     ) {
         let image = Image::Raster {
             handle,
             filter_method,
             bounds: bounds * transformation,
+            rotation,
+            scale,
         };
 
         self.images.push(image);
@@ -133,11 +138,15 @@ impl Layer {
         color: Option<Color>,
         bounds: Rectangle,
         transformation: Transformation,
+        rotation: f32,
+        scale: Size,
     ) {
         let svg = Image::Vector {
             handle,
             color,
             bounds: bounds * transformation,
+            rotation,
+            scale,
         };
 
         self.images.push(svg);
diff --git a/wgpu/src/lib.rs b/wgpu/src/lib.rs
index 178522de..a42d71c3 100644
--- a/wgpu/src/lib.rs
+++ b/wgpu/src/lib.rs
@@ -517,9 +517,18 @@ impl core::image::Renderer for Renderer {
         handle: Self::Handle,
         filter_method: core::image::FilterMethod,
         bounds: Rectangle,
+        rotation: f32,
+        scale: Size,
     ) {
         let (layer, transformation) = self.layers.current_mut();
-        layer.draw_image(handle, filter_method, bounds, transformation);
+        layer.draw_image(
+            handle,
+            filter_method,
+            bounds,
+            transformation,
+            rotation,
+            scale,
+        );
     }
 }
 
@@ -534,9 +543,18 @@ impl core::svg::Renderer for Renderer {
         handle: core::svg::Handle,
         color_filter: Option<Color>,
         bounds: Rectangle,
+        rotation: f32,
+        scale: Size,
     ) {
         let (layer, transformation) = self.layers.current_mut();
-        layer.draw_svg(handle, color_filter, bounds, transformation);
+        layer.draw_svg(
+            handle,
+            color_filter,
+            bounds,
+            transformation,
+            rotation,
+            scale,
+        );
     }
 }
 
diff --git a/wgpu/src/shader/image.wgsl b/wgpu/src/shader/image.wgsl
index 7b2e5238..de962098 100644
--- a/wgpu/src/shader/image.wgsl
+++ b/wgpu/src/shader/image.wgsl
@@ -9,10 +9,13 @@ struct Globals {
 struct VertexInput {
     @builtin(vertex_index) vertex_index: u32,
     @location(0) pos: vec2<f32>,
-    @location(1) scale: vec2<f32>,
-    @location(2) atlas_pos: vec2<f32>,
-    @location(3) atlas_scale: vec2<f32>,
-    @location(4) layer: i32,
+    @location(1) center: vec2<f32>,
+    @location(2) image_size: vec2<f32>,
+    @location(3) rotation: f32,
+    @location(4) scale: vec2<f32>,
+    @location(5) atlas_pos: vec2<f32>,
+    @location(6) atlas_scale: vec2<f32>,
+    @location(7) layer: i32,
 }
 
 struct VertexOutput {
@@ -25,24 +28,42 @@ struct VertexOutput {
 fn vs_main(input: VertexInput) -> VertexOutput {
     var out: VertexOutput;
 
-    let v_pos = vertex_position(input.vertex_index);
+    // Generate a vertex position in the range [0, 1] from the vertex index.
+    var v_pos = vertex_position(input.vertex_index);
 
+    // Map the vertex position to the atlas texture.
     out.uv = vec2<f32>(v_pos * input.atlas_scale + input.atlas_pos);
     out.layer = f32(input.layer);
 
-    var transform: mat4x4<f32> = mat4x4<f32>(
+    // Calculate the vertex position and move the center to the origin
+    v_pos = input.pos + v_pos * input.image_size - input.center;
+
+    // Apply the rotation around the center of the image
+    let cos_rot = cos(input.rotation);
+    let sin_rot = sin(input.rotation);
+    let rotate = mat4x4<f32>(
+        vec4<f32>(cos_rot, sin_rot, 0.0, 0.0),
+        vec4<f32>(-sin_rot, cos_rot, 0.0, 0.0),
+        vec4<f32>(0.0, 0.0, 1.0, 0.0),
+        vec4<f32>(0.0, 0.0, 0.0, 1.0)
+    );
+
+    // Scale the image and then translate to the final position by moving the center to the position
+    let scale_translate = mat4x4<f32>(
         vec4<f32>(input.scale.x, 0.0, 0.0, 0.0),
         vec4<f32>(0.0, input.scale.y, 0.0, 0.0),
         vec4<f32>(0.0, 0.0, 1.0, 0.0),
-        vec4<f32>(input.pos, 0.0, 1.0)
+        vec4<f32>(input.center, 0.0, 1.0)
     );
 
-    out.position = globals.transform * transform * vec4<f32>(v_pos, 0.0, 1.0);
+    // Calculate the final position of the vertex
+    out.position = globals.transform * scale_translate * rotate * vec4<f32>(v_pos, 0.0, 1.0);
 
     return out;
 }
 
 @fragment
 fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> {
+    // Sample the texture at the given UV coordinate and layer.
     return textureSample(u_texture, u_sampler, input.uv, i32(input.layer));
 }