RobustToolbox/Robust.Shared/Serialization/Markdown/DataNodeParser.cs

using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using Robust.Shared.Collections;
using Robust.Shared.Serialization.Markdown.Mapping;
using Robust.Shared.Serialization.Markdown.Sequence;
using Robust.Shared.Serialization.Markdown.Value;
using YamlDotNet.Core;
using YamlDotNet.Core.Events;

namespace Robust.Shared.Serialization.Markdown;

// YDN has broken nullable annotations. Yeppers.
#nullable disable

public static class DataNodeParser
{
    public static IEnumerable<DataNodeDocument> ParseYamlStream(TextReader reader)
    {
        return ParseYamlStream(reader, internStrings: false);
    }

    internal static IEnumerable<DataNodeDocument> ParseYamlStream(TextReader reader, bool internStrings)
    {
        return ParseYamlStream(new Parser(reader), internStrings);
    }

    internal static IEnumerable<DataNodeDocument> ParseYamlStream(Parser parser, bool internStrings = false)
    {
        var state = new ParserState(internStrings);

        parser.Consume<StreamStart>();

        while (!parser.TryConsume<StreamEnd>(out _))
        {
            yield return ParseDocument(parser, state);
        }

        // System.Console.WriteLine(state.TotalStringsSaved);
    }

    private static DataNodeDocument ParseDocument(Parser parser, ParserState parserState)
    {
        var state = new DocumentState(parserState);

        parser.Consume<DocumentStart>();

        var root = Parse(parser, state);

        parser.Consume<DocumentEnd>();

        ResolveAliases(state);

        return new DataNodeDocument(root);
    }

    private static DataNode Parse(Parser parser, DocumentState state)
    {
        if (parser.Current is Scalar)
            return ParseValue(parser, state);

        if (parser.Current is SequenceStart)
            return ParseSequence(parser, state);

        if (parser.Current is MappingStart)
            return ParseMapping(parser, state);

        if (parser.Current is AnchorAlias)
            return ParseAlias(parser, state);

        throw new NotSupportedException();
    }

    private static DataNode ParseAlias(Parser parser, DocumentState state)
    {
        var alias = parser.Consume<AnchorAlias>();

        if (!state.Anchors.TryGetValue(alias.Value, out var node))
        {
            // Don't have this anchor yet. It may be defined later in the document.
            return new DataNodeAlias(alias.Value);
        }

        return node;
    }

    private static ValueDataNode ParseValue(Parser parser, DocumentState state)
    {
        var ev = parser.Consume<Scalar>();
        var node = new ValueDataNode(ev)
        {
            Tag = ConvertTag(ev.Tag, state.ParserState),
            Value = state.ParserState.InternString(ev.Value)
        };

        NodeParsed(node, ev, false, state);

        return node;
    }

    private static string ParseKey(Parser parser)
    {
        var ev = parser.Consume<Scalar>();

        if (!ev.Anchor.IsEmpty)
            throw new NotSupportedException();

        return ev.Value;
    }

    private static SequenceDataNode ParseSequence(Parser parser, DocumentState state)
    {
        var ev = parser.Consume<SequenceStart>();

        var node = new SequenceDataNode();
        node.Tag = ConvertTag(ev.Tag, state.ParserState);
        node.Start = ev.Start;

        var unresolvedAlias = false;

        SequenceEnd seqEnd;
        while (!parser.TryConsume(out seqEnd))
        {
            var value = Parse(parser, state);

            node.Add(value);

            unresolvedAlias |= value is DataNodeAlias;
        }

        node.End = seqEnd.End;

        NodeParsed(node, ev, unresolvedAlias, state);

        return node;
    }

    private static MappingDataNode ParseMapping(Parser parser, DocumentState state)
    {
        var ev = parser.Consume<MappingStart>();

        var node = new MappingDataNode();
        node.Tag = ConvertTag(ev.Tag, state.ParserState);

        var unresolvedAlias = false;

        MappingEnd mapEnd;
        while (!parser.TryConsume(out mapEnd))
        {
            var key = state.ParserState.InternString(ParseKey(parser));
            var value = Parse(parser, state);

            node.Add(key, value);

            unresolvedAlias |= value is DataNodeAlias;
        }

        node.End = mapEnd.End;

        NodeParsed(node, ev, unresolvedAlias, state);

        return node;
    }

