// Some basic operations defined on AST nodes, mainly converting to
// strings for debugging.

#include <stdarg.h>

#include "ast.h"
#include "stdlib/datatypes.h"
#include "stdlib/integers.h"
#include "stdlib/num.h"
#include "stdlib/number.h"
#include "stdlib/optionals.h"
#include "stdlib/table.h"
#include "stdlib/text.h"

const int op_tightness[NUM_AST_TAGS] = {
    [Power] = 10,
    // Negation is an operator like any other, and it sits here: `-x ^ 2` is
    // `-(x ^ 2)`, while `-x * y` is `(-x) * y`.
    [Negative] = 9,
    [Multiply] = 8,
    [Divide] = 8,
    [FloorDivide] = 8,
    [Mod] = 8,
    [Mod1] = 8,
    [Plus] = 7,
    [Minus] = 7,
    [Concat] = 6,
    [LeftShift] = 5,
    [RightShift] = 5,
    [UnsignedLeftShift] = 5,
    [UnsignedRightShift] = 5,
    [Min] = 4,
    [Max] = 4,
    [Equals] = 3,
    [NotEquals] = 3,
    [LessThan] = 2,
    [LessThanOrEquals] = 2,
    [GreaterThan] = 2,
    [GreaterThanOrEquals] = 2,
    [Compare] = 2,
    [And] = 1,
    [Or] = 1,
    [Xor] = 1,
};

// Associativity, the companion to op_tightness above: everything groups to the
// left (10-3-2 is (10-3)-2) except exponentiation, which groups to the right
// (2^3^2 is 2^(3^2)).
const bool op_is_right_associative[NUM_AST_TAGS] = {
    [Power] = true,
};

// How tightly this expression's outermost operator binds: op_tightness for its
// tag, except that a literal the parser folded a leading `-` into (see
// negate_literal) binds like the negation it is written as -- `-2 ^ 2` is
// `-(2 ^ 2)`, and `(-2) ^ 2` needs its parentheses. That's a question about the
// written sign, not the value's: `-0` is zero, but a suffix still binds
// tighter than its `-`.
PUREFUNC int expr_tightness(ast_t *ast) {
    const char *str = ast->tag == Int ? Match(ast, Int)->str : ast->tag == Num ? Match(ast, Num)->str : NULL;
    if (str && str[0] == '-') return op_tightness[Negative];
    return op_tightness[ast->tag];
}

// Whether `outer_op` absorbs an expression of the given tightness on its right
// into itself, rather than leaving it for whatever encloses `outer_op` in
// turn. It absorbs the expression when it binds more tightly, or exactly as
// tightly if `outer_op` is right-associative, which is what groups
// `a ^ b ^ c` as `a ^ (b ^ c)`. The parser asks about an operator it has yet
// to parse the operands of (op_tightness[op]), and the formatter about a
// finished expression (expr_tightness).
//
// `outer_op` is Unknown for an expression with nothing around it. That absorbs
// every real operator, since Unknown's tightness of 0 is below all of them.
CONSTFUNC bool absorbs_rhs(ast_e outer_op, int tightness) {
    if (op_is_right_associative[outer_op]) return tightness >= op_tightness[outer_op];
    return tightness > op_tightness[outer_op];
}

// The same question for an expression on the *left* of `outer_op`, where the
// tie goes the other way: left-associativity is what groups `a - b - c` as
// `(a - b) - c`. Only the formatter needs to ask, since the parser builds its
// left-hand side out of terms, and it asks in order to know whether an operand
// has to be parenthesized to be read back the way it was written.
CONSTFUNC bool absorbs_lhs(ast_e outer_op, int tightness) {
    if (op_is_right_associative[outer_op]) return tightness > op_tightness[outer_op];
    return tightness >= op_tightness[outer_op];
}

const binop_info_t binop_info[NUM_AST_TAGS] = {
    [Power] = {"power", "^"},
    [PowerUpdate] = {"power", "^="},
    [Multiply] = {"times", "*"},
    [MultiplyUpdate] = {"times", "*="},
    [Divide] = {"divided_by", "/"},
    [DivideUpdate] = {"divided_by", "/="},
    [FloorDivide] = {"floor_divided_by", "//"},
    [FloorDivideUpdate] = {"floor_divided_by", "//="},
    [Mod] = {"modulo", "mod"},
    [ModUpdate] = {"modulo", "mod="},
    [Mod1] = {"modulo1", "mod1"},
    [Mod1Update] = {"modulo1", "mod1="},
    [Plus] = {"plus", "+"},
    [PlusUpdate] = {"plus", "+="},
    [Minus] = {"minus", "-"},
    [MinusUpdate] = {"minus", "-="},
    [Concat] = {"concatenated_with", "++"},
    [ConcatUpdate] = {"concatenated_with", "++="},
    [LeftShift] = {"left_shifted", "<<"},
    [LeftShiftUpdate] = {"left_shifted", "<<="},
    [RightShift] = {"right_shifted", ">>"},
    [RightShiftUpdate] = {"right_shifted", ">>="},
    [UnsignedLeftShift] = {"unsigned_left_shifted", "<<<"},
    [UnsignedLeftShiftUpdate] = {"unsigned_left_shifted", "<<<="},
    [UnsignedRightShift] = {"unsigned_right_shifted", ">>>"},
    [UnsignedRightShiftUpdate] = {"unsigned_right_shifted", ">>>="},
    [And] = {"bit_and", "and"},
    [AndUpdate] = {"bit_and", "and="},
    [Or] = {"bit_or", "or"},
    [OrUpdate] = {"bit_or", "or="},
    [Xor] = {"bit_xor", "xor"},
    [XorUpdate] = {"bit_xor", "xor="},
    [Equals] = {NULL, "=="},
    [NotEquals] = {NULL, "!="},
    [LessThan] = {NULL, "<"},
    [LessThanOrEquals] = {NULL, "<="},
    [GreaterThan] = {NULL, ">"},
    [GreaterThanOrEquals] = {NULL, ">="},
    [Compare] = {NULL, "<>"},
    [Min] = {NULL, "_min_"},
    [Max] = {NULL, "_max_"},
};

