ZuneUIXTools/libs/UIX.DecompXml/Decompiler.cs

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.Iris.Asm;
using Microsoft.Iris.Debug.Symbols;
using Microsoft.Iris.DecompXml.Mock;
using Microsoft.Iris.Markup;
using System;
using System.Collections.Generic;
using System.Globalization;
using System.Linq;
using System.Text;
using System.Xml;
using System.Xml.Linq;

namespace Microsoft.Iris.DecompXml;

public partial class Decompiler
{
    private static readonly XNamespace _nsUix = XNamespace.Get("http://schemas.microsoft.com/2007/uix");
    private readonly DecompileContext _context;

    private Decompiler(DecompileContext context)
    {
        _context = context;
    }

    public static Decompiler Load(LoadResult loadResult, LoadResult dataTableLoadResult = null)
    {
        if (loadResult is not MarkupLoadResult markupLoadResult)
            throw new ArgumentException($"Disassembly can only be performed on markup. Expected '{nameof(MarkupLoadResult)}', got '{loadResult?.GetType().Name}'.", nameof(loadResult));

        if (dataTableLoadResult is not MarkupLoadResult and not null)
            throw new ArgumentException($"Data table must be markup. Expected '{nameof(MarkupLoadResult)}', got '{dataTableLoadResult?.GetType().Name}'.", nameof(dataTableLoadResult));
        
        DecompileContext context = new(markupLoadResult, (MarkupLoadResult)dataTableLoadResult);
        return new(context);
    }

    public XDocument Decompile()
    {
        XElement xRoot = new(_nsUix + "UIX", new XAttribute("xmlns", _nsUix));

        foreach (var export in _context.LoadResult.ExportTable.Cast<MarkupTypeSchema>())
        {
            var name = export.Name;

            XElement xExport = new(_nsUix + export.MarkupType.ToString(),
                new XAttribute("Name", name));
            
            var baseType = export.MarkupTypeBase;
            if (baseType is not null)
            {
                var baseTypeName = _context.GetQualifiedName(baseType);
                xExport.SetAttributeValue("Base", baseTypeName);
            }

            if (export is ClassTypeSchema classExport)
            {
                if (classExport.IsShared)
                    xExport.SetAttributeValue("Shared", true);
            }

            // TODO: Combine this with InitializePropertiesOffset
            if (export.Properties is { Length: > 0 })
            {
                XElement xProperties = GetOrCreateElement(xExport, _nsUix + "Properties");

                foreach (var property in export.Properties.Where(p => p.RequiredForCreation))
                {
                    IrisObject value = new("$Required", property.PropertyType);
                    PropertyDictionaryAddOnXElement(xProperties, value, property.Name);
                }
            }

            if (export.InitializePropertiesOffset is not uint.MaxValue)
                AnalyzeMethodForInit(export.InitializePropertiesOffset, xExport, export, name + "_prop");

            if (export.InitializeLocalsInputOffset is not uint.MaxValue)
                AnalyzeMethodForInit(export.InitializeLocalsInputOffset, xExport, export, name + "_locl");

            foreach (var offset in export.InitialEvaluateOffsets ?? [])
            {
                var syntaxTree = DecompileScript(offset, export);
                _context.SetScriptContent(export, offset, syntaxTree);
            }

            foreach (var offset in export.FinalEvaluateOffsets ?? [])
            {
                var statements = DecompileMethod(offset, export);
                AddMethodAttribute(statements, "FinalEvaluate");
                _context.SetScriptContent(export, offset, CreateTree(statements, offset));
            }

            foreach (var offset in export.RefreshGroupOffsets ?? [])
            {
                AnalyzeRefreshMethod(offset, export, $"{name}_rfsh_0x{offset:X}");
            }
            
            XElement xScripts = GetOrCreateElement(xExport, _nsUix + "Scripts");

            foreach (var syntaxTree in _context.GetScriptContents(export))
            {
                var scriptText = FormatScript(syntaxTree);
                XElement xScript = new(_nsUix + "Script", new XCData(Environment.NewLine + scriptText + Environment.NewLine));
                xScripts.Add(xScript);
            }

            if (!xScripts.HasElements)
                xScripts.Remove();

            List<string> methodsContent = [];

            foreach (var method in export.Methods.OfType<MarkupMethodSchema>())
            {
                var methodSyntax = DecompileMethodDeclaration(method, export);
                methodsContent.Add(FormatSyntaxNode(methodSyntax));
            }

            if (methodsContent.Count > 0)
            {
                var methodsContentStr = Environment.NewLine
                    + string.Join(Environment.NewLine + Environment.NewLine, methodsContent)
                    + Environment.NewLine;
                XElement xMethods = new(_nsUix + "Methods", new XCData(methodsContentStr));
                xExport.Add(xMethods);
            }

            if (export.InitializeContentOffset is not uint.MaxValue)
                AnalyzeMethodForInit(export.InitializeContentOffset, xExport, export, name + "_cont");

            xRoot.Add(xExport);
        }

        XDocument xDoc = new(xRoot);

        // Add all namespaces to root element
        var xNamespaceDeclarations = _context.GetUsedNamespaces()
            .Select(p => new XAttribute(XNamespace.Xmlns + p.Key, p.Value))
            .ToArray();

        xDoc.Root.Add(xNamespaceDeclarations);

        return xDoc;
    }

