compile.c

#include <unistd.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <assert.h>

#include "misc.h"
#include "mpconfig.h"
#include "mpqstr.h"
#include "lexer.h"
#include "parse.h"
#include "scope.h"
#include "runtime0.h"
#include "emit.h"
#include "obj.h"
#include "compile.h"
#include "runtime.h"

// TODO need to mangle __attr names

#define MICROPY_EMIT_NATIVE (MICROPY_EMIT_X64 || MICROPY_EMIT_THUMB)

typedef enum {
    PN_none = 0,
#define DEF_RULE(rule, comp, kind, arg...) PN_##rule,
#include "grammar.h"
#undef DEF_RULE
    PN_maximum_number_of,
} pn_kind_t;

#define EMIT(fun, arg...) (comp->emit_method_table->fun(comp->emit, ##arg))
#define EMIT_INLINE_ASM(fun, arg...) (comp->emit_inline_asm_method_table->fun(comp->emit_inline_asm, ##arg))

#define EMIT_OPT_NONE           (0)
#define EMIT_OPT_BYTE_CODE      (1)
#define EMIT_OPT_NATIVE_PYTHON  (2)
#define EMIT_OPT_VIPER          (3)
#define EMIT_OPT_ASM_THUMB      (4)

typedef struct _compiler_t {
    bool is_repl;
    pass_kind_t pass;
    bool had_error; // try to keep compiler clean from nlr

    int next_label;

    int break_label;
    int continue_label;
    int except_nest_level;

    int n_arg_keyword;
    bool have_star_arg;
    bool have_dbl_star_arg;
    bool have_bare_star;
    int param_pass;
    int param_pass_num_dict_params;
    int param_pass_num_default_params;

    scope_t *scope_head;
    scope_t *scope_cur;

    emit_t *emit;                                   // current emitter
    const emit_method_table_t *emit_method_table;   // current emit method table

    emit_inline_asm_t *emit_inline_asm;                                   // current emitter for inline asm
    const emit_inline_asm_method_table_t *emit_inline_asm_method_table;   // current emit method table for inline asm
} compiler_t;

mp_parse_node_t fold_constants(mp_parse_node_t pn) {
    if (MP_PARSE_NODE_IS_STRUCT(pn)) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        int n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns);

        // fold arguments first
        for (int i = 0; i < n; i++) {
            pns->nodes[i] = fold_constants(pns->nodes[i]);
        }

        switch (MP_PARSE_NODE_STRUCT_KIND(pns)) {
            case PN_shift_expr:
                if (n == 3 && MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[0]) && MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[2])) {
                    int arg0 = MP_PARSE_NODE_LEAF_ARG(pns->nodes[0]);
                    int arg1 = MP_PARSE_NODE_LEAF_ARG(pns->nodes[2]);
                    if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[1], MP_TOKEN_OP_DBL_LESS)) {
#if MICROPY_EMIT_CPYTHON
                        // can overflow; enabled only to compare with CPython
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, arg0 << arg1);
#endif
                    } else if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[1], MP_TOKEN_OP_DBL_MORE)) {
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, arg0 >> arg1);
                    } else {
                        // shouldn't happen
                        assert(0);
                    }
                }
                break;

            case PN_arith_expr:
                // overflow checking here relies on SMALL_INT being strictly smaller than machine_int_t
                if (n == 3 && MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[0]) && MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[2])) {
                    machine_int_t arg0 = MP_PARSE_NODE_LEAF_ARG(pns->nodes[0]);
                    machine_int_t arg1 = MP_PARSE_NODE_LEAF_ARG(pns->nodes[2]);
                    machine_int_t res;
                    if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[1], MP_TOKEN_OP_PLUS)) {
                        res = arg0 + arg1;
                    } else if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[1], MP_TOKEN_OP_MINUS)) {
                        res = arg0 - arg1;
                    } else {
                        // shouldn't happen
                        assert(0);
                        res = 0;
                    }
                    if (MP_FIT_SMALL_INT(res)) {
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, res);
                    }
                }
                break;

            case PN_term:
                if (n == 3 && MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[0]) && MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[2])) {
                    int arg0 = MP_PARSE_NODE_LEAF_ARG(pns->nodes[0]);
                    int arg1 = MP_PARSE_NODE_LEAF_ARG(pns->nodes[2]);
                    if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[1], MP_TOKEN_OP_STAR)) {
#if MICROPY_EMIT_CPYTHON
                        // can overflow; enabled only to compare with CPython
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, arg0 * arg1);
#endif
                    } else if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[1], MP_TOKEN_OP_SLASH)) {
                        ; // pass
                    } else if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[1], MP_TOKEN_OP_PERCENT)) {
                        // XXX implement this properly as Python's % operator acts differently to C's
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, arg0 % arg1);
                    } else if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[1], MP_TOKEN_OP_DBL_SLASH)) {
                        // XXX implement this properly as Python's // operator acts differently to C's
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, arg0 / arg1);
                    } else {
                        // shouldn't happen
                        assert(0);
                    }
                }
                break;

            case PN_factor_2:
                if (MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[1])) {
                    machine_int_t arg = MP_PARSE_NODE_LEAF_ARG(pns->nodes[1]);
                    if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[0], MP_TOKEN_OP_PLUS)) {
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, arg);
                    } else if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[0], MP_TOKEN_OP_MINUS)) {
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, -arg);
                    } else if (MP_PARSE_NODE_IS_TOKEN_KIND(pns->nodes[0], MP_TOKEN_OP_TILDE)) {
                        pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, ~arg);
                    } else {
                        // shouldn't happen
                        assert(0);
                    }
                }
                break;

