pickle.c 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730
  1. #include "pocketpy/common/vector.h"
  2. #include "pocketpy/interpreter/typeinfo.h"
  3. #include "pocketpy/pocketpy.h"
  4. #include "pocketpy/common/utils.h"
  5. #include "pocketpy/common/sstream.h"
  6. #include "pocketpy/interpreter/vm.h"
  7. #include "pocketpy/interpreter/array2d.h"
  8. #include <stdint.h>
  9. typedef enum {
  10. // clang-format off
  11. PKL_MEMO_GET,
  12. PKL_MEMO_SET,
  13. PKL_NIL, PKL_NONE, PKL_ELLIPSIS,
  14. PKL_INT_0, PKL_INT_1, PKL_INT_2, PKL_INT_3, PKL_INT_4, PKL_INT_5, PKL_INT_6, PKL_INT_7,
  15. PKL_INT_8, PKL_INT_9, PKL_INT_10, PKL_INT_11, PKL_INT_12, PKL_INT_13, PKL_INT_14, PKL_INT_15,
  16. PKL_INT8, PKL_INT16, PKL_INT32, PKL_INT64,
  17. PKL_FLOAT32, PKL_FLOAT64,
  18. PKL_TRUE, PKL_FALSE,
  19. PKL_STRING, PKL_BYTES,
  20. PKL_BUILD_LIST,
  21. PKL_BUILD_TUPLE,
  22. PKL_BUILD_DICT,
  23. PKL_VEC2, PKL_VEC3,
  24. PKL_VEC2I, PKL_VEC3I,
  25. PKL_TYPE,
  26. PKL_ARRAY2D,
  27. PKL_TVALUE,
  28. PKL_CALL,
  29. PKL_OBJECT,
  30. PKL_EOF,
  31. // clang-format on
  32. } PickleOp;
  33. typedef struct {
  34. bool* used_types;
  35. int used_types_length;
  36. c11_smallmap_p2i memo;
  37. c11_vector /*T=char*/ codes;
  38. } PickleObject;
  39. static void PickleObject__ctor(PickleObject* self) {
  40. self->used_types_length = pk_current_vm->types.length;
  41. self->used_types = PK_MALLOC(self->used_types_length);
  42. memset(self->used_types, 0, self->used_types_length);
  43. c11_smallmap_p2i__ctor(&self->memo);
  44. c11_vector__ctor(&self->codes, sizeof(char));
  45. }
  46. static void PickleObject__dtor(PickleObject* self) {
  47. PK_FREE(self->used_types);
  48. c11_smallmap_p2i__dtor(&self->memo);
  49. c11_vector__dtor(&self->codes);
  50. }
  51. static bool PickleObject__py_submit(PickleObject* self, py_OutRef out);
  52. static void PickleObject__write_bytes(PickleObject* buf, const void* data, int size) {
  53. c11_vector__extend(char, &buf->codes, data, size);
  54. }
  55. static void c11_sbuf__write_type_path(c11_sbuf* path_buf, py_Type type) {
  56. py_TypeInfo* ti = pk__type_info(type);
  57. if(py_isnil(&ti->module)) {
  58. c11_sbuf__write_cstr(path_buf, py_name2str(ti->name));
  59. return;
  60. }
  61. const char* mod_path = py_tostr(py_getdict(&ti->module, __path__));
  62. c11_sbuf__write_cstr(path_buf, mod_path);
  63. c11_sbuf__write_char(path_buf, '.');
  64. c11_sbuf__write_cstr(path_buf, py_name2str(ti->name));
  65. }
  66. static void pkl__emit_op(PickleObject* buf, PickleOp op) {
  67. c11_vector__push(char, &buf->codes, op);
  68. }
  69. static void pkl__emit_int(PickleObject* buf, py_i64 val) {
  70. if(val >= 0 && val <= 15) {
  71. pkl__emit_op(buf, PKL_INT_0 + val);
  72. return;
  73. }
  74. if(INT8_MIN <= val && val <= INT8_MAX) {
  75. pkl__emit_op(buf, PKL_INT8);
  76. PickleObject__write_bytes(buf, &val, 1);
  77. } else if(INT16_MIN <= val && val <= INT16_MAX) {
  78. pkl__emit_op(buf, PKL_INT16);
  79. PickleObject__write_bytes(buf, &val, 2);
  80. } else if(INT32_MIN <= val && val <= INT32_MAX) {
  81. pkl__emit_op(buf, PKL_INT32);
  82. PickleObject__write_bytes(buf, &val, 4);
  83. } else {
  84. pkl__emit_op(buf, PKL_INT64);
  85. PickleObject__write_bytes(buf, &val, 8);
  86. }
  87. }
  88. #define UNALIGNED_READ(p_val, p_buf) \
  89. do { \
  90. memcpy((p_val), (p_buf), sizeof(*(p_val))); \
  91. (p_buf) += sizeof(*(p_val)); \
  92. } while(0)
  93. static py_i64 pkl__read_int(const unsigned char** p) {
  94. PickleOp op = (PickleOp) * *p;
  95. (*p)++;
  96. switch(op) {
  97. // clang-format off
  98. case PKL_INT_0: return 0; case PKL_INT_1: return 1; case PKL_INT_2: return 2; case PKL_INT_3: return 3;
  99. case PKL_INT_4: return 4; case PKL_INT_5: return 5; case PKL_INT_6: return 6; case PKL_INT_7: return 7;
  100. case PKL_INT_8: return 8; case PKL_INT_9: return 9; case PKL_INT_10: return 10; case PKL_INT_11: return 11;
  101. case PKL_INT_12: return 12; case PKL_INT_13: return 13; case PKL_INT_14: return 14; case PKL_INT_15: return 15;
  102. // clang-format on
  103. case PKL_INT8: {
  104. int8_t val;
  105. UNALIGNED_READ(&val, *p);
  106. return val;
  107. }
  108. case PKL_INT16: {
  109. int16_t val;
  110. UNALIGNED_READ(&val, *p);
  111. return val;
  112. }
  113. case PKL_INT32: {
  114. int32_t val;
  115. UNALIGNED_READ(&val, *p);
  116. return val;
  117. }
  118. case PKL_INT64: {
  119. int64_t val;
  120. UNALIGNED_READ(&val, *p);
  121. return val;
  122. }
  123. default: c11__abort("pkl__read_int(): invalid op: %d", op);
  124. }
  125. }
  126. static bool pickle_loads(int argc, py_Ref argv) {
  127. PY_CHECK_ARGC(1);
  128. PY_CHECK_ARG_TYPE(0, tp_bytes);
  129. int size;
  130. const unsigned char* data = py_tobytes(argv, &size);
  131. return py_pickle_loads(data, size);
  132. }
  133. static bool pickle_dumps(int argc, py_Ref argv) {
  134. PY_CHECK_ARGC(1);
  135. return py_pickle_dumps(argv);
  136. }
  137. void pk__add_module_pickle() {
  138. py_Ref mod = py_newmodule("pickle");
  139. py_bindfunc(mod, "loads", pickle_loads);
  140. py_bindfunc(mod, "dumps", pickle_dumps);
  141. }
  142. static bool pkl__write_object(PickleObject* buf, py_TValue* obj);
  143. static bool pkl__write_array(PickleObject* buf, PickleOp op, py_TValue* arr, int length) {
  144. for(int i = 0; i < length; i++) {
  145. bool ok = pkl__write_object(buf, arr + i);
  146. if(!ok) return false;
  147. }
  148. pkl__emit_op(buf, op);
  149. pkl__emit_int(buf, length);
  150. return true;
  151. }
  152. static bool pkl__write_dict_kv(py_Ref k, py_Ref v, void* ctx) {
  153. PickleObject* buf = (PickleObject*)ctx;
  154. if(!pkl__write_object(buf, k)) return false;
  155. if(!pkl__write_object(buf, v)) return false;
  156. return true;
  157. }
  158. static bool pkl__try_memo(PickleObject* buf, PyObject* memo_key) {
  159. int index = c11_smallmap_p2i__get(&buf->memo, memo_key, -1);
  160. if(index != -1) {
  161. pkl__emit_op(buf, PKL_MEMO_GET);
  162. pkl__emit_int(buf, index);
  163. return true;
  164. }
  165. return false;
  166. }
  167. static void pkl__store_memo(PickleObject* buf, PyObject* memo_key) {
  168. int index = buf->memo.length;
  169. c11_smallmap_p2i__set(&buf->memo, memo_key, index);
  170. pkl__emit_op(buf, PKL_MEMO_SET);
  171. pkl__emit_int(buf, index);
  172. }
  173. static bool pkl__write_object(PickleObject* buf, py_TValue* obj) {
  174. switch(obj->type) {
  175. case tp_nil: {
  176. return ValueError("'nil' object is not picklable");
  177. }
  178. case tp_NoneType: {
  179. pkl__emit_op(buf, PKL_NONE);
  180. return true;
  181. }
  182. case tp_ellipsis: {
  183. pkl__emit_op(buf, PKL_ELLIPSIS);
  184. return true;
  185. }
  186. case tp_int: {
  187. py_i64 val = obj->_i64;
  188. pkl__emit_int(buf, val);
  189. return true;
  190. }
  191. case tp_float: {
  192. py_f64 val = obj->_f64;
  193. float val32 = (float)val;
  194. if(val == val32) {
  195. pkl__emit_op(buf, PKL_FLOAT32);
  196. PickleObject__write_bytes(buf, &val32, 4);
  197. } else {
  198. pkl__emit_op(buf, PKL_FLOAT64);
  199. PickleObject__write_bytes(buf, &val, 8);
  200. }
  201. return true;
  202. }
  203. case tp_bool: {
  204. bool val = obj->_bool;
  205. pkl__emit_op(buf, val ? PKL_TRUE : PKL_FALSE);
  206. return true;
  207. }
  208. case tp_str: {
  209. if(pkl__try_memo(buf, obj->_obj))
  210. return true;
  211. else {
  212. pkl__emit_op(buf, PKL_STRING);
  213. c11_sv sv = py_tosv(obj);
  214. pkl__emit_int(buf, sv.size);
  215. PickleObject__write_bytes(buf, sv.data, sv.size);
  216. }
  217. pkl__store_memo(buf, obj->_obj);
  218. return true;
  219. }
  220. case tp_bytes: {
  221. if(pkl__try_memo(buf, obj->_obj))
  222. return true;
  223. else {
  224. pkl__emit_op(buf, PKL_BYTES);
  225. int size;
  226. unsigned char* data = py_tobytes(obj, &size);
  227. pkl__emit_int(buf, size);
  228. PickleObject__write_bytes(buf, data, size);
  229. }
  230. pkl__store_memo(buf, obj->_obj);
  231. return true;
  232. }
  233. case tp_list: {
  234. if(pkl__try_memo(buf, obj->_obj))
  235. return true;
  236. else {
  237. bool ok =
  238. pkl__write_array(buf, PKL_BUILD_LIST, py_list_data(obj), py_list_len(obj));
  239. if(!ok) return false;
  240. }
  241. pkl__store_memo(buf, obj->_obj);
  242. return true;
  243. }
  244. case tp_tuple: {
  245. if(pkl__try_memo(buf, obj->_obj))
  246. return true;
  247. else {
  248. bool ok =
  249. pkl__write_array(buf, PKL_BUILD_TUPLE, py_tuple_data(obj), py_tuple_len(obj));
  250. if(!ok) return false;
  251. }
  252. pkl__store_memo(buf, obj->_obj);
  253. return true;
  254. }
  255. case tp_dict: {
  256. if(pkl__try_memo(buf, obj->_obj))
  257. return true;
  258. else {
  259. bool ok = py_dict_apply(obj, pkl__write_dict_kv, (void*)buf);
  260. if(!ok) return false;
  261. pkl__emit_op(buf, PKL_BUILD_DICT);
  262. pkl__emit_int(buf, py_dict_len(obj));
  263. }
  264. pkl__store_memo(buf, obj->_obj);
  265. return true;
  266. }
  267. case tp_vec2: {
  268. c11_vec2 val = py_tovec2(obj);
  269. pkl__emit_op(buf, PKL_VEC2);
  270. PickleObject__write_bytes(buf, &val, sizeof(c11_vec2));
  271. return true;
  272. }
  273. case tp_vec3: {
  274. c11_vec3 val = py_tovec3(obj);
  275. pkl__emit_op(buf, PKL_VEC3);
  276. PickleObject__write_bytes(buf, &val, sizeof(c11_vec3));
  277. return true;
  278. }
  279. case tp_vec2i: {
  280. c11_vec2i val = py_tovec2i(obj);
  281. pkl__emit_op(buf, PKL_VEC2I);
  282. pkl__emit_int(buf, val.x);
  283. pkl__emit_int(buf, val.y);
  284. return true;
  285. }
  286. case tp_vec3i: {
  287. c11_vec3i val = py_tovec3i(obj);
  288. pkl__emit_op(buf, PKL_VEC3I);
  289. pkl__emit_int(buf, val.x);
  290. pkl__emit_int(buf, val.y);
  291. pkl__emit_int(buf, val.z);
  292. return true;
  293. }
  294. case tp_type: {
  295. pkl__emit_op(buf, PKL_TYPE);
  296. py_Type type = py_totype(obj);
  297. buf->used_types[type] = true;
  298. pkl__emit_int(buf, type);
  299. return true;
  300. }
  301. case tp_array2d: {
  302. if(pkl__try_memo(buf, obj->_obj))
  303. return true;
  304. else {
  305. c11_array2d* arr = py_touserdata(obj);
  306. for(int i = 0; i < arr->header.numel; i++) {
  307. if(arr->data[i].is_ptr)
  308. return TypeError(
  309. "'array2d' object is not picklable because it contains heap-allocated objects");
  310. buf->used_types[arr->data[i].type] = true;
  311. }
  312. pkl__emit_op(buf, PKL_ARRAY2D);
  313. pkl__emit_int(buf, arr->header.n_cols);
  314. pkl__emit_int(buf, arr->header.n_rows);
  315. PickleObject__write_bytes(buf, arr->data, arr->header.numel * sizeof(py_TValue));
  316. }
  317. pkl__store_memo(buf, obj->_obj);
  318. return true;
  319. }
  320. default: {
  321. if(!obj->is_ptr) {
  322. pkl__emit_op(buf, PKL_TVALUE);
  323. PickleObject__write_bytes(buf, obj, sizeof(py_TValue));
  324. buf->used_types[obj->type] = true;
  325. return true;
  326. }
  327. // try memo for `is_ptr=true` objects
  328. if(pkl__try_memo(buf, obj->_obj)) return true;
  329. py_TypeInfo* ti = pk__type_info(obj->type);
  330. py_Ref f_reduce = py_tpfindmagic(obj->type, __reduce__);
  331. if(f_reduce != NULL) {
  332. if(!py_call(f_reduce, 1, obj)) return false;
  333. // expected: (callable, args)
  334. py_Ref reduced = py_retval();
  335. if(!py_istuple(reduced)) { return TypeError("__reduce__ must return a tuple"); }
  336. if(py_tuple_len(reduced) != 2) {
  337. return TypeError("__reduce__ must return a tuple of length 2");
  338. }
  339. if(!pkl__write_object(buf, py_tuple_getitem(reduced, 0))) return false;
  340. pkl__emit_op(buf, PKL_NIL);
  341. py_Ref args_tuple = py_tuple_getitem(reduced, 1);
  342. int args_length = py_tuple_len(args_tuple);
  343. for(int i = 0; i < args_length; i++) {
  344. if(!pkl__write_object(buf, py_tuple_getitem(args_tuple, i))) return false;
  345. }
  346. pkl__emit_op(buf, PKL_CALL);
  347. pkl__emit_int(buf, args_length);
  348. // store memo
  349. pkl__store_memo(buf, obj->_obj);
  350. return true;
  351. }
  352. if(ti->is_python) {
  353. NameDict* dict = PyObject__dict(obj->_obj);
  354. for(int i = dict->length - 1; i >= 0; i--) {
  355. NameDict_KV* kv = c11__at(NameDict_KV, dict, i);
  356. if(!pkl__write_object(buf, &kv->value)) return false;
  357. }
  358. pkl__emit_op(buf, PKL_OBJECT);
  359. pkl__emit_int(buf, obj->type);
  360. buf->used_types[obj->type] = true;
  361. pkl__emit_int(buf, dict->length);
  362. for(int i = 0; i < dict->length; i++) {
  363. NameDict_KV* kv = c11__at(NameDict_KV, dict, i);
  364. c11_sv field = py_name2sv(kv->key);
  365. // include '\0'
  366. PickleObject__write_bytes(buf, field.data, field.size + 1);
  367. }
  368. // store memo
  369. pkl__store_memo(buf, obj->_obj);
  370. return true;
  371. }
  372. return TypeError("'%t' object is not picklable", obj->type);
  373. }
  374. }
  375. c11__unreachable();
  376. }
  377. bool py_pickle_dumps(py_Ref val) {
  378. PickleObject buf;
  379. PickleObject__ctor(&buf);
  380. bool ok = pkl__write_object(&buf, val);
  381. if(!ok) {
  382. PickleObject__dtor(&buf);
  383. return false;
  384. }
  385. pkl__emit_op(&buf, PKL_EOF);
  386. return PickleObject__py_submit(&buf, py_retval());
  387. }
  388. static py_Type pkl__header_find_type(c11_sv path) {
  389. int sep_index = c11_sv__rindex(path, '.');
  390. if(sep_index == -1) return py_gettype(NULL, py_namev(path));
  391. c11_sv mod_name = c11_sv__slice2(path, 0, sep_index);
  392. c11_sv name = c11_sv__slice(path, sep_index + 1);
  393. char buf[PK_MAX_MODULE_PATH_LEN + 1];
  394. memcpy(buf, mod_name.data, mod_name.size);
  395. buf[mod_name.size] = '\0';
  396. return py_gettype(buf, py_namev(name));
  397. }
  398. static c11_sv pkl__header_read_sv(const unsigned char** p, char sep) {
  399. c11_sv text;
  400. text.data = (const char*)*p;
  401. const char* p_end = strchr(text.data, sep);
  402. assert(p_end != NULL);
  403. text.size = p_end - text.data;
  404. *p = (const unsigned char*)p_end + 1;
  405. return text;
  406. }
  407. static py_i64 pkl__header_read_int(const unsigned char** p, char sep) {
  408. c11_sv text = pkl__header_read_sv(p, sep);
  409. py_i64 out;
  410. IntParsingResult res = c11__parse_uint(text, &out, 10);
  411. assert(res == IntParsing_SUCCESS);
  412. return out;
  413. }
  414. bool py_pickle_loads_body(const unsigned char* p, int memo_length, c11_smallmap_n2i* type_mapping);
  415. bool py_pickle_loads(const unsigned char* data, int size) {
  416. const unsigned char* p = data;
  417. c11_smallmap_n2i type_mapping;
  418. c11_smallmap_n2i__ctor(&type_mapping);
  419. while(true) {
  420. if(*p == '\n') {
  421. p++;
  422. break;
  423. }
  424. py_Type type = pkl__header_read_int(&p, '(');
  425. c11_sv path = pkl__header_read_sv(&p, ')');
  426. py_Type new_type = pkl__header_find_type(path);
  427. if(new_type == 0) {
  428. c11_smallmap_n2i__dtor(&type_mapping);
  429. return ImportError("cannot find type '%v'", path);
  430. }
  431. if(type != new_type) c11_smallmap_n2i__set(&type_mapping, type, new_type);
  432. }
  433. int memo_length = pkl__header_read_int(&p, '\n');
  434. bool ok = py_pickle_loads_body(p, memo_length, &type_mapping);
  435. c11_smallmap_n2i__dtor(&type_mapping);
  436. return ok;
  437. }
  438. static py_Type pkl__fix_type(py_Type type, c11_smallmap_n2i* type_mapping) {
  439. int new_type = c11_smallmap_n2i__get(type_mapping, type, -1);
  440. if(new_type != -1) return (py_Type)new_type;
  441. return type;
  442. }
  443. bool py_pickle_loads_body(const unsigned char* p, int memo_length, c11_smallmap_n2i* type_mapping) {
  444. py_StackRef p0 = py_peek(0);
  445. py_Ref p_memo = py_newtuple(py_pushtmp(), memo_length);
  446. while(true) {
  447. PickleOp op = (PickleOp)*p;
  448. p++;
  449. switch(op) {
  450. case PKL_MEMO_GET: {
  451. int index = pkl__read_int(&p);
  452. py_Ref val = &p_memo[index];
  453. assert(!py_isnil(val));
  454. py_push(val);
  455. break;
  456. }
  457. case PKL_MEMO_SET: {
  458. int index = pkl__read_int(&p);
  459. p_memo[index] = *py_peek(-1);
  460. break;
  461. }
  462. case PKL_NIL: {
  463. py_pushnil();
  464. break;
  465. }
  466. case PKL_NONE: {
  467. py_pushnone();
  468. break;
  469. }
  470. case PKL_ELLIPSIS: {
  471. py_newellipsis(py_pushtmp());
  472. break;
  473. }
  474. // clang-format off
  475. case PKL_INT_0: case PKL_INT_1: case PKL_INT_2: case PKL_INT_3:
  476. case PKL_INT_4: case PKL_INT_5: case PKL_INT_6: case PKL_INT_7:
  477. case PKL_INT_8: case PKL_INT_9: case PKL_INT_10: case PKL_INT_11:
  478. case PKL_INT_12: case PKL_INT_13: case PKL_INT_14: case PKL_INT_15: {
  479. py_newint(py_pushtmp(), op - PKL_INT_0);
  480. break;
  481. }
  482. // clang-format on
  483. case PKL_INT8: {
  484. int8_t val;
  485. UNALIGNED_READ(&val, p);
  486. py_newint(py_pushtmp(), val);
  487. break;
  488. }
  489. case PKL_INT16: {
  490. int16_t val;
  491. UNALIGNED_READ(&val, p);
  492. py_newint(py_pushtmp(), val);
  493. break;
  494. }
  495. case PKL_INT32: {
  496. int32_t val;
  497. UNALIGNED_READ(&val, p);
  498. py_newint(py_pushtmp(), val);
  499. break;
  500. }
  501. case PKL_INT64: {
  502. int64_t val;
  503. UNALIGNED_READ(&val, p);
  504. py_newint(py_pushtmp(), val);
  505. break;
  506. }
  507. case PKL_FLOAT32: {
  508. float val;
  509. UNALIGNED_READ(&val, p);
  510. py_newfloat(py_pushtmp(), val);
  511. break;
  512. }
  513. case PKL_FLOAT64: {
  514. double val;
  515. UNALIGNED_READ(&val, p);
  516. py_newfloat(py_pushtmp(), val);
  517. break;
  518. }
  519. case PKL_TRUE: {
  520. py_newbool(py_pushtmp(), true);
  521. break;
  522. }
  523. case PKL_FALSE: {
  524. py_newbool(py_pushtmp(), false);
  525. break;
  526. }
  527. case PKL_STRING: {
  528. int size = pkl__read_int(&p);
  529. char* dst = py_newstrn(py_pushtmp(), size);
  530. memcpy(dst, p, size);
  531. p += size;
  532. break;
  533. }
  534. case PKL_BYTES: {
  535. int size = pkl__read_int(&p);
  536. unsigned char* dst = py_newbytes(py_pushtmp(), size);
  537. memcpy(dst, p, size);
  538. p += size;
  539. break;
  540. }
  541. case PKL_BUILD_LIST: {
  542. int length = pkl__read_int(&p);
  543. py_OutRef val = py_retval();
  544. py_newlistn(val, length);
  545. for(int i = length - 1; i >= 0; i--) {
  546. py_StackRef item = py_peek(-1);
  547. py_list_setitem(val, i, item);
  548. py_pop();
  549. }
  550. py_push(val);
  551. break;
  552. }
  553. case PKL_BUILD_TUPLE: {
  554. int length = pkl__read_int(&p);
  555. py_OutRef val = py_retval();
  556. py_Ref p = py_newtuple(val, length);
  557. for(int i = length - 1; i >= 0; i--) {
  558. p[i] = *py_peek(-1);
  559. py_pop();
  560. }
  561. py_push(val);
  562. break;
  563. }
  564. case PKL_BUILD_DICT: {
  565. int length = pkl__read_int(&p);
  566. py_OutRef val = py_pushtmp();
  567. py_newdict(val);
  568. py_StackRef begin = py_peek(-1) - 2 * length;
  569. py_StackRef end = py_peek(-1);
  570. for(py_StackRef i = begin; i < end; i += 2) {
  571. py_StackRef k = i;
  572. py_StackRef v = i + 1;
  573. bool ok = py_dict_setitem(val, k, v);
  574. if(!ok) return false;
  575. }
  576. py_assign(py_retval(), val);
  577. py_shrink(2 * length + 1);
  578. py_push(py_retval());
  579. break;
  580. }
  581. case PKL_VEC2: {
  582. c11_vec2 val;
  583. UNALIGNED_READ(&val, p);
  584. py_newvec2(py_pushtmp(), val);
  585. break;
  586. }
  587. case PKL_VEC3: {
  588. c11_vec3 val;
  589. UNALIGNED_READ(&val, p);
  590. py_newvec3(py_pushtmp(), val);
  591. break;
  592. }
  593. case PKL_VEC2I: {
  594. c11_vec2i val;
  595. val.x = pkl__read_int(&p);
  596. val.y = pkl__read_int(&p);
  597. py_newvec2i(py_pushtmp(), val);
  598. break;
  599. }
  600. case PKL_VEC3I: {
  601. c11_vec3i val;
  602. val.x = pkl__read_int(&p);
  603. val.y = pkl__read_int(&p);
  604. val.z = pkl__read_int(&p);
  605. py_newvec3i(py_pushtmp(), val);
  606. break;
  607. }
  608. case PKL_TYPE: {
  609. py_Type type = (py_Type)pkl__read_int(&p);
  610. type = pkl__fix_type(type, type_mapping);
  611. py_push(py_tpobject(type));
  612. break;
  613. }
  614. case PKL_ARRAY2D: {
  615. int n_cols = pkl__read_int(&p);
  616. int n_rows = pkl__read_int(&p);
  617. c11_array2d* arr = py_newarray2d(py_pushtmp(), n_cols, n_rows);
  618. int total_size = arr->header.numel * sizeof(py_TValue);
  619. memcpy(arr->data, p, total_size);
  620. for(int i = 0; i < arr->header.numel; i++) {
  621. arr->data[i].type = pkl__fix_type(arr->data[i].type, type_mapping);
  622. }
  623. p += total_size;
  624. break;
  625. }
  626. case PKL_TVALUE: {
  627. py_TValue* tmp = py_pushtmp();
  628. memcpy(tmp, p, sizeof(py_TValue));
  629. tmp->type = pkl__fix_type(tmp->type, type_mapping);
  630. p += sizeof(py_TValue);
  631. break;
  632. }
  633. case PKL_CALL: {
  634. int argc = pkl__read_int(&p);
  635. if(!py_vectorcall(argc, 0)) return false;
  636. py_push(py_retval());
  637. break;
  638. }
  639. case PKL_OBJECT: {
  640. py_Type type = (py_Type)pkl__read_int(&p);
  641. type = pkl__fix_type(type, type_mapping);
  642. py_newobject(py_retval(), type, -1, 0);
  643. NameDict* dict = PyObject__dict(py_retval()->_obj);
  644. int dict_length = pkl__read_int(&p);
  645. for(int i = 0; i < dict_length; i++) {
  646. py_StackRef value = py_peek(-1);
  647. c11_sv field = {(const char*)p, strlen((const char*)p)};
  648. NameDict__set(dict, py_namev(field), *value);
  649. py_pop();
  650. p += field.size + 1;
  651. }
  652. py_push(py_retval());
  653. break;
  654. }
  655. case PKL_EOF: {
  656. // [memo, obj]
  657. if(py_peek(0) - p0 != 2) return ValueError("invalid pickle data");
  658. py_assign(py_retval(), py_peek(-1));
  659. py_shrink(2);
  660. return true;
  661. }
  662. default: c11__unreachable();
  663. }
  664. }
  665. c11__unreachable();
  666. }
  667. static bool PickleObject__py_submit(PickleObject* self, py_OutRef out) {
  668. c11_sbuf cleartext;
  669. c11_sbuf__ctor(&cleartext);
  670. // line 1: type mappping
  671. for(py_Type type = 0; type < self->used_types_length; type++) {
  672. if(self->used_types[type]) {
  673. c11_sbuf__write_int(&cleartext, type);
  674. c11_sbuf__write_char(&cleartext, '(');
  675. c11_sbuf__write_type_path(&cleartext, type);
  676. c11_sbuf__write_char(&cleartext, ')');
  677. }
  678. }
  679. c11_sbuf__write_char(&cleartext, '\n');
  680. // line 2: memo length
  681. c11_sbuf__write_int(&cleartext, self->memo.length);
  682. c11_sbuf__write_char(&cleartext, '\n');
  683. // -------------------------------------------------- //
  684. c11_string* header = c11_sbuf__submit(&cleartext);
  685. int total_size = header->size + self->codes.length;
  686. unsigned char* p = py_newbytes(py_retval(), total_size);
  687. memcpy(p, header->data, header->size);
  688. memcpy(p + header->size, self->codes.data, self->codes.length);
  689. c11_string__delete(header);
  690. PickleObject__dtor(self);
  691. return true;
  692. }
  693. #undef UNALIGNED_READ