pocketpy/src/modules/array2d.c

#include "pocketpy/interpreter/array2d.h"
#include "pocketpy/interpreter/vm.h"
#include "pocketpy/pocketpy.h"
#include <limits.h>

static bool c11_array2d_like_is_valid(c11_array2d_like* self, unsigned int col, unsigned int row) {
    return col < self->n_cols && row < self->n_rows;
}

static py_Ref c11_array2d__get(c11_array2d* self, int col, int row) {
    return self->data + row * self->header.n_cols + col;
}

static bool c11_array2d__set(c11_array2d* self, int col, int row, py_Ref value) {
    self->data[row * self->header.n_cols + col] = *value;
    return true;
}

c11_array2d* py_newarray2d(py_OutRef out, int n_cols, int n_rows) {
    int numel = n_cols * n_rows;
    c11_array2d* ud = py_newobject(out, tp_array2d, numel, sizeof(c11_array2d));
    ud->header.n_cols = n_cols;
    ud->header.n_rows = n_rows;
    ud->header.numel = numel;
    ud->header.f_get = (py_Ref(*)(c11_array2d_like*, int, int))c11_array2d__get;
    ud->header.f_set = (bool (*)(c11_array2d_like*, int, int, py_Ref))c11_array2d__set;
    ud->data = py_getslot(out, 0);
    return ud;
}

/* array2d_like bindings */
static bool array2d_like_n_cols(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    py_newint(py_retval(), self->n_cols);
    return true;
}

static bool array2d_like_n_rows(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    py_newint(py_retval(), self->n_rows);
    return true;
}

static bool array2d_like_shape(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    c11_vec2i shape;
    shape.x = self->n_cols;
    shape.y = self->n_rows;
    py_newvec2i(py_retval(), shape);
    return true;
}

static bool array2d_like_numel(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    py_newint(py_retval(), self->numel);
    return true;
}

static bool array2d_like_is_valid(int argc, py_Ref argv) {
    c11_array2d_like* self = py_touserdata(argv);
    int col, row;
    if(argc == 2) {
        PY_CHECK_ARG_TYPE(1, tp_vec2i);
        c11_vec2i pos = py_tovec2i(py_arg(1));
        col = pos.x;
        row = pos.y;
    } else if(argc == 3) {
        PY_CHECK_ARG_TYPE(1, tp_int);
        PY_CHECK_ARG_TYPE(2, tp_int);
        col = py_toint(py_arg(1));
        row = py_toint(py_arg(2));
    } else {
        return TypeError("is_valid() expected 2 or 3 arguments");
    }
    py_newbool(py_retval(), c11_array2d_like_is_valid(self, col, row));
    return true;
}

static bool array2d_like_get(int argc, py_Ref argv) {
    PY_CHECK_ARG_TYPE(1, tp_int);
    PY_CHECK_ARG_TYPE(2, tp_int);
    py_Ref default_;
    c11_array2d_like* self = py_touserdata(argv);
    if(argc == 3) {
        default_ = py_None();
    } else if(argc == 4) {
        default_ = py_arg(3);
    } else {
        return TypeError("get() expected 2 or 3 arguments");
    }
    int col = py_toint(py_arg(1));
    int row = py_toint(py_arg(2));
    if(c11_array2d_like_is_valid(self, col, row)) {
        py_assign(py_retval(), self->f_get(self, col, row));
    } else {
        py_assign(py_retval(), default_);
    }
    return true;
}

static bool array2d_like_render(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_sbuf buf;
    c11_sbuf__ctor(&buf);
    c11_array2d_like* self = py_touserdata(argv);
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_Ref item = self->f_get(self, i, j);
            if(!py_str(item)) return false;
            c11_sbuf__write_sv(&buf, py_tosv(py_retval()));
        }
        if(j < self->n_rows - 1) c11_sbuf__write_char(&buf, '\n');
    }
    c11_sbuf__py_submit(&buf, py_retval());
    return true;
}

static bool array2d_like_all(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_Ref item = self->f_get(self, i, j);
            if(!py_checkbool(item)) return false;
            if(!py_tobool(item)) {
                py_newbool(py_retval(), false);
                return true;
            }
        }
    }
    py_newbool(py_retval(), true);
    return true;
}

static bool array2d_like_any(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_Ref item = self->f_get(self, i, j);
            if(!py_checkbool(item)) return false;
            if(py_tobool(item)) {
                py_newbool(py_retval(), true);
                return true;
            }
        }
    }
    py_newbool(py_retval(), false);
    return true;
}

static bool array2d_like_map(int argc, py_Ref argv) {
    // def map(self, f: Callable[[T], Any]) -> 'array2d': ...
    PY_CHECK_ARGC(2);
    c11_array2d_like* self = py_touserdata(argv);
    py_Ref f = py_arg(1);
    c11_array2d* res = py_newarray2d(py_pushtmp(), self->n_cols, self->n_rows);
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_Ref item = self->f_get(self, i, j);
            if(!py_call(f, 1, item)) return false;
            res->data[j * self->n_cols + i] = *py_retval();
        }
    }
    py_assign(py_retval(), py_peek(-1));
    py_pop();
    return true;
}

static bool array2d_like_apply(int argc, py_Ref argv) {
    // def apply_(self, f: Callable[[T], T]) -> None: ...
    PY_CHECK_ARGC(2);
    c11_array2d_like* self = py_touserdata(argv);
    py_Ref f = py_arg(1);
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_Ref item = self->f_get(self, i, j);
            if(!py_call(f, 1, item)) return false;
            bool ok = self->f_set(self, i, j, py_retval());
            if(!ok) return false;
        }
    }
    py_newnone(py_retval());
    return true;
}