static Text_t ast_list_to_sexp(ast_list_t *asts);
static Text_t arg_list_to_sexp(arg_ast_t *args);
static Text_t arg_defs_to_sexp(arg_ast_t *args);
static Text_t match_clauses_to_sexp(match_clause_t *clauses);
static Text_t tags_to_sexp(tag_ast_t *tags);
static Text_t optional_sexp(const char *tag, ast_t *ast);
static Text_t optional_type_sexp(const char *tag, type_ast_t *ast);

static Text_t quoted_text(const char *text) {
    return Text$quoted(Text$from_str(text), false, Text("\""));
}

Text_t ast_list_to_sexp(ast_list_t *asts) {
    Text_t c = EMPTY_TEXT;
    for (; asts; asts = asts->next) {
        c = Texts(c, " ", ast_to_sexp(asts->ast));
    }
    return c;
}

// A `; name` flag written on a definition, for the s-expression below. These
// are part of what the source says and have to appear in it: `tomo format
// --verify` compares two of these to decide that formatting was faithful, and
// what it cannot see it cannot check.
static Text_t flag_sexp(const char *name, bool present) {
    return present ? Texts(" :", name) : EMPTY_TEXT;
}

Text_t arg_defs_to_sexp(arg_ast_t *args) {
    Text_t c = Text("(args");
    for (arg_ast_t *arg = args; arg; arg = arg->next) {
        c = Texts(c, " (arg ", arg->name ? quoted_text(arg->name) : Text("nil"),
                  arg->alias ? Texts(" :alias ", quoted_text(arg->alias)) : EMPTY_TEXT, " ",
                  type_ast_to_sexp(arg->type), " ", ast_to_sexp(arg->value), ")");
    }
    return Texts(c, ")");
}

Text_t arg_list_to_sexp(arg_ast_t *args) {
    Text_t c = EMPTY_TEXT;
    for (arg_ast_t *arg = args; arg; arg = arg->next) {
        assert(arg->value && !arg->type);
        if (arg->name) c = Texts(c, " :", arg->name);
        c = Texts(c, " ", ast_to_sexp(arg->value));
    }
    return c;
}

Text_t match_clauses_to_sexp(match_clause_t *clauses) {
    Text_t c = EMPTY_TEXT;
    for (; clauses; clauses = clauses->next) {
        c = Texts(c, " (case ", ast_to_sexp(clauses->pattern), " ", ast_to_sexp(clauses->body), ")");
    }
    return c;
}

Text_t tags_to_sexp(tag_ast_t *tags) {
    Text_t c = EMPTY_TEXT;
    for (; tags; tags = tags->next) {
        c = Texts(c, "(tag \"", tags->name, "\"", flag_sexp("secret", tags->secret),
                  flag_sexp("packed_bools", tags->packed_bools), " ", arg_defs_to_sexp(tags->fields), ")");
    }
    return c;
}

Text_t type_ast_to_sexp(type_ast_t *t) {
    if (!t) return Text("nil");

    switch (t->tag) {
#define T(type, ...)                                                                                                   \
    case type: {                                                                                                       \
        __typeof(t->__data.type) data = t->__data.type;                                                                \
        (void)data;                                                                                                    \
        return Texts(__VA_ARGS__);                                                                                     \
    }
        T(UnknownTypeAST, "(UnknownType)");
        T(VarTypeAST, "(VarType \"", data.name, "\")");
        T(PointerTypeAST, "(PointerType \"", data.is_stack ? "stack" : "heap", "\" ", type_ast_to_sexp(data.pointed),
          ")");
        T(ListTypeAST, "(ListType ", type_ast_to_sexp(data.item), ")");
        T(TableTypeAST, "(TableType ", type_ast_to_sexp(data.key), " ", type_ast_to_sexp(data.value),
          data.default_value ? Texts(" :default ", ast_to_sexp(data.default_value)) : EMPTY_TEXT, ")");
        T(FunctionTypeAST, "(FunctionType ", arg_defs_to_sexp(data.args), " ", type_ast_to_sexp(data.ret), ")");
        T(OptionalTypeAST, "(OptionalType ", type_ast_to_sexp(data.type), ")");
        T(EnumTypeAST, "(EnumType ", data.name, " ", tags_to_sexp(data.tags), ")");
#undef T
    default: return EMPTY_TEXT;
    }
}

Text_t optional_sexp(const char *name, ast_t *ast) {
    return ast ? Texts(" :", name, " ", ast_to_sexp(ast)) : EMPTY_TEXT;
}

Text_t optional_type_sexp(const char *name, type_ast_t *ast) {
    return ast ? Texts(" :", name, " ", type_ast_to_sexp(ast)) : EMPTY_TEXT;
}

