<commit_before>#include <setjmp.h>
#include <stdarg.h>
#include <string.h>
#include <stdbool.h>
#include <assert.h>
#include <errno.h>
#include <math.h>
#include <stdnoreturn.h>

#include "vec.h"
#include "token.h"
#include "test.h"
#include "ast.h"
#include "util.h"
#include "alloc.h"
#include "lex.h"
#include "operators.h"
#include "value.h"
#include "log.h"
#include "vm.h"

#define BINARY_OPERATOR(name, token, prec, right_assoc) \
        static struct expression * \
        infix_ ## name(struct expression *left) \
        { \
                consume(TOKEN_ ## token); \
                struct expression *e = mkexpr(); \
                e->type = EXPRESSION_ ## token; \
                e->left = left; \
                e->right = parse_expr(prec - (right_assoc ? 1 : 0)); \
                return e; \
        } \

#define BINARY_LVALUE_OPERATOR(name, token, prec, right_assoc) \
        static struct expression * \
        infix_ ## name(struct expression *left) \
        { \
                consume(TOKEN_ ## token); \
                struct expression *e = mkexpr(); \
                e->type = EXPRESSION_ ## token; \
                e->target = assignment_lvalue(left); \
                e->value = parse_expr(prec - (right_assoc ? 1 : 0)); \
                return e; \
        } \

#define PREFIX_OPERATOR(name, token, prec) \
        static struct expression * \
        prefix_ ## name(void) \
        { \
                consume(TOKEN_ ## token); \
                struct expression *e = mkexpr(); \
                e->type = EXPRESSION_PREFIX_ ## token; \
                e->operand = parse_expr(prec); \
                return e; \
        } \

#define PREFIX_LVALUE_OPERATOR(name, token, prec) \
        static struct expression * \
        prefix_ ## name(void) \
        { \
                consume(TOKEN_ ## token); \
                struct expression *e = mkexpr(); \
                e->type = EXPRESSION_PREFIX_ ## token; \
                e->operand = assignment_lvalue(parse_expr(prec)); \
                return e; \
        } \

typedef struct expression *parse_fn();

enum {
        MAX_ERR_LEN  = 2048
};

/*
 * Meaningful names for different contexts in which we can parse lvalues.
 *
 * In a let declaration, the lvalue must be '='.
 * In a for-each loop, it must be followed by 'in'.
 *
 * Passing these values to parse_definition_lvalue determines its behaviour.
 */
enum {
        LV_LET,
        LV_EACH,
        LV_SUB,
        LV_ANY
};

static jmp_buf jb;
static char errbuf[MAX_ERR_LEN + 1];

static vec(struct token) tokens;
static int tokidx = 0;
enum lex_context lex_ctx = LEX_PREFIX;

static int depth;
static int dot_lambda_depth;
static bool new_dot_lambda;
static bool NoEquals = false;

static char const *filename;

static struct statement BREAK_STATEMENT    = { .type = STATEMENT_BREAK,    .loc = {42, 42}  };
static struct statement CONTINUE_STATEMENT = { .type = STATEMENT_CONTINUE, .loc = {42, 42}  };
static struct statement NULL_STATEMENT     = { .type = STATEMENT_NULL,     .loc = {42, 42}  };

static struct statement *
parse_statement(void);

static struct expression *
parse_expr(int);

static struct statement *
parse_match_statement(void);

static struct statement *
parse_let_definition(void);

static struct expression *
parse_target_list(void);

static struct statement *
parse_block(void);

static struct expression *
assignment_lvalue(struct expression *e);

static struct expression *
definition_lvalue(struct expression *e);

static struct expression *
infix_member_access(struct expression *e);

inline static struct token *
tok(void);

/*
 * Get a unique identifier name.
 * This sucks.
 */
char *
gensym(void)
{
        static int sym = 0;
        char buf[24];

        sprintf(buf, ":%d", sym++);

        return sclone(buf);
}

inline static struct expression *
mkexpr(void)
{
        struct expression *e = alloc(sizeof *e);
        e->loc = tok()->loc;
        return e;
}

inline static struct statement *
mkstmt(void)
{
        struct statement *s = alloc(sizeof *s);
        s->loc = tok()->loc;
        return s;
}

inline static struct statement *
mkret(struct expression *value)
{
        struct statement *s = mkstmt();
        s->type = STATEMENT_RETURN;
        vec_init(s->returns);
        vec_push(s->returns, value);
        return s;
}

inline static struct statement *
mkdef(struct expression *lvalue, char *name)
{
        struct expression *value = mkexpr();
        value->type = EXPRESSION_IDENTIFIER;
        value->identifier = name;
        value->module = NULL;

        struct statement *s = mkstmt();
        s->type = STATEMENT_DEFINITION;
        s->target = lvalue;
        s->value = value;

        return s;
}

noreturn static void
error(char const *fmt, ...)
{
        if (fmt == NULL) {
                strcpy(errbuf, lex_error());
        } else {
                va_list ap;
                va_start(ap, fmt);

                int n = 0;

                char *err = errbuf;
                n += sprintf(err, "ParseError: %s:%d:%d: ", filename, tok()->loc.line + 1, tok()->loc.col + 1);
                vsnprintf(err + n, MAX_ERR_LEN - n, fmt, ap);

                va_end(ap);
        }

        LOG("Parse Error: %s", errbuf);

        longjmp(jb, 1);
}

inline static void
skip(int n)
{
        tokidx += n;
}

inline static void
next(void)
{
        skip(1);
}

inline static struct token *
token(int i)
{
        struct token t;
        while (tokens.count <= i + tokidx) {
                t = lex_token(lex_ctx);
                if (t.type == TOKEN_ERROR) {
                        error(NULL);
                }
                vec_push(tokens, t);
        }

        return &tokens.items[tokidx + i];
}

inline static struct token *
tok(void)
{
        return token(0);
}

/*
 * Push a token into the token stream, so that it will be returned by the next call
 * to tok().
 */
inline static void
unconsume(int type)
{
        vec_insert(tokens, ((struct token){ .type = type, .loc = { 42, 42 } }), tokidx);
}

static void
expect(int type)
{
        if (tok()->type != type) {
                error("expected %s but found %s", token_show_type(type), token_show(tok()));
        }
}


static void
expect_keyword(int type)
{
        if (tok()->type != TOKEN_KEYWORD || tok()->keyword != type) {
                error("expected %s but found %s", token_show(&(struct token){ .type = TOKEN_KEYWORD, .keyword = type }), token_show(tok()));
        }
}

inline static void
consume(int type)
{
        expect(type);
        tokidx += 1;
}

inline static void
consume_keyword(int type)
{
        expect_keyword(type);
        tokidx += 1;

}

/* * * * | prefix parsers | * * * */
static struct expression *
prefix_integer(void)
{
        expect(TOKEN_INTEGER);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_INTEGER;
        e->integer = tok()->integer;

        consume(TOKEN_INTEGER);

        return e;
}

static struct expression *
prefix_real(void)
{
        expect(TOKEN_REAL);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_REAL;
        e->real = tok()->real;

        consume(TOKEN_REAL);

        return e;
}

static struct expression *
prefix_string(void)
{
        expect(TOKEN_STRING);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_STRING;
        e->string = tok()->string;

        consume(TOKEN_STRING);

        return e;
}

static struct expression *
prefix_special_string(void)
{
        expect(TOKEN_SPECIAL_STRING);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_SPECIAL_STRING;
        vec_init(e->expressions);
        
        e->strings.items = tok()->strings.items;
        e->strings.count = tok()->strings.count;

        char **exprs = tok()->expressions.items;
        int count = tok()->expressions.count;

        struct location *locs = tok()->locations.items;

        consume(TOKEN_SPECIAL_STRING);

        for (int i = 0; i < count; ++i) {
                lex_start(exprs[i]);
                vec_push(e->expressions, parse_expr(0));
                if (tok()->type != TOKEN_END) {
                        /*
                         * This results in a more accurate location in the error message.
                         */
                        tok()->loc = locs[i];
                        error("expression in interpolated string has trailing tokens: %s", exprs[i]);
                }
                consume(TOKEN_END);
                lex_end();
        }

        return e;
}

static struct expression *
prefix_dollar(void)
{
        consume('$');

        struct expression *e = mkexpr();

        expect(TOKEN_IDENTIFIER);

        e->type = EXPRESSION_MATCH_NOT_NIL;
        e->identifier = tok()->identifier;
        e->module = tok()->module;

        if (e->module != NULL)
                error("unpexpected module in lvalue");

        consume(TOKEN_IDENTIFIER);

        return e;
}

static struct expression *
prefix_identifier(void)
{
        expect(TOKEN_IDENTIFIER);

        struct expression *e = mkexpr();

        e->type = EXPRESSION_IDENTIFIER;
        e->identifier = tok()->identifier;
        e->module = tok()->module;

        consume(TOKEN_IDENTIFIER);

        return e;
}

static struct expression *
prefix_function(void)
{
        consume_keyword(KEYWORD_FUNCTION);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_FUNCTION;
        e->rest = false;
        vec_init(e->params);

        if (tok()->type == TOKEN_IDENTIFIER) {
                e->name = tok()->identifier;
                next();
        } else {
                e->name = NULL;
        }

        consume('(');

        while (tok()->type != ')') {
                bool rest = false;
                if (tok()->type == TOKEN_STAR) {
                        rest = true;
                        next();
                }

                expect(TOKEN_IDENTIFIER);
                vec_push(e->params, sclone(tok()->identifier));
                next();

                if (rest) {
                        e->params.count -= 1;
                        e->rest = 1;
                        expect(')');
                } else if (tok()->type == ',') {
                        next();
                }
        }

        consume(')');

        e->body = parse_statement();

        return e;
}

static struct expression *
prefix_star(void)
{
        consume(TOKEN_STAR);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_MATCH_REST;

        expect(TOKEN_IDENTIFIER);
        e->identifier = tok()->identifier;
        if (tok()->module != NULL)
                error("unexpected module qualifier in lvalue");

        consume(TOKEN_IDENTIFIER);

        return e;
}

static struct expression *
prefix_match(void)
{
        consume_keyword(KEYWORD_MATCH);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_MATCH;

        e->subject = parse_expr(0);

        consume('{');

        vec_init(e->patterns);
        vec_init(e->conds);
        vec_init(e->thens);

        vec_push(e->patterns, parse_expr(0));
        if (tok()->type == TOKEN_BIT_OR) {
                next();
                vec_push(e->conds, parse_expr(0));
        } else {
                vec_push(e->conds, NULL);
        }

        consume(TOKEN_FAT_ARROW);
        vec_push(e->thens, parse_expr(0));

        while (tok()->type == ',') {
                next();
                vec_push(e->patterns, parse_expr(0));
                if (tok()->type == TOKEN_BIT_OR) {
                        next();
                        vec_push(e->conds, parse_expr(0));
                } else {
                        vec_push(e->conds, NULL);
                }

                consume(TOKEN_FAT_ARROW);
                vec_push(e->thens, parse_expr(0));
        }

        consume('}');

        return e;
}

static struct expression *
prefix_parenthesis(void)
{
        /*
         * This can either be a plain old parenthesized expression, e.g., (4 + 4)
         * or it can be an identifier list for an arrow function, e.g., (a, b, c).
         */

        consume('(');

        /*
         * () is an empty identifier list.
         */
        if (tok()->type == ')') {
                next();
                struct expression *list = mkexpr();
                list->type = EXPRESSION_LIST;
                list->only_identifiers = true;
                vec_init(list->es);
                return list;
        }

        struct expression *e = parse_expr(0);

        if (tok()->type == ',') {
                struct expression *list = mkexpr();
                list->only_identifiers = true;

                /*
                 * It _must_ be an identifier list.
                 */
                if (e->type != EXPRESSION_IDENTIFIER && e->type != EXPRESSION_MATCH_REST) {
                        list->only_identifiers = false;
                }

                list->type = EXPRESSION_LIST;
                vec_init(list->es);
                vec_push(list->es, e);

                while (tok()->type == ',') {
                        next();
                        struct expression *e = parse_expr(0);
                        if (e->type == EXPRESSION_MATCH_REST) {
                                expect(')');
                        } else if (e->type != EXPRESSION_IDENTIFIER) {
                                list->only_identifiers = false;
                        }
                        vec_push(list->es, e);
                }

                consume(')');

                return list;
        } else {
                consume(')');
                return e;
        }
}

static struct expression *
prefix_true(void)
{
        consume_keyword(KEYWORD_TRUE);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_BOOLEAN;
        e->boolean = true;

        return e;
}

static struct expression *
prefix_false(void)
{
        consume_keyword(KEYWORD_FALSE);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_BOOLEAN;
        e->boolean = false;

        return e;
}

static struct expression *
prefix_nil(void)
{
        consume_keyword(KEYWORD_NIL);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_NIL;

        return e;
}

static struct expression *
prefix_regex(void)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_REGEX;
        e->regex = tok()->regex;
        e->extra = tok()->extra;
        e->pattern = tok()->pattern;

        consume(TOKEN_REGEX);

        return e;
}


static struct expression *
prefix_array(void)
{
        consume('[');

        struct expression *e = mkexpr();
        e->type = EXPRESSION_ARRAY;
        vec_init(e->elements);

        if (tok()->type == ']') {
                next();
                return e;
        } else {
                vec_push(e->elements, parse_expr(0));
        }

        while (tok()->type != ']') {
                if (tok()->type == TOKEN_KEYWORD && tok()->keyword == KEYWORD_FOR) {
                        next();
                        e->type = EXPRESSION_ARRAY_COMPR;
                        e->compr.pattern = parse_target_list();
                        consume_keyword(KEYWORD_IN);
                        e->compr.iter = parse_expr(0);
                        if (tok()->type == TOKEN_BIT_OR) {
                                next();
                                e->compr.cond = parse_expr(0);
                        } else {
                                e->compr.cond = NULL;
                        }
                        expect(']');
                } else {
                        consume(',');
                        vec_push(e->elements, parse_expr(0));
                }
        }

        consume(']');

        return e;
}

static struct expression *
prefix_tick(void)
{
        consume('`');

        expect(TOKEN_IDENTIFIER);

        if (tok()->module != NULL)
                error("unexpected module qualifier in ` pattern");

        struct expression *e = mkexpr();

        e->type = EXPRESSION_TICK;
        e->identifier = tok()->identifier;
        e->module = NULL;

        next();

        return e;
}

static struct expression *
prefix_dot(void)
{
        if (dot_lambda_depth > 0 && !new_dot_lambda) {
                struct expression *e = mkexpr();
                e->type = EXPRESSION_IDENTIFIER;
                e->identifier = "<object>";
                e->module = NULL;
                return infix_member_access(e);
        } else {
                new_dot_lambda = false;
                ++dot_lambda_depth;
                unconsume(TOKEN_ARROW);
                unconsume(TOKEN_IDENTIFIER);
                tok()->identifier = "<object>";
                tok()->module = NULL;
                return prefix_identifier();
        }
}

static struct expression *
prefix_incrange(void)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_DOT_DOT_DOT;

        struct expression *zero = mkexpr();
        zero->type = EXPRESSION_INTEGER;
        zero->integer = 0;

        consume(TOKEN_DOT_DOT_DOT);

        e->left = zero;
        e->right = parse_expr(0);

        return e;
}

static struct expression *
prefix_range(void)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_DOT_DOT;

        struct expression *zero = mkexpr();
        zero->type = EXPRESSION_INTEGER;
        zero->integer = 0;

        consume(TOKEN_DOT_DOT);

        e->left = zero;
        e->right = parse_expr(0);

        return e;
}

static struct expression *
prefix_hash(void)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_IDENTIFIER;
        e->identifier = "#";
        e->module = NULL;

        consume('#');

        return e;
}