static bool array2d_like_copy(int argc, py_Ref argv) {
    // def copy(self) -> 'array2d': ...
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    c11_array2d* res = py_newarray2d(py_retval(), self->n_cols, self->n_rows);
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_Ref item = self->f_get(self, i, j);
            res->data[j * self->n_cols + i] = *item;
        }
    }
    return true;
}

static bool array2d_like_tolist(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    py_newlistn(py_retval(), self->n_rows);
    for(int j = 0; j < self->n_rows; j++) {
        py_Ref row_j = py_list_getitem(py_retval(), j);
        py_newlistn(row_j, self->n_cols);
        for(int i = 0; i < self->n_cols; i++) {
            py_Ref item = self->f_get(self, i, j);
            py_list_setitem(row_j, i, item);
        }
    }
    return true;
}

static bool _check_same_shape(int colA, int rowA, int colB, int rowB) {
    if(colA != colB || rowA != rowB) {
        const char* fmt = "expected the same shape: (%d, %d) != (%d, %d)";
        return ValueError(fmt, colA, rowA, colB, rowB);
    }
    return true;
}

static bool _array2d_like_check_same_shape(c11_array2d_like* self, c11_array2d_like* other) {
    return _check_same_shape(self->n_cols, self->n_rows, other->n_cols, other->n_rows);
}

static bool array2d_like__eq__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    c11_array2d_like* self = py_touserdata(argv);
    c11_array2d* res = py_newarray2d(py_pushtmp(), self->n_cols, self->n_rows);
    if(py_isinstance(py_arg(1), tp_array2d_like)) {
        c11_array2d_like* other = py_touserdata(py_arg(1));
        if(!_array2d_like_check_same_shape(self, other)) return false;
        for(int j = 0; j < self->n_rows; j++) {
            for(int i = 0; i < self->n_cols; i++) {
                py_Ref lhs = self->f_get(self, i, j);
                py_Ref rhs = other->f_get(other, i, j);
                int code = py_equal(lhs, rhs);
                if(code == -1) return false;
                py_newbool(&res->data[j * self->n_cols + i], (bool)code);
            }
        }
    } else {
        // broadcast
        for(int j = 0; j < self->n_rows; j++) {
            for(int i = 0; i < self->n_cols; i++) {
                py_Ref lhs = self->f_get(self, i, j);
                int code = py_equal(lhs, py_arg(1));
                if(code == -1) return false;
                py_newbool(&res->data[j * self->n_cols + i], (bool)code);
            }
        }
    }
    py_assign(py_retval(), py_peek(-1));
    py_pop();
    return true;
}

static bool array2d_like__ne__(int argc, py_Ref argv) {
    bool ok = array2d_like__eq__(argc, argv);
    if(!ok) return false;
    assert(py_istype(py_retval(), tp_array2d));
    c11_array2d* res = py_touserdata(py_retval());
    py_TValue* data = res->data;
    for(int i = 0; i < res->header.numel; i++) {
        assert(py_isbool(&data[i]));
        py_newbool(&data[i], !py_tobool(&data[i]));
    }
    return true;
}

static bool array2d_like__iter__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    c11_array2d_like_iterator* ud =
        py_newobject(py_retval(), tp_array2d_like_iterator, 1, sizeof(c11_array2d_like_iterator));
    py_setslot(py_retval(), 0, argv);  // keep the array alive
    ud->array = self;
    ud->j = 0;
    ud->i = 0;
    return true;
}

static bool array2d_like__repr__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like* self = py_touserdata(argv);
    char buf[256];
    snprintf(buf,
             sizeof(buf),
             "%s(%d, %d)",
             py_tpname(py_typeof(argv)),
             self->n_cols,
             self->n_rows);
    py_newstr(py_retval(), buf);
    return true;
}

#define HANDLE_SLICE()                                                                             \
    int start_col, stop_col, step_col;                                                             \
    int start_row, stop_row, step_row;                                                             \
    if(!pk__parse_int_slice(x, self->n_cols, &start_col, &stop_col, &step_col)) return false;      \
    if(!pk__parse_int_slice(y, self->n_rows, &start_row, &stop_row, &step_row)) return false;      \
    if(step_col != 1 || step_row != 1) return ValueError("slice step must be 1");                  \
    int slice_width = stop_col - start_col;                                                        \
    int slice_height = stop_row - start_row;

static bool _array2d_like_IndexError(c11_array2d_like* self, int col, int row) {
    return IndexError("(%d, %d) is not a valid index of array2d_like(%d, %d)",
                      col,
                      row,
                      self->n_cols,
                      self->n_rows);
}

static py_Ref c11_array2d_view__get(c11_array2d_view* self, int col, int row) {
    return self->f_get(self->ctx, col + self->origin.x, row + self->origin.y);
}

static bool c11_array2d_view__set(c11_array2d_view* self, int col, int row, py_Ref value) {
    return self->f_set(self->ctx, col + self->origin.x, row + self->origin.y, value);
}

static c11_array2d_view* _array2d_view__new(py_OutRef out,
                                            py_Ref keepalive,
                                            int start_col,
                                            int start_row,
                                            int width,
                                            int height) {
    c11_array2d_view* res = py_newobject(out, tp_array2d_view, 1, sizeof(c11_array2d_view));
    if(width <= 0 || height <= 0) {
        ValueError("width and height must be positive");
        return NULL;
    }
    res->header.n_cols = width;
    res->header.n_rows = height;
    res->header.numel = width * height;
    res->header.f_get = (py_Ref(*)(c11_array2d_like*, int, int))c11_array2d_view__get;
    res->header.f_set = (bool (*)(c11_array2d_like*, int, int, py_Ref))c11_array2d_view__set;
    res->origin.x = start_col;
    res->origin.y = start_row;
    py_setslot(out, 0, keepalive);
    return res;
}

