bak

include/pocketpy/interpreter/array2d.h

@@ -19,3 +19,34 @@ typedef struct c11_array2d_iterator {
 } c11_array2d_iterator;
 
 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;
+
+    struct {
+        py_Ref missing_chunk;
+    } callbacks;
+
+    c11_chunked_array2d_chunks_KV last_visited;
+} 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);
+
+/* array2d_view */

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/typings/array2d.pyi

@@ -84,3 +84,32 @@ 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]:
+    mask: array2d[bool] | None
+
+    @property
+    def n_cols(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 get[R](self, col: int, row: int, default: R = None) -> T | R: ...
+
+    def view(self, pos: vec2i, width: int, height: int) -> array2d_view[T]:
+        """Return a view of the grid in the given rectangle."""

src/modules/array2d.c

@@ -736,3 +736,105 @@ void pk__add_module_array2d() {
 
 #undef INC_COUNT
 #undef HANDLE_SLICE
+
+/* chunked_array2d */
+#define SMALLMAP_T__SOURCE
+#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__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
+    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,
+                                                     int col,
+                                                     int row,
+                                                     c11_vec2i* chunk_pos,
+                                                     c11_vec2i* local_pos) {
+    if(col >= 0) {
+        chunk_pos->x = col >> self->chunk_size_log2;
+        local_pos->x = col & self->chunk_size_mask;
+    } else {
+        chunk_pos->x = -((-col) >> self->chunk_size_log2);
+        local_pos->x = (-col) & self->chunk_size_mask;
+    }
+    if(row >= 0) {
+        chunk_pos->y = row >> self->chunk_size_log2;
+        local_pos->y = row & self->chunk_size_mask;
+    } else {
+        chunk_pos->y = -((-row) >> self->chunk_size_log2);
+        local_pos->y = (-row) & self->chunk_size_mask;
+    }
+    py_TValue* data;
+    if(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);
+    }
+    if(data != NULL) {
+        self->last_visited.key = *chunk_pos;
+        self->last_visited.value = data;
+    }
+    return data;
+}
+
+void c11_chunked_array2d__ctor(c11_chunked_array2d* self, int chunk_size) {
+    c11_chunked_array2d_chunks__ctor(&self->chunks);
+    self->chunk_size = chunk_size;
+    switch(chunk_size) {
+        case 2: self->chunk_size_log2 = 1; break;
+        case 4: self->chunk_size_log2 = 2; break;
+        case 8: self->chunk_size_log2 = 3; break;
+        case 16: self->chunk_size_log2 = 4; break;
+        case 32: self->chunk_size_log2 = 5; break;
+        case 64: self->chunk_size_log2 = 6; break;
+        case 128: self->chunk_size_log2 = 7; break;
+        case 256: self->chunk_size_log2 = 8; break;
+        case 512: self->chunk_size_log2 = 9; break;
+        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);
+    }
+    self->chunk_size_mask = chunk_size - 1;
+    c11_chunked_array2d__new_chunk(self,
+                                   (c11_vec2i){
+                                       {0, 0}
+    });
+}
+
+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_chunked_array2d_chunks__dtor(&self->chunks);
+}
+
+py_Ref c11_chunked_array2d__get(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) return NULL;
+    py_Ref retval = &data[local_pos.y * self->chunk_size + local_pos.x];
+    if(py_isnil(retval)) return NULL;
+    return retval;
+}
+
+void 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);
+    data[local_pos.y * self->chunk_size + local_pos.x] = *value;
+}