#if MICROPY_EMIT_CPYTHON
            case PN_power:
                // can overflow; enabled only to compare with CPython
                if (MP_PARSE_NODE_IS_SMALL_INT(pns->nodes[0]) && MP_PARSE_NODE_IS_NULL(pns->nodes[1]) && !MP_PARSE_NODE_IS_NULL(pns->nodes[2])) {
                    mp_parse_node_struct_t* pns2 = (mp_parse_node_struct_t*)pns->nodes[2];
                    if (MP_PARSE_NODE_IS_SMALL_INT(pns2->nodes[0])) {
                        int power = MP_PARSE_NODE_LEAF_ARG(pns2->nodes[0]);
                        if (power >= 0) {
                            int ans = 1;
                            int base = MP_PARSE_NODE_LEAF_ARG(pns->nodes[0]);
                            for (; power > 0; power--) {
                                ans *= base;
                            }
                            pn = mp_parse_node_new_leaf(MP_PARSE_NODE_SMALL_INT, ans);
                        }
                    }
                }
                break;
#endif
        }
    }

    return pn;
}

void compile_node(compiler_t *comp, mp_parse_node_t pn);

static int comp_next_label(compiler_t *comp) {
    return comp->next_label++;
}

static scope_t *scope_new_and_link(compiler_t *comp, scope_kind_t kind, mp_parse_node_t pn, uint emit_options) {
    scope_t *scope = scope_new(kind, pn, rt_get_unique_code_id(), emit_options);
    scope->parent = comp->scope_cur;
    scope->next = NULL;
    if (comp->scope_head == NULL) {
        comp->scope_head = scope;
    } else {
        scope_t *s = comp->scope_head;
        while (s->next != NULL) {
            s = s->next;
        }
        s->next = scope;
    }
    return scope;
}

static int list_len(mp_parse_node_t pn, int pn_kind) {
    if (MP_PARSE_NODE_IS_NULL(pn)) {
        return 0;
    } else if (MP_PARSE_NODE_IS_LEAF(pn)) {
        return 1;
    } else {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        if (MP_PARSE_NODE_STRUCT_KIND(pns) != pn_kind) {
            return 1;
        } else {
            return MP_PARSE_NODE_STRUCT_NUM_NODES(pns);
        }
    }
}

static void apply_to_single_or_list(compiler_t *comp, mp_parse_node_t pn, int pn_list_kind, void (*f)(compiler_t*, mp_parse_node_t)) {
    if (MP_PARSE_NODE_IS_STRUCT(pn) && MP_PARSE_NODE_STRUCT_KIND((mp_parse_node_struct_t*)pn) == pn_list_kind) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        int num_nodes = MP_PARSE_NODE_STRUCT_NUM_NODES(pns);
        for (int i = 0; i < num_nodes; i++) {
            f(comp, pns->nodes[i]);
        }
    } else if (!MP_PARSE_NODE_IS_NULL(pn)) {
        f(comp, pn);
    }
}

static int list_get(mp_parse_node_t *pn, int pn_kind, mp_parse_node_t **nodes) {
    if (MP_PARSE_NODE_IS_NULL(*pn)) {
        *nodes = NULL;
        return 0;
    } else if (MP_PARSE_NODE_IS_LEAF(*pn)) {
        *nodes = pn;
        return 1;
    } else {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)(*pn);
        if (MP_PARSE_NODE_STRUCT_KIND(pns) != pn_kind) {
            *nodes = pn;
            return 1;
        } else {
            *nodes = pns->nodes;
            return MP_PARSE_NODE_STRUCT_NUM_NODES(pns);
        }
    }
}

void compile_do_nothing(compiler_t *comp, mp_parse_node_struct_t *pns) {
}

void compile_generic_all_nodes(compiler_t *comp, mp_parse_node_struct_t *pns) {
    int num_nodes = MP_PARSE_NODE_STRUCT_NUM_NODES(pns);
    for (int i = 0; i < num_nodes; i++) {
        compile_node(comp, pns->nodes[i]);
    }
}

#if MICROPY_EMIT_CPYTHON
static bool cpython_c_tuple_is_const(mp_parse_node_t pn) {
    if (!MP_PARSE_NODE_IS_LEAF(pn)) {
        return false;
    }
    if (MP_PARSE_NODE_IS_ID(pn)) {
        return false;
    }
    return true;
}

static void cpython_c_print_quoted_str(vstr_t *vstr, qstr qstr, bool bytes) {
    const char *str = qstr_str(qstr);
    int len = strlen(str);
    bool has_single_quote = false;
    bool has_double_quote = false;
    for (int i = 0; i < len; i++) {
        if (str[i] == '\'') {
            has_single_quote = true;
        } else if (str[i] == '"') {
            has_double_quote = true;
        }
    }
    if (bytes) {
        vstr_printf(vstr, "b");
    }
    bool quote_single = false;
    if (has_single_quote && !has_double_quote) {
        vstr_printf(vstr, "\"");
    } else {
        quote_single = true;
        vstr_printf(vstr, "'");
    }
    for (int i = 0; i < len; i++) {
        if (str[i] == '\n') {
            vstr_printf(vstr, "\\n");
        } else if (str[i] == '\\') {
            vstr_printf(vstr, "\\\\");
        } else if (str[i] == '\'' && quote_single) {
            vstr_printf(vstr, "\\'");
        } else {
            vstr_printf(vstr, "%c", str[i]);
        }
    }
    if (has_single_quote && !has_double_quote) {
        vstr_printf(vstr, "\"");
    } else {
        vstr_printf(vstr, "'");
    }
}