    public string DecompileToSource(bool generateDebugSymbols)
    {
        var xmlDoc = Decompile();

        XmlWriterSettings writerSettings = new()
        {
            Indent = true,
            NamespaceHandling = NamespaceHandling.OmitDuplicates,
            Encoding = Encoding.UTF8,
        };

        StringBuilder sb = new();
        using (XmlWriter writer = XmlWriter.Create(sb, writerSettings))
        {
            xmlDoc.WriteTo(writer);
        }

        var xmlStr = sb.ToString();

        if (generateDebugSymbols)
            PopulateDebugSymbolsWithXml(xmlDoc, xmlStr);

        return xmlStr;
    }

    public void PopulateDebugSymbolsWithXml(XDocument xmlDoc, string xmlStr)
    {
        // Reparse the file to get line info
        var xmlDocWithLineInfo = XDocument.Parse(xmlStr, LoadOptions.SetLineInfo);
        var lineInfoNodes = xmlDocWithLineInfo.DescendantNodes();
        var offsetInfoNodes = xmlDoc.DescendantNodes();

        foreach (var (liNode, oiNode) in lineInfoNodes.Zip(offsetInfoNodes, (l, o) => (l, o)))
        {
            var offsetInfo = oiNode.Annotation<UIBOffsetXmlAnnotation>();
            if (offsetInfo is null)
                continue;
            
            if (liNode is not IXmlLineInfo lineInfo || !lineInfo.HasLineInfo())
                continue;

            var offset = offsetInfo.Offset;
            var start = new SourcePosition(lineInfo.LineNumber, lineInfo.LinePosition);

            var elementStr = oiNode.ToString();
            var endLine = start.Line + elementStr.Count(c => c == '\n');
            var endColumn = elementStr.Length;
            var idxLastLine = elementStr.LastIndexOf('\n');
            if (idxLastLine > 0)
                endColumn -= idxLastLine;

            var end = new SourcePosition(endLine, endColumn);

            var span = new SourceSpan(start, end);

            DebugSymbols.SourceMap.Xml[offset] = span;
        }
    }

    public FileDebugSymbols DebugSymbols => _context.DebugSymbols;