static bool _array2d_view(py_OutRef out,
                          py_Ref keepalive,
                          c11_array2d_like* array,
                          int start_col,
                          int start_row,
                          int width,
                          int height) {
    c11_array2d_view* res = _array2d_view__new(out, keepalive, start_col, start_row, width, height);
    if(res == NULL) return false;
    res->ctx = array;
    res->f_get = (py_Ref(*)(void*, int, int))array->f_get;
    res->f_set = (bool (*)(void*, int, int, py_Ref))array->f_set;
    return true;
}

static bool _chunked_array2d_view(py_OutRef out,
                                  py_Ref keepalive,
                                  c11_chunked_array2d* array,
                                  int start_col,
                                  int start_row,
                                  int width,
                                  int height) {
    c11_array2d_view* res = _array2d_view__new(out, keepalive, start_col, start_row, width, height);
    if(res == NULL) return false;
    res->ctx = array;
    res->f_get = (py_Ref(*)(void*, int, int))c11_chunked_array2d__get;
    res->f_set = (bool (*)(void*, int, int, py_Ref))c11_chunked_array2d__set;
    return true;
}

static bool array2d_like__getitem__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    c11_array2d_like* self = py_touserdata(argv);
    if(argv[1].type == tp_vec2i) {
        c11_vec2i pos = py_tovec2i(&argv[1]);
        if(c11_array2d_like_is_valid(self, pos.x, pos.y)) {
            py_assign(py_retval(), self->f_get(self, pos.x, pos.y));
            return true;
        }
        return _array2d_like_IndexError(self, pos.x, pos.y);
    }

    if(py_isinstance(&argv[1], tp_array2d_like)) {
        c11_array2d_like* mask = py_touserdata(&argv[1]);
        if(!_array2d_like_check_same_shape(self, mask)) return false;
        py_newlist(py_retval());
        for(int j = 0; j < self->n_rows; j++) {
            for(int i = 0; i < self->n_cols; i++) {
                py_Ref item = self->f_get(self, i, j);
                py_Ref cond = mask->f_get(mask, i, j);
                if(!py_checkbool(cond)) return false;
                if(py_tobool(cond)) py_list_append(py_retval(), item);
            }
        }
        return true;
    }

    PY_CHECK_ARG_TYPE(1, tp_tuple);
    if(py_tuple_len(&argv[1]) != 2) return TypeError("expected a tuple of 2 elements");
    py_Ref x = py_tuple_getitem(&argv[1], 0);
    py_Ref y = py_tuple_getitem(&argv[1], 1);
    if(py_isint(x) && py_isint(y)) {
        int col = py_toint(x);
        int row = py_toint(y);
        if(c11_array2d_like_is_valid(self, col, row)) {
            py_assign(py_retval(), self->f_get(self, col, row));
            return true;
        }
        return _array2d_like_IndexError(self, col, row);
    }

    if(py_istype(x, tp_slice) && py_istype(y, tp_slice)) {
        HANDLE_SLICE();
        return _array2d_view(py_retval(),
                             argv,
                             self,
                             start_col,
                             start_row,
                             slice_width,
                             slice_height);
    }

    return TypeError("expected `tuple[int, int]` or `tuple[slice, slice]`");
}

static bool array2d_like__setitem__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(3);
    c11_array2d_like* self = py_touserdata(argv);
    py_Ref value = &argv[2];
    if(argv[1].type == tp_vec2i) {
        c11_vec2i pos = py_tovec2i(&argv[1]);
        if(c11_array2d_like_is_valid(self, pos.x, pos.y)) {
            bool ok = self->f_set(self, pos.x, pos.y, value);
            if(!ok) return false;
            py_newnone(py_retval());
            return true;
        }
        return _array2d_like_IndexError(self, pos.x, pos.y);
    }

    if(py_isinstance(&argv[1], tp_array2d_like)) {
        c11_array2d_like* mask = py_touserdata(&argv[1]);
        if(!_array2d_like_check_same_shape(self, mask)) return false;
        for(int j = 0; j < self->n_rows; j++) {
            for(int i = 0; i < self->n_cols; i++) {
                py_Ref cond = mask->f_get(mask, i, j);
                if(!py_checkbool(cond)) return false;
                if(py_tobool(cond)) {
                    bool ok = self->f_set(self, i, j, value);
                    if(!ok) return false;
                }
            }
        }
        py_newnone(py_retval());
        return true;
    }

    PY_CHECK_ARG_TYPE(1, tp_tuple);
    if(py_tuple_len(py_arg(1)) != 2) return TypeError("expected a tuple of 2 elements");
    py_Ref x = py_tuple_getitem(py_arg(1), 0);
    py_Ref y = py_tuple_getitem(py_arg(1), 1);
    if(py_isint(x) && py_isint(y)) {
        int col = py_toint(x);
        int row = py_toint(y);
        if(c11_array2d_like_is_valid(self, col, row)) {
            bool ok = self->f_set(self, col, row, value);
            if(!ok) return false;
            py_newnone(py_retval());
            return true;
        }
        return _array2d_like_IndexError(self, col, row);
    }

    if(py_istype(x, tp_slice) && py_istype(y, tp_slice)) {
        HANDLE_SLICE();
        if(py_isinstance(value, tp_array2d_like)) {
            c11_array2d_like* values = py_touserdata(value);
            if(!_check_same_shape(slice_width, slice_height, values->n_cols, values->n_rows))
                return false;
            for(int j = 0; j < slice_height; j++) {
                for(int i = 0; i < slice_width; i++) {
                    py_Ref item = values->f_get(values, i, j);
                    bool ok = self->f_set(self, start_col + i, start_row + j, item);
                    if(!ok) return false;
                }
            }
        } else {
            for(int j = 0; j < slice_height; j++) {
                for(int i = 0; i < slice_width; i++) {
                    bool ok = self->f_set(self, start_col + i, start_row + j, value);
                    if(!ok) return false;
                }
            }
        }
        py_newnone(py_retval());
        return true;
    }

    return TypeError("expected `tuple[int, int]` or `tuple[slice, slice]");
}

