From ea8a9e9d361ff489afece85db01babf2cc8ac437 Mon Sep 17 00:00:00 2001
From: David Gileadi <gileadis@gmail.com>
Date: Sat, 19 Aug 2017 21:57:27 -0700
Subject: [PATCH] Preserve sprite modifications in more situations

---
 Editor/UIExtensionsMenuOptions.cs    |  1 +
 Scripts/Controls/SegmentedControl.cs | 30 +++++++++++++++++++---------
 Scripts/Controls/Stepper.cs          |  5 +++++
 3 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/Editor/UIExtensionsMenuOptions.cs b/Editor/UIExtensionsMenuOptions.cs
index b8acb64..310a0bf 100644
--- a/Editor/UIExtensionsMenuOptions.cs
+++ b/Editor/UIExtensionsMenuOptions.cs
@@ -1666,6 +1666,7 @@ namespace UnityEditor.UI
             for (int i = 0; i < 3; i++)
 			{
                 var button = AddButtonAsChild(go);
+				button.gameObject.AddComponent<Segment>();
                 button.name = "Segment " + (i + 1);
                 var text = button.GetComponentInChildren<Text>();
                 text.text = labels[i];
diff --git a/Scripts/Controls/SegmentedControl.cs b/Scripts/Controls/SegmentedControl.cs
index 2b25a2a..2e1aba7 100644
--- a/Scripts/Controls/SegmentedControl.cs
+++ b/Scripts/Controls/SegmentedControl.cs
@@ -74,21 +74,24 @@ namespace UnityEngine.UI.Extensions
             {
                 value = Math.Max(value, -1);
                 value = Math.Min(value, segments.Length - 1);
+                if (value == m_selectedSegmentIndex)
+                    return;
+
                 m_selectedSegmentIndex = value;
-                if (value == -1)
+
+                if (selectedSegment)
                 {
-                    if (selectedSegment)
+                    var segment = selectedSegment.GetComponent<Segment>();
+                    if (segment)
                     {
-                        var segment = selectedSegment.GetComponent<Segment>();
-                        if (segment)
-                        {
-                            segment.selected = false;
-                        }
-                        selectedSegment = null;
+                        segment.selected = false;
                     }
+                    selectedSegment = null;
                 }
-                else
+
+                if (value != -1)
                 {
+                    selectedSegment = segments[value];
                     var segment = selectedSegment.GetComponent<Segment>();
                     if (segment)
                     {
@@ -314,6 +317,13 @@ namespace UnityEngine.UI.Extensions
             MaintainSelection();
         }
 
+        protected override void OnEnable()
+        {
+            base.OnEnable();
+            if (segmentedControl)
+                MaintainSelection();
+        }
+
         public virtual void OnSubmit(BaseEventData eventData)
         {
             selected = true;
@@ -388,6 +398,7 @@ namespace UnityEngine.UI.Extensions
             switch (button.transition)
             {
                 case Selectable.Transition.ColorTint:
+                    button.image.overrideSprite = cutSprite;
                     StartColorTween(tintColor * button.colors.colorMultiplier, instant);
                     ChangeTextColor(textColor * button.colors.colorMultiplier);
                     break;
@@ -397,6 +408,7 @@ namespace UnityEngine.UI.Extensions
                     DoSpriteSwap(transitionSprite);
                     break;
                 case Selectable.Transition.Animation:
+                    button.image.overrideSprite = cutSprite;
                     TriggerAnimation(triggerName);
                     break;
             }
diff --git a/Scripts/Controls/Stepper.cs b/Scripts/Controls/Stepper.cs
index 1da6af9..721b890 100644
--- a/Scripts/Controls/Stepper.cs
+++ b/Scripts/Controls/Stepper.cs
@@ -104,6 +104,11 @@ namespace UnityEngine.UI.Extensions
         }
 #endif
 
+        protected override void Start()
+        {
+            RecreateSprites(sides);
+        }
+
         private Selectable[] GetSides()
         {
             var buttons = GetComponentsInChildren<Selectable>();