From 8dfc7216860cb0fd5a7b0a282d4997cb970e0eef Mon Sep 17 00:00:00 2001
From: Danny McGee <dannymcgee@gmail.com>
Date: Thu, 21 Mar 2024 22:34:44 -0400
Subject: [PATCH] Fix selection highlight for primitives

---
 studio/src/interaction/selection.rs | 40 +++++++++++++++++------------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/studio/src/interaction/selection.rs b/studio/src/interaction/selection.rs
index 6483d22..985fd2e 100644
--- a/studio/src/interaction/selection.rs
+++ b/studio/src/interaction/selection.rs
@@ -213,22 +213,30 @@ fn draw_selection_outline(
         cmd.entity(ent).remove::<OutlineBundle>();
     }
 
-    if let Some(selected_children) = r_selection
-        .as_ref()
-        .and_then(|selection| q_children.get(***selection).ok())
-    {
-        for &child in selected_children.iter() {
-            if q_meshes.contains(child) {
-                cmd.entity(child).insert(OutlineBundle {
-                    outline: OutlineVolume {
-                        visible: true,
-                        width: 4.,
-                        colour: Color::YELLOW,
-                    },
-                    ..default()
-                });
-            }
-        }
+    let Some(selection) = r_selection else {
+        return;
+    };
+
+    let target = if q_meshes.contains(**selection) {
+        Some(**selection)
+    } else {
+        q_children.get(**selection).ok().and_then(|children| {
+            children
+                .iter()
+                .copied()
+                .find(|&child| q_meshes.contains(child))
+        })
+    };
+
+    if let Some(target) = target {
+        cmd.entity(target).insert(OutlineBundle {
+            outline: OutlineVolume {
+                visible: true,
+                width: 4.,
+                colour: Color::YELLOW,
+            },
+            ..default()
+        });
     }
 }