// Logic defining ASTs (abstract syntax trees) to represent code

#pragma once

#include <err.h>
#include <gc.h>
#include <stdbool.h>
#include <stdint.h>

#include "stdlib/datatypes.h"
#include "stdlib/files.h"
#include "util.h"

#define NewAST(_file, _start, _end, ast_tag, ...)                                                                      \
    (new (ast_t, .file = _file, .start = _start, .end = _end, .tag = ast_tag, .__data.ast_tag = {__VA_ARGS__}))
#define NewTypeAST(_file, _start, _end, ast_tag, ...)                                                                  \
    (new (type_ast_t, .file = _file, .start = _start, .end = _end, .tag = ast_tag, .__data.ast_tag = {__VA_ARGS__}))
#define FakeAST(ast_tag, ...) (new (ast_t, .tag = ast_tag, .__data.ast_tag = {__VA_ARGS__}))
#define WrapAST(ast, ast_tag, ...)                                                                                     \
    (new (ast_t, .file = (ast)->file, .start = (ast)->start, .end = (ast)->end, .tag = ast_tag,                        \
          .__data.ast_tag = {__VA_ARGS__}))
#define TextAST(ast, _str) WrapAST(ast, TextLiteral, .str = GC_strdup(_str))
#define LiteralCode(code, ...)                                                                                         \
    new (ast_t, .tag = InlineCCode,                                                                                    \
         .__data.InlineCCode = {.chunks = new (ast_list_t, .ast = FakeAST(TextLiteral, code)), __VA_ARGS__})
#define WrapLiteralCode(_ast, code, ...)                                                                               \
    new (ast_t, .tag = InlineCCode, .file = (_ast)->file, .start = (_ast)->start, .end = (_ast)->end,                  \
         .__data.InlineCCode = {.chunks = new (ast_list_t, .ast = WrapAST(_ast, TextLiteral, code)), __VA_ARGS__})
