Handle function returns

Upsilon/Binder/Binder.cs

@@ -55,6 +55,8 @@ namespace Upsilon.Binder
                     return BindIfStatement((IfStatementSyntax) s);
                 case SyntaxKind.FunctionStatement:
                     return BindFunctionStatement((FunctionStatementSyntax) s);
+                case SyntaxKind.ReturnStatement:
+                    return BindReturnStatement((ReturnStatementSyntax) s);
             }
 
             throw new NotImplementedException(s.Kind.ToString());
@@ -338,5 +340,11 @@ namespace Upsilon.Binder
             }
         }
 
+        private BoundStatement BindReturnStatement(ReturnStatementSyntax e)
+        {
+            var expression = BindExpression(e.Expression);
+            return new BoundReturnStatement(expression);
+        }
+
     }
 }

Upsilon/Binder/BoundKind.cs

@@ -17,6 +17,7 @@ namespace Upsilon.Binder
         BoundIfStatement,
         BoundElseStatement,
         BoundFunctionStatement,
-        BoundPromise
+        BoundPromise,
+        BoundReturnStatement
     }
 }

Upsilon/Binder/BoundStatements/BoundReturnStatement.cs

@@ -0,0 +1,13 @@
+namespace Upsilon.Binder
+{
+    public class BoundReturnStatement : BoundStatement
+    {
+        public BoundReturnStatement(BoundExpression expression)
+        {
+            Expression = expression;
+        }
+
+        public BoundExpression Expression { get; }
+        public override BoundKind Kind => BoundKind.BoundReturnStatement;
+    }
+}

Upsilon/Evaluator/Evaluator.cs

