improve `array2d_like`

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;

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)

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: ...
 

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);

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) {

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

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']