#define Match(x, _tag)                                                                                                 \
    ((x)->tag == _tag ? &(x)->__data._tag                                                                              \
                      : (errx(1, __FILE__ ":%d This was supposed to be a " #_tag "\n", __LINE__), &(x)->__data._tag))
#define DeclareMatch(var, x, _tag) __typeof((x)->__data._tag) *var = Match(x, _tag)
#define BINARY_OPERANDS(ast)                                                                                           \
    ({                                                                                                                 \
        if (!is_binary_operation(ast)) errx(1, __FILE__ ":%d This is not a binary operation!", __LINE__);              \
        (ast)->__data.Plus;                                                                                            \
    })
#define UPDATE_OPERANDS(ast)                                                                                           \
    ({                                                                                                                 \
        if (!is_update_assignment(ast)) errx(1, __FILE__ ":%d This is not an update assignment!", __LINE__);           \
        (ast)->__data.PlusUpdate;                                                                                      \
    })

#define REVERSE_LIST(list)                                                                                             \
    do {                                                                                                               \
        __typeof(list) _prev = NULL;                                                                                   \
        __typeof(list) _next = NULL;                                                                                   \
        __typeof(list) _current = list;                                                                                \
        while (_current != NULL) {                                                                                     \
            _next = _current->next;                                                                                    \
            _current->next = _prev;                                                                                    \
            _prev = _current;                                                                                          \
            _current = _next;                                                                                          \
        }                                                                                                              \
        list = _prev;                                                                                                  \
    } while (0)

struct binding_s;
typedef struct type_ast_s type_ast_t;
typedef struct ast_s ast_t;
typedef struct {
    ast_t *lhs, *rhs;
} binary_operands_t;

typedef struct ast_list_s {
    ast_t *ast;
    struct ast_list_s *next;
} ast_list_t;

typedef struct arg_ast_s {
    file_t *file;
    const char *start, *end;
    const char *name, *alias;
    // The comments belonging to this argument, joined into one paragraph. The
    // compiler reads a parameter's as its description on the command line,
    // where that is what a description is.
    Text_t comment;
    // The same comments as source positions, spanning from where the previous
    // argument's line ended to where this one's does. The formatter reads
    // these instead, so that a comment written on three lines is written back
    // out on three lines rather than joined into one long one.
    const char *comments_start, *comments_end;
    type_ast_t *type;
    ast_t *value;
    struct arg_ast_s *next;
} arg_ast_t;

typedef struct match_clause_s {
    ast_t *pattern, *body;
    struct match_clause_s *next;
} match_clause_t;

typedef enum {
    UnknownTypeAST,
    VarTypeAST,
    PointerTypeAST,
    ListTypeAST,
    TableTypeAST,
    FunctionTypeAST,
    OptionalTypeAST,
    EnumTypeAST,
} type_ast_e;

typedef struct tag_ast_s {
    file_t *file;
    const char *start, *end;
    const char *name;
    arg_ast_t *fields;
    struct tag_ast_s *next;
    bool secret : 1, packed_bools : 1;
} tag_ast_t;

struct type_ast_s {
    type_ast_e tag;
    file_t *file;
    const char *start, *end;
    union {
        struct {
        } UnknownTypeAST;
        struct {
            const char *name;
        } VarTypeAST;
        struct {
            type_ast_t *pointed;
            bool is_stack : 1;
        } PointerTypeAST;
        struct {
            type_ast_t *item;
        } ListTypeAST;
        struct {
            type_ast_t *key, *value /* value may be NULL */;
            ast_t *default_value;
        } TableTypeAST;
        struct {
            arg_ast_t *args;
            type_ast_t *ret;
        } FunctionTypeAST;
        struct {
            type_ast_t *type;
        } OptionalTypeAST;
        struct {
            Text_t name;
            tag_ast_t *tags;
        } EnumTypeAST;
    } __data;
};

#define BINOP_CASES                                                                                                    \
    Power:                                                                                                             \
    case Multiply:                                                                                                     \
    case Divide:                                                                                                       \
    case FloorDivide:                                                                                                  \
    case Mod:                                                                                                          \
    case Mod1:                                                                                                         \
    case Plus:                                                                                                         \
    case Minus:                                                                                                        \
    case Concat:                                                                                                       \
    case LeftShift:                                                                                                    \
    case UnsignedLeftShift:                                                                                            \
    case RightShift:                                                                                                   \
    case UnsignedRightShift:                                                                                           \
    case Equals:                                                                                                       \
    case NotEquals:                                                                                                    \
    case LessThan:                                                                                                     \
    case LessThanOrEquals:                                                                                             \
    case GreaterThan:                                                                                                  \
    case GreaterThanOrEquals:                                                                                          \
    case Compare:                                                                                                      \
    case And:                                                                                                          \
    case Or:                                                                                                           \
    case Xor:                                                                                                          \
    case PowerUpdate:                                                                                                  \
    case MultiplyUpdate:                                                                                               \
    case DivideUpdate:                                                                                                 \
    case FloorDivideUpdate:                                                                                            \
    case ModUpdate:                                                                                                    \
    case Mod1Update:                                                                                                   \
    case PlusUpdate:                                                                                                   \
    case MinusUpdate:                                                                                                  \
    case ConcatUpdate:                                                                                                 \
    case LeftShiftUpdate:                                                                                              \
    case UnsignedLeftShiftUpdate:                                                                                      \
    case RightShiftUpdate:                                                                                             \
    case UnsignedRightShiftUpdate:                                                                                     \
    case AndUpdate:                                                                                                    \
    case OrUpdate:                                                                                                     \
    case XorUpdate
#define UPDATE_CASES                                                                                                   \
    PowerUpdate:                                                                                                       \
    case MultiplyUpdate:                                                                                               \
    case DivideUpdate:                                                                                                 \
    case FloorDivideUpdate:                                                                                            \
    case ModUpdate:                                                                                                    \
    case Mod1Update:                                                                                                   \
    case PlusUpdate:                                                                                                   \
    case MinusUpdate:                                                                                                  \
    case ConcatUpdate:                                                                                                 \
    case LeftShiftUpdate:                                                                                              \
    case UnsignedLeftShiftUpdate:                                                                                      \
    case RightShiftUpdate:                                                                                             \
    case UnsignedRightShiftUpdate:                                                                                     \
    case AndUpdate:                                                                                                    \
    case OrUpdate:                                                                                                     \
    case XorUpdate

typedef enum {
    Unknown = 0,
    None,
    Bool,
    Var,
    Int,
    Num,
    TextLiteral,
    TextJoin,
    Path,
    Declare,
    Assign,
    Power,
    Multiply,
    Divide,
    FloorDivide,
    Mod,
    Mod1,
    Plus,
    Minus,
    Concat,
    LeftShift,
    UnsignedLeftShift,
    RightShift,
    UnsignedRightShift,
    Equals,
    NotEquals,
    LessThan,
    LessThanOrEquals,
    GreaterThan,
    GreaterThanOrEquals,
    Compare,
    And,
    Or,
    Xor,
    PowerUpdate,
    MultiplyUpdate,
    DivideUpdate,
    FloorDivideUpdate,
    ModUpdate,
    Mod1Update,
    PlusUpdate,
    MinusUpdate,
    ConcatUpdate,
    LeftShiftUpdate,
    UnsignedLeftShiftUpdate,
    RightShiftUpdate,
    UnsignedRightShiftUpdate,
    AndUpdate,
    OrUpdate,
    XorUpdate,
    Not,
    Negative,
    HeapAllocate,
    StackReference,
    Min,
    Max,
    List,
    Table,
    TableEntry,
    Comprehension,
    FunctionDef,
    Lambda,
    ConvertDef,
    FunctionCall,
    MethodCall,
    RecordLiteral,
    Block,
    For,
    While,
    If,
    Match,
    Repeat,
    Reduction,
    Continue,
    Break,
    Pass,
    Defer,
    Return,
    StructDef,
    EnumDef,
    LangDef,
    Index,
    FieldAccess,
    NonOptional,
    DebugLog,
    Assert,
    Test,
    Use,
    InlineCCode,
    Embed,
    Serialize,
    Deserialize,
    ExplicitlyTyped,
    Metadata,
} ast_e;
#define NUM_AST_TAGS (ExplicitlyTyped + 1)

struct ast_s {
    ast_e tag;
    file_t *file;
    const char *start, *end;
    union {
        struct {
        } Unknown;
        struct {
            struct type_s *type;
        } None;
        struct {
            bool b;
        } Bool;
        struct {
            const char *name;
        } Var;
        struct {
            // The value the digits denote, and the digits themselves: the
            // literal as it was written, including its base, `_` separators and
            // sign. It is kept here rather than read back out of the node's
            // span, because a parenthesized term's span covers its parentheses
            // (see parse_parens), which are no part of the literal. NULL for a
            // literal the parser synthesized rather than read.
            Int_t i;
            const char *str;
        } Int;
        struct {
            // The exact value, computed at parse time: `3.15` is 63/20, not
            // the nearest double to it (3.14999999999999991...). Codegen picks
            // the tightest encoding this value admits, usually a
            // NUMBER_SMALL immediate, costing nothing at runtime. As with an
            // Int literal, `str` is the digits as written, along with any
            // `%`/`deg` suffix.
            Num_t n;
            const char *str;
        } Num;
        struct {
            Text_t text;
        } TextLiteral;
        struct {
            type_ast_t *lang;
            ast_list_t *children;
        } TextJoin;
        struct {
            const char *path;
        } Path;
        struct {
            ast_t *path; // A `Path` AST node
        } Embed;
        struct {
            ast_t *var;
            type_ast_t *type;
            ast_t *value;
            bool top_level : 1;
        } Declare;
        struct {
            ast_list_t *targets, *values;
        } Assign;
        binary_operands_t Power, Multiply, Divide, FloorDivide, Mod, Mod1, Plus, Minus, Concat, LeftShift,
            UnsignedLeftShift, RightShift, UnsignedRightShift, Equals, NotEquals, LessThan, LessThanOrEquals,
            GreaterThan, GreaterThanOrEquals, Compare, And, Or, Xor, PowerUpdate, MultiplyUpdate, DivideUpdate,
            FloorDivideUpdate, ModUpdate, Mod1Update, PlusUpdate, MinusUpdate, ConcatUpdate, LeftShiftUpdate,
            UnsignedLeftShiftUpdate, RightShiftUpdate, UnsignedRightShiftUpdate, AndUpdate, OrUpdate, XorUpdate;
        struct {
            ast_t *value;
        } Not, Negative, HeapAllocate, StackReference;
        struct {
            ast_t *lhs, *rhs, *key;
        } Min, Max;
        struct {
            ast_list_t *items;
        } List;
        struct {
            ast_t *default_value;
            ast_t *fallback;
            ast_list_t *entries;
        } Table;
        struct {
            ast_t *key, *value /* value may be NULL */;
        } TableEntry;
        struct {
            ast_list_t *vars;
            ast_t *at; // optional `at i` iteration-counter variable
            ast_t *expr, *filter;
            ast_list_t *iters; // one or more iterables (lockstep iteration when >1)
        } Comprehension;
        struct {
            ast_t *name;
            arg_ast_t *args;
            type_ast_t *ret_type;
            ast_t *body;
            ast_t *cache;
            bool is_inline;
            Text_t comment; // Doc comment lines directly above the definition
        } FunctionDef;
        struct {
            arg_ast_t *args;
            type_ast_t *ret_type;
            ast_t *body;
            ast_t *cache;
            bool is_inline;
        } ConvertDef;
        struct {
            arg_ast_t *args;
            type_ast_t *ret_type;
            ast_t *body;
            int64_t id;
        } Lambda;
        struct {
            ast_t *fn;
            arg_ast_t *args;
        } FunctionCall;
        struct {
            const char *name;
            ast_t *self;
            arg_ast_t *args;
        } MethodCall;
        // `Foo{...}` / `Baz.A{...}`: field-wise construction of a struct or an
        // enum variant. Distinct from FunctionCall so that parens on a type name
        // mean "call a convert/constructor function" and braces mean "build a
        // record". `type` is the Var or FieldAccess naming the struct/variant.
        struct {
            ast_t *type;
            arg_ast_t *args;
        } RecordLiteral;
        struct {
            ast_list_t *statements;
        } Block;
        struct {
            ast_list_t *vars;
            ast_t *at; // optional `at i` iteration-counter variable
            ast_list_t *iters; // one or more iterables (lockstep iteration when >1)
            ast_t *body, *empty;
        } For;
        struct {
            ast_t *condition, *body;
        } While;
        struct {
            ast_t *body;
        } Repeat;
        struct {
            ast_t *condition, *body, *else_body;
            bool postfix;
        } If;
        struct {
            ast_t *subject;
            match_clause_t *clauses;
            ast_t *else_body;
        } Match;
        struct {
            ast_t *iter, *key;
            ast_e op;
        } Reduction;
        struct {
            const char *target;
        } Continue, Break;
        struct {
        } Pass;
        struct {
            ast_t *body;
        } Defer;
        struct {
            ast_t *value;
        } Return;
        struct {
            const char *name;
            arg_ast_t *fields;
            ast_t *namespace;
            bool secret : 1, external : 1, opaque : 1, packed_bools : 1;
        } StructDef;
        struct {
            const char *name;
            tag_ast_t *tags;
            ast_t *namespace;
        } EnumDef;
        struct {
            const char *name;
            ast_t *namespace;
        } LangDef;
        struct {
            ast_t *indexed, *index;
        } Index;
        struct {
            ast_t *fielded;
            const char *field;
        } FieldAccess;
        struct {
            ast_t *value;
        } NonOptional;
        struct {
            ast_list_t *values;
        } DebugLog;
        struct {
            ast_t *expr, *message;
        } Assert;
        struct {
            const char *label; // test "this label"
            ast_t *body; // Block
            // At most one of these is non-NULL (both NULL => the test should pass). The string is the substring the
            // failure/error must contain; "" means "match any failure".
            const char *expected_failure; // `fails "..."`: expected runtime panic/abort
            const char *expected_compile_error; // `fails_compile "..."`: expected typecheck error
        } Test;
        struct {
            ast_t *var;
            const char *path;
            enum { USE_LOCAL, USE_PACKAGE, USE_SHARED_OBJECT, USE_HEADER, USE_C_CODE, USE_ASM } what;
        } Use;
        struct {
            ast_list_t *chunks;
            struct type_s *type;
            type_ast_t *type_ast;
        } InlineCCode;
        struct {
            ast_t *value;
        } Serialize;
        struct {
            type_ast_t *type_ast;
            ast_t *value;
        } Deserialize;
        struct {
            ast_t *ast;
            struct type_s *type;
        } ExplicitlyTyped;
        struct {
            Text_t key, value;
        } Metadata;
    } __data;
};

extern const int op_tightness[NUM_AST_TAGS];
extern const bool op_is_right_associative[NUM_AST_TAGS];

typedef struct {
    const char *method_name;
    const char *operator;
} binop_info_t;

extern const binop_info_t binop_info[NUM_AST_TAGS];

OptionalText_t ast_source(ast_t *ast);
double num_literal_double(ast_t *ast);
bool fold_num_constant(ast_t *ast, Num_t *out);

Text_t ast_to_sexp(ast_t *ast);
const char *ast_to_sexp_str(ast_t *ast);
Text_t type_ast_to_sexp(type_ast_t *ast);

PUREFUNC bool is_idempotent(ast_t *ast);
void visit_topologically(ast_list_t *ast, Closure_t fn);
CONSTFUNC bool is_update_assignment(ast_t *ast);
CONSTFUNC ast_e binop_tag(ast_e tag);
CONSTFUNC bool is_binary_operation(ast_t *ast);
PUREFUNC bool is_operation(ast_t *ast);
bool is_int_literal(ast_t *ast, Int_t *value);
PUREFUNC int expr_tightness(ast_t *ast);
CONSTFUNC bool absorbs_rhs(ast_e outer_op, int tightness);
CONSTFUNC bool absorbs_lhs(ast_e outer_op, int tightness);
typedef enum { VISIT_STOP, VISIT_PROCEED } visit_behavior_t;
void ast_visit(ast_t *ast, visit_behavior_t (*visitor)(ast_t *, void *), void *userdata);
void type_ast_visit(ast_t *ast, visit_behavior_t (*visitor)(type_ast_t *, void *), void *userdata);
OptionalText_t ast_metadata(ast_t *ast, const char *key);