diff --git a/py/lexer.c b/py/lexer.c
index ce6aa78..f4aedb4 100644
--- a/py/lexer.c
+++ b/py/lexer.c
@@ -27,6 +27,7 @@
 #include <stdio.h>
 #include <assert.h>
 
+#include "py/mpstate.h"
 #include "py/lexer.h"
 
 #define TAB_SIZE (8)
@@ -34,8 +35,6 @@
 // TODO seems that CPython allows NULL byte in the input stream
 // don't know if that's intentional or not, but we don't allow it
 
-mp_uint_t mp_optimise_value;
-
 // TODO replace with a call to a standard function
 STATIC bool str_strn_equal(const char *str, const char *strn, mp_uint_t len) {
     mp_uint_t i = 0;
@@ -662,7 +661,7 @@
             if (str_strn_equal(tok_kw[i], lex->vstr.buf, lex->vstr.len)) {
                 if (i == MP_ARRAY_SIZE(tok_kw) - 1) {
                     // tok_kw[MP_ARRAY_SIZE(tok_kw) - 1] == "__debug__"
-                    lex->tok_kind = (mp_optimise_value == 0 ? MP_TOKEN_KW_TRUE : MP_TOKEN_KW_FALSE);
+                    lex->tok_kind = (MP_STATE_VM(mp_optimise_value) == 0 ? MP_TOKEN_KW_TRUE : MP_TOKEN_KW_FALSE);
                 } else {
                     lex->tok_kind = MP_TOKEN_KW_FALSE + i;
                 }