From e2b0ed98dd0a48435d222a8442ff6364e7c057db Mon Sep 17 00:00:00 2001
From: EliteMasterEric <ericmyllyoja@gmail.com>
Date: Wed, 13 Dec 2023 23:11:19 -0500
Subject: [PATCH] Fix a bug where making one selection and then making a
 different one would leave the original highlighted

---
 .../ui/debug/charting/ChartEditorState.hx     | 21 +++++---
 .../components/ChartEditorNotePreview.hx      | 17 +++++-
 source/funkin/util/tools/ArrayTools.hx        | 52 +++++++++++++++++++
 3 files changed, 82 insertions(+), 8 deletions(-)

diff --git a/source/funkin/ui/debug/charting/ChartEditorState.hx b/source/funkin/ui/debug/charting/ChartEditorState.hx
index 78de08fdf..b24bc1d95 100644
--- a/source/funkin/ui/debug/charting/ChartEditorState.hx
+++ b/source/funkin/ui/debug/charting/ChartEditorState.hx
@@ -753,15 +753,23 @@ class ChartEditorState extends UIState // UIState derives from MusicBeatState
 
   function set_currentNoteSelection(value:Array<SongNoteData>):Array<SongNoteData>
   {
+    // This value is true if all elements of the current selection are also in the new selection.
+    var isSuperset:Bool = currentNoteSelection.isSubset(value);
+    var isEqual:Bool = currentNoteSelection.isEqualUnordered(value);
+
     currentNoteSelection = value;
 
-    if (currentNoteSelection.length > 0)
+    if (!isEqual)
     {
-      notePreview.addNotes(currentNoteSelection, Std.int(songLengthInMs), true);
-    }
-    else
-    {
-      notePreviewDirty = true;
+      if (currentNoteSelection.length > 0 && isSuperset)
+      {
+        notePreview.addSelectedNotes(currentNoteSelection, Std.int(songLengthInMs));
+      }
+      else
+      {
+        // The new selection removes elements from the old selection, so we have to redraw the note preview.
+        notePreviewDirty = true;
+      }
     }
 
     return currentNoteSelection;
@@ -5295,6 +5303,7 @@ class ChartEditorState extends UIState // UIState derives from MusicBeatState
       // TODO: Only update the notes that have changed.
       notePreview.erase();
       notePreview.addNotes(currentSongChartNoteData, Std.int(songLengthInMs));
+      notePreview.addSelectedNotes(currentNoteSelection, Std.int(songLengthInMs));
       notePreview.addEvents(currentSongChartEventData, Std.int(songLengthInMs));
     }
 
diff --git a/source/funkin/ui/debug/charting/components/ChartEditorNotePreview.hx b/source/funkin/ui/debug/charting/components/ChartEditorNotePreview.hx
index 09c99531d..598cbb544 100644
--- a/source/funkin/ui/debug/charting/components/ChartEditorNotePreview.hx
+++ b/source/funkin/ui/debug/charting/components/ChartEditorNotePreview.hx
@@ -80,11 +80,24 @@ class ChartEditorNotePreview extends FlxSprite
    * @param notes The data for the notes.
    * @param songLengthInMs The total length of the song in milliseconds.
    */
-  public function addNotes(notes:Array<SongNoteData>, songLengthInMs:Int, ?isSelection:Bool = false):Void
+  public function addNotes(notes:Array<SongNoteData>, songLengthInMs:Int):Void
   {
     for (note in notes)
     {
-      addNote(note, songLengthInMs, isSelection);
+      addNote(note, songLengthInMs, false);
+    }
+  }
+
+  /**
+   * Add an array of selected notes to the preview.
+   * @param notes The data for the notes.
+   * @param songLengthInMs The total length of the song in milliseconds.
+   */
+  public function addSelectedNotes(notes:Array<SongNoteData>, songLengthInMs:Int):Void
+  {
+    for (note in notes)
+    {
+      addNote(note, songLengthInMs, true);
     }
   }
 
diff --git a/source/funkin/util/tools/ArrayTools.hx b/source/funkin/util/tools/ArrayTools.hx
index a88f8a861..925bb67a5 100644
--- a/source/funkin/util/tools/ArrayTools.hx
+++ b/source/funkin/util/tools/ArrayTools.hx
@@ -76,4 +76,56 @@ class ArrayTools
     while (array.length > 0)
       array.pop();
   }
+
+  /**
+   * Return true only if both arrays contain the same elements (possibly in a different order).
+   * @param a The first array to compare.
+   * @param b The second array to compare.
+   * @return Weather both arrays contain the same elements.
+   */
+  public static function isEqualUnordered<T>(a:Array<T>, b:Array<T>):Bool
+  {
+    if (a.length != b.length) return false;
+    for (element in a)
+    {
+      if (!b.contains(element)) return false;
+    }
+    for (element in b)
+    {
+      if (!a.contains(element)) return false;
+    }
+    return true;
+  }
+
+  /**
+   * Returns true if `superset` contains all elements of `subset`.
+   * @param superset The array to query for each element.
+   * @param subset The array containing the elements to query for.
+   * @return Weather `superset` contains all elements of `subset`.
+   */
+  public static function isSuperset<T>(superset:Array<T>, subset:Array<T>):Bool
+  {
+    // Shortcuts.
+    if (subset.length == 0) return true;
+    if (subset.length > superset.length) return false;
+
+    // Check each element.
+    for (element in subset)
+    {
+      if (!superset.contains(element)) return false;
+    }
+    return true;
+  }
+
+  /**
+   * Returns true if `superset` contains all elements of `subset`.
+   * @param subset The array containing the elements to query for.
+   * @param superset The array to query for each element.
+   * @return Weather `superset` contains all elements of `subset`.
+   */
+  public static function isSubset<T>(subset:Array<T>, superset:Array<T>):Bool
+  {
+    // Switch it around.
+    return isSuperset(superset, subset);
+  }
 }