    private Stack<object> AnalyzeMethodForInit(uint startOffset, XElement elemToInit, MarkupTypeSchema initType, string methodName = "")
    {
        var methodBody = _context.GetMethodBody(startOffset).ToArray();

        Stack<object> stack = new([elemToInit]);

        for (int i = 0; i < methodBody.Length; i++)
        {
            var instruction = methodBody[i];

            try
            {
                switch (instruction.OpCode)
                {
                    case OpCode.PushConstant:
                        var constant = _context.GetConstant(instruction.Operands.First());
                        stack.Push(constant);
                        break;

                    case OpCode.PushNull:
                        stack.Push(null);
                        break;

                    case OpCode.ConstructObject:
                        var typeToCtor = _context.GetImportedType(instruction.Operands.ElementAt(0));
                        var xObj = new XElement(_context.GetXName(typeToCtor)).WithOffset(instruction);
                        stack.Push(new IrisObject(xObj, typeToCtor));
                        break;

                    case OpCode.ConstructFromString:
                        var typeFromStringSchema = _context.GetImportedType(instruction.Operands.ElementAt(0));
                        var fromString = _context.GetConstant(instruction.Operands.ElementAt(1));

                        var fromStringObj = IrisObject.Create(fromString, typeFromStringSchema, _context, initType);
                        stack.Push(fromStringObj);
                        break;

                    case OpCode.LookupSymbol:
                        var symbolIndex = (ushort)instruction.Operands.ElementAt(0).Value;
                        stack.Push(initType.SymbolReferenceTable[symbolIndex]);
                        break;

                    case OpCode.PropertyInitialize:
                        var propertyToInit = _context.GetImportedProperty(instruction.Operands.ElementAt(0));
                        var newPropValue = stack.Pop();

                        var target = stack.Pop();
                        var xTarget = (XElement)ToXmlFriendlyObject(target);

                        var xPropSetter = PropertyAssignOnXElement(xTarget, propertyToInit, IrisObject.Create(newPropValue, propertyToInit.PropertyType, _context, initType));
                        xPropSetter.SetOffset(instruction);

                        stack.Push(new IrisObject(xTarget, propertyToInit.Owner));
                        break;

                    case OpCode.PropertyDictionaryAdd:
                        var keyReference = instruction.Operands.ElementAt(1);
                        var key = _context.GetConstant(keyReference).Value.ToString();

                        var dictValue = stack.Pop();
                        var dictValueType = initType.InheritableSymbolsTable?
                            .FirstOrDefault(s => s.Name == key)?
                            .Type;

                        var targetInstance = (XElement)ToXmlFriendlyObject(stack.Peek());
                        var dictValueObj = IrisObject.Create(dictValue, dictValueType, _context, initType);

                        var targetDictPropertyIndex = (ushort)instruction.Operands.ElementAt(0).Value;
                        var xKeyValue = targetDictPropertyIndex is ushort.MaxValue
                            ? PropertyDictionaryAddOnXElement(targetInstance, dictValueObj, key)
                            : PropertyDictionaryAddOnXElement(targetInstance, _context.ImportTables.PropertyImports[targetDictPropertyIndex], dictValueObj, key);
                        xKeyValue.SetOffset(instruction);
                        break;

                    case OpCode.PropertyListAdd:
                        var valueToAdd = stack.Pop();
                        TypeSchema valueToAddType = null;

                        if (valueToAdd is SymbolReference symRef)
                        {
                            valueToAddType = initType.InheritableSymbolsTable?
                                .FirstOrDefault(s => s.Name == symRef.Symbol)?
                                .Type;
                        }

                        var valueToAddObj = IrisObject.Create(valueToAdd, valueToAddType, _context, initType);

                        var targetInstance2 = (XElement)ToXmlFriendlyObject(stack.Peek());

                        XObject xListItem;
                        var targetListPropertyIndex = (ushort)instruction.Operands.First().Value;
                        if (targetListPropertyIndex != ushort.MaxValue)
                        {
                            var targetListProperty = _context.ImportTables.PropertyImports[targetListPropertyIndex];

                            if (valueToAddObj.Type is null)
                            {
                                var valueRuntimeType = targetListProperty.PropertyType.RuntimeType.GetGenericArguments().FirstOrDefault();
                                valueToAddType = _context.ImportTables.TypeImports.FirstOrDefault(t => t.RuntimeType == valueRuntimeType);
                                valueToAddObj = valueToAddObj with { Type = valueToAddType };
                            }

                            xListItem = PropertyListAddOnXElement(targetInstance2, targetListProperty, valueToAddObj);
                        }
                        else
                        {
                            xListItem = PropertyListAddOnXElement(targetInstance2, valueToAddObj);
                        }

                        xListItem.SetOffset(instruction);
                        break;

                    case OpCode.InitializeInstance:
                    case OpCode.JumpIfDictionaryContains:
                    case OpCode.ConstructListenerStorage:
                        // These instructions are inconsequential for determining how objects are initialized
                        break;

                    case OpCode.Jump:
                    case OpCode.JumpIfFalse:
                    case OpCode.JumpIfFalsePeek:
                    case OpCode.JumpIfTruePeek:
                        Console.WriteLine($"Suspicious instruction: {instruction}");

                        // Ensure stack is in valid state
                        var isPeek = instruction.OpCode is OpCode.JumpIfFalsePeek or OpCode.JumpIfTruePeek or OpCode.JumpIfNullPeek;
                        if (isPeek)
                            stack.Peek();
                        else
                            stack.Pop();
                        break;

                    case OpCode.ConstructObjectParam:
                    case OpCode.PushThis:
                    case OpCode.MethodInvoke:
                    case OpCode.MethodInvokePeek:
                    case OpCode.MethodInvokeStatic:
                    case OpCode.MethodInvokePushLastParam:
                    case OpCode.MethodInvokeStaticPushLastParam:
                    case OpCode.PropertyGet:
                    case OpCode.PropertyGetPeek:
                    case OpCode.PropertyGetStatic:
                    case OpCode.Operation:
                    case OpCode.TypeOf:
                    case OpCode.ConvertType:
                        // These instructions only appear in initializers as inline expressions
                        if (!TryDecompileExpression(instruction, stack))
                            throw new NotImplementedException();
                        break;

                    case not OpCode.ReturnVoid:
                        Console.WriteLine($"Unsupported instruction: {instruction}");
                        break;
                }
            }
            catch (Exception ex)
            {
                throw new Exception($"Failed to analyze instruction `{instruction}` @ 0x{instruction.Offset:X}, {methodName}[{i}]", ex);
            }
        }

        return stack;
    }

