CatStation/Content.Client/Effects/ScreechShockWaveSystem.cs

using Content.Client.Overlays;
using Content.Shared.Screech;
using Robust.Client.Graphics;

namespace Content.Client.Effects;

/// <summary>
/// This system ensures that <see cref="ScreechShockWaveOverlay"/> does not use costly queries.
/// </summary>
public sealed partial class ScreechShockWaveSystem : EntitySystem
{
    [Dependency] private IOverlayManager _overlayMan = default!;

    private readonly HashSet<EntityUid> _registered = [];

    public override void Initialize()
    {
        SubscribeLocalEvent<ScreechShockWaveComponent, ComponentStartup>(OnScreechShockWaveStartup);
        SubscribeLocalEvent<ScreechShockWaveComponent, ComponentRemove>(OnScreechShockWaveRemoved);
    }

    private void OnScreechShockWaveStartup(Entity<ScreechShockWaveComponent> ent, ref ComponentStartup args)
    {
        var overlay = _overlayMan.GetOverlay<ScreechShockWaveOverlay>();
        // we must only pass here once
        if (!_registered.Add(ent.Owner))
            return;

        overlay.Register(ent);
    }

    private void OnScreechShockWaveRemoved(Entity<ScreechShockWaveComponent> ent, ref ComponentRemove args)
    {
        _registered.Remove(ent.Owner);
    }
}