[WIP] Data table imports

UIXC/Commands/DecompileCommand.cs

@@ -1,12 +1,9 @@
 using Microsoft.Iris.Asm;
 using Microsoft.Iris.Debug;
 using Microsoft.Iris.Markup;
-using Microsoft.Iris.Session;
 using Spectre.Console;
 using Spectre.Console.Cli;
 using System.ComponentModel;
-using System.Diagnostics.CodeAnalysis;
-using System.Runtime.Loader;
 
 namespace UIXC.Commands;
 
@@ -47,23 +44,24 @@ public class DecompileCommand : CompilerCommandBase<DecompileCommand.Settings>
         MarkupSystem.Startup(true);
         bool success = false;
 
+        // Load the shared data table if one was specified
+        LoadResult? dataTableLoadResult = null;
+        if (settings.DataTable is not null)
+        {
+            dataTableLoadResult = LoadIrisFile(settings.DataTable, settings);
+            if (dataTableLoadResult.Status != LoadResultStatus.Success)
+                throw new Exception($"Failed to load data table from {dataTableLoadResult.ErrorContextUri}");
+        }
+
         if (settings.Language == SourceLanguage.Asm)
         {
             foreach (var input in settings.Inputs)
             {
                 try
                 {
-                    if (!TryResolvePath(input, GetSearchPaths(settings), out var inputPath))
-                        inputPath = input;
-
-                    if (!inputPath.Contains("://"))
-                        inputPath = $"file://{inputPath}";
-
-                    var uibLoadResult = MarkupSystem.Load(inputPath, (uint)Random.Shared.Next());
-                    uibLoadResult.FullLoad();
-
+                    var uibLoadResult = LoadIrisFile(input, settings);
                     if (uibLoadResult.Status != LoadResultStatus.Success)
-                        throw new Exception($"Failed to load UIB source ({uibLoadResult.ErrorContextUri})");
+                        throw new Exception($"Failed to load UIB source from {uibLoadResult.ErrorContextUri}");
 
                     var disassembler = Disassembler.Load(uibLoadResult);
                     var asm = disassembler.Write();
@@ -111,6 +109,19 @@ public class DecompileCommand : CompilerCommandBase<DecompileCommand.Settings>
         return outputFile;
     }
 
+    private LoadResult LoadIrisFile(string input, Settings settings)
+    {
+        if (!TryResolvePath(input, GetSearchPaths(settings), out var inputPath))
+            inputPath = input;
+
+        if (!inputPath.Contains("://"))
+            inputPath = $"file://{inputPath}";
+
+        var uibLoadResult = MarkupSystem.Load(inputPath, (uint)Random.Shared.Next());
+        uibLoadResult.FullLoad();
+        return uibLoadResult;
+    }
+
     public sealed class Settings : CompilerSettings
     {
         [Description("The UIB files to decompile.")]

libs/UIX.Asm/AsmMarkupLoader.cs

@@ -55,7 +55,7 @@ internal class AsmMarkupLoader
 
     public bool HasErrors { get; protected set; }
 
-    private Program Program { get; set; }
+    public Program Program { get; private set; }
 
     public LoadResult FindDependency(string prefix)
     {

libs/UIX.Asm/Disassembler.cs

@@ -10,14 +10,24 @@ namespace Microsoft.Iris.Asm;
 
 public class Disassembler
 {
+    private Dictionary<int, string> _constantsTable = [];
+    private Program _program = null;
+
     private readonly MarkupLoadResult _loadResult;
+    private readonly MarkupLoadResult _dataTableLoadResult;
     private readonly Dictionary<string, string> _importedUris;
     private readonly Dictionary<uint, List<Label>> _offsetLabelMap = new();
+    
     private static readonly TypeSchema _stringTypeSchema = UIXTypes.MapIDToType(UIXTypeID.String);
 
-    private Disassembler(MarkupLoadResult loadResult)
+    private Disassembler(MarkupLoadResult loadResult, MarkupLoadResult dataTableLoadResult = null)
     {
         _loadResult = loadResult;
+
+        _dataTableLoadResult = dataTableLoadResult is not null
+            ? dataTableLoadResult
+            : loadResult.BinaryDataTable?.SharedDependenciesTableWithBinaryDataTable?.FirstOrDefault() as MarkupLoadResult;
+
         _importedUris = new()
         {
             [_loadResult.Uri] = "me",
@@ -25,15 +35,58 @@ public class Disassembler
         };
     }
 
-    public static Disassembler Load(LoadResult loadResult)
+    public static Disassembler 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}'.");
-        
+            throw new ArgumentException($"Disassembly can only be performed on markup. Expected '{nameof(MarkupLoadResult)}', got '{loadResult?.GetType().Name}'.", nameof(loadResult));
+
+        if (dataTableLoadResult is MarkupLoadResult dataTableMarkupLoadResult)
+            return new(markupLoadResult, dataTableMarkupLoadResult);
+
+        if (dataTableLoadResult is not null)
+            throw new ArgumentException($"Data table must be markup. Expected '{nameof(MarkupLoadResult)}', got '{dataTableLoadResult?.GetType().Name}'.", nameof(dataTableLoadResult));
+
         return new(markupLoadResult);
     }
 
-    public IEnumerable<Directive> GetExports()
+    public Program Disassemble()
+    {
+        if (_program is null)
+        {
+            _loadResult.FullLoad();
+            if (_loadResult.Status == LoadResultStatus.Error)
+                throw new Exception($"Failed to load '{_loadResult.ErrorContextUri}'");
+
+            if (_dataTableLoadResult is not null)
+            {
+                _dataTableLoadResult.FullLoad();
+                if (_dataTableLoadResult.Status == LoadResultStatus.Error)
+                    throw new Exception($"Failed to load '{_dataTableLoadResult.ErrorContextUri}'");
+            }
+
+            List<IEnumerable<IBodyItem>> segments = [
+                GetExports(),
+                GetImports(),
+                GetConstants(),
+                GetCode(),
+            ];
+
+            List<IBodyItem> body = [];
+            foreach (var segment in segments)
+                body.AddRange(segment);
+
+            if (_dataTableLoadResult is not null)
+                body.Add(new SharedDataTableDirective(_dataTableLoadResult.Uri));
+
+            _program = new(body);
+        }
+
+        return _program;
+    }
+
+    public string Write() => Disassemble().ToString();
+
+    private IEnumerable<ExportDirective> GetExports()
     {
         foreach (var typeSchema in _loadResult.ExportTable)
         {
@@ -91,9 +144,9 @@ public class Disassembler
         }
     }
 
-    public IEnumerable<IImportDirective> GetImports() => EnumerateImports().Distinct();
+    private IEnumerable<IImportDirective> GetImports() => EnumerateImports().Distinct();
 
-    public IEnumerable<IImportDirective> EnumerateImports()
+    private IEnumerable<IImportDirective> EnumerateImports()
     {
         // Ues _importedUris to keep track of what has already been imported.
         // Skip self and default UIX namespace.
@@ -101,9 +154,9 @@ public class Disassembler
         foreach (var typeImport in _loadResult.ImportTables.TypeImports)
         {
             var uri = typeImport.Owner.Uri;
-            if (!_importedUris.TryGetValue(uri, out var namespacePrefix))
+            if (!_importedUris.TryGetValue(uri, out var baseNamespacePrefix))
             {
-                namespacePrefix = uri;
+                baseNamespacePrefix = uri;
                 var schemeLength = uri.IndexOf("://");
                 if (schemeLength > 0)
                 {
@@ -116,7 +169,7 @@ public class Disassembler
                         if (nsIndex >= 0)
                         {
                             var importedNamespace = path[(nsIndex + 1)..];
-                            namespacePrefix = importedNamespace.Split('.', '/', '\\', '!')[^1];
+                            baseNamespacePrefix = importedNamespace.Split('.', '/', '\\', '!')[^1];
 
                             System.Reflection.AssemblyName assemblyName = new(path[..nsIndex]);
                             uri = $"assembly://{assemblyName.Name}/{importedNamespace}";
@@ -125,16 +178,16 @@ public class Disassembler
                         {
                             // No namespace was specified, assume we're importing the whole assembly
                             System.Reflection.AssemblyName assemblyName = new(path);
-                            namespacePrefix = assemblyName.Name;
+                            baseNamespacePrefix = assemblyName.Name;
 
-                            uri = $"assembly://{assemblyName.Name}";
+                            uri = $"assembly://{assemblyName.Name}/";
                         }
                     }
                     else
                     {
                         // Assume the URI represents a file,
                         // skip the extension
-                        namespacePrefix = uri.Split('.', '/', '\\', '!')[^2];
+                        baseNamespacePrefix = uri.Split('.', '/', '\\', '!')[^2];
                     }
                 }
 
@@ -143,13 +196,20 @@ public class Disassembler
                 if (_importedUris.ContainsKey(uri))
                     continue;
 
-                namespacePrefix = namespacePrefix.Camelize();
+                baseNamespacePrefix = baseNamespacePrefix.Camelize();
+                var namespacePrefix = baseNamespacePrefix;
+                
+                // Prevent similar imports from generating the same prefix
+                int duplicateCount = 0;
+                while (_importedUris.ContainsValue(namespacePrefix))
+                    namespacePrefix = $"{baseNamespacePrefix}{++duplicateCount}";
+
                 _importedUris.Add(uri, namespacePrefix);
 
                 yield return new NamespaceImport(uri, namespacePrefix);
             }
 
-            yield return new TypeImport(new(namespacePrefix, typeImport.Name));
+            yield return new TypeImport(new(baseNamespacePrefix, typeImport.Name));
         };
 
         foreach (var constructorImport in _loadResult.ImportTables.ConstructorImports)
@@ -187,63 +247,7 @@ public class Disassembler
             yield return new NamedMemberImport(groupedImport.Key, groupedImport.Value);
     }
 
-    public IEnumerable<ConstantDirective> GetConstants()
-    {
-        List<(int c, TypeSchema typeSchema, object constantValue)> constants = new();
-
-        var constantsTable = _loadResult.ConstantsTable;
-        bool hasPersistList = constantsTable.PersistList is not null;
-        bool hasSharedBinaryTable = _loadResult.BinaryDataTable?.SharedDependenciesTableWithBinaryDataTable is not null;
-
-        if (hasPersistList)
-        {
-            var persistedList = _loadResult.ConstantsTable.PersistList;
-
-            for (int c = 0; c < persistedList.Length; c++)
-            {
-                var persistedConstant = persistedList[c];
-                var typeSchema = persistedConstant.Type;
-
-                constants.Add((c, typeSchema, persistedConstant.Data));
-            }
-        }
-        else if (hasSharedBinaryTable)
-        {
-            // Constants need to be imported from the shared binary table
-
-        }
-        else
-        {
-            // UIB doesn't persist constants, so we have to use an alternate, slower method
-            for (int c = 0; ; c++)
-            {
-                object constantValue;
-                try
-                {
-                    constantValue = constantsTable.Get(c);
-                }
-                catch
-                {
-                    break;
-                }
-
-                var runtimeType = constantValue.GetType();
-                var typeSchema = _loadResult.ImportTables.TypeImports.FirstOrDefault(t => t.RuntimeType == runtimeType)
-                    ?? _loadResult.ImportTables.TypeImports.FirstOrDefault(t => t.RuntimeType == runtimeType.BaseType)
-                    ?? throw new Exception($"Failed to find type schema for '{runtimeType.Name}' in {_loadResult.Uri}");
-
-                constants.Add((c, typeSchema, constantValue));
-            }
-        }
-
-        foreach (var (c, typeSchema, constantValue) in constants)
-        {
-            var constantName = $"const{c:D}";
-            yield return EncodeConstant(constantValue, constantName, typeSchema);
-        }
-    }
-
-    public IEnumerable<IBodyItem> GetCode()
+    private IEnumerable<IBodyItem> GetCode()
     {
         var reader = _loadResult.ObjectSection;
 
@@ -297,27 +301,6 @@ public class Disassembler
         yield break;
     }
 
-    public string Write()
-    {
-        _loadResult.FullLoad();
-        if (_loadResult.Status == LoadResultStatus.Error)
-            throw new Exception($"Failed to load '{_loadResult.ErrorContextUri}'");
-
-        List<IEnumerable<IBodyItem>> segments = [
-            GetExports(),
-            GetImports(),
-            GetConstants(),
-            GetCode(),
-        ];
-
-        List<IBodyItem> body = [];
-        foreach (var segment in segments)
-            body.AddRange(segment);
-
-        Program asmProgram = new(body);
-        return asmProgram.ToString();
-    }
-
     private QualifiedTypeName GetQualifiedName(TypeSchema schema)
     {
         _importedUris.TryGetValue(schema.Owner.Uri, out string prefix);
@@ -331,6 +314,78 @@ public class Disassembler
         labels.Add(new(labelName));
     }
 
+    private IEnumerable<ConstantDirective> GetConstants()
+    {
+        if (_dataTableLoadResult is not null)
+        {
+            // Constants have already been disassembled from the shared binary table
+            var asmConstants = _dataTableLoadResult is AsmMarkupLoadResult asmDataTableLoadResult
+                ? asmDataTableLoadResult.Loader.Program.Directives.OfType<ConstantDirective>().ToList()
+                : null;
+
+            foreach (var info in EnumerateConstantInfo(_dataTableLoadResult))
+            {
+                var constantName = asmConstants?[info.Index]?.Name
+                    ?? $"sharedConst{info.Index:D}";
+
+                _constantsTable.Add(info.Index, constantName);
+            }
+
+            yield break;
+        }
+        else
+        {
+            foreach (var info in EnumerateConstantInfo(_loadResult))
+            {
+                var constantName = $"const{info.Index:D}";
+                _constantsTable[info.Index] = constantName;
+                yield return EncodeConstant(info.Value, constantName, info.Type);
+            }
+        }
+    }
+
+    private static IEnumerable<RawConstantInfo> EnumerateConstantInfo(MarkupLoadResult loadResult)
+    {
+        var constantsTable = loadResult.ConstantsTable;
+        bool hasPersistList = constantsTable.PersistList is not null;
+
+        if (hasPersistList)
+        {
+            var persistedList = loadResult.ConstantsTable.PersistList;
+
+            for (int c = 0; c < persistedList.Length; c++)
+            {
+                var persistedConstant = persistedList[c];
+                var typeSchema = persistedConstant.Type;
+
+               yield return new(c, typeSchema, persistedConstant.Data);
+            }
+        }
+        else
+        {
+            // UIB doesn't persist constants, so we have to use an alternate, slower method
+            for (int c = 0; ; c++)
+            {
+                object constantValue;
+                try
+                {
+                    constantValue = constantsTable.Get(c);
+                }
+                catch
+                {
+                    break;
+                }
+
+                var runtimeType = constantValue.GetType();
+                var typeSchema = loadResult.ImportTables.TypeImports.FirstOrDefault(t => t.RuntimeType == runtimeType)
+                    ?? loadResult.ImportTables.TypeImports.FirstOrDefault(t => t.RuntimeType == runtimeType.BaseType)
+                    ?? throw new Exception($"Failed to find type schema for '{runtimeType.Name}' in {loadResult.Uri}");
+
+                yield return new(c, typeSchema, constantValue);
+            }
+        }
+    }
+
     private ConstantDirective EncodeConstant(object constantValue, string constantName, TypeSchema typeSchema)
     {
         var qualifiedTypeName = GetQualifiedName(typeSchema);
@@ -393,4 +448,6 @@ public class Disassembler
 
         throw new NotSupportedException($"Unable to encode constant value '{constantValue}' of type '{qualifiedTypeName}'");
     }
+
+    private record RawConstantInfo(int Index, TypeSchema Type, object Value);
 }

libs/UIX.Asm/Extensions.cs

@@ -15,7 +15,78 @@ public static class Extensions
         result.Load(LoadPass.Done);
     }
 
+    public static string Unescape(this string input)
+    {
+        // https://stackoverflow.com/a/6736653/6232957
+
+        if (input.Length <= 1) return input;
+
+        // The input string can only get shorter,
+        // so init the buffer so we won't have to reallocate later
+        char[] buffer = new char[input.Length];
+        int outIdx = 0;
+        for (int i = 0; i < input.Length; i++)
+        {
+            char c = input[i];
+            if (c == '\\')
+            {
+                if (i < input.Length - 1)
+                {
+                    // https://learn.microsoft.com/en-us/dotnet/csharp/language-reference/language-specification/lexical-structure#unicode-character-escape-sequences
+                    var escapedChar = input[i + 1];
+                    var unescapedChar = escapedChar switch
+                    {
+                        '0' => '\0',
+                        'b' => '\b',
+                        'n' => '\n',
+                        'r' => '\r',
+                        't' => '\t',
+                        _ => escapedChar
+                    };
+
+                    buffer[outIdx++] = unescapedChar;
+                    i++;
+                    continue;
+                }
+            }
+
+            buffer[outIdx++] = c;
+        }
+
+        return new string(buffer, 0, outIdx);
+    }
+
+    public static string Escape(this string input)
+    {
+        StringBuilder sb = new();
+
+        var inputSpan = input.AsSpan();
+        for (int i = 0; i < input.Length; i++)
+        {
+            var ch = inputSpan[i];
+            var escapedCh = ch switch
+            {
+                '\0' => @"\0",
+                '\b' => @"\b",
+                '\n' => @"\n",
+                '\r' => @"\r",
+                '\t' => @"\t",
+                '\'' => @"\'",
+                '"' => "\"",
+                _ => null
+            };
+
+            if (escapedCh is not null)
+                sb.Append(escapedCh);
+            else
+                sb.Append(ch);
+        }
+
+        return sb.ToString();
+    }
+
 #if NETSTANDARD
+
     internal static StringBuilder AppendJoin<T>(this StringBuilder sb, string? separator, IEnumerable<T> values)
     {
         _ = values ?? throw new ArgumentNullException(nameof(values));
@@ -46,5 +117,6 @@ public static class Extensions
         }
         return sb;
     }