static struct expression *
prefix_implicit_lambda(void)
{
        consume(TOKEN_BIT_OR);

        struct expression *f = mkexpr();
        f->type = EXPRESSION_FUNCTION;
        f->name = NULL;
        vec_init(f->params);
        vec_push(f->params, "#");

        struct expression *e = parse_expr(0);
        f->body = mkret(e);

        consume(TOKEN_BIT_OR);

        return f;
}

static struct expression *
prefix_object(void)
{
        consume('{');

        struct expression *e = mkexpr();
        e->type = EXPRESSION_DICT;

        vec_init(e->keys);
        vec_init(e->values);

        if (tok()->type == '}') {
                next();
                return e;
        } else {
                vec_push(e->keys, parse_expr(0));
                if (tok()->type == ':') {
                        next();
                        vec_push(e->values, parse_expr(0));
                } else {
                        vec_push(e->values, NULL);
                }
        }

        while (tok()->type != '}') {
                if (tok()->type == TOKEN_KEYWORD && tok()->keyword == KEYWORD_FOR) {
                        next();
                        e->type = EXPRESSION_DICT_COMPR;
                        e->dcompr.pattern = parse_target_list();
                        consume_keyword(KEYWORD_IN);
                        e->dcompr.iter = parse_expr(0);
                        if (tok()->type == TOKEN_BIT_OR) {
                                next();
                                e->dcompr.cond = parse_expr(0);
                        } else {
                                e->dcompr.cond = NULL;
                        }
                        expect('}');
                } else {
                        consume(',');
                        vec_push(e->keys, parse_expr(0));
                        if (tok()->type == ':') {
                                next();
                                vec_push(e->values, parse_expr(0));
                        } else {
                                vec_push(e->values, NULL);
                        }
                }
        }

        consume('}');

        return e;
}

