RobustToolbox/Robust.Client/UserInterface/DevWindow/ProfTree.xaml.cs

using System.Collections.Generic;
using System.Runtime.InteropServices;
using Robust.Client.AutoGenerated;
using Robust.Client.Profiling;
using Robust.Client.UserInterface.XAML;
using Robust.Shared.Collections;
using Robust.Shared.IoC;
using Robust.Shared.Log;
using Robust.Shared.Maths;
using Robust.Shared.Profiling;

namespace Robust.Client.UserInterface;

[GenerateTypedNameReferences]
internal sealed partial class ProfTree : Control
{
    public const float TreeLevelMargin = 16;

    [Dependency] private readonly ProfManager _prof = default!;
    [Dependency] private readonly ProfViewManager _profViewMgr = default!;

    private ProfViewManager.Snapshot? _snapshot;
    public long Frame;

    private readonly TreeExpand _treeExpand = new();

    public ProfTree()
    {
        RobustXamlLoader.Load(this);
        IoCManager.InjectDependencies(this);

        Header.Text.Text = "Thing";
        Header.SideLabel.Text = "Misc";
        Header.TimeLabel.Text = "Time";
        Header.PercentTimeLabel.Text = "Time%";
        Header.AllocLabel.Text = "Alloc";
        Header.PercentAllocLabel.Text = "Alloc%";
    }

    public void LoadSnapshot(ProfViewManager.Snapshot snapshot)
    {
        _snapshot = snapshot;
    }

    public void RebuildTree()
    {
        if (_snapshot == null)
            return;

        TreeRoot.RemoveAllChildren();

        ref var buf = ref _snapshot.Buffer;
        var indexIdx = _profViewMgr.GetIndexOfFrame(Frame, _snapshot);
        if (indexIdx == 0)
        {
            Logger.WarningS("prof.ui", $"Unable to find index for frame: {Frame}");
            return;
        }

        // Index for this frame's data.
        ref var index = ref _snapshot.Buffer.Index(indexIdx);

        // The frame must be wrapped in a whole group. We display this group.
        var i = index.EndPos - 1;
        ref var logEnd = ref buf.Log(i);
        var controls = new ValueList<Control>();

        // Create child controls.
        var countDict = new Dictionary<int, int>();
        var totalFrameTime = ProfGraphView.GetFrameTime(buf, index);
        var (totalCounts, _) = GetTotalGroupCounts(buf, i, logEnd.GroupEnd.StartIndex);

        BuildData data = default;
        data.Buffer = buf;
        data.Index = index;
        data.TotalFrameTime = totalFrameTime;

        // We traverse the log backwards because we can use GroupEnd nodes to skip over data entirely.

        for (; i >= index.StartPos; i--)
        {
            // Recursively build the tree.
            ref var log = ref buf.Log(i);
            RebuildTreeAddControls(
                in data,
                ref controls,
                ref i,
                in log,
                totalCounts,
                countDict,
                _treeExpand,
                2);
        }

        // Because we traversed backwards, the control list is backwards too.
        // Insert backwards again to fix that.
        for (var c = controls.Count - 1; c >= 0; c--)
        {
            TreeRoot.AddChild(controls[c]);
        }

        // Do a final pass to assign the alternating background colors.
        // Easier to do here than to do it while we're creating the controls.
        var alt = false;
        foreach (var child in TreeRoot.Children)
        {
            DoAltBackgrounds(child, ref alt);
        }
    }

    private static void DoAltBackgrounds(Control control, ref bool alt)
    {
        if (control is not ProfAltBackground altBg)
            return;

        altBg.IsAltBackground = alt;
        alt = !alt;
    }

    private void RebuildTreeAddControls(
        in BuildData data,
        ref ValueList<Control> insertInto,
        ref long i,
        in ProfLog log,
        Dictionary<int, int> totalCounts,
        Dictionary<int, int> countDict,
        TreeExpand expandParent,
        float margin)
    {
        switch (log.Type)
        {
            case ProfLogType.Value:
                TreeInsertSample(in data, ref insertInto, log.Value.StringId, log.Value.Value, margin);
                break;

            case ProfLogType.GroupEnd:
                var stringId = log.GroupEnd.StringId;

                ref var count = ref CollectionsMarshal.GetValueRefOrAddDefault(countDict, stringId, out _);
                count += 1;
                var totalCount = totalCounts[stringId];

                var stringI = totalCount - count;

                TreeInsertGroup(
                    in data,
                    ref insertInto,
                    ref i,
                    in log.GroupEnd,
                    (stringId, stringI),
                    expandParent,
                    margin);
                break;
        }
    }

