Reaching definition analysis (#17)

Co-authored-by: Santo Cariotti <santo@dcariotti.me>
Co-authored-by: geno <gabriele.genovese2@studio.unibo.it>
Co-authored-by: geno <gabrigeno@gmail.com>

46 files changed, 1151 insertions, 204 deletions

diff --git a/.gitignore b/.gitignore
index b89c3af..6730f0b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,5 @@ trees/
 .project
 .vscode/
 src/parser/.antlr/
-code.asm
\ No newline at end of file
+code.asm
+optimized.py
\ No newline at end of file
diff --git a/README.md b/README.md
index d8b2780..ec7323c 100644
--- a/README.md
+++ b/README.md
@@ -2,12 +2,130 @@
 
 Project for the AA 2023/2024 **Complementi di programmazione**'s course of the Master Degree in Computer Science.
 
-# How to start
+## How to build
+
+You first need to build the project using the Makefile.
+
+```
+make build
+```
+
+## Run
+
+You can run the software using the following command
+
+```
+java -cp lib/antlr-4.13.1-complete.jar:out Main <file.py>
+```
+
+you also can use three flags on the execution:
+
+- `--optimize`: runs the optimizer for the file `<file.py>`. The optimized
+  version is auto saved at `./optimized.py`.
+
+- `--tree`: shows the AST on video and on JFrame.
+
+- `--exec`: executes the SVM for the Assembly code generated in `code.asm`.
+
+You can use all the three flags togheter if you want to.
+
+## Example
+
+- Optimizer 
 
-```shell
-make
-make run
-make runall
 ```
+$ java -cp lib/antlr-4.13.1-complete.jar:out Main test/2a.py --optimize
+n=int(input())
+x=int(input())
+m=1
+tmp=2 * x
+while n>1:
 
-For archlinux users, add the following env variable to visualize the window: `export _JAVA_AWT_WM_NONREPARENTING=1`
+        m=m*tmp
+        n=n-1
+
+
+print(m+tmp)
+
+Saving optimized file...
+Everything is OK!
+```
+
+- Tree
+
+```
+$ java -cp lib/antlr-4.13.1-complete.jar:out Main test/2a.py --tree
+Root
+  SimpleStmts
+    SimpleStmt
+      Assignment
+        ExprList
+          Expr
+            AtomNode: n
+        Augassign(=)
+        ExprList
+          Expr
+            AtomNode: int
+            TrailerNode
+              ArglistNode
+                Expr
+      Assignment                                                                                                                                                                                             [40/225]
+    [...] 
+    SimpleStmt
+      Expr
+        AtomNode: print
+        TrailerNode
+          ArglistNode
+            Expr
+              Expr
+                AtomNode: m
+              Expr
+                AtomNode: tmp
+              Op(+)
+
+Everything is OK!
+```
+
+- Exec
+
+```
+$ java -cp lib/antlr-4.13.1-complete.jar:out Main test/3-1.py --exec
+Creating VM code...
+Executing assemply code...
+0: 17 ----- SP = 999, FP = 999, AL = 998, RA = 0, A0 = 0, T1 = 0
+1: 17	999 ----- SP = 998, FP = 999, AL = 998, RA = 0, A0 = 0, T1 = 0
+2: 17	999	998 ----- SP = 997, FP = 999, AL = 998, RA = 0, A0 = 0, T1 = 0
+3: 7	999	998	999 ----- SP = 996, FP = 999, AL = 998, RA = 0, A0 = 0, T1 = 0
+4: 9	999	998	999 ----- SP = 996, FP = 996, AL = 998, RA = 0, A0 = 0, T1 = 0
+5: 7	999	998	999 ----- SP = 996, FP = 997, AL = 998, RA = 0, A0 = 0, T1 = 0
+[...]
+216: 18	999	998	999	998	29	2	14	997	996	29	3	140	992	991 ----- SP = 985, FP = 987, AL = 986, RA = 140, A0 = 1, T1 = 1
+217: 5	999	998	999	998	29	2	14	997	996	29	3	140	992 ----- SP = 986, FP = 987, AL = 986, RA = 140, A0 = 1, T1 = 1
+218: 7	999	998	999	998	29	2	14	997	996	29	3	140	992 ----- SP = 986, FP = 992, AL = 986, RA = 140, A0 = 1, T1 = 1
+219: 11	999	998	999	998	29	2	14	997	996	29	3	140	992 ----- SP = 986, FP = 992, AL = 992, RA = 140, A0 = 1, T1 = 1
+220: 18	999	998	999	998	29	2	14	997	996	29	3	140	992 ----- SP = 986, FP = 992, AL = 991, RA = 140, A0 = 1, T1 = 1
+221: 24	999	998	999	998	29	2	14	997	996	29	3	140 ----- SP = 987, FP = 992, AL = 991, RA = 140, A0 = 1, T1 = 1
+222: 19	999	998	999	998	29	2	14	997	996	29	3	140 ----- SP = 987, FP = 992, AL = 991, RA = 140, A0 = 1, T1 = 1
+223: 9	999	998	999	998	29	2	14	997	996	29	3 ----- SP = 988, FP = 992, AL = 991, RA = 140, A0 = 1, T1 = 1
+224: 18	999	998	999	998	29	2	14	997	996 ----- SP = 990, FP = 992, AL = 991, RA = 140, A0 = 1, T1 = 1
+225: 5	999	998	999	998	29	2	14	997 ----- SP = 991, FP = 992, AL = 991, RA = 140, A0 = 1, T1 = 1
+226: 7	999	998	999	998	29	2	14	997 ----- SP = 991, FP = 997, AL = 991, RA = 140, A0 = 1, T1 = 1
+227: 11	999	998	999	998	29	2	14	997 ----- SP = 991, FP = 997, AL = 997, RA = 140, A0 = 1, T1 = 1
+228: 18	999	998	999	998	29	2	14	997 ----- SP = 991, FP = 997, AL = 996, RA = 140, A0 = 1, T1 = 1
+229: 24	999	998	999	998	29	2	14 ----- SP = 992, FP = 997, AL = 996, RA = 140, A0 = 1, T1 = 1
+230: 19	999	998	999	998	29	2	14 ----- SP = 992, FP = 997, AL = 996, RA = 140, A0 = 1, T1 = 1
+231: 9	999	998	999	998	29	2 ----- SP = 993, FP = 997, AL = 996, RA = 14, A0 = 1, T1 = 1
+232: 18	999	998	999	998 ----- SP = 995, FP = 997, AL = 996, RA = 14, A0 = 1, T1 = 1
+233: 5	999	998	999 ----- SP = 996, FP = 997, AL = 996, RA = 14, A0 = 1, T1 = 1
+234: 7	999	998	999 ----- SP = 996, FP = 999, AL = 996, RA = 14, A0 = 1, T1 = 1
+235: 11	999	998	999 ----- SP = 996, FP = 999, AL = 999, RA = 14, A0 = 1, T1 = 1
+236: 18	999	998	999 ----- SP = 996, FP = 999, AL = 998, RA = 14, A0 = 1, T1 = 1
+237: 24	999	998 ----- SP = 997, FP = 999, AL = 998, RA = 14, A0 = 1, T1 = 1
+238: 18	999	998 ----- SP = 997, FP = 999, AL = 998, RA = 14, A0 = 1, T1 = 1
+239: 18	999 ----- SP = 998, FP = 999, AL = 998, RA = 14, A0 = 1, T1 = 1
+240: 25 ----- SP = 999, FP = 999, AL = 998, RA = 14, A0 = 1, T1 = 1
+
+Result: 1
+
+Everything is OK!
+```
diff --git a/progs/test.py b/progs/test.py
index 4ccfb1a..ef59445 100644
--- a/progs/test.py
+++ b/progs/test.py
@@ -1,3 +1,8 @@
+n = int(input())
+x = int(input())
 m = 1
-for i in range(10):
-    m = m + 1
+while n > 1:
+    tmp = 2 * x
+    m = m * tmp
+    n = n - 1
+print(m + tmp)
diff --git a/src/Main.java b/src/Main.java
index 15d472d..5970b3e 100644
--- a/src/Main.java
+++ b/src/Main.java
@@ -1,6 +1,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+
 import javax.swing.*;
 import org.antlr.v4.gui.TreeViewer;
 import org.antlr.v4.runtime.*;
@@ -11,15 +12,18 @@ import java.nio.file.Path;
 import java.nio.file.Paths;
 import ast.*;
 import ast.nodes.*;
-import parser.Python3Lexer;
-import parser.Python3Parser;
+import parser.*;
 import semanticanalysis.*;
 import svm.*;
 
 public class Main {
 
+    private static boolean execute = false;
+    private static boolean treeFlag = false;
+    private static boolean optimize = false;
+
     public static void main(String[] args) {
-        if (args.length != 1) {
+        if (args.length < 1) {
             System.err
                     .println(
                             "You must execute this program with a file parameter.\nUsage: java -cp lib/antlr-4.13.1-complete.jar:out Main <file.py>");
@@ -27,36 +31,58 @@ public class Main {
         }
 
         try {
+            for (int i = 0; i < args.length; i++) {
+                if (args[i].equals("--exec")) {
+                    execute = true;
+                }
+
+                if (args[i].equals("--tree")) {
+                    treeFlag = true;
+                }
+
+                if (args[i].equals("--optimize")) {
+                    optimize = true;
+                }
+            }
+
             String fileStr = args[0];
-            System.out.println(fileStr);
-            System.out.println(Share.readFile(fileStr));
             CharStream cs = CharStreams.fromFileName(fileStr);
             Python3Lexer lexer = new Python3Lexer(cs);
             CommonTokenStream tokens = new CommonTokenStream(lexer);
             Python3Parser parser = new Python3Parser(tokens);
             Python3Parser.RootContext tree = parser.root();
 
-            JFrame frame = new JFrame("Parse Tree");
-            JPanel panel = new JPanel();
-            TreeViewer viewer = new TreeViewer(Arrays.asList(parser.getRuleNames()),
-                    tree);
-            viewer.setScale(1); // Zoom factor
-            panel.add(viewer);
-            frame.add(panel);
-            frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
-            frame.setSize(800, 600);
-            frame.setVisible(true);
             if (tree == null) {
                 System.err.println("The tree is null.");
                 return;
             }
+
             if (parser.getNumberOfSyntaxErrors() > 0) {
                 System.err.println("Error on program parsing.");
                 return;
             }
-            Python3VisitorImpl visitor = new Python3VisitorImpl();
-            SymbolTable ST = new SymbolTable();
+
+            Python3VisitorImpl visitor = new Python3VisitorImpl(tokens, optimize);
             Node ast = visitor.visit(tree);
+            SymbolTable ST = new SymbolTable();
+
+            if (optimize) {
+                // first visit to optimize
+                cs = CharStreams.fromString(visitor.getRewriter());
+                System.out.println(cs);
+                lexer = new Python3Lexer(cs);
+                tokens = new CommonTokenStream(lexer);
+                parser = new Python3Parser(tokens);
+                tree = parser.root();
+                ST = new SymbolTable();
+
+                ast = visitor.visit(tree);
+
+                System.out.println("Saving optimized file...");
+                String astToPrint = ast.toPrint("");
+                Share.saveFile("optimized.py", astToPrint);
+            }
+
             ArrayList<SemanticError> errorsWithDup = ast.checkSemantics(ST, 0, null);
             ArrayList<SemanticError> errors = Share.removeDuplicates(errorsWithDup);
             if (!errors.isEmpty()) {
@@ -65,35 +91,57 @@ public class Main {
                     System.out.println("\t" + e);
                 }
             } else {
-                System.out.println("Visualizing AST...");
-                System.out.println(ast.toPrint(""));
-                System.out.println("Creating VM code...");
-                String prog = ast.codeGeneration();
-                String asmFile = "code.asm";
-                Path file = Paths.get(asmFile);
-                if (!Files.exists(file)) {
-                    Files.createFile(file);
+                if (treeFlag) {
+                    JFrame frame = new JFrame("Parse Tree");
+                    JPanel panel = new JPanel();
+                    TreeViewer viewer = new TreeViewer(Arrays.asList(parser.getRuleNames()),
+                            tree);
+                    viewer.setScale(1); // Zoom factor
+                    panel.add(viewer);
+                    frame.add(panel);
+                    frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
+                    frame.setSize(800, 600);
+                    frame.setVisible(true);
+
+                    System.out.println(ast.printAST(""));
+                }
+
+                if (execute) {
+                    System.out.println("Creating VM code...");
+                    String asmFile = "code.asm";
+                    String prog = ast.codeGeneration();
+                    Share.saveFile(asmFile, prog);
+                    executeVMFile(asmFile);
                 }
-                Files.write(file, prog.getBytes(), StandardOpenOption.TRUNCATE_EXISTING);
-                System.out.println("Done!");
 
-                CharStream inputASM = CharStreams.fromFileName(asmFile);
-                SVMLexer lexerASM = new SVMLexer(inputASM);
-                CommonTokenStream tokensASM = new CommonTokenStream(lexerASM);
-                SVMParser parserASM = new SVMParser(tokensASM);
+                System.out.println("Everything is OK!");
+            }
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+    }
 
-                SVMVisitorImpl visitorSVM = new SVMVisitorImpl();
-                visitorSVM.visit(parserASM.assembly());
+    private static void executeVMFile(String fileName) {
+        try {
+            CharStream inputASM = CharStreams.fromFileName(fileName);
+            SVMLexer lexerASM = new SVMLexer(inputASM);
+            CommonTokenStream tokensASM = new CommonTokenStream(lexerASM);
+            SVMParser parserASM = new SVMParser(tokensASM);
 
-                System.out.println("You had: " + lexerASM.lexicalErrors + " lexical errors and "
-                        + parserASM.getNumberOfSyntaxErrors() + " syntax errors.");
-                if (lexerASM.lexicalErrors > 0 || parserASM.getNumberOfSyntaxErrors() > 0)
-                    System.exit(1);
+            SVMVisitorImpl visitorSVM = new SVMVisitorImpl();
+            visitorSVM.visit(parserASM.assembly());
 
-                System.out.println("Starting Virtual Machine...");
-                ExecuteVM vm = new ExecuteVM(visitorSVM.code);
-                vm.cpu();
+            int nLexicalErrors = lexerASM.lexicalErrors;
+            int nSyntaxErrors = parserASM.getNumberOfSyntaxErrors();
+            if (nLexicalErrors > 0 || nSyntaxErrors > 0) {
+                System.out.println("You had " + nLexicalErrors + " lexical errors and "
+                        + nSyntaxErrors + " syntax errors in the asm file.");
+                System.exit(1);
             }
+
+            System.out.println("Executing assemply code...");
+            ExecuteVM vm = new ExecuteVM(visitorSVM.code);
+            vm.cpu();
         } catch (Exception e) {
             e.printStackTrace();
         }
diff --git a/src/ParseAll.java b/src/ParseAll.java
index 06b99ad..091a2f4 100644
--- a/src/ParseAll.java
+++ b/src/ParseAll.java
@@ -35,7 +35,7 @@ public class ParseAll {
                 Python3Parser.RootContext tree = parser.root();
                 String treeStr = tree.toStringTree();
 
-                Python3VisitorImpl visitor = new Python3VisitorImpl();
+                Python3VisitorImpl visitor = new Python3VisitorImpl(tokenStream, false);
                 SymbolTable ST = new SymbolTable();
                 Node ast = visitor.visit(tree);
                 ArrayList<SemanticError> errorsWithDup = ast.checkSemantics(ST, 0, null);
diff --git a/src/ast/Python3VisitorImpl.java b/src/ast/Python3VisitorImpl.java
index 98ede6d..6b3d7d0 100644
--- a/src/ast/Python3VisitorImpl.java
+++ b/src/ast/Python3VisitorImpl.java
@@ -1,10 +1,19 @@
 package ast;
 
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 import ast.nodes.*;
+import ast.types.*;
+import codegen.Label;
+import parser.Python3Lexer;
 import parser.Python3ParserBaseVisitor;
 import parser.Python3Parser.*;
+
+import org.antlr.v4.runtime.*;
+import org.antlr.v4.runtime.tree.*;
 import org.antlr.v4.runtime.tree.TerminalNode;
 
 /**
@@ -13,6 +22,25 @@ import org.antlr.v4.runtime.tree.TerminalNode;
  */
 public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
 
+    Map<String, Integer> R;
+    private TokenStreamRewriter rewriter;
+    private boolean optimize;
+    private boolean optimizationDone;
+
+    public Python3VisitorImpl(CommonTokenStream tokens, boolean optimize) {
+        rewriter = new TokenStreamRewriter(tokens);
+        this.optimize = optimize;
+        optimizationDone = false;
+    }
+
+    public String getRewriter() {
+        return rewriter.getText();
+    }
+
+    public boolean getOptimizationDone() {
+        return optimizationDone;
+    }
+
     /**
      * Since a root can be a simple_stmts or a compound_stmt, this method
      * returns a new `RootNode` with a list of them.
@@ -20,7 +48,9 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ``` root : NEWLINE* (simple_stmts | compound_stmt)* EOF; ```
      */
     public Node visitRoot(RootContext ctx) {
-        ArrayList<Node> childs = new ArrayList<Node>();
+        ArrayList<Node> childs = new ArrayList<>();
+
+        R = new HashMap<>();
 
         for (int i = 0; i < ctx.getChildCount(); i++) {
             var child = ctx.getChild(i);
@@ -32,6 +62,7 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
             }
         }
 
+        // cfg.addEdge(cfg.getExitNode());
         return new RootNode(childs);
     }
 
@@ -41,7 +72,7 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ``` simple_stmts : simple_stmt (';' simple_stmt)* ';'? NEWLINE ; ```
      */
     public Node visitSimple_stmts(Simple_stmtsContext ctx) {
-        ArrayList<Node> stmts = new ArrayList<Node>();
+        ArrayList<Node> stmts = new ArrayList<>();
 
         for (Simple_stmtContext stm : ctx.simple_stmt()) {
             stmts.add(visit(stm));
@@ -64,21 +95,17 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
 
         if (ctx.if_stmt() != null) {
             ifStmt = visit(ctx.if_stmt());
-        }
-
-        if (ctx.funcdef() != null) {
+        } else if (ctx.funcdef() != null) {
             funcDef = visit(ctx.funcdef());
-        }
-
-        if (ctx.for_stmt() != null) {
+        } else if (ctx.for_stmt() != null) {
             forStmt = visit(ctx.for_stmt());
-        }
-
-        if (ctx.while_stmt() != null) {
+        } else if (ctx.while_stmt() != null) {
             whileStmt = visit(ctx.while_stmt());
         }
 
-        return new CompoundNode(ifStmt, funcDef, forStmt, whileStmt);
+        CompoundNode compoundNode = new CompoundNode(ifStmt, funcDef, forStmt, whileStmt);
+
+        return compoundNode;
     }
 
     /**
@@ -123,7 +150,13 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
         Node assign = visit(ctx.augassign());
         Node rhr = visit(ctx.exprlist(1));
 
-        return new AssignmentNode(lhr, assign, rhr);
+        AssignmentNode assignmentNode = new AssignmentNode(lhr, assign, rhr,
+                ctx.exprlist(0).getStart().getTokenIndex(),
+                ctx.exprlist(1).getStop().getTokenIndex());
+
+        R.put(((ExprNode) ((ExprListNode) lhr).getElem(0)).getId(), ctx.getStart().getLine());
+
+        return assignmentNode;
     }
 
     /**
@@ -154,7 +187,7 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
 
         Node dottedName = visit(ctx.dotted_name());
 
-        ArrayList<String> names = new ArrayList<String>();
+        ArrayList<String> names = new ArrayList<>();
 
         for (var s : ctx.NAME()) {
             names.add(s.toString());
@@ -169,7 +202,7 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ``` dotted_name : NAME ('.' NAME)* ; ```
      */
     public Node visitDotted_name(Dotted_nameContext ctx) {
-        ArrayList<TerminalNode> names = new ArrayList<TerminalNode>();
+        ArrayList<TerminalNode> names = new ArrayList<>();
 
         for (var name : ctx.NAME()) {
             names.add(name);
@@ -191,6 +224,8 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
         }
 
         Node block = visit(ctx.block());
+        rewriter.insertAfter(ctx.getStart().getTokenIndex(), " ");
+        rewriter.insertBefore(ctx.CLOSE_PAREN().getSymbol().getStartIndex() - 1, "\n    ");
 
         return new FuncdefNode(ctx.NAME(), paramlist, block);
     }
@@ -204,7 +239,7 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ```
      */
     public Node visitParamlist(ParamlistContext ctx) {
-        ArrayList<Node> params = new ArrayList<Node>();
+        ArrayList<Node> params = new ArrayList<>();
 
         for (ParamdefContext s : ctx.paramdef()) {
             params.add(visit(s));
@@ -220,6 +255,7 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ``` paramdef : NAME (':' expr)? ; ```
      */
     public Node visitParamdef(ParamdefContext ctx) {
+        R.remove(ctx.NAME().toString());
         return new ParamdefNode(ctx.NAME().toString());
     }
 
@@ -267,7 +303,7 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
     }
 
     /**
-     * Returns a `IfNode`. FIXME: add support for elif statement.
+     * Returns a `IfNode`.
      *
      * ``` if_stmt : 'if' expr ':' block ('elif' expr ':' block)* ('else' ':'
      * block)? ; ```
@@ -275,12 +311,15 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
     public Node visitIf_stmt(If_stmtContext ctx) {
         var blocks = ctx.block();
         Node condExp = visit(ctx.expr(0));
+
         Node thenExp = visit(blocks.get(0));
+
         Node elseExp = null;
         if (blocks.size() > 1) {
             elseExp = visit(blocks.get(1));
         }
 
+        rewriter.insertAfter(ctx.getStart().getTokenIndex(), " ");
         return new IfNode(condExp, thenExp, elseExp);
     }
 
@@ -290,12 +329,132 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ``` while_stmt : 'while' expr ':' block ('else' ':' block)? ; ```
      */
     public Node visitWhile_stmt(While_stmtContext ctx) {
-        Node expr = visit(ctx.expr());
+        // Do the same for the while expression and the block
+        ExprNode expr = (ExprNode) visit(ctx.expr());
 
         // Block 1 is for the while-else statement
-        Node block = visit(ctx.block(0));
+        BlockNode block = (BlockNode) visit(ctx.block(0));
+
+        WhileStmtNode whileStmt = new WhileStmtNode(expr, block);
+        if (!optimize) {
+            return whileStmt;
+        }
+
+        rewriter.insertAfter(ctx.COLON(0).getSymbol().getTokenIndex(), "\n");
+
+        int lineStart = ctx.getStart().getLine();
+        int lineStop = ctx.getStop().getLine();
+        int index = ctx.getStart().getTokenIndex();
+        optimizeWithSecond(block, lineStart, lineStop, index);
+
+        // optimize while's guard
+        int counter = 0;
+        var exprs = expr.getExprs();
+        for (var e : exprs) {
+            if (e instanceof ExprNode) {
+                ExprNode exprNode = (ExprNode) e;
+                if (exprNode.typeCheck() instanceof AtomType) {
+                    continue;
+                }
+            }
+            ArrayList<String> al = findAtomPresent(e, new ArrayList<>());
+            if (!al.isEmpty()) {
+                boolean constant = true;
+                for (String a : al) {
+                    int n = R.get(a);
+                    if (n > lineStart && n <= lineStop) {
+                        constant = false;
+                        break;
+                    }
+                }
+                if (constant) {
+                    String newVar = Label.newVar();
+                    rewriter.insertBefore(index, newVar + "=" + e.toPrint("") + "\n");
+                    int lastToken = ctx.expr().expr(counter).getStop().getTokenIndex();
+                    int firstToken = ctx.expr().expr(counter).getStart().getTokenIndex();
+                    rewriter.replace(firstToken, lastToken, newVar);
+                }
+            }
+            counter++;
+        }
 
-        return new WhileStmtNode(expr, block);
+        optimizeWithThird(block, lineStart, lineStop, index);
+
+        return whileStmt;
+    }
+
+    private ArrayList<String> findAtomPresent(Node e, ArrayList<String> Acc) {
+        if (e instanceof ExprNode) {
+            ExprNode expNode = (ExprNode) e;
+            ArrayList<Node> exprs = expNode.getExprs();
+            if (!exprs.isEmpty()) {
+                for (Node i : exprs) {
+                    findAtomPresent(i, Acc);
+                }
+            } else {
+                AtomNode a = (AtomNode) expNode.getAtom();
+                if (a.typeCheck() instanceof AtomType) {
+                    Acc.add(a.getId());
+                }
+            }
+        }
+        return Acc;
+    }
+
+    private void optimizeWithSecond(BlockNode block, int lineStart, int lineStop, int index) {
+        rewriter.insertAfter(index, " ");
+        ArrayList<AssignmentNode> assignments = new ArrayList<>();
+        for (var child : block.getChilds()) {
+            if (child instanceof SimpleStmtsNode) {
+                var stmts = (SimpleStmtsNode) child;
+                for (var stmt : stmts.getStmts()) {
+                    var assignment = (AssignmentNode) ((SimpleStmtNode) stmt).getAssignment();
+                    if (assignment != null) {
+                        assignments.add(assignment);
+                    }
+                }
+            }
+        }
+
+        // g , x + 2 * y
+        // m , m + n + g
+        // n , n + 1
+        for (var assignment : assignments) {
+
+            var lhr = (ExprNode) assignment.getLhr().getElem(0);
+            var rhr = (ExprNode) assignment.getRhr().getElem(0);
+            ArrayList<String> al = findAtomPresent(rhr, new ArrayList<>());
+            if (!al.isEmpty()) {
+                boolean constant = true;
+                for (String a : al) {
+                    if (R.get(a) == null) {
+                        rewriter.insertBefore(assignment.getLhrIndex(), "\n");
+                        rewriter.insertAfter(assignment.getLhrIndex() - 1, "\t");
+                        constant = false;
+                        break;
+                    }
+                    int n = R.get(a);
+                    if (n > lineStart && n <= lineStop) {
+                        constant = false;
+                        break;
+                    }
+                }
+
+                rewriter.insertAfter(assignment.getRhrIndex(), "\n");
+                if (constant) {
+                    rewriter.insertBefore(index, lhr.toPrint("") + "=" + rhr.toPrint("") + "\n");
+                    rewriter.replace(assignment.getLhrIndex(), assignment.getRhrIndex(), "");
+                    optimizationDone = true;
+                } else {
+                    rewriter.insertBefore(assignment.getLhrIndex(), "\t");
+                }
+            } else {
+                String newVar = Label.newVar();
+                rewriter.insertBefore(index, newVar + "=" + rhr.toPrint("") + "\n");
+                rewriter.replace(assignment.getLhrIndex(), assignment.getRhrIndex(),
+                        "\t" + lhr.toPrint("") + "=" + newVar + "\n");
+            }
+        }
     }
 
     /**
@@ -304,12 +463,89 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * ``` for_stmt : 'for' exprlist ':' block ('else' ':' block)? ; ```
      */
     public Node visitFor_stmt(For_stmtContext ctx) {
+
+        // Do the same for the for expression and the block
         Node exprList = visit(ctx.exprlist());
 
+        int dimGetExpr = ((ExprListNode) exprList).getExprs().size();
+
+        for (int e = 0; e < dimGetExpr; e++) {
+            if (e == dimGetExpr - 1) {
+                R.remove(((ExprNode) ((ExprNode) ((ExprListNode) exprList).getElem(e)).getExpr(0)).getId());
+            } else {
+                R.remove(((ExprNode) ((ExprListNode) exprList).getElem(e)).getId());
+            }
+        }
+
         // Block 1 is for the for-else statement
-        Node block = visit(ctx.block(0));
+        BlockNode block = (BlockNode) visit(ctx.block(0));
+
+        Node forNode = new ForStmtNode(exprList, block);
+        if (!optimize) {
+            return forNode;
+        }
+
+        int lineStart = ctx.getStart().getLine();
+        int lineStop = ctx.getStop().getLine();
+        int index = ctx.getStart().getTokenIndex();
 
-        return new ForStmtNode(exprList, block);
+        rewriter.insertAfter(index, " ");
+        // NOTE: works only for one argument
+        rewriter.insertAfter(index + 1, " ");
+        rewriter.insertAfter(index + 2, " ");
+        optimizeWithSecond(block, lineStart, lineStop, index);
+        optimizeWithThird(block, lineStart, lineStop, index);
+
+        return forNode;
+    }
+
+    private void optimizeWithThird(BlockNode block, int lineStart, int lineStop, int index) {
+        int counter = 0;
+        ArrayList<Node> stms = block.getChilds();
+        for (var e : stms) {
+            if (e instanceof SimpleStmtsNode) {
+                SimpleStmtsNode stmss = (SimpleStmtsNode) e;
+                for (Node stm : stmss.getStmts()) {
+                    SimpleStmtNode singleStm = (SimpleStmtNode) stm;
+                    AssignmentNode ass = (AssignmentNode) singleStm.getAssignment();
+                    if (ass != null) {
+                        ExprListNode rhr = ass.getRhr();
+                        ExprNode rExpr = (ExprNode) rhr.getElem(0);
+                        ArrayList<Node> exprsList = rExpr.getExprs();
+                        if (exprsList.size() > 1) {
+                            List<Node> exprsLists = exprsList.subList(0, exprsList.size() - 1);
+                            for (var elem : exprsLists) {
+                                if (elem instanceof ExprNode) {
+                                    ExprNode exprNode = (ExprNode) elem;
+                                    if (exprNode.typeCheck() instanceof AtomType) {
+                                        continue;
+                                    }
+                                }
+                                ArrayList<String> al = findAtomPresent(elem, new ArrayList<>());
+                                if (!al.isEmpty()) {
+                                    boolean constant = true;
+                                    for (String a : al) {
+                                        int n = R.get(a);
+                                        if (n > lineStart && n <= lineStop) {
+                                            constant = false;
+                                            break;
+                                        }
+                                    }
+                                    if (constant) {
+                                        String newVar = Label.newVar();
+                                        rewriter.insertBefore(index, newVar + "=" + elem.toPrint("") + "\n");
+                                        int firstToken = ass.getLhrIndex() + 2;
+                                        int lastToken = ass.getRhrIndex() - 2;
+                                        rewriter.replace(firstToken, lastToken, newVar);
+                                    }
+                                }
+                                counter++;
+                            }
+                        }
+                    }
+                }
+            }
+        }
     }
 
     /**
@@ -418,18 +654,35 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
 
         if (ctx.atom() != null) {
             atom = visit(ctx.atom());
-        }
 
-        if (ctx.comp_op() != null) {
-            compOp = visit(ctx.comp_op());
-        }
+            for (TrailerContext s : ctx.trailer()) {
+                trailers.add(visit(s));
+            }
+        } else {
+            if (ctx.ADD(0) != null) {
+                op = ctx.ADD(0).toString();
 
-        for (ExprContext s : ctx.expr()) {
-            exprs.add(visit(s));
-        }
+            } else if (ctx.MINUS(0) != null) {
+                op = ctx.MINUS(0).toString();
 
-        for (TrailerContext s : ctx.trailer()) {
-            trailers.add(visit(s));
+            } else if (ctx.NOT() != null) {
+                op = ctx.NOT().toString();
+
+            } else if (ctx.STAR() != null) {
+                op = ctx.STAR().toString();
+
+            } else if (ctx.DIV() != null) {
+                op = ctx.DIV().toString();
+
+            }
+
+            if (ctx.comp_op() != null) {
+                compOp = visit(ctx.comp_op());
+            }
+
+            for (ExprContext s : ctx.expr()) {
+                exprs.add(visit(s));
+            }
         }
 
         return new ExprNode(atom, compOp, exprs, op, trailers);
@@ -461,11 +714,11 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
             }
             return new AtomNode(varName, null);
         } else if (ctx.OPEN_BRACE() != null && ctx.CLOSE_BRACE() != null) {
-            return manageCompListContext(tlc);
+            return manageCompListContext(tlc, "{", "}");
         } else if (ctx.OPEN_BRACK() != null && ctx.CLOSE_BRACK() != null) {
-            return manageCompListContext(tlc);
+            return manageCompListContext(tlc, "[", "]");
         } else if (ctx.OPEN_PAREN() != null && ctx.CLOSE_PAREN() != null) {
-            return manageCompListContext(tlc);
+            return manageCompListContext(tlc, "(", ")");
         }
         return new AtomNode(null, null);
     }
@@ -475,12 +728,12 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      * `testlist_comp` set if the context is not null. Otherwise, returns an
      * `AtomNode` with nulls.
      */
-    public AtomNode manageCompListContext(Testlist_compContext tlc) {
+    public AtomNode manageCompListContext(Testlist_compContext tlc, String prefix, String suffix) {
         if (tlc != null) {
             Node testlist_comp = visit(tlc);
-            return new AtomNode(null, testlist_comp);
+            return new AtomNode(null, testlist_comp, prefix, suffix);
         }
-        return new AtomNode(null, null);
+        return new AtomNode(null, null, prefix, suffix);
     }
 
     /**
@@ -491,6 +744,19 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
      */
     public Node visitTrailer(TrailerContext ctx) {
         Node arglist = null;
+        String prefix = "";
+        String suffix = "";
+
+        if (ctx.OPEN_BRACK() != null) {
+            prefix = "[";
+            suffix = "]";
+        } else if (ctx.OPEN_PAREN() != null) {
+            prefix = "(";
+            suffix = ")";
+        } else if (ctx.DOT() != null) {
+            prefix = ".";
+        }
+
         if (ctx.arglist() != null) {
             arglist = visit(ctx.arglist());
         }
@@ -506,11 +772,11 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
             methodCall = ctx.NAME();
         }
 
-        return new TrailerNode(arglist, exprs, methodCall, ctx.OPEN_PAREN() != null);
+        return new TrailerNode(arglist, exprs, methodCall, prefix, suffix);
     }
 
     /**
-     * Returns a `Node`. FIXME: what to do in case of list??
+     * Returns a `Node`.
      *
      * ``` exprlist : expr (',' expr )* ','? ; ```
      */
@@ -577,12 +843,11 @@ public class Python3VisitorImpl extends Python3ParserBaseVisitor<Node> {
 
     /**
      * Returns a `CompIterNode`.
+     * NOTE: We ignore `comp_if`.
      *
      * ``` comp_iter : comp_for | comp_if ; ;```
      */
     public Node visitComp_iter(Comp_iterContext ctx) {
-        // TODO: Implement comp_if
-        // Node iter = visit(ctx.comp_if());
         Comp_forContext cfc = ctx.comp_for();
         Node forNode = null;
         if (cfc != null) {
diff --git a/src/ast/nodes/ArglistNode.java b/src/ast/nodes/ArglistNode.java
index 5da7e2c..4b7b292 100644
--- a/src/ast/nodes/ArglistNode.java
+++ b/src/ast/nodes/ArglistNode.java
@@ -1,6 +1,7 @@
 package ast.nodes;
 
 import ast.types.*;
+
 import java.util.ArrayList;
 import java.util.Arrays;
 import semanticanalysis.SemanticError;
@@ -27,7 +28,6 @@ public class ArglistNode implements Node {
                 String argName = argExpr.getId();
                 errors.addAll(arg.checkSemantics(ST, _nesting, ft));
 
-                // TODO: check IntType for params
                 if (argName != null) {
                     if (Arrays.asList(bif).contains(argName)) {
                         continue;
@@ -66,12 +66,26 @@ public class ArglistNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "ArglistNode\n";
 
         prefix += "  ";
         for (Node arg : arguments) {
-            str += arg.toPrint(prefix);
+            str += arg.printAST(prefix);
+        }
+
+        return str;
+    }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix;
+        str += arguments.get(0).toPrint("");
+
+        if (arguments.size() > 1) {
+            for (int i = 1; i < arguments.size(); i++) {
+                str += ", " + arguments.get(i).toPrint("");
+            }
         }
 
         return str;
diff --git a/src/ast/nodes/AssignmentNode.java b/src/ast/nodes/AssignmentNode.java
index 13e7a33..688be76 100644
--- a/src/ast/nodes/AssignmentNode.java
+++ b/src/ast/nodes/AssignmentNode.java
@@ -16,15 +16,20 @@ public class AssignmentNode implements Node {
     private final Node assign;
     private final ExprListNode rhr;
 
+    private final int startIndex;
+    private final int endIndex;
+
     // Useful for code gen
     private int offset;
     private boolean alreadyDef;
 
-    public AssignmentNode(Node lhr, Node assign, Node rhr) {
+    public AssignmentNode(Node lhr, Node assign, Node rhr, int startIndex, int endIndex) {
         this.lhr = (ExprListNode) lhr;
         this.assign = assign;
         this.rhr = (ExprListNode) rhr;
         this.alreadyDef = false;
+        this.startIndex = startIndex;
+        this.endIndex = endIndex;
     }
 
     @Override
@@ -42,11 +47,10 @@ public class AssignmentNode implements Node {
             STentry e = ST.lookup(leftAtom.getId());
             if (ft != null) {
                 ft.addLocalVar();
-            } else {
-                Label.addGlobalVar();
             }
 
             if (e == null) {
+                Label.addGlobalVar();
                 ST.insert(leftAtom.getId(), new AtomType(), _nesting, "");
                 e = ST.lookup(leftAtom.getId());
             } else {
@@ -58,9 +62,9 @@ public class AssignmentNode implements Node {
             offset = e.getOffset();
         }
         return errors;
+
     }
 
-    // TODO: check it out for this type
     @Override
     public Type typeCheck() {
         return rhr.typeCheck();
@@ -92,9 +96,31 @@ public class AssignmentNode implements Node {
     }
 
     @Override
+    public String printAST(String prefix) {
+        return prefix + "Assignment\n" + lhr.printAST(prefix + "  ") + assign.printAST(prefix + "  ")
+                + rhr.printAST(prefix + "  ");
+    }
+
+    @Override
     public String toPrint(String prefix) {
-        return prefix + "Assignment\n" + lhr.toPrint(prefix + "  ") + assign.toPrint(prefix + "  ")
-                + rhr.toPrint(prefix + "  ");
+        String str = prefix;
+        str += lhr.toPrint("") + " " + assign.toPrint("") + " " + rhr.toPrint("");
+        return str;
     }
 
+    public ExprListNode getLhr() {
+        return lhr;
+    }
+
+    public ExprListNode getRhr() {
+        return rhr;
+    }
+
+    public int getLhrIndex() {
+        return this.startIndex;
+    }
+
+    public int getRhrIndex() {
+        return this.endIndex;
+    }
 }
diff --git a/src/ast/nodes/AtomNode.java b/src/ast/nodes/AtomNode.java
index ea99808..204c3b3 100644
--- a/src/ast/nodes/AtomNode.java
+++ b/src/ast/nodes/AtomNode.java
@@ -13,6 +13,8 @@ import semanticanalysis.SymbolTable;
  */
 public class AtomNode implements Node {
 
+    private String prefix;
+    private String suffix;
     protected String val;
     protected TestlistCompNode exprlist;
 
@@ -21,8 +23,14 @@ public class AtomNode implements Node {
     protected int offset;
 
     public AtomNode(String val, Node exprlist) {
+        this(val, exprlist, "", "");
+    }
+
+    public AtomNode(String val, Node exprlist, String prefix, String suffix) {
         this.val = val;
         this.exprlist = (TestlistCompNode) exprlist;
+        this.prefix = prefix;
+        this.suffix = suffix;
     }
 
     /**
@@ -33,6 +41,10 @@ public class AtomNode implements Node {
         return val;
     }
 
+    public void setId(String id) {
+        this.val = id;
+    }
+
     @Override
     public ArrayList<SemanticError> checkSemantics(SymbolTable ST, int _nesting, FunctionType ft) {
         ArrayList<SemanticError> errors = new ArrayList<>();
@@ -134,13 +146,25 @@ public class AtomNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
-        // FIXME: can be a testlist_comp with two expr and two atoms
+    public String printAST(String prefix) {
         if (val != null) {
-            return prefix + "Atom(" + val + ")\n";
+            return prefix + "AtomNode: " + val + "\n";
+        } else {
+            return prefix + "AtomNode\n" + exprlist.printAST(prefix + "  ");
         }
+    }
 
-        return prefix + "Atom(null)\n";
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix + this.prefix;
+
+        if (val != null) {
+            str += val;
+        } else {
+            str += this.exprlist.toPrint("");
+        }
 
+        str += this.suffix;
+        return str;
     }
 }
diff --git a/src/ast/nodes/AugassignNode.java b/src/ast/nodes/AugassignNode.java
index 906dcd7..8b07c67 100644
--- a/src/ast/nodes/AugassignNode.java
+++ b/src/ast/nodes/AugassignNode.java
@@ -36,7 +36,13 @@ public class AugassignNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         return prefix + "Augassign(" + val + ")\n";
     }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix + val.toString();
+        return str;
+    }
 }
diff --git a/src/ast/nodes/BlockNode.java b/src/ast/nodes/BlockNode.java
index 2ea0a47..86a0db6 100644
--- a/src/ast/nodes/BlockNode.java
+++ b/src/ast/nodes/BlockNode.java
@@ -43,15 +43,25 @@ public class BlockNode extends RootNode {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "Block\n";
 
         prefix += "  ";
         for (Node child : childs) {
-            str += child.toPrint(prefix);
+            str += child.printAST(prefix);
         }
 
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+        String str = "";
+
+        for (Node child : childs) {
+            str += child.toPrint(prefix);
+        }
+
+        return str;
+    }
 }
diff --git a/src/ast/nodes/CompForNode.java b/src/ast/nodes/CompForNode.java
index 1f51af5..df86a74 100644
--- a/src/ast/nodes/CompForNode.java
+++ b/src/ast/nodes/CompForNode.java
@@ -12,12 +12,12 @@ import semanticanalysis.SymbolTable;
 public class CompForNode implements Node {
 
     protected ExprListNode exprlist;
-    protected ExprNode single_expr;
+    protected ExprNode expr;
     protected CompIterNode comp_iter;
 
-    public CompForNode(Node exprlist, Node single_expr, Node comp_iter) {
+    public CompForNode(Node exprlist, Node expr, Node comp_iter) {
         this.exprlist = (ExprListNode) exprlist;
-        this.single_expr = (ExprNode) single_expr;
+        this.expr = (ExprNode) expr;
         this.comp_iter = (CompIterNode) comp_iter;
     }
 
@@ -26,7 +26,7 @@ public class CompForNode implements Node {
         ArrayList<SemanticError> errors = new ArrayList<>();
 
         errors.addAll(exprlist.checkSemantics(ST, _nesting, ft));
-        errors.addAll(single_expr.checkSemantics(ST, _nesting, ft));
+        errors.addAll(expr.checkSemantics(ST, _nesting, ft));
         if (comp_iter != null) {
             errors.addAll(comp_iter.checkSemantics(ST, _nesting, ft));
         }
@@ -47,13 +47,27 @@ public class CompForNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "CompForNode\n";
 
         prefix += "  ";
-        str += exprlist.toPrint(prefix);
-        str += single_expr.toPrint(prefix);
-        str += comp_iter.toPrint(prefix);
+        str += exprlist.printAST(prefix);
+        str += expr.printAST(prefix);
+        str += comp_iter.printAST(prefix);
+        return str;
+    }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix + "for";
+
+        for (int i = 0; i < exprlist.getExprs().size(); i++) {
+            str += exprlist.getExprs().get(i).toPrint("");
+        }
+
+        str += " in " + expr.toPrint("") + "\n";
+        str += comp_iter.toPrint(prefix + "\t");
+
         return str;
     }
 
diff --git a/src/ast/nodes/CompIterNode.java b/src/ast/nodes/CompIterNode.java
index de443fa..4f09ec5 100644
--- a/src/ast/nodes/CompIterNode.java
+++ b/src/ast/nodes/CompIterNode.java
@@ -40,10 +40,17 @@ public class CompIterNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "CompIterNode\n";
 
         prefix += "  ";
+        str += comp_for.printAST(prefix);
+        return str;
+    }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = "";
         str += comp_for.toPrint(prefix);
         return str;
     }
diff --git a/src/ast/nodes/CompOpNode.java b/src/ast/nodes/CompOpNode.java
index 7997052..761a49c 100644
--- a/src/ast/nodes/CompOpNode.java
+++ b/src/ast/nodes/CompOpNode.java
@@ -27,7 +27,6 @@ public class CompOpNode implements Node {
         return new ArrayList<>();
     }
 
-    // TODO: it should be boolean, right?
     @Override
     public Type typeCheck() {
         return new VoidType();
@@ -42,7 +41,12 @@ public class CompOpNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         return prefix + "CompOpNode(" + op + ")\n";
     }
+
+    @Override
+    public String toPrint(String prefix) {
+        return prefix + op;
+    }
 }
diff --git a/src/ast/nodes/CompoundNode.java b/src/ast/nodes/CompoundNode.java
index f0f63f2..ad3420d 100644
--- a/src/ast/nodes/CompoundNode.java
+++ b/src/ast/nodes/CompoundNode.java
@@ -22,6 +22,7 @@ public class CompoundNode implements Node {
         this.whileStmt = whileStmt;
     }
 
+
     @Override
     public ArrayList<SemanticError> checkSemantics(SymbolTable ST, int _nesting, FunctionType ft) {
         ArrayList<SemanticError> errors = new ArrayList<>();
@@ -72,12 +73,35 @@ public class CompoundNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "CompoundNode\n";
 
         prefix += "  ";
 
         if (ifNode != null) {
+            str += ifNode.printAST(prefix);
+        }
+
+        if (funcDef != null) {
+            str += funcDef.printAST(prefix);
+        }
+
+        if (forStmt != null) {
+            str += forStmt.printAST(prefix);
+        }
+
+        if (whileStmt != null) {
+            str += whileStmt.printAST(prefix);
+        }
+
+        return str;
+    }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = "";
+
+        if (ifNode != null) {
             str += ifNode.toPrint(prefix);
         }
 
@@ -94,6 +118,22 @@ public class CompoundNode implements Node {
         }
 
         return str;
+
     }
 
+    public Node getForStmt() {
+        return forStmt;
+    }
+
+    public Node getFuncDef() {
+        return funcDef;
+    }
+
+    public Node getIfNode() {
+        return ifNode;
+    }
+
+    public Node getWhileStmt() {
+        return whileStmt;
+    }
 }
diff --git a/src/ast/nodes/DottedNameNode.java b/src/ast/nodes/DottedNameNode.java
index 827945a..ce00c00 100644
--- a/src/ast/nodes/DottedNameNode.java
+++ b/src/ast/nodes/DottedNameNode.java
@@ -44,7 +44,7 @@ public class DottedNameNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "DottedName\n";
 
         prefix += "  ";
@@ -55,4 +55,15 @@ public class DottedNameNode implements Node {
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix + names.get(0);
+        
+        for (int i = 1; i < names.size(); ++i) {
+            str += "." + names.get(i);
+        }
+
+        return str;
+    }
+
 }
diff --git a/src/ast/nodes/ExprListNode.java b/src/ast/nodes/ExprListNode.java
index a3ac237..44feb6b 100644
--- a/src/ast/nodes/ExprListNode.java
+++ b/src/ast/nodes/ExprListNode.java
@@ -2,6 +2,8 @@ package ast.nodes;
 
 import ast.types.*;
 import java.util.ArrayList;
+import java.util.List;
+
 import semanticanalysis.SemanticError;
 import semanticanalysis.SymbolTable;
 
@@ -10,9 +12,9 @@ import semanticanalysis.SymbolTable;
  */
 public class ExprListNode implements Node {
 
-    private final ArrayList<Node> exprs;
+    private final List<Node> exprs;
 
-    public ExprListNode(ArrayList<Node> exprs) {
+    public ExprListNode(List<Node> exprs) {
         this.exprs = exprs;
     }
 
@@ -59,15 +61,31 @@ public class ExprListNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "ExprList\n";
 
         prefix += "  ";
         for (var param : exprs) {
-            str += param.toPrint(prefix);
+            str += param.printAST(prefix);
         }
 
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+
+        String str = prefix + exprs.get(0).toPrint("");
+
+        for (int i = 1; i < exprs.size(); ++i) {
+            str += ", " + exprs.get(i).toPrint("");
+        }
+
+        return str;
+    }
+
+    public List<Node> getExprs() {
+        return exprs;
+    }
+
 }
diff --git a/src/ast/nodes/ExprNode.java b/src/ast/nodes/ExprNode.java
index 3d05eea..a750768 100644
--- a/src/ast/nodes/ExprNode.java
+++ b/src/ast/nodes/ExprNode.java
@@ -43,6 +43,10 @@ public class ExprNode implements Node {
         return this.exprs.get(i);
     }
 
+    public void setExpr(int i, Node expr) {
+        this.exprs.set(i, expr);
+    }
+
     /**
      * Returns the identifier of the `AtomNode` if it's not `null`, otherwise
      * returns `null`.
@@ -55,6 +59,26 @@ public class ExprNode implements Node {
         }
     }
 
+    public Node getCompOp() {
+        return compOp;
+    }
+
+    public String getOp() {
+        return op;
+    }
+
+    public ArrayList<Node> getTrailers() {
+        return trailers;
+    }
+
+    public ArrayList<Node> getExprs() {
+        return exprs;
+    }
+
+    public AtomNode getAtom() {
+        return atom;
+    }
+
     @Override
     public ArrayList<SemanticError> checkSemantics(SymbolTable ST, int _nesting, FunctionType ft) {
         ArrayList<SemanticError> errors = new ArrayList<>();
@@ -69,7 +93,6 @@ public class ExprNode implements Node {
                 TrailerNode trailer = (TrailerNode) trailers.get(0);
                 String funName = atom.getId();
 
-                // TODO: it isnt a function, it could be a variable
                 STentry fun = ST.lookup(funName);
 
                 if (fun != null && !(fun.getType() instanceof ImportType)) {
@@ -108,7 +131,6 @@ public class ExprNode implements Node {
         return errors;
     }
 
-    // FIXME: type for the expr
     @Override
     public Type typeCheck() {
         if (this.atom != null) {
@@ -156,8 +178,8 @@ public class ExprNode implements Node {
                 case "+":
                 case "-":
                 case "*":
-                // In real Python `/` is a float division but we'll consider the
-                // int division here below.
+                    // In real Python `/` is a float division but we'll consider the
+                    // int division here below.
                 case "/":
                     return intOpCodeGen(exprs.get(0), exprs.get(1), op);
                 case "%":
@@ -211,31 +233,55 @@ public class ExprNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "Expr\n";
 
         prefix += "  ";
         if (atom != null) {
-            str += atom.toPrint(prefix);
-        }
-        if (compOp != null) {
-            str += compOp.toPrint(prefix);
-        }
+            str += atom.printAST(prefix);
 
-        if (exprs != null) {
-            for (var expr : exprs) {
-                str += expr.toPrint(prefix);
+            for (var trailer : trailers) {
+                str += trailer.printAST(prefix);
+            }
+        } else {
+            if (compOp != null) {
+                str += compOp.printAST(prefix);
             }
-        }
 
-        if (trailers != null) {
-            for (var trailer : trailers) {
-                str += trailer.toPrint(prefix);
+            if (exprs != null) {
+                for (var expr : exprs) {
+                    str += expr.printAST(prefix);
+                }
+            }
+
+            if (op != null) {
+                str += prefix + "Op(" + op + ")\n";
             }
         }
 
-        if (op != null) {
-            str += prefix + "Op(" + op + ")\n";
+        return str;
+    }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix;
+
+        if (atom != null) {
+            str += atom.toPrint("");
+            for (var trailer : trailers) {
+                str += trailer.toPrint("");
+            }
+        } else {
+            if (op == "+" || op == "-" || op == "~") {
+                str += prefix + op + exprs.get(0).toPrint("");
+            } else {
+                str += prefix + exprs.get(0).toPrint("") + " ";
+                if (compOp != null) {
+                    str += compOp.toPrint("") + " " + exprs.get(1).toPrint("");
+                } else {
+                    str += op + " " + exprs.get(1).toPrint("");
+                }
+            }
         }
 
         return str;
diff --git a/src/ast/nodes/ForStmtNode.java b/src/ast/nodes/ForStmtNode.java
index 507d314..629d67e 100644
--- a/src/ast/nodes/ForStmtNode.java
+++ b/src/ast/nodes/ForStmtNode.java
@@ -73,8 +73,23 @@ public class ForStmtNode implements Node {
     }
 
     @Override
+    public String printAST(String prefix) {
+        return prefix + "For\n" + exprList.printAST(prefix + "  ") + block.printAST(prefix + "  ");
+    }
+
+    @Override
     public String toPrint(String prefix) {
-        return prefix + "For\n" + exprList.toPrint(prefix + "  ") + block.toPrint(prefix + "  ");
+        String str = prefix + "for ";
+        str += exprList.toPrint("") + ":\n";
+        str += block.toPrint(prefix + "\t");
+        return str;
     }
 
+    public Node getBlock() {
+        return block;
+    }
+
+    public Node getExprList() {
+        return exprList;
+    }
 }
diff --git a/src/ast/nodes/FuncdefNode.java b/src/ast/nodes/FuncdefNode.java
index c2671b2..e99f3a3 100644
--- a/src/ast/nodes/FuncdefNode.java
+++ b/src/ast/nodes/FuncdefNode.java
@@ -22,7 +22,11 @@ public class FuncdefNode implements Node {
 
     public FuncdefNode(TerminalNode name, Node paramlist, Node block) {
         this.name = name;
-        this.paramlist = paramlist;
+        if (paramlist != null) {
+            this.paramlist = paramlist;
+        } else {
+            this.paramlist = new ParamlistNode(new ArrayList<Node>());
+        }
         this.block = block;
         this.funLabel = Label.newFun("FUN");
     }
@@ -62,7 +66,6 @@ public class FuncdefNode implements Node {
         return errors;
     }
 
-    // FIXME: this type must be the same of the return stmt variable
     @Override
     public Type typeCheck() {
         return new VoidType();
@@ -84,17 +87,29 @@ public class FuncdefNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "Funcdef(" + name + ")\n";
 
         prefix += "  ";
 
         if (paramlist != null) {
-            str += paramlist.toPrint(prefix);
+            str += paramlist.printAST(prefix);
         }
 
-        str += block.toPrint(prefix);
+        str += block.printAST(prefix);
+
+        return str;
+    }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = "\n" + prefix + "def " + this.name + '(';
+        if (paramlist != null) {
+            str += paramlist.toPrint("");
+        }
 
+        str += "):\n";
+        str += block.toPrint(prefix + "\t") + "\n";
         return str;
     }
 
diff --git a/src/ast/nodes/IfNode.java b/src/ast/nodes/IfNode.java
index ca42b7a..fd878f8 100644
--- a/src/ast/nodes/IfNode.java
+++ b/src/ast/nodes/IfNode.java
@@ -34,7 +34,6 @@ public class IfNode implements Node {
         return errors;
     }
 
-    // FIXME: fix the if statement
     @Override
     public Type typeCheck() {
         if (guard.typeCheck() instanceof BoolType) {
@@ -78,14 +77,26 @@ public class IfNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
-        String str = prefix + "If\n" + guard.toPrint(prefix + "  ") + thenBranch.toPrint(prefix + "  ");
+    public String printAST(String prefix) {
+        String str = prefix + "If\n" + guard.printAST(prefix + "  ") + thenBranch.printAST(prefix + "  ");
 
         if (elseBranch != null) {
-            str += elseBranch.toPrint(prefix + "  ");
+            str += elseBranch.printAST(prefix + "  ");
         }
 
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix + "if ";
+        str += guard.toPrint("") + ":\n";
+        str += thenBranch.toPrint(prefix + "\t") + "\n";
+        if (elseBranch != null) {
+            str += prefix + "else:\n";
+            str += elseBranch.toPrint(prefix + "\t") + "\n";
+        }
+        return str;
+    }
+
 }
diff --git a/src/ast/nodes/ImportNode.java b/src/ast/nodes/ImportNode.java
index 97e3404..1f797f8 100644
--- a/src/ast/nodes/ImportNode.java
+++ b/src/ast/nodes/ImportNode.java
@@ -58,14 +58,14 @@ public class ImportNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "Import\n";
 
         prefix += "  ";
         if (isFrom) {
-            str += prefix + "  From\n" + dottedName.toPrint(prefix + "  ");
+            str += prefix + "  From\n" + dottedName.printAST(prefix + "  ");
         } else {
-            str += dottedName.toPrint(prefix);
+            str += dottedName.printAST(prefix);
         }
 
         if (importAs) {
@@ -88,4 +88,27 @@ public class ImportNode implements Node {
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix;
+
+        if (isFrom) {
+            str += "from " + dottedName.toPrint("") + " import ";
+            if (importAll) {
+                str += "*";
+            } else {
+                str += names.get(0);
+                for (int i = 1; i < names.size(); ++i) {
+                    str += ", " + names.get(i);
+                }
+            }
+        } else {
+            str += "import " + dottedName.toPrint("");
+            if (importAs) {
+                str += " as " + names.get(0);
+            }
+        }
+        return str;
+    }
+
 }
diff --git a/src/ast/nodes/Node.java b/src/ast/nodes/Node.java
index 0e58093..5bf5a33 100644
--- a/src/ast/nodes/Node.java
+++ b/src/ast/nodes/Node.java
@@ -105,5 +105,11 @@ public interface Node {
      * Returns a string for a given node with a prefix. It used when an AST
      * wants to be visualized on screen.
      */
+    String printAST(String prefix);
+
+    /**
+     * Returns a string for a given node with a prefix. It used when the code
+     * wants to be visualized on screen.
+     */
     String toPrint(String prefix);
 }
diff --git a/src/ast/nodes/ParamdefNode.java b/src/ast/nodes/ParamdefNode.java
index 4bba772..b7786eb 100644
--- a/src/ast/nodes/ParamdefNode.java
+++ b/src/ast/nodes/ParamdefNode.java
@@ -29,14 +29,18 @@ public class ParamdefNode extends AtomNode {
         return errors;
     }
 
-    // FIXME: it should returns the param' type
     @Override
     public Type typeCheck() {
         return new VoidType();
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         return prefix + "Paramdef(" + val + ")\n";
     }
+
+    @Override
+    public String toPrint(String prefix) {
+        return prefix + val;
+    }
 }
diff --git a/src/ast/nodes/ParamlistNode.java b/src/ast/nodes/ParamlistNode.java
index 9e8d75c..2cb6c89 100644
--- a/src/ast/nodes/ParamlistNode.java
+++ b/src/ast/nodes/ParamlistNode.java
@@ -42,15 +42,29 @@ public class ParamlistNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "Paramlist\n";
 
         prefix += "  ";
-        for (var param : params) {
-            str += param.toPrint(prefix);
+        for (Node param : params) {
+            str += param.printAST(prefix);
         }
 
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix;
+
+        for (int i = 0; i < params.size(); i++) {
+            str += params.get(i).toPrint("");
+            if (i != params.size() - 1) {
+                str += ", ";
+            }
+        }
+        
+        return str;
+    }
+
 }
diff --git a/src/ast/nodes/ReturnStmtNode.java b/src/ast/nodes/ReturnStmtNode.java
index 15c8d69..1264fa9 100644
--- a/src/ast/nodes/ReturnStmtNode.java
+++ b/src/ast/nodes/ReturnStmtNode.java
@@ -63,15 +63,24 @@ public class ReturnStmtNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "ReturnStmt\n";
 
         prefix += "  ";
         if (this.exprList != null) {
-            str += this.exprList.toPrint(prefix);
+            str += this.exprList.printAST(prefix);
         }
 
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix + "return ";
+        if (this.exprList != null) {
+            str += this.exprList.toPrint("");
+        }
+        return str + "\n";
+    }
+
 }
diff --git a/src/ast/nodes/RootNode.java b/src/ast/nodes/RootNode.java
index 2c99164..ebb786a 100644
--- a/src/ast/nodes/RootNode.java
+++ b/src/ast/nodes/RootNode.java
@@ -54,7 +54,7 @@ public class RootNode implements Node {
             str += child.codeGeneration();
         }
 
-        for (int i = 0; i < Label.getGlobalVarNum(); i++) {
+        for (int i = 0; i < Label.getGlobalVarNum() - 1; i++) {
             str += "pop\n";
         }
 
@@ -62,16 +62,32 @@ public class RootNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = "Root\n";
 
         prefix += "  ";
 
         for (Node child : childs) {
-            str += child.toPrint(prefix);
+            str += child.printAST(prefix);
         }
 
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix;
+        for (Node child : childs) {
+            str += child.toPrint("");
+        }
+        return str;
+    }
+
+    public ArrayList<Node> getChilds() {
+        return childs;
+    }
+
+    public Node getChild(int i) {
+        return childs.get(i);
+    }
 }
diff --git a/src/ast/nodes/SimpleStmtNode.java b/src/ast/nodes/SimpleStmtNode.java
index 31a935f..3952242 100644
--- a/src/ast/nodes/SimpleStmtNode.java
+++ b/src/ast/nodes/SimpleStmtNode.java
@@ -73,28 +73,56 @@ public class SimpleStmtNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "SimpleStmt\n";
 
         prefix += "  ";
 
         if (assignment != null) {
-            str += assignment.toPrint(prefix);
+            str += assignment.printAST(prefix);
+        } else if (expr != null) {
+            str += expr.printAST(prefix);
+        } else if (returnStmt != null) {
+            str += returnStmt.printAST(prefix);
+        } else if (importStmt != null) {
+            str += importStmt.printAST(prefix);
         }
 
-        if (expr != null) {
-            str += expr.toPrint(prefix);
-        }
+        return str;
+    }
 
-        if (returnStmt != null) {
-            str += returnStmt.toPrint(prefix);
-        }
+    @Override
+    public String toPrint(String prefix) {
+        
+        String str = prefix;
 
-        if (importStmt != null) {
-            str += importStmt.toPrint(prefix);
+        if (assignment != null) {
+            str += assignment.toPrint("");
+        } else if (expr != null) {
+            str += expr.toPrint("");
+        } else if (returnStmt != null) {
+            str += returnStmt.toPrint("");
+        } else if (importStmt != null) {
+            str += importStmt.toPrint("");
         }
 
         return str;
     }
 
+    public Node getImportStmt() {
+        return importStmt;
+    }
+
+    public Node getReturnStmt() {
+        return returnStmt;
+    }
+    
+    public Node getAssignment() {
+        return assignment;
+    }
+
+    public Node getExpr() {
+        return expr;
+    }
+
 }
diff --git a/src/ast/nodes/SimpleStmtsNode.java b/src/ast/nodes/SimpleStmtsNode.java
index 8e5c924..e7a5557 100644
--- a/src/ast/nodes/SimpleStmtsNode.java
+++ b/src/ast/nodes/SimpleStmtsNode.java
@@ -44,16 +44,31 @@ public class SimpleStmtsNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "SimpleStmts\n";
 
         prefix += "  ";
 
         for (Node stmt : stmts) {
-            str += stmt.toPrint(prefix);
+            str += stmt.printAST(prefix);
+        }
+
+        return str;
+    }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = "";
+        
+        for (Node stmt : stmts) {
+            str += stmt.toPrint(prefix) + "\n";
         }
 
         return str;
     }
 
+    public ArrayList<Node> getStmts() {
+        return stmts;
+    }
+
 }
diff --git a/src/ast/nodes/TestlistCompNode.java b/src/ast/nodes/TestlistCompNode.java
index e55f854..024d742 100644
--- a/src/ast/nodes/TestlistCompNode.java
+++ b/src/ast/nodes/TestlistCompNode.java
@@ -72,12 +72,29 @@ public class TestlistCompNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "Testlist_comp\n";
 
         prefix += "  ";
         for (var param : exprs) {
-            str += param.toPrint(prefix);
+            str += param.printAST(prefix);
+        }
+
+        return str;
+    }
+
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix;
+
+        str += exprs.get(0).toPrint("");
+
+        if (comp != null) {
+            str += comp.toPrint("");
+        } else {
+            for (int i = 1; i < exprs.size(); i++) {
+                str += ", " + exprs.get(i).toPrint("");
+            }
         }
 
         return str;
diff --git a/src/ast/nodes/TrailerNode.java b/src/ast/nodes/TrailerNode.java
index d70e962..27eb166 100644
--- a/src/ast/nodes/TrailerNode.java
+++ b/src/ast/nodes/TrailerNode.java
@@ -12,17 +12,19 @@ import org.antlr.v4.runtime.tree.TerminalNode;
  */
 public class TrailerNode implements Node {
 
+    private String prefix;
+    private String suffix;
     private final Node arglist;
     private final ArrayList<Node> exprs;
     private final TerminalNode methodCall;
-    private final boolean isParenthesis;
     private final boolean isEmpty;
 
-    public TrailerNode(Node arglist, ArrayList<Node> exprs, TerminalNode methodCall, boolean isParenthesis) {
+    public TrailerNode(Node arglist, ArrayList<Node> exprs, TerminalNode methodCall, String prefix, String suffix){
         this.arglist = arglist;
         this.exprs = exprs;
         this.methodCall = methodCall;
-        this.isParenthesis = isParenthesis;
+        this.prefix = prefix;
+        this.suffix = suffix;
 
         this.isEmpty = (this.arglist == null && this.exprs.isEmpty() && this.methodCall == null);
     }
@@ -50,10 +52,6 @@ public class TrailerNode implements Node {
         return ((ArglistNode) arglist).getArgumentNumber();
     }
 
-    public boolean isParenthesis() {
-        return this.isParenthesis;
-    }
-
     @Override
     public Type typeCheck() {
         return new VoidType();
@@ -68,17 +66,17 @@ public class TrailerNode implements Node {
     }
 
     @Override
-    public String toPrint(String prefix) {
+    public String printAST(String prefix) {
         String str = prefix + "TrailerNode\n";
 
         prefix += "  ";
 
         if (arglist != null) {
-            str += arglist.toPrint(prefix);
+            str += arglist.printAST(prefix);
         }
 
         for (var expr : exprs) {
-            str += expr.toPrint(prefix);
+            str += expr.printAST(prefix);
         }
 
         if (methodCall != null) {
@@ -92,4 +90,25 @@ public class TrailerNode implements Node {
         return str;
     }
 
+    @Override
+    public String toPrint(String prefix) {
+        String str = prefix + this.prefix;
+
+        if (arglist != null) {
+            str += arglist.toPrint("");
+        } else if (methodCall != null) {
+            str += methodCall.getText();
+        } else {
+            for (var expr : exprs) {
+                str += expr.toPrint("");
+                if (exprs.indexOf(expr) != exprs.size() - 1) {
+                    str += ", ";
+                }
+            }
+        }
+
+        str += this.suffix;
+        return str;
+    }
+
 }
diff --git a/src/ast/nodes/WhileStmtNode.java b/src/ast/nodes/WhileStmtNode.java
index d371046..43bf5ff 100644
--- a/src/ast/nodes/WhileStmtNode.java
+++ b/src/ast/nodes/WhileStmtNode.java
@@ -33,14 +33,33 @@ public class WhileStmtNode implements Node {
         return new VoidType();
     }
 
-    // TODO: add cgen per while (but it's not requested from the exercise)
+    /**
+     * NOTE: It is not a part for this project.
+     */
     @Override
     public String codeGeneration() {
         return "";
     }
 
     @Override
+    public String printAST(String prefix) {
+        return prefix + "While\n" + expr.printAST(prefix + "  ") + block.printAST(prefix + "  ");
+    }
+
+    @Override
     public String toPrint(String prefix) {
-        return prefix + "While\n" + expr.toPrint(prefix + "  ") + block.toPrint(prefix + "  ");
+        String str = prefix + "while ";
+        str += expr.toPrint("") + ":\n";
+        str += block.toPrint(prefix + "\t") + "\n";
+        return str;
+    }
+
+    public Node getBlock() {
+        return block;
+    }
+
+    public Node getExpr() {
+        return expr;
     }
+
 }
diff --git a/src/ast/types/Type.java b/src/ast/types/Type.java
index c5b8312..fc9c672 100644
--- a/src/ast/types/Type.java
+++ b/src/ast/types/Type.java
@@ -15,6 +15,11 @@ public class Type implements Node {
     }
 
     @Override
+    public String printAST(String s) {
+        return s;
+    }
+
+    @Override
     public String toPrint(String s) {
         return s;
     }
diff --git a/src/codegen/Label.java b/src/codegen/Label.java
index f0a4c6e..a301ede 100644
--- a/src/codegen/Label.java
+++ b/src/codegen/Label.java
@@ -6,6 +6,7 @@ public class Label {
     private static int labelCounter = 0;
 	private static int globalVarNum = 0;
     private static int functionLabelCounter = 0;
+    private static int varDefCount = 0;
 
     public static void addFunDef(String s) {
         funDef += s;
@@ -38,4 +39,8 @@ public class Label {
         return base + (functionLabelCounter++);
     }
 
+
+    public static String newVar() {
+        return "_tmp" + (varDefCount++);
+    }
 }
diff --git a/src/parser/Python3Lexer.g4 b/src/parser/Python3Lexer.g4
index 8d53301..182f3ef 100644
--- a/src/parser/Python3Lexer.g4
+++ b/src/parser/Python3Lexer.g4
@@ -128,4 +128,4 @@ fragment SPACES: [ \t]+;
 
 fragment COMMENT: '#' ~[\r\n\f]*;
 
-fragment LINE_JOINING: '\\' SPACES? ( '\r'? '\n' | '\r' | '\f');
\ No newline at end of file
+fragment LINE_JOINING: '\\' SPACES? ( '\r'? '\n' | '\r' | '\f');
diff --git a/src/semanticanalysis/Share.java b/src/semanticanalysis/Share.java
index a05c4c5..5bb8099 100644
--- a/src/semanticanalysis/Share.java
+++ b/src/semanticanalysis/Share.java
@@ -1,6 +1,10 @@
 package semanticanalysis;
 
 import java.io.*;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
 import java.util.*;
 
 public class Share {
@@ -55,4 +59,16 @@ public class Share {
         }
         return content.toString();
     }
+
+    public static void saveFile(String fileName, String content) {
+        try {
+            Path file = Paths.get(fileName);
+            if (!Files.exists(file)) {
+                Files.createFile(file);
+            }
+            Files.write(file, content.getBytes(), StandardOpenOption.TRUNCATE_EXISTING);
+        } catch (Exception e) {
+            e.printStackTrace();
+        }
+    }
 }
diff --git a/src/svm/ExecuteVM.java b/src/svm/ExecuteVM.java
index d3dec33..d2e42b5 100755
--- a/src/svm/ExecuteVM.java
+++ b/src/svm/ExecuteVM.java
@@ -1,7 +1,5 @@
 package svm;
 
-import java.util.concurrent.TimeUnit;
-
 public class ExecuteVM {
     
     public static final int CODESIZE = 1000 ;
diff --git a/src/svm/SVMBaseListener.java b/src/svm/SVMBaseListener.java
index 19d7c60..c12f809 100644
--- a/src/svm/SVMBaseListener.java
+++ b/src/svm/SVMBaseListener.java
@@ -1,9 +1,6 @@
 // Generated from /home/gabri/Desktop/clp_project/src/svm/SVM.g4 by ANTLR 4.13.1
 package svm;
 
-import java.util.HashMap;
-
-
 import org.antlr.v4.runtime.ParserRuleContext;
 import org.antlr.v4.runtime.tree.ErrorNode;
 import org.antlr.v4.runtime.tree.TerminalNode;
diff --git a/src/svm/SVMBaseVisitor.java b/src/svm/SVMBaseVisitor.java
index ba371ce..3944d4a 100644
--- a/src/svm/SVMBaseVisitor.java
+++ b/src/svm/SVMBaseVisitor.java
@@ -1,8 +1,6 @@
 // Generated from /home/gabri/Desktop/clp_project/src/svm/SVM.g4 by ANTLR 4.13.1
 package svm;
 
-import java.util.HashMap;
-
 import org.antlr.v4.runtime.tree.AbstractParseTreeVisitor;
 
 /**
diff --git a/src/svm/SVMListener.java b/src/svm/SVMListener.java
index 770e56e..18d7174 100644
--- a/src/svm/SVMListener.java
+++ b/src/svm/SVMListener.java
@@ -1,8 +1,6 @@
 // Generated from /home/gabri/Desktop/clp_project/src/svm/SVM.g4 by ANTLR 4.13.1
 package svm;
 
-import java.util.HashMap;
-
 import org.antlr.v4.runtime.tree.ParseTreeListener;
 
 /**
diff --git a/src/svm/SVMVisitor.java b/src/svm/SVMVisitor.java
index d4e132d..0e4193e 100644
--- a/src/svm/SVMVisitor.java
+++ b/src/svm/SVMVisitor.java
@@ -1,8 +1,6 @@
 // Generated from /home/gabri/Desktop/clp_project/src/svm/SVM.g4 by ANTLR 4.13.1
 package svm;
 
-import java.util.HashMap;
-
 import org.antlr.v4.runtime.tree.ParseTreeVisitor;
 
 /**
diff --git a/test/2a.py b/test/2a.py
new file mode 100644
index 0000000..f45041a
--- /dev/null
+++ b/test/2a.py
@@ -0,0 +1,12 @@
+n = int(input())
+x = int(input())
+m = 1
+while n > 1:
+    tmp = 2 * x
+    m = m * tmp
+    n = n - 1
+c = 0
+for i in range(m):
+    g = 2 * m
+    c = c + i + g
+print(m + c)
diff --git a/test/2b.py b/test/2b.py
new file mode 100644
index 0000000..d6d0bb2
--- /dev/null
+++ b/test/2b.py
@@ -0,0 +1,9 @@
+n = 1
+x = 2
+y = 3
+m = 1
+while n + 2 < 2 * x - 3 * y + 50:
+    m = m + n
+    n = n + 1
+
+print(m)
diff --git a/test/2c.py b/test/2c.py
new file mode 100644
index 0000000..459cf8c
--- /dev/null
+++ b/test/2c.py
@@ -0,0 +1,13 @@
+n = [0,0,0,0,0]
+x = 3
+y = 7
+c = 0
+m = 1
+while y > 1:
+    tmp = 2 * x + y
+    m = m * tmp
+    y = y - 1
+for i in n:
+    g = 2 * m + i
+    c = c + i + g
+print(m + c)
diff --git a/test/2d.py b/test/2d.py
new file mode 100644
index 0000000..a418aab
--- /dev/null
+++ b/test/2d.py
@@ -0,0 +1,13 @@
+n = 2
+y = 7
+c = 0
+while y > 0:
+    g = 2 * n
+    c = y * n + g
+    y = y - 1
+
+y = 5
+for i in range(y):
+    g = 2 * y
+    c = c + i + g
+print(g)
diff --git a/test/2e.py b/test/2e.py
new file mode 100644
index 0000000..f7b111a
--- /dev/null
+++ b/test/2e.py
@@ -0,0 +1,17 @@
+n = 2
+y = 7
+c = 0
+while y > 0:
+    g = 2 * n
+    c = y * n + g
+    y = y - 1
+y = 10
+g = 1
+for i in range(y):
+    g = i * y
+    c = c + i + g
+while c > 0:
+    g = 4 * y
+    n = n + c + g
+    c = c - 1
+print(g)