From 830ac83a912777237d4107614a4502c280be8542 Mon Sep 17 00:00:00 2001
From: felix de Maneville <felix.maneville@gmail.com>
Date: Tue, 9 Aug 2022 19:05:52 +0200
Subject: [PATCH] Ui Texture Atlas

---
 crates/bevy_ui/src/render/mod.rs | 29 ++++++++++++++++----
 examples/ui/button.rs            | 47 +++++++++++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_ui/src/render/mod.rs b/crates/bevy_ui/src/render/mod.rs
index 7fd783d362bde8..cf379dd0500990 100644
--- a/crates/bevy_ui/src/render/mod.rs
+++ b/crates/bevy_ui/src/render/mod.rs
@@ -23,7 +23,7 @@ use bevy_render::{
     view::{ComputedVisibility, ExtractedView, ViewUniforms},
     Extract, RenderApp, RenderStage,
 };
-use bevy_sprite::{Rect, SpriteAssetEvents, TextureAtlas};
+use bevy_sprite::{Rect, SpriteAssetEvents, TextureAtlas, TextureSheet};
 use bevy_text::{DefaultTextPipeline, Text};
 use bevy_transform::components::GlobalTransform;
 use bevy_utils::FloatOrd;
@@ -176,6 +176,7 @@ pub struct ExtractedUiNodes {
 pub fn extract_uinodes(
     mut extracted_uinodes: ResMut<ExtractedUiNodes>,
     images: Extract<Res<Assets<Image>>>,
+    texture_atlases: Extract<Res<Assets<TextureAtlas>>>,
     uinode_query: Extract<
         Query<(
             &Node,
@@ -183,13 +184,15 @@ pub fn extract_uinodes(
             &UiColor,
             &UiImage,
             &ComputedVisibility,
+            Option<&TextureSheet>,
             Option<&CalculatedClip>,
         )>,
     >,
 ) {
     extracted_uinodes.uinodes.clear();
-    for (uinode, transform, color, image, visibility, clip) in uinode_query.iter() {
-        if !visibility.is_visible() {
+    for (uinode, transform, color, image, visibility, sheet, clip) in uinode_query.iter() {
+        // Skip if the node is invisible or if its size is set to zero (e.g. when a parent is set to `Display::None`)
+        if !visibility.is_visible() || uinode.size == Vec2::ZERO {
             continue;
         }
         let image = image.0.clone_weak();
@@ -201,15 +204,29 @@ pub fn extract_uinodes(
         if color.0.a() == 0.0 {
             continue;
         }
+        let (atlas_size, rect_min) = sheet
+            .and_then(|s| {
+                texture_atlases
+                    .get(&s.texture_atlas)
+                    .map(|a| (a, s.texture_index))
+            })
+            .and_then(|(atlas, index)| {
+                atlas
+                    .textures
+                    .get(index)
+                    .map(|rect| (Some(atlas.size), rect.min))
+            })
+            .unwrap_or((None, Vec2::ZERO));
+
         extracted_uinodes.uinodes.push(ExtractedUiNode {
             transform: transform.compute_matrix(),
             color: color.0,
             rect: bevy_sprite::Rect {
-                min: Vec2::ZERO,
-                max: uinode.size,
+                min: rect_min,
+                max: rect_min + uinode.size,
             },
             image,
-            atlas_size: None,
+            atlas_size,
             clip: clip.map(|clip| clip.clip),
         });
     }
diff --git a/examples/ui/button.rs b/examples/ui/button.rs
index 8007ee86c0b358..c9ddb1feef9144 100644
--- a/examples/ui/button.rs
+++ b/examples/ui/button.rs
@@ -10,6 +10,7 @@ fn main() {
         .insert_resource(WinitSettings::desktop_app())
         .add_startup_system(setup)
         .add_system(button_system)
+        .add_system(sheet_button_system)
         .run();
 }
 
@@ -43,9 +44,35 @@ fn button_system(
     }
 }
 
-fn setup(mut commands: Commands, asset_server: Res<AssetServer>) {
+fn sheet_button_system(
+    mut interaction_query: Query<
+        (&Interaction, &mut TextureSheet),
+        (Changed<Interaction>, With<Button>),
+    >,
+) {
+    for (interaction, mut sheet) in interaction_query.iter_mut() {
+        match *interaction {
+            Interaction::Clicked => {
+                sheet.texture_index = 2;
+            }
+            Interaction::Hovered => {
+                sheet.texture_index = 1;
+            }
+            Interaction::None => {
+                sheet.texture_index = 0;
+            }
+        }
+    }
+}
+
+fn setup(
+    mut commands: Commands,
+    asset_server: Res<AssetServer>,
+    mut texture_atlases: ResMut<Assets<TextureAtlas>>,
+) {
     // ui camera
     commands.spawn_bundle(Camera2dBundle::default());
+    // Classic button
     commands
         .spawn_bundle(ButtonBundle {
             style: Style {
@@ -71,4 +98,22 @@ fn setup(mut commands: Commands, asset_server: Res<AssetServer>) {
                 },
             ));
         });
+    // Texture sheet button
+    let handle = asset_server.load("textures/array_texture.png");
+    let texture_atlas = texture_atlases.add(TextureAtlas::from_grid(Vec2::splat(250.0), 1, 4));
+    commands
+        .spawn_bundle(ButtonBundle {
+            style: Style {
+                size: Size::new(Val::Px(150.0), Val::Px(65.0)),
+                // center button
+                margin: UiRect::all(Val::Auto),
+                ..default()
+            },
+            image: handle.into(),
+            ..default()
+        })
+        .insert(TextureSheet {
+            texture_atlas,
+            texture_index: 0,
+        });
 }