From bb61dbcd5425f806d30645fa5e549addbc08e07e Mon Sep 17 00:00:00 2001
From: Danila Fedorin <danila.fedorin@gmail.com>
Date: Thu, 8 Mar 2018 21:31:24 -0800
Subject: [PATCH] Add support for the return keyword.

---
 include/lexer.h |  1 +
 include/tree.h  |  3 ++-
 src/debug.c     |  3 ++-
 src/lexer.c     |  6 ++++--
 src/parser.c    | 27 +++++++++++++++++++++++++++
 src/tree.c      |  3 ++-
 6 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/include/lexer.h b/include/lexer.h
index e321beb..b51386f 100644
--- a/include/lexer.h
+++ b/include/lexer.h
@@ -64,6 +64,7 @@ enum libab_lexer_token_e {
     TOKEN_KW_DO,
     TOKEN_KW_ARROW,
     TOKEN_KW_FUN,
+    TOKEN_KW_RETURN,
     TOKEN_LAST
 };
 
diff --git a/include/tree.h b/include/tree.h
index fb8f9d6..9d51dc0 100644
--- a/include/tree.h
+++ b/include/tree.h
@@ -22,7 +22,8 @@ enum libab_tree_variant_e {
     TREE_DOWHILE,
     TREE_CALL,
     TREE_FUN,
-    TREE_FUN_PARAM
+    TREE_FUN_PARAM,
+    TREE_RETURN
 };
 
 /**
diff --git a/src/debug.c b/src/debug.c
index 1a73d31..b9c43f1 100644
--- a/src/debug.c
+++ b/src/debug.c
@@ -17,7 +17,8 @@ const char* _debug_node_name(libab_tree_variant var) {
         "dowhile",
         "call",
         "fun",
-        "fun_param"
+        "fun_param",
+        "return"
     };
     return names[var];
 }
diff --git a/src/lexer.c b/src/lexer.c
index 6547fab..1113d03 100644
--- a/src/lexer.c
+++ b/src/lexer.c
@@ -16,7 +16,8 @@ libab_result libab_lexer_init(libab_lexer* lexer) {
         "while",
         "do",
         "->",
-        "fun"
+        "fun",
+        "return"
     };
     libab_lexer_token tokens[] = {
         TOKEN_CHAR,
@@ -27,7 +28,8 @@ libab_result libab_lexer_init(libab_lexer* lexer) {
         TOKEN_KW_WHILE,
         TOKEN_KW_DO,
         TOKEN_KW_ARROW,
-        TOKEN_KW_FUN
+        TOKEN_KW_FUN,
+        TOKEN_KW_RETURN
     };
     const size_t count = sizeof(tokens)/sizeof(libab_lexer_token);
 
diff --git a/src/parser.c b/src/parser.c
index 05b4efc..b1600f4 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -467,6 +467,31 @@ libab_result _parse_fun(struct parser_state* state, libab_tree** store_into) {
     return result;
 }
 
+libab_result _parse_return(struct parser_state* state, libab_tree** store_into) {
+    libab_result result = LIBAB_SUCCESS;
+    libab_tree* child = NULL;
+    if(_parser_is_type(state, TOKEN_KW_RETURN)) {
+        result = _parser_construct_node_vec(state->current_match, store_into);
+        if(result == LIBAB_SUCCESS) {
+            (*store_into)->variant = TREE_RETURN;
+            _parser_state_step(state);
+        }
+    } else {
+        result = LIBAB_UNEXPECTED;
+    }
+
+    if(result == LIBAB_SUCCESS) {
+        PARSE_CHILD(result, state, _parse_expression, child, &(*store_into)->children);
+    }
+
+    if(result != LIBAB_SUCCESS) {
+        if(*store_into) libab_tree_free_recursive(*store_into);
+        *store_into = NULL;
+    }
+
+    return result;
+}
+
 libab_result _parse_while(struct parser_state* state, libab_tree** store_into) {
     libab_result result = LIBAB_SUCCESS;
     libab_tree* condition = NULL;
@@ -610,6 +635,8 @@ libab_result _parse_atom(struct parser_state* state, libab_tree** store_into) {
         result = _parse_braced_block(state, store_into);
     } else if(_parser_is_type(state, TOKEN_KW_FUN)) {
         result = _parse_fun(state, store_into);
+    } else if(_parser_is_type(state, TOKEN_KW_RETURN)) {
+        result = _parse_return(state, store_into);
     } else {
         result = LIBAB_UNEXPECTED;
     }
diff --git a/src/tree.c b/src/tree.c
index d0b38e4..bc09e40 100644
--- a/src/tree.c
+++ b/src/tree.c
@@ -5,7 +5,8 @@ int libab_tree_has_vector(libab_tree_variant variant) {
     return variant == TREE_BASE || variant == TREE_OP ||
             variant == TREE_UNARY_OP || variant == TREE_BLOCK ||
             variant == TREE_IF || variant == TREE_CALL || variant == TREE_WHILE ||
-            variant == TREE_DOWHILE || variant == TREE_FUN;
+            variant == TREE_DOWHILE || variant == TREE_FUN ||
+            variant == TREE_RETURN;
 }
 
 int libab_tree_has_string(libab_tree_variant variant) {