// count(self, value: T) -> int
static bool array2d_like_count(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    c11_array2d_like* self = py_touserdata(argv);
    int count = 0;
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            int code = py_equal(self->f_get(self, i, j), py_arg(1));
            if(code == -1) return false;
            count += code;
        }
    }
    py_newint(py_retval(), count);
    return true;
}

// get_bounding_rect(self, value: T) -> tuple[int, int, int, int]
static bool array2d_like_get_bounding_rect(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    c11_array2d_like* self = py_touserdata(argv);
    py_Ref value = py_arg(1);
    int left = self->n_cols;
    int top = self->n_rows;
    int right = 0;
    int bottom = 0;
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_Ref item = self->f_get(self, i, j);
            int res = py_equal(item, value);
            if(res == -1) return false;
            if(res == 1) {
                left = c11__min(left, i);
                top = c11__min(top, j);
                right = c11__max(right, i);
                bottom = c11__max(bottom, j);
            }
        }
    }
    int width = right - left + 1;
    int height = bottom - top + 1;
    if(width <= 0 || height <= 0) {
        return ValueError("value not found");
    } else {
        py_newtuple(py_retval(), 4);
        py_TValue* data = py_tuple_data(py_retval());
        py_newint(&data[0], left);
        py_newint(&data[1], top);
        py_newint(&data[2], width);
        py_newint(&data[3], height);
    }
    return true;
}

// count_neighbors(self, value: T, neighborhood: Neighborhood) -> array2d[int]
static bool array2d_like_count_neighbors(int argc, py_Ref argv) {
    PY_CHECK_ARGC(3);
    c11_array2d_like* self = py_touserdata(argv);
    c11_array2d* res = py_newarray2d(py_pushtmp(), self->n_cols, self->n_rows);
    py_Ref value = py_arg(1);
    const char* neighborhood = py_tostr(py_arg(2));

    const static c11_vec2i Moore[] = {
        {{-1, -1}},
        {{0, -1}},
        {{1, -1}},
        {{-1, 0}},
        {{1, 0}},
        {{-1, 1}},
        {{0, 1}},
        {{1, 1}},
    };

    const static c11_vec2i von_Neumann[] = {
        {{0, -1}},
        {{-1, 0}},
        {{1, 0}},
        {{0, 1}},
    };

    const c11_vec2i* offsets;
    int n_offsets;
    if(strcmp(neighborhood, "Moore") == 0) {
        offsets = Moore;
        n_offsets = c11__count_array(Moore);
    } else if(strcmp(neighborhood, "von Neumann") == 0) {
        offsets = von_Neumann;
        n_offsets = c11__count_array(von_Neumann);
    } else {
        return ValueError("neighborhood must be 'Moore' or 'von Neumann'");
    }
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_i64 count = 0;
            for(int k = 0; k < n_offsets; k++) {
                int x = i + offsets[k].x;
                int y = j + offsets[k].y;
                if(x >= 0 && x < self->n_cols && y >= 0 && y < self->n_rows) {
                    py_Ref item = self->f_get(self, x, y);
                    int code = py_equal(item, value);
                    if(code == -1) return false;
                    count += code;
                }
            }
            py_newint(c11_array2d__get(res, i, j), count);
        }
    }
    py_assign(py_retval(), py_peek(-1));
    py_pop();
    return true;
}

// convolve(self: array2d_like[int], kernel: array2d_like[int], padding: int) -> array2d[int]
static bool array2d_like_convolve(int argc, py_Ref argv) {
    PY_CHECK_ARGC(3);
    if(!py_checkinstance(&argv[1], tp_array2d_like)) return false;
    PY_CHECK_ARG_TYPE(2, tp_int);
    c11_array2d_like* self = py_touserdata(&argv[0]);
    c11_array2d_like* kernel = py_touserdata(&argv[1]);
    int padding = py_toint(py_arg(2));
    if(kernel->n_cols != kernel->n_rows) return ValueError("kernel must be square");
    int ksize = kernel->n_cols;
    if(ksize % 2 == 0) return ValueError("kernel size must be odd");
    int ksize_half = ksize / 2;
    c11_array2d* res = py_newarray2d(py_pushtmp(), self->n_cols, self->n_rows);
    for(int j = 0; j < self->n_rows; j++) {
        for(int i = 0; i < self->n_cols; i++) {
            py_i64 sum = 0;
            for(int jj = 0; jj < ksize; jj++) {
                for(int ii = 0; ii < ksize; ii++) {
                    int x = i + ii - ksize_half;
                    int y = j + jj - ksize_half;
                    py_i64 _0, _1;
                    if(x < 0 || x >= self->n_cols || y < 0 || y >= self->n_rows) {
                        _0 = padding;
                    } else {
                        py_Ref item = self->f_get(self, x, y);
                        if(!py_checkint(item)) return false;
                        _0 = py_toint(item);
                    }
                    py_Ref kitem = kernel->f_get(kernel, ii, jj);
                    if(!py_checkint(kitem)) return false;
                    _1 = py_toint(kitem);
                    sum += _0 * _1;
                }
            }
            py_newint(c11_array2d__get(res, i, j), sum);
        }
    }
    py_assign(py_retval(), py_peek(-1));
    py_pop();
    return true;
}

#undef HANDLE_SLICE

