Oops, all IBodyItems!

libs/UIX.Asm/CachedEnumerable.cs

@@ -0,0 +1,108 @@
+using System;
+using System.Collections;
+using System.Collections.Generic;
+
+namespace Microsoft.Iris.Asm;
+
+// Yoinked from https://www.meziantou.net/caching-an-ienumerable-t-instance.htm
+
+internal static class CachedEnumerable
+{
+    public static CachedEnumerable<T> Create<T>(IEnumerable<T> enumerable) => new(enumerable);
+
+    /// <summary>
+    /// Wraps this <see cref="IEnumerable{T}"/> such that it only has to be enumerated once.
+    /// </summary>
+    public static CachedEnumerable<T> Cached<T>(this IEnumerable<T> enumerable) => new(enumerable);
+}
+
+internal sealed class CachedEnumerable<T> : IEnumerable<T>, IDisposable
+{
+    private readonly List<T> _cache = [];
+    private readonly IEnumerable<T> _enumerable;
+    private IEnumerator<T> _enumerator;
+    private bool _enumerated = false;
+
+    public CachedEnumerable(IEnumerable<T> enumerable)
+    {
+        _enumerable = enumerable ?? throw new ArgumentNullException(nameof(enumerable));
+    }
+
+    public IEnumerator<T> GetEnumerator()
+    {
+        var index = 0;
+        while (true)
+        {
+            if (TryGetItem(index, out var result))
+            {
+                yield return result;
+                index++;
+            }
+            else
+            {
+                // There are no more items
+                yield break;
+            }
+        }
+    }
+
+    private bool TryGetItem(int index, out T result)
+    {
+        // if the item is in the cache, use it
+        if (index < _cache.Count)
+        {
+            result = _cache[index];
+            return true;
+        }
+
+        lock (_cache)
+        {
+            if (_enumerator == null && !_enumerated)
+            {
+                _enumerator = _enumerable.GetEnumerator();
+            }
+
+            // Another thread may have get the item while we were acquiring the lock
+            if (index < _cache.Count)
+            {
+                result = _cache[index];
+                return true;
+            }
+
+            // If we have already enumerate the whole stream, there is nothing else to do
+            if (_enumerated)
+            {
+                result = default;
+                return false;
+            }
+
+            // Get the next item and store it to the cache
+            if (_enumerator.MoveNext())
+            {
+                result = _enumerator.Current;
+                _cache.Add(result);
+                return true;
+            }
+            else
+            {
+                // There are no more items, we can dispose the underlying enumerator
+                _enumerator.Dispose();
+                _enumerator = null;
+                _enumerated = true;
+                result = default;
+                return false;
+            }
+        }
+    }
+
+    public void Dispose()
+    {
+        if (_enumerator != null)
+        {
+            _enumerator.Dispose();
+            _enumerator = null;
+        }
+    }
+
+    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+}

libs/UIX.Asm/Disassembler.cs

@@ -83,7 +83,7 @@ public class Disassembler
         }
     }
 
-    public IEnumerable<IImport> GetImports()
+    public IEnumerable<IImportDirective> GetImports()
     {
         // Ues _importedUris to keep track of what has already been imported.
         // Skip self and default UIX namespace.
@@ -117,6 +117,11 @@ public class Disassembler
                     }
                 }
 
+                // Some imports, such as assembly imports, require additional parsing
+                // and might change the URI that actually gets imported.
+                if (_importedUris.ContainsKey(uri))
+                    continue;
+
                 namespacePrefix = namespacePrefix.Camelize();
                 _importedUris.Add(uri, namespacePrefix);
 
@@ -127,7 +132,7 @@ public class Disassembler
         };
     }
 
-    public IEnumerable<IBodyItem> GetBody()
+    public IEnumerable<IBodyItem> GetCode()
     {
         var reader = _loadResult.ObjectSection;
 
@@ -254,12 +259,13 @@ public class Disassembler
         _loadResult.Load(LoadPass.Full);
         _loadResult.Load(LoadPass.Done);
 
-        List<IDirective> directives = GetImports().Cast<IDirective>()
+        IEnumerable<IEnumerable<IBodyItem>> segments = [
-            .Concat(GetExports())
+            GetExports(),
-            .ToList();
+            GetImports(),
-        List<IBodyItem> body = new(GetBody());
+            GetCode(),
+        ];
 
-        Program asmProgram = new(directives, body);
+        Program asmProgram = new(segments.SelectMany(e => e));
         return asmProgram.ToString();
     }
 