static void cpython_c_tuple_emit_const(compiler_t *comp, mp_parse_node_t pn, vstr_t *vstr) {
    assert(MP_PARSE_NODE_IS_LEAF(pn));
    int arg = MP_PARSE_NODE_LEAF_ARG(pn);
    switch (MP_PARSE_NODE_LEAF_KIND(pn)) {
        case MP_PARSE_NODE_ID: assert(0);
        case MP_PARSE_NODE_SMALL_INT: vstr_printf(vstr, "%d", arg); break;
        case MP_PARSE_NODE_INTEGER: vstr_printf(vstr, "%s", qstr_str(arg)); break;
        case MP_PARSE_NODE_DECIMAL: vstr_printf(vstr, "%s", qstr_str(arg)); break;
        case MP_PARSE_NODE_STRING: cpython_c_print_quoted_str(vstr, arg, false); break;
        case MP_PARSE_NODE_BYTES: cpython_c_print_quoted_str(vstr, arg, true); break;
        case MP_PARSE_NODE_TOKEN:
            switch (arg) {
                case MP_TOKEN_KW_FALSE: vstr_printf(vstr, "False"); break;
                case MP_TOKEN_KW_NONE: vstr_printf(vstr, "None"); break;
                case MP_TOKEN_KW_TRUE: vstr_printf(vstr, "True"); break;
                default: assert(0);
            }
            break;
        default: assert(0);
    }
}

static void cpython_c_tuple(compiler_t *comp, mp_parse_node_t pn, mp_parse_node_struct_t *pns_list) {
    int n = 0;
    if (pns_list != NULL) {
        n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns_list);
    }
    int total = n;
    bool is_const = true;
    if (!MP_PARSE_NODE_IS_NULL(pn)) {
        total += 1;
        if (!cpython_c_tuple_is_const(pn)) {
            is_const = false;
        }
    }
    for (int i = 0; i < n; i++) {
        if (!cpython_c_tuple_is_const(pns_list->nodes[i])) {
            is_const = false;
            break;
        }
    }
    if (total > 0 && is_const) {
        bool need_comma = false;
        vstr_t *vstr = vstr_new();
        vstr_printf(vstr, "(");
        if (!MP_PARSE_NODE_IS_NULL(pn)) {
            cpython_c_tuple_emit_const(comp, pn, vstr);
            need_comma = true;
        }
        for (int i = 0; i < n; i++) {
            if (need_comma) {
                vstr_printf(vstr, ", ");
            }
            cpython_c_tuple_emit_const(comp, pns_list->nodes[i], vstr);
            need_comma = true;
        }
        if (total == 1) {
            vstr_printf(vstr, ",)");
        } else {
            vstr_printf(vstr, ")");
        }
        EMIT(load_const_verbatim_str, vstr_str(vstr));
        vstr_free(vstr);
    } else {
        if (!MP_PARSE_NODE_IS_NULL(pn)) {
            compile_node(comp, pn);
        }
        for (int i = 0; i < n; i++) {
            compile_node(comp, pns_list->nodes[i]);
        }
        EMIT(build_tuple, total);
    }
}
#endif

// funnelling all tuple creations through this function is purely so we can optionally agree with CPython
void c_tuple(compiler_t *comp, mp_parse_node_t pn, mp_parse_node_struct_t *pns_list) {
#if MICROPY_EMIT_CPYTHON
    cpython_c_tuple(comp, pn, pns_list);
#else
    int total = 0;
    if (!MP_PARSE_NODE_IS_NULL(pn)) {
        compile_node(comp, pn);
        total += 1;
    }
    if (pns_list != NULL) {
        int n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns_list);
        for (int i = 0; i < n; i++) {
            compile_node(comp, pns_list->nodes[i]);
        }
        total += n;
    }
    EMIT(build_tuple, total);
#endif
}

void compile_generic_tuple(compiler_t *comp, mp_parse_node_struct_t *pns) {
    // a simple tuple expression
    c_tuple(comp, MP_PARSE_NODE_NULL, pns);
}

static bool node_is_const_false(mp_parse_node_t pn) {
    return MP_PARSE_NODE_IS_TOKEN_KIND(pn, MP_TOKEN_KW_FALSE);
    // untested: || (MP_PARSE_NODE_IS_SMALL_INT(pn) && MP_PARSE_NODE_LEAF_ARG(pn) == 1);
}

static bool node_is_const_true(mp_parse_node_t pn) {
    return MP_PARSE_NODE_IS_TOKEN_KIND(pn, MP_TOKEN_KW_TRUE) || (MP_PARSE_NODE_IS_SMALL_INT(pn) && MP_PARSE_NODE_LEAF_ARG(pn) == 1);
}

