Support enums in rust-transient-0

4 files changed, 83 insertions, 3 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rust/RustGenerator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rust/RustGenerator.java
index db572fb..0a0290d 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rust/RustGenerator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rust/RustGenerator.java
@@ -46,7 +46,10 @@ public class RustGenerator {
             }
 
             case EnumType enumType -> {
-                throw new UnsupportedOperationException("TODO");
+                yield generateTypePath(
+                    enumType.namespaceName(),
+                    enumType.name()
+                );
             }
 
             case SimpleNativeType simpleNativeType -> {
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rusttransient0/RustTransient0Generator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rusttransient0/RustTransient0Generator.java
index 9dd4108..6295f17 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rusttransient0/RustTransient0Generator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rusttransient0/RustTransient0Generator.java
@@ -16,8 +16,11 @@ import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
+import static org.zwobble.hobgoblin.compiler.util.ListCollectors.toArrayList;
+
 public class RustTransient0Generator implements Generator {
     public static final String NAME = "rust-transient-0";
 
@@ -295,7 +298,66 @@ public class RustTransient0Generator implements Generator {
     ) {
         return switch (statement) {
             case TypedEnumDefinitionNode enumDefinition -> {
-                yield Stream.of();
+                var rustType = this.rustGenerator.generateRustTypeExpression(enumDefinition.type());
+
+                var encodeFunction = generateEncodeFunction(
+                    enumDefinition.type(),
+                    generateEncode(new RustPrefixExpression(
+                        RustPrefixOperator.BORROW,
+                        new RustMatchExpression(
+                            RustPath.of(VALUE_NAME),
+                            IntStream.range(0, enumDefinition.variants().size())
+                                .mapToObj(variantIndex -> {
+                                    var variant = enumDefinition.variants().get(variantIndex);
+                                    var variantPath = this.rustGenerator.generateRustTypeExpression(enumDefinition.type())
+                                        .addSegment(RustPathSegment.of(this.rustGenerator.generateVariantName(variant.name())));
+
+                                    return new RustMatchArm(
+                                        new RustPathPattern(variantPath),
+                                        new RustIntegerLiteral(variantIndex, Optional.empty())
+                                    );
+                                })
+                                .toList()
+                        )
+                    ), NativeTypes.INT_32)
+                );
+
+                var decodeMatchArms = IntStream.range(0, enumDefinition.variants().size())
+                    .mapToObj(variantIndex -> {
+                        var variant = enumDefinition.variants().get(variantIndex);
+                        var variantPath = this.rustGenerator.generateRustTypeExpression(enumDefinition.type())
+                            .addSegment(RustPathSegment.of(this.rustGenerator.generateVariantName(variant.name())));
+
+                        return new RustMatchArm(
+                            new RustLiteralPattern(new RustIntegerLiteral(variantIndex, Optional.empty())),
+                            variantPath
+                        );
+                    })
+                    .collect(toArrayList());
+
+                decodeMatchArms.add(new RustMatchArm(
+                    new RustWildcardPattern(),
+                    generateTodo()
+                ));
+
+                var decodeFunction = generateDecodeFunction(
+                    enumDefinition.type(),
+                    new RustBlockExpression(
+                        List.of(),
+                        Optional.of(new RustMatchExpression(
+                            new RustPrefixExpression(
+                                RustPrefixOperator.BORROW,
+                                generateDecode(NativeTypes.INT_32)
+                            ),
+                            decodeMatchArms
+                        ))
+                    )
+                );
+
+                yield Stream.of(
+                    encodeFunction,
+                    decodeFunction
+                );
             }
 
             case TypedNativeTypeDefinitionNode nativeTypeDefinition -> {
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustPath.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustPath.java
index 8390293..5647001 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustPath.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustPath.java
@@ -115,7 +115,7 @@ public record RustPath(java.util.List<org.zwobble.hobgoblin.compiler.output.lang
         return new RustPath(segments);
     }
 
-    public RustExpression addSegment(RustPathSegment segment) {
+    public RustPath addSegment(RustPathSegment segment) {
         var segments = new ArrayList<>(this.segments);
         segments.add(segment);
         return new RustPath(segments);
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/util/ListCollectors.java b/src/main/java/org/zwobble/hobgoblin/compiler/util/ListCollectors.java
new file mode 100644
index 0000000..cf38abd
--- /dev/null
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/util/ListCollectors.java
@@ -0,0 +1,15 @@
+package org.zwobble.hobgoblin.compiler.util;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collector;
+
+public class ListCollectors {
+    public static <T> Collector<T, ?, List<T>> toArrayList() {
+        return Collector.of(
+            ArrayList::new,
+            List::add,
+            (left, right) -> { left.addAll(right); return left; }
+        );
+    }
+}