Text_t ast_to_sexp(ast_t *ast) {
    if (!ast) return Text("nil");

    switch (ast->tag) {
#define T(type, ...)                                                                                                   \
    case type: {                                                                                                       \
        __typeof(ast->__data.type) data = ast->__data.type;                                                            \
        (void)data;                                                                                                    \
        return Texts(__VA_ARGS__);                                                                                     \
    }
        T(Unknown, "(Unknown)");
        T(None, "(None)");
        T(Bool, "(Bool ", data.b ? "yes" : "no", ")");
        T(Var, "(Var ", quoted_text(data.name), ")");
        T(Int, "(Int ", data.str ? quoted_text(data.str) : Int$value_as_text(data.i), ")");
        T(Num, "(Num ", data.str ? quoted_text(data.str) : Text$from_str(number_to_symbolic(data.n)), ")");
        T(TextLiteral, Text$quoted(data.text, false, Text("\"")));
        T(TextJoin, "(Text", data.lang ? Texts(" :lang ", type_ast_to_sexp(data.lang)) : EMPTY_TEXT,
          ast_list_to_sexp(data.children), ")");
        T(Path, "(Path ", quoted_text(data.path), ")");
        T(Embed, "(Embed ", ast_to_sexp(data.path), ")");
        T(Declare, "(Declare ", ast_to_sexp(data.var), " ", type_ast_to_sexp(data.type), " ", ast_to_sexp(data.value),
          ")");
        T(Assign, "(Assign (targets ", ast_list_to_sexp(data.targets), ") (values ", ast_list_to_sexp(data.values),
          "))");
#define BINOP(name) T(name, "(" #name " ", ast_to_sexp(data.lhs), " ", ast_to_sexp(data.rhs), ")")
        BINOP(Power);
        BINOP(PowerUpdate);
        BINOP(Multiply);
        BINOP(MultiplyUpdate);
        BINOP(Divide);
        BINOP(DivideUpdate);
        BINOP(FloorDivide);
        BINOP(FloorDivideUpdate);
        BINOP(Mod);
        BINOP(ModUpdate);
        BINOP(Mod1);
        BINOP(Mod1Update);
        BINOP(Plus);
        BINOP(PlusUpdate);
        BINOP(Minus);
        BINOP(MinusUpdate);
        BINOP(Concat);
        BINOP(ConcatUpdate);
        BINOP(LeftShift);
        BINOP(LeftShiftUpdate);
        BINOP(RightShift);
        BINOP(RightShiftUpdate);
        BINOP(UnsignedLeftShift);
        BINOP(UnsignedLeftShiftUpdate);
        BINOP(UnsignedRightShift);
        BINOP(UnsignedRightShiftUpdate);
        BINOP(And);
        BINOP(AndUpdate);
        BINOP(Or);
        BINOP(OrUpdate);
        BINOP(Xor);
        BINOP(XorUpdate);
        BINOP(Compare);
        BINOP(Equals);
        BINOP(NotEquals);
        BINOP(LessThan);
        BINOP(LessThanOrEquals);
        BINOP(GreaterThan);
        BINOP(GreaterThanOrEquals);
#undef BINOP
        T(Negative, "(Negative ", ast_to_sexp(data.value), ")");
        T(Not, "(Not ", ast_to_sexp(data.value), ")");
        T(HeapAllocate, "(HeapAllocate ", ast_to_sexp(data.value), ")");
        T(StackReference, "(StackReference ", ast_to_sexp(data.value), ")");
        T(Min, "(Min ", ast_to_sexp(data.lhs), " ", ast_to_sexp(data.rhs), optional_sexp("key", data.key), ")");
        T(Max, "(Max ", ast_to_sexp(data.lhs), " ", ast_to_sexp(data.rhs), optional_sexp("key", data.key), ")");
        T(List, "(List", ast_list_to_sexp(data.items), ")");
        T(Table, "(Table", optional_sexp("default", data.default_value), optional_sexp("fallback", data.fallback),
          ast_list_to_sexp(data.entries), ")");
        T(TableEntry, "(TableEntry ", ast_to_sexp(data.key), " ", ast_to_sexp(data.value), ")");
        T(Comprehension, "(Comprehension ", ast_to_sexp(data.expr), " (vars", ast_list_to_sexp(data.vars), ") (iters",
          ast_list_to_sexp(data.iters), ") ", optional_sexp("at", data.at), optional_sexp("filter", data.filter), ")");
        T(FunctionDef, "(FunctionDef ", ast_to_sexp(data.name), " ", arg_defs_to_sexp(data.args),
          optional_type_sexp("return", data.ret_type), flag_sexp("inline", data.is_inline),
          optional_sexp("cache", data.cache), " ", ast_to_sexp(data.body), ")");
        T(ConvertDef, "(ConvertDef ", arg_defs_to_sexp(data.args), " ", type_ast_to_sexp(data.ret_type),
          flag_sexp("inline", data.is_inline), optional_sexp("cache", data.cache), " ", ast_to_sexp(data.body), ")");
        T(Lambda, "(Lambda ", arg_defs_to_sexp(data.args), optional_type_sexp("return", data.ret_type), " ",
          ast_to_sexp(data.body), ")");
        T(FunctionCall, "(FunctionCall ", ast_to_sexp(data.fn), arg_list_to_sexp(data.args), ")");
        T(MethodCall, "(MethodCall ", ast_to_sexp(data.self), " ", quoted_text(data.name), arg_list_to_sexp(data.args),
          ")")
        T(RecordLiteral, "(RecordLiteral ", ast_to_sexp(data.type), arg_list_to_sexp(data.args), ")");
        T(Block, "(Block", ast_list_to_sexp(data.statements), ")");
        T(For, "(For (vars", ast_list_to_sexp(data.vars), ") (iters", ast_list_to_sexp(data.iters), ") ",
          optional_sexp("at", data.at), ast_to_sexp(data.body), " ", ast_to_sexp(data.empty), ")");
        T(While, "(While ", ast_to_sexp(data.condition), " ", ast_to_sexp(data.body), ")");
        T(Repeat, "(Repeat ", ast_to_sexp(data.body), ")");
        T(If, "(If ", ast_to_sexp(data.condition), " ", ast_to_sexp(data.body), optional_sexp("else", data.else_body),
          ")");
        T(Match, "(Match ", ast_to_sexp(data.subject), match_clauses_to_sexp(data.clauses),
          optional_sexp("else", data.else_body), ")");
        T(Reduction, "(Reduction ", quoted_text(binop_info[data.op].operator), " ", ast_to_sexp(data.key), " ",
          ast_to_sexp(data.iter), ")");
        T(Continue, "(Continue ", quoted_text(data.target), ")");
        T(Break, "(Break ", quoted_text(data.target), ")");
        T(Pass, "(Pass)");
        T(Defer, "(Defer ", ast_to_sexp(data.body), ")");
        T(Return, "(Return ", ast_to_sexp(data.value), ")");
        T(StructDef, "(StructDef \"", data.name, "\"", flag_sexp("secret", data.secret),
          flag_sexp("external", data.external), flag_sexp("opaque", data.opaque),
          flag_sexp("packed_bools", data.packed_bools), " ", arg_defs_to_sexp(data.fields), " ",
          ast_to_sexp(data.namespace), ")");
        T(EnumDef, "(EnumDef \"", data.name, "\" (tags ", tags_to_sexp(data.tags), ") ", ast_to_sexp(data.namespace),
          ")");
        T(LangDef, "(LangDef \"", data.name, "\" ", ast_to_sexp(data.namespace), ")");
        T(Index, "(Index ", ast_to_sexp(data.indexed), " ", ast_to_sexp(data.index), ")");
        T(FieldAccess, "(FieldAccess ", ast_to_sexp(data.fielded), " \"", data.field, "\")");
        T(NonOptional, "(NonOptional ", ast_to_sexp(data.value), ")");
        T(DebugLog, "(DebugLog ", ast_list_to_sexp(data.values), ")");
        T(Assert, "(Assert ", ast_to_sexp(data.expr), " ", optional_sexp("message", data.message), ")");
        T(Test, "(Test ", quoted_text(data.label),
          data.expected_failure ? Texts(" fails ", quoted_text(data.expected_failure)) : EMPTY_TEXT,
          data.expected_compile_error ? Texts(" fails_compile ", quoted_text(data.expected_compile_error)) : EMPTY_TEXT,
          " ", ast_to_sexp(data.body), ")");
        T(Use, "(Use ", optional_sexp("var", data.var), " ", quoted_text(data.path), ")");
        T(InlineCCode, "(InlineCCode ", ast_list_to_sexp(data.chunks), optional_type_sexp("type", data.type_ast), ")");
        T(Serialize, "(Serialize ", ast_to_sexp(data.value), ")");
        T(Deserialize, "(Deserialize ", optional_type_sexp("type", data.type_ast), " ", ast_to_sexp(data.value), ")");
        T(Metadata, "((Metadata ", Text$quoted(data.key, false, Text("\"")), " ",
          Text$quoted(data.value, false, Text("\"")), ")");
    default: errx(1, "S-expressions are not implemented for this AST");
#undef T
    }
}

