union definitions

compiler/Generator.cs

@@ -44,9 +44,14 @@ public class Generator
         {
             foreach (var (name, info) in module.GetTypes())
             {
-                if (info is Module.TypeInfoStruct s)
+                switch (info)
                 {
-                    writer.WriteLine($"struct {NameMangler.Mangle(module.Name, name, NubTypeStruct.Get(module.Name, name))};");
+                    case Module.TypeInfoStruct s:
+                        writer.WriteLine($"struct {NameMangler.Mangle(module.Name, name, NubTypeStruct.Get(module.Name, name))};");
+                        break;
+                    case Module.TypeInfoEnum e:
+                        writer.WriteLine($"struct {NameMangler.Mangle(module.Name, name, NubTypeStruct.Get(module.Name, name))};");
+                        break;
                 }
             }
         }
@@ -55,24 +60,59 @@ public class Generator
         {
             foreach (var (name, info) in module.GetTypes())
             {
-                if (info is Module.TypeInfoStruct s)
+                switch (info)
                 {
-                    writer.WriteLine();
+                    case Module.TypeInfoStruct s:
-                    writer.Write("struct ");
-
-                    if (s.Packed)
-                        writer.Write("__attribute__((__packed__)) ");
-
-                    writer.WriteLine(NameMangler.Mangle(module.Name, name, NubTypeStruct.Get(module.Name, name)));
-                    writer.WriteLine("{");
-                    using (writer.Indent())
                     {
-                        foreach (var field in s.Fields)
+                        writer.WriteLine();
+                        writer.Write("struct ");
+
+                        if (s.Packed)
+                            writer.Write("__attribute__((__packed__)) ");
+
+                        writer.WriteLine(NameMangler.Mangle(module.Name, name, NubTypeStruct.Get(module.Name, name)));
+                        writer.WriteLine("{");
+                        using (writer.Indent())
                         {
-                            writer.WriteLine($"{CType(field.Type, field.Name)};");
+                            foreach (var field in s.Fields)
+                            {
+                                writer.WriteLine($"{CType(field.Type, field.Name)};");
+                            }
                         }
+                        writer.WriteLine("};");
+                        break;
+                    }
+                    case Module.TypeInfoEnum e:
+                    {
+                        writer.WriteLine();
+                        writer.Write($"struct {NameMangler.Mangle(module.Name, name, NubTypeStruct.Get(module.Name, name))}");
+                        writer.WriteLine("{");
+                        using (writer.Indent())
+                        {
+                            writer.WriteLine("uint32_t tag;");
+                            writer.WriteLine("union");
+                            writer.WriteLine("{");
+                            using (writer.Indent())
+                            {
+                                foreach (var variant in e.Variants)
+                                {
+                                    writer.WriteLine($"struct {variant.Name}");
+                                    writer.WriteLine("{");
+                                    using (writer.Indent())
+                                    {
+                                        foreach (var field in variant.Fields)
+                                        {
+                                            writer.WriteLine($"{CType(field.Type, field.Name)};");
+                                        }
+                                    }
+                                    writer.WriteLine("};");
+                                }
+                            }
+                            writer.WriteLine("};");
+                        }
+                        writer.WriteLine("};");
+                        break;
                     }
-                    writer.WriteLine("};");
                 }
             }
         }

compiler/ModuleGraph.cs

@@ -85,11 +85,21 @@ public class ModuleGraph
                         switch (type)
                         {
                             case Manifest.Module.TypeInfoStruct s:
+                            {
                                 var info = new Module.TypeInfoStruct(Module.DefinitionKind.External, s.Packed);
                                 var fields = s.Fields.Select(x => new Module.TypeInfoStruct.Field(x.Name, x.Type)).ToList();
                                 info.SetFields(fields);
                                 module.AddType(name, info);
                                 break;
+                            }
+                            case Manifest.Module.TypeInfoEnum s:
+                            {
+                                var info = new Module.TypeInfoEnum(Module.DefinitionKind.External);
+                                var variants = s.Variants.Select(v => new Module.TypeInfoEnum.Variant(v.Name, v.Fields.Select(x => new Module.TypeInfoEnum.Variant.Field(x.Name, x.Type)).ToList())).ToList();
+                                info.SetVariants(variants);
+                                module.AddType(name, info);
+                                break;
+                            }
                             default:
                                 throw new ArgumentOutOfRangeException(nameof(type));
                         }