#if MICROPY_EMIT_CPYTHON
// the is_nested variable is purely to match with CPython, which doesn't fully optimise not's
static void cpython_c_if_cond(compiler_t *comp, mp_parse_node_t pn, bool jump_if, int label, bool is_nested) {
    if (node_is_const_false(pn)) {
        if (jump_if == false) {
            EMIT(jump, label);
        }
        return;
    } else if (node_is_const_true(pn)) {
        if (jump_if == true) {
            EMIT(jump, label);
        }
        return;
    } else if (MP_PARSE_NODE_IS_STRUCT(pn)) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        int n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns);
        if (MP_PARSE_NODE_STRUCT_KIND(pns) == PN_or_test) {
            if (jump_if == false) {
                int label2 = comp_next_label(comp);
                for (int i = 0; i < n - 1; i++) {
                    cpython_c_if_cond(comp, pns->nodes[i], true, label2, true);
                }
                cpython_c_if_cond(comp, pns->nodes[n - 1], false, label, true);
                EMIT(label_assign, label2);
            } else {
                for (int i = 0; i < n; i++) {
                    cpython_c_if_cond(comp, pns->nodes[i], true, label, true);
                }
            }
            return;
        } else if (MP_PARSE_NODE_STRUCT_KIND(pns) == PN_and_test) {
            if (jump_if == false) {
                for (int i = 0; i < n; i++) {
                    cpython_c_if_cond(comp, pns->nodes[i], false, label, true);
                }
            } else {
                int label2 = comp_next_label(comp);
                for (int i = 0; i < n - 1; i++) {
                    cpython_c_if_cond(comp, pns->nodes[i], false, label2, true);
                }
                cpython_c_if_cond(comp, pns->nodes[n - 1], true, label, true);
                EMIT(label_assign, label2);
            }
            return;
        } else if (!is_nested && MP_PARSE_NODE_STRUCT_KIND(pns) == PN_not_test_2) {
            cpython_c_if_cond(comp, pns->nodes[0], !jump_if, label, true);
            return;
        }
    }

    // nothing special, fall back to default compiling for node and jump
    compile_node(comp, pn);
    if (jump_if == false) {
        EMIT(pop_jump_if_false, label);
    } else {
        EMIT(pop_jump_if_true, label);
    }
}
#endif

static void c_if_cond(compiler_t *comp, mp_parse_node_t pn, bool jump_if, int label) {
#if MICROPY_EMIT_CPYTHON
    cpython_c_if_cond(comp, pn, jump_if, label, false);
#else
    if (node_is_const_false(pn)) {
        if (jump_if == false) {
            EMIT(jump, label);
        }
        return;
    } else if (node_is_const_true(pn)) {
        if (jump_if == true) {
            EMIT(jump, label);
        }
        return;
    } else if (MP_PARSE_NODE_IS_STRUCT(pn)) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        int n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns);
        if (MP_PARSE_NODE_STRUCT_KIND(pns) == PN_or_test) {
            if (jump_if == false) {
                int label2 = comp_next_label(comp);
                for (int i = 0; i < n - 1; i++) {
                    c_if_cond(comp, pns->nodes[i], true, label2);
                }
                c_if_cond(comp, pns->nodes[n - 1], false, label);
                EMIT(label_assign, label2);
            } else {
                for (int i = 0; i < n; i++) {
                    c_if_cond(comp, pns->nodes[i], true, label);
                }
            }
            return;
        } else if (MP_PARSE_NODE_STRUCT_KIND(pns) == PN_and_test) {
            if (jump_if == false) {
                for (int i = 0; i < n; i++) {
                    c_if_cond(comp, pns->nodes[i], false, label);
                }
            } else {
                int label2 = comp_next_label(comp);
                for (int i = 0; i < n - 1; i++) {
                    c_if_cond(comp, pns->nodes[i], false, label2);
                }
                c_if_cond(comp, pns->nodes[n - 1], true, label);
                EMIT(label_assign, label2);
            }
            return;
        } else if (MP_PARSE_NODE_STRUCT_KIND(pns) == PN_not_test_2) {
            c_if_cond(comp, pns->nodes[0], !jump_if, label);
            return;
        }
    }

    // nothing special, fall back to default compiling for node and jump
    compile_node(comp, pn);
    if (jump_if == false) {
        EMIT(pop_jump_if_false, label);
    } else {
        EMIT(pop_jump_if_true, label);
    }
#endif
}

typedef enum { ASSIGN_STORE, ASSIGN_AUG_LOAD, ASSIGN_AUG_STORE } assign_kind_t;
void c_assign(compiler_t *comp, mp_parse_node_t pn, assign_kind_t kind);

void c_assign_power(compiler_t *comp, mp_parse_node_struct_t *pns, assign_kind_t assign_kind) {
    if (assign_kind != ASSIGN_AUG_STORE) {
        compile_node(comp, pns->nodes[0]);
    }

    if (MP_PARSE_NODE_IS_STRUCT(pns->nodes[1])) {
        mp_parse_node_struct_t *pns1 = (mp_parse_node_struct_t*)pns->nodes[1];
        if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_power_trailers) {
            int n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns1);
            if (assign_kind != ASSIGN_AUG_STORE) {
                for (int i = 0; i < n - 1; i++) {
                    compile_node(comp, pns1->nodes[i]);
                }
            }
            assert(MP_PARSE_NODE_IS_STRUCT(pns1->nodes[n - 1]));
            pns1 = (mp_parse_node_struct_t*)pns1->nodes[n - 1];
        }
        if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_trailer_paren) {
            printf("SyntaxError: can't assign to function call\n");
            return;
        } else if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_trailer_bracket) {
            if (assign_kind == ASSIGN_AUG_STORE) {
                EMIT(rot_three);
                EMIT(store_subscr);
            } else {
                compile_node(comp, pns1->nodes[0]);
                if (assign_kind == ASSIGN_AUG_LOAD) {
                    EMIT(dup_top_two);
                    EMIT(binary_op, RT_BINARY_OP_SUBSCR);
                } else {
                    EMIT(store_subscr);
                }
            }
        } else if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_trailer_period) {
            assert(MP_PARSE_NODE_IS_ID(pns1->nodes[0]));
            if (assign_kind == ASSIGN_AUG_LOAD) {
                EMIT(dup_top);
                EMIT(load_attr, MP_PARSE_NODE_LEAF_ARG(pns1->nodes[0]));
            } else {
                if (assign_kind == ASSIGN_AUG_STORE) {
                    EMIT(rot_two);
                }
                EMIT(store_attr, MP_PARSE_NODE_LEAF_ARG(pns1->nodes[0]));
            }
        } else {
            // shouldn't happen
            assert(0);
        }
    } else {
        // shouldn't happen
        assert(0);
    }

    if (!MP_PARSE_NODE_IS_NULL(pns->nodes[2])) {
        // SyntaxError, cannot assign
        assert(0);
    }
}

