support union generation

Author: StunxFS

rivetc/src/codegen/__init__.py

@@ -3024,12 +3024,20 @@ def gen_types(self):
                 # TODO: in the self-hosted compiler calculate the enum value here
                 # not in register nor resolver.
                 if ts.info.is_tagged:
+                    mangled_name = cg_utils.mangle_symbol(ts)
+                    fields = []
+                    for v in ts.info.variants:
+                        if v.has_typ:
+                            typ_sym = v.typ.symbol()
+                            fields.append(ir.Field(f"v{typ_sym.id}", self.ir_type(v.typ)))
+                    union_name = mangled_name + "5Union"
+                    self.out_rir.unions.append(ir.Union(union_name, fields))
                     self.out_rir.structs.append(
                         ir.Struct(
-                            False, cg_utils.mangle_symbol(ts), [
+                            False, mangled_name, [
                                 ir.Field("_rc_", ir.UINT_T),
                                 ir.Field("_idx_", ir.UINT_T),
-                                ir.Field("obj", ir.VOID_PTR_T)
+                                ir.Field("obj", ir.Type(union_name))
                             ]
                         )
                     )

rivetc/src/codegen/c.py

@@ -29,12 +29,15 @@ class CGen:
     def __init__(self, comp):
         self.comp = comp
         self.typedefs = utils.Builder()
+        self.unions = utils.Builder()
         self.structs = utils.Builder()
         self.protos = utils.Builder()
         self.globals = utils.Builder()
         self.out = utils.Builder()
 
     def gen(self, out_rir):
+        self.comp.vlog("cgen: generating unions...")
+        self.gen_unions(out_rir.unions)
         self.comp.vlog("cgen: generating structs...")
         self.gen_structs(out_rir.structs)
         self.comp.vlog("cgen: generating externs...")
@@ -51,6 +54,7 @@ def gen(self, out_rir):
             if self.comp.prefs.build_mode != prefs.BuildMode.Release:
                 out.write(c_headers.RIVET_BREAKPOINT)
             out.write(str(self.typedefs).strip() + "\n\n")
+            out.write(str(self.unions).strip() + "\n\n")
             out.write(str(self.structs).strip() + "\n\n")
             out.write(str(self.protos).strip() + "\n\n")
             out.write(str(self.globals).strip() + "\n\n")
@@ -96,6 +100,19 @@ def write(self, txt):
     def writeln(self, txt = ""):
         self.out.writeln(txt)
 
+    def gen_unions(self, unions):
+        for u in unions:
+            self.typedefs.writeln(f"typedef union {u.name} {u.name};")
+            self.unions.writeln(f"union {u.name} {{")
+            for i, f in enumerate(u.fields):
+                f_name = c_escape(f.name)
+                self.unions.write("  ")
+                self.unions.write(self.gen_type(f.typ, f_name))
+                if not isinstance(f.typ, (ir.Array, ir.Function)):
+                    self.unions.write(f" {f_name}")
+                self.unions.writeln(";")
+            self.unions.writeln("};\n")
+
     def gen_structs(self, structs):
         for s in structs:
             self.typedefs.writeln(f"typedef struct {s.name} {s.name};")
@@ -107,10 +124,7 @@ def gen_structs(self, structs):
                     self.structs.write(self.gen_type(f.typ, f_name))
                     if not isinstance(f.typ, (ir.Array, ir.Function)):
                         self.structs.write(f" {f_name}")
-                    if i < len(s.fields) - 1:
-                        self.structs.writeln(";")
-                    else:
-                        self.structs.writeln(";")
+                    self.structs.writeln(";")
                 self.structs.writeln("};")
             self.structs.writeln()
 

rivetc/src/codegen/ir.py

@@ -128,6 +128,7 @@ def __eq__(self, other):
 class RIRFile:
     def __init__(self, mod_name):
         self.mod_name = mod_name
+        self.unions = []
         self.structs = []
         self.externs = []
         self.globals = []
@@ -146,6 +147,11 @@ def __repr__(self):
             "// and is subject to change without notice. Knock yourself out."
         )
         sb.writeln()
+        for i, u in enumerate(self.unions):
+            sb.writeln(str(u))
+            if i < len(self.unions) - 1:
+                sb.writeln()
+        sb.writeln()
         for i, s in enumerate(self.structs):
             sb.writeln(str(s))
             if i < len(self.structs) - 1:
@@ -189,6 +195,23 @@ def __str__(self):
         sb.write("}")
         return str(sb)
 
+class Union:
+    def __init__(self, name, fields):
+        self.name = name
+        self.fields = fields
+
+    def __str__(self):
+        sb = utils.Builder()
+        sb.writeln(f'union {self.name} {{')
+        for i, f in enumerate(self.fields):
+            sb.write(f'  {f.name}: {f.typ}')
+            if i < len(self.fields) - 1:
+                sb.writeln(",")
+            else:
+                sb.writeln()
+        sb.write("}")
+        return str(sb)
+
 class Struct:
     def __init__(self, is_opaque, name, fields):
         self.is_opaque = is_opaque
@@ -198,9 +221,9 @@ def __init__(self, is_opaque, name, fields):
     def __str__(self):
         sb = utils.Builder()
         if self.is_opaque:
-            sb.write(f'type {self.name} opaque')
+            sb.write(f'struct {self.name} opaque')
         else:
-            sb.writeln(f'type {self.name} {{')
+            sb.writeln(f'struct {self.name} {{')
             for i, f in enumerate(self.fields):
                 sb.write(f'  {f.name}: {f.typ}')
                 if i < len(self.fields) - 1: