Add parsing of if statements

5 files changed, 86 insertions, 2 deletions

diff --git a/src/lexer.c b/src/lexer.c
index 729ea14..b4dcaf5 100644
--- a/src/lexer.c
+++ b/src/lexer.c
@@ -29,6 +29,11 @@ static const char* token_str[] = {
     [TOKEN_NOT]     = "TOKEN_NOT",
     [TOKEN_AND]     = "TOKEN_AND",
     [TOKEN_OR]      = "TOKEN_OR",
+    [TOKEN_EQ]      = "TOKEN_EQ",
+    [TOKEN_LT]      = "TOKEN_LT",
+    [TOKEN_GT]      = "TOKEN_GT",
+    [TOKEN_LE]      = "TOKEN_LE",
+    [TOKEN_GE]      = "TOKEN_GE",
     [TOKEN_COLON]   = "TOKEN_COLON",
     [TOKEN_DOT]     = "TOKEN_DOT",
     [TOKEN_AT]      = "TOKEN_AT",
@@ -61,6 +66,11 @@ static const Keyword keywords[] = {
     KEYWORD("not",    TOKEN_NOT),
     KEYWORD("and",    TOKEN_AND),
     KEYWORD("or",     TOKEN_OR),
+    KEYWORD("=",      TOKEN_EQ),
+    KEYWORD("<",      TOKEN_LT),
+    KEYWORD(">",      TOKEN_GT),
+    KEYWORD("<=",     TOKEN_LE),
+    KEYWORD(">=",     TOKEN_GE),
 };
 
 void
diff --git a/src/lexer.h b/src/lexer.h
index fa43ebb..949abaf 100644
--- a/src/lexer.h
+++ b/src/lexer.h
@@ -41,6 +41,11 @@ typedef enum TokenType {
     TOKEN_NOT,
     TOKEN_AND,
     TOKEN_OR,
+    TOKEN_EQ,
+    TOKEN_LT,
+    TOKEN_GT,
+    TOKEN_LE,
+    TOKEN_GE,
 
     // Special operators.
     TOKEN_COLON,
diff --git a/src/nodes.c b/src/nodes.c
index 3a63246..af3b772 100644
--- a/src/nodes.c
+++ b/src/nodes.c
@@ -11,6 +11,9 @@ alloc_node(NodeType type) {
 
 void
 print_node(Node *node) {
+    if (node == NULL) {
+        return;
+    }
     switch (node->type) {
         case NODE_NUMBER: {
             if (node->number.negative) {
@@ -73,6 +76,17 @@ print_node(Node *node) {
             }
             printf(" }");
         } break;
+        case NODE_IF: {
+            printf("(if ");
+            print_node(node->ifexpr.cond);
+            printf(" ");
+            print_node(node->ifexpr.expr_true);
+            if (node->ifexpr.expr_false != NULL) {
+                printf(" ");
+                print_node(node->ifexpr.expr_false);
+            }
+            printf(")");
+        } break;
         case NODE_FUN: {
             printf("(fun ");
             print_node(node->fun.name);
diff --git a/src/nodes.h b/src/nodes.h
index c1520b3..4fb48a1 100644
--- a/src/nodes.h
+++ b/src/nodes.h
@@ -12,6 +12,7 @@ typedef enum NodeType {
     NODE_SET,
     NODE_FUN,
     NODE_BLOCK,
+    NODE_IF,
 } NodeType;
 
 typedef struct Node {
@@ -63,6 +64,13 @@ typedef struct Node {
         struct {
             struct Node **expr;
         } block;
+
+        // If statement.
+        struct {
+            struct Node *cond;
+            struct Node *expr_true;
+            struct Node *expr_false;
+        } ifexpr;
     };
 } Node;
 
diff --git a/src/parser.c b/src/parser.c
index 58ad722..b76e32f 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -291,7 +291,46 @@ parse_fun(Parser *parser) {
         return NULL;
     }
     node->fun.body = body;
-    consume_rparen(parser);
+    if (!consume_rparen(parser)) {
+        return NULL;
+    }
+
+    return node;
+}
+
+Node *
+parse_if(Parser *parser) {
+    next_token(parser); // Skip keyword.
+
+    Node *node = alloc_node(NODE_IF);
+    node->ifexpr.cond = NULL;
+    node->ifexpr.expr_true = NULL;
+    node->ifexpr.expr_false = NULL;
+
+    Node *cond = parse_next(parser);
+    if (cond == NULL) {
+        return NULL;
+    }
+    Node *expr_true = parse_next(parser);
+    if (expr_true == NULL) {
+        return NULL;
+    }
+    node->ifexpr.cond = cond;
+    node->ifexpr.expr_true = expr_true;
+    Token tok = peek_token(parser);
+
+    // Optional else statement.
+    if (tok.type != TOKEN_RPAREN) {
+        Node *expr_false = parse_next(parser);
+        if (expr_false == NULL) {
+            return NULL;
+        }
+        node->ifexpr.expr_false = expr_false;
+    }
+
+    if (!consume_rparen(parser)) {
+        return NULL;
+    }
 
     return node;
 }
@@ -299,6 +338,7 @@ parse_fun(Parser *parser) {
 Node *
 parse_paren(Parser *parser) {
     next_token(parser); // Skip paren.
+
     Token tok = peek_token(parser);
 
     switch (tok.type) {
@@ -310,11 +350,17 @@ parse_paren(Parser *parser) {
         case TOKEN_MOD:
         case TOKEN_NOT:
         case TOKEN_AND:
-        case TOKEN_OR: { return parse_builtin(parser); } break;
+        case TOKEN_OR:
+        case TOKEN_EQ:
+        case TOKEN_LT:
+        case TOKEN_GT:
+        case TOKEN_LE:
+        case TOKEN_GE: { return parse_builtin(parser); } break;
         // Special functions.
         case TOKEN_DEF: { return parse_def(parser); } break;
         case TOKEN_SET: { return parse_set(parser); } break;
         case TOKEN_FUN: { return parse_fun(parser); } break;
+        case TOKEN_IF: { return parse_if(parser); } break;
         default: break;
     }
 
@@ -326,6 +372,7 @@ parse_paren(Parser *parser) {
 Node *
 parse_block(Parser *parser) {
     next_token(parser); // Skip paren.
+
     Node *node = alloc_node(NODE_BLOCK);
     array_init(node->block.expr, 0);
     while (true) {