1 files changed, 65 insertions, 19 deletions

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;