static void register_array2d_like(py_Ref mod) {
    py_Type type = py_newtype("array2d_like", tp_object, mod, NULL);
    assert(type == tp_array2d_like);

    py_bindproperty(type, "n_cols", array2d_like_n_cols, NULL);
    py_bindproperty(type, "n_rows", array2d_like_n_rows, NULL);
    py_bindproperty(type, "width", array2d_like_n_cols, NULL);
    py_bindproperty(type, "height", array2d_like_n_rows, NULL);
    py_bindproperty(type, "shape", array2d_like_shape, NULL);
    py_bindproperty(type, "numel", array2d_like_numel, NULL);

    py_bindmethod(type, "is_valid", array2d_like_is_valid);
    py_bindmethod(type, "get", array2d_like_get);

    py_bindmethod(type, "render", array2d_like_render);

    py_bindmethod(type, "all", array2d_like_all);
    py_bindmethod(type, "any", array2d_like_any);

    py_bindmethod(type, "map", array2d_like_map);
    py_bindmethod(type, "apply", array2d_like_apply);
    py_bindmethod(type, "copy", array2d_like_copy);
    py_bindmethod(type, "tolist", array2d_like_tolist);

    py_bindmagic(type, __eq__, array2d_like__eq__);
    py_bindmagic(type, __ne__, array2d_like__ne__);
    py_bindmagic(type, __iter__, array2d_like__iter__);
    py_bindmagic(type, __repr__, array2d_like__repr__);

    py_bindmagic(type, __getitem__, array2d_like__getitem__);
    py_bindmagic(type, __setitem__, array2d_like__setitem__);

    py_bindmethod(type, "count", array2d_like_count);
    py_bindmethod(type, "get_bounding_rect", array2d_like_get_bounding_rect);
    py_bindmethod(type, "count_neighbors", array2d_like_count_neighbors);
    py_bindmethod(type, "convolve", array2d_like_convolve);

    const char* scc =
        "\ndef get_connected_components(self, value: T, neighborhood: Neighborhood) -> tuple[array2d[int], int]:\n    from collections import deque\n    from linalg import vec2i\n\n    DIRS = [vec2i.LEFT, vec2i.RIGHT, vec2i.UP, vec2i.DOWN]\n    assert neighborhood in ['Moore', 'von Neumann']\n\n    if neighborhood == 'Moore':\n        DIRS.extend([\n            vec2i.LEFT+vec2i.UP,\n            vec2i.RIGHT+vec2i.UP,\n            vec2i.LEFT+vec2i.DOWN,\n            vec2i.RIGHT+vec2i.DOWN\n            ])\n\n    visited = array2d[int](self.width, self.height, default=0)\n    queue = deque()\n    count = 0\n    for y in range(self.height):\n        for x in range(self.width):\n            if visited[x, y] or self[x, y] != value:\n                continue\n            count += 1\n            queue.append((x, y))\n            visited[x, y] = count\n            while queue:\n                cx, cy = queue.popleft()\n                for dx, dy in DIRS:\n                    nx, ny = cx+dx, cy+dy\n                    if self.is_valid(nx, ny) and not visited[nx, ny] and self[nx, ny] == value:\n                        queue.append((nx, ny))\n                        visited[nx, ny] = count\n    return visited, count\n\narray2d_like.get_connected_components = get_connected_components\ndel get_connected_components\n";
    if(!py_exec(scc, "array2d.py", EXEC_MODE, mod)) {
        py_printexc();
        c11__abort("failed to execute array2d.py");
    }
}

static bool array2d_like_iterator__next__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_like_iterator* self = py_touserdata(argv);
    if(self->j >= self->array->n_rows) return StopIteration();
    py_newtuple(py_retval(), 2);
    py_TValue* data = py_tuple_data(py_retval());
    py_newvec2i(&data[0],
                (c11_vec2i){
                    {self->i, self->j}
    });
    py_assign(&data[1], self->array->f_get(self->array, self->i, self->j));
    self->i++;
    if(self->i >= self->array->n_cols) {
        self->i = 0;
        self->j++;
    }
    return true;
}

static void register_array2d_like_iterator(py_Ref mod) {
    py_Type type = py_newtype("array2d_like_iterator", tp_object, mod, NULL);
    assert(type == tp_array2d_like_iterator);
    py_bindmagic(type, __iter__, pk_wrapper__self);
    py_bindmagic(type, __next__, array2d_like_iterator__next__);
}

static bool array2d__new__(int argc, py_Ref argv) {
    // __new__(cls, n_cols: int, n_rows: int, default: Callable[[vec2i], T] = None)
    py_Ref default_ = py_arg(3);
    PY_CHECK_ARG_TYPE(0, tp_type);
    PY_CHECK_ARG_TYPE(1, tp_int);
    PY_CHECK_ARG_TYPE(2, tp_int);
    int n_cols = argv[1]._i64;
    int n_rows = argv[2]._i64;
    if(n_cols <= 0 || n_rows <= 0) return ValueError("array2d() expected positive dimensions");
    c11_array2d* ud = py_newarray2d(py_pushtmp(), n_cols, n_rows);
    // setup initial values
    if(py_callable(default_)) {
        for(int j = 0; j < n_rows; j++) {
            for(int i = 0; i < n_cols; i++) {
                py_TValue tmp;
                py_newvec2i(&tmp,
                            (c11_vec2i){
                                {i, j}
                });
                if(!py_call(default_, 1, &tmp)) return false;
                ud->data[j * n_cols + i] = *py_retval();
            }
        }
    } else {
        for(int i = 0; i < ud->header.numel; i++) {
            ud->data[i] = *default_;
        }
    }
    py_assign(py_retval(), py_peek(-1));
    py_pop();
    return true;
}