@@ -111,6 +121,12 @@ public class ModuleGraph
                     var kind = structDef.Exported ? Module.DefinitionKind.Exported : Module.DefinitionKind.Internal;
                     module.AddType(structDef.Name.Ident, new Module.TypeInfoStruct(kind, structDef.Packed));
                 }
+
+                foreach (var enumDef in ast.Definitions.OfType<NodeDefinitionEnum>())
+                {
+                    var kind = enumDef.Exported ? Module.DefinitionKind.Exported : Module.DefinitionKind.Internal;
+                    module.AddType(enumDef.Name.Ident, new Module.TypeInfoEnum(kind));
+                }
             }
 
             foreach (var ast in asts)
@@ -128,6 +144,18 @@ public class ModuleGraph
                         structType.SetFields(fields);
                     }
                 }
+
+                foreach (var enumDef in ast.Definitions.OfType<NodeDefinitionEnum>())
+                {
+                    if (!module.TryResolveType(enumDef.Name.Ident, true, out var typeInfo))
+                        throw new UnreachableException($"{nameof(typeInfo)} should always be registered");
+
+                    if (typeInfo is Module.TypeInfoEnum enumType)
+                    {
+                        var variants = enumDef.Variants.Select(v => new Module.TypeInfoEnum.Variant(v.Name.Ident, v.Fields.Select(f => new Module.TypeInfoEnum.Variant.Field(f.Name.Ident, ResolveType(f.Type, module.Name))).ToList())).ToList();
+                        enumType.SetVariants(variants);
+                    }
+                }
             }
 
             foreach (var ast in asts)
@@ -188,6 +216,7 @@ public class ModuleGraph
                 return customType switch
                 {
                     Module.TypeInfoStruct => NubTypeStruct.Get(type.Module.Ident, type.Name.Ident),
+                    Module.TypeInfoEnum => NubTypeEnum.Get(type.Module.Ident, type.Name.Ident),
                     _ => throw new ArgumentOutOfRangeException(nameof(customType))
                 };
             }
@@ -301,4 +330,28 @@ public class Module(string name)
             public NubType Type { get; } = type;
         }
     }
+
+    public class TypeInfoEnum(DefinitionKind kind) : TypeInfo(kind)
+    {
+        private IReadOnlyList<Variant>? variants;
+
+        public IReadOnlyList<Variant> Variants => variants ?? throw new InvalidOperationException("Fields has not been set yet");
+
+        public void SetVariants(IReadOnlyList<Variant> variants)
+        {
+            this.variants = variants;
+        }
+
+        public class Variant(string name, List<Variant.Field> fields)
+        {
+            public string Name { get; } = name;
+            public List<Field> Fields { get; } = fields;
+
+            public class Field(string name, NubType type)
+            {
+                public string Name { get; } = name;
+                public NubType Type { get; } = type;
+            }
+        }
+    }
 }

compiler/NubLib.cs

@@ -99,6 +99,7 @@ public record Manifest(Dictionary<string, Manifest.Module> Modules)
             return typeInfo switch
             {
                 Compiler.Module.TypeInfoStruct s => new Module.TypeInfoStruct(s.Packed, s.Fields.Select(x => new Module.TypeInfoStruct.Field(x.Name, x.Type)).ToList()),
+                Compiler.Module.TypeInfoEnum e => new Module.TypeInfoEnum(e.Variants.Select(v => new Module.TypeInfoEnum.Variant(v.Name, v.Fields.Select(x => new Module.TypeInfoEnum.Variant.Field(x.Name, x.Type)).ToList())).ToList()),
                 _ => throw new ArgumentOutOfRangeException(nameof(typeInfo))
             };
         }
