From 9eb5fb1cadfb3a4fb64453d394a362e419d6ef51 Mon Sep 17 00:00:00 2001
From: "K. Lange" <klange@toaruos.org>
Date: Thu, 7 Jan 2021 13:46:40 +0900
Subject: [PATCH] Add ternary (a) if (cond) else (b)

---
 compiler.c                  | 45 +++++++++++++++++++++++++++++++++++--
 test/testTernary.krk        | 19 ++++++++++++++++
 test/testTernary.krk.expect |  7 ++++++
 3 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 test/testTernary.krk
 create mode 100644 test/testTernary.krk.expect

diff --git a/compiler.c b/compiler.c
index a2e0b3f..8968925 100644
--- a/compiler.c
+++ b/compiler.c
@@ -49,6 +49,7 @@ typedef struct {
 typedef enum {
 	PREC_NONE,
 	PREC_ASSIGNMENT, /* = */
+	PREC_TERNARY,
 	PREC_OR,         /* or */
 	PREC_AND,        /* and */
 	PREC_BITOR,      /* | */
@@ -178,6 +179,7 @@ static void expression();
 static void statement();
 static void declaration();
 static void or_(int canAssign);
+static void ternary(int canAssign);
 static void and_(int canAssign);
 static void classDeclaration();
 static void declareVariable();
@@ -1772,7 +1774,7 @@ ParseRule rules[] = {
 	RULE(TOKEN_FALSE,         literal,  NULL,   PREC_NONE),
 	RULE(TOKEN_FOR,           NULL,     NULL,   PREC_NONE),
 	RULE(TOKEN_DEF,           NULL,     NULL,   PREC_NONE),
-	RULE(TOKEN_IF,            NULL,     NULL,   PREC_NONE),
+	RULE(TOKEN_IF,            NULL,     ternary,PREC_TERNARY),
 	RULE(TOKEN_IN,            NULL,     in_,    PREC_COMPARISON),
 	RULE(TOKEN_LET,           NULL,     NULL,   PREC_NONE),
 	RULE(TOKEN_NONE,          literal,  NULL,   PREC_NONE),
@@ -1809,7 +1811,38 @@ ParseRule rules[] = {
 	RULE(TOKEN_EOF,           NULL,     NULL,   PREC_NONE),
 };
 
+static void actualTernary(size_t count, KrkScanner oldScanner, Parser oldParser) {
+	currentChunk()->count = count;
+
+	parsePrecedence(PREC_OR);
+
+	int thenJump = emitJump(OP_JUMP_IF_TRUE);
+	emitByte(OP_POP); /* Pop the condition */
+	consume(TOKEN_ELSE, "Expected 'else' after ternary condition");
+
+	parsePrecedence(PREC_OR);
+
+	KrkScanner outScanner = krk_tellScanner();
+	Parser outParser = parser;
+
+	int elseJump = emitJump(OP_JUMP);
+	patchJump(thenJump);
+	emitByte(OP_POP);
+
+	krk_rewindScanner(oldScanner);
+	parser = oldParser;
+	parsePrecedence(PREC_OR);
+	patchJump(elseJump);
+
+	krk_rewindScanner(outScanner);
+	parser = outParser;
+}
+
 static void parsePrecedence(Precedence precedence) {
+	size_t count = currentChunk()->count;
+	KrkScanner oldScanner = krk_tellScanner();
+	Parser oldParser = parser;
+
 	advance();
 	ParseFn prefixRule = getRule(parser.previous.type)->prefix;
 	if (prefixRule == NULL) {
@@ -1821,7 +1854,11 @@ static void parsePrecedence(Precedence precedence) {
 	while (precedence <= getRule(parser.current.type)->precedence) {
 		advance();
 		ParseFn infixRule = getRule(parser.previous.type)->infix;
-		infixRule(canAssign);
+		if (infixRule == ternary) {
+			actualTernary(count, oldScanner, oldParser);
+		} else {
+			infixRule(canAssign);
+		}
 	}
 
 	if (canAssign && matchAssignment()) {
@@ -1988,6 +2025,10 @@ static void and_(int canAssign) {
 	patchJump(endJump);
 }
 
+static void ternary(int canAssign) {
+	error("This function should not run.");
+}
+
 static void or_(int canAssign) {
 	int endJump = emitJump(OP_JUMP_IF_TRUE);
 	emitByte(OP_POP);
diff --git a/test/testTernary.krk b/test/testTernary.krk
new file mode 100644
index 0000000..8023c5a
--- /dev/null
+++ b/test/testTernary.krk
@@ -0,0 +1,19 @@
+def foo():
+    print("Called foo")
+    return 42
+
+def bar():
+    print("Called bar")
+    return "a string"
+
+let x = foo() if True else bar()
+print(x)
+
+let y = foo() if False else bar()
+print(y)
+
+let z = foo() if y == "a string" else bar()
+print(z)
+
+let l = (lambda x: x * 72) if x == 42 else (lambda y: y - 48)
+print(l(x))
diff --git a/test/testTernary.krk.expect b/test/testTernary.krk.expect
new file mode 100644
index 0000000..13d39c8
--- /dev/null
+++ b/test/testTernary.krk.expect
@@ -0,0 +1,7 @@
+Called foo
+42
+Called bar
+a string
+Called foo
+42
+3024