// fromlist(data: list[list[T]]) -> array2d[T]
static bool array2d_fromlist_STATIC(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    if(!py_checktype(argv, tp_list)) return false;
    int n_rows = py_list_len(argv);
    if(n_rows == 0) return ValueError("fromlist() expected a non-empty list");
    int n_cols = -1;
    for(int j = 0; j < n_rows; j++) {
        py_Ref row_j = py_list_getitem(argv, j);
        if(!py_checktype(row_j, tp_list)) return false;
        int n_cols_j = py_list_len(row_j);
        if(n_cols == -1) {
            if(n_cols_j == 0) return ValueError("fromlist() expected a non-empty list");
            n_cols = n_cols_j;
        } else if(n_cols != n_cols_j) {
            return ValueError("fromlist() expected a list of lists with the same length");
        }
    }
    c11_array2d* res = py_newarray2d(py_retval(), n_cols, n_rows);
    for(int j = 0; j < n_rows; j++) {
        py_Ref row_j = py_list_getitem(argv, j);
        for(int i = 0; i < n_cols; i++) {
            c11_array2d__set(res, i, j, py_list_getitem(row_j, i));
        }
    }
    return true;
}

static void register_array2d(py_Ref mod) {
    py_Type type = py_newtype("array2d", tp_array2d_like, mod, NULL);
    assert(type == tp_array2d);
    py_bind(py_tpobject(type),
            "__new__(cls, n_cols: int, n_rows: int, default=None)",
            array2d__new__);
    py_bindstaticmethod(type, "fromlist", array2d_fromlist_STATIC);
}

static bool array2d_view_origin(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_array2d_view* self = py_touserdata(argv);
    py_newvec2i(py_retval(), self->origin);
    return true;
}

static void register_array2d_view(py_Ref mod) {
    py_Type type = py_newtype("array2d_view", tp_array2d_like, mod, NULL);
    assert(type == tp_array2d_view);
    py_bindproperty(type, "origin", array2d_view_origin, NULL);
}

/* chunked_array2d */
#define SMALLMAP_T__SOURCE
#define K c11_vec2i
#define V py_TValue*
#define NAME c11_chunked_array2d_chunks
#define less(a, b) (a._i64 < b._i64)
#define equal(a, b) (a._i64 == b._i64)
#include "pocketpy/xmacros/smallmap.h"
#undef SMALLMAP_T__SOURCE

static py_TValue* c11_chunked_array2d__new_chunk(c11_chunked_array2d* self, c11_vec2i pos) {
#ifndef NDEBUG
    bool exists = c11_chunked_array2d_chunks__contains(&self->chunks, pos);
    assert(!exists);
#endif
    int chunk_numel = self->chunk_size * self->chunk_size + 1;
    py_TValue* data = PK_MALLOC(sizeof(py_TValue) * chunk_numel);
    if(!py_isnone(&self->context_builder)) {
        py_newvec2i(&data[0], pos);
        bool ok = py_call(&self->context_builder, 1, &data[0]);
        if(!ok) return NULL;
        data[0] = *py_retval();
    } else {
        data[0] = *py_None();
    }
    memset(&data[1], 0, sizeof(py_TValue) * (chunk_numel - 1));
    c11_chunked_array2d_chunks__set(&self->chunks, pos, data);
    self->last_visited.key = pos;
    self->last_visited.value = data;
    return data;
}

static void c11_chunked_array2d__world_to_chunk(c11_chunked_array2d* self,
                                                int col,
                                                int row,
                                                c11_vec2i* chunk_pos,
                                                c11_vec2i* local_pos) {
    if(col >= 0) {
        chunk_pos->x = col >> self->chunk_size_log2;
        local_pos->x = col & self->chunk_size_mask;
    } else {
        chunk_pos->x = -((-col) >> self->chunk_size_log2);
        local_pos->x = (-col) & self->chunk_size_mask;
    }
    if(row >= 0) {
        chunk_pos->y = row >> self->chunk_size_log2;
        local_pos->y = row & self->chunk_size_mask;
    } else {
        chunk_pos->y = -((-row) >> self->chunk_size_log2);
        local_pos->y = (-row) & self->chunk_size_mask;
    }
}

static py_TValue* c11_chunked_array2d__parse_col_row(c11_chunked_array2d* self,
                                                     int col,
                                                     int row,
                                                     c11_vec2i* chunk_pos,
                                                     c11_vec2i* local_pos) {
    c11_chunked_array2d__world_to_chunk(self, col, row, chunk_pos, local_pos);
    py_TValue* data;
    if(self->last_visited.value != NULL && chunk_pos->_i64 == self->last_visited.key._i64) {
        data = self->last_visited.value;
    } else {
        data = c11_chunked_array2d_chunks__get(&self->chunks, *chunk_pos, NULL);
    }
    if(data != NULL) {
        self->last_visited.key = *chunk_pos;
        self->last_visited.value = data;
    }
    return data;
}

py_Ref c11_chunked_array2d__get(c11_chunked_array2d* self, int col, int row) {
    c11_vec2i chunk_pos, local_pos;
    py_TValue* data = c11_chunked_array2d__parse_col_row(self, col, row, &chunk_pos, &local_pos);
    if(data == NULL) return &self->default_T;
    py_Ref retval = &data[1 + local_pos.y * self->chunk_size + local_pos.x];
    if(py_isnil(retval)) return &self->default_T;
    return retval;
}

bool c11_chunked_array2d__set(c11_chunked_array2d* self, int col, int row, py_Ref value) {
    c11_vec2i chunk_pos, local_pos;
    py_TValue* data = c11_chunked_array2d__parse_col_row(self, col, row, &chunk_pos, &local_pos);
    if(data == NULL) {
        data = c11_chunked_array2d__new_chunk(self, chunk_pos);
        if(data == NULL) return false;
    }
    data[1 + local_pos.y * self->chunk_size + local_pos.x] = *value;
    return true;
}

static void c11_chunked_array2d__del(c11_chunked_array2d* self, int col, int row) {
    c11_vec2i chunk_pos, local_pos;
    py_TValue* data = c11_chunked_array2d__parse_col_row(self, col, row, &chunk_pos, &local_pos);
    if(data != NULL) data[1 + local_pos.y * self->chunk_size + local_pos.x] = *py_NIL();
}

static bool chunked_array2d__new__(int argc, py_Ref argv) {
    py_Type cls = py_totype(argv);
    py_newobject(py_retval(), cls, 0, sizeof(c11_chunked_array2d));
    return true;
}

static bool chunked_array2d__init__(int argc, py_Ref argv) {
    c11_chunked_array2d* self = py_touserdata(&argv[0]);
    PY_CHECK_ARG_TYPE(1, tp_int);
    int chunk_size = py_toint(&argv[1]);
    self->default_T = argv[2];
    self->context_builder = argv[3];

    c11_chunked_array2d_chunks__ctor(&self->chunks);
    self->chunk_size = chunk_size;
    switch(chunk_size) {
        case 2: self->chunk_size_log2 = 1; break;
        case 4: self->chunk_size_log2 = 2; break;
        case 8: self->chunk_size_log2 = 3; break;
        case 16: self->chunk_size_log2 = 4; break;
        case 32: self->chunk_size_log2 = 5; break;
        case 64: self->chunk_size_log2 = 6; break;
        case 128: self->chunk_size_log2 = 7; break;
        case 256: self->chunk_size_log2 = 8; break;
        case 512: self->chunk_size_log2 = 9; break;
        case 1024: self->chunk_size_log2 = 10; break;
        case 2048: self->chunk_size_log2 = 11; break;
        case 4096: self->chunk_size_log2 = 12; break;
        default: return ValueError("invalid chunk_size: %d", chunk_size);
    }
    self->chunk_size_mask = chunk_size - 1;
    memset(&self->last_visited, 0, sizeof(c11_chunked_array2d_chunks_KV));
    py_newnone(py_retval());
    return true;
}

static bool chunked_array2d__chunk_size(int argc, py_Ref argv) {
    c11_chunked_array2d* self = py_touserdata(argv);
    py_newint(py_retval(), self->chunk_size);
    return true;
}

static bool chunked_array2d__getitem__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    c11_chunked_array2d* self = py_touserdata(argv);
    c11_vec2i pos = py_tovec2i(&argv[1]);
    py_Ref res = c11_chunked_array2d__get(self, pos.x, pos.y);
    py_assign(py_retval(), res);
    return true;
}

static bool chunked_array2d__setitem__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(3);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    c11_chunked_array2d* self = py_touserdata(argv);
    c11_vec2i pos = py_tovec2i(&argv[1]);
    bool ok = c11_chunked_array2d__set(self, pos.x, pos.y, &argv[2]);
    if(!ok) return false;
    py_newnone(py_retval());
    return true;
}

static bool chunked_array2d__delitem__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    c11_chunked_array2d* self = py_touserdata(argv);
    c11_vec2i pos = py_tovec2i(&argv[1]);
    c11_chunked_array2d__del(self, pos.x, pos.y);
    py_newnone(py_retval());
    return true;
}

static bool chunked_array2d__iter__(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_chunked_array2d* self = py_touserdata(argv);
    py_newtuple(py_retval(), self->chunks.length);
    for(int i = 0; i < self->chunks.length; i++) {
        py_TValue* data = c11__getitem(c11_chunked_array2d_chunks_KV, &self->chunks, i).value;
        py_tuple_setitem(py_retval(), i, &data[0]);
    }
    return py_iter(py_retval());
}

static bool chunked_array2d__clear(int argc, py_Ref argv) {
    c11_chunked_array2d* self = py_touserdata(argv);
    c11_chunked_array2d_chunks__clear(&self->chunks);
    self->last_visited.value = NULL;
    py_newnone(py_retval());
    return true;
}

static bool chunked_array2d__world_to_chunk(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    PY_CHECK_ARG_TYPE(1, tp_vec2);
    c11_chunked_array2d* self = py_touserdata(argv);
    c11_vec2i pos = py_tovec2i(&argv[1]);
    c11_vec2i chunk_pos, local_pos;
    c11_chunked_array2d__world_to_chunk(self, pos.x, pos.y, &chunk_pos, &local_pos);
    py_newtuple(py_retval(), 2);
    py_TValue* data = py_tuple_data(py_retval());
    py_newvec2i(&data[0], chunk_pos);
    py_newvec2i(&data[1], local_pos);
    return true;
}

static bool chunked_array2d__add_chunk(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    c11_chunked_array2d* self = py_touserdata(argv);
    c11_vec2i pos = py_tovec2i(&argv[1]);
    py_TValue* data = c11_chunked_array2d__new_chunk(self, pos);
    if(data == NULL) return false;
    py_assign(py_retval(), &data[0]);  // context
    return true;
}

static bool chunked_array2d__remove_chunk(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    c11_chunked_array2d* self = py_touserdata(argv);
    c11_vec2i pos = py_tovec2i(&argv[1]);
    bool ok = c11_chunked_array2d_chunks__del(&self->chunks, pos);
    self->last_visited.value = NULL;
    py_newbool(py_retval(), ok);
    return true;
}

static bool chunked_array2d__get_context(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    c11_chunked_array2d* self = py_touserdata(argv);
    c11_vec2i pos = py_tovec2i(&argv[1]);
    py_TValue* data = c11_chunked_array2d_chunks__get(&self->chunks, pos, NULL);
    if(data == NULL) {
        py_newnone(py_retval());
    } else {
        py_assign(py_retval(), &data[0]);
    }
    return true;
}

void c11_chunked_array2d__dtor(c11_chunked_array2d* self) {
    c11__foreach(c11_chunked_array2d_chunks_KV, &self->chunks, p_kv) PK_FREE(p_kv->value);
    c11_chunked_array2d_chunks__dtor(&self->chunks);
}