PREFIX_OPERATOR(at,    AT,    9)
PREFIX_OPERATOR(minus, MINUS, 9)
PREFIX_OPERATOR(bang,  BANG,  10)

PREFIX_LVALUE_OPERATOR(inc,   INC,   6)
PREFIX_LVALUE_OPERATOR(dec,   DEC,   6)
/* * * * | end of prefix parsers | * * * */

/* * * * | infix parsers | * * * */
static struct expression *
infix_function_call(struct expression *left)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_FUNCTION_CALL;
        e->function = left;
        vec_init(e->args);

        consume('(');

        lex_ctx = LEX_PREFIX;

        if (tok()->type == ')') {
                next();
                return e;
        } else {
                vec_push(e->args, parse_expr(0));
        }

        while (tok()->type == ',') {
                next();
                vec_push(e->args, parse_expr(0));
        }

        consume(')');

        return e;
}

static struct expression *
infix_eq(struct expression *left)
{
        consume(TOKEN_EQ);
        struct expression *e = mkexpr();
        e->type = EXPRESSION_EQ;
        e->target = assignment_lvalue(left);
        if (left->type == EXPRESSION_LIST) {
                e->value = parse_expr(-1);
        } else {
                e->value = parse_expr(1);
        }

        return e;
}

static struct expression *
infix_user_op(struct expression *left)
{
        struct expression *e = mkexpr();

        e->type = EXPRESSION_FUNCTION_CALL;
        e->function = mkexpr();
        e->function->type = EXPRESSION_IDENTIFIER;
        e->function->identifier = tok()->identifier;
        e->function->module = NULL;
        consume(TOKEN_USER_OP);

        vec_init(e->args);
        vec_push(e->args, left);
        vec_push(e->args, parse_expr(2));

        return e;
}

