implement `array2d.chunked_array2d[T, TContext]` (#332)

* bak

* backup

* ...

* Update array2d.pyi

* backup

* backup

* backup

* backup

* backup

* backup

* backup

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

include/pocketpy/linalg.h

@@ -1,8 +1,11 @@
 #pragma once
 
+#include <stdint.h>
+
 typedef union c11_vec2i {
     struct { int x, y; };
     int data[2];
+    int64_t _i64;
 } c11_vec2i;
 
 typedef union c11_vec3i {

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

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,54 +13,52 @@ 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."""
+        """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_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: tuple[int, int]) -> T: ...
     @overload
     def __getitem__(self, index: vec2i) -> T: ...
     @overload
-    def __getitem__(self, index: tuple[slice, slice]) -> array2d[T]: ...
+    def __getitem__(self, index: tuple[int, int]) -> T: ...
     @overload
-    def __getitem__(self, mask: array2d[bool]) -> list[T]: ...
+    def __getitem__(self, index: tuple[slice, slice]) -> array2d_view[T]: ...
     @overload
-    def __setitem__(self, index: tuple[int, int], value: T): ...
+    def __getitem__(self, mask: array2d_like[bool]) -> list[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]'): ...
+    def __setitem__(self, index: tuple[int, int], value: 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: ...
-
-    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 tolist(self) -> list[list[T]]: ...
+    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:
@@ -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]: ...

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 <stdint.h>
 
@@ -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;

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

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:

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