    private static void NodeParsed(DataNode node, NodeEvent ev, bool unresolvedAlias, DocumentState state)
    {
        if (unresolvedAlias)
            state.UnresolvedAliasOwners.Add(node);

        if (ev.Anchor.IsEmpty)
            return;

        if (state.Anchors.ContainsKey(ev.Anchor))
            throw new DataParseException($"Duplicate anchor defined in document: {ev.Anchor}");

        state.Anchors[ev.Anchor] = node;
    }

    private static void ResolveAliases(DocumentState state)
    {
        foreach (var node in state.UnresolvedAliasOwners)
        {
            switch (node)
            {
                case MappingDataNode mapping:
                    ResolveMappingAliases(mapping, state);
                    break;

                case SequenceDataNode sequence:
                    ResolveSequenceAliases(sequence, state);
                    break;
            }
        }
    }

    private static void ResolveMappingAliases(MappingDataNode mapping, DocumentState state)
    {
        var swaps = new ValueList<(string key, DataNode value)>();

        foreach (var (key, value) in mapping)
        {
            if (value is not DataNodeAlias valueAlias)
                return;

            var newValue = ResolveAlias(valueAlias, state);

            swaps.Add((key, newValue));
            mapping.Remove(key);
        }

        foreach (var (key, value) in swaps)
        {
            mapping[key] = value;
        }
    }

    private static void ResolveSequenceAliases(SequenceDataNode sequence, DocumentState state)
    {
        for (var i = 0; i < sequence.Count; i++)
        {
            if (sequence[i] is DataNodeAlias alias)
                sequence[i] = ResolveAlias(alias, state);
        }
    }

    private static DataNode ResolveAlias(DataNodeAlias alias, DocumentState state)
    {
        if (!state.Anchors.TryGetValue(alias.Anchor, out var node))
            throw new DataParseException($"Unable to resolve alias '{alias.Anchor}'");

        return node;
    }

    private static string ConvertTag(TagName tag, ParserState state)
    {
        return (tag.IsNonSpecific || tag.IsEmpty) ? null : state.InternString(tag.Value);
    }

    private sealed class DocumentState(ParserState parserState)
    {
        public readonly ParserState ParserState = parserState;
        public readonly Dictionary<AnchorName, DataNode> Anchors = new();
        public ValueList<DataNode> UnresolvedAliasOwners;
    }

    private sealed class DataNodeAlias : DataNode
    {
        public readonly AnchorName Anchor;

        public DataNodeAlias(AnchorName anchor) : base(default, default)
        {
            Anchor = anchor;
        }

        public override bool IsEmpty => true;

        public override DataNode Copy()
        {
            throw new NotSupportedException();
        }

        public override DataNode Except(DataNode node)
        {
            throw new NotSupportedException();
        }

        [Obsolete("Use SerializationManager.PushComposition()")]
        public override DataNode PushInheritance(DataNode parent)
        {
            throw new NotSupportedException();
        }
    }

#nullable enable

    private sealed class ParserState(bool internStrings)
    {
        public readonly HashSet<string>? StringInternIndex = internStrings ? [] : null;
        //public int TotalStringsSaved = 0;

        [return: NotNullIfNotNull(nameof(str))]
        public string? InternString(string? str)
        {
            if (StringInternIndex == null)
                return str;

            if (str == null)
                return null;

            // Use a basic string interning system to avoid releasing a bunch of equivalent strings.
            // This avoids having thousands of identical strings for stuff like "type" in prototypes stored in memory.
            if (StringInternIndex.TryGetValue(str, out var indexedString))
            {
                // if (!ReferenceEquals(str, indexedString))
                //     TotalStringsSaved += 1;

                return indexedString;
            }

            StringInternIndex.Add(str);
            return str;
        }
    }
}

public sealed class DataParseException : Exception
{
    public DataParseException()
    {
    }

    public DataParseException(string message) : base(message)
    {
    }

    public DataParseException(string message, Exception inner) : base(message, inner)
    {
    }
}

public sealed record DataNodeDocument(DataNode Root);