Generate builders for generic structs

2 files changed, 132 insertions, 39 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
index 6d1e4a8..1034a98 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/generators/javatypes/JavaTypesGenerator.java
@@ -160,36 +160,73 @@ public class JavaTypesGenerator implements Generator {
             body.add(generateStructSingletonConstructor(javaRawTypeRef));
         }
 
-        // TODO: handle generic types
-        if (structDefinition.typeOrConstructor() instanceof TypeOrConstructor.Type<SimpleStructType>) {
-            var builderJavaTypeRef = generateBuilderTypeRef(structDefinition.typeOrThrow());
+        var builderJavaTypeRef = switch (structDefinition.typeOrConstructor()) {
+            case TypeOrConstructor.Constructor<SimpleStructType> constructor ->
+                generateBuilderTypeRef(constructor.value().genericType()).withArgs(
+                    javaTypeParams.stream()
+                        .map(javaTypeParam -> JavaTypeArg.invariant(JavaTypeRef.local(javaTypeParam.name())))
+                        .toList()
+                );
 
-            var arbitraryValueExpression = new JavaNewExpression(
-                builderJavaTypeRef,
-                structDefinition.fields().orElse(List.of()).stream()
-                    .map(field -> arbitraryValue(field.type(), context))
-                    .toList()
-            );
+            case TypeOrConstructor.Type<SimpleStructType> type ->
+                generateBuilderTypeRef(type.value());
+        };
 
-            body.add(new JavaMethodDeclaration(
-                ARBITRARY_METHOD_NAME,
-                JavaVisibility.PUBLIC,
-                JavaMemberKind.STATIC,
-                List.of(),
-                builderJavaTypeRef,
-                List.of(),
-                List.of(),
-                new JavaBlock(List.of(
-                    new JavaReturn(arbitraryValueExpression)
-                ))
-            ));
+        var structJavaTypeRef = switch (structDefinition.typeOrConstructor()) {
+            case TypeOrConstructor.Constructor<SimpleStructType> constructor ->
+                generateTypeRef(constructor.value().genericType(), context).withArgs(
+                    javaTypeParams.stream()
+                        .map(javaTypeParam -> JavaTypeArg.invariant(JavaTypeRef.local(javaTypeParam.name())))
+                        .toList()
+                );
 
-            body.add(generateStructDefinitionBuilder(
-                structDefinition,
-                components,
-                context
-            ));
-        }
+            case TypeOrConstructor.Type<SimpleStructType> type ->
+                generateTypeRef(type.value(), context);
+        };
+
+        var arbitraryValueExpression = new JavaNewExpression(
+            builderJavaTypeRef,
+            structDefinition.fields().orElse(List.of()).stream()
+                .map(field -> arbitraryValue(field.type(), context))
+                .toList()
+        );
+
+        // TODO: don't require arbitrary values for option, list, etc.
+        var arbitraryValueParams = switch (structDefinition.typeOrConstructor()) {
+            case TypeOrConstructor.Constructor<SimpleStructType> constructor ->
+                constructor.value().params().stream()
+                    .map(typeParam -> new JavaParam(
+                        JavaTypeRef.supplier(this.javaGenerator.generateTypeRef(typeParam)),
+                        arbitraryValueSupplier(typeParam)
+                    ))
+                    .toList();
+
+            case TypeOrConstructor.Type<SimpleStructType> _ ->
+                List.<JavaParam>of();
+        };
+
+
+        body.add(new JavaMethodDeclaration(
+            ARBITRARY_METHOD_NAME,
+            JavaVisibility.PUBLIC,
+            JavaMemberKind.STATIC,
+            javaTypeParams,
+            builderJavaTypeRef,
+            arbitraryValueParams,
+            List.of(),
+            new JavaBlock(List.of(
+                new JavaReturn(arbitraryValueExpression)
+            ))
+        ));
+
+        body.add(generateStructDefinitionBuilder(
+            structDefinition,
+            structJavaTypeRef,
+            builderJavaTypeRef,
+            javaTypeParams,
+            components,
+            context
+        ));
 
         if (structDefinition.fields().isPresent()) {
             return new JavaRecordDeclaration(
@@ -234,16 +271,22 @@ public class JavaTypesGenerator implements Generator {
 
     private JavaRecordDeclaration generateStructDefinitionBuilder(
         TypedStructDefinitionNode structDefinition,
+        JavaTypeRef structJavaTypeRef,
+        JavaTypeRef builderJavaTypeRef,
+        List<JavaTypeParameter> javaTypeParams,
         List<JavaRecordComponent> components,
         Context context
     ) {
-        var structType = structDefinition.typeOrThrow();
-
-        var builderJavaTypeRef = generateBuilderTypeRef(structType);
+        // TODO: handle generics
+        var implementsTypes = switch (structDefinition.typeOrConstructor()) {
+            case TypeOrConstructor.Constructor<SimpleStructType> _ ->
+                List.<JavaTypeRef>of();
 
-        var implementsTypes = context.variantOf(structType).stream()
-            .map(sumType -> this.generateBuilderTypeRef(sumType))
-            .toList();
+            case TypeOrConstructor.Type<SimpleStructType> constructor ->
+                context.variantOf(constructor.value()).stream()
+                    .map(sumType -> this.generateBuilderTypeRef(sumType))
+                    .toList();
+        };
 
         var body = new ArrayList<JavaClassBodyDeclaration>();
         body.add(new JavaMethodDeclaration(
@@ -251,17 +294,17 @@ public class JavaTypesGenerator implements Generator {
             JavaVisibility.PUBLIC,
             JavaMemberKind.INSTANCE,
             List.of(),
-            generateTypeRef(structType, context),
+            structJavaTypeRef,
             List.of(),
             List.of(),
             new JavaBlock(List.of(
                 new JavaReturn(
                     structDefinition.fields().isPresent()
                     ? new JavaNewExpression(
-                        generateTypeRef(structType, context),
+                        structJavaTypeRef,
                         components.stream().<JavaExpression>map(component -> new JavaRef(component.name())).toList()
                     )
-                    : new JavaStaticFieldAccess(this.generateTypeRef(structType, context), JavaGenerator.INSTANCE_FIELD_NAME)
+                    : new JavaStaticFieldAccess(structJavaTypeRef, JavaGenerator.INSTANCE_FIELD_NAME)
                 )
             ))
         ));
@@ -272,7 +315,7 @@ public class JavaTypesGenerator implements Generator {
 
         return new JavaRecordDeclaration(
             builderJavaTypeRef,
-            List.of(),
+            javaTypeParams,
             components,
             implementsTypes,
             body,
@@ -594,11 +637,31 @@ public class JavaTypesGenerator implements Generator {
     }
 
     private JavaExpression arbitraryValue(Type type, Context context) {
+        // TODO: unify logic for constructed types.
         return switch (this.javaGenerator.collapseType(type)) {
-            case ConstructedType constructedType -> {
+            case ConstructedNativeType constructedType -> {
                 yield arbitraryValue(constructedType.constructor().genericType(), context);
             }
 
+            case ConstructedStructType constructedType -> {
+                yield new JavaMethodCall(
+                    new JavaStaticMethodCall(
+                        generateTypeRef(constructedType.constructor().genericType(), context),
+                        ARBITRARY_METHOD_NAME,
+                        constructedType.args().stream()
+                            .<JavaExpression>map(typeArg -> new JavaLambdaExpression(
+                                List.of(),
+                                new JavaBlock(List.of(
+                                    new JavaReturn(arbitraryValue(typeArg, context))
+                                ))
+                            ))
+                            .toList()
+                    ),
+                    BUILD_METHOD_NAME,
+                    List.of()
+                );
+            }
+
             case EnumType enumType -> {
                 yield new JavaStaticFieldAccess(
                     generateTypeRef(enumType, context),
@@ -638,11 +701,20 @@ public class JavaTypesGenerator implements Generator {
             }
 
             case TypeParam typeParam -> {
-                throw new UnsupportedOperationException("TODO");
+                yield new JavaMethodCall(
+                    new JavaRef(arbitraryValueSupplier(typeParam)),
+                    JavaIdentifier.of("get"),
+                    List.of()
+                );
             }
         };
     }
 
+    private JavaIdentifier arbitraryValueSupplier(TypeParam typeParam) {
+        var typeName = this.javaGenerator.generateTypeName(typeParam.name()).value();
+        return JavaIdentifier.of("arbitraryValue$" + typeName);
+    }
+
     private JavaIdentifier generateEnumConstantName(String name) {
         return this.javaGenerator.generateEnumConstantName(name);
     }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaTypeRef.java b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaTypeRef.java
index 734ee95..00a6c1d 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaTypeRef.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/output/lang/java/ast/JavaTypeRef.java
@@ -51,11 +51,24 @@ public record JavaTypeRef(
         );
     }
 
+    public JavaTypeRef withArgs(List<JavaTypeArg> args) {
+        if (!this.args.isEmpty()) {
+            throw new UnsupportedOperationException("type ref already has args");
+        }
+
+        if (arrayDimensions != 0) {
+            throw new UnsupportedOperationException("type ref is array");
+        }
+
+        return new JavaTypeRef(packageName, typeNames, args, 0);
+    }
+
     private static class JavaPackages {
         private static final JavaPackageName JAVA_IO = JavaPackageName.of("java", "io");
         private static final JavaPackageName JAVA_LANG = JavaPackageName.of("java", "lang");
         private static final JavaPackageName JAVA_NIO_CHARSET = JavaPackageName.of("java", "nio", "charset");
         private static final JavaPackageName JAVA_UTIL = JavaPackageName.of("java", "util");
+        private static final JavaPackageName JAVA_UTIL_FUNCTION = JavaPackageName.of("java", "util", "function");
         private static final JavaPackageName JAVA_UTIL_STREAM = JavaPackageName.of("java", "util", "stream");
     }
 
@@ -136,5 +149,13 @@ public record JavaTypeRef(
 
     public static final JavaTypeRef STREAM = topLevel(JavaPackages.JAVA_UTIL_STREAM, JavaIdentifier.of("Stream"));
 
+    public static JavaTypeRef supplier(JavaTypeRef returnType) {
+        return topLevelGeneric(
+            JavaPackages.JAVA_UTIL_FUNCTION,
+            JavaIdentifier.of("Supplier"),
+            List.of(JavaTypeArg.invariant(returnType))
+        );
+    }
+
     public static final JavaTypeRef SYSTEM = topLevel(JavaPackages.JAVA_LANG, JavaIdentifier.of("System"));
 }