// Constant-folds an arithmetic expression of numeric literals into the exact
// Num value it denotes, so `1/3` compiles to NUMBER_SMALL(1, 3) and
// `0.1 + 0.1` to NUMBER_SMALL(1, 5) rather than a runtime call chain. The
// compiler links the same number library the runtime uses, so "fold" here
// means "run the actual arithmetic", and the folded value is bit-identical to
// what the runtime would have produced, Euclidean division/modulus included.
//
// Deliberately narrow: literals, negation, and the binary operators, folding
// only when every step stays exact and rational. An error value (1/0: the
// runtime should report it, at the site, with its own message) or an
// irrational result (2^0.5: no constant-expression form exists) returns
// false and leaves the expression to compile the ordinary way.
bool fold_num_constant(ast_t *ast, Num_t *out) {
    switch (ast->tag) {
    case Num: *out = Match(ast, Num)->n; return true;
    case Int: *out = Num$from_int(Match(ast, Int)->i); return true;
    case Negative: {
        Num_t inner;
        if (!fold_num_constant(Match(ast, Negative)->value, &inner)) return false;
        *out = number_neg(inner);
        return true;
    }
    case Plus:
    case Minus:
    case Multiply:
    case Divide:
    case FloorDivide:
    case Mod:
    case Power: {
        binary_operands_t binop = BINARY_OPERANDS(ast);
        Num_t lhs, rhs;
        if (!fold_num_constant(binop.lhs, &lhs) || !fold_num_constant(binop.rhs, &rhs)) return false;
        Num_t result;
        switch (ast->tag) {
        case Plus: result = number_add(lhs, rhs); break;
        case Minus: result = number_sub(lhs, rhs); break;
        case Multiply: result = number_mul(lhs, rhs); break;
        case Divide: result = number_div(lhs, rhs); break;
        case FloorDivide: {
            // The Euclidean quotient, exactly as Num$floor_divided_by
            // computes it at runtime (see num.c).
            Num_t q = number_div(lhs, rhs);
            if (NUMBER_IS_ERROR(q)) return false;
            result = number_is_negative(rhs) ? number_ceil(q) : number_floor(q);
            break;
        }
        case Mod: {
            // Euclidean, matching Num$modulo: x - y*(x//y).
            Num_t q = number_div(lhs, rhs);
            if (NUMBER_IS_ERROR(q)) return false;
            q = number_is_negative(rhs) ? number_ceil(q) : number_floor(q);
            if (NUMBER_IS_ERROR(q)) return false;
            result = number_sub(lhs, number_mul(rhs, q));
            break;
        }
        default: result = number_pow(lhs, rhs); break;
        }
        if (NUMBER_IS_ERROR(result) || !number_is_rational(result)) return false;
        *out = result;
        return true;
    }
    default: return false;
    }
}