libs/UIX.Asm/Lexer.Body.cs

@@ -1,7 +1,6 @@
 using Microsoft.Iris.Asm.Models;
 using Sprache;
 using System.Collections.Generic;
-using System.Linq;
 
 namespace Microsoft.Iris.Asm;
 

libs/UIX.Asm/Lexer.Directives.cs

@@ -30,12 +30,12 @@ partial class Lexer
 
             case "SECTION":
                 if (StatementEnd(input).WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid section directive", ["Expected a section name"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid section directive", ["Expected a section name"]);
 
                 var sectionNameResult = Parse.Letter.AtLeastOnce().Token().Text()(input);
                 input = sectionNameResult.Remainder;
                 if (!sectionNameResult.WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid section name", ["Expected a section name containing only letters"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid section name", ["Expected a section name containing only letters"]);
 
                 directive = new SectionDirective(sectionNameResult.Value)
                 {
@@ -46,25 +46,25 @@ partial class Lexer
 
             case "EXPORT":
                 if (StatementEnd(input).WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid export directive", ["Expected export information"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid export directive", ["Expected export information"]);
 
                 var labelPrefixResult = Identifier.Token()(input);
                 input = labelPrefixResult.Remainder;
                 if (!labelPrefixResult.WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid export directive", ["Expected prefix of labels to export"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid export directive", ["Expected prefix of labels to export"]);
 
                 var listenerCountResult = WholeNumber.Token()(input);
                 input = listenerCountResult.Remainder;
                 if (!listenerCountResult.WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid export directive", ["Expected listener count"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid export directive", ["Expected listener count"]);
 
                 if (!uint.TryParse(listenerCountResult.Value, out var listenerCount))
-                    return Result.Failure<IImport>(input, "Invalid export directive", ["Expected export listener count to be an unsigned integer"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid export directive", ["Expected export listener count to be an unsigned integer"]);
 
                 var baseTypeNameResult = AlphanumericText.Token()(input);
                 input = baseTypeNameResult.Remainder;
                 if (!baseTypeNameResult.WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid export directive", ["Expected base type name"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid export directive", ["Expected base type name"]);
 
                 var labelPrefix = labelPrefixResult.Value;
                 var baseTypeName = baseTypeNameResult.Value;
@@ -76,7 +76,7 @@ partial class Lexer
                 break;
 
             default:
-                return Result.Failure<IImport>(input, $"Unknown import type '{directiveIdResult.Value}'", ["Expected 'export', 'import', or 'section'"]);
+                return Result.Failure<IImportDirective>(input, $"Unknown import type '{directiveIdResult.Value}'", ["Expected 'export', 'import', or 'section'"]);
         }
 
         return Result.Success(directive, input);

libs/UIX.Asm/Lexer.Imports.cs

@@ -5,44 +5,44 @@ namespace Microsoft.Iris.Asm;
 
 partial class Lexer
 {
-    private static IResult<IImport> ParseImport(IInput input)
+    private static IResult<IImportDirective> ParseImport(IInput input)
     {
         input = ConsumeWhitespace(input);
 
         var importDirectiveResult = Parse.String(".import")(input);
         input = importDirectiveResult.Remainder;
         if (!importDirectiveResult.WasSuccessful)
-            return Result.Failure<IImport>(input, "Invalid import directive", ["Expected '.import'"]);
+            return Result.Failure<IImportDirective>(input, "Invalid import directive", ["Expected '.import'"]);
 
         return ParseImportAsDirective(input);
     }
 
-    private static IResult<IImport> ParseImportAsDirective(IInput input)
+    private static IResult<IImportDirective> ParseImportAsDirective(IInput input)
     {
         if (input.Current != '-' || input.AtEnd)
-            return Result.Failure<IImport>(input, "Invalid import type", ["An import type must be specified"]);
+            return Result.Failure<IImportDirective>(input, "Invalid import type", ["An import type must be specified"]);
         input = input.Advance();
 
         var importTypeResult = WordText(input);
         input = importTypeResult.Remainder;
         if (!importTypeResult.WasSuccessful)
-            return Result.Failure<IImport>(input, "Invalid import type", ["Expected 'ns'"]);
+            return Result.Failure<IImportDirective>(input, "Invalid import type", ["Expected 'ns'"]);
 
-        IImport import;
+        IImportDirective import;
         switch (importTypeResult.Value.ToUpperInvariant())
         {
             case "NS":
                 var uriResult = Uri.Token()(input);
                 input = uriResult.Remainder;
                 if (!uriResult.WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid URI", ["Expected a valid URI"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid URI", ["Expected a valid URI"]);
 
                 input = Parse.String("as").Token()(input).Remainder;
 
                 var nameResult = AlphanumericText(input);
                 input = nameResult.Remainder;
                 if (!nameResult.WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid namespace alias", ["Expected a valid namespace alias"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid namespace alias", ["Expected a valid namespace alias"]);
 
                 import = new NamespaceImport(uriResult.Value, nameResult.Value);
                 break;
@@ -51,7 +51,7 @@ partial class Lexer
                 var typePrefixResult = Identifier.Token()(input);
                 input = typePrefixResult.Remainder;
                 if (!typePrefixResult.WasSuccessful)
-                    return Result.Failure<IImport>(input, "Invalid type import", ["Expected a valid namespace prefix"]);
+                    return Result.Failure<IImportDirective>(input, "Invalid type import", ["Expected a valid namespace prefix"]);
 
                 var typeNamespaceDelimitterResult = Parse.Char(':')(input);
                 input = typeNamespaceDelimitterResult.Remainder;
@@ -62,7 +62,7 @@ partial class Lexer
                     var typeNameResult = Identifier(input);
                     input = typeNameResult.Remainder;
                     if (!typeNameResult.WasSuccessful)
-                        return Result.Failure<IImport>(input, "Invalid type import", ["Expected a valid type name"]);
+                        return Result.Failure<IImportDirective>(input, "Invalid type import", ["Expected a valid type name"]);
 
                     typePrefix = typePrefixResult.Value;
                     typeName = typeNameResult.Value;
@@ -77,7 +77,7 @@ partial class Lexer
                 break;
 
             default:
-                return Result.Failure<IImport>(input, $"Unknown import type '{importTypeResult.Value}'", ["Expected 'ns'"]);
+                return Result.Failure<IImportDirective>(input, $"Unknown import type '{importTypeResult.Value}'", ["Expected 'ns'"]);
         }
 
         return Result.Success(import, input);

libs/UIX.Asm/Lexer.cs

@@ -18,16 +18,15 @@ public static partial class Lexer
 
     public static readonly Parser<string> StatementEnd = Parse.Char(';').Return(";").Or(Parse.LineTerminator);
 
-    public static readonly Parser<IImport> Import = ParseImport;
+    public static readonly Parser<IImportDirective> Import = ParseImport;
 
     public static readonly Parser<IDirective> Directive = ParseDirective;
 
     public static readonly Parser<IBodyItem> BodyItem = ParseBodyItem;
 
     public static readonly Parser<Program> Program =
-        from directives in Directive.Many()
         from body in BodyItem.Many()
-        select new Program(directives, body);
+        select new Program(body);
 
     private static IInput ConsumeWhitespace(IInput input) => Parse.WhiteSpace.Many()(input).Remainder;
 }

libs/UIX.Asm/Models/Directives.cs

@@ -1,6 +1,6 @@
 namespace Microsoft.Iris.Asm.Models;
 
-public record SectionDirective : Directive, IBodyItem
+public record SectionDirective : Directive
 {
     public SectionDirective(string name) : base("section")
     {

libs/UIX.Asm/Models/Imports.cs

@@ -1,6 +1,6 @@
 namespace Microsoft.Iris.Asm.Models;
 
-public record NamespaceImport : Import
+public record NamespaceImport : ImportDirective
 {
     public NamespaceImport(string uri, string name) : base("ns")
     {
@@ -14,7 +14,7 @@ public record NamespaceImport : Import
     public override string ToString() => $"{base.ToString()} {Uri} as {Name}";
 }
 
-public record TypeImport : Import
+public record TypeImport : ImportDirective
 {
     public TypeImport(string namespacePrefix, string name) : base("type")
     {

libs/UIX.Asm/Models/Instructions.cs

@@ -7,7 +7,7 @@ using System.Linq;
 namespace Microsoft.Iris.Asm.Models;
 
 [DebuggerDisplay("{ToString()} " + DebuggerDisplay)]
-public record Instruction(string Mnemonic, IEnumerable<Operand> Operands) : BodyItem
+public record Instruction(string Mnemonic, IEnumerable<Operand> Operands) : CodeItem
 {
     public Instruction(OpCode opCode, OperationType? operationType, IEnumerable<Operand> Operands)
         : this(InstructionSet.GetMnemonic(opCode, operationType), Operands)

libs/UIX.Asm/Models/Interfaces.cs

@@ -14,7 +14,10 @@ public abstract record AsmItem : IAsmItem
     internal const string DebuggerDisplay = "({Line}, {Column})";
 }
 
-public interface IDirective : IAsmItem
+public interface IBodyItem : IAsmItem;
+public abstract record BodyItem : AsmItem, IBodyItem;
+
+public interface IDirective : IBodyItem
 {
     string Identifier { get; init; }
 }
@@ -23,15 +26,15 @@ public abstract record Directive(string Identifier) : AsmItem, IDirective
     public override string ToString() => $".{Identifier}";
 }
 
-public interface IImport : IDirective;
+public interface IImportDirective : IDirective;
-public abstract record Import : Directive, IImport
+public abstract record ImportDirective : Directive, IImportDirective
 {
-    public Import(string Type) : base($"import-{Type}")
+    public ImportDirective(string Type) : base($"import-{Type}")
     {
     }
 
     public override string ToString() => base.ToString();
 }
 
-public interface IBodyItem : IAsmItem;
+public interface ICodeItem : IBodyItem;
-public abstract record BodyItem : AsmItem, IBodyItem;
+public abstract record CodeItem : BodyItem, ICodeItem;

libs/UIX.Asm/Models/SyntaxModels.cs

@@ -6,7 +6,7 @@ using System.Text;
 namespace Microsoft.Iris.Asm.Models;
 
 [DebuggerDisplay("{Name} " + DebuggerDisplay)]
-public record Label(string Name) : BodyItem
+public record Label(string Name) : CodeItem
 {
     public override string ToString() => $"{Name}:";
 }
@@ -25,19 +25,41 @@ public record Operand(object Value, OperandDataType DataType, string Content = n
     public override string ToString() => Content ?? Value.ToString();
 }
 
-public record Program(IEnumerable<IDirective> Directives, IEnumerable<IBodyItem> Body)
+public record Program
 {
+    public Program(IEnumerable<IBodyItem> body)
+    {
+        Body = body.Cached();
+
+        Directives = body.OfType<IDirective>().Cached();
+        Imports = Directives.OfType<IImportDirective>().Cached();
+        Exports = Directives.OfType<ExportDirective>().Cached();
+
+        Code = body.OfType<ICodeItem>().Cached();
+    }
+
+    public IEnumerable<IBodyItem> Body { get; }
+
+    public IEnumerable<IDirective> Directives { get; }
+    public IEnumerable<IImportDirective> Imports { get; }
+    public IEnumerable<ExportDirective> Exports { get; }
+
+    public IEnumerable<ICodeItem> Code { get; }
+
     public override string ToString()
     {
         const string lineEnding = "\r\n";
         const string indent = "    ";
         StringBuilder sb = new();
 
-        sb.AppendJoin(lineEnding, Directives.Select(i => i.ToString()));
+        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);
 
-        foreach (var bodyItem in Body)
+        foreach (var bodyItem in Code)
         {
             if (bodyItem is Instruction)
                 sb.Append(indent);

libs/UIX.Test/Assembly.cs

@@ -43,7 +43,7 @@ main:
 
         Assert.NotNull(ast);
         Assert.Equal(3 + 2 + 2, ast.Directives.Count());
-        Assert.Equal(9, ast.Body.Count());
+        Assert.Equal(9, ast.Code.Count());
     }
 
     [Theory]
@@ -109,7 +109,10 @@ main:
         ErrorManager.OnErrors += (errors) =>
         {
             foreach (ErrorRecord error in errors)
-                output.WriteLine($"Error at (L{error.Line}, C{error.Column}): {error.Message}");
+            {
+                var errorTypeText = error.Warning ? "Warning" : "Error";
+                output.WriteLine($"{errorTypeText} at (L{error.Line}, C{error.Column}): {error.Message}");
+            }
         };
 
         var success = MarkupCompiler.Compile(compilerInputs, default);