1 files changed, 27 insertions, 10 deletions

diff --git a/src/tokenizer.c b/src/tokenizer.c
index 742f9ca..b9f0ad3 100644
--- a/src/tokenizer.c
+++ b/src/tokenizer.c
@@ -25,6 +25,8 @@ int tokenizer_init(Tokenizer *self, BufferView code, BufferView code_name) {
     self->index = 0;
     self->prev_index = 0;
     self->line = 1;
+    self->token = TOK_NONE;
+    self->needs_update = bool_true;
     self->code_name = code_name.data ? code_name : create_buffer_view("<buffer>", 8);
     return 0;
 }
@@ -49,10 +51,12 @@ static Token tokenizer_skip_whitespace(Tokenizer *self) {
             case '\t':
                 break;
             default:
+                self->prev_index = self->index;
                 return TOK_NONE;
         }
         ++self->index;
     }
+    self->prev_index = self->index;
 }
 
 /* Returns -1 if end of string can't be found */
@@ -73,7 +77,9 @@ static int find_end_of_string(BufferView buf, int index) {
     return -1;
 }
 
-int tokenizer_next(Tokenizer *self, Token *token) {
+static CHECK_RESULT int tokenizer_next(Tokenizer *self, Token *token);
+
+static CHECK_RESULT int __tokenizer_next(Tokenizer *self, Token *token) {
     Token last_token;
     int c;
     int result;
@@ -84,7 +90,6 @@ int tokenizer_next(Tokenizer *self, Token *token) {
         return TOKENIZER_OK;
     }
 
-    self->prev_index = self->index;
     c = tokenizer_get_char(self);
     if(isAlpha(c) || c == '_') {
         int identifier_start;
@@ -181,11 +186,21 @@ int tokenizer_next(Tokenizer *self, Token *token) {
     return TOKENIZER_OK;
 }
 
+/* Wrapper around __tokenizer_next to store last parsed token */
+int tokenizer_next(Tokenizer *self, Token *token) {
+    int result;
+    result = __tokenizer_next(self, token);
+    self->token = *token;
+    return result;
+}
+
 int tokenizer_accept(Tokenizer *self, Token expected_token) {
     Token actual_token;
     return_if_error(tokenizer_next(self, &actual_token));
-    if(actual_token == expected_token)
+    if(actual_token == expected_token) {
+        self->needs_update = bool_true;
         return TOKENIZER_OK;
+    }
 
     /* Todo: convert token to string */
     tokenizer_print_error(self, "Expected %d, got %d", expected_token, actual_token);
@@ -193,19 +208,21 @@ int tokenizer_accept(Tokenizer *self, Token expected_token) {
 }
 
 int tokenizer_consume_if(Tokenizer *self, Token expected_token, bool *result) {
-    int index;
-    int line;
     Token actual_token;
 
-    index = self->index;
-    line = self->line;
+    if(!self->needs_update) {
+        *result = (self->token == expected_token);
+        if(*result)
+            self->needs_update = bool_true;
+        return TOKENIZER_OK;
+    }
+
     return_if_error(tokenizer_next(self, &actual_token));
     if(actual_token == expected_token) {
+        self->needs_update = bool_true;
         *result = bool_true;
     } else {
-        self->index = index;
-        self->prev_index = index;
-        self->line = line;
+        self->needs_update = bool_false;
         *result = bool_false;
     }
     return TOKENIZER_OK;