Simplify builder usage

1 files changed, 8 insertions, 0 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java b/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
index cd324b9..65d2fda 100644
--- a/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
+++ b/src/main/java/org/zwobble/couscous/ast/ClassNodeBuilder.java
@@ -77,10 +77,12 @@ public class ClassNodeBuilder {
     
     public class MethodBuilder {
         private final String methodName;
+        private final ImmutableList.Builder<FormalArgumentNode> arguments;
         private final ImmutableList.Builder<StatementNode> statements;
 
         public MethodBuilder(String name) {
             this.methodName = name;
+            this.arguments = ImmutableList.builder();
             this.statements = ImmutableList.builder();
         }
         
@@ -88,6 +90,11 @@ public class ClassNodeBuilder {
             return ThisReferenceNode.thisReference(name);
         }
         
+        public MethodBuilder argument(FormalArgumentNode argument) {
+            arguments.add(argument);
+            return this;
+        }
+        
         public MethodBuilder statement(StatementNode statement) {
             statements.add(statement);
             return this;
@@ -96,6 +103,7 @@ public class ClassNodeBuilder {
         public MethodNode build() {
             return MethodNode.builder()
                 .name(methodName)
+                .arguments(arguments.build())
                 .body(statements.build())
                 .build();
         }