static struct expression *
infix_list(struct expression *left)
{

        struct expression *e = mkexpr();
        e->type = EXPRESSION_LIST;
        vec_init(e->es);
        vec_push(e->es, left);

        NoEquals = true;

        while (tok()->type == ',') {
                next();
                vec_push(e->es, parse_expr(1));
        }

        NoEquals = false;

        return e;
}

static struct expression *
infix_subscript(struct expression *left)
{
        consume('[');
        
        struct expression *e = mkexpr();
        e->type = EXPRESSION_SUBSCRIPT;
        e->container = left;
        e->subscript = parse_expr(0);

        consume(']');

        return e;
}

static struct expression *
infix_member_access(struct expression *left)
{
        consume('.');
        
        struct expression *e = mkexpr();
        e->object = left;

        expect(TOKEN_IDENTIFIER);

        if (token(1)->type != '(') {
                e->type = EXPRESSION_MEMBER_ACCESS;
                e->member_name = sclone(tok()->identifier);
                consume(TOKEN_IDENTIFIER);
                return e;
        }

        new_dot_lambda = true;
        e->type = EXPRESSION_METHOD_CALL;
        e->method_name = sclone(tok()->identifier);
        consume(TOKEN_IDENTIFIER);
        vec_init(e->method_args);

        consume('(');

        lex_ctx = LEX_PREFIX;

        if (tok()->type == ')')
                goto end;
        else
                vec_push(e->method_args, parse_expr(0));

        while (tok()->type == ',') {
                next();
                vec_push(e->method_args, parse_expr(0));
        }

end:
        consume(')');
        new_dot_lambda = false;
        return e;
}

static struct expression *
infix_squiggly_arrow(struct expression *left)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_VIEW_PATTERN;

        consume(TOKEN_SQUIGGLY_ARROW);

        e->left = left;
        e->right = parse_expr(0);

        return e;
}

static struct expression *
infix_arrow_function(struct expression *left)
{
        /*
         * If it's not a list, we need to make sure it's a valid definition lvalue.
         * For example, ([a, b, c] -> a + b * c) is allowed, but not ([a[0], b, c] -> b * c).
         */
        if (left->type != EXPRESSION_LIST)
                left = definition_lvalue(left);

        consume(TOKEN_ARROW);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_FUNCTION;
        e->rest = false;
        e->name = NULL;
        vec_init(e->params);


        /*
         * If the left hand side consists of nothing more than identifiers, we can use a regular function.
         *
         * e.g., (a, b, c) -> (a + b + c) can become function (a, b, c) { return (a + b + c); }
         *
         * but
         *
         *       ([a, b], [c, d]) -> (a + b + c + d) must be transformed into something more complicated.
         */
        if (left->type == EXPRESSION_MATCH_REST || left->type == EXPRESSION_IDENTIFIER || (left->type == EXPRESSION_LIST && left->only_identifiers)) {
                if (left->type == EXPRESSION_MATCH_REST) {
                        vec_push(e->params, left->identifier);
                        e->params.count -= 1;
                        e->rest = true;
                } else if (left->type == EXPRESSION_IDENTIFIER) {
                        vec_push(e->params, left->identifier);
                } else {
                        for (int i = 0; i < left->es.count; ++i) {
                                vec_push(e->params, left->es.items[i]->identifier);
                        }
                        if (left->es.items[left->es.count - 1]->type == EXPRESSION_MATCH_REST) {
                                e->params.count -= 1;
                                e->rest = true;
                        }
                }

                e->body = mkret(parse_expr(0));

                return e;
        } else {
                /*
                 * In this case, we will generate unique names for each parameter, and then use
                 * 'let' statements in the function body to unpack them.
                 */
                struct statement *body = mkstmt();
                body->type = STATEMENT_BLOCK;
                vec_init(body->statements);

                if (left->type == EXPRESSION_LIST) {
                        for (int i = 0; i < left->es.count; ++i) {
                                char *name = gensym();
                                vec_push(e->params, name);
                                vec_push(body->statements, mkdef(left->es.items[i], name));
                        }
                } else {
                        char *name = gensym();
                        vec_push(e->params, name);
                        vec_push(body->statements, mkdef(left, name));
                }

                vec_push(body->statements, mkret(parse_expr(0)));

                e->body = body;

                return e;
        }
}