+
 #endif
 }

libs/UIX.Asm/Lexer.Directives.cs

@@ -27,6 +27,22 @@ partial class Lexer
         IDirective directive;
         switch (directiveIdResult.Value.ToUpperInvariant())
         {
+            case "DATATABLE":
+                if (StatementEnd(input).WasSuccessful)
+                    return Result.Failure<IDirective>(input, "Invalid data table import", ["Expected a URI"]);
+
+                var dataTableUriResult = Uri.Token()(input);
+                input = dataTableUriResult.Remainder;
+                if (!dataTableUriResult.WasSuccessful)
+                    return Result.Failure<IDirective>(input, "Invalid data table import", ["Expected a valid data table URI"]);
+
+                directive = new SharedDataTableDirective(dataTableUriResult.Value)
+                {
+                    Line = line,
+                    Column = column,
+                };
+                break;
+
             case "IMPORT":
                 return ParseImportAsDirective(input);
 

libs/UIX.Asm/Models/Directives.cs

@@ -1,7 +1,4 @@
-using System;
-using System.Linq;
-
-namespace Microsoft.Iris.Asm.Models;
+namespace Microsoft.Iris.Asm.Models;
 
 public record SectionDirective : Directive
 {
@@ -15,6 +12,18 @@ public record SectionDirective : Directive
     public override string ToString() => $"{base.ToString()} {Name}";
 }
 
+public record SharedDataTableDirective : Directive
+{
+    public SharedDataTableDirective(string dataTableUri) : base("dataTable")
+    {
+        Uri = dataTableUri;
+    }
+
+    public string Uri { get; init; }
+
+    public override string ToString() => $"{base.ToString()} {Uri}";
+}
+
 public record ExportDirective : Directive
 {
     public ExportDirective(string labelPrefix, uint listenerCount, string baseTypeName) : base("export")

libs/UIX.Asm/Models/Interfaces.cs

@@ -26,13 +26,12 @@ public abstract record Directive(string Identifier) : AsmItem, IDirective
     public override string ToString() => $".{Identifier}";
 }
 
-public interface IImportDirective : IDirective;
-public abstract record ImportDirective : Directive, IImportDirective
+public interface IImportDirective : IDirective
+{
+    string Type { get; init; }
+}
+public abstract record ImportDirective(string Type) : Directive($"import-{Type}"), IImportDirective
 {
-    public ImportDirective(string Type) : base($"import-{Type}")
-    {
-    }
-
     public override string ToString() => base.ToString();
 }
 

libs/UIX.Asm/Models/SyntaxModels.cs

@@ -29,6 +29,7 @@ public record Program
         Body = body.Cached();
 
         Directives = body.OfType<IDirective>().Cached();
+        DataTableDirective = Directives.OfType<SharedDataTableDirective>().SingleOrDefault();
         Imports = Directives.OfType<IImportDirective>().Cached();
         Exports = Directives.OfType<ExportDirective>().Cached();
 
@@ -38,6 +39,7 @@ public record Program
     public IEnumerable<IBodyItem> Body { get; }
 
     public IEnumerable<IDirective> Directives { get; }
+    public SharedDataTableDirective DataTableDirective { get; }
     public IEnumerable<IImportDirective> Imports { get; }
     public IEnumerable<ExportDirective> Exports { get; }
 
@@ -49,15 +51,46 @@ public record Program
         const string indent = "    ";
         StringBuilder sb = new();
 
-        sb.AppendJoin(lineEnding, Exports.Select(i => i.ToString()));
-        sb.Append(lineEnding);
-        sb.Append(lineEnding);
-        sb.AppendJoin(lineEnding, Imports.Select(i => i.ToString()));
-        sb.Append(lineEnding);
-        sb.Append(lineEnding);
-        sb.AppendJoin(lineEnding, Directives.Where(d => d is not (IImportDirective or ExportDirective)).Select(i => i.ToString()));
-        sb.Append(lineEnding);
-        sb.Append(lineEnding);
+        if (DataTableDirective is not null)
+        {
+            sb.Append(DataTableDirective);
+            sb.Append(lineEnding);
+            sb.Append(lineEnding);
+        }
+
+        if (Exports.Any())
+        {
+            sb.AppendJoin(lineEnding, Exports.Select(i => i.ToString()));
+            sb.Append(lineEnding);
+            sb.Append(lineEnding);
+        }
+
+        if (Imports.Any())
+        {
+            var sortedImports = Imports
+                .OrderBy(import => import.Type switch
+                {
+                    "ns" => 0,
+                    "type" => 1,
+                    "ctor" => 2,
+                    "mthd" => 3,
+                    "mbrs" => 4,
+                    _ => int.MaxValue
+                })
+                .Select(i => i.ToString());
+
+            sb.AppendJoin(lineEnding, sortedImports);
+            sb.Append(lineEnding);
+            sb.Append(lineEnding);
+        }
+
+        var otherDirectives = Directives.Where(d => d is not (IImportDirective or ExportDirective or SharedDataTableDirective));
+        if (otherDirectives.Any())
+        {
+            sb.AppendJoin(lineEnding, otherDirectives.Select(i => i.ToString()));
+            sb.Append(lineEnding);
+            sb.Append(lineEnding);
+        }
 
         foreach (var bodyItem in Code)
         {

libs/UIX.Test/Assembly.cs

@@ -83,6 +83,52 @@ Alt_locl:
         Assert.Equal(24, ast.Code.Count());
     }
 
+    [Fact]
+    public void ParseWithDataTable()
+    {
+        const string code =
+"""
+.export Default 0 UI
+.export Alt 0 UI
+
+.dataTable res://ZuneShellResources.dll!_DataTable.uib
+
+.section object
+
+Default_cont:
+    COBJ 5
+    PSHC @const0
+    PINI 2
+    PSHC @const1
+    PINI 3
+    PINI 1
+    RETV
+Default_locl:
+    COBJ 2
+    PDAD 0, @const2
+    RETV
+Alt_cont:
+    COBJ 5
+    PSHC @const3
+    PINI 2
+    PSHC @const4
+    PINI 4
+    PSHC @const5
+    PINI 3
+    PINI 1
+    RETV
+Alt_locl:
+    RETV
+""";
+
+        var ast = Lexer.Program.Parse(code);
+        _output.WriteLine(ast.ToString());
+
+        Assert.NotNull(ast);
+        Assert.Equal(4, ast.Directives.Count());
+        Assert.Equal(24, ast.Code.Count());
+    }
+
     [Theory]
     [InlineData("text")]
     public async Task ReassembleFromUIB(string fileNameWithoutExtension)