void c_assign_tuple(compiler_t *comp, int n, mp_parse_node_t *nodes) {
    assert(n >= 0);
    int have_star_index = -1;
    for (int i = 0; i < n; i++) {
        if (MP_PARSE_NODE_IS_STRUCT_KIND(nodes[i], PN_star_expr)) {
            if (have_star_index < 0) {
                EMIT(unpack_ex, i, n - i - 1);
                have_star_index = i;
            } else {
                printf("SyntaxError: two starred expressions in assignment\n");
                return;
            }
        }
    }
    if (have_star_index < 0) {
        EMIT(unpack_sequence, n);
    }
    for (int i = 0; i < n; i++) {
        if (i == have_star_index) {
            c_assign(comp, ((mp_parse_node_struct_t*)nodes[i])->nodes[0], ASSIGN_STORE);
        } else {
            c_assign(comp, nodes[i], ASSIGN_STORE);
        }
    }
}

// assigns top of stack to pn
void c_assign(compiler_t *comp, mp_parse_node_t pn, assign_kind_t assign_kind) {
    tail_recursion:
    if (MP_PARSE_NODE_IS_NULL(pn)) {
        assert(0);
    } else if (MP_PARSE_NODE_IS_LEAF(pn)) {
        if (MP_PARSE_NODE_IS_ID(pn)) {
            int arg = MP_PARSE_NODE_LEAF_ARG(pn);
            switch (assign_kind) {
                case ASSIGN_STORE:
                case ASSIGN_AUG_STORE:
                    EMIT(store_id, arg);
                    break;
                case ASSIGN_AUG_LOAD:
                    EMIT(load_id, arg);
                    break;
            }
        } else {
            printf("SyntaxError: can't assign to literal\n");
            return;
        }
    } else {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        switch (MP_PARSE_NODE_STRUCT_KIND(pns)) {
            case PN_power:
                // lhs is an index or attribute
                c_assign_power(comp, pns, assign_kind);
                break;

            case PN_testlist_star_expr:
            case PN_exprlist:
                // lhs is a tuple
                if (assign_kind != ASSIGN_STORE) {
                    goto bad_aug;
                }
                c_assign_tuple(comp, MP_PARSE_NODE_STRUCT_NUM_NODES(pns), pns->nodes);
                break;

            case PN_atom_paren:
                // lhs is something in parenthesis
                if (MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
                    // empty tuple
                    printf("SyntaxError: can't assign to ()\n");
                    return;
                } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[0], PN_testlist_comp)) {
                    pns = (mp_parse_node_struct_t*)pns->nodes[0];
                    goto testlist_comp;
                } else {
                    // parenthesis around 1 item, is just that item
                    pn = pns->nodes[0];
                    goto tail_recursion;
                }
                break;

            case PN_atom_bracket:
                // lhs is something in brackets
                if (assign_kind != ASSIGN_STORE) {
                    goto bad_aug;
                }
                if (MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
                    // empty list, assignment allowed
                    c_assign_tuple(comp, 0, NULL);
                } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[0], PN_testlist_comp)) {
                    pns = (mp_parse_node_struct_t*)pns->nodes[0];
                    goto testlist_comp;
                } else {
                    // brackets around 1 item
                    c_assign_tuple(comp, 1, &pns->nodes[0]);
                }
                break;

            default:
                printf("unknown assign, %u\n", (uint)MP_PARSE_NODE_STRUCT_KIND(pns));
                assert(0);
        }
        return;

        testlist_comp:
        // lhs is a sequence
        if (MP_PARSE_NODE_IS_STRUCT(pns->nodes[1])) {
            mp_parse_node_struct_t *pns2 = (mp_parse_node_struct_t*)pns->nodes[1];
            if (MP_PARSE_NODE_STRUCT_KIND(pns2) == PN_testlist_comp_3b) {
                // sequence of one item, with trailing comma
                assert(MP_PARSE_NODE_IS_NULL(pns2->nodes[0]));
                c_assign_tuple(comp, 1, &pns->nodes[0]);
            } else if (MP_PARSE_NODE_STRUCT_KIND(pns2) == PN_testlist_comp_3c) {
                // sequence of many items
                // TODO call c_assign_tuple instead
                int n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns2);
                EMIT(unpack_sequence, 1 + n);
                c_assign(comp, pns->nodes[0], ASSIGN_STORE);
                for (int i = 0; i < n; i++) {
                    c_assign(comp, pns2->nodes[i], ASSIGN_STORE);
                }
            } else if (MP_PARSE_NODE_STRUCT_KIND(pns) == PN_comp_for) {
                // TODO not implemented
                assert(0);
            } else {
                // sequence with 2 items
                goto sequence_with_2_items;
            }
        } else {
            // sequence with 2 items
            sequence_with_2_items:
            c_assign_tuple(comp, 2, pns->nodes);
        }
        return;
    }
    return;

    bad_aug:
    printf("SyntaxError: illegal expression for augmented assignment\n");
}