static struct expression *
infix_conditional(struct expression *left)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_CONDITIONAL;

        e->then = left;
        consume_keyword(KEYWORD_IF);
        e->cond = parse_expr(0);
        consume_keyword(KEYWORD_ELSE);<commit_msg>refactor infix arrow function parser<commit_after>#include <setjmp.h>
#include <stdarg.h>
#include <string.h>
#include <stdbool.h>
#include <assert.h>
#include <errno.h>
#include <math.h>
#include <stdnoreturn.h>

#include "vec.h"
#include "token.h"
#include "test.h"
#include "ast.h"
#include "util.h"
#include "alloc.h"
#include "lex.h"
#include "operators.h"
#include "value.h"
#include "log.h"
#include "vm.h"

#define BINARY_OPERATOR(name, token, prec, right_assoc) \
        static struct expression * \
        infix_ ## name(struct expression *left) \
        { \
                consume(TOKEN_ ## token); \
                struct expression *e = mkexpr(); \
                e->type = EXPRESSION_ ## token; \
                e->left = left; \
                e->right = parse_expr(prec - (right_assoc ? 1 : 0)); \
                return e; \
        } \

#define BINARY_LVALUE_OPERATOR(name, token, prec, right_assoc) \
        static struct expression * \
        infix_ ## name(struct expression *left) \
        { \
                consume(TOKEN_ ## token); \
                struct expression *e = mkexpr(); \
                e->type = EXPRESSION_ ## token; \
                e->target = assignment_lvalue(left); \
                e->value = parse_expr(prec - (right_assoc ? 1 : 0)); \
                return e; \
        } \

#define PREFIX_OPERATOR(name, token, prec) \
        static struct expression * \
        prefix_ ## name(void) \
        { \
                consume(TOKEN_ ## token); \
                struct expression *e = mkexpr(); \
                e->type = EXPRESSION_PREFIX_ ## token; \
                e->operand = parse_expr(prec); \
                return e; \
        } \

#define PREFIX_LVALUE_OPERATOR(name, token, prec) \
        static struct expression * \
        prefix_ ## name(void) \
        { \
                consume(TOKEN_ ## token); \
                struct expression *e = mkexpr(); \
                e->type = EXPRESSION_PREFIX_ ## token; \
                e->operand = assignment_lvalue(parse_expr(prec)); \
                return e; \
        } \

typedef struct expression *parse_fn();

enum {
        MAX_ERR_LEN  = 2048
};

/*
 * Meaningful names for different contexts in which we can parse lvalues.
 *
 * In a let declaration, the lvalue must be '='.
 * In a for-each loop, it must be followed by 'in'.
 *
 * Passing these values to parse_definition_lvalue determines its behaviour.
 */
enum {
        LV_LET,
        LV_EACH,
        LV_SUB,
        LV_ANY
};

static jmp_buf jb;
static char errbuf[MAX_ERR_LEN + 1];

static vec(struct token) tokens;
static int tokidx = 0;
enum lex_context lex_ctx = LEX_PREFIX;

static int depth;
static int dot_lambda_depth;
static bool new_dot_lambda;
static bool NoEquals = false;

static char const *filename;

static struct statement BREAK_STATEMENT    = { .type = STATEMENT_BREAK,    .loc = {42, 42}  };
static struct statement CONTINUE_STATEMENT = { .type = STATEMENT_CONTINUE, .loc = {42, 42}  };
static struct statement NULL_STATEMENT     = { .type = STATEMENT_NULL,     .loc = {42, 42}  };

static struct statement *
parse_statement(void);

static struct expression *
parse_expr(int);

static struct statement *
parse_match_statement(void);

static struct statement *
parse_let_definition(void);

static struct expression *
parse_target_list(void);

static struct statement *
parse_block(void);

static struct expression *
assignment_lvalue(struct expression *e);

static struct expression *
definition_lvalue(struct expression *e);

static struct expression *
infix_member_access(struct expression *e);

inline static struct token *
tok(void);

/*
 * Get a unique identifier name.
 * This sucks.
 */
char *
gensym(void)
{
        static int sym = 0;
        char buf[24];

        sprintf(buf, ":%d", sym++);

        return sclone(buf);
}

inline static struct expression *
mkexpr(void)
{
        struct expression *e = alloc(sizeof *e);
        e->loc = tok()->loc;
        return e;
}

inline static struct statement *
mkstmt(void)
{
        struct statement *s = alloc(sizeof *s);
        s->loc = tok()->loc;
        return s;
}

inline static struct statement *
mkret(struct expression *value)
{
        struct statement *s = mkstmt();
        s->type = STATEMENT_RETURN;
        vec_init(s->returns);
        vec_push(s->returns, value);
        return s;
}

inline static struct statement *
mkdef(struct expression *lvalue, char *name)
{
        struct expression *value = mkexpr();
        value->type = EXPRESSION_IDENTIFIER;
        value->identifier = name;
        value->module = NULL;

        struct statement *s = mkstmt();
        s->type = STATEMENT_DEFINITION;
        s->target = lvalue;
        s->value = value;

        return s;
}

noreturn static void
error(char const *fmt, ...)
{
        if (fmt == NULL) {
                strcpy(errbuf, lex_error());
        } else {
                va_list ap;
                va_start(ap, fmt);

                int n = 0;

                char *err = errbuf;
                n += sprintf(err, "ParseError: %s:%d:%d: ", filename, tok()->loc.line + 1, tok()->loc.col + 1);
                vsnprintf(err + n, MAX_ERR_LEN - n, fmt, ap);

                va_end(ap);
        }

        LOG("Parse Error: %s", errbuf);

        longjmp(jb, 1);
}

inline static void
skip(int n)
{
        tokidx += n;
}

inline static void
next(void)
{
        skip(1);
}

inline static struct token *
token(int i)
{
        struct token t;
        while (tokens.count <= i + tokidx) {
                t = lex_token(lex_ctx);
                if (t.type == TOKEN_ERROR) {
                        error(NULL);
                }
                vec_push(tokens, t);
        }

        return &tokens.items[tokidx + i];
}

inline static struct token *
tok(void)
{
        return token(0);
}

/*
 * Push a token into the token stream, so that it will be returned by the next call
 * to tok().
 */
inline static void
unconsume(int type)
{
        vec_insert(tokens, ((struct token){ .type = type, .loc = { 42, 42 } }), tokidx);
}

static void
expect(int type)
{
        if (tok()->type != type) {
                error("expected %s but found %s", token_show_type(type), token_show(tok()));
        }
}


static void
expect_keyword(int type)
{
        if (tok()->type != TOKEN_KEYWORD || tok()->keyword != type) {
                error("expected %s but found %s", token_show(&(struct token){ .type = TOKEN_KEYWORD, .keyword = type }), token_show(tok()));
        }
}

inline static void
consume(int type)
{
        expect(type);
        tokidx += 1;
}

inline static void
consume_keyword(int type)
{
        expect_keyword(type);
        tokidx += 1;

}

/* * * * | prefix parsers | * * * */
static struct expression *
prefix_integer(void)
{
        expect(TOKEN_INTEGER);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_INTEGER;
        e->integer = tok()->integer;

        consume(TOKEN_INTEGER);

        return e;
}

static struct expression *
prefix_real(void)
{
        expect(TOKEN_REAL);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_REAL;
        e->real = tok()->real;

        consume(TOKEN_REAL);

        return e;
}

static struct expression *
prefix_string(void)
{
        expect(TOKEN_STRING);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_STRING;
        e->string = tok()->string;

        consume(TOKEN_STRING);

        return e;
}

static struct expression *
prefix_special_string(void)
{
        expect(TOKEN_SPECIAL_STRING);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_SPECIAL_STRING;
        vec_init(e->expressions);
        
        e->strings.items = tok()->strings.items;
        e->strings.count = tok()->strings.count;

        char **exprs = tok()->expressions.items;
        int count = tok()->expressions.count;

        struct location *locs = tok()->locations.items;

        consume(TOKEN_SPECIAL_STRING);

        for (int i = 0; i < count; ++i) {
                lex_start(exprs[i]);
                vec_push(e->expressions, parse_expr(0));
                if (tok()->type != TOKEN_END) {
                        /*
                         * This results in a more accurate location in the error message.
                         */
                        tok()->loc = locs[i];
                        error("expression in interpolated string has trailing tokens: %s", exprs[i]);
                }
                consume(TOKEN_END);
                lex_end();
        }

        return e;
}

static struct expression *
prefix_dollar(void)
{
        consume('$');

        struct expression *e = mkexpr();

        expect(TOKEN_IDENTIFIER);

        e->type = EXPRESSION_MATCH_NOT_NIL;
        e->identifier = tok()->identifier;
        e->module = tok()->module;

        if (e->module != NULL)
                error("unpexpected module in lvalue");

        consume(TOKEN_IDENTIFIER);

        return e;
}

static struct expression *
prefix_identifier(void)
{
        expect(TOKEN_IDENTIFIER);

        struct expression *e = mkexpr();

        e->type = EXPRESSION_IDENTIFIER;
        e->identifier = tok()->identifier;
        e->module = tok()->module;

        consume(TOKEN_IDENTIFIER);

        return e;
}

static struct expression *
prefix_function(void)
{
        consume_keyword(KEYWORD_FUNCTION);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_FUNCTION;
        e->rest = false;
        vec_init(e->params);

        if (tok()->type == TOKEN_IDENTIFIER) {
                e->name = tok()->identifier;
                next();
        } else {
                e->name = NULL;
        }

        consume('(');

        while (tok()->type != ')') {
                bool rest = false;
                if (tok()->type == TOKEN_STAR) {
                        rest = true;
                        next();
                }

                expect(TOKEN_IDENTIFIER);
                vec_push(e->params, sclone(tok()->identifier));
                next();

                if (rest) {
                        e->params.count -= 1;
                        e->rest = 1;
                        expect(')');
                } else if (tok()->type == ',') {
                        next();
                }
        }

        consume(')');

        e->body = parse_statement();

        return e;
}

static struct expression *
prefix_star(void)
{
        consume(TOKEN_STAR);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_MATCH_REST;

        expect(TOKEN_IDENTIFIER);
        e->identifier = tok()->identifier;
        if (tok()->module != NULL)
                error("unexpected module qualifier in lvalue");

        consume(TOKEN_IDENTIFIER);

        return e;
}

static struct expression *
prefix_match(void)
{
        consume_keyword(KEYWORD_MATCH);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_MATCH;

        e->subject = parse_expr(0);

        consume('{');

        vec_init(e->patterns);
        vec_init(e->conds);
        vec_init(e->thens);

        vec_push(e->patterns, parse_expr(0));
        if (tok()->type == TOKEN_BIT_OR) {
                next();
                vec_push(e->conds, parse_expr(0));
        } else {
                vec_push(e->conds, NULL);
        }

        consume(TOKEN_FAT_ARROW);
        vec_push(e->thens, parse_expr(0));

        while (tok()->type == ',') {
                next();
                vec_push(e->patterns, parse_expr(0));
                if (tok()->type == TOKEN_BIT_OR) {
                        next();
                        vec_push(e->conds, parse_expr(0));
                } else {
                        vec_push(e->conds, NULL);
                }

                consume(TOKEN_FAT_ARROW);
                vec_push(e->thens, parse_expr(0));
        }

        consume('}');

        return e;
}

static struct expression *
prefix_parenthesis(void)
{
        /*
         * This can either be a plain old parenthesized expression, e.g., (4 + 4)
         * or it can be an identifier list for an arrow function, e.g., (a, b, c).
         */

        consume('(');

        /*
         * () is an empty identifier list.
         */
        if (tok()->type == ')') {
                next();
                struct expression *list = mkexpr();
                list->type = EXPRESSION_LIST;
                list->only_identifiers = true;
                vec_init(list->es);
                return list;
        }

        struct expression *e = parse_expr(0);

        if (tok()->type == ',') {
                struct expression *list = mkexpr();
                list->only_identifiers = true;

                /*
                 * It _must_ be an identifier list.
                 */
                if (e->type != EXPRESSION_IDENTIFIER && e->type != EXPRESSION_MATCH_REST) {
                        list->only_identifiers = false;
                }

                list->type = EXPRESSION_LIST;
                vec_init(list->es);
                vec_push(list->es, e);

                while (tok()->type == ',') {
                        next();
                        struct expression *e = parse_expr(0);
                        if (e->type == EXPRESSION_MATCH_REST) {
                                expect(')');
                        } else if (e->type != EXPRESSION_IDENTIFIER) {
                                list->only_identifiers = false;
                        }
                        vec_push(list->es, e);
                }

                consume(')');

                return list;
        } else {
                consume(')');
                return e;
        }
}

static struct expression *
prefix_true(void)
{
        consume_keyword(KEYWORD_TRUE);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_BOOLEAN;
        e->boolean = true;

        return e;
}

static struct expression *
prefix_false(void)
{
        consume_keyword(KEYWORD_FALSE);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_BOOLEAN;
        e->boolean = false;

        return e;
}

static struct expression *
prefix_nil(void)
{
        consume_keyword(KEYWORD_NIL);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_NIL;

        return e;
}

static struct expression *
prefix_regex(void)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_REGEX;
        e->regex = tok()->regex;
        e->extra = tok()->extra;
        e->pattern = tok()->pattern;

        consume(TOKEN_REGEX);

        return e;
}


static struct expression *
prefix_array(void)
{
        consume('[');

        struct expression *e = mkexpr();
        e->type = EXPRESSION_ARRAY;
        vec_init(e->elements);

        if (tok()->type == ']') {
                next();
                return e;
        } else {
                vec_push(e->elements, parse_expr(0));
        }

        while (tok()->type != ']') {
                if (tok()->type == TOKEN_KEYWORD && tok()->keyword == KEYWORD_FOR) {
                        next();
                        e->type = EXPRESSION_ARRAY_COMPR;
                        e->compr.pattern = parse_target_list();
                        consume_keyword(KEYWORD_IN);
                        e->compr.iter = parse_expr(0);
                        if (tok()->type == TOKEN_BIT_OR) {
                                next();
                                e->compr.cond = parse_expr(0);
                        } else {
                                e->compr.cond = NULL;
                        }
                        expect(']');
                } else {
                        consume(',');
                        vec_push(e->elements, parse_expr(0));
                }
        }

        consume(']');

        return e;
}

static struct expression *
prefix_tick(void)
{
        consume('`');

        expect(TOKEN_IDENTIFIER);

        if (tok()->module != NULL)
                error("unexpected module qualifier in ` pattern");

        struct expression *e = mkexpr();

        e->type = EXPRESSION_TICK;
        e->identifier = tok()->identifier;
        e->module = NULL;

        next();

        return e;
}

static struct expression *
prefix_dot(void)
{
        if (dot_lambda_depth > 0 && !new_dot_lambda) {
                struct expression *e = mkexpr();
                e->type = EXPRESSION_IDENTIFIER;
                e->identifier = "<object>";
                e->module = NULL;
                return infix_member_access(e);
        } else {
                new_dot_lambda = false;
                ++dot_lambda_depth;
                unconsume(TOKEN_ARROW);
                unconsume(TOKEN_IDENTIFIER);
                tok()->identifier = "<object>";
                tok()->module = NULL;
                return prefix_identifier();
        }
}

static struct expression *
prefix_incrange(void)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_DOT_DOT_DOT;

        struct expression *zero = mkexpr();
        zero->type = EXPRESSION_INTEGER;
        zero->integer = 0;

        consume(TOKEN_DOT_DOT_DOT);

        e->left = zero;
        e->right = parse_expr(0);

        return e;
}

static struct expression *
prefix_range(void)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_DOT_DOT;

        struct expression *zero = mkexpr();
        zero->type = EXPRESSION_INTEGER;
        zero->integer = 0;

        consume(TOKEN_DOT_DOT);

        e->left = zero;
        e->right = parse_expr(0);

        return e;
}

static struct expression *
prefix_hash(void)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_IDENTIFIER;
        e->identifier = "#";
        e->module = NULL;

        consume('#');

        return e;
}

