mirror of
https://github.com/pocketpy/pocketpy
synced 2025-10-20 03:20:18 +00:00
improve array2d_like
This commit is contained in:
parent
d2ebadb93f
commit
eb414df4cd
@ -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;
|
||||
|
@ -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)
|
||||
|
@ -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: ...
|
||||
|
||||
|
@ -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);
|
||||
|
@ -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) {
|
||||
|
@ -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
|
||||
|
@ -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']
|
||||
|
Loading…
x
Reference in New Issue
Block a user