1 files changed, 71 insertions, 90 deletions
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
index e13c3dd..beef2ca 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
@@ -8,6 +8,7 @@ import org.zwobble.hobgoblin.compiler.output.lang.java.ast.*;
 import java.io.IOException;
 import java.io.StringWriter;
 import java.util.List;
+import java.util.Optional;
 
 import static org.zwobble.precisely.AssertThat.assertThat;
 import static org.zwobble.precisely.Matchers.equalTo;
@@ -18,7 +19,7 @@ public class JavaWriterTests {
         var java = new JavaClassDeclaration(
             "Rectangle",
             List.of(),
-            JavaCustomArea.EMPTY,
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Rectangle"), "body"),
             DocComment.EMPTY
         );
 
@@ -26,8 +27,8 @@ public class JavaWriterTests {
 
         assertThat(string, equalTo("""
             public class Rectangle {
-                // Custom area start: body
-                // Custom area end: body
+                // Custom area start: Rectangle body
+                // Custom area end: Rectangle body
             }"""
         ));
     }
@@ -52,7 +53,7 @@ public class JavaWriterTests {
                     JavaBlock.EMPTY
                 )
             ),
-            JavaCustomArea.EMPTY,
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Rectangle"), "body"),
             DocComment.EMPTY
         );
 
@@ -66,28 +67,8 @@ public class JavaWriterTests {
                 public void b() {
                 }
 
-                // Custom area start: body
-                // Custom area end: body
-            }"""
-        ));
-    }
-
-    @Test
-    public void classDeclarationWithCustomArea() throws IOException {
-        var java = new JavaClassDeclaration(
-            "Point",
-            List.of(),
-            new JavaCustomArea(List.of("// <custom area>")),
-            DocComment.EMPTY
-        );
-
-        var string = write(writer -> writer.writeTypeDeclaration(java));
-
-        assertThat(string, equalTo("""
-            public class Point {
-                // Custom area start: body
-                // <custom area>
-                // Custom area end: body
+                // Custom area start: Rectangle body
+                // Custom area end: Rectangle body
             }"""
         ));
     }
@@ -97,7 +78,7 @@ public class JavaWriterTests {
         var java = new JavaClassDeclaration(
             "Point",
             List.of(),
-            JavaCustomArea.EMPTY,
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Point"), "body"),
             new DocComment("A 2D point.")
         );
 
@@ -106,8 +87,8 @@ public class JavaWriterTests {
         assertThat(string, equalTo("""
             /// A 2D point.
             public class Point {
-                // Custom area start: body
-                // Custom area end: body
+                // Custom area start: Point body
+                // Custom area end: Point body
             }"""
         ));
     }
@@ -117,7 +98,7 @@ public class JavaWriterTests {
         var java = new JavaInterfaceDeclaration(
             "Shape",
             List.of(),
-            JavaCustomArea.EMPTY,
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Shape"), "body"),
             DocComment.EMPTY
         );
 
@@ -125,8 +106,8 @@ public class JavaWriterTests {
 
         assertThat(string, equalTo("""
             public sealed interface Shape {
-                // Custom area start: body
-                // Custom area end: body
+                // Custom area start: Shape body
+                // Custom area end: Shape body
             }"""
         ));
     }
@@ -139,7 +120,7 @@ public class JavaWriterTests {
                 JavaTypeRef.topLevel(List.of("abc", "def"), "One"),
                 JavaTypeRef.topLevel(List.of("abc", "def"), "Two")
             ),
-            JavaCustomArea.EMPTY,
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Shape"), "body"),
             DocComment.EMPTY
         );
 
@@ -147,28 +128,8 @@ public class JavaWriterTests {
 
         assertThat(string, equalTo("""
             public sealed interface Shape permits abc.def.One, abc.def.Two {
-                // Custom area start: body
-                // Custom area end: body
-            }"""
-        ));
-    }
-
-    @Test
-    public void interfaceDeclarationWithCustomArea() throws IOException {
-        var java = new JavaInterfaceDeclaration(
-            "Shape",
-            List.of(),
-            new JavaCustomArea(List.of("// <custom area>")),
-            DocComment.EMPTY
-        );
-
-        var string = write(writer -> writer.writeTypeDeclaration(java));
-
-        assertThat(string, equalTo("""
-            public sealed interface Shape {
-                // Custom area start: body
-                // <custom area>
-                // Custom area end: body
+                // Custom area start: Shape body
+                // Custom area end: Shape body
             }"""
         ));
     }
@@ -178,7 +139,7 @@ public class JavaWriterTests {
         var java = new JavaInterfaceDeclaration(
             "Shape",
             List.of(),
-            JavaCustomArea.EMPTY,
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Shape"), "body"),
             new DocComment("A 2D shape.")
         );
 
@@ -187,8 +148,8 @@ public class JavaWriterTests {
         assertThat(string, equalTo("""
             /// A 2D shape.
             public sealed interface Shape {
-                // Custom area start: body
-                // Custom area end: body
+                // Custom area start: Shape body
+                // Custom area end: Shape body
             }"""
         ));
     }
@@ -200,28 +161,7 @@ public class JavaWriterTests {
             List.of(),
             List.of(),
             List.of(),
-            JavaCustomArea.EMPTY,
-            DocComment.EMPTY
-        );
-
-        var string = write(writer -> writer.writeTypeDeclaration(java));
-
-        assertThat(string, equalTo("""
-            public record Point() {
-                // Custom area start: body
-                // Custom area end: body
-            }"""
-        ));
-    }
-
-    @Test
-    public void recordDeclarationWithCustomArea() throws IOException {
-        var java = new JavaRecordDeclaration(
-            "Point",
-            List.of(),
-            List.of(),
-            List.of(),
-            new JavaCustomArea(List.of("// <custom area>")),
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Point"), "body"),
             DocComment.EMPTY
         );
 
@@ -229,9 +169,8 @@ public class JavaWriterTests {
 
         assertThat(string, equalTo("""
             public record Point() {
-                // Custom area start: body
-                // <custom area>
-                // Custom area end: body
+                // Custom area start: Point body
+                // Custom area end: Point body
             }"""
         ));
     }
@@ -243,7 +182,7 @@ public class JavaWriterTests {
             List.of(),
             List.of(),
             List.of(),
-            JavaCustomArea.EMPTY,
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Point"), "body"),
             new DocComment("A 2D point.")
         );
 
@@ -252,8 +191,8 @@ public class JavaWriterTests {
         assertThat(string, equalTo("""
             /// A 2D point.
             public record Point() {
-                // Custom area start: body
-                // Custom area end: body
+                // Custom area start: Point body
+                // Custom area end: Point body
             }"""
         ));
     }
@@ -280,7 +219,7 @@ public class JavaWriterTests {
                     JavaBlock.EMPTY
                 )
             ),
-            JavaCustomArea.EMPTY,
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Rectangle"), "body"),
             DocComment.EMPTY
         );
 
@@ -294,8 +233,8 @@ public class JavaWriterTests {
                 public void b() {
                 }
 
-                // Custom area start: body
-                // Custom area end: body
+                // Custom area start: Rectangle body
+                // Custom area end: Rectangle body
             }"""
         ));
     }
