From 93cd5e48a72e75c6c725f93e304f8fb6cb7c980b Mon Sep 17 00:00:00 2001 From: BLUELOVETH Date: Thu, 13 Feb 2025 16:08:47 +0800 Subject: [PATCH] implement `array2d.chunked_array2d[T, TContext]` (#332) * bak * backup * ... * Update array2d.pyi * backup * backup * backup * backup * backup * backup * backup --- include/pocketpy/interpreter/array2d.h | 57 +- include/pocketpy/linalg.h | 3 + include/pocketpy/pocketpy.h | 11 +- include/typings/array2d.pyi | 122 +- src/modules/array2d.c | 1459 +++++++++++++++--------- src/modules/pickle.c | 13 +- src/public/cast.c | 5 + tests/90_array2d.py | 31 +- tests/90_chunked_array2d.py | 8 + 9 files changed, 1125 insertions(+), 584 deletions(-) create mode 100644 tests/90_chunked_array2d.py diff --git a/include/pocketpy/interpreter/array2d.h b/include/pocketpy/interpreter/array2d.h index ac040e67..6bf052a2 100644 --- a/include/pocketpy/interpreter/array2d.h +++ b/include/pocketpy/interpreter/array2d.h @@ -1,21 +1,58 @@ #pragma once #include "pocketpy/pocketpy.h" +#include "pocketpy/common/smallmap.h" +#include "pocketpy/objects/base.h" -#include "pocketpy/common/utils.h" -#include "pocketpy/common/sstream.h" -#include "pocketpy/interpreter/vm.h" - -typedef struct c11_array2d { - py_TValue* data; // slots +typedef struct c11_array2d_like { int n_cols; int n_rows; int numel; + py_Ref (*f_get)(struct c11_array2d_like* self, int col, int row); + bool (*f_set)(struct c11_array2d_like* self, int col, int row, py_Ref value); +} c11_array2d_like; + +typedef struct c11_array2d_like_iterator { + c11_array2d_like* array; + int j; + int i; +} c11_array2d_like_iterator; + +typedef struct c11_array2d { + c11_array2d_like header; + py_TValue* data; // slots } c11_array2d; -typedef struct c11_array2d_iterator { - c11_array2d* array; - int index; -} c11_array2d_iterator; +typedef struct c11_array2d_view { + c11_array2d_like header; + void* ctx; + py_Ref (*f_get)(void* ctx, int col, int row); + bool (*f_set)(void* ctx, int col, int row, py_Ref value); + c11_vec2i origin; +} c11_array2d_view; c11_array2d* py_newarray2d(py_OutRef out, int n_cols, int n_rows); + +/* chunked_array2d */ +#define SMALLMAP_T__HEADER +#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__HEADER + +typedef struct c11_chunked_array2d { + c11_chunked_array2d_chunks chunks; + int chunk_size; + int chunk_size_log2; + int chunk_size_mask; + c11_chunked_array2d_chunks_KV last_visited; + + py_TValue default_T; + py_TValue context_builder; +} c11_chunked_array2d; + +py_Ref c11_chunked_array2d__get(c11_chunked_array2d* self, int col, int row); +bool c11_chunked_array2d__set(c11_chunked_array2d* self, int col, int row, py_Ref value); diff --git a/include/pocketpy/linalg.h b/include/pocketpy/linalg.h index a77907f3..38e37228 100644 --- a/include/pocketpy/linalg.h +++ b/include/pocketpy/linalg.h @@ -1,8 +1,11 @@ #pragma once +#include + typedef union c11_vec2i { struct { int x, y; }; int data[2]; + int64_t _i64; } c11_vec2i; typedef union c11_vec3i { diff --git a/include/pocketpy/pocketpy.h b/include/pocketpy/pocketpy.h index fff7920e..d8a70302 100644 --- a/include/pocketpy/pocketpy.h +++ b/include/pocketpy/pocketpy.h @@ -302,10 +302,14 @@ PK_API const char* py_tpname(py_Type type); /// Call a type to create a new instance. PK_API bool py_tpcall(py_Type type, int argc, py_Ref argv) PY_RAISE PY_RETURN; -/// Check if the object is an instance of the given type. +/// Check if the object is an instance of the given type exactly. /// Raise `TypeError` if the check fails. PK_API bool py_checktype(py_Ref self, py_Type type) PY_RAISE; +/// Check if the object is an instance of the given type or its subclass. +/// Raise `TypeError` if the check fails. +PK_API bool py_checkinstance(py_Ref self, py_Type type) PY_RAISE; + #define py_checkint(self) py_checktype(self, tp_int) #define py_checkfloat(self) py_checktype(self, tp_float) #define py_checkbool(self) py_checktype(self, tp_bool) @@ -737,8 +741,11 @@ enum py_PredefinedTypes { tp_vec3i, tp_mat3x3, /* array2d */ + tp_array2d_like, + tp_array2d_like_iterator, tp_array2d, - tp_array2d_iterator, + tp_array2d_view, + tp_chunked_array2d, }; #ifdef __cplusplus diff --git a/include/typings/array2d.pyi b/include/typings/array2d.pyi index 50feb3d0..48d2cb62 100644 --- a/include/typings/array2d.pyi +++ b/include/typings/array2d.pyi @@ -1,9 +1,9 @@ -from typing import Callable, Any, Generic, TypeVar, Literal, overload, Iterator +from typing import Callable, Literal, overload, Iterator from linalg import vec2i Neighborhood = Literal['Moore', 'von Neumann'] -class array2d[T]: +class array2d_like[T]: @property def n_cols(self) -> int: ... @property @@ -13,55 +13,53 @@ class array2d[T]: @property def height(self) -> int: ... @property + def shape(self) -> vec2i: ... + @property def numel(self) -> int: ... - def __new__(cls, n_cols: int, n_rows: int, default: T | Callable[[vec2i], T] | None = None): ... - def __eq__(self, other: object) -> array2d[bool]: ... # type: ignore - def __ne__(self, other: object) -> array2d[bool]: ... # type: ignore - def __repr__(self) -> str: ... - def __iter__(self) -> Iterator[tuple[vec2i, T]]: ... - @overload def is_valid(self, col: int, row: int) -> bool: ... @overload def is_valid(self, pos: vec2i) -> bool: ... def get[R](self, col: int, row: int, default: R = None) -> T | R: - """Gets the value at the given position. If the position is out of bounds, return the default value.""" - - @overload - def __getitem__(self, index: tuple[int, int]) -> T: ... - @overload - def __getitem__(self, index: vec2i) -> T: ... - @overload - def __getitem__(self, index: tuple[slice, slice]) -> array2d[T]: ... - @overload - def __getitem__(self, mask: array2d[bool]) -> list[T]: ... - @overload - def __setitem__(self, index: tuple[int, int], value: T): ... - @overload - def __setitem__(self, index: vec2i, value: T): ... - @overload - def __setitem__(self, index: tuple[slice, slice], value: int | float | str | bool | None | 'array2d[T]'): ... - @overload - def __setitem__(self, mask: array2d[bool], value: T): ... - - def map[R](self, f: Callable[[T], R]) -> array2d[R]: ... - def copy(self) -> 'array2d[T]': ... - - def fill_(self, value: T) -> None: ... - def apply_(self, f: Callable[[T], T]) -> None: ... - def copy_(self, other: array2d[T] | list[T]) -> None: ... + """Get the value at the given position. + + If the position is out of bounds, return the default value. + """ def render(self) -> str: ... - def all(self: array2d[bool]) -> bool: ... - def any(self: array2d[bool]) -> bool: ... - - @staticmethod - def fromlist(data: list[list[T]]) -> array2d[T]: ... + def all(self: array2d_like[bool]) -> bool: ... + def any(self: array2d_like[bool]) -> bool: ... + + def map[R](self, f: Callable[[T], R]) -> array2d[R]: ... + def apply(self, f: Callable[[T], T]) -> None: ... + def copy(self) -> 'array2d[T]': ... def tolist(self) -> list[list[T]]: ... + def __eq__(self, other: object) -> array2d[bool]: ... # type: ignore + def __ne__(self, other: object) -> array2d[bool]: ... # type: ignore + def __iter__(self) -> Iterator[tuple[vec2i, T]]: ... + def __repr__(self) -> str: ... + + @overload + def __getitem__(self, index: vec2i) -> T: ... + @overload + def __getitem__(self, index: tuple[int, int]) -> T: ... + @overload + def __getitem__(self, index: tuple[slice, slice]) -> array2d_view[T]: ... + @overload + def __getitem__(self, mask: array2d_like[bool]) -> list[T]: ... + @overload + def __setitem__(self, index: vec2i, value: T): ... + @overload + def __setitem__(self, index: tuple[int, int], value: T): ... + @overload + def __setitem__(self, index: tuple[slice, slice], value: T | 'array2d_like[T]'): ... + @overload + def __setitem__(self, mask: array2d_like[bool], value: T): ... + # algorithms def count(self, value: T) -> int: """Counts the number of cells with the given value.""" @@ -75,7 +73,7 @@ class array2d[T]: Returns a tuple `(x, y, width, height)` or raise `ValueError` if the value is not found. """ - def convolve(self: array2d[int], kernel: array2d[int], padding: int) -> array2d[int]: + def convolve(self: array2d_like[int], kernel: array2d_like[int], padding: int) -> array2d[int]: """Convolves the array with the given kernel.""" def get_connected_components(self, value: T, neighborhood: Neighborhood) -> tuple[array2d[int], int]: @@ -84,3 +82,49 @@ class array2d[T]: Returns the `visited` array and the number of connected components, where `0` means unvisited, and non-zero means the index of the connected component. """ + + +class array2d_view[T](array2d_like[T]): + @property + def origin(self) -> vec2i: ... + + +class array2d[T](array2d_like[T]): + def __new__( + cls, + n_cols: int, + n_rows: int, + default: T | Callable[[vec2i], T] | None = None + ): ... + + @staticmethod + def fromlist(data: list[list[T]]) -> array2d[T]: ... + + +class chunked_array2d[T, TContext]: + def __init__( + self, + chunk_size: int, + default: T = None, + context_builder: Callable[[vec2i], TContext] | None = None, + ): ... + + @property + def chunk_size(self) -> int: ... + + def __getitem__(self, index: vec2i) -> T: ... + def __setitem__(self, index: vec2i, value: T): ... + def __delitem__(self, index: vec2i): ... + def __iter__(self) -> Iterator[tuple[vec2i, TContext]]: ... + + def clear(self) -> None: ... + + def world_to_chunk(self, world_pos: vec2i) -> tuple[vec2i, vec2i]: ... + def add_chunk(self, chunk_pos: vec2i) -> TContext: ... + def remove_chunk(self, chunk_pos: vec2i) -> bool: ... + def get_context(self, chunk_pos: vec2i) -> TContext | None: ... + + def view(self) -> array2d_view[T]: ... + def view_rect(self, pos: vec2i, width: int, height: int) -> array2d_view[T]: ... + def view_chunk(self, chunk_pos: vec2i) -> array2d_view[T]: ... + def view_chunks(self, chunk_pos: vec2i, width: int, height: int) -> array2d_view[T]: ... diff --git a/src/modules/array2d.c b/src/modules/array2d.c index 9a5bdae7..6ef39894 100644 --- a/src/modules/array2d.c +++ b/src/modules/array2d.c @@ -1,86 +1,67 @@ #include "pocketpy/interpreter/array2d.h" +#include "pocketpy/interpreter/vm.h" +#include "pocketpy/pocketpy.h" +#include -static bool py_array2d_is_valid(c11_array2d* self, int col, int row) { - return col >= 0 && col < self->n_cols && row >= 0 && row < self->n_rows; +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_ObjectRef py_array2d__get(c11_array2d* self, int col, int row) { - return self->data + row * self->n_cols + col; +static py_Ref c11_array2d__get(c11_array2d* self, int col, int row) { + return self->data + row * self->header.n_cols + col; } -static void py_array2d__set(c11_array2d* self, int col, int row, py_Ref value) { - self->data[row * self->n_cols + col] = *value; +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); - ud->n_cols = n_cols; - ud->n_rows = n_rows; - ud->numel = numel; return ud; } -/* bindings */ -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; - int numel = n_cols * n_rows; - 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} - }); - bool ok = py_call(default_, 1, &tmp); - if(!ok) return false; - ud->data[j * n_cols + i] = *py_retval(); - } - } - } else { - for(int i = 0; i < numel; i++) { - ud->data[i] = *default_; - } - } - py_assign(py_retval(), py_peek(-1)); - py_pop(); - return true; -} - -static bool array2d_n_cols(int argc, py_Ref argv) { +/* array2d_like bindings */ +static bool array2d_like_n_cols(int argc, py_Ref argv) { PY_CHECK_ARGC(1); - c11_array2d* self = py_touserdata(argv); + c11_array2d_like* self = py_touserdata(argv); py_newint(py_retval(), self->n_cols); return true; } -static bool array2d_n_rows(int argc, py_Ref argv) { +static bool array2d_like_n_rows(int argc, py_Ref argv) { PY_CHECK_ARGC(1); - c11_array2d* self = py_touserdata(argv); + c11_array2d_like* self = py_touserdata(argv); py_newint(py_retval(), self->n_rows); return true; } -static bool array2d_numel(int argc, py_Ref argv) { +static bool array2d_like_shape(int argc, py_Ref argv) { PY_CHECK_ARGC(1); - c11_array2d* self = py_touserdata(argv); + 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_is_valid(int argc, py_Ref argv) { - c11_array2d* self = py_touserdata(argv); +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); @@ -95,15 +76,15 @@ static bool array2d_is_valid(int argc, py_Ref argv) { } else { return TypeError("is_valid() expected 2 or 3 arguments"); } - py_newbool(py_retval(), py_array2d_is_valid(self, col, row)); + py_newbool(py_retval(), c11_array2d_like_is_valid(self, col, row)); return true; } -static bool array2d_get(int argc, py_Ref argv) { - py_Ref default_; - c11_array2d* self = py_touserdata(argv); +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) { @@ -113,20 +94,124 @@ static bool array2d_get(int argc, py_Ref argv) { } int col = py_toint(py_arg(1)); int row = py_toint(py_arg(2)); - if(py_array2d_is_valid(self, col, row)) { - py_assign(py_retval(), py_array2d__get(self, col, row)); + 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_check_all_type(c11_array2d* self, py_Type type) { - for(int i = 0; i < self->numel; i++) { - py_Type item_type = self->data[i].type; - if(item_type != type) { - const char* fmt = "expected array2d[%t], got %t"; - return TypeError(fmt, type, item_type); +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; @@ -140,56 +225,35 @@ static bool _check_same_shape(int colA, int rowA, int colB, int rowB) { return true; } -static bool _array2d_check_same_shape(c11_array2d* self, c11_array2d* other) { +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_all(int argc, py_Ref argv) { - PY_CHECK_ARGC(1); - c11_array2d* self = py_touserdata(argv); - if(!_array2d_check_all_type(self, tp_bool)) return false; - for(int i = 0; i < self->numel; i++) { - if(!py_tobool(self->data + i)) { - py_newbool(py_retval(), false); - return true; - } - } - py_newbool(py_retval(), true); - return true; -} - -static bool array2d_any(int argc, py_Ref argv) { - PY_CHECK_ARGC(1); - c11_array2d* self = py_touserdata(argv); - if(!_array2d_check_all_type(self, tp_bool)) return false; - for(int i = 0; i < self->numel; i++) { - if(py_tobool(self->data + i)) { - py_newbool(py_retval(), true); - return true; - } - } - py_newbool(py_retval(), false); - return true; -} - -static bool array2d__eq__(int argc, py_Ref argv) { +static bool array2d_like__eq__(int argc, py_Ref argv) { PY_CHECK_ARGC(2); - c11_array2d* self = py_touserdata(argv); + c11_array2d_like* self = py_touserdata(argv); c11_array2d* res = py_newarray2d(py_pushtmp(), self->n_cols, self->n_rows); - if(py_istype(py_arg(1), tp_array2d)) { - c11_array2d* other = py_touserdata(py_arg(1)); - if(!_array2d_check_same_shape(self, other)) return false; - for(int i = 0; i < self->numel; i++) { - int code = py_equal(self->data + i, other->data + i); - if(code == -1) return false; - py_newbool(res->data + i, (bool)code); + 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 i = 0; i < self->numel; i++) { - int code = py_equal(self->data + i, py_arg(1)); - if(code == -1) return false; - py_newbool(res->data + i, (bool)code); + 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)); @@ -197,209 +261,273 @@ static bool array2d__eq__(int argc, py_Ref argv) { return true; } -static bool array2d__ne__(int argc, py_Ref argv) { - bool ok = array2d__eq__(argc, argv); +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->numel; i++) { + 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__repr__(int argc, py_Ref argv) { +static bool array2d_like__iter__(int argc, py_Ref argv) { PY_CHECK_ARGC(1); - c11_array2d* self = py_touserdata(argv); + 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), "array2d(%d, %d)", self->n_cols, self->n_rows); + 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; } -static bool array2d__iter__(int argc, py_Ref argv) { - PY_CHECK_ARGC(1); - c11_array2d* self = py_touserdata(argv); - c11_array2d_iterator* ud = - py_newobject(py_retval(), tp_array2d_iterator, 1, sizeof(c11_array2d_iterator)); - py_setslot(py_retval(), 0, argv); // keep the array alive - ud->array = self; - ud->index = 0; +#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; } -// __iter__(self) -> Iterator[tuple[int, int, T]] -static bool array2d_iterator__next__(int argc, py_Ref argv) { - PY_CHECK_ARGC(1); - c11_array2d_iterator* self = py_touserdata(argv); - if(self->index < self->array->numel) { - div_t res = div(self->index, self->array->n_cols); - py_newtuple(py_retval(), 2); - py_TValue* data = py_tuple_data(py_retval()); - py_newvec2i(&data[0], - (c11_vec2i){ - {res.rem, res.quot} - }); - py_assign(&data[1], self->array->data + self->index); - self->index++; +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; } - return StopIteration(); -} -static bool array2d_map(int argc, py_Ref argv) { - // def map(self, f: Callable[[T], Any]) -> 'array2d': ... - PY_CHECK_ARGC(2); - c11_array2d* 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 i = 0; i < self->numel; i++) { - bool ok = py_call(f, 1, self->data + i); - if(!ok) return false; - res->data[i] = *py_retval(); + 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); } - py_assign(py_retval(), py_peek(-1)); - py_pop(); - return true; -} -static bool array2d_copy(int argc, py_Ref argv) { - // def copy(self) -> 'array2d': ... - PY_CHECK_ARGC(1); - c11_array2d* self = py_touserdata(argv); - c11_array2d* res = py_newarray2d(py_retval(), self->n_cols, self->n_rows); - memcpy(res->data, self->data, self->numel * sizeof(py_TValue)); - return true; -} - -static bool array2d_fill_(int argc, py_Ref argv) { - // def fill_(self, value: T) -> None: ... - PY_CHECK_ARGC(2); - c11_array2d* self = py_touserdata(argv); - for(int i = 0; i < self->numel; i++) - self->data[i] = argv[1]; - py_newnone(py_retval()); - return true; -} - -static bool array2d_apply_(int argc, py_Ref argv) { - // def apply_(self, f: Callable[[T], T]) -> None: ... - PY_CHECK_ARGC(2); - c11_array2d* self = py_touserdata(argv); - py_Ref f = py_arg(1); - for(int i = 0; i < self->numel; i++) { - bool ok = py_call(f, 1, self->data + i); - if(!ok) return false; - self->data[i] = *py_retval(); + 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); } - py_newnone(py_retval()); - return true; + + return TypeError("expected `tuple[int, int]` or `tuple[slice, slice]`"); } -static bool array2d_copy_(int argc, py_Ref argv) { - // def copy_(self, src: 'array2d') -> None: ... - PY_CHECK_ARGC(2); - c11_array2d* self = py_touserdata(argv); +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); + } - py_Type src_type = py_typeof(py_arg(1)); - if(src_type == tp_array2d) { - c11_array2d* src = py_touserdata(py_arg(1)); - if(!_array2d_check_same_shape(self, src)) return false; - memcpy(self->data, src->data, self->numel * sizeof(py_TValue)); - } else { - py_TValue* data; - int length = pk_arrayview(py_arg(1), &data); - if(length != -1) { - if(self->numel != length) { - return ValueError("copy_() expected the same numel: %d != %d", self->numel, length); + 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; + } } - memcpy(self->data, data, self->numel * sizeof(py_TValue)); } else { - return TypeError("copy_() expected `array2d`, `list` or `tuple`, got '%t", src_type); + 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; } - py_newnone(py_retval()); - 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++) { - py_array2d__set(res, i, j, py_list_getitem(row_j, i)); - } - } - return true; -} - -// tolist(self) -> list[list[T]] -static bool array2d_tolist(int argc, py_Ref argv) { - PY_CHECK_ARGC(1); - c11_array2d* 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_list_setitem(row_j, i, py_array2d__get(self, i, j)); - } - } - return true; -} - -static bool array2d_render(int argc, py_Ref argv) { - PY_CHECK_ARGC(1); - c11_sbuf buf; - c11_sbuf__ctor(&buf); - c11_array2d* 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 = py_array2d__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; + return TypeError("expected `tuple[int, int]` or `tuple[slice, slice]"); } // count(self, value: T) -> int -static bool array2d_count(int argc, py_Ref argv) { +static bool array2d_like_count(int argc, py_Ref argv) { PY_CHECK_ARGC(2); - c11_array2d* self = py_touserdata(argv); + c11_array2d_like* self = py_touserdata(argv); int count = 0; - for(int i = 0; i < self->numel; i++) { - int res = py_equal(self->data + i, &argv[1]); - if(res == -1) return false; - count += res; + 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_get_bounding_rect(int argc, py_Ref argv) { +static bool array2d_like_get_bounding_rect(int argc, py_Ref argv) { PY_CHECK_ARGC(2); - c11_array2d* self = py_touserdata(argv); + c11_array2d_like* self = py_touserdata(argv); py_Ref value = py_arg(1); int left = self->n_cols; int top = self->n_rows; @@ -407,7 +535,8 @@ static bool array2d_get_bounding_rect(int argc, py_Ref argv) { int bottom = 0; for(int j = 0; j < self->n_rows; j++) { for(int i = 0; i < self->n_cols; i++) { - int res = py_equal(py_array2d__get(self, i, j), value); + 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); @@ -433,217 +562,75 @@ static bool array2d_get_bounding_rect(int argc, py_Ref argv) { } // count_neighbors(self, value: T, neighborhood: Neighborhood) -> array2d[int] -static bool array2d_count_neighbors(int argc, py_Ref argv) { +static bool array2d_like_count_neighbors(int argc, py_Ref argv) { PY_CHECK_ARGC(3); - c11_array2d* self = py_touserdata(argv); + 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)); -#define INC_COUNT(i, j) \ - do { \ - if(py_array2d_is_valid(self, i, j)) { \ - int res = py_equal(py_array2d__get(self, i, j), value); \ - if(res == -1) return false; \ - count += res; \ - } \ - } while(0) + 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) { - for(int j = 0; j < res->n_rows; j++) { - for(int i = 0; i < res->n_cols; i++) { - int count = 0; - INC_COUNT(i - 1, j - 1); - INC_COUNT(i, j - 1); - INC_COUNT(i + 1, j - 1); - INC_COUNT(i - 1, j); - INC_COUNT(i + 1, j); - INC_COUNT(i - 1, j + 1); - INC_COUNT(i, j + 1); - INC_COUNT(i + 1, j + 1); - py_newint(py_array2d__get(res, i, j), count); - } - } + offsets = Moore; + n_offsets = c11__count_array(Moore); } else if(strcmp(neighborhood, "von Neumann") == 0) { - for(int j = 0; j < res->n_rows; j++) { - for(int i = 0; i < res->n_cols; i++) { - int count = 0; - INC_COUNT(i, j - 1); - INC_COUNT(i - 1, j); - INC_COUNT(i + 1, j); - INC_COUNT(i, j + 1); - py_newint(py_array2d__get(res, i, j), count); - } - } + 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; } -#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; \ - if(slice_width <= 0 || slice_height <= 0) \ - return ValueError("slice width and height must be positive"); - -static bool _array2d_IndexError(c11_array2d* self, int col, int row) { - return IndexError("(%d, %d) is not a valid index of array2d(%d, %d)", - col, - row, - self->n_cols, - self->n_rows); -} - -static bool array2d__getitem__(int argc, py_Ref argv) { - PY_CHECK_ARGC(2); - c11_array2d* self = py_touserdata(argv); - if(argv[1].type == tp_vec2i) { - // fastpath for vec2i - c11_vec2i pos = py_tovec2i(&argv[1]); - if(py_array2d_is_valid(self, pos.x, pos.y)) { - py_assign(py_retval(), py_array2d__get(self, pos.x, pos.y)); - return true; - } - return _array2d_IndexError(self, pos.x, pos.y); - } - - if(argv[1].type == tp_array2d) { - c11_array2d* mask = py_touserdata(&argv[1]); - if(!_array2d_check_same_shape(self, mask)) return false; - if(!_array2d_check_all_type(mask, tp_bool)) return false; - py_newlist(py_retval()); - for(int i = 0; i < self->numel; i++) { - if(py_tobool(mask->data + i)) py_list_append(py_retval(), self->data + i); - } - 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(py_array2d_is_valid(self, col, row)) { - py_assign(py_retval(), py_array2d__get(self, col, row)); - return true; - } - return _array2d_IndexError(self, col, row); - } else if(py_istype(x, tp_slice) && py_istype(y, tp_slice)) { - HANDLE_SLICE(); - c11_array2d* res = py_newarray2d(py_retval(), slice_width, slice_height); - for(int j = start_row; j < stop_row; j++) { - for(int i = start_col; i < stop_col; i++) { - py_array2d__set(res, i - start_col, j - start_row, py_array2d__get(self, i, j)); - } - } - return true; - } else { - return TypeError("expected a tuple of 2 ints or slices"); - } -} - -static bool array2d__setitem__(int argc, py_Ref argv) { +// 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); - c11_array2d* self = py_touserdata(argv); - py_Ref value = py_arg(2); - if(argv[1].type == tp_vec2i) { - // fastpath for vec2i - c11_vec2i pos = py_tovec2i(&argv[1]); - if(py_array2d_is_valid(self, pos.x, pos.y)) { - py_array2d__set(self, pos.x, pos.y, value); - py_newnone(py_retval()); - return true; - } - return _array2d_IndexError(self, pos.x, pos.y); - } - - if(argv[1].type == tp_array2d) { - c11_array2d* mask = py_touserdata(&argv[1]); - if(!_array2d_check_same_shape(self, mask)) return false; - if(!_array2d_check_all_type(mask, tp_bool)) return false; - for(int i = 0; i < self->numel; i++) { - if(py_tobool(mask->data + i)) self->data[i] = *value; - } - 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(py_array2d_is_valid(self, col, row)) { - py_array2d__set(self, col, row, value); - py_newnone(py_retval()); - return true; - } - return _array2d_IndexError(self, col, row); - } else if(py_istype(x, tp_slice) && py_istype(y, tp_slice)) { - HANDLE_SLICE(); - bool is_basic_type = false; - switch(value->type) { - case tp_int: - case tp_float: - case tp_str: - case tp_NoneType: - case tp_bool: is_basic_type = true; break; - default: { - if(!py_istype(value, tp_array2d)) { - return TypeError("expected int/float/str/bool/None or an array2d instance"); - } - } - } - - if(is_basic_type) { - for(int j = start_row; j < stop_row; j++) { - for(int i = start_col; i < stop_col; i++) { - py_array2d__set(self, i, j, py_arg(2)); - } - } - } else { - c11_array2d* src = py_touserdata(value); - if(!_check_same_shape(slice_width, slice_height, src->n_cols, src->n_rows)) - return false; - for(int j = 0; j < slice_height; j++) { - for(int i = 0; i < slice_width; i++) { - py_array2d__set(self, i + start_col, j + start_row, py_array2d__get(src, i, j)); - } - } - } - py_newnone(py_retval()); - return true; - } else { - return TypeError("expected a tuple of 2 ints or slices"); - } -} - -// convolve(self: array2d[int], kernel: array2d[int], padding: int) -> array2d[int] -static bool array2d_convolve(int argc, py_Ref argv) { - PY_CHECK_ARGC(3); - PY_CHECK_ARG_TYPE(1, tp_array2d); + if(!py_checkinstance(&argv[1], tp_array2d_like)) return false; PY_CHECK_ARG_TYPE(2, tp_int); - c11_array2d* self = py_touserdata(argv); - c11_array2d* kernel = py_touserdata(py_arg(1)); + 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"); } + 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; - if(!_array2d_check_all_type(self, tp_int)) return false; - if(!_array2d_check_all_type(kernel, tp_int)) return false; 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++) { @@ -656,13 +643,17 @@ static bool array2d_convolve(int argc, py_Ref argv) { if(x < 0 || x >= self->n_cols || y < 0 || y >= self->n_rows) { _0 = padding; } else { - _0 = py_toint(py_array2d__get(self, x, y)); + py_Ref item = self->f_get(self, x, y); + if(!py_checkint(item)) return false; + _0 = py_toint(item); } - _1 = py_toint(py_array2d__get(kernel, ii, jj)); + py_Ref kitem = kernel->f_get(kernel, ii, jj); + if(!py_checkint(kitem)) return false; + _1 = py_toint(kitem); sum += _0 * _1; } } - py_newint(py_array2d__get(res, i, j), sum); + py_newint(c11_array2d__get(res, i, j), sum); } } py_assign(py_retval(), py_peek(-1)); @@ -670,69 +661,519 @@ static bool array2d_convolve(int argc, py_Ref argv) { return true; } -void pk__add_module_array2d() { - py_GlobalRef mod = py_newmodule("array2d"); - py_Type array2d = pk_newtype("array2d", tp_object, mod, NULL, false, true); - py_Type array2d_iterator = pk_newtype("array2d_iterator", tp_object, mod, NULL, false, true); - assert(array2d == tp_array2d); - assert(array2d_iterator == tp_array2d_iterator); +#undef HANDLE_SLICE - py_setdict(mod, py_name("array2d"), py_tpobject(array2d)); +static void register_array2d_like(py_Ref mod) { + py_Type type = py_newtype("array2d_like", tp_object, mod, NULL); + assert(type == tp_array2d_like); - // array2d is unhashable - py_setdict(py_tpobject(array2d), __hash__, py_None()); + 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_bindmagic(array2d_iterator, __iter__, pk_wrapper__self); - py_bindmagic(array2d_iterator, __next__, array2d_iterator__next__); - py_bind(py_tpobject(array2d), - "__new__(cls, n_cols: int, n_rows: int, default=None)", - array2d__new__); + py_bindmethod(type, "is_valid", array2d_like_is_valid); + py_bindmethod(type, "get", array2d_like_get); - py_bindmagic(array2d, __eq__, array2d__eq__); - py_bindmagic(array2d, __ne__, array2d__ne__); - py_bindmagic(array2d, __repr__, array2d__repr__); - py_bindmagic(array2d, __iter__, array2d__iter__); + py_bindmethod(type, "render", array2d_like_render); - py_bindmagic(array2d, __getitem__, array2d__getitem__); - py_bindmagic(array2d, __setitem__, array2d__setitem__); + py_bindmethod(type, "all", array2d_like_all); + py_bindmethod(type, "any", array2d_like_any); - py_bindproperty(array2d, "n_cols", array2d_n_cols, NULL); - py_bindproperty(array2d, "n_rows", array2d_n_rows, NULL); - py_bindproperty(array2d, "width", array2d_n_cols, NULL); - py_bindproperty(array2d, "height", array2d_n_rows, NULL); - py_bindproperty(array2d, "numel", array2d_numel, NULL); + 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_bindmethod(array2d, "is_valid", array2d_is_valid); - py_bindmethod(array2d, "get", array2d_get); + 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_bindmethod(array2d, "map", array2d_map); - py_bindmethod(array2d, "copy", array2d_copy); + py_bindmagic(type, __getitem__, array2d_like__getitem__); + py_bindmagic(type, __setitem__, array2d_like__setitem__); - py_bindmethod(array2d, "fill_", array2d_fill_); - py_bindmethod(array2d, "apply_", array2d_apply_); - py_bindmethod(array2d, "copy_", array2d_copy_); - - py_bindmethod(array2d, "render", array2d_render); - - py_bindmethod(array2d, "all", array2d_all); - py_bindmethod(array2d, "any", array2d_any); - - py_bindstaticmethod(array2d, "fromlist", array2d_fromlist_STATIC); - py_bindmethod(array2d, "tolist", array2d_tolist); - - py_bindmethod(array2d, "count", array2d_count); - py_bindmethod(array2d, "get_bounding_rect", array2d_get_bounding_rect); - py_bindmethod(array2d, "count_neighbors", array2d_count_neighbors); - py_bindmethod(array2d, "convolve", array2d_convolve); + 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.get_connected_components = get_connected_components\ndel get_connected_components\n"; - + "\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"); } } -#undef INC_COUNT -#undef HANDLE_SLICE +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); + 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); +} \ No newline at end of file diff --git a/src/modules/pickle.c b/src/modules/pickle.c index c19fabfa..b5afae1e 100644 --- a/src/modules/pickle.c +++ b/src/modules/pickle.c @@ -4,6 +4,7 @@ #include "pocketpy/common/utils.h" #include "pocketpy/common/sstream.h" +#include "pocketpy/interpreter/vm.h" #include "pocketpy/interpreter/array2d.h" #include @@ -324,16 +325,16 @@ static bool pkl__write_object(PickleObject* buf, py_TValue* obj) { return true; else { c11_array2d* arr = py_touserdata(obj); - for(int i = 0; i < arr->numel; i++) { + for(int i = 0; i < arr->header.numel; i++) { if(arr->data[i].is_ptr) return TypeError( "'array2d' object is not picklable because it contains heap-allocated objects"); buf->used_types[arr->data[i].type] = true; } pkl__emit_op(buf, PKL_ARRAY2D); - pkl__emit_int(buf, arr->n_cols); - pkl__emit_int(buf, arr->n_rows); - PickleObject__write_bytes(buf, arr->data, arr->numel * sizeof(py_TValue)); + pkl__emit_int(buf, arr->header.n_cols); + pkl__emit_int(buf, arr->header.n_rows); + PickleObject__write_bytes(buf, arr->data, arr->header.numel * sizeof(py_TValue)); } pkl__store_memo(buf, obj->_obj); return true; @@ -651,9 +652,9 @@ bool py_pickle_loads_body(const unsigned char* p, int memo_length, c11_smallmap_ int n_cols = pkl__read_int(&p); int n_rows = pkl__read_int(&p); c11_array2d* arr = py_newarray2d(py_pushtmp(), n_cols, n_rows); - int total_size = arr->numel * sizeof(py_TValue); + int total_size = arr->header.numel * sizeof(py_TValue); memcpy(arr->data, p, total_size); - for(int i = 0; i < arr->numel; i++) { + for(int i = 0; i < arr->header.numel; i++) { arr->data[i].type = pkl__fix_type(arr->data[i].type, type_mapping); } p += total_size; diff --git a/src/public/cast.c b/src/public/cast.c index 6cda6c66..ea5169ca 100644 --- a/src/public/cast.c +++ b/src/public/cast.c @@ -61,6 +61,11 @@ bool py_checktype(py_Ref self, py_Type type) { return TypeError("expected '%t', got '%t'", type, self->type); } +bool py_checkinstance(py_Ref self, py_Type type) { + if(py_isinstance(self, type)) return true; + return TypeError("expected '%t' or its subclass, got '%t'", type, self->type); +} + bool py_isinstance(py_Ref obj, py_Type type) { return py_issubclass(obj->type, type); } bool py_issubclass(py_Type derived, py_Type base) { diff --git a/tests/90_array2d.py b/tests/90_array2d.py index 4b5cbc71..a762602a 100644 --- a/tests/90_array2d.py +++ b/tests/90_array2d.py @@ -1,10 +1,13 @@ from array2d import array2d from linalg import vec2i +def exit_on_error(): + raise KeyboardInterrupt + # test error args for __init__ try: a = array2d(0, 0) - exit(0) + exit_on_error() except ValueError: pass @@ -39,7 +42,7 @@ assert a[0, 0] == (0, 0) assert a[1, 3] == (1, 3) try: a[2, 0] - exit(1) + exit_on_error() except IndexError: pass @@ -51,7 +54,7 @@ a[1, 3] = 6 assert a[1, 3] == 6 try: a[0, -1] = 7 - exit(1) + exit_on_error() except IndexError: pass @@ -83,21 +86,19 @@ d = c.copy() assert (d == c).all() and d is not c # test fill_ -d.fill_(-3) +d[:, :] = -3 # d.fill_(-3) assert (d == array2d(2, 4, default=-3)).all() -# test apply_ -d.apply_(lambda x: x + 3) +# test apply +d.apply(lambda x: x + 3) assert (d == array2d(2, 4, default=0)).all() # test copy_ -a.copy_(d) +a[:, :] = d assert (a == d).all() and a is not d x = array2d(2, 4, default=0) -x.copy_(d) +x[:, :] = d assert (x == d).all() and x is not d -x.copy_([1, 2, 3, 4, 5, 6, 7, 8]) -assert x.tolist() == [[1, 2], [3, 4], [5, 6], [7, 8]] # test alive_neighbors a = array2d[int](3, 3, default=0) @@ -148,7 +149,7 @@ assert a.get_bounding_rect(0) == (0, 0, 5, 5) try: a.get_bounding_rect(2) - exit(1) + exit_on_error() except ValueError: pass @@ -165,16 +166,10 @@ assert a == array2d(3, 2, default=3) try: a[:, :] = array2d(1, 1) - exit(1) + exit_on_error() except ValueError: pass -try: - a[:, :] = ... - exit(1) -except TypeError: - pass - # test __iter__ a = array2d(3, 4, default=1) for xy, val in a: diff --git a/tests/90_chunked_array2d.py b/tests/90_chunked_array2d.py new file mode 100644 index 00000000..257a7ffe --- /dev/null +++ b/tests/90_chunked_array2d.py @@ -0,0 +1,8 @@ +from array2d import chunked_array2d +from linalg import vec2i + +a = chunked_array2d(4, default=0) + +a[vec2i.ONE] = 1 + +print(a.view().render())