// stuff for lambda and comprehensions and generators
void close_over_variables_etc(compiler_t *comp, scope_t *this_scope, int n_dict_params, int n_default_params) {
    // make closed over variables, if any
    // ensure they are closed over in the order defined in the outer scope (mainly to agree with CPython)
    int nfree = 0;
    if (comp->scope_cur->kind != SCOPE_MODULE) {
        for (int i = 0; i < comp->scope_cur->id_info_len; i++) {
            id_info_t *id = &comp->scope_cur->id_info[i];
            if (id->kind == ID_INFO_KIND_CELL || id->kind == ID_INFO_KIND_FREE) {
                for (int j = 0; j < this_scope->id_info_len; j++) {
                    id_info_t *id2 = &this_scope->id_info[j];
                    if (id2->kind == ID_INFO_KIND_FREE && id->qstr == id2->qstr) {
#if MICROPY_EMIT_CPYTHON
                        EMIT(load_closure, id->qstr, id->local_num);
#else
                        // in Micro Python we load closures using LOAD_FAST
                        EMIT(load_fast, id->qstr, id->local_num);
#endif
                        nfree += 1;
                    }
                }
            }
        }
    }
    if (nfree > 0) {
        EMIT(build_tuple, nfree);
    }

    // make the function/closure
    if (nfree == 0) {
        EMIT(make_function, this_scope, n_dict_params, n_default_params);
    } else {
        EMIT(make_closure, this_scope, n_dict_params, n_default_params);
    }
}

void compile_funcdef_param(compiler_t *comp, mp_parse_node_t pn) {
    if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_typedargslist_name)) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        if (!MP_PARSE_NODE_IS_NULL(pns->nodes[2])) {
            // this parameter has a default value
            // in CPython, None (and True, False?) as default parameters are loaded with LOAD_NAME; don't understandy why
            if (comp->have_bare_star) {
                comp->param_pass_num_dict_params += 1;
                if (comp->param_pass == 1) {
                    EMIT(load_const_id, MP_PARSE_NODE_LEAF_ARG(pns->nodes[0]));
                    compile_node(comp, pns->nodes[2]);
                }
            } else {
                comp->param_pass_num_default_params += 1;
                if (comp->param_pass == 2) {
                    compile_node(comp, pns->nodes[2]);
                }
            }
        }
    } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_typedargslist_star)) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        if (MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
            // bare star
            comp->have_bare_star = true;
        }
    }
}

// leaves function object on stack
// returns function name
qstr compile_funcdef_helper(compiler_t *comp, mp_parse_node_struct_t *pns, uint emit_options) {
    if (comp->pass == PASS_1) {
        // create a new scope for this function
        scope_t *s = scope_new_and_link(comp, SCOPE_FUNCTION, (mp_parse_node_t)pns, emit_options);
        // store the function scope so the compiling function can use it at each pass
        pns->nodes[4] = (mp_parse_node_t)s;
    }

    // save variables (probably don't need to do this, since we can't have nested definitions..?)
    bool old_have_bare_star = comp->have_bare_star;
    int old_param_pass = comp->param_pass;
    int old_param_pass_num_dict_params = comp->param_pass_num_dict_params;
    int old_param_pass_num_default_params = comp->param_pass_num_default_params;

    // compile default parameters
    comp->have_bare_star = false;
    comp->param_pass = 1; // pass 1 does any default parameters after bare star
    comp->param_pass_num_dict_params = 0;
    comp->param_pass_num_default_params = 0;
    apply_to_single_or_list(comp, pns->nodes[1], PN_typedargslist, compile_funcdef_param);
    comp->have_bare_star = false;
    comp->param_pass = 2; // pass 2 does any default parameters before bare star
    comp->param_pass_num_dict_params = 0;
    comp->param_pass_num_default_params = 0;
    apply_to_single_or_list(comp, pns->nodes[1], PN_typedargslist, compile_funcdef_param);

    // get the scope for this function
    scope_t *fscope = (scope_t*)pns->nodes[4];

    // make the function
    close_over_variables_etc(comp, fscope, comp->param_pass_num_dict_params, comp->param_pass_num_default_params);

    // restore variables
    comp->have_bare_star = old_have_bare_star;
    comp->param_pass = old_param_pass;
    comp->param_pass_num_dict_params = old_param_pass_num_dict_params;
    comp->param_pass_num_default_params = old_param_pass_num_default_params;

    // return its name (the 'f' in "def f(...):")
    return fscope->simple_name;
}

// leaves class object on stack
// returns class name
qstr compile_classdef_helper(compiler_t *comp, mp_parse_node_struct_t *pns, uint emit_options) {
    if (comp->pass == PASS_1) {
        // create a new scope for this class
        scope_t *s = scope_new_and_link(comp, SCOPE_CLASS, (mp_parse_node_t)pns, emit_options);
        // store the class scope so the compiling function can use it at each pass
        pns->nodes[3] = (mp_parse_node_t)s;
    }

    EMIT(load_build_class);

    // scope for this class
    scope_t *cscope = (scope_t*)pns->nodes[3];

    // compile the class
    close_over_variables_etc(comp, cscope, 0, 0);

    // get its name
    EMIT(load_const_id, cscope->simple_name);

    // nodes[1] has parent classes, if any
    if (MP_PARSE_NODE_IS_NULL(pns->nodes[1])) {
        // no parent classes
        EMIT(call_function, 2, 0, false, false);
    } else {
        // have a parent class or classes
        // TODO what if we have, eg, *a or **a in the parent list?
        compile_node(comp, pns->nodes[1]);
        EMIT(call_function, 2 + list_len(pns->nodes[1], PN_arglist), 0, false, false);
    }

    // return its name (the 'C' in class C(...):")
    return cscope->simple_name;
}