@@ -463,9 +402,51 @@ public class JavaWriterTests {
         assertThat(string, equalTo("abc.def.One.Two"));
     }
 
+    @Test
+    public void customAreaContentsArePreserved() throws IOException {
+        var java = new JavaRecordDeclaration(
+            "Point",
+            List.of(),
+            List.of(),
+            List.of(),
+            new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Point"), "body"),
+            DocComment.EMPTY
+        );
+
+        var string = write(
+            writer -> writer.writeTypeDeclaration(java),
+            """
+                public class Point {
+                // Custom area start: Point body
+                // <custom area>
+                // Custom area end: Point body
+            }
+            """
+        );
+
+        assertThat(string, equalTo("""
+            public record Point() {
+                // Custom area start: Point body
+                // <custom area>
+                // Custom area end: Point body
+            }"""
+        ));
+    }
+
     private String write(Write write) throws IOException {
+        return write(write, Optional.empty());
+    }
+
+    private String write(Write write, String originalContents) throws IOException {
+        return write(write, Optional.of(originalContents));
+    }
+
+    private String write(Write write, Optional<String> originalContents) throws IOException {
         var stringWriter = new StringWriter();
-        var javaWriter = new JavaWriter(new CodeWriter(stringWriter));
+        var javaWriter = new JavaWriter(new CodeWriter(
+            stringWriter,
+            CodeWriter.CustomAreas.read(originalContents.orElse(""), "// Custom area start: ", "// Custom area end: ")
+        ));
         write.write(javaWriter);
         return stringWriter.toString();
     }