    private static XElement GetOrCreateElement(XElement parent, XName name)
    {
        var elem = parent.Element(name);

        if (elem is null)
        {
            elem = new XElement(name);
            parent.Add(elem);
        }

        return elem;
    }

    private object ToXmlFriendlyObject(object obj)
    {
        if (obj is Disassembler.RawConstantInfo rci)
        {
            obj = rci.Value;
        }
        else if (obj is IrisObject irisObj)
        {
            obj = irisObj.Object;
        }

        return obj switch
        {
            string str => str,
            null => "{null}",
            bool b => b ? "true" : "false",
            IStringEncodable strEnc => strEnc.EncodeString(),
            ExpressionSyntax expr => FormatInlineExpression(expr),
            SymbolReference symRef => '{' + symRef.Symbol + '}',
            IFormattable formattable => formattable.ToString(null, CultureInfo.InvariantCulture),

            Layout.ILayout layoutObj
                when Layout.PredefinedLayouts.TryConvertToString(layoutObj, out var layoutName)
                => layoutName,

            XElement xElem => xElem,

            _ => SerializeToXml(obj)
        };
    }

    private XElement SerializeToXml(object obj)
    {
        var type = Disassembler.GuessTypeSchema(obj.GetType(), _context.LoadResult);

        XElement xObj = new(_context.GetXName(type));
        var defaultObj = type.ConstructDefault();

        foreach (var prop in type.Properties)
        {
            var defaultPropValue = prop.GetValue(defaultObj);
            var propValue = prop.GetValue(obj);

            if (propValue == defaultPropValue || propValue.Equals(defaultPropValue))
                continue;

            PropertyAssignOnXElement(xObj, prop, new(propValue, prop.PropertyType));
        }

        return xObj;
    }

    private XObject PropertyAssignOnXElement(XElement xTarget, PropertySchema property, IrisObject value)
    {
        object xfValue = ToXmlFriendlyObject(value);

        switch (xfValue)
        {
            case XElement xValue:
                var xProperty = GetOrCreateElement(xTarget, _nsUix + property.Name);

                // Flatten collections
                if (typeof(System.Collections.IList).IsAssignableFrom(value.Type.RuntimeType)
                    || typeof(System.Collections.IDictionary).IsAssignableFrom(value.Type.RuntimeType))
                {
                    xProperty.Add(xValue.Elements());
                }
                else
                {
                    xProperty.Add(xValue);
                }

                return xProperty;

            case string strValue:
                var xAttr = new XAttribute(property.Name, strValue);
                xTarget.Add(xAttr);
                return xAttr;

            default:
                throw new InvalidOperationException();
        }
    }

    private XElement PropertyListAddOnXElement(XElement xTarget, PropertySchema property, IrisObject value)
    {
        var xList = GetOrCreateElement(xTarget, _nsUix + property.Name);
        return PropertyListAddOnXElement(xList, value);
    }

    private XElement PropertyListAddOnXElement(XElement xList, IrisObject value)
    {
        object xValue = ToXmlFriendlyObject(value);

        XElement xListEntry;

        switch (xValue)
        {
            case string strValue:
                xListEntry = new(_context.GetXName(value.Type));
                xListEntry.SetAttributeValue(value.Type.Name, strValue);
                break;
            case XElement xValueELem:
                xListEntry = xValueELem;
                break;
            default:
                throw new InvalidOperationException();
        }

        xList.Add(xListEntry);

        return xListEntry;
    }

    private XElement PropertyDictionaryAddOnXElement(XElement xDictionary, IrisObject value, string key)
    {
        var xDictionaryEntry = PropertyListAddOnXElement(xDictionary, value);

        // Insert the Name attribute (make sure it's always first)
        var attributes = xDictionaryEntry.Attributes().ToArray();
        xDictionaryEntry.RemoveAttributes();
        xDictionaryEntry.Add(new XAttribute("Name", key));
        xDictionaryEntry.Add(attributes);

        return xDictionaryEntry;
    }

    private XElement PropertyDictionaryAddOnXElement(XElement xTarget, PropertySchema property, IrisObject value, string key)
    {
        var xDictionary = GetOrCreateElement(xTarget, _nsUix + property.Name);
        return PropertyDictionaryAddOnXElement(xDictionary, value, key);
    }
}