Assign tags to structs

2 files changed, 54 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index a1baec0f..4eb221c1 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -32,8 +32,26 @@ public record PatioCodeGenContext(
     FunctionContext functionContext
 ) {
     public record GlobalContext(
-        TypeInfos typeInfos
+        TypeInfos typeInfos,
+        PMap<StructType, Integer> tags
     ) {
+        public static final GlobalContext EMPTY = forTypeInfos(TypeInfos.EMPTY);
+
+        public static GlobalContext forTypeInfos(TypeInfos typeInfos) {
+            var tags = P.<StructType, Integer>map();
+            var structTypes = typeInfos.structTypes();
+            var nextTag = 0;
+            for (var structType : structTypes) {
+                tags = tags.plus(structType, nextTag);
+                nextTag++;
+            }
+
+            return new GlobalContext(
+                typeInfos,
+                tags
+            );
+        }
+
         public NamespaceTypeInfo namespaceTypeInfo(NamespaceName namespaceName) {
             return typeInfos.namespaceTypeInfo(namespaceName)
                 // TODO: throw better error
@@ -46,12 +64,14 @@ public record PatioCodeGenContext(
                 .orElseThrow(() -> new RuntimeException("No type info for: " + type));
         }
 
-        public static final GlobalContext EMPTY = forTypeInfos(TypeInfos.EMPTY);
-
-        public static GlobalContext forTypeInfos(TypeInfos typeInfos) {
-            return new GlobalContext(
-                typeInfos
-            );
+        public int tag(StructType type) {
+            var tag = this.tags.get(type);
+            if (tag == null) {
+                // TODO: throw better error
+                throw new RuntimeException("No tag for: " + type);
+            } else {
+                return tag;
+            }
         }
     }
 
@@ -259,6 +279,10 @@ public record PatioCodeGenContext(
         return this.globalContext.structTypeInfo(type);
     }
 
+    public int tag(StructType type) {
+        return globalContext.tag(type);
+    }
+
     public PatioCodeGenResult<String> allocateCompilationUnitElementName(
         PatioCodeGenCompilationUnitElementSpace space,
         String name
diff --git a/src/test/java/org/shedlang/compiler/patio/PatioCodeGenDefineStructTests.java b/src/test/java/org/shedlang/compiler/patio/PatioCodeGenDefineStructTests.java
index 6464dfe0..932f626a 100644
--- a/src/test/java/org/shedlang/compiler/patio/PatioCodeGenDefineStructTests.java
+++ b/src/test/java/org/shedlang/compiler/patio/PatioCodeGenDefineStructTests.java
@@ -16,6 +16,8 @@ import java.util.List;
 import java.util.Map;
 
 import static org.shedlang.compiler.patio.PatioTesting.serialize;
+import static org.zwobble.precisely.AssertThat.assertThat;
+import static org.zwobble.precisely.Matchers.equalTo;
 
 @ExtendWith({SnapshotResolver.class})
 public class PatioCodeGenDefineStructTests {
@@ -40,4 +42,25 @@ public class PatioCodeGenDefineStructTests {
 
         snapshotter.assertSnapshot(serialize(result));
     }
+
+    @Test
+    public void structsAreAssignedSequentialTags(Snapshotter snapshotter) throws IOException {
+        // TODO: we can probably optimise this to ensure that (for instance)
+        // variants of the same sum type have sequential tags, but source order
+        // will suffice for now
+        var structType1 = Types.struct(NamespaceName.of("Example", "Project"), "User");
+        var structType2 = Types.struct(NamespaceName.of("Example", "Project"), "Id");
+        var instruction1 = Instructions.defineStruct(structType1, List.of());
+        var instruction2 = Instructions.defineStruct(structType1, List.of());
+        var typeInfos = TypeInfos.EMPTY
+            .addStructTypeInfo(structType1, StructTypeInfo.fromFields(P.orderedMap()))
+            .addStructTypeInfo(structType2, StructTypeInfo.fromFields(P.orderedMap()));
+        var context = PatioCodeGenContext.forTypeInfos(typeInfos);
+
+        context = PatioCodeGen.generateInstruction(instruction1, context);
+        context = PatioCodeGen.generateInstruction(instruction2, context);
+
+        assertThat(context.tag(structType1), equalTo(0));
+        assertThat(context.tag(structType2), equalTo(1));
+    }
 }