From 96674ce01a3ca1e8d5eaebceccf047f14de69394 Mon Sep 17 00:00:00 2001
From: Acid Bubbles <acidbubbles.anon@gmail.com>
Date: Wed, 15 May 2024 23:33:09 -0400
Subject: [PATCH] Fix ternery operators

---
 Scripter.Plugin/src/Lib/Parsing/Parser.cs | 27 +++++++++++++++--------
 Scripter.Tests/LanguageTests.cs           |  6 +++--
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/Scripter.Plugin/src/Lib/Parsing/Parser.cs b/Scripter.Plugin/src/Lib/Parsing/Parser.cs
index 4be2f0e..26b9bdc 100644
--- a/Scripter.Plugin/src/Lib/Parsing/Parser.cs
+++ b/Scripter.Plugin/src/Lib/Parsing/Parser.cs
@@ -346,6 +346,7 @@ private Expression ParseForStatement(ScopeLexicalContext lexicalContext)
                 initializer = ParseVariableDeclaration(loopLexicalContext, false);
             else
                 initializer = ParseValueStatementExpression(loopLexicalContext);
+            Consume().Expect(TokenType.SemiColon);
             var condition = ParseValueStatementExpression(loopLexicalContext);
             Consume().Expect(TokenType.SemiColon);
             var increment = ParseValueStatementExpression(loopLexicalContext);
@@ -370,7 +371,9 @@ private Expression ParseWhileStatement(ScopeLexicalContext lexicalContext)
         private Expression ParseReturnStatement(ScopeLexicalContext lexicalContext)
         {
             MoveNext();
-            var value = !Peek().Match(TokenType.SemiColon) ? ParseValueStatementExpression(lexicalContext) : UndefinedExpression.Instance;
+            var value = !Peek().Match(TokenType.SemiColon)
+                ? ParseValueOrTernaryExpression(lexicalContext)
+                : UndefinedExpression.Instance;
             Consume().Expect(TokenType.SemiColon);
             return new ReturnExpression(value, lexicalContext);
         }
@@ -386,13 +389,7 @@ private VariableDeclarationExpression ParseVariableDeclaration(ScopeLexicalConte
             if (next.Match(TokenType.Assignment))
             {
                 MoveNext();
-                var initialValue = ParseValueStatementExpression(lexicalContext);
-                if (Peek().Match(TokenType.Ternary))
-                {
-                    MoveNext();
-                    initialValue = ParseTernaryRightExpression(lexicalContext, initialValue);
-                }
-                Consume().Expect(TokenType.SemiColon);
+                var initialValue = ParseValueOrTernaryExpression(lexicalContext);
                 return new VariableDeclarationExpression(nameToken.value, initialValue, lexicalContext);
             }
 
@@ -400,6 +397,18 @@ private VariableDeclarationExpression ParseVariableDeclaration(ScopeLexicalConte
             return new VariableDeclarationExpression(nameToken.value, UndefinedExpression.Instance, lexicalContext);
         }
 
+        private Expression ParseValueOrTernaryExpression(ScopeLexicalContext lexicalContext)
+        {
+            // TODO: This whole function is workaround for bad parsing of return a ? b : c
+            var initialValue = ParseValueStatementExpression(lexicalContext);
+            if (Peek().Match(TokenType.Ternary))
+            {
+                MoveNext();
+                initialValue = ParseTernaryRightExpression(lexicalContext, initialValue);
+            }
+            return initialValue;
+        }
+
         private Expression ParseValueStatementExpression(ScopeLexicalContext lexicalContext, int precedence = 0)
         {
             var left = ParsePureValueExpression(lexicalContext);
@@ -649,7 +658,7 @@ private Expression ParseArrowFunctionExpression(ScopeLexicalContext lexicalConte
             if (Peek().Match(TokenType.LeftBrace))
                 body = ParseFunctionBody(functionLexicalContext);
             else
-                body = new CodeBlockExpression(new List<Expression> { ParseValueStatementExpression(functionLexicalContext) }, lexicalContext);
+                body = new CodeBlockExpression(new List<Expression> { ParseValueOrTernaryExpression(functionLexicalContext) }, lexicalContext);
             var function = new FunctionDeclarationExpression(name, arguments, body, functionLexicalContext);
             return function;
         }
diff --git a/Scripter.Tests/LanguageTests.cs b/Scripter.Tests/LanguageTests.cs
index 78e0ab5..4f6d1fa 100644
--- a/Scripter.Tests/LanguageTests.cs
+++ b/Scripter.Tests/LanguageTests.cs
@@ -225,10 +225,12 @@ public void Ternary()
             var value = "true";
             var first = value ? 'first' + '' : 'error';
             var second = (!value) ? 'error' : ('second');
-            return first + ' ' + second;
+            function third() { return true ? 'third' : 'error'; }
+            const fourth = () => true ? 'fourth' : 'error';
+            return first + ' ' + second + ' ' + third() + ' ' + fourth();
             """);
         var result = _program.Run();
-        Assert.That(result.ToString(), Is.EqualTo("first second"));
+        Assert.That(result.ToString(), Is.EqualTo("first second third fourth"));
     }
 
     [Test]