static void c11_chunked_array2d__mark(void* ud) {
    c11_chunked_array2d* self = ud;
    pk__mark_value(&self->default_T);
    pk__mark_value(&self->context_builder);
    int chunk_numel = self->chunk_size * self->chunk_size + 1;
    for(int i = 0; i < self->chunks.length; i++) {
        py_TValue* data = c11__getitem(c11_chunked_array2d_chunks_KV, &self->chunks, i).value;
        for(int j = 0; j < chunk_numel; j++) {
            pk__mark_value(data + j);
        }
    }
}

static bool chunked_array2d_view(int argc, py_Ref argv) {
    PY_CHECK_ARGC(1);
    c11_chunked_array2d* self = py_touserdata(&argv[0]);
    if(self->chunks.length == 0) { return ValueError("chunked_array2d is empty"); }
    int min_chunk_x = INT_MAX;
    int min_chunk_y = INT_MAX;
    int max_chunk_x = INT_MIN;
    int max_chunk_y = INT_MIN;
    for(int i = 0; i < self->chunks.length; i++) {
        c11_vec2i chunk_pos = c11__getitem(c11_chunked_array2d_chunks_KV, &self->chunks, i).key;
        min_chunk_x = c11__min(min_chunk_x, chunk_pos.x);
        min_chunk_y = c11__min(min_chunk_y, chunk_pos.y);
        max_chunk_x = c11__max(max_chunk_x, chunk_pos.x);
        max_chunk_y = c11__max(max_chunk_y, chunk_pos.y);
    }
    int start_col = min_chunk_x << self->chunk_size_log2;
    int start_row = min_chunk_y << self->chunk_size_log2;
    int width = (max_chunk_x - min_chunk_x + 1) * self->chunk_size;
    int height = (max_chunk_y - min_chunk_y + 1) * self->chunk_size;
    return _chunked_array2d_view(py_retval(), argv, self, start_col, start_row, width, height);
}

static bool chunked_array2d_view_rect(int argc, py_Ref argv) {
    PY_CHECK_ARGC(4);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    PY_CHECK_ARG_TYPE(2, tp_int);
    PY_CHECK_ARG_TYPE(3, tp_int);
    c11_chunked_array2d* self = py_touserdata(&argv[0]);
    c11_vec2i pos = py_tovec2i(&argv[1]);
    int width = py_toint(&argv[2]);
    int height = py_toint(&argv[3]);
    return _chunked_array2d_view(py_retval(), argv, self, pos.x, pos.y, width, height);
}

static bool chunked_array2d_view_chunk(int argc, py_Ref argv) {
    PY_CHECK_ARGC(2);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    c11_chunked_array2d* self = py_touserdata(&argv[0]);
    c11_vec2i chunk_pos = py_tovec2i(&argv[1]);
    int start_col = chunk_pos.x << self->chunk_size_log2;
    int start_row = chunk_pos.y << self->chunk_size_log2;
    return _chunked_array2d_view(py_retval(),
                                 argv,
                                 self,
                                 start_col,
                                 start_row,
                                 self->chunk_size,
                                 self->chunk_size);
}

static bool chunked_array2d_view_chunks(int argc, py_Ref argv) {
    PY_CHECK_ARGC(4);
    PY_CHECK_ARG_TYPE(1, tp_vec2i);
    PY_CHECK_ARG_TYPE(2, tp_int);
    PY_CHECK_ARG_TYPE(3, tp_int);
    c11_chunked_array2d* self = py_touserdata(&argv[0]);
    c11_vec2i chunk_pos = py_tovec2i(&argv[1]);
    int width = py_toint(&argv[2]) * self->chunk_size;
    int height = py_toint(&argv[3]) * self->chunk_size;
    int start_col = chunk_pos.x << self->chunk_size_log2;
    int start_row = chunk_pos.y << self->chunk_size_log2;
    return _chunked_array2d_view(py_retval(), argv, self, start_col, start_row, width, height);
}

static void register_chunked_array2d(py_Ref mod) {
    py_Type type =
        py_newtype("chunked_array2d", tp_object, mod, (py_Dtor)c11_chunked_array2d__dtor);
    pk__tp_set_marker(type, c11_chunked_array2d__mark);
    assert(type == tp_chunked_array2d);

    py_bind(py_tpobject(type), "__new__(cls, *args, **kwargs)", chunked_array2d__new__);
    py_bind(py_tpobject(type),
            "__init__(self, chunk_size, default=None, context_builder=None)",
            chunked_array2d__init__);

    py_bindproperty(type, "chunk_size", chunked_array2d__chunk_size, NULL);

    py_bindmagic(type, __getitem__, chunked_array2d__getitem__);
    py_bindmagic(type, __setitem__, chunked_array2d__setitem__);
    py_bindmagic(type, __delitem__, chunked_array2d__delitem__);
    py_bindmagic(type, __iter__, chunked_array2d__iter__);

    py_bindmethod(type, "clear", chunked_array2d__clear);
    py_bindmethod(type, "world_to_chunk", chunked_array2d__world_to_chunk);
    py_bindmethod(type, "add_chunk", chunked_array2d__add_chunk);
    py_bindmethod(type, "remove_chunk", chunked_array2d__remove_chunk);
    py_bindmethod(type, "get_context", chunked_array2d__get_context);

    py_bindmethod(type, "view", chunked_array2d_view);
    py_bindmethod(type, "view_rect", chunked_array2d_view_rect);
    py_bindmethod(type, "view_chunk", chunked_array2d_view_chunk);
    py_bindmethod(type, "view_chunks", chunked_array2d_view_chunks);
}

void pk__add_module_array2d() {
    py_GlobalRef mod = py_newmodule("array2d");

    register_array2d_like(mod);
    register_array2d_like_iterator(mod);
    register_array2d(mod);
    register_array2d_view(mod);
    register_chunked_array2d(mod);
}