diff --git a/sdk/java/dagger-java-annotation-processor/src/main/java/io/dagger/annotation/processor/DaggerModuleAnnotationProcessor.java b/sdk/java/dagger-java-annotation-processor/src/main/java/io/dagger/annotation/processor/DaggerModuleAnnotationProcessor.java
index 83dbfd8495..014ae43ee8 100644
--- a/sdk/java/dagger-java-annotation-processor/src/main/java/io/dagger/annotation/processor/DaggerModuleAnnotationProcessor.java
+++ b/sdk/java/dagger-java-annotation-processor/src/main/java/io/dagger/annotation/processor/DaggerModuleAnnotationProcessor.java
@@ -236,7 +236,7 @@ static JavaFile generate(ModuleInfo moduleInfo) {
                 "clazz.getMethod(\"setClient\", $T.class).invoke(obj, dag)", Client.class);
         var firstFn = true;
         for (var fnInfo : objectInfo.functions()) {
-          var fnReturnClass = classForName(fnInfo.returnType());
+          CodeBlock fnReturnType = typeName(fnInfo.returnType());
           if (firstFn) {
             firstFn = false;
             im.beginControlFlow("if (fnName.equals($S))", fnInfo.name());
@@ -246,32 +246,40 @@ static JavaFile generate(ModuleInfo moduleInfo) {
           var fnBlock =
               CodeBlock.builder().add("$T fn = clazz.getMethod($S", Method.class, fnInfo.qName());
           var invokeBlock =
-              CodeBlock.builder().add("$T res = ($T) fn.invoke(obj", fnReturnClass, fnReturnClass);
+              CodeBlock.builder()
+                  .add(fnReturnType)
+                  .add(" res = (")
+                  .add(fnReturnType)
+                  .add(") fn.invoke(obj");
           for (var parameterInfo : fnInfo.parameters()) {
-            Class<?> paramClazz = classForName(parameterInfo.type());
-            fnBlock.add(", $T.class", paramClazz);
+            CodeBlock paramType = typeName(parameterInfo.type());
+            fnBlock.add(", ").add(paramType).add(".class");
 
             invokeBlock.add(", $L", parameterInfo.name());
 
-            if (paramClazz.isPrimitive()) {
-              if (paramClazz.isAssignableFrom(boolean.class)) {
-                im.addStatement("$T $L = false", paramClazz, parameterInfo.name());
-              } else if (paramClazz.isAssignableFrom(int.class)) {
-                im.addStatement("$T $L = 0", paramClazz, parameterInfo.name());
-              } else {
-                im.addStatement("$T $L", paramClazz, parameterInfo.name());
-              }
-            } else {
-              im.addStatement("$T $L = null", paramClazz, parameterInfo.name());
+            String defaultValue = "null";
+            if (parameterInfo.type().kindName().equals(TypeKind.INT.name())) {
+              defaultValue = "0";
+            } else if (parameterInfo.type().kindName().equals(TypeKind.BOOLEAN.name())) {
+              defaultValue = "false";
             }
+            im.addStatement(
+                CodeBlock.builder()
+                    .add(paramType)
+                    .add(" $L = $L", parameterInfo.name(), defaultValue)
+                    .build());
             im.beginControlFlow("if (inputArgs.get($S) != null)", parameterInfo.name());
             im.addStatement(
-                "$L = ($T) $T.fromJSON(dag, inputArgs.get($S), $T.class)",
-                parameterInfo.name(),
-                paramClazz,
-                JsonConverter.class,
-                parameterInfo.name(),
-                classForName(parameterInfo.type()));
+                CodeBlock.builder()
+                    .add("$L = (", parameterInfo.name())
+                    .add(paramType)
+                    .add(
+                        ") $T.fromJSON(dag, inputArgs.get($S), ",
+                        JsonConverter.class,
+                        parameterInfo.name())
+                    .add(paramType)
+                    .add(".class)")
+                    .build());
             im.endControlFlow();
           }
           fnBlock.add(")");
@@ -391,14 +399,12 @@ static CodeBlock typeDef(TypeInfo ti) throws ClassNotFoundException {
           "dag.typeDef().withKind($T.$L)", TypeDefKind.class, TypeDefKind.BOOLEAN_KIND.name());
     } else if (name.startsWith("java.util.List<")) {
       name = name.substring("java.util.List<".length(), name.length() - 1);
-      return CodeBlock.of(
-          "dag.typeDef().withListOf($L)", typeDef(new TypeInfo(name, "")).toString());
+      return CodeBlock.of("dag.typeDef().withListOf($L)", typeDef(tiFromName(name)).toString());
     } else if (!ti.kindName().isEmpty() && TypeKind.valueOf(ti.kindName()) == TypeKind.ARRAY) {
       // in that case the type name is com.example.Type[]
       // so we remove the [] to get the underlying type
       name = name.substring(0, name.length() - 2);
-      return CodeBlock.of(
-          "dag.typeDef().withListOf($L)", typeDef(new TypeInfo(name, "")).toString());
+      return CodeBlock.of("dag.typeDef().withListOf($L)", typeDef(tiFromName(name)).toString());
     }
 
     try {
@@ -425,18 +431,48 @@ static CodeBlock typeDef(TypeInfo ti) throws ClassNotFoundException {
     }
   }
 
-  static Class<?> classForName(TypeInfo ti) throws ClassNotFoundException {
-    String name = ti.typeName();
+  static TypeInfo tiFromName(String name) {
     if (name.equals("int")) {
-      return int.class;
+      return new TypeInfo(name, TypeKind.INT.name());
     } else if (name.equals("boolean")) {
-      return boolean.class;
-    } else if (name.startsWith("java.util.List<")) {
-      return List.class;
-    } else if (TypeKind.valueOf(ti.kindName()) == TypeKind.ARRAY) {
-      return Class.forName("[L" + name.substring(0, name.length() - 2) + ";");
+      return new TypeInfo(name, TypeKind.BOOLEAN.name());
+    } else if (name.equals("void")) {
+      return new TypeInfo(name, TypeKind.VOID.name());
+    } else {
+      return new TypeInfo(name, "");
+    }
+  }
+
+  static CodeBlock typeName(TypeInfo ti) {
+    try {
+      TypeKind tk = TypeKind.valueOf(ti.kindName());
+      if (tk == TypeKind.INT) {
+        return CodeBlock.of("$T", int.class);
+      } else if (tk == TypeKind.BOOLEAN) {
+        return CodeBlock.of("$T", boolean.class);
+      } else if (tk == TypeKind.VOID) {
+        return CodeBlock.of("$T", void.class);
+      } else if (tk == TypeKind.ARRAY) {
+        return CodeBlock.builder()
+            .add(typeName(tiFromName(ti.typeName().substring(0, ti.typeName().length() - 2))))
+            .add("[]")
+            .build();
+      }
+    } catch (IllegalArgumentException ignored) {
+    }
+    String name = ti.typeName();
+    if (name.startsWith("java.util.List<")) {
+      return CodeBlock.of("$T", List.class);
+    }
+    try {
+      Class<?> clazz = Class.forName(name);
+      return CodeBlock.of("$T", clazz);
+    } catch (ClassNotFoundException e) {
+      return CodeBlock.of(
+          "$T",
+          ClassName.get(
+              name.substring(0, name.lastIndexOf(".")), name.substring(name.lastIndexOf(".") + 1)));
     }
-    return Class.forName(name);
   }
 
   private String trimDoc(String doc) {
diff --git a/sdk/java/dagger-java-annotation-processor/src/test/resources/io/dagger/gen/entrypoint/entrypoint.java b/sdk/java/dagger-java-annotation-processor/src/test/resources/io/dagger/gen/entrypoint/entrypoint.java
index b1c283ad94..408b3f277f 100644
--- a/sdk/java/dagger-java-annotation-processor/src/test/resources/io/dagger/gen/entrypoint/entrypoint.java
+++ b/sdk/java/dagger-java-annotation-processor/src/test/resources/io/dagger/gen/entrypoint/entrypoint.java
@@ -15,6 +15,7 @@
 import io.dagger.client.ModuleID;
 import io.dagger.client.TypeDef;
 import io.dagger.client.TypeDefKind;
+import io.dagger.java.module.DaggerJava;
 import jakarta.json.bind.JsonbBuilder;
 import java.lang.Class;
 import java.lang.Exception;
@@ -81,7 +82,20 @@ private ModuleID register() throws ExecutionException, DaggerQueryException,
                         dag.typeDef().withKind(TypeDefKind.STRING_KIND))
                         .withDescription("Returns lines that match a pattern in the files of the provided Directory")
                         .withArg("directoryArg", dag.typeDef().withObject("Directory").withOptional(false), new Function.WithArgArguments().withDescription("Directory to grep"))
-                        .withArg("pattern", dag.typeDef().withKind(TypeDefKind.STRING_KIND).withOptional(false), new Function.WithArgArguments().withDescription("Pattern to search for in the directory"))));
+                        .withArg("pattern", dag.typeDef().withKind(TypeDefKind.STRING_KIND).withOptional(false), new Function.WithArgArguments().withDescription("Pattern to search for in the directory")))
+                .withFunction(
+                    dag.function("itself",
+                        dag.typeDef().withObject("DaggerJava")))
+                .withFunction(
+                    dag.function("isZero",
+                        dag.typeDef().withKind(TypeDefKind.BOOLEAN_KIND))
+                        .withArg("value", dag.typeDef().withKind(TypeDefKind.INTEGER_KIND).withOptional(false)))
+                .withFunction(
+                    dag.function("doThings",
+                        dag.typeDef().withListOf(dag.typeDef().withKind(io.dagger.client.TypeDefKind.INTEGER_KIND)))
+                        .withArg("stringArray", dag.typeDef().withListOf(dag.typeDef().withKind(io.dagger.client.TypeDefKind.STRING_KIND)).withOptional(false))
+                        .withArg("ints", dag.typeDef().withListOf(dag.typeDef().withKind(io.dagger.client.TypeDefKind.INTEGER_KIND)).withOptional(false))
+                        .withArg("containers", dag.typeDef().withListOf(dag.typeDef().withObject("Container")).withOptional(false))));
     return module.id();
   }
 