// returns true if it was a built-in decorator (even if the built-in had an error)
static bool compile_built_in_decorator(compiler_t *comp, int name_len, mp_parse_node_t *name_nodes, uint *emit_options) {
    if (MP_PARSE_NODE_LEAF_ARG(name_nodes[0]) != MP_QSTR_micropython) {
        return false;
    }

    if (name_len != 2) {
        printf("SyntaxError: invalid micropython decorator\n");
        return true;
    }

    qstr attr = MP_PARSE_NODE_LEAF_ARG(name_nodes[1]);
    if (attr == MP_QSTR_byte_code) {
        *emit_options = EMIT_OPT_BYTE_CODE;
#if MICROPY_EMIT_NATIVE
    } else if (attr == MP_QSTR_native) {
        *emit_options = EMIT_OPT_NATIVE_PYTHON;
    } else if (attr == MP_QSTR_viper) {
        *emit_options = EMIT_OPT_VIPER;
#endif
#if MICROPY_EMIT_INLINE_THUMB
    } else if (attr == MP_QSTR_asm_thumb) {
        *emit_options = EMIT_OPT_ASM_THUMB;
#endif
    } else {
        printf("SyntaxError: invalid micropython decorator '%s'\n", qstr_str(attr));
    }

    return true;
}

void compile_decorated(compiler_t *comp, mp_parse_node_struct_t *pns) {
    // get the list of decorators
    mp_parse_node_t *nodes;
    int n = list_get(&pns->nodes[0], PN_decorators, &nodes);

    // inherit emit options for this function/class definition
    uint emit_options = comp->scope_cur->emit_options;

    // compile each decorator
    int num_built_in_decorators = 0;
    for (int i = 0; i < n; i++) {
        assert(MP_PARSE_NODE_IS_STRUCT_KIND(nodes[i], PN_decorator)); // should be
        mp_parse_node_struct_t *pns_decorator = (mp_parse_node_struct_t*)nodes[i];

        // nodes[0] contains the decorator function, which is a dotted name
        mp_parse_node_t *name_nodes;
        int name_len = list_get(&pns_decorator->nodes[0], PN_dotted_name, &name_nodes);

        // check for built-in decorators
        if (compile_built_in_decorator(comp, name_len, name_nodes, &emit_options)) {
            // this was a built-in
            num_built_in_decorators += 1;

        } else {
            // not a built-in, compile normally

            // compile the decorator function
            compile_node(comp, name_nodes[0]);
            for (int i = 1; i < name_len; i++) {
                assert(MP_PARSE_NODE_IS_ID(name_nodes[i])); // should be
                EMIT(load_attr, MP_PARSE_NODE_LEAF_ARG(name_nodes[i]));
            }

            // nodes[1] contains arguments to the decorator function, if any
            if (!MP_PARSE_NODE_IS_NULL(pns_decorator->nodes[1])) {
                // call the decorator function with the arguments in nodes[1]
                compile_node(comp, pns_decorator->nodes[1]);
            }
        }
    }

    // compile the body (funcdef or classdef) and get its name
    mp_parse_node_struct_t *pns_body = (mp_parse_node_struct_t*)pns->nodes[1];
    qstr body_name = 0;
    if (MP_PARSE_NODE_STRUCT_KIND(pns_body) == PN_funcdef) {
        body_name = compile_funcdef_helper(comp, pns_body, emit_options);
    } else if (MP_PARSE_NODE_STRUCT_KIND(pns_body) == PN_classdef) {
        body_name = compile_classdef_helper(comp, pns_body, emit_options);
    } else {
        // shouldn't happen
        assert(0);
    }

    // call each decorator
    for (int i = 0; i < n - num_built_in_decorators; i++) {
        EMIT(call_function, 1, 0, false, false);
    }

    // store func/class object into name
    EMIT(store_id, body_name);
}

void compile_funcdef(compiler_t *comp, mp_parse_node_struct_t *pns) {
    qstr fname = compile_funcdef_helper(comp, pns, comp->scope_cur->emit_options);
    // store function object into function name
    EMIT(store_id, fname);
}