// The value of a numeric literal as a double, for the paths that compile one
// into a Float64/Float32 rather than a Num. Rounds the exact value once, so
// `90deg` is the nearest double to pi/180*90 rather than the accumulation of
// rounding a truncated pi/180 constant and then multiplying.
double num_literal_double(ast_t *ast) {
    return number_to_double(Match(ast, Num)->n);
}

const char *ast_to_sexp_str(ast_t *ast) {
    return Text$as_c_string(ast_to_sexp(ast));
}

OptionalText_t ast_source(ast_t *ast) {
    if (ast == NULL || ast->start == NULL || ast->end == NULL) return NONE_TEXT;
    return Text$from_strn(ast->start, (size_t)(ast->end - ast->start));
}

PUREFUNC bool is_idempotent(ast_t *ast) {
    switch (ast->tag) {
    case Int:
    case Bool:
    case Num:
    case Var:
    case None:
    case TextLiteral: return true;
    case Index: {
        DeclareMatch(index, ast, Index);
        return is_idempotent(index->indexed) && index->index != NULL && is_idempotent(index->index);
    }
    case FieldAccess: {
        DeclareMatch(access, ast, FieldAccess);
        return is_idempotent(access->fielded);
    }
    default: return false;
    }
}

void _visit_topologically(ast_t *ast, Table_t definitions, Table_t *visited, Closure_t fn) {
    void (*visit)(void *, ast_t *) = (void *)fn.fn;
    if (ast->tag == StructDef) {
        DeclareMatch(def, ast, StructDef);
        if (Table$str_get(*visited, def->name)) return;

        Table$str_set(visited, def->name, (void *)_visit_topologically);
        for (arg_ast_t *field = def->fields; field; field = field->next) {
            if (field->type && field->type->tag == VarTypeAST) {
                const char *field_type_name = Match(field->type, VarTypeAST)->name;
                ast_t *dependency = Table$str_get(definitions, field_type_name);
                if (dependency) {
                    _visit_topologically(dependency, definitions, visited, fn);
                }
            }
        }
        visit(fn.userdata, ast);
    } else if (ast->tag == EnumDef) {
        DeclareMatch(def, ast, EnumDef);
        if (Table$str_get(*visited, def->name)) return;

        Table$str_set(visited, def->name, (void *)_visit_topologically);
        for (tag_ast_t *tag = def->tags; tag; tag = tag->next) {
            for (arg_ast_t *field = tag->fields; field; field = field->next) {
                if (field->type && field->type->tag == VarTypeAST) {
                    const char *field_type_name = Match(field->type, VarTypeAST)->name;
                    ast_t *dependency = Table$str_get(definitions, field_type_name);
                    if (dependency) {
                        _visit_topologically(dependency, definitions, visited, fn);
                    }
                }
            }
        }
        visit(fn.userdata, ast);
    } else if (ast->tag == LangDef) {
        DeclareMatch(def, ast, LangDef);
        if (Table$str_get(*visited, def->name)) return;
        visit(fn.userdata, ast);
    } else {
        visit(fn.userdata, ast);
    }
}

void visit_topologically(ast_list_t *asts, Closure_t fn) {
    // Visit each top-level statement in topological order:
    // - 'use' statements first
    // - then typedefs
    //   - visiting typedefs' dependencies first
    // - then function/variable declarations

    Table_t definitions = EMPTY_TABLE;
    for (ast_list_t *stmt = asts; stmt; stmt = stmt->next) {
        if (stmt->ast->tag == StructDef) {
            DeclareMatch(def, stmt->ast, StructDef);
            Table$str_set(&definitions, def->name, stmt->ast);
        } else if (stmt->ast->tag == EnumDef) {
            DeclareMatch(def, stmt->ast, EnumDef);
            Table$str_set(&definitions, def->name, stmt->ast);
        } else if (stmt->ast->tag == LangDef) {
            DeclareMatch(def, stmt->ast, LangDef);
            Table$str_set(&definitions, def->name, stmt->ast);
        }
    }

    void (*visit)(void *, ast_t *) = (void *)fn.fn;
    Table_t visited = EMPTY_TABLE;
    // First: 'use' statements in order:
    for (ast_list_t *stmt = asts; stmt; stmt = stmt->next) {
        if (stmt->ast->tag == Use) visit(fn.userdata, stmt->ast);
    }
    // Then typedefs in topological order:
    for (ast_list_t *stmt = asts; stmt; stmt = stmt->next) {
        if (stmt->ast->tag == StructDef || stmt->ast->tag == EnumDef || stmt->ast->tag == LangDef)
            _visit_topologically(stmt->ast, definitions, &visited, fn);
    }
    // Then everything else in order:
    for (ast_list_t *stmt = asts; stmt; stmt = stmt->next) {
        if (!(stmt->ast->tag == StructDef || stmt->ast->tag == EnumDef || stmt->ast->tag == LangDef
              || stmt->ast->tag == Use)) {
            visit(fn.userdata, stmt->ast);
        }
    }
}

