array2d.c 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754
  1. #include "pocketpy/pocketpy.h"
  2. #include "pocketpy/common/utils.h"
  3. #include "pocketpy/common/sstream.h"
  4. #include "pocketpy/interpreter/vm.h"
  5. typedef struct c11_array2d {
  6. py_TValue* data; // slots
  7. int n_cols;
  8. int n_rows;
  9. int numel;
  10. } c11_array2d;
  11. typedef struct c11_array2d_iterator {
  12. c11_array2d* array;
  13. int index;
  14. } c11_array2d_iterator;
  15. static bool py_array2d_is_valid(c11_array2d* self, int col, int row) {
  16. return col >= 0 && col < self->n_cols && row >= 0 && row < self->n_rows;
  17. }
  18. static py_ObjectRef py_array2d__get(c11_array2d* self, int col, int row) {
  19. return self->data + row * self->n_cols + col;
  20. }
  21. static void py_array2d__set(c11_array2d* self, int col, int row, py_Ref value) {
  22. self->data[row * self->n_cols + col] = *value;
  23. }
  24. static c11_array2d* py_array2d(py_OutRef out, int n_cols, int n_rows) {
  25. int numel = n_cols * n_rows;
  26. c11_array2d* ud = py_newobject(out, tp_array2d, numel, sizeof(c11_array2d));
  27. ud->data = py_getslot(out, 0);
  28. ud->n_cols = n_cols;
  29. ud->n_rows = n_rows;
  30. ud->numel = numel;
  31. return ud;
  32. }
  33. /* bindings */
  34. static bool array2d__new__(int argc, py_Ref argv) {
  35. // __new__(cls, n_cols: int, n_rows: int, default: Callable[[vec2i], T] = None)
  36. py_Ref default_ = py_arg(3);
  37. PY_CHECK_ARG_TYPE(0, tp_type);
  38. PY_CHECK_ARG_TYPE(1, tp_int);
  39. PY_CHECK_ARG_TYPE(2, tp_int);
  40. int n_cols = argv[1]._i64;
  41. int n_rows = argv[2]._i64;
  42. int numel = n_cols * n_rows;
  43. if(n_cols <= 0 || n_rows <= 0) return ValueError("array2d() expected positive dimensions");
  44. c11_array2d* ud = py_array2d(py_pushtmp(), n_cols, n_rows);
  45. // setup initial values
  46. if(py_callable(default_)) {
  47. for(int j = 0; j < n_rows; j++) {
  48. for(int i = 0; i < n_cols; i++) {
  49. py_TValue tmp;
  50. py_newvec2i(&tmp,
  51. (c11_vec2i){
  52. {i, j}
  53. });
  54. bool ok = py_call(default_, 1, &tmp);
  55. if(!ok) return false;
  56. ud->data[j * n_cols + i] = *py_retval();
  57. }
  58. }
  59. } else {
  60. for(int i = 0; i < numel; i++) {
  61. ud->data[i] = *default_;
  62. }
  63. }
  64. py_assign(py_retval(), py_peek(-1));
  65. py_pop();
  66. return true;
  67. }
  68. static bool array2d_n_cols(int argc, py_Ref argv) {
  69. PY_CHECK_ARGC(1);
  70. c11_array2d* self = py_touserdata(argv);
  71. py_newint(py_retval(), self->n_cols);
  72. return true;
  73. }
  74. static bool array2d_n_rows(int argc, py_Ref argv) {
  75. PY_CHECK_ARGC(1);
  76. c11_array2d* self = py_touserdata(argv);
  77. py_newint(py_retval(), self->n_rows);
  78. return true;
  79. }
  80. static bool array2d_numel(int argc, py_Ref argv) {
  81. PY_CHECK_ARGC(1);
  82. c11_array2d* self = py_touserdata(argv);
  83. py_newint(py_retval(), self->numel);
  84. return true;
  85. }
  86. static bool array2d_is_valid(int argc, py_Ref argv) {
  87. c11_array2d* self = py_touserdata(argv);
  88. int col, row;
  89. if(argc == 2) {
  90. PY_CHECK_ARG_TYPE(1, tp_vec2i);
  91. c11_vec2i pos = py_tovec2i(py_arg(1));
  92. col = pos.x;
  93. row = pos.y;
  94. } else if(argc == 3) {
  95. PY_CHECK_ARG_TYPE(1, tp_int);
  96. PY_CHECK_ARG_TYPE(2, tp_int);
  97. col = py_toint(py_arg(1));
  98. row = py_toint(py_arg(2));
  99. } else {
  100. return TypeError("is_valid() expected 2 or 3 arguments");
  101. }
  102. py_newbool(py_retval(), py_array2d_is_valid(self, col, row));
  103. return true;
  104. }
  105. static bool array2d_get(int argc, py_Ref argv) {
  106. py_Ref default_;
  107. c11_array2d* self = py_touserdata(argv);
  108. PY_CHECK_ARG_TYPE(1, tp_int);
  109. PY_CHECK_ARG_TYPE(2, tp_int);
  110. if(argc == 3) {
  111. default_ = py_None();
  112. } else if(argc == 4) {
  113. default_ = py_arg(3);
  114. } else {
  115. return TypeError("get() expected 2 or 3 arguments");
  116. }
  117. int col = py_toint(py_arg(1));
  118. int row = py_toint(py_arg(2));
  119. if(py_array2d_is_valid(self, col, row)) {
  120. py_assign(py_retval(), py_array2d__get(self, col, row));
  121. } else {
  122. py_assign(py_retval(), default_);
  123. }
  124. return true;
  125. }
  126. static bool _array2d_check_all_type(c11_array2d* self, py_Type type) {
  127. for(int i = 0; i < self->numel; i++) {
  128. py_Type item_type = self->data[i].type;
  129. if(item_type != type) {
  130. const char* fmt = "expected array2d[%t], got %t";
  131. return TypeError(fmt, type, item_type);
  132. }
  133. }
  134. return true;
  135. }
  136. static bool _check_same_shape(int colA, int rowA, int colB, int rowB) {
  137. if(colA != colB || rowA != rowB) {
  138. const char* fmt = "expected the same shape: (%d, %d) != (%d, %d)";
  139. return ValueError(fmt, colA, rowA, colB, rowB);
  140. }
  141. return true;
  142. }
  143. static bool _array2d_check_same_shape(c11_array2d* self, c11_array2d* other) {
  144. return _check_same_shape(self->n_cols, self->n_rows, other->n_cols, other->n_rows);
  145. }
  146. static bool array2d_all(int argc, py_Ref argv) {
  147. PY_CHECK_ARGC(1);
  148. c11_array2d* self = py_touserdata(argv);
  149. if(!_array2d_check_all_type(self, tp_bool)) return false;
  150. for(int i = 0; i < self->numel; i++) {
  151. if(!py_tobool(self->data + i)) {
  152. py_newbool(py_retval(), false);
  153. return true;
  154. }
  155. }
  156. py_newbool(py_retval(), true);
  157. return true;
  158. }
  159. static bool array2d_any(int argc, py_Ref argv) {
  160. PY_CHECK_ARGC(1);
  161. c11_array2d* self = py_touserdata(argv);
  162. if(!_array2d_check_all_type(self, tp_bool)) return false;
  163. for(int i = 0; i < self->numel; i++) {
  164. if(py_tobool(self->data + i)) {
  165. py_newbool(py_retval(), true);
  166. return true;
  167. }
  168. }
  169. py_newbool(py_retval(), false);
  170. return true;
  171. }
  172. static bool array2d__eq__(int argc, py_Ref argv) {
  173. PY_CHECK_ARGC(2);
  174. c11_array2d* self = py_touserdata(argv);
  175. c11_array2d* res = py_array2d(py_pushtmp(), self->n_cols, self->n_rows);
  176. if(py_istype(py_arg(1), tp_array2d)) {
  177. c11_array2d* other = py_touserdata(py_arg(1));
  178. if(!_array2d_check_same_shape(self, other)) return false;
  179. for(int i = 0; i < self->numel; i++) {
  180. int code = py_equal(self->data + i, other->data + i);
  181. if(code == -1) return false;
  182. py_newbool(res->data + i, (bool)code);
  183. }
  184. } else {
  185. // broadcast
  186. for(int i = 0; i < self->numel; i++) {
  187. int code = py_equal(self->data + i, py_arg(1));
  188. if(code == -1) return false;
  189. py_newbool(res->data + i, (bool)code);
  190. }
  191. }
  192. py_assign(py_retval(), py_peek(-1));
  193. py_pop();
  194. return true;
  195. }
  196. static bool array2d__ne__(int argc, py_Ref argv) {
  197. bool ok = array2d__eq__(argc, argv);
  198. if(!ok) return false;
  199. c11_array2d* res = py_touserdata(py_retval());
  200. py_TValue* data = res->data;
  201. for(int i = 0; i < res->numel; i++) {
  202. py_newbool(&data[i], !py_tobool(&data[i]));
  203. }
  204. return true;
  205. }
  206. static bool array2d__repr__(int argc, py_Ref argv) {
  207. PY_CHECK_ARGC(1);
  208. c11_array2d* self = py_touserdata(argv);
  209. char buf[256];
  210. snprintf(buf, sizeof(buf), "array2d(%d, %d)", self->n_cols, self->n_rows);
  211. py_newstr(py_retval(), buf);
  212. return true;
  213. }
  214. static bool array2d__iter__(int argc, py_Ref argv) {
  215. PY_CHECK_ARGC(1);
  216. c11_array2d* self = py_touserdata(argv);
  217. c11_array2d_iterator* ud =
  218. py_newobject(py_retval(), tp_array2d_iterator, 1, sizeof(c11_array2d_iterator));
  219. py_setslot(py_retval(), 0, argv); // keep the array alive
  220. ud->array = self;
  221. ud->index = 0;
  222. return true;
  223. }
  224. // __iter__(self) -> Iterator[tuple[int, int, T]]
  225. static bool array2d_iterator__next__(int argc, py_Ref argv) {
  226. PY_CHECK_ARGC(1);
  227. c11_array2d_iterator* self = py_touserdata(argv);
  228. if(self->index < self->array->numel) {
  229. div_t res = div(self->index, self->array->n_cols);
  230. py_newtuple(py_retval(), 2);
  231. py_TValue* data = py_tuple_data(py_retval());
  232. py_newvec2i(&data[0],
  233. (c11_vec2i){
  234. {res.rem, res.quot}
  235. });
  236. py_assign(&data[1], self->array->data + self->index);
  237. self->index++;
  238. return true;
  239. }
  240. return StopIteration();
  241. }
  242. static bool array2d_map(int argc, py_Ref argv) {
  243. // def map(self, f: Callable[[T], Any]) -> 'array2d': ...
  244. PY_CHECK_ARGC(2);
  245. c11_array2d* self = py_touserdata(argv);
  246. py_Ref f = py_arg(1);
  247. c11_array2d* res = py_array2d(py_pushtmp(), self->n_cols, self->n_rows);
  248. for(int i = 0; i < self->numel; i++) {
  249. bool ok = py_call(f, 1, self->data + i);
  250. if(!ok) return false;
  251. res->data[i] = *py_retval();
  252. }
  253. py_assign(py_retval(), py_peek(-1));
  254. py_pop();
  255. return true;
  256. }
  257. static bool array2d_copy(int argc, py_Ref argv) {
  258. // def copy(self) -> 'array2d': ...
  259. PY_CHECK_ARGC(1);
  260. c11_array2d* self = py_touserdata(argv);
  261. c11_array2d* res = py_array2d(py_retval(), self->n_cols, self->n_rows);
  262. memcpy(res->data, self->data, self->numel * sizeof(py_TValue));
  263. return true;
  264. }
  265. static bool array2d_fill_(int argc, py_Ref argv) {
  266. // def fill_(self, value: T) -> None: ...
  267. PY_CHECK_ARGC(2);
  268. c11_array2d* self = py_touserdata(argv);
  269. for(int i = 0; i < self->numel; i++)
  270. self->data[i] = argv[1];
  271. py_newnone(py_retval());
  272. return true;
  273. }
  274. static bool array2d_apply_(int argc, py_Ref argv) {
  275. // def apply_(self, f: Callable[[T], T]) -> None: ...
  276. PY_CHECK_ARGC(2);
  277. c11_array2d* self = py_touserdata(argv);
  278. py_Ref f = py_arg(1);
  279. for(int i = 0; i < self->numel; i++) {
  280. bool ok = py_call(f, 1, self->data + i);
  281. if(!ok) return false;
  282. self->data[i] = *py_retval();
  283. }
  284. py_newnone(py_retval());
  285. return true;
  286. }
  287. static bool array2d_copy_(int argc, py_Ref argv) {
  288. // def copy_(self, src: 'array2d') -> None: ...
  289. PY_CHECK_ARGC(2);
  290. c11_array2d* self = py_touserdata(argv);
  291. py_Type src_type = py_typeof(py_arg(1));
  292. if(src_type == tp_array2d) {
  293. c11_array2d* src = py_touserdata(py_arg(1));
  294. if(!_array2d_check_same_shape(self, src)) return false;
  295. memcpy(self->data, src->data, self->numel * sizeof(py_TValue));
  296. } else {
  297. py_TValue* data;
  298. int length = pk_arrayview(py_arg(1), &data);
  299. if(length != -1) {
  300. if(self->numel != length) {
  301. return ValueError("copy_() expected the same numel: %d != %d", self->numel, length);
  302. }
  303. memcpy(self->data, data, self->numel * sizeof(py_TValue));
  304. } else {
  305. return TypeError("copy_() expected `array2d`, `list` or `tuple`, got '%t", src_type);
  306. }
  307. }
  308. py_newnone(py_retval());
  309. return true;
  310. }
  311. // fromlist(data: list[list[T]]) -> array2d[T]
  312. static bool array2d_fromlist_STATIC(int argc, py_Ref argv) {
  313. PY_CHECK_ARGC(1);
  314. if(!py_checktype(argv, tp_list)) return false;
  315. int n_rows = py_list_len(argv);
  316. if(n_rows == 0) return ValueError("fromlist() expected a non-empty list");
  317. int n_cols = -1;
  318. for(int j = 0; j < n_rows; j++) {
  319. py_Ref row_j = py_list_getitem(argv, j);
  320. if(!py_checktype(row_j, tp_list)) return false;
  321. int n_cols_j = py_list_len(row_j);
  322. if(n_cols == -1) {
  323. if(n_cols_j == 0) return ValueError("fromlist() expected a non-empty list");
  324. n_cols = n_cols_j;
  325. } else if(n_cols != n_cols_j) {
  326. return ValueError("fromlist() expected a list of lists with the same length");
  327. }
  328. }
  329. c11_array2d* res = py_array2d(py_retval(), n_cols, n_rows);
  330. for(int j = 0; j < n_rows; j++) {
  331. py_Ref row_j = py_list_getitem(argv, j);
  332. for(int i = 0; i < n_cols; i++) {
  333. py_array2d__set(res, i, j, py_list_getitem(row_j, i));
  334. }
  335. }
  336. return true;
  337. }
  338. // tolist(self) -> list[list[T]]
  339. static bool array2d_tolist(int argc, py_Ref argv) {
  340. PY_CHECK_ARGC(1);
  341. c11_array2d* self = py_touserdata(argv);
  342. py_newlistn(py_retval(), self->n_rows);
  343. for(int j = 0; j < self->n_rows; j++) {
  344. py_Ref row_j = py_list_getitem(py_retval(), j);
  345. py_newlistn(row_j, self->n_cols);
  346. for(int i = 0; i < self->n_cols; i++) {
  347. py_list_setitem(row_j, i, py_array2d__get(self, i, j));
  348. }
  349. }
  350. return true;
  351. }
  352. static bool array2d_render(int argc, py_Ref argv) {
  353. PY_CHECK_ARGC(1);
  354. c11_sbuf buf;
  355. c11_sbuf__ctor(&buf);
  356. c11_array2d* self = py_touserdata(argv);
  357. for(int j = 0; j < self->n_rows; j++) {
  358. for(int i = 0; i < self->n_cols; i++) {
  359. py_Ref item = py_array2d__get(self, i, j);
  360. if(!py_str(item)) return false;
  361. c11_sbuf__write_sv(&buf, py_tosv(py_retval()));
  362. }
  363. if(j < self->n_rows - 1) c11_sbuf__write_char(&buf, '\n');
  364. }
  365. c11_sbuf__py_submit(&buf, py_retval());
  366. return true;
  367. }
  368. // count(self, value: T) -> int
  369. static bool array2d_count(int argc, py_Ref argv) {
  370. PY_CHECK_ARGC(2);
  371. c11_array2d* self = py_touserdata(argv);
  372. int count = 0;
  373. for(int i = 0; i < self->numel; i++) {
  374. int res = py_equal(self->data + i, &argv[1]);
  375. if(res == -1) return false;
  376. count += res;
  377. }
  378. py_newint(py_retval(), count);
  379. return true;
  380. }
  381. // get_bounding_rect(self, value: T) -> tuple[int, int, int, int]
  382. static bool array2d_get_bounding_rect(int argc, py_Ref argv) {
  383. PY_CHECK_ARGC(2);
  384. c11_array2d* self = py_touserdata(argv);
  385. py_Ref value = py_arg(1);
  386. int left = self->n_cols;
  387. int top = self->n_rows;
  388. int right = 0;
  389. int bottom = 0;
  390. for(int j = 0; j < self->n_rows; j++) {
  391. for(int i = 0; i < self->n_cols; i++) {
  392. int res = py_equal(py_array2d__get(self, i, j), value);
  393. if(res == -1) return false;
  394. if(res == 1) {
  395. left = c11__min(left, i);
  396. top = c11__min(top, j);
  397. right = c11__max(right, i);
  398. bottom = c11__max(bottom, j);
  399. }
  400. }
  401. }
  402. int width = right - left + 1;
  403. int height = bottom - top + 1;
  404. if(width <= 0 || height <= 0) {
  405. return ValueError("value not found");
  406. } else {
  407. py_newtuple(py_retval(), 4);
  408. py_TValue* data = py_tuple_data(py_retval());
  409. py_newint(&data[0], left);
  410. py_newint(&data[1], top);
  411. py_newint(&data[2], width);
  412. py_newint(&data[3], height);
  413. }
  414. return true;
  415. }
  416. // count_neighbors(self, value: T, neighborhood: Neighborhood) -> array2d[int]
  417. static bool array2d_count_neighbors(int argc, py_Ref argv) {
  418. PY_CHECK_ARGC(3);
  419. c11_array2d* self = py_touserdata(argv);
  420. c11_array2d* res = py_array2d(py_pushtmp(), self->n_cols, self->n_rows);
  421. py_Ref value = py_arg(1);
  422. const char* neighborhood = py_tostr(py_arg(2));
  423. #define INC_COUNT(i, j) \
  424. do { \
  425. if(py_array2d_is_valid(self, i, j)) { \
  426. int res = py_equal(py_array2d__get(self, i, j), value); \
  427. if(res == -1) return false; \
  428. count += res; \
  429. } \
  430. } while(0)
  431. if(strcmp(neighborhood, "Moore") == 0) {
  432. for(int j = 0; j < res->n_rows; j++) {
  433. for(int i = 0; i < res->n_cols; i++) {
  434. int count = 0;
  435. INC_COUNT(i - 1, j - 1);
  436. INC_COUNT(i, j - 1);
  437. INC_COUNT(i + 1, j - 1);
  438. INC_COUNT(i - 1, j);
  439. INC_COUNT(i + 1, j);
  440. INC_COUNT(i - 1, j + 1);
  441. INC_COUNT(i, j + 1);
  442. INC_COUNT(i + 1, j + 1);
  443. py_newint(py_array2d__get(res, i, j), count);
  444. }
  445. }
  446. } else if(strcmp(neighborhood, "von Neumann") == 0) {
  447. for(int j = 0; j < res->n_rows; j++) {
  448. for(int i = 0; i < res->n_cols; i++) {
  449. int count = 0;
  450. INC_COUNT(i, j - 1);
  451. INC_COUNT(i - 1, j);
  452. INC_COUNT(i + 1, j);
  453. INC_COUNT(i, j + 1);
  454. py_newint(py_array2d__get(res, i, j), count);
  455. }
  456. }
  457. } else {
  458. return ValueError("neighborhood must be 'Moore' or 'von Neumann'");
  459. }
  460. py_assign(py_retval(), py_peek(-1));
  461. py_pop();
  462. return true;
  463. }
  464. #define HANDLE_SLICE() \
  465. int start_col, stop_col, step_col; \
  466. int start_row, stop_row, step_row; \
  467. if(!pk__parse_int_slice(x, self->n_cols, &start_col, &stop_col, &step_col)) return false; \
  468. if(!pk__parse_int_slice(y, self->n_rows, &start_row, &stop_row, &step_row)) return false; \
  469. if(step_col != 1 || step_row != 1) return ValueError("slice step must be 1"); \
  470. int slice_width = stop_col - start_col; \
  471. int slice_height = stop_row - start_row; \
  472. if(slice_width <= 0 || slice_height <= 0) \
  473. return ValueError("slice width and height must be positive");
  474. static bool _array2d_IndexError(c11_array2d* self, int col, int row) {
  475. return IndexError("(%d, %d) is not a valid index of array2d(%d, %d)",
  476. col,
  477. row,
  478. self->n_cols,
  479. self->n_rows);
  480. }
  481. static bool array2d__getitem__(int argc, py_Ref argv) {
  482. PY_CHECK_ARGC(2);
  483. c11_array2d* self = py_touserdata(argv);
  484. if(argv[1].type == tp_vec2i) {
  485. // fastpath for vec2i
  486. c11_vec2i pos = py_tovec2i(&argv[1]);
  487. if(py_array2d_is_valid(self, pos.x, pos.y)) {
  488. py_assign(py_retval(), py_array2d__get(self, pos.x, pos.y));
  489. return true;
  490. }
  491. return _array2d_IndexError(self, pos.x, pos.y);
  492. }
  493. if(argv[1].type == tp_array2d) {
  494. c11_array2d* mask = py_touserdata(&argv[1]);
  495. if(!_array2d_check_same_shape(self, mask)) return false;
  496. if(!_array2d_check_all_type(mask, tp_bool)) return false;
  497. py_newlist(py_retval());
  498. for(int i = 0; i < self->numel; i++) {
  499. if(py_tobool(mask->data + i)) py_list_append(py_retval(), self->data + i);
  500. }
  501. return true;
  502. }
  503. PY_CHECK_ARG_TYPE(1, tp_tuple);
  504. if(py_tuple_len(py_arg(1)) != 2) return TypeError("expected a tuple of 2 elements");
  505. py_Ref x = py_tuple_getitem(py_arg(1), 0);
  506. py_Ref y = py_tuple_getitem(py_arg(1), 1);
  507. if(py_isint(x) && py_isint(y)) {
  508. int col = py_toint(x);
  509. int row = py_toint(y);
  510. if(py_array2d_is_valid(self, col, row)) {
  511. py_assign(py_retval(), py_array2d__get(self, col, row));
  512. return true;
  513. }
  514. return _array2d_IndexError(self, col, row);
  515. } else if(py_istype(x, tp_slice) && py_istype(y, tp_slice)) {
  516. HANDLE_SLICE();
  517. c11_array2d* res = py_array2d(py_retval(), slice_width, slice_height);
  518. for(int j = start_row; j < stop_row; j++) {
  519. for(int i = start_col; i < stop_col; i++) {
  520. py_array2d__set(res, i - start_col, j - start_row, py_array2d__get(self, i, j));
  521. }
  522. }
  523. return true;
  524. } else {
  525. return TypeError("expected a tuple of 2 ints or slices");
  526. }
  527. }
  528. static bool array2d__setitem__(int argc, py_Ref argv) {
  529. PY_CHECK_ARGC(3);
  530. c11_array2d* self = py_touserdata(argv);
  531. py_Ref value = py_arg(2);
  532. if(argv[1].type == tp_vec2i) {
  533. // fastpath for vec2i
  534. c11_vec2i pos = py_tovec2i(&argv[1]);
  535. if(py_array2d_is_valid(self, pos.x, pos.y)) {
  536. py_array2d__set(self, pos.x, pos.y, value);
  537. py_newnone(py_retval());
  538. return true;
  539. }
  540. return _array2d_IndexError(self, pos.x, pos.y);
  541. }
  542. if(argv[1].type == tp_array2d) {
  543. c11_array2d* mask = py_touserdata(&argv[1]);
  544. if(!_array2d_check_same_shape(self, mask)) return false;
  545. if(!_array2d_check_all_type(mask, tp_bool)) return false;
  546. for(int i = 0; i < self->numel; i++) {
  547. if(py_tobool(mask->data + i)) self->data[i] = *value;
  548. }
  549. py_newnone(py_retval());
  550. return true;
  551. }
  552. PY_CHECK_ARG_TYPE(1, tp_tuple);
  553. if(py_tuple_len(py_arg(1)) != 2) return TypeError("expected a tuple of 2 elements");
  554. py_Ref x = py_tuple_getitem(py_arg(1), 0);
  555. py_Ref y = py_tuple_getitem(py_arg(1), 1);
  556. if(py_isint(x) && py_isint(y)) {
  557. int col = py_toint(x);
  558. int row = py_toint(y);
  559. if(py_array2d_is_valid(self, col, row)) {
  560. py_array2d__set(self, col, row, value);
  561. py_newnone(py_retval());
  562. return true;
  563. }
  564. return _array2d_IndexError(self, col, row);
  565. } else if(py_istype(x, tp_slice) && py_istype(y, tp_slice)) {
  566. HANDLE_SLICE();
  567. bool is_basic_type = false;
  568. switch(value->type) {
  569. case tp_int:
  570. case tp_float:
  571. case tp_str:
  572. case tp_NoneType:
  573. case tp_bool: is_basic_type = true; break;
  574. default: {
  575. if(!py_istype(value, tp_array2d)) {
  576. return TypeError("expected int/float/str/bool/None or an array2d instance");
  577. }
  578. }
  579. }
  580. if(is_basic_type) {
  581. for(int j = start_row; j < stop_row; j++) {
  582. for(int i = start_col; i < stop_col; i++) {
  583. py_array2d__set(self, i, j, py_arg(2));
  584. }
  585. }
  586. } else {
  587. c11_array2d* src = py_touserdata(value);
  588. if(!_check_same_shape(slice_width, slice_height, src->n_cols, src->n_rows))
  589. return false;
  590. for(int j = 0; j < slice_height; j++) {
  591. for(int i = 0; i < slice_width; i++) {
  592. py_array2d__set(self, i + start_col, j + start_row, py_array2d__get(src, i, j));
  593. }
  594. }
  595. }
  596. py_newnone(py_retval());
  597. return true;
  598. } else {
  599. return TypeError("expected a tuple of 2 ints or slices");
  600. }
  601. }
  602. // convolve(self: array2d[int], kernel: array2d[int], padding: int) -> array2d[int]
  603. static bool array2d_convolve(int argc, py_Ref argv) {
  604. PY_CHECK_ARGC(3);
  605. PY_CHECK_ARG_TYPE(1, tp_array2d);
  606. PY_CHECK_ARG_TYPE(2, tp_int);
  607. c11_array2d* self = py_touserdata(argv);
  608. c11_array2d* kernel = py_touserdata(py_arg(1));
  609. int padding = py_toint(py_arg(2));
  610. if(kernel->n_cols != kernel->n_rows) { return ValueError("kernel must be square"); }
  611. int ksize = kernel->n_cols;
  612. if(ksize % 2 == 0) return ValueError("kernel size must be odd");
  613. int ksize_half = ksize / 2;
  614. if(!_array2d_check_all_type(self, tp_int)) return false;
  615. if(!_array2d_check_all_type(kernel, tp_int)) return false;
  616. c11_array2d* res = py_array2d(py_pushtmp(), self->n_cols, self->n_rows);
  617. for(int j = 0; j < self->n_rows; j++) {
  618. for(int i = 0; i < self->n_cols; i++) {
  619. py_i64 sum = 0;
  620. for(int jj = 0; jj < ksize; jj++) {
  621. for(int ii = 0; ii < ksize; ii++) {
  622. int x = i + ii - ksize_half;
  623. int y = j + jj - ksize_half;
  624. py_i64 _0, _1;
  625. if(x < 0 || x >= self->n_cols || y < 0 || y >= self->n_rows) {
  626. _0 = padding;
  627. } else {
  628. _0 = py_toint(py_array2d__get(self, x, y));
  629. }
  630. _1 = py_toint(py_array2d__get(kernel, ii, jj));
  631. sum += _0 * _1;
  632. }
  633. }
  634. py_newint(py_array2d__get(res, i, j), sum);
  635. }
  636. }
  637. py_assign(py_retval(), py_peek(-1));
  638. py_pop();
  639. return true;
  640. }
  641. void pk__add_module_array2d() {
  642. py_GlobalRef mod = py_newmodule("array2d");
  643. py_Type array2d = pk_newtype("array2d", tp_object, mod, NULL, false, true);
  644. py_Type array2d_iterator = pk_newtype("array2d_iterator", tp_object, mod, NULL, false, true);
  645. assert(array2d == tp_array2d);
  646. assert(array2d_iterator == tp_array2d_iterator);
  647. py_setdict(mod, py_name("array2d"), py_tpobject(array2d));
  648. // array2d is unhashable
  649. py_setdict(py_tpobject(array2d), __hash__, py_None());
  650. py_bindmagic(array2d_iterator, __iter__, pk_wrapper__self);
  651. py_bindmagic(array2d_iterator, __next__, array2d_iterator__next__);
  652. py_bind(py_tpobject(array2d),
  653. "__new__(cls, n_cols: int, n_rows: int, default=None)",
  654. array2d__new__);
  655. py_bindmagic(array2d, __eq__, array2d__eq__);
  656. py_bindmagic(array2d, __ne__, array2d__ne__);
  657. py_bindmagic(array2d, __repr__, array2d__repr__);
  658. py_bindmagic(array2d, __iter__, array2d__iter__);
  659. py_bindmagic(array2d, __getitem__, array2d__getitem__);
  660. py_bindmagic(array2d, __setitem__, array2d__setitem__);
  661. py_bindproperty(array2d, "n_cols", array2d_n_cols, NULL);
  662. py_bindproperty(array2d, "n_rows", array2d_n_rows, NULL);
  663. py_bindproperty(array2d, "width", array2d_n_cols, NULL);
  664. py_bindproperty(array2d, "height", array2d_n_rows, NULL);
  665. py_bindproperty(array2d, "numel", array2d_numel, NULL);
  666. py_bindmethod(array2d, "is_valid", array2d_is_valid);
  667. py_bindmethod(array2d, "get", array2d_get);
  668. py_bindmethod(array2d, "map", array2d_map);
  669. py_bindmethod(array2d, "copy", array2d_copy);
  670. py_bindmethod(array2d, "fill_", array2d_fill_);
  671. py_bindmethod(array2d, "apply_", array2d_apply_);
  672. py_bindmethod(array2d, "copy_", array2d_copy_);
  673. py_bindmethod(array2d, "render", array2d_render);
  674. py_bindmethod(array2d, "all", array2d_all);
  675. py_bindmethod(array2d, "any", array2d_any);
  676. py_bindstaticmethod(array2d, "fromlist", array2d_fromlist_STATIC);
  677. py_bindmethod(array2d, "tolist", array2d_tolist);
  678. py_bindmethod(array2d, "count", array2d_count);
  679. py_bindmethod(array2d, "get_bounding_rect", array2d_get_bounding_rect);
  680. py_bindmethod(array2d, "count_neighbors", array2d_count_neighbors);
  681. py_bindmethod(array2d, "convolve", array2d_convolve);
  682. const char* scc =
  683. "\ndef get_connected_components(self, value: T, neighborhood: Neighborhood) -> tuple[array2d[int], int]:\n from collections import deque\n from linalg import vec2i\n\n DIRS = [vec2i.LEFT, vec2i.RIGHT, vec2i.UP, vec2i.DOWN]\n assert neighborhood in ['Moore', 'von Neumann']\n\n if neighborhood == 'Moore':\n DIRS.extend([\n vec2i.LEFT+vec2i.UP,\n vec2i.RIGHT+vec2i.UP,\n vec2i.LEFT+vec2i.DOWN,\n vec2i.RIGHT+vec2i.DOWN\n ])\n\n visited = array2d[int](self.width, self.height, default=0)\n queue = deque()\n count = 0\n for y in range(self.height):\n for x in range(self.width):\n if visited[x, y] or self[x, y] != value:\n continue\n count += 1\n queue.append((x, y))\n visited[x, y] = count\n while queue:\n cx, cy = queue.popleft()\n for dx, dy in DIRS:\n nx, ny = cx+dx, cy+dy\n if self.is_valid(nx, ny) and not visited[nx, ny] and self[nx, ny] == value:\n queue.append((nx, ny))\n visited[nx, ny] = count\n return visited, count\n\narray2d.get_connected_components = get_connected_components\ndel get_connected_components\n";
  684. if(!py_exec(scc, "array2d.py", EXEC_MODE, mod)) {
  685. py_printexc();
  686. c11__abort("failed to execute array2d.py");
  687. }
  688. }
  689. #undef INC_COUNT
  690. #undef HANDLE_SLICE