void c_del_stmt(compiler_t *comp, mp_parse_node_t pn) {
    if (MP_PARSE_NODE_IS_ID(pn)) {
        EMIT(delete_id, MP_PARSE_NODE_LEAF_ARG(pn));
    } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_power)) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;

        compile_node(comp, pns->nodes[0]); // base of the power node

        if (MP_PARSE_NODE_IS_STRUCT(pns->nodes[1])) {
            mp_parse_node_struct_t *pns1 = (mp_parse_node_struct_t*)pns->nodes[1];
            if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_power_trailers) {
                int n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns1);
                for (int i = 0; i < n - 1; i++) {
                    compile_node(comp, pns1->nodes[i]);
                }
                assert(MP_PARSE_NODE_IS_STRUCT(pns1->nodes[n - 1]));
                pns1 = (mp_parse_node_struct_t*)pns1->nodes[n - 1];
            }
            if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_trailer_paren) {
                // SyntaxError: can't delete a function call
                assert(0);
            } else if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_trailer_bracket) {
                compile_node(comp, pns1->nodes[0]);
                EMIT(delete_subscr);
            } else if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_trailer_period) {
                assert(MP_PARSE_NODE_IS_ID(pns1->nodes[0]));
                EMIT(delete_attr, MP_PARSE_NODE_LEAF_ARG(pns1->nodes[0]));
            } else {
                // shouldn't happen
                assert(0);
            }
        } else {
            // shouldn't happen
            assert(0);
        }

        if (!MP_PARSE_NODE_IS_NULL(pns->nodes[2])) {
            // SyntaxError, cannot delete
            assert(0);
        }
    } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_atom_paren)) {
        pn = ((mp_parse_node_struct_t*)pn)->nodes[0];
        if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_testlist_comp)) {
            mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
            // TODO perhaps factorise testlist_comp code with other uses of PN_testlist_comp

            if (MP_PARSE_NODE_IS_STRUCT(pns->nodes[1])) {
                mp_parse_node_struct_t *pns1 = (mp_parse_node_struct_t*)pns->nodes[1];
                if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_testlist_comp_3b) {
                    // sequence of one item, with trailing comma
                    assert(MP_PARSE_NODE_IS_NULL(pns1->nodes[0]));
                    c_del_stmt(comp, pns->nodes[0]);
                } else if (MP_PARSE_NODE_STRUCT_KIND(pns1) == PN_testlist_comp_3c) {
                    // sequence of many items
                    int n = MP_PARSE_NODE_STRUCT_NUM_NODES(pns1);
                    c_del_stmt(comp, pns->nodes[0]);
                    for (int i = 0; i < n; i++) {
                        c_del_stmt(comp, pns1->nodes[i]);
                    }
                } else if (MP_PARSE_NODE_STRUCT_KIND(pns) == PN_comp_for) {
                    // TODO not implemented; can't del comprehension?
                    assert(0);
                } else {
                    // sequence with 2 items
                    goto sequence_with_2_items;
                }
            } else {
                // sequence with 2 items
                sequence_with_2_items:
                c_del_stmt(comp, pns->nodes[0]);
                c_del_stmt(comp, pns->nodes[1]);
            }
        } else {
            // tuple with 1 element
            c_del_stmt(comp, pn);
        }
    } else {
        // not implemented
        assert(0);
    }
}

void compile_del_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
    apply_to_single_or_list(comp, pns->nodes[0], PN_exprlist, c_del_stmt);
}

void compile_break_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
    if (comp->break_label == 0) {
        printf("ERROR: cannot break from here\n");
    }
    EMIT(break_loop, comp->break_label);
}

void compile_continue_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
    if (comp->continue_label == 0) {
        printf("ERROR: cannot continue from here\n");
    }
    if (comp->except_nest_level > 0) {
        EMIT(continue_loop, comp->continue_label);
    } else {
        EMIT(jump, comp->continue_label);
    }
}

void compile_return_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
    if (comp->scope_cur->kind != SCOPE_FUNCTION) {
        printf("SyntaxError: 'return' outside function\n");
        comp->had_error = true;
        return;
    }
    if (MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
        // no argument to 'return', so return None
        EMIT(load_const_tok, MP_TOKEN_KW_NONE);
    } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[0], PN_test_if_expr)) {
        // special case when returning an if-expression; to match CPython optimisation
        mp_parse_node_struct_t *pns_test_if_expr = (mp_parse_node_struct_t*)pns->nodes[0];
        mp_parse_node_struct_t *pns_test_if_else = (mp_parse_node_struct_t*)pns_test_if_expr->nodes[1];

        int l_fail = comp_next_label(comp);
        c_if_cond(comp, pns_test_if_else->nodes[0], false, l_fail); // condition
        compile_node(comp, pns_test_if_expr->nodes[0]); // success value
        EMIT(return_value);
        EMIT(label_assign, l_fail);
        compile_node(comp, pns_test_if_else->nodes[1]); // failure value
    } else {
        compile_node(comp, pns->nodes[0]);
    }
    EMIT(return_value);
}

void compile_yield_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
    compile_node(comp, pns->nodes[0]);
    EMIT(pop_top);
}

void compile_raise_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
    if (MP_PARSE_NODE_IS_NULL(pns->nodes[0])) {
        // raise
        EMIT(raise_varargs, 0);
    } else if (MP_PARSE_NODE_IS_STRUCT_KIND(pns->nodes[0], PN_raise_stmt_arg)) {
        // raise x from y
        pns = (mp_parse_node_struct_t*)pns->nodes[0];
        compile_node(comp, pns->nodes[0]);
        compile_node(comp, pns->nodes[1]);
        EMIT(raise_varargs, 2);
    } else {
        // raise x
        compile_node(comp, pns->nodes[0]);
        EMIT(raise_varargs, 1);
    }
}

// q1 holds the base, q2 the full name
// eg   a -> q1=q2=a
//      a.b.c -> q1=a, q2=a.b.c
void do_import_name(compiler_t *comp, mp_parse_node_t pn, qstr *q1, qstr *q2) {
    bool is_as = false;
    if (MP_PARSE_NODE_IS_STRUCT_KIND(pn, PN_dotted_as_name)) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        // a name of the form x as y; unwrap it
        *q1 = MP_PARSE_NODE_LEAF_ARG(pns->nodes[1]);
        pn = pns->nodes[0];
        is_as = true;
    }
    if (MP_PARSE_NODE_IS_ID(pn)) {
        // just a simple name
        *q2 = MP_PARSE_NODE_LEAF_ARG(pn);
        if (!is_as) {
            *q1 = *q2;
        }
        EMIT(import_name, *q2);
    } else if (MP_PARSE_NODE_IS_STRUCT(pn)) {
        mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)pn;
        if (MP_PARSE_NODE_STRUCT_KIND(pns) == PN_dotted_name) {