// Whether this expression is an operator application that expr_tightness
// describes the grouping of: the binary operators, negation, and a literal
// with a folded `-`.
PUREFUNC bool is_operation(ast_t *ast) {
    return expr_tightness(ast) > 0;
}

// An integer literal, seeing through a `-` written apart from its digits. `-2`
// folds the sign into the literal as it's parsed, but `- 2` and `-(2)` are a
// negation wrapping one, and they denote the same value, so anything that
// treats a literal specially, whether that's compiling it straight to a sized
// type or letting it take its type from the other operand, should treat all
// three alike. The value is written to `*value` when one is asked for.
bool is_int_literal(ast_t *ast, Int_t *value) {
    switch (ast->tag) {
    case Int:
        if (value) *value = Match(ast, Int)->i;
        return true;
    case Negative: {
        Int_t inner;
        if (!is_int_literal(Match(ast, Negative)->value, &inner)) return false;
        if (value) *value = Int$negative(inner);
        return true;
    }
    default: return false;
    }
}

CONSTFUNC bool is_binary_operation(ast_t *ast) {
    switch (ast->tag) {
    case Min:
    case Max:
    case BINOP_CASES: return true;
    default: return false;
    }
}

CONSTFUNC bool is_update_assignment(ast_t *ast) {
    switch (ast->tag) {
    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: return true;
    default: return false;
    }
}

CONSTFUNC ast_e binop_tag(ast_e tag) {
    switch (tag) {
    case PowerUpdate: return Power;
    case MultiplyUpdate: return Multiply;
    case DivideUpdate: return Divide;
    case FloorDivideUpdate: return FloorDivide;
    case ModUpdate: return Mod;
    case Mod1Update: return Mod1;
    case PlusUpdate: return Plus;
    case MinusUpdate: return Minus;
    case ConcatUpdate: return Concat;
    case LeftShiftUpdate: return LeftShift;
    case UnsignedLeftShiftUpdate: return UnsignedLeftShift;
    case RightShiftUpdate: return RightShift;
    case UnsignedRightShiftUpdate: return UnsignedRightShift;
    case AndUpdate: return And;
    case OrUpdate: return Or;
    case XorUpdate: return Xor;
    default: return Unknown;
    }
}

static void ast_visit_list(ast_list_t *ast_list, visit_behavior_t (*visitor)(ast_t *, void *), void *userdata) {
    for (ast_list_t *ast = ast_list; ast; ast = ast->next)
        ast_visit(ast->ast, visitor, userdata);
}

static void ast_visit_args(arg_ast_t *args, visit_behavior_t (*visitor)(ast_t *, void *), void *userdata) {
    for (arg_ast_t *arg = args; arg; arg = arg->next)
        ast_visit(arg->value, visitor, userdata);
}