static struct expression *
prefix_implicit_lambda(void)
{
        consume(TOKEN_BIT_OR);

        struct expression *f = mkexpr();
        f->type = EXPRESSION_FUNCTION;
        f->name = NULL;
        vec_init(f->params);
        vec_push(f->params, "#");

        struct expression *e = parse_expr(0);
        f->body = mkret(e);

        consume(TOKEN_BIT_OR);

        return f;
}

static struct expression *
prefix_object(void)
{
        consume('{');

        struct expression *e = mkexpr();
        e->type = EXPRESSION_DICT;

        vec_init(e->keys);
        vec_init(e->values);

        if (tok()->type == '}') {
                next();
                return e;
        } else {
                vec_push(e->keys, parse_expr(0));
                if (tok()->type == ':') {
                        next();
                        vec_push(e->values, parse_expr(0));
                } else {
                        vec_push(e->values, NULL);
                }
        }

        while (tok()->type != '}') {
                if (tok()->type == TOKEN_KEYWORD && tok()->keyword == KEYWORD_FOR) {
                        next();
                        e->type = EXPRESSION_DICT_COMPR;
                        e->dcompr.pattern = parse_target_list();
                        consume_keyword(KEYWORD_IN);
                        e->dcompr.iter = parse_expr(0);
                        if (tok()->type == TOKEN_BIT_OR) {
                                next();
                                e->dcompr.cond = parse_expr(0);
                        } else {
                                e->dcompr.cond = NULL;
                        }
                        expect('}');
                } else {
                        consume(',');
                        vec_push(e->keys, parse_expr(0));
                        if (tok()->type == ':') {
                                next();
                                vec_push(e->values, parse_expr(0));
                        } else {
                                vec_push(e->values, NULL);
                        }
                }
        }

        consume('}');

        return e;
}

