diff --git a/include/pocketpy/common/utils.h b/include/pocketpy/common/utils.h index 8610e0f3..cd842498 100644 --- a/include/pocketpy/common/utils.h +++ b/include/pocketpy/common/utils.h @@ -28,11 +28,6 @@ #define c11__count_array(a) (sizeof(a) / sizeof(a[0])) -// NARGS -#define PK_NARGS_SEQ(_1, _2, _3, _4, N, ...) N -#define PK_NARGS(...) PK_NARGS_SEQ(__VA_ARGS__, 4, 3, 2, 1, 0) -#define PK_NPTRS(...) PK_NARGS_SEQ(__VA_ARGS__, int****, int***, int**, int*, int) - // ref counting typedef struct RefCounted { int count; diff --git a/include/pocketpy/interpreter/vm.h b/include/pocketpy/interpreter/vm.h index 62ac4f40..0080b8fe 100644 --- a/include/pocketpy/interpreter/vm.h +++ b/include/pocketpy/interpreter/vm.h @@ -52,7 +52,7 @@ void VM__dtor(VM* self); void VM__push_frame(VM* self, Frame* frame); void VM__pop_frame(VM* self); -bool pk__parse_int_slice(py_Ref slice, int length, int* start, int* stop, int* step); +bool pk__parse_int_slice(py_Ref slice, int length, int* restrict start, int* restrict stop, int* restrict step); bool pk__normalize_index(int* index, int length); #define pk__mark_value(val) if((val)->is_ptr && !(val)->_obj->gc_marked) PyObject__mark((val)->_obj) diff --git a/include/typings/array2d.pyi b/include/typings/array2d.pyi index 79debad9..1ed417c7 100644 --- a/include/typings/array2d.pyi +++ b/include/typings/array2d.pyi @@ -39,8 +39,26 @@ class array2d_like[T]: 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 __le__(self, other: T | array2d_like[T]) -> array2d[bool]: ... + def __lt__(self, other: T | array2d_like[T]) -> array2d[bool]: ... + def __ge__(self, other: T | array2d_like[T]) -> array2d[bool]: ... + def __gt__(self, other: T | array2d_like[T]) -> array2d[bool]: ... + def __eq__(self, other: T | array2d_like[T]) -> array2d[bool]: ... # type: ignore + def __ne__(self, other: T | array2d_like[T]) -> array2d[bool]: ... # type: ignore + + def __add__(self, other: T | array2d_like[T]) -> array2d: ... + def __sub__(self, other: T | array2d_like[T]) -> array2d: ... + def __mul__(self, other: T | array2d_like[T]) -> array2d: ... + def __truediv__(self, other: T | array2d_like[T]) -> array2d: ... + def __floordiv__(self, other: T | array2d_like[T]) -> array2d: ... + def __mod__(self, other: T | array2d_like[T]) -> array2d: ... + def __pow__(self, other: T | array2d_like[T]) -> array2d: ... + + def __and__(self, other: T | array2d_like[T]) -> array2d: ... + def __or__(self, other: T | array2d_like[T]) -> array2d: ... + def __xor__(self, other: T | array2d_like[T]) -> array2d: ... + def __invert__(self) -> array2d: ... + def __iter__(self) -> Iterator[tuple[vec2i, T]]: ... def __repr__(self) -> str: ... diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c index 88fc7794..5ef35650 100644 --- a/src/interpreter/vm.c +++ b/src/interpreter/vm.c @@ -281,7 +281,7 @@ static void _clip_int(int* value, int min, int max) { if(*value > max) *value = max; } -bool pk__parse_int_slice(py_Ref slice, int length, int* start, int* stop, int* step) { +bool pk__parse_int_slice(py_Ref slice, int length, int* restrict start, int* restrict stop, int* restrict step) { if(py_isint(slice)) { int index = py_toint(slice); bool ok = pk__normalize_index(&index, length); diff --git a/src/modules/array2d.c b/src/modules/array2d.c index 4eb9247e..756ff4f6 100644 --- a/src/modules/array2d.c +++ b/src/modules/array2d.c @@ -200,6 +200,35 @@ static bool _array2d_like_check_same_shape(c11_array2d_like* self, c11_array2d_l return _check_same_shape(self->n_cols, self->n_rows, other->n_cols, other->n_rows); } +static bool _array2d_like_broadcasted_zip_with(int argc, py_Ref argv, py_Name op, py_Name rop) { + PY_CHECK_ARGC(2); + c11_array2d_like* self = py_touserdata(argv); + c11_array2d_like* other; + if(py_isinstance(py_arg(1), tp_array2d_like)) { + other = py_touserdata(py_arg(1)); + if(!_array2d_like_check_same_shape(self, other)) return false; + } else { + other = NULL; + } + 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 lhs = self->f_get(self, i, j); + py_Ref rhs; + if(other != NULL) { + rhs = other->f_get(other, i, j); + } else { + rhs = py_arg(1); // broadcast + } + if(!py_binaryop(lhs, rhs, op, rop)) return false; + c11_array2d__set(res, i, j, py_retval()); + } + } + py_assign(py_retval(), py_peek(-1)); + py_pop(); + return true; +} + static bool array2d_like_zip_with(int argc, py_Ref argv) { PY_CHECK_ARGC(3); c11_array2d_like* self = py_touserdata(argv); @@ -223,6 +252,48 @@ static bool array2d_like_zip_with(int argc, py_Ref argv) { return true; } +#define DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(name, op, rop) \ + static bool array2d_like##name(int argc, py_Ref argv) { \ + return _array2d_like_broadcasted_zip_with(argc, argv, op, rop); \ + } + +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__le__, __le__, __ge__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__lt__, __lt__, __gt__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__ge__, __ge__, __le__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__gt__, __gt__, __lt__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__eq__, __eq__, __eq__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__ne__, __ne__, __ne__) + +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__add__, __add__, __radd__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__sub__, __sub__, __rsub__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__mul__, __mul__, __rmul__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__truediv__, __truediv__, __rtruediv__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__floordiv__, __floordiv__, __rfloordiv__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__mod__, __mod__, __rmod__) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__pow__, __pow__, __rpow__) + +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__and__, __and__, 0) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__or__, __or__, 0) +DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH(__xor__, __xor__, 0) + +#undef DEF_ARRAY2D_LIKE__MAGIC_ZIP_WITH + +static bool array2d_like__invert__(int argc, py_Ref argv) { + PY_CHECK_ARGC(1); + c11_array2d_like* self = py_touserdata(argv); + 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(!pk_callmagic(__invert__, 1, item)) return false; + c11_array2d__set(res, i, j, py_retval()); + } + } + py_assign(py_retval(), py_peek(-1)); + py_pop(); + return true; +} + static bool array2d_like_copy(int argc, py_Ref argv) { // def copy(self) -> 'array2d': ... PY_CHECK_ARGC(1); @@ -252,51 +323,6 @@ static bool array2d_like_tolist(int argc, py_Ref argv) { return true; } -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); @@ -717,8 +743,26 @@ static void register_array2d_like(py_Ref mod) { py_bindmethod(type, "copy", array2d_like_copy); py_bindmethod(type, "tolist", array2d_like_tolist); + py_bindmagic(type, __le__, array2d_like__le__); + py_bindmagic(type, __lt__, array2d_like__lt__); + py_bindmagic(type, __ge__, array2d_like__ge__); + py_bindmagic(type, __gt__, array2d_like__gt__); py_bindmagic(type, __eq__, array2d_like__eq__); py_bindmagic(type, __ne__, array2d_like__ne__); + + py_bindmagic(type, __add__, array2d_like__add__); + py_bindmagic(type, __sub__, array2d_like__sub__); + py_bindmagic(type, __mul__, array2d_like__mul__); + py_bindmagic(type, __truediv__, array2d_like__truediv__); + py_bindmagic(type, __floordiv__, array2d_like__floordiv__); + py_bindmagic(type, __mod__, array2d_like__mod__); + py_bindmagic(type, __pow__, array2d_like__pow__); + + py_bindmagic(type, __and__, array2d_like__and__); + py_bindmagic(type, __or__, array2d_like__or__); + py_bindmagic(type, __xor__, array2d_like__xor__); + py_bindmagic(type, __invert__, array2d_like__invert__); + py_bindmagic(type, __iter__, array2d_like__iter__); py_bindmagic(type, __repr__, array2d_like__repr__); @@ -879,7 +923,8 @@ static py_TValue* c11_chunked_array2d__new_chunk(c11_chunked_array2d* self, c11_ return data; } -static void cpy11__divmod_int_uint(int a, int b_log2, int b_mask, int* q, int* r) { +static void + cpy11__divmod_int_uint(int a, int b_log2, int b_mask, int* restrict q, int* restrict r) { if(a >= 0) { *q = a >> b_log2; *r = a & b_mask; @@ -892,8 +937,8 @@ static void cpy11__divmod_int_uint(int a, int b_log2, int b_mask, int* q, int* r static void c11_chunked_array2d__world_to_chunk(c11_chunked_array2d* self, int col, int row, - c11_vec2i* chunk_pos, - c11_vec2i* local_pos) { + c11_vec2i* restrict chunk_pos, + c11_vec2i* restrict local_pos) { cpy11__divmod_int_uint(col, self->chunk_size_log2, self->chunk_size_mask, @@ -909,8 +954,8 @@ static void c11_chunked_array2d__world_to_chunk(c11_chunked_array2d* self, 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_vec2i* restrict chunk_pos, + c11_vec2i* restrict 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) { diff --git a/src/public/py_number.c b/src/public/py_number.c index f02f50d7..c7769aa8 100644 --- a/src/public/py_number.c +++ b/src/public/py_number.c @@ -463,6 +463,13 @@ DEF_BOOL_BITWISE(__and__, &&) DEF_BOOL_BITWISE(__or__, ||) DEF_BOOL_BITWISE(__xor__, !=) +static bool bool__invert__(int argc, py_Ref argv) { + PY_CHECK_ARGC(1); + bool val = py_tobool(&argv[0]); + py_newbool(py_retval(), !val); + return true; +} + void pk_number__register() { /****** tp_int & tp_float ******/ py_bindmagic(tp_int, __add__, int__add__); @@ -539,6 +546,7 @@ void pk_number__register() { py_bindmagic(tp_bool, __and__, bool__and__); py_bindmagic(tp_bool, __or__, bool__or__); py_bindmagic(tp_bool, __xor__, bool__xor__); + py_bindmagic(tp_bool, __invert__, bool__invert__); } #undef DEF_NUM_BINARY_OP diff --git a/tests/90_array2d.py b/tests/90_array2d.py index b683bcb8..adc34cd0 100644 --- a/tests/90_array2d.py +++ b/tests/90_array2d.py @@ -247,6 +247,34 @@ b = array2d[int].fromlist([[5, 6], [7, 8]]) c = a.zip_with(b, lambda x, y: x + y) assert c.tolist() == [[6, 8], [10, 12]] +# test magic op +a = array2d[int].fromlist([[1, 2], [3, 4]]) +assert (a <= 2).tolist() == [[True, True], [False, False]] +assert (a < 2).tolist() == [[True, False], [False, False]] +assert (a >= 2).tolist() == [[False, True], [True, True]] +assert (a > 2).tolist() == [[False, False], [True, True]] +assert (a == 2).tolist() == [[False, True], [False, False]] +assert (a != 2).tolist() == [[True, False], [True, True]] +assert (a + 1).tolist() == [[2, 3], [4, 5]] +assert (a - 1).tolist() == [[0, 1], [2, 3]] +assert (a * 2).tolist() == [[2, 4], [6, 8]] +assert (a / 1).tolist() == [[1.0, 2.0], [3.0, 4.0]] +assert (a // 2).tolist() == [[0, 1], [1, 2]] +assert (a % 2).tolist() == [[1, 0], [1, 0]] +assert (a ** 2).tolist() == [[1, 4], [9, 16]] + +a = array2d[bool].fromlist([[True, False], [False, True]]) +assert (a & True).tolist() == [[True, False], [False, True]] +assert (a | True).tolist() == [[True, True], [True, True]] +assert (a ^ True).tolist() == [[False, True], [True, False]] + +b = array2d[bool].fromlist([[True, True], [False, False]]) +assert (a & b).tolist() == [[True, False], [False, False]] +assert (a | b).tolist() == [[True, True], [False, True]] +assert (a ^ b).tolist() == [[False, True], [False, True]] +assert (~a).tolist() == [[False, True], [True, False]] +assert (~b).tolist() == [[False, False], [True, True]] + # stackoverflow bug due to recursive mark-and-sweep # class Cell: # neighbors: list['Cell']