void ast_visit(ast_t *ast, visit_behavior_t (*visitor)(ast_t *, void *), void *userdata) {
    if (!ast) return;
    if (visitor(ast, userdata) == VISIT_STOP) return;

    switch (ast->tag) {
    case Unknown:
    case None:
    case Bool:
    case Var:
    case Int:
    case Num:
    case Path:
    case TextLiteral:
    case Metadata: return;
    case Embed: ast_visit(Match(ast, Embed)->path, visitor, userdata); return;
    case TextJoin: ast_visit_list(Match(ast, TextJoin)->children, visitor, userdata); return;
    case Declare: {
        DeclareMatch(decl, ast, Declare);
        ast_visit(decl->var, visitor, userdata);
        ast_visit(decl->value, visitor, userdata);
        return;
    }
    case Assign: {
        DeclareMatch(assign, ast, Assign);
        ast_visit_list(assign->targets, visitor, userdata);
        ast_visit_list(assign->values, visitor, userdata);
        return;
    }
    case BINOP_CASES: {
        binary_operands_t op = BINARY_OPERANDS(ast);
        ast_visit(op.lhs, visitor, userdata);
        ast_visit(op.rhs, visitor, userdata);
        return;
    }
    case Negative: {
        ast_visit(Match(ast, Negative)->value, visitor, userdata);
        return;
    }
    case Not: {
        ast_visit(Match(ast, Not)->value, visitor, userdata);
        return;
    }
    case HeapAllocate: {
        ast_visit(Match(ast, HeapAllocate)->value, visitor, userdata);
        return;
    }
    case StackReference: {
        ast_visit(Match(ast, StackReference)->value, visitor, userdata);
        return;
    }
    case Min: {
        DeclareMatch(min, ast, Min);
        ast_visit(min->lhs, visitor, userdata);
        ast_visit(min->key, visitor, userdata);
        ast_visit(min->rhs, visitor, userdata);
        return;
    }
    case Max: {
        DeclareMatch(max, ast, Max);
        ast_visit(max->lhs, visitor, userdata);
        ast_visit(max->key, visitor, userdata);
        ast_visit(max->rhs, visitor, userdata);
        return;
    }
    case List: {
        ast_visit_list(Match(ast, List)->items, visitor, userdata);
        return;
    }
    case Table: {
        DeclareMatch(table, ast, Table);
        ast_visit_list(table->entries, visitor, userdata);
        ast_visit(table->default_value, visitor, userdata);
        ast_visit(table->fallback, visitor, userdata);
        return;
    }
    case TableEntry: {
        DeclareMatch(entry, ast, TableEntry);
        ast_visit(entry->key, visitor, userdata);
        ast_visit(entry->value, visitor, userdata);
        return;
    }
    case Comprehension: {
        DeclareMatch(comp, ast, Comprehension);
        ast_visit(comp->expr, visitor, userdata);
        ast_visit_list(comp->vars, visitor, userdata);
        ast_visit(comp->at, visitor, userdata);
        ast_visit_list(comp->iters, visitor, userdata);
        ast_visit(comp->filter, visitor, userdata);
        return;
    }
    case FunctionDef: {
        DeclareMatch(def, ast, FunctionDef);
        ast_visit(def->name, visitor, userdata);
        ast_visit_args(def->args, visitor, userdata);
        ast_visit(def->body, visitor, userdata);
        return;
    }
    case ConvertDef: {
        DeclareMatch(def, ast, ConvertDef);
        ast_visit_args(def->args, visitor, userdata);
        ast_visit(def->body, visitor, userdata);
        return;
    }
    case Lambda: {
        DeclareMatch(lambda, ast, Lambda);
        ast_visit_args(lambda->args, visitor, userdata);
        ast_visit(lambda->body, visitor, userdata);
        return;
    }
    case FunctionCall: {
        DeclareMatch(call, ast, FunctionCall);
        ast_visit(call->fn, visitor, userdata);
        ast_visit_args(call->args, visitor, userdata);
        return;
    }
    case MethodCall: {
        DeclareMatch(call, ast, MethodCall);
        ast_visit(call->self, visitor, userdata);
        ast_visit_args(call->args, visitor, userdata);
        return;
    }
    case RecordLiteral: {
        DeclareMatch(record, ast, RecordLiteral);
        ast_visit(record->type, visitor, userdata);
        ast_visit_args(record->args, visitor, userdata);
        return;
    }
    case Block: {
        ast_visit_list(Match(ast, Block)->statements, visitor, userdata);
        return;
    }
    case For: {
        DeclareMatch(for_, ast, For);
        ast_visit_list(for_->vars, visitor, userdata);
        ast_visit(for_->at, visitor, userdata);
        ast_visit_list(for_->iters, visitor, userdata);
        ast_visit(for_->body, visitor, userdata);
        ast_visit(for_->empty, visitor, userdata);
        return;
    }
    case While: {
        DeclareMatch(while_, ast, While);
        ast_visit(while_->condition, visitor, userdata);
        ast_visit(while_->body, visitor, userdata);
        return;
    }
    case Repeat: {
        ast_visit(Match(ast, Repeat)->body, visitor, userdata);
        return;
    }
    case If: {
        DeclareMatch(if_, ast, If);
        ast_visit(if_->condition, visitor, userdata);
        ast_visit(if_->body, visitor, userdata);
        ast_visit(if_->else_body, visitor, userdata);
        return;
    }
    case Match: {
        DeclareMatch(match, ast, Match);
        ast_visit(match->subject, visitor, userdata);
        for (match_clause_t *clause = match->clauses; clause; clause = clause->next) {
            ast_visit(clause->pattern, visitor, userdata);
            ast_visit(clause->body, visitor, userdata);
        }
        ast_visit(match->else_body, visitor, userdata);
        return;
    }
    case Reduction: {
        DeclareMatch(reduction, ast, Reduction);
        ast_visit(reduction->key, visitor, userdata);
        ast_visit(reduction->iter, visitor, userdata);
        return;
    }
    case Continue:
    case Break:
    case Pass: return;
    case Defer: {
        ast_visit(Match(ast, Defer)->body, visitor, userdata);
        return;
    }
    case Return: {
        ast_visit(Match(ast, Return)->value, visitor, userdata);
        return;
    }
    case StructDef: {
        DeclareMatch(def, ast, StructDef);
        ast_visit_args(def->fields, visitor, userdata);
        ast_visit(def->namespace, visitor, userdata);
        return;
    }
    case EnumDef: {
        DeclareMatch(def, ast, EnumDef);
        for (tag_ast_t *tag = def->tags; tag; tag = tag->next)
            ast_visit_args(tag->fields, visitor, userdata);
        ast_visit(def->namespace, visitor, userdata);
        return;
    }
    case LangDef: {
        ast_visit(Match(ast, LangDef)->namespace, visitor, userdata);
        return;
    }
    case Index: {
        DeclareMatch(index, ast, Index);
        ast_visit(index->indexed, visitor, userdata);
        ast_visit(index->index, visitor, userdata);
        return;
    }
    case FieldAccess: {
        ast_visit(Match(ast, FieldAccess)->fielded, visitor, userdata);
        return;
    }
    case NonOptional: {
        ast_visit(Match(ast, NonOptional)->value, visitor, userdata);
        return;
    }
    case DebugLog: {
        DeclareMatch(show, ast, DebugLog);
        ast_visit_list(show->values, visitor, userdata);
        return;
    }
    case Assert: {
        DeclareMatch(assert, ast, Assert);
        ast_visit(assert->expr, visitor, userdata);
        ast_visit(assert->message, visitor, userdata);
        return;
    }
    case Test: {
        ast_visit(Match(ast, Test)->body, visitor, userdata);
        return;
    }
    case Use: {
        ast_visit(Match(ast, Use)->var, visitor, userdata);
        return;
    }
    case InlineCCode: {
        ast_visit_list(Match(ast, InlineCCode)->chunks, visitor, userdata);
        return;
    }
    case Serialize: {
        ast_visit(Match(ast, Serialize)->value, visitor, userdata);
        return;
    }
    case Deserialize: {
        ast_visit(Match(ast, Deserialize)->value, visitor, userdata);
        return;
    }
    default: errx(1, "Visiting is not supported for this AST: %s", Text$as_c_string(ast_to_sexp(ast)));
#undef T
    }
}

