From fda65c7d930b6961c877e14486a6de494a6735ea Mon Sep 17 00:00:00 2001
From: nitrix <me@nitrix.one>
Date: Mon, 25 Nov 2019 14:45:35 +0100
Subject: [PATCH] implement return statement

---
 .../hsrm/compiler/Klang/ContextAnalysis.java  |  6 +++++
 .../compiler/Klang/visitors/EvalVisitor.java  | 25 +++++++++++++------
 .../Klang/visitors/PrettyPrintVisitor.java    |  8 ++++++
 .../hsrm/compiler/Klang/visitors/Visitor.java |  1 +
 4 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/src/main/java/de/hsrm/compiler/Klang/ContextAnalysis.java b/src/main/java/de/hsrm/compiler/Klang/ContextAnalysis.java
index 83fc6e4..8cee8ee 100644
--- a/src/main/java/de/hsrm/compiler/Klang/ContextAnalysis.java
+++ b/src/main/java/de/hsrm/compiler/Klang/ContextAnalysis.java
@@ -59,6 +59,12 @@ public class ContextAnalysis extends KlangBaseVisitor<Node> {
     return new VariableAssignment(name, expression);
   }
 
+  @Override
+  public Node visitReturn_statement(KlangParser.Return_statementContext ctx) {
+    Expression expression = (Expression) this.visit(ctx.expression()); 
+    return new ReturnStatement(expression);
+  }
+
   @Override
   public Node visitMultiplicationExpression(KlangParser.MultiplicationExpressionContext ctx) {
     Node left = this.visit(ctx.atom(0));
diff --git a/src/main/java/de/hsrm/compiler/Klang/visitors/EvalVisitor.java b/src/main/java/de/hsrm/compiler/Klang/visitors/EvalVisitor.java
index 8b0bdda..300360f 100644
--- a/src/main/java/de/hsrm/compiler/Klang/visitors/EvalVisitor.java
+++ b/src/main/java/de/hsrm/compiler/Klang/visitors/EvalVisitor.java
@@ -64,15 +64,17 @@ public class EvalVisitor implements Visitor<Value> {
     // value is actually a type that we can use as boolean
     Value condition = e.cond.welcome(this);
 
+    Value result = null;
+
     if (condition.asInteger() != 0) {
-      e.then.welcome(this);
+      result = e.then.welcome(this);
     } else if (e.alt != null) {
-      e.alt.welcome(this);
+      result = e.alt.welcome(this);
     } else if (e.elif != null) {
-      e.elif.welcome(this);
+      result = e.elif.welcome(this);
     }
 
-    return null;
+    return result;
   }
 
   @Override
@@ -89,16 +91,23 @@ public class EvalVisitor implements Visitor<Value> {
   public Value visit(VariableAssignment e) {
     Value result = e.expression.welcome(this);
     this.env.put(e.name, result);
-    return result;
+    return null;
+  }
+
+  @Override
+  public Value visit(ReturnStatement e) {
+    return e.expression.welcome(this);
   }
 
   @Override
   public Value visit(Block e) {
-    Value result = null;
     for (var stmt : e.statements) {
-      result = stmt.welcome(this);
+      Value result = stmt.welcome(this);
+      if (result != null) {
+        return result;
+      }
     }
-    return result;
+    return null;
   }
 
   @Override
diff --git a/src/main/java/de/hsrm/compiler/Klang/visitors/PrettyPrintVisitor.java b/src/main/java/de/hsrm/compiler/Klang/visitors/PrettyPrintVisitor.java
index 508013b..fed3c01 100644
--- a/src/main/java/de/hsrm/compiler/Klang/visitors/PrettyPrintVisitor.java
+++ b/src/main/java/de/hsrm/compiler/Klang/visitors/PrettyPrintVisitor.java
@@ -134,6 +134,14 @@ public class PrettyPrintVisitor implements Visitor<Void> {
         return null;
     }
 
+    @Override
+    public Void visit(ReturnStatement e) {
+        ex.write("return ");
+        e.expression.welcome(this);
+        ex.write(";");
+        return null;
+    }
+
     @Override
     public Void visit(Block e) {
         ex.write("{");
diff --git a/src/main/java/de/hsrm/compiler/Klang/visitors/Visitor.java b/src/main/java/de/hsrm/compiler/Klang/visitors/Visitor.java
index 28cbbd5..456c8cd 100644
--- a/src/main/java/de/hsrm/compiler/Klang/visitors/Visitor.java
+++ b/src/main/java/de/hsrm/compiler/Klang/visitors/Visitor.java
@@ -16,6 +16,7 @@ public interface Visitor<R> {
     R visit(IfStatement e);
     R visit(PrintStatement e);
     R visit(VariableAssignment e);
+    R visit(ReturnStatement e);
     R visit(Block e);
     R visit(FunctionDefinition e);
     R visit(FunctionCall e);