    private void TreeInsertSample(
        in BuildData data,
        ref ValueList<Control> insertInto,
        int stringId,
        in ProfValue value,
        float margin)
    {
        var treeLine = new ProfTreeLine();
        treeLine.Margin = new Thickness(treeLine.Margin.Left + margin + 12 + 2, 0, 0, 0);
        FillTreeLine(in data, treeLine, stringId, value);
        insertInto.Add(new ProfAltBackground { Children = { treeLine } });
    }

    private void TreeInsertGroup(
        in BuildData data,
        ref ValueList<Control> insertInto,
        ref long i,
        in ProfLogGroupEnd logEnd,
        (int str, int i) expandId,
        TreeExpand expandParent,
        float margin)
    {
        i -= 1;

        var (totalCounts, anyChildren) = GetTotalGroupCounts(data.Buffer, i, logEnd.StartIndex);

        if (!anyChildren)
        {
            // Node has no children we can display. Just insert it as if it's a sample and move along.
            TreeInsertSample(in data, ref insertInto, logEnd.StringId, logEnd.Value, margin);
            i = logEnd.StartIndex;
            return;
        }

        var groupControl = new ProfTreeEntry(this, expandParent, expandId, margin);
        var altBg = new ProfAltBackground { Children = { groupControl } };
        FillTreeLine(in data, groupControl.TextLine, logEnd.StringId, logEnd.Value);


        // Look up in the parent's expand data whether we're expanded.
        var expand = expandParent.ExpandedItems.GetValueOrDefault(expandId);
        if (expand is not { Enabled: true })
        {
            // Note expanded. Skip!
            i = logEnd.StartIndex;
            insertInto.Add(altBg);
            return;
        }

        groupControl.Arrow.Rotated = true;

        // Create child controls.
        var countDict = new Dictionary<int, int>();

        for (; i >= data.Index.StartPos; i--)
        {
            ref var log = ref data.Buffer.Log(i);
            if (log.Type == ProfLogType.GroupStart)
                break;

            RebuildTreeAddControls(
                in data,
                ref insertInto,
                ref i,
                in log,
                totalCounts,
                countDict,
                expand,
                margin + TreeLevelMargin);
        }

        insertInto.Add(altBg);
    }

    private static (Dictionary<int, int> counts, bool anyChildren) GetTotalGroupCounts(
        in ProfBuffer buffer,
        long i,
        long startIdx)
    {
        var anyChildren = false;
        var dict = new Dictionary<int, int>();

        for (; i > startIdx; i--)
        {
            ref var log = ref buffer.Log(i);

            anyChildren |= log.Type is ProfLogType.GroupEnd or ProfLogType.Value;

            if (log.Type != ProfLogType.GroupEnd)
                continue;

            i = log.GroupEnd.StartIndex;
            ref var val = ref CollectionsMarshal.GetValueRefOrAddDefault(dict, log.GroupEnd.StringId, out _);
            val += 1;
        }

        return (dict, anyChildren);
    }

    private void FillTreeLine(
        in BuildData data,
        ProfTreeLine line,
        int stringId,
        in ProfValue value)
    {
        line.Text.Text = _prof.GetString(stringId);

        switch (value.Type)
        {
            case ProfValueType.TimeAllocSample:
                line.TimeLabel.Text = $"{value.TimeAllocSample.Time * 1000:N2} ms";
                line.PercentTimeLabel.Text = $"{value.TimeAllocSample.Time / data.TotalFrameTime.Time:P2}";
                line.AllocLabel.Text = $"{value.TimeAllocSample.Alloc} B";
                line.PercentAllocLabel.Text = $"{value.TimeAllocSample.Alloc / (float)data.TotalFrameTime.Alloc:P2}";
                break;
            case ProfValueType.Int32:
                line.SideLabel.Text = value.Int32.ToString();
                break;
            case ProfValueType.Int64:
                line.SideLabel.Text = value.Int64.ToString();
                break;
            default:
                line.SideLabel.Text = "???";
                break;
        }
    }

    internal sealed class TreeExpand
    {
        public bool Enabled = true;
        public readonly Dictionary<(int str, int i), TreeExpand> ExpandedItems = new();
    }

    private struct BuildData
    {
        public ProfBuffer Buffer;
        public ProfIndex Index;
        public TimeAndAllocSample TotalFrameTime;
    }
}