Update array2d.pyi

include/pocketpy/interpreter/array2d.h

@@ -35,18 +35,18 @@ typedef struct c11_chunked_array2d {
     int chunk_size;
     int chunk_size_log2;
     int chunk_size_mask;
-
-    struct {
-        py_Ref missing_chunk;
-    } callbacks;
-
     c11_chunked_array2d_chunks_KV last_visited;
+
+    py_TValue default_T;
+    py_TValue context_builder;
 } c11_chunked_array2d;
 
-void c11_chunked_array2d__ctor(c11_chunked_array2d* self, int chunk_size);
 void c11_chunked_array2d__dtor(c11_chunked_array2d* self);
 
 py_Ref c11_chunked_array2d__get(c11_chunked_array2d* self, int col, int row);
-void c11_chunked_array2d__set(c11_chunked_array2d* self, int col, int row, py_Ref value);
+bool c11_chunked_array2d__set(c11_chunked_array2d* self, int col, int row, py_Ref value) PY_RAISE;
+void c11_chunked_array2d__del(c11_chunked_array2d* self, int col, int row);
+
+void pk__register_chunked_array2d(py_Ref mod);
 
 /* array2d_view */

include/typings/array2d.pyi

@@ -3,7 +3,7 @@ from linalg import vec2i
 
 Neighborhood = Literal['Moore', 'von Neumann']
 
-class array2d[T]:
+class array2d_like[T]:
     @property
     def n_cols(self) -> int: ...
     @property
@@ -15,33 +15,39 @@ class array2d[T]:
     @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: ...
 
+    @overload
+    def __getitem__(self, index: vec2i) -> T: ...
+    @overload
+    def __getitem__(self, index: tuple[int, int]) -> T: ...
+    @overload
+    def __setitem__(self, index: vec2i, value: T): ...
+    @overload
+    def __setitem__(self, index: tuple[int, int], value: T): ...
+
+
+class array2d_view[T](array2d_like[T]):
+    origin: vec2i
+
+
+class array2d[T](array2d_like[T]):
+    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 __iter__(self) -> Iterator[tuple[vec2i, T]]: ...
+
     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): ...
@@ -86,30 +92,30 @@ class array2d[T]:
         """
 
 
-class array2d_view[T]:
+class chunked_array2d[T, TContext]:
-    mask: array2d[bool] | None
+    def __init__(
+            self,
+            chunk_size: int,
+            default: T = None,
+            context_builder: Callable[[vec2i], TContext] | None = None,
+            ): ...
     
     @property
-    def n_cols(self) -> int: ...
+    def chunk_size(self) -> int: ...
-    @property
-    def n_rows(self) -> int: ...
-    @property
-    def width(self) -> int: ...
-    @property
-    def height(self) -> int: ...
-
-    def __getitem__(self, index: vec2i) -> T: ...
-    def __setitem__(self, index: vec2i, value: T): ...
-
-
-class chunked_array2d[T]:
-    def __init__(self, chunk_size: 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 get[R](self, col: int, row: int, default: R = None) -> T | R: ...
+    def clear(self) -> None: ...
 
-    def view(self, pos: vec2i, width: int, height: int) -> array2d_view[T]:
+    def world_to_chunk(self, world_pos: vec2i) -> tuple[vec2i, vec2i]: ...
-        """Return a view of the grid in the given rectangle."""
+    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/array2d.c

@@ -732,6 +732,8 @@ void pk__add_module_array2d() {
         py_printexc();
         c11__abort("failed to execute array2d.py");
     }
+
+    pk__register_chunked_array2d(mod);
 }
 
 #undef INC_COUNT
@@ -748,24 +750,32 @@ void pk__add_module_array2d() {
 #undef SMALLMAP_T__SOURCE
 
 static py_TValue* c11_chunked_array2d__new_chunk(c11_chunked_array2d* self, c11_vec2i pos) {
-    int chunk_numel = self->chunk_size * self->chunk_size;
-    py_TValue* data = PK_MALLOC(sizeof(py_TValue) * chunk_numel);
-    memset(data, 0, sizeof(py_TValue) * chunk_numel);
 #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 py_TValue* c11_chunked_array2d__parse_col_row(c11_chunked_array2d* self,
+void c11_chunked_array2d__world_to_chunk(c11_chunked_array2d* self,
-                                                     int col,
+                                         int col,
-                                                     int row,
+                                         int row,
-                                                     c11_vec2i* chunk_pos,
+                                         c11_vec2i* chunk_pos,
-                                                     c11_vec2i* local_pos) {
+                                         c11_vec2i* local_pos) {
     if(col >= 0) {
         chunk_pos->x = col >> self->chunk_size_log2;
         local_pos->x = col & self->chunk_size_mask;
@@ -780,8 +790,16 @@ static py_TValue* c11_chunked_array2d__parse_col_row(c11_chunked_array2d* self,
         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(chunk_pos->_i64 == self->last_visited.key._i64) {
+    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);
@@ -790,10 +808,22 @@ static py_TValue* c11_chunked_array2d__parse_col_row(c11_chunked_array2d* self,
         self->last_visited.key = *chunk_pos;
         self->last_visited.value = data;
     }
-    return data;
+    return data + 1;  // skip context
 }
 
-void c11_chunked_array2d__ctor(c11_chunked_array2d* self, int chunk_size) {
+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) {
@@ -809,17 +839,125 @@ void c11_chunked_array2d__ctor(c11_chunked_array2d* self, int chunk_size) {
         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: c11__abort("invalid chunk_size: %d", chunk_size);
+        default: return ValueError("invalid chunk_size: %d", chunk_size);
     }
     self->chunk_size_mask = chunk_size - 1;
-    c11_chunked_array2d__new_chunk(self,
+    memset(&self->last_visited, 0, sizeof(c11_chunked_array2d_chunks_KV));
-                                   (c11_vec2i){
+    py_newnone(py_retval());
-                                       {0, 0}
+    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);
+    if(res != NULL) {
+        py_assign(py_retval(), res);
+    } else {
+        py_assign(py_retval(), &self->default_T);
+    }
+    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__foreach(c11_chunked_array2d_chunks_KV, &self->chunks, p_kv) PK_FREE(p_kv->value);
     c11_chunked_array2d_chunks__dtor(&self->chunks);
 }
 
@@ -832,9 +970,55 @@ py_Ref c11_chunked_array2d__get(c11_chunked_array2d* self, int col, int row) {
     return retval;
 }
 
-void c11_chunked_array2d__set(c11_chunked_array2d* self, int col, int row, py_Ref value) {
+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) {
+        data = c11_chunked_array2d__new_chunk(self, chunk_pos);
+        if(data == NULL) return false;
+    }
     data[local_pos.y * self->chunk_size + local_pos.x] = *value;
+    return true;
+}
+
+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[local_pos.y * self->chunk_size + local_pos.x] = *py_NIL();
+}
+
+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);
+        }
+    }
+}
+
+void pk__register_chunked_array2d(py_Ref mod) {
+    py_Type cls = py_newtype("chunked_array2d", tp_object, mod, (py_Dtor)c11_chunked_array2d__dtor);
+    pk__tp_set_marker(cls, c11_chunked_array2d__mark);
+
+    py_bindmagic(cls, __new__, chunked_array2d__new__);
+    py_bind(py_tpobject(cls),
+            "__init__(self, chunk_size, default=None, context_builder=None)",
+            chunked_array2d__init__);
+
+    py_bindproperty(cls, "chunk_size", chunked_array2d__chunk_size, NULL);
+
+    py_bindmagic(cls, __getitem__, chunked_array2d__getitem__);
+    py_bindmagic(cls, __setitem__, chunked_array2d__setitem__);
+    py_bindmagic(cls, __delitem__, chunked_array2d__delitem__);
+    py_bindmagic(cls, __iter__, chunked_array2d__iter__);
+
+    py_bindmethod(cls, "clear", chunked_array2d__clear);
+    py_bindmethod(cls, "world_to_chunk", chunked_array2d__world_to_chunk);
+    py_bindmethod(cls, "add_chunk", chunked_array2d__add_chunk);
+    py_bindmethod(cls, "remove_chunk", chunked_array2d__remove_chunk);
+    py_bindmethod(cls, "get_context", chunked_array2d__get_context);
 }

src/public/modules.c

@@ -165,7 +165,7 @@ bool py_importlib_reload(py_GlobalRef module) {
     c11_sv path = py_tosv(py_getdict(module, __path__));
     c11_string* slashed_path = c11_sv__replace(path, '.', PK_PLATFORM_SEP);
     c11_string* filename = c11_string__new3("%s.py", slashed_path->data);
-    const char* data = vm->callbacks.importfile(filename->data);
+    char* data = vm->callbacks.importfile(filename->data);
     if(data == NULL) {
         c11_string__delete(filename);
         filename = c11_string__new3("%s%c__init__.py", slashed_path->data, PK_PLATFORM_SEP);
@@ -175,6 +175,7 @@ bool py_importlib_reload(py_GlobalRef module) {
     if(data == NULL) return ImportError("module '%v' not found", path);
     bool ok = py_exec(data, filename->data, EXEC_MODE, module);
     c11_string__delete(filename);
+    PK_FREE(data);
     py_assign(py_retval(), module);
     return ok;
 }