1 files changed, 8 insertions, 1 deletions
diff --git a/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java b/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
index b41c972..f1240f6 100644
--- a/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
+++ b/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
@@ -2,6 +2,7 @@ package org.zwobble.couscous.ast;
 
 import com.google.common.collect.ImmutableList;
 import org.zwobble.couscous.ast.identifiers.Identifier;
+import org.zwobble.couscous.values.UnitValue;
 
 import java.util.Optional;
 import java.util.function.Function;
@@ -21,7 +22,6 @@ public class ClassNodeBuilder {
         this.fields = ImmutableList.builder();
         this.constructor = Optional.empty();
         this.methods = ImmutableList.builder();
-
     }
 
     public ClassNodeBuilder(String name) {
@@ -76,6 +76,7 @@ public class ClassNodeBuilder {
             .isStatic(isStatic)
             .annotations(builder.annotations.build())
             .arguments(builder.arguments.build())
+            .returns(builder.returnType)
             .body(builder.statements.build())
             .build());
     }
@@ -93,6 +94,7 @@ public class ClassNodeBuilder {
         private final Function<MethodBuilder<?>, T> build;
         private final ImmutableList.Builder<AnnotationNode> annotations;
         private final ImmutableList.Builder<FormalArgumentNode> arguments;
+        private TypeName returnType = UnitValue.REF;
         private final ImmutableList.Builder<StatementNode> statements;
 
         public MethodBuilder(Function<MethodBuilder<?>, T> build) {
@@ -115,6 +117,11 @@ public class ClassNodeBuilder {
             arguments.add(formalArg(var(id, name, type)));
             return this;
         }
+
+        public MethodBuilder<T> returns(TypeName returnType) {
+            this.returnType = returnType;
+            return this;
+        }
         
         public MethodBuilder<T> argument(FormalArgumentNode argument) {
             arguments.add(argument);