PREFIX_OPERATOR(at,    AT,    9)
PREFIX_OPERATOR(minus, MINUS, 9)
PREFIX_OPERATOR(bang,  BANG,  10)

PREFIX_LVALUE_OPERATOR(inc,   INC,   6)
PREFIX_LVALUE_OPERATOR(dec,   DEC,   6)
/* * * * | end of prefix parsers | * * * */

/* * * * | infix parsers | * * * */
static struct expression *
infix_function_call(struct expression *left)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_FUNCTION_CALL;
        e->function = left;
        vec_init(e->args);

        consume('(');

        lex_ctx = LEX_PREFIX;

        if (tok()->type == ')') {
                next();
                return e;
        } else {
                vec_push(e->args, parse_expr(0));
        }

        while (tok()->type == ',') {
                next();
                vec_push(e->args, parse_expr(0));
        }

        consume(')');

        return e;
}

static struct expression *
infix_eq(struct expression *left)
{
        consume(TOKEN_EQ);
        struct expression *e = mkexpr();
        e->type = EXPRESSION_EQ;
        e->target = assignment_lvalue(left);
        if (left->type == EXPRESSION_LIST) {
                e->value = parse_expr(-1);
        } else {
                e->value = parse_expr(1);
        }

        return e;
}

static struct expression *
infix_user_op(struct expression *left)
{
        struct expression *e = mkexpr();

        e->type = EXPRESSION_FUNCTION_CALL;
        e->function = mkexpr();
        e->function->type = EXPRESSION_IDENTIFIER;
        e->function->identifier = tok()->identifier;
        e->function->module = NULL;
        consume(TOKEN_USER_OP);

        vec_init(e->args);
        vec_push(e->args, left);
        vec_push(e->args, parse_expr(2));

        return e;
}

