diff --git a/src/main/java/ASTGenerator.java b/src/main/java/ASTGenerator.java
index b629df4..0d9570f 100644
--- a/src/main/java/ASTGenerator.java
+++ b/src/main/java/ASTGenerator.java
@@ -2,14 +2,16 @@ import abstractSyntaxTree.Class.FieldDecl;
 import abstractSyntaxTree.Class.IClass;
 import abstractSyntaxTree.Class.MethodDecl;
 import abstractSyntaxTree.Class.RefType;
+import abstractSyntaxTree.Expression.BinaryExpression;
 import abstractSyntaxTree.Node;
 import abstractSyntaxTree.Parameter.Parameter;
 import abstractSyntaxTree.Parameter.ParameterList;
 import abstractSyntaxTree.Program;
-import abstractSyntaxTree.Statement.BlockStatement;
+import abstractSyntaxTree.Statement.*;
 import gen.DecafBaseVisitor;
 import gen.DecafParser;
 
+import java.beans.Expression;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -25,6 +27,7 @@ public class ASTGenerator extends DecafBaseVisitor<Node> {
 
     @Override
     public Node visitClassdecl(DecafParser.ClassdeclContext ctx) {
+        String name = ctx.Identifier().getText();
         List<FieldDecl> fieldDecls = new ArrayList<>();
         List<MethodDecl> methodDecls = new ArrayList<>();
         boolean hasMain;
@@ -37,12 +40,16 @@ public class ASTGenerator extends DecafBaseVisitor<Node> {
             fieldDecls.add((FieldDecl) visit(fieldDecl));
         }
         for (DecafParser.ConstuctorDeclContext constDecl: ctx.constuctorDecl()) {
-            methodDecls.add((MethodDecl) visit(constDecl));
+            MethodDecl constructor = ((MethodDecl) visit(constDecl));
+            constructor.classThatContainsMethod = name;
+            methodDecls.add(constructor);
         }
         for (DecafParser.MethodDeclContext methodDecl: ctx.methodDecl()) {
-            methodDecls.add((MethodDecl) visit(methodDecl));
+            MethodDecl method = (MethodDecl) visit(methodDecl);
+            method.classThatContainsMethod = name;
+            methodDecls.add(method);
         }
-        return new RefType(ctx.Identifier().getText(),fieldDecls, methodDecls, hasMain);
+        return new RefType(name,fieldDecls, methodDecls, hasMain);
     }
 
     @Override
@@ -83,4 +90,48 @@ public class ASTGenerator extends DecafBaseVisitor<Node> {
         }
         return new ParameterList(parameters);
     }
+
+    @Override
+    public Node visitStatement(DecafParser.StatementContext ctx) {
+        if (ctx.returnStmt() != null) {
+            return visitReturnStmt(ctx.returnStmt());
+        } else if (ctx.ifElseStmt() != null) {
+            return visitIfElseStmt(ctx.ifElseStmt());
+        } else if (ctx.whileStmt() != null) {
+            return visitWhileStmt(ctx.whileStmt());
+        } else if (ctx.block() != null) {
+            return visitBlock(ctx.block());
+        } else if (ctx.stmtExpr() != null) {
+            return visitStmtExpr(ctx.stmtExpr());
+        }
+        return null;
+    }
+
+    @Override
+    public Node visitReturnStmt(DecafParser.ReturnStmtContext ctx) {
+        return new ReturnStatement(new BinaryExpression());
+    }
+
+    @Override
+    public Node visitIfElseStmt(DecafParser.IfElseStmtContext ctx) {
+        if (ctx.elseStmt() != null) {
+            return visitIfStmt(ctx.ifStmt());
+        } else {
+            Node ifStatement = visitStatement(ctx.ifStmt().statement());
+            Node elseStatement = visitStatement(ctx.elseStmt().statement());
+            return new IfElseStatement(new BinaryExpression(), (IStatement) ifStatement, (IStatement) elseStatement);
+        }
+    }
+
+    @Override
+    public Node visitIfStmt(DecafParser.IfStmtContext ctx) {
+        Node statement = visitStatement(ctx.statement());
+        return new IfStatement(new BinaryExpression(), (IStatement) statement);
+    }
+
+    @Override
+    public Node visitWhileStmt(DecafParser.WhileStmtContext ctx) {
+        Node statement = visitStatement(ctx.statement());
+        return new WhileStatement(new BinaryExpression(), (IStatement) statement);
+    }
 }
diff --git a/src/main/java/Compiler.java b/src/main/java/Compiler.java
index 320d575..d3c6153 100644
--- a/src/main/java/Compiler.java
+++ b/src/main/java/Compiler.java
@@ -75,7 +75,7 @@ public class Compiler {
 //        MethodDecl methodDecl2 = new MethodDecl("ClassA", "int", "m", new ArrayList<>(), new ArrayList<>());
 //        methodDecls.add(methodDecl2);
 
-        abstractSyntaxTree.classes.add(new RefType("MyClass", fieldDecls, methodDecls));
+        abstractSyntaxTree.classes.add(new RefType("MyClass", fieldDecls, methodDecls, false));
 
         System.out.println("Parsed " + abstractSyntaxTree.classes.size() + " classes with identifiers/names:");
         for (RefType refType : abstractSyntaxTree.classes) {