From 71ce764c194d6960f4c71cb2fc467ba463486506 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Mon, 10 Apr 2023 19:39:15 +0800
Subject: [PATCH] fix https://github.com/blueloveTH/pocketpy/issues/35

---
 src/compiler.h          | 41 ++++++++++++++++++++++++++++-------------
 src/lexer.h             |  1 +
 tests/20_controlflow.py |  6 ++++++
 3 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/src/compiler.h b/src/compiler.h
index f717499b..beaed8f2 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -25,10 +25,10 @@ class Compiler {
     int i = 0;
     std::vector<Token> tokens;
 
-    const Token& prev() { return tokens.at(i-1); }
-    const Token& curr() { return tokens.at(i); }
-    const Token& next() { return tokens.at(i+1); }
-    const Token& err() {
+    const Token& prev() const{ return tokens.at(i-1); }
+    const Token& curr() const{ return tokens.at(i); }
+    const Token& next() const{ return tokens.at(i+1); }
+    const Token& err() const{
         if(i >= tokens.size()) return prev();
         return curr();
     }
@@ -89,6 +89,7 @@ class Compiler {
         rules[TK("|")] =        { nullptr,               METHOD(exprBinaryOp),       PREC_BITWISE_OR };
         rules[TK("^")] =        { nullptr,               METHOD(exprBinaryOp),       PREC_BITWISE_XOR };
         rules[TK("?")] =        { nullptr,               METHOD(exprTernary),        PREC_TERNARY };
+        rules[TK("if")] =       { nullptr,               METHOD(exprTernary),        PREC_TERNARY };
         rules[TK(",")] =        { nullptr,               METHOD(exprTuple),          PREC_TUPLE };
         rules[TK("not in")] =   { nullptr,               METHOD(exprBinaryOp),       PREC_TEST };
         rules[TK("is not")] =   { nullptr,               METHOD(exprBinaryOp),       PREC_TEST };
@@ -195,7 +196,8 @@ class Compiler {
             consume(TK(":"));
         }
         e->decl->code = push_context(lexer->src, e->decl->name.sv());
-        EXPR(false); // https://github.com/blueloveTH/pocketpy/issues/37
+        // https://github.com/blueloveTH/pocketpy/issues/37
+        parse_expression(PREC_LAMBDA + 1, false);
         ctx()->emit(OP_RETURN_VALUE, BC_NOARG, BC_KEEPLINE);
         pop_context();
         ctx()->s_expr.push(std::move(e));
@@ -235,12 +237,25 @@ class Compiler {
     
     void exprTernary(){
         auto e = make_expr<TernaryExpr>();
-        e->cond = ctx()->s_expr.popx();
-        EXPR();         // if true
-        e->true_expr = ctx()->s_expr.popx();
-        consume(TK(":"));
-        EXPR();         // if false
-        e->false_expr = ctx()->s_expr.popx();
+        if(prev().type == TK("if")){
+            e->true_expr = ctx()->s_expr.popx();
+            // cond
+            parse_expression(PREC_TERNARY + 1);
+            e->cond = ctx()->s_expr.popx();
+            consume(TK("else"));
+            // if false
+            parse_expression(PREC_TERNARY + 1);
+            e->false_expr = ctx()->s_expr.popx();
+        }else{  // ?:
+            e->cond = ctx()->s_expr.popx();
+            // if true
+            parse_expression(PREC_TERNARY + 1);
+            e->true_expr = ctx()->s_expr.popx();
+            consume(TK(":"));
+            // if false
+            parse_expression(PREC_TERNARY + 1);
+            e->false_expr = ctx()->s_expr.popx();
+        }
         ctx()->s_expr.push(std::move(e));
     }
 
@@ -291,11 +306,11 @@ class Compiler {
         ce->expr = std::move(expr);
         ce->vars = EXPR_VARS();
         consume(TK("in"));
-        EXPR();
+        parse_expression(PREC_TERNARY + 1);
         ce->iter = ctx()->s_expr.popx();
         match_newlines_repl();
         if(match(TK("if"))){
-            EXPR();
+            parse_expression(PREC_TERNARY + 1);
             ce->cond = ctx()->s_expr.popx();
         }
         ctx()->s_expr.push(std::move(ce));
diff --git a/src/lexer.h b/src/lexer.h
index 89f05147..726a5a36 100644
--- a/src/lexer.h
+++ b/src/lexer.h
@@ -72,6 +72,7 @@ struct Token{
 enum Precedence {
   PREC_NONE,
   PREC_TUPLE,         // ,
+  PREC_LAMBDA,        // lambda
   PREC_TERNARY,       // ?:
   PREC_LOGICAL_OR,    // or
   PREC_LOGICAL_AND,   // and
diff --git a/tests/20_controlflow.py b/tests/20_controlflow.py
index 40ebf24a..7c3dd8a0 100644
--- a/tests/20_controlflow.py
+++ b/tests/20_controlflow.py
@@ -65,3 +65,9 @@ while (count < 1000):
    count = count + 1
 assert count == 1000
 
+# ternary operator
+d = 1 if 2 > 1 else 2
+assert d == 1
+d = 1 if 2 < 1 else 2
+assert d == 2
+