static struct expression *
infix_list(struct expression *left)
{

        struct expression *e = mkexpr();
        e->type = EXPRESSION_LIST;
        vec_init(e->es);
        vec_push(e->es, left);

        NoEquals = true;

        while (tok()->type == ',') {
                next();
                vec_push(e->es, parse_expr(1));
        }

        NoEquals = false;

        return e;
}

static struct expression *
infix_subscript(struct expression *left)
{
        consume('[');
        
        struct expression *e = mkexpr();
        e->type = EXPRESSION_SUBSCRIPT;
        e->container = left;
        e->subscript = parse_expr(0);

        consume(']');

        return e;
}

static struct expression *
infix_member_access(struct expression *left)
{
        consume('.');
        
        struct expression *e = mkexpr();
        e->object = left;

        expect(TOKEN_IDENTIFIER);

        if (token(1)->type != '(') {
                e->type = EXPRESSION_MEMBER_ACCESS;
                e->member_name = sclone(tok()->identifier);
                consume(TOKEN_IDENTIFIER);
                return e;
        }

        new_dot_lambda = true;
        e->type = EXPRESSION_METHOD_CALL;
        e->method_name = sclone(tok()->identifier);
        consume(TOKEN_IDENTIFIER);
        vec_init(e->method_args);

        consume('(');

        lex_ctx = LEX_PREFIX;

        if (tok()->type == ')')
                goto end;
        else
                vec_push(e->method_args, parse_expr(0));

        while (tok()->type == ',') {
                next();
                vec_push(e->method_args, parse_expr(0));
        }

end:
        consume(')');
        new_dot_lambda = false;
        return e;
}

static struct expression *
infix_squiggly_arrow(struct expression *left)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_VIEW_PATTERN;

        consume(TOKEN_SQUIGGLY_ARROW);

        e->left = left;
        e->right = parse_expr(0);

        return e;
}

static struct expression *
infix_arrow_function(struct expression *left)
{
        /*
         * If it's not a list, we need to make sure it's a valid definition lvalue.
         * For example, ([a, b, c] -> a + b * c) is allowed, but not ([a[0], b, c] -> b * c).
         */
        if (left->type != EXPRESSION_LIST)
                left = definition_lvalue(left);

        consume(TOKEN_ARROW);

        struct expression *e = mkexpr();
        e->type = EXPRESSION_FUNCTION;
        e->rest = false;
        e->name = NULL;
        vec_init(e->params);

        if (left->type != EXPRESSION_LIST) {
                struct expression *l = mkexpr();
                l->type = EXPRESSION_LIST;
                vec_init(l->es);
                vec_push(l->es, left);
                left = l;
        }

        struct statement *body = mkstmt();
        body->type = STATEMENT_BLOCK;
        vec_init(body->statements);

        for (int i = 0; i < left->es.count; ++i) {
                struct expression *p = left->es.items[i];
                if (p->type == EXPRESSION_IDENTIFIER) {
                        vec_push(e->params, p->identifier);
                } else if (p->type == EXPRESSION_MATCH_REST) {
                        vec_push(e->params, p->identifier);
                        e->params.count -= 1;
                        e->rest = true;
                } else {
                        char *name = gensym();
                        vec_push(e->params, name);
                        vec_push(body->statements, mkdef(p, name));
                }
        }

        struct statement *ret = mkret(parse_expr(0));

        if (body->statements.count == 0) {
                free(body);
                e->body = ret;
        } else {
                vec_push(body->statements, ret);
                e->body = body;
        }

        return e;
}

static struct expression *
infix_conditional(struct expression *left)
{
        struct expression *e = mkexpr();
        e->type = EXPRESSION_CONDITIONAL;

        e->then = left;
        consume_keyword(KEYWORD_IF);
        e->cond = parse_expr(0);
        consume_keyword(KEYWORD_ELSE);
{'fraction_non_alphanumeric': 0.12739998108389294, 'fraction_numerical': 0.003499479807055708, 'mean_word_length': 1.5473907386883825, 'pattern_counts': {'":': 2, '<': 30, '<?xml version=': 0, '>': 572, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2164641', 'n_tokens_mistral': 16067, 'n_tokens_neox': 15198, 'n_words': 3429}