From b7ff72ba6efb6c7b6b878b0f0c41d91649853ac2 Mon Sep 17 00:00:00 2001
From: mirrorcult <lunarautomaton6@gmail.com>
Date: Mon, 4 Apr 2022 23:08:36 -0700
Subject: [PATCH] Optimize instruments (#7425)

---
 .../Instruments/InstrumentComponent.cs        |  3 ++
 .../Instruments/InstrumentSystem.cs           | 38 ++++++++++++-------
 Content.Server/PAI/PAISystem.cs               |  9 +++--
 3 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/Content.Server/Instruments/InstrumentComponent.cs b/Content.Server/Instruments/InstrumentComponent.cs
index 36538a75f7..5cb4721b74 100644
--- a/Content.Server/Instruments/InstrumentComponent.cs
+++ b/Content.Server/Instruments/InstrumentComponent.cs
@@ -32,3 +32,6 @@ public sealed class InstrumentComponent : SharedInstrumentComponent
 
     [ViewVariables] public BoundUserInterface? UserInterface => Owner.GetUIOrNull(InstrumentUiKey.Key);
 }
+
+[RegisterComponent]
+public sealed class ActiveInstrumentComponent : Component {}
diff --git a/Content.Server/Instruments/InstrumentSystem.cs b/Content.Server/Instruments/InstrumentSystem.cs
index 137d633e72..ece7c1252d 100644
--- a/Content.Server/Instruments/InstrumentSystem.cs
+++ b/Content.Server/Instruments/InstrumentSystem.cs
@@ -1,15 +1,12 @@
-using System;
 using System.Linq;
 using Content.Server.Stunnable;
 using Content.Server.UserInterface;
 using Content.Shared.Instruments;
 using Content.Shared.Popups;
 using JetBrains.Annotations;
+using Robust.Server.GameObjects;
 using Robust.Server.Player;
 using Robust.Shared.Configuration;
-using Robust.Shared.GameObjects;
-using Robust.Shared.IoC;
-using Robust.Shared.Localization;
 
 namespace Content.Server.Instruments;
 
@@ -18,6 +15,7 @@ public sealed partial class InstrumentSystem : SharedInstrumentSystem
 {
     [Dependency] private readonly IConfigurationManager _cfg = default!;
     [Dependency] private readonly StunSystem _stunSystem = default!;
+    [Dependency] private readonly UserInterfaceSystem _userInterfaceSystem = default!;
 
     public override void Initialize()
     {
@@ -29,7 +27,8 @@ public sealed partial class InstrumentSystem : SharedInstrumentSystem
         SubscribeNetworkEvent<InstrumentStartMidiEvent>(OnMidiStart);
         SubscribeNetworkEvent<InstrumentStopMidiEvent>(OnMidiStop);
 
-        SubscribeLocalEvent<InstrumentComponent, ActivatableUIPlayerChangedEvent>(InstrumentNeedsClean);
+        SubscribeLocalEvent<InstrumentComponent, BoundUIClosedEvent>(OnBoundUIClosed);
+        SubscribeLocalEvent<InstrumentComponent, BoundUIOpenedEvent>(OnBoundUIOpened);
     }
 
     private void OnMidiStart(InstrumentStartMidiEvent msg, EntitySessionEventArgs args)
@@ -66,6 +65,24 @@ public sealed partial class InstrumentSystem : SharedInstrumentSystem
         ShutdownCVars();
     }
 
+    private void OnBoundUIClosed(EntityUid uid, InstrumentComponent component, BoundUIClosedEvent args)
+    {
+        if (HasComp<ActiveInstrumentComponent>(uid)
+            && _userInterfaceSystem.TryGetUi(uid, args.UiKey, out var bui)
+            && bui.SubscribedSessions.Count == 0)
+        {
+            RemComp<ActiveInstrumentComponent>(uid);
+        }
+
+        Clean(uid, component);
+    }
+
+    private void OnBoundUIOpened(EntityUid uid, InstrumentComponent component, BoundUIOpenedEvent args)
+    {
+        EnsureComp<ActiveInstrumentComponent>(uid);
+        Clean(uid, component);
+    }
+
     public void Clean(EntityUid uid, InstrumentComponent? instrument = null)
     {
         if (!Resolve(uid, ref instrument))
@@ -83,11 +100,6 @@ public sealed partial class InstrumentSystem : SharedInstrumentSystem
         instrument.Dirty();
     }
 
-    private void InstrumentNeedsClean(EntityUid uid, InstrumentComponent component, ActivatableUIPlayerChangedEvent ev)
-    {
-        Clean(uid, component);
-    }
-
     private void OnMidiEventRx(InstrumentMidiEventEvent msg, EntitySessionEventArgs args)
     {
         var uid = msg.Uid;
@@ -148,15 +160,15 @@ public sealed partial class InstrumentSystem : SharedInstrumentSystem
     {
         base.Update(frameTime);
 
-        foreach (var instrument in EntityManager.EntityQuery<InstrumentComponent>(true))
+        foreach (var (_, instrument) in EntityManager.EntityQuery<ActiveInstrumentComponent, InstrumentComponent>(true))
         {
             if (instrument.DirtyRenderer)
             {
-                instrument.Dirty();
+                Dirty(instrument);
                 instrument.DirtyRenderer = false;
             }
 
-            if (instrument.RespectMidiLimits && instrument.InstrumentPlayer != null &&
+            if (instrument.RespectMidiLimits &&
                 (instrument.BatchesDropped >= MaxMidiBatchesDropped
                  || instrument.LaggedBatches >= MaxMidiLaggedBatches))
             {
diff --git a/Content.Server/PAI/PAISystem.cs b/Content.Server/PAI/PAISystem.cs
index eba8a44659..6796a24055 100644
--- a/Content.Server/PAI/PAISystem.cs
+++ b/Content.Server/PAI/PAISystem.cs
@@ -119,12 +119,13 @@ namespace Content.Server.PAI
         private void PAITurningOff(EntityUid uid)
         {
             UpdatePAIAppearance(uid, PAIStatus.Off);
+
             //  Close the instrument interface if it was open
             //  before closing
-            if (EntityManager.TryGetComponent<ServerUserInterfaceComponent>(uid, out var serverUi))
-                if (EntityManager.TryGetComponent<ActorComponent>(uid, out var actor))
-                    if (serverUi.TryGetBoundUserInterface(InstrumentUiKey.Key,out var bui))
-                        bui.Close(actor.PlayerSession);
+            if (HasComp<ActiveInstrumentComponent>(uid) && TryComp<ActorComponent>(uid, out var actor))
+            {
+                _instrumentSystem.ToggleInstrumentUi(uid, actor.PlayerSession);
+            }
 
             //  Stop instrument
             if (EntityManager.TryGetComponent<InstrumentComponent>(uid, out var instrument)) _instrumentSystem.Clean(uid, instrument);