@@ -110,11 +111,20 @@ public record Manifest(Dictionary<string, Manifest.Module> Modules)
         public record IdentifierInfo(NubType Type, string MangledName);
 
         [JsonDerivedType(typeof(TypeInfoStruct), "struct")]
+        [JsonDerivedType(typeof(TypeInfoEnum), "enum")]
         public abstract record TypeInfo;
 
         public record TypeInfoStruct(bool Packed, IReadOnlyList<TypeInfoStruct.Field> Fields) : TypeInfo
         {
             public record Field(string Name, NubType Type);
         }
+
+        public record TypeInfoEnum(IReadOnlyList<TypeInfoEnum.Variant> Variants) : TypeInfo
+        {
+            public record Variant(string Name, List<Variant.Field> Fields)
+            {
+                public record Field(string Name, NubType Type);
+            }
+        }
     }
 }

compiler/NubType.cs

@@ -111,6 +111,30 @@ public class NubTypeStruct : NubType
     public override string ToString() => $"struct {Module}::{Name}";
 }
 
+public class NubTypeEnum : NubType
+{
+    private static readonly Dictionary<(string Module, string Name), NubTypeEnum> Cache = new();
+
+    public static NubTypeEnum Get(string module, string name)
+    {
+        if (!Cache.TryGetValue((module, name), out var enumType))
+            Cache[(module, name)] = enumType = new NubTypeEnum(module, name);
+
+        return enumType;
+    }
+
+    private NubTypeEnum(string module, string name)
+    {
+        Module = module;
+        Name = name;
+    }
+
+    public string Module { get; }
+    public string Name { get; }
+
+    public override string ToString() => $"enum {Module}::{Name}";
+}
+
 public class NubTypePointer : NubType
 {
     private static readonly Dictionary<NubType, NubTypePointer> Cache = new();
@@ -172,8 +196,6 @@ public class TypeEncoder
     {
     }
 
-    private Dictionary<NubTypeStruct, Definition> structDefinitions = new();
-
     private string EncodeRoot(NubType type)
     {
         var sb = new StringBuilder();
@@ -223,6 +245,14 @@ public class TypeEncoder
                 sb.Append(')');
                 break;
 
+            case NubTypeEnum st:
+                sb.Append("E(");
+                sb.Append(st.Module);
+                sb.Append(':');
+                sb.Append(st.Name);
+                sb.Append(')');
+                break;
+
             case NubTypeFunc fn:
                 sb.Append("F(");
                 for (int i = 0; i < fn.Parameters.Count; i++)
@@ -273,6 +303,7 @@ public class TypeDecoder
             'P' => DecodePointer(),
             'F' => DecodeFunc(),
             'T' => DecodeStruct(),
+            'E' => DecodeEnum(),
             _ => throw new Exception($"'{start}' is not a valid start to a type")
         };
     }
@@ -333,6 +364,25 @@ public class TypeDecoder
         return NubTypeStruct.Get(module, name);
     }
 
+    private NubTypeEnum DecodeEnum()
+    {
+        var sb = new StringBuilder();
+
+        Expect('(');
+        while (!TryExpect(':'))
+            sb.Append(Consume());
+
+        var module = sb.ToString();
+        sb.Clear();
+
+        while (!TryExpect(')'))
+            sb.Append(Consume());
+
+        var name = sb.ToString();
+
+        return NubTypeEnum.Get(module, name);
+    }
+
     private bool TryPeek(out char c)
     {
         if (index >= encoded.Length)

compiler/Parser.cs

@@ -129,6 +129,45 @@ public class Parser
             return new NodeDefinitionStruct(TokensFrom(startIndex), exported, packed, name, fields);
         }
 
