Add attributes to Rust enums

5 files changed, 59 insertions, 11 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rusttypes/RustTypesGenerator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rusttypes/RustTypesGenerator.java
index 28926a3..19f3365 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rusttypes/RustTypesGenerator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/rusttypes/RustTypesGenerator.java
@@ -102,6 +102,7 @@ public class RustTypesGenerator implements Generator {
             .toList();
 
         var rustEnum = new RustEnum(
+            List.of(),
             rustEnumName,
             rustVariants,
             enumDefinition.docComment()
@@ -154,7 +155,12 @@ public class RustTypesGenerator implements Generator {
             })
             .toList();
 
-        var rustEnum = new RustEnum(rustEnumName, rustVariants, sumDefinition.docComment());
+        var rustEnum = new RustEnum(
+            List.of(),
+            rustEnumName,
+            rustVariants,
+            sumDefinition.docComment()
+        );
         items.add(rustEnum);
 
         for (var variant : sumDefinition.variants()) {
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriter.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriter.java
index bdf834e..e0b0810 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriter.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriter.java
@@ -146,6 +146,7 @@ public class RustWriter implements AutoCloseable {
 
     private void writeEnum(RustEnum enum_) throws IOException {
         this.writeDocComment(enum_.docComment());
+        this.writerOuterAttributes(enum_.attributes());
         this.writer.write("pub enum ");
         this.writeIdentifier(enum_.name());
         this.writer.write(" {");
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustEnum.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustEnum.java
index be3ae6a..eee5e08 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustEnum.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustEnum.java
@@ -5,42 +5,58 @@ package org.zwobble.hobgoblin.compiler.output.lang.rust.ast;
 // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum imports
 // Custom area end: org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum imports
 
-public record RustEnum(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier name, java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant> variants, org.zwobble.hobgoblin.compiler.ast.DocComment docComment) implements org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustItem {
+public record RustEnum(java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustOuterAttribute> attributes, org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier name, java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant> variants, org.zwobble.hobgoblin.compiler.ast.DocComment docComment) implements org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustItem {
     public static org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder arbitrary() {
-        return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier.arbitrary().build(), java.util.List.of(), org.zwobble.hobgoblin.compiler.ast.HobgoblinNativeAst.arbitraryDocComment());
+        return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(java.util.List.of(), org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier.arbitrary().build(), java.util.List.of(), org.zwobble.hobgoblin.compiler.ast.HobgoblinNativeAst.arbitraryDocComment());
     }
 
-    public record Builder(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier name, java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant> variants, org.zwobble.hobgoblin.compiler.ast.DocComment docComment) implements org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustItem.Builder {
+    public record Builder(java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustOuterAttribute> attributes, org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier name, java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant> variants, org.zwobble.hobgoblin.compiler.ast.DocComment docComment) implements org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustItem.Builder {
         public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum build() {
-            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum(name, variants, docComment);
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum(attributes, name, variants, docComment);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder withAttributes(java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustOuterAttribute> attributes) {
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name, variants, docComment);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder addAttribute(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustOuterAttribute attribute) {
+            var attributes = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustOuterAttribute>(this.attributes);
+            attributes.add(attribute);
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name, variants, docComment);
+        }
+
+        public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder addAttribute(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustOuterAttribute.Builder attribute) {
+            var attributes = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustOuterAttribute>(this.attributes);
+            attributes.add(attribute.build());
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name, variants, docComment);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder withName(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier name) {
-            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(name, variants, docComment);
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name, variants, docComment);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder withName(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier.Builder name) {
-            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(name.build(), variants, docComment);
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name.build(), variants, docComment);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder withVariants(java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant> variants) {
-            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(name, variants, docComment);
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name, variants, docComment);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder addVariant(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant variant) {
             var variants = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant>(this.variants);
             variants.add(variant);
-            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(name, variants, docComment);
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name, variants, docComment);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder addVariant(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant.Builder variant) {
             var variants = new java.util.ArrayList<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumVariant>(this.variants);
             variants.add(variant.build());
-            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(name, variants, docComment);
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name, variants, docComment);
         }
 
         public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder withDocComment(org.zwobble.hobgoblin.compiler.ast.DocComment docComment) {
-            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(name, variants, docComment);
+            return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder(attributes, name, variants, docComment);
         }
 
         // Custom area start: org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.Builder body
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriterTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriterTests.java
index 62b361a..776d51d 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriterTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/rust/RustWriterTests.java
@@ -176,6 +176,25 @@ public class RustWriterTests {
     }
 
     @Test
+    public void enumWithAttributes() throws IOException {
+        var rust = RustEnum.arbitrary()
+            .withName(new RustIdentifier("Shape"))
+            .withAttributes(List.of(
+                new RustOuterAttribute(RustTypePath.of("a", "b", "C"), Optional.empty()),
+                new RustOuterAttribute(RustTypePath.of("derive"), Optional.of("(Debug)"))
+            ))
+            .build();
+
+        var string = write(writer -> writer.writeItem(rust));
+
+        assertThat(string, equalTo("""
+            #[a::b::C]
+            #[derive(Debug)]
+            pub enum Shape {
+            }"""));
+    }
+
+    @Test
     public void enumWithFieldlessVariants() throws IOException {
         var rust = RustEnum.arbitrary()
             .withName(new RustIdentifier("Shape"))
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustEnumMatcher.java b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustEnumMatcher.java
index 547bdf7..9d1a8cd 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustEnumMatcher.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/rust/ast/RustEnumMatcher.java
@@ -28,6 +28,12 @@ public class RustEnumMatcher implements org.zwobble.precisely.Matcher<java.lang.
         return org.zwobble.precisely.Matchers.instanceOf(org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum.class, this.submatchers);
     }
 
+    public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumMatcher withAttributes(org.zwobble.precisely.Matcher<? super java.util.List<org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustOuterAttribute>> attributes) {
+        var submatchers = new java.util.ArrayList<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum>>(this.submatchers);
+        submatchers.add(org.zwobble.precisely.Matchers.has("attributes", org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum::attributes, attributes));
+        return new org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumMatcher(submatchers);
+    }
+
     public org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnumMatcher withName(org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustIdentifier> name) {
         var submatchers = new java.util.ArrayList<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum>>(this.submatchers);
         submatchers.add(org.zwobble.precisely.Matchers.has("name", org.zwobble.hobgoblin.compiler.output.lang.rust.ast.RustEnum::name, name));