Add ternary (a) if (cond) else (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,8 +1854,12 @@ static void parsePrecedence(Precedence precedence) {
 	while (precedence <= getRule(parser.current.type)->precedence) {
 		advance();
 		ParseFn infixRule = getRule(parser.previous.type)->infix;
+		if (infixRule == ternary) {
+			actualTernary(count, oldScanner, oldParser);
+		} else {
 			infixRule(canAssign);
 		}
+	}
 
 	if (canAssign && matchAssignment()) {
 		error("invalid assignment target");
@@ -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);

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))

test/testTernary.krk.expect

@@ -0,0 +1,7 @@
+Called foo
+42
+Called bar
+a string
+Called foo
+42
+3024