diff --git a/src/main/java/parser/astBuilder/ASTBuilder.java b/src/main/java/parser/astBuilder/ASTBuilder.java
index bf84fd4..b6d8283 100644
--- a/src/main/java/parser/astBuilder/ASTBuilder.java
+++ b/src/main/java/parser/astBuilder/ASTBuilder.java
@@ -60,13 +60,16 @@ public class ASTBuilder extends SimpleJavaBaseVisitor<ASTNode> {
                 constructorNode.addParameter((ParameterNode) visit(parameter));
             }
         }
+        constructorNode.block.addStatement(new ReturnNode(null));
         return constructorNode;
     }
 
     @Override
     public ASTNode visitMethodDeclaration(SimpleJavaParser.MethodDeclarationContext ctx) {
         if(ctx.MainMethodDeclaration() != null) {
-            return new MainMethodNode((BlockNode) visit(ctx.blockStatement()));
+            MainMethodNode mainMethod = new MainMethodNode((BlockNode) visit(ctx.blockStatement()));
+            mainMethod.block.addStatement(new ReturnNode(null));
+            return mainMethod;
         } else {
             if(ctx.type() != null) {
                 MethodNode methodNode = new MethodNode(ctx.AccessModifier().getText(), createTypeNode(ctx.type().getText()), false, ctx.Identifier().getText(), (BlockNode) visit(ctx.blockStatement()));
@@ -83,6 +86,7 @@ public class ASTBuilder extends SimpleJavaBaseVisitor<ASTNode> {
                         methodNode.addParameter((ParameterNode) visit(parameter));
                     }
                 }
+                methodNode.block.addStatement(new ReturnNode(null));
                 return methodNode;
             }
         }
@@ -133,16 +137,8 @@ public class ASTBuilder extends SimpleJavaBaseVisitor<ASTNode> {
     @Override
     public ASTNode visitBlockStatement(SimpleJavaParser.BlockStatementContext ctx) {
         BlockNode blockNode = new BlockNode();
-        boolean hasReturnStatement = false;
         for(SimpleJavaParser.StatementContext statement : ctx.statement()) {
-            IStatementNode statementNode = (IStatementNode) visit(statement);
-            if(statementNode instanceof ReturnNode) {
-                hasReturnStatement = true;
-            }
-            blockNode.addStatement(statementNode);
-        }
-        if(!hasReturnStatement) {
-            blockNode.addStatement(new ReturnNode(null));
+            blockNode.addStatement((IStatementNode) visit(statement));
         }
         return blockNode;
     }