+        if (TryExpectKeyword(Keyword.Enum))
+        {
+            var exported = modifiers.Remove(Keyword.Export);
+
+            foreach (var modifier in modifiers)
+                // todo(nub31): Add to diagnostics instead of throwing
+                throw new CompileException(Diagnostic.Error("Invalid modifier for struct").At(fileName, modifier.Value).Build());
+
+            var name = ExpectIdent();
+            var variants = new List<NodeDefinitionEnum.Variant>();
+
+            ExpectSymbol(Symbol.OpenCurly);
+            while (!TryExpectSymbol(Symbol.CloseCurly))
+            {
+                var variantsStartIndex = index;
+                var variantName = ExpectIdent();
+
+                var variantFields = new List<NodeDefinitionEnum.Variant.Field>();
+
+                if (TryExpectSymbol(Symbol.OpenCurly))
+                {
+                    while (!TryExpectSymbol(Symbol.CloseCurly))
+                    {
+                        var fieldStartIndex = index;
+
+                        var fieldName = ExpectIdent();
+                        ExpectSymbol(Symbol.Colon);
+                        var fieldType = ParseType();
+
+                        variantFields.Add(new NodeDefinitionEnum.Variant.Field(TokensFrom(fieldStartIndex), fieldName, fieldType));
+                    }
+                }
+
+                variants.Add(new NodeDefinitionEnum.Variant(TokensFrom(variantsStartIndex), variantName, variantFields));
+            }
+
+            return new NodeDefinitionEnum(TokensFrom(startIndex), exported, name, variants);
+        }
+
         if (TryExpectKeyword(Keyword.Let))
         {
             var exported = modifiers.Remove(Keyword.Export);
@@ -645,6 +684,25 @@ public class NodeDefinitionStruct(List<Token> tokens, bool exported, bool packed
     }
 }
 
+public class NodeDefinitionEnum(List<Token> tokens, bool exported, TokenIdent name, List<NodeDefinitionEnum.Variant> variants) : NodeDefinition(tokens)
+{
+    public bool Exported { get; } = exported;
+    public TokenIdent Name { get; } = name;
+    public List<Variant> Variants { get; } = variants;
+
+    public class Variant(List<Token> tokens, TokenIdent name, List<Variant.Field> fields) : Node(tokens)
+    {
+        public TokenIdent Name { get; } = name;
+        public List<Field> Fields { get; } = fields;
+
+        public class Field(List<Token> tokens, TokenIdent name, NodeType type) : Node(tokens)
+        {
+            public TokenIdent Name { get; } = name;
+            public NodeType Type { get; } = type;
+        }
+    }
+}
+
 public class NodeDefinitionGlobalVariable(List<Token> tokens, bool exported, TokenIdent name, NodeType type) : NodeDefinition(tokens)
 {
     public bool Exported { get; } = exported;

compiler/Tokenizer.cs

@@ -388,6 +388,7 @@ public class Tokenizer
                         "func" => new TokenKeyword(line, startColumn, column - startColumn, Keyword.Func),
                         "struct" => new TokenKeyword(line, startColumn, column - startColumn, Keyword.Struct),
                         "packed" => new TokenKeyword(line, startColumn, column - startColumn, Keyword.Packed),
+                        "enum" => new TokenKeyword(line, startColumn, column - startColumn, Keyword.Enum),
                         "let" => new TokenKeyword(line, startColumn, column - startColumn, Keyword.Let),
                         "if" => new TokenKeyword(line, startColumn, column - startColumn, Keyword.If),
                         "else" => new TokenKeyword(line, startColumn, column - startColumn, Keyword.Else),
@@ -534,6 +535,7 @@ public enum Keyword
     Func,
     Struct,
     Packed,
+    Enum,
     Let,
     If,
     Else,
@@ -598,6 +600,7 @@ public static class TokenExtensions
             Keyword.Func => "func",
             Keyword.Struct => "struct",
             Keyword.Packed => "packed",
+            Keyword.Enum => "enum",
             Keyword.Let => "let",
             Keyword.If => "if",
             Keyword.Else => "else",

examples/math/math.nub

@@ -5,6 +5,14 @@ struct vec3 { x: i32 y: i32 z: i32 }
 struct color { r: i32 g: i32 b: i32 a: i32 }
 struct example { a: math::vec2 b: math::vec3 c: math::color }
 
+export enum message {
+    quit
+    move {
+        x: i32
+        y: i32
+    }
+}
+
 export func add(a: i32 b: i32): i32
 {
     return math::add_internal(a b)