@@ -112,9 +126,37 @@ private JSON invoke(JSON parentJson, String parentName, String fnName,
           Method fn = clazz.getMethod("grepDir", Directory.class, String.class);
           String res = (String) fn.invoke(obj, directoryArg, pattern);
           return JsonConverter.toJSON(res);
+        } else if (fnName.equals("itself")) {
+          Method fn = clazz.getMethod("itself");
+          DaggerJava res = (DaggerJava) fn.invoke(obj);
+          return JsonConverter.toJSON(res);
+        } else if (fnName.equals("isZero")) {
+          int value = 0;
+          if (inputArgs.get("value") != null) {
+            value = (int) JsonConverter.fromJSON(dag, inputArgs.get("value"), int.class);
+          }
+          Method fn = clazz.getMethod("isZero", int.class);
+          boolean res = (boolean) fn.invoke(obj, value);
+          return JsonConverter.toJSON(res);
+        } else if (fnName.equals("doThings")) {
+          String[] stringArray = null;
+          if (inputArgs.get("stringArray") != null) {
+            stringArray = (String[]) JsonConverter.fromJSON(dag, inputArgs.get("stringArray"), String[].class);
+          }
+          List ints = null;
+          if (inputArgs.get("ints") != null) {
+            ints = (List) JsonConverter.fromJSON(dag, inputArgs.get("ints"), List.class);
+          }
+          List containers = null;
+          if (inputArgs.get("containers") != null) {
+            containers = (List) JsonConverter.fromJSON(dag, inputArgs.get("containers"), List.class);
+          }
+          Method fn = clazz.getMethod("doThings", String[].class, List.class, List.class);
+          int[] res = (int[]) fn.invoke(obj, stringArray, ints, containers);
+          return JsonConverter.toJSON(res);
         }
       }
     }
     return null;
   }
-}
+}
\ No newline at end of file
diff --git a/sdk/java/dagger-java-annotation-processor/src/test/resources/io/dagger/java/module/DaggerJava.java b/sdk/java/dagger-java-annotation-processor/src/test/resources/io/dagger/java/module/DaggerJava.java
index 16c7a5486b..74efc714a4 100644
--- a/sdk/java/dagger-java-annotation-processor/src/test/resources/io/dagger/java/module/DaggerJava.java
+++ b/sdk/java/dagger-java-annotation-processor/src/test/resources/io/dagger/java/module/DaggerJava.java
@@ -49,4 +49,20 @@ public String grepDir(Directory directoryArg, String pattern)
         .withExec(List.of("grep", "-R", pattern, "."))
         .stdout();
   }
+
+  @Function
+  public DaggerJava itself() {
+    return this;
+  }
+
+  @Function
+  public boolean isZero(int value) {
+    return value == 0;
+  }
+
+  @Function
+  public int[] doThings(String[] stringArray, List<Integer> ints, List<Container> containers) {
+    int[] intsArray = {stringArray.length, ints.size()};
+    return intsArray;
+  }
 }