static void _recursive_type_ast_visit(type_ast_t *ast, void *userdata) {
    if (ast == NULL) return;

    visit_behavior_t (*visit)(type_ast_t *, void *) = ((Closure_t *)userdata)->fn;
    void *visitor_userdata = ((Closure_t *)userdata)->userdata;
    if (visit(ast, visitor_userdata) == VISIT_STOP) return;

    switch (ast->tag) {
    case UnknownTypeAST:
    case VarTypeAST: break;
    case PointerTypeAST: {
        _recursive_type_ast_visit(Match(ast, PointerTypeAST)->pointed, userdata);
        break;
    }
    case ListTypeAST: {
        _recursive_type_ast_visit(Match(ast, ListTypeAST)->item, userdata);
        break;
    }
    case TableTypeAST: {
        DeclareMatch(table, ast, TableTypeAST);
        _recursive_type_ast_visit(table->key, userdata);
        _recursive_type_ast_visit(table->value, userdata);
        break;
    }
    case FunctionTypeAST: {
        DeclareMatch(fn, ast, FunctionTypeAST);
        for (arg_ast_t *arg = fn->args; arg; arg = arg->next)
            _recursive_type_ast_visit(arg->type, userdata);
        _recursive_type_ast_visit(fn->ret, userdata);
        break;
    }
    case OptionalTypeAST: {
        _recursive_type_ast_visit(Match(ast, OptionalTypeAST)->type, userdata);
        break;
    }
    case EnumTypeAST: {
        for (tag_ast_t *tag = Match(ast, EnumTypeAST)->tags; tag; tag = tag->next) {
            for (arg_ast_t *field = tag->fields; field; field = field->next) {
                _recursive_type_ast_visit(field->type, userdata);
            }
        }
        break;
    }
    default: errx(1, "Invalid type AST");
    }
}

static visit_behavior_t _type_ast_visit(ast_t *ast, void *userdata) {
    switch (ast->tag) {
    case Declare: {
        _recursive_type_ast_visit(Match(ast, Declare)->type, userdata);
        break;
    }
    case FunctionDef: {
        for (arg_ast_t *arg = Match(ast, FunctionDef)->args; arg; arg = arg->next)
            _recursive_type_ast_visit(arg->type, userdata);
        _recursive_type_ast_visit(Match(ast, FunctionDef)->ret_type, userdata);
        break;
    }
    case Lambda: {
        for (arg_ast_t *arg = Match(ast, Lambda)->args; arg; arg = arg->next)
            _recursive_type_ast_visit(arg->type, userdata);
        _recursive_type_ast_visit(Match(ast, Lambda)->ret_type, userdata);
        break;
    }
    case ConvertDef: {
        for (arg_ast_t *arg = Match(ast, ConvertDef)->args; arg; arg = arg->next)
            _recursive_type_ast_visit(arg->type, userdata);
        _recursive_type_ast_visit(Match(ast, ConvertDef)->ret_type, userdata);
        break;
    }
    case StructDef: {
        for (arg_ast_t *field = Match(ast, StructDef)->fields; field; field = field->next)
            _recursive_type_ast_visit(field->type, userdata);
        break;
    }
    case EnumDef: {
        for (tag_ast_t *tag = Match(ast, EnumDef)->tags; tag; tag = tag->next) {
            for (arg_ast_t *field = tag->fields; field; field = field->next) {
                _recursive_type_ast_visit(field->type, userdata);
            }
        }
        break;
    }
    case InlineCCode: {
        _recursive_type_ast_visit(Match(ast, InlineCCode)->type_ast, userdata);
        break;
    }
    case Deserialize: {
        _recursive_type_ast_visit(Match(ast, Deserialize)->type_ast, userdata);
        break;
    }
    default: break;
    }
    return VISIT_PROCEED;
}

void type_ast_visit(ast_t *ast, visit_behavior_t (*visitor)(type_ast_t *, void *), void *userdata) {
    Closure_t fn = {.fn = visitor, .userdata = userdata};
    ast_visit(ast, _type_ast_visit, &fn);
}

OptionalText_t ast_metadata(ast_t *ast, const char *key) {
    if (ast->tag != Block) return NONE_TEXT;
    Text_t key_text = Text$from_str(key);
    for (ast_list_t *stmt = Match(ast, Block)->statements; stmt; stmt = stmt->next) {
        if (stmt->ast->tag == Metadata) {
            DeclareMatch(m, stmt->ast, Metadata);
            if (Text$equal_values(m->key, key_text)) return m->value;
        }
    }
    return NONE_TEXT;
}