@@ -10,7 +10,8 @@ namespace Upsilon.Evaluator
     {
         private readonly Diagnostics _diagnostics;
         private LuaType _value;
-        public EvaluationScope Scope { get; private set; }
+        public EvaluationScope Scope { get; }
+        private bool HasReturned { get; set; }
 
         internal Evaluator(Diagnostics diagnostics, Dictionary<VariableSymbol, LuaType> variables)
         {
@@ -18,14 +19,51 @@ namespace Upsilon.Evaluator
             Scope       = new EvaluationScope(variables);
         }
 
+        private Evaluator(Diagnostics diagnostics, EvaluationScope parentScope)
+        {
+            _diagnostics = diagnostics;
+            Scope        = new EvaluationScope(parentScope);
+        }
+
         public LuaType Evaluate(BoundScript e)
         {
-            EvaluateStatement(e.Statement);
+            Evaluate(e.Statement);
+            return _value;
+        }
+
+        private LuaType Evaluate(BoundNode b)
+        {
+            switch (b.Kind)
+            {
+                case BoundKind.BoundScript:
+                    EvaluateStatement(((BoundScript)b).Statement);
+                    break;
+                case BoundKind.BoundLiteralExpression:
+                case BoundKind.BoundBinaryExpression:
+                case BoundKind.BoundUnaryExpression:
+                case BoundKind.VariableExpression:
+                case BoundKind.BoundFunctionCallExpression:
+                    _value = EvaluateExpression((BoundExpression) b);
+                    break;
+                case BoundKind.BoundAssignmentStatement:
+                case BoundKind.BoundExpressionStatement:
+                case BoundKind.BoundBlockStatement:
+                case BoundKind.BoundIfStatement:
+                case BoundKind.BoundElseStatement:
+                case BoundKind.BoundFunctionStatement:
+                case BoundKind.BoundPromise:
+                    EvaluateStatement((BoundStatement) b);
+                    break;
+                default:
+                    throw new ArgumentOutOfRangeException();
+            }
             return _value;
         }
 
         private void EvaluateStatement(BoundStatement e)
         {
+            if (HasReturned)
+                return;
             switch (e.Kind)
             {
                 case BoundKind.BoundAssignmentStatement:
@@ -43,6 +81,9 @@ namespace Upsilon.Evaluator
                 case BoundKind.BoundPromise:
                     EvaluateUnboundFunctionStatement((UnboundFunctionStatement) e);
                     break;
+                case BoundKind.BoundReturnStatement:
+                    EvaluateReturnStatement((BoundReturnStatement) e);
+                    break;
                 default:
                     EvaluateExpressionStatement((BoundExpressionStatement) e);
                     break;
@@ -140,28 +181,30 @@ namespace Upsilon.Evaluator
         {
             foreach (var boundStatement in boundBlockStatement.Statements)
             {
+                if (HasReturned)
+                    return;
                 EvaluateStatement(boundStatement);
             }
         }
 
         private void EvaluateBoundIfStatement(BoundIfStatement boundBlockStatement)
         {
-            Scope = new EvaluationScope(Scope);
+            var innerEvaluator = new Evaluator(_diagnostics, Scope);
-            var condition = EvaluateExpression(boundBlockStatement.Condition.Expression);
+            var condition = innerEvaluator.EvaluateExpression(boundBlockStatement.Condition.Expression);
             if ((LuaBoolean) condition)
             {
-                EvaluateBoundBlockStatement(boundBlockStatement.Block);
+                innerEvaluator.EvaluateStatement(boundBlockStatement.Block);
             }
             else if (boundBlockStatement.NextElseIf != null)
             {
-                EvaluateBoundIfStatement(boundBlockStatement.NextElseIf);
+                innerEvaluator.EvaluateStatement(boundBlockStatement.NextElseIf);
             }
             else if (boundBlockStatement.ElseStatement != null)
             {
-                EvaluateBoundBlockStatement(boundBlockStatement.ElseStatement.Block);
+                innerEvaluator.EvaluateStatement(boundBlockStatement.ElseStatement.Block);
             }
-
+            HasReturned = innerEvaluator.HasReturned;
-            Scope = Scope.ParentScope;
+            _value = innerEvaluator._value;
         }
 
         private void EvaluateBoundFunctionStatement(BoundFunctionStatement boundFunctionStatement)
@@ -193,18 +236,22 @@ namespace Upsilon.Evaluator
 
             var function = (LuaFunction) functionObj;
 
-            Scope = new EvaluationScope(Scope);
+            var innerEvaluator = new Evaluator(_diagnostics, Scope);
             for (var i = 0; i < function.Parameters.Length; i++)
             {
                 var parameterVariable = function.Parameters[i];
-                var parameterValue = EvaluateExpression(boundFunctionCallExpression.Parameters[i]);
+                var parameterValue = innerEvaluator.EvaluateExpression(boundFunctionCallExpression.Parameters[i]);
-                Scope.Set(parameterVariable, parameterValue);
+                innerEvaluator.Scope.Set(parameterVariable, parameterValue);
             }
-            EvaluateBoundBlockStatement(function.Block);
+            _value = innerEvaluator.Evaluate(function.Block);
-            Scope = Scope.ParentScope;
             return _value;
         }
 
+        private void EvaluateReturnStatement(BoundReturnStatement b)
+        {
+            _value = Evaluate(b.Expression);
+            HasReturned = true;
+        }
 
     }
 }

Upsilon/Parser/Parser.cs

@@ -79,6 +79,10 @@ namespace Upsilon.Parser
             {
                 return ParseFunctionStatement();
             }
+            if (Current.Kind == SyntaxKind.ReturnKeyword)
+            {
+                return ParseReturnStatement();
+            }
 
             return ParseExpressionStatement();
         }
@@ -154,6 +158,13 @@ namespace Upsilon.Parser
             return new ExpressionStatementSyntax(expression);
         }
 
+        private StatementSyntax ParseReturnStatement()
+        {
+            var returnToken = MatchToken(SyntaxKind.ReturnKeyword);
+            var expression = ParseExpression();
+            return new ReturnStatementSyntax(returnToken, expression);
+        }
+
         private ExpressionSyntax ParseExpression()
         {
             return ParseBinaryExpression();

Upsilon/Parser/StatementSyntax/ReturnStatementSyntax.cs

@@ -0,0 +1,21 @@
+using System.Collections.Generic;
+
+namespace Upsilon.Parser
+{
+    public class ReturnStatementSyntax : StatementSyntax
+    {
+        public ReturnStatementSyntax(SyntaxToken returnToken, ExpressionSyntax expression)
+        {
+            ReturnToken = returnToken;
+            Expression = expression;
+        }
+
+        public SyntaxToken ReturnToken { get; }
+        public ExpressionSyntax Expression { get; }
+        public override SyntaxKind Kind => SyntaxKind.ReturnStatement;
+        public override IEnumerable<SyntaxNode> ChildNodes()
+        {
+            yield return Expression;
+        }
+    }
+}

Upsilon/Parser/SyntaxKeyWords.cs

@@ -32,6 +32,8 @@ namespace Upsilon.Parser
                     return SyntaxKind.NilKeyword;
                 case "function":
                     return SyntaxKind.FunctionKeyword;
+                case "return":
+                    return SyntaxKind.ReturnKeyword;
                 default:
                     return SyntaxKind.Identifier;
             }

Upsilon/Parser/SyntaxKind.cs

@@ -35,6 +35,7 @@ namespace Upsilon.Parser
         ElseKeyword,
         NilKeyword,
         FunctionKeyword,
+        ReturnKeyword,
 
         Identifier,
 
@@ -58,5 +59,6 @@ namespace Upsilon.Parser
         ElseIfStatement,
         ElseStatement,
         FunctionStatement,
+        ReturnStatement
     }
 }

UpsilonTests/FunctionTests.cs

@@ -81,5 +81,75 @@ end
             Assert.NotNull(castType.Block);
         }
 
+        [Fact]
+        public void ReturnFromFunction()
+        {
+            const string input = @"
+function testFunc ()
+    return 5
+end
+a = testFunc()
+";
+            var script = new Script(input);
+            Assert.Empty(script.Diagnostics.Messages);
+            script.Evaluate();
+            Assert.Empty(script.Diagnostics.Messages);
+            Assert.True(script.Scope.TryGet("a", out var result));
+            Assert.Equal(5, (long)(NumberLong)result);
+        }
+
+        [Fact]
+        public void ReturnFromFunctionOnce()
+        {
+            const string input = @"
+function testFunc ()
+    return 5
+    return 10
+end
+a = testFunc()
+";
+            var script = new Script(input);
+            Assert.Empty(script.Diagnostics.Messages);
+            script.Evaluate();
+            Assert.Empty(script.Diagnostics.Messages);
+            Assert.True(script.Scope.TryGet("a", out var result));
+            Assert.Equal(5, (long)(NumberLong)result);
+        }
+
+        [Fact]
+        public void ReturnFromFunctionNested()
+        {
+            const string input = @"
+function testFunc ()
+    if true then
+        return 5
+    end
+    return 10
+end
+a = testFunc()
+";
+            var script = new Script(input);
+            Assert.Empty(script.Diagnostics.Messages);
+            script.Evaluate();
+            Assert.Empty(script.Diagnostics.Messages);
+            Assert.True(script.Scope.TryGet("a", out var result));
+            Assert.Equal(5, (long)(NumberLong)result);
+        }
+
+        [Fact]
+        public void ReturnFromScriptAsFunction()
+        {
+            const string input = @"
+a = 100
+return 60
+a = 87
+";
+            var script = new Script(input);
+            Assert.Empty(script.Diagnostics.Messages);
+            var result = script.Evaluate<NumberLong>();
+            Assert.Empty(script.Diagnostics.Messages);
+            Assert.Equal(60, (long)(NumberLong)result);
+        }
+
     }
 }