From a1797869dbe3637bd826524fcbecf3cdf38c9e17 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Wed, 7 Feb 2024 00:18:09 +0800
Subject: [PATCH] some update

---
 include/typings/array2d.pyi | 18 +++++++++++++++
 src/array2d.cpp             | 46 ++++++++++++++++++++++++++++++-------
 tests/80_array2d.py         |  5 ++++
 3 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/include/typings/array2d.pyi b/include/typings/array2d.pyi
index 95ef98bc..4cfc51b7 100644
--- a/include/typings/array2d.pyi
+++ b/include/typings/array2d.pyi
@@ -89,3 +89,21 @@ class array2d(Generic[T]):
         self.n_cols = other.n_cols
         self.n_rows = other.n_rows
         self.data = other.data.copy()
+
+    # for cellular automata
+    def count_neighbors(self, value) -> array2d[int]:
+        new_a: array2d[int] = array2d(self.n_cols, self.n_rows)
+        for j in range(self.n_rows):
+            for i in range(self.n_cols):
+                count = 0
+                count += int(self.is_valid(i-1, j-1) and self[i-1, j-1] == value)
+                count += int(self.is_valid(i, j-1) and self[i, j-1] == value)
+                count += int(self.is_valid(i+1, j-1) and self[i+1, j-1] == value)
+                count += int(self.is_valid(i-1, j) and self[i-1, j] == value)
+                count += int(self.is_valid(i+1, j) and self[i+1, j] == value)
+                count += int(self.is_valid(i-1, j+1) and self[i-1, j+1] == value)
+                count += int(self.is_valid(i, j+1) and self[i, j+1] == value)
+                count += int(self.is_valid(i+1, j+1) and self[i+1, j+1] == value)
+                new_a[i, j] = count
+        return new_a
+
diff --git a/src/array2d.cpp b/src/array2d.cpp
index a07206b6..a91c5dab 100644
--- a/src/array2d.cpp
+++ b/src/array2d.cpp
@@ -31,6 +31,14 @@ struct Array2d{
         return 0 <= col && col < n_cols && 0 <= row && row < n_rows;
     }
 
+    PyObject* _get(int col, int row){
+        return data[row * n_cols + col];
+    }
+
+    void _set(int col, int row, PyObject* value){
+        data[row * n_cols + col] = value;
+    }
+
     static void _register(VM* vm, PyObject* mod, PyObject* type){
         vm->bind(type, "__new__(cls, *args, **kwargs)", [](VM* vm, ArgsView args){
             Type cls = PK_OBJ_GET(Type, args[0]);
@@ -71,7 +79,7 @@ struct Array2d{
             int col = CAST(int, args[1]);
             int row = CAST(int, args[2]);
             if(!self.is_valid(col, row)) return args[3];
-            return self.data[row * self.n_cols + col];
+            return self._get(col, row);
         });
 
         vm->bind__getitem__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0, PyObject* _1){
@@ -82,7 +90,7 @@ struct Array2d{
             if(!self.is_valid(col, row)){
                 vm->IndexError(_S('(', col, ", ", row, ')', " is not a valid index for array2d(", self.n_cols, ", ", self.n_rows, ')'));
             }
-            return self.data[row * self.n_cols + col];
+            return self._get(col, row);
         });
 
         vm->bind__setitem__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0, PyObject* _1, PyObject* _2){
@@ -93,18 +101,16 @@ struct Array2d{
             if(!self.is_valid(col, row)){
                 vm->IndexError(_S('(', col, ", ", row, ')', " is not a valid index for array2d(", self.n_cols, ", ", self.n_rows, ')'));
             }
-            self.data[row * self.n_cols + col] = _2;
+            self._set(col, row, _2);
         });
 
         vm->bind__iter__(PK_OBJ_GET(Type, type), [](VM* vm, PyObject* _0){
             Array2d& self = PK_OBJ_GET(Array2d, _0);
             List t(self.n_rows);
             List row(self.n_cols);
-            for(int i = 0; i < self.n_rows; i++){
-                for(int j = 0; j < self.n_cols; j++){
-                    row[j] = self.data[i * self.n_cols + j];
-                }
-                t[i] = VAR(row);    // copy
+            for(int j = 0; j < self.n_rows; j++){
+                for(int i = 0; i < self.n_cols; i++) row[i] = self._get(i, j);
+                t[j] = VAR(row);    // copy
             }
             return vm->py_iter(VAR(std::move(t)));
         });
@@ -180,6 +186,30 @@ struct Array2d{
             }
             return vm->True;
         });
+
+        // for cellular automata
+        vm->bind(type, "count_neighbors(self, value) -> array2d[int]", [](VM* vm, ArgsView args){
+            Array2d& self = PK_OBJ_GET(Array2d, args[0]);
+            PyObject* new_array_obj = vm->heap.gcnew<Array2d>(Array2d::_type(vm));
+            Array2d& new_array = PK_OBJ_GET(Array2d, new_array_obj);
+            new_array.init(self.n_cols, self.n_rows);
+            PyObject* value = args[1];
+            for(int j = 0; j < new_array.n_rows; j++){
+                for(int i = 0; i < new_array.n_cols; i++){
+                    int count = 0;
+                    count += self.is_valid(i-1, j-1) && vm->py_eq(self._get(i-1, j-1), value);
+                    count += self.is_valid(i, j-1) && vm->py_eq(self._get(i, j-1), value);
+                    count += self.is_valid(i+1, j-1) && vm->py_eq(self._get(i+1, j-1), value);
+                    count += self.is_valid(i-1, j) && vm->py_eq(self._get(i-1, j), value);
+                    count += self.is_valid(i+1, j) && vm->py_eq(self._get(i+1, j), value);
+                    count += self.is_valid(i-1, j+1) && vm->py_eq(self._get(i-1, j+1), value);
+                    count += self.is_valid(i, j+1) && vm->py_eq(self._get(i, j+1), value);
+                    count += self.is_valid(i+1, j+1) && vm->py_eq(self._get(i+1, j+1), value);
+                    new_array._set(i, j, VAR(count));
+                }
+            }
+            return new_array_obj; 
+        });
     }
 
     void _gc_mark() const{
diff --git a/tests/80_array2d.py b/tests/80_array2d.py
index 946058fd..d526f441 100644
--- a/tests/80_array2d.py
+++ b/tests/80_array2d.py
@@ -99,3 +99,8 @@ assert A().width == 2
 assert A().height == 4
 assert A().numel == 8
 assert A().get(0, 0, default=2) == 0
+
+# test alive_neighbors
+a = array2d(3, 3, default=0)
+a.count_neighbors(0) == a
+