Skip to content

Commit ae8d867

Browse files
dpgeorgedpgeorge
committed
py: Add iter_buf to getiter type method.
Allows to iterate over the following without allocating on the heap: - tuple - list - string, bytes - bytearray, array - dict (not dict.keys, dict.values, dict.items) - set, frozenset Allows to call the following without heap memory: - all, any, min, max, sum TODO: still need to allocate stack memory in bytecode for iter_buf.
1 parent 101886f commit ae8d867

36 files changed

Lines changed: 162 additions & 106 deletions

‎cc3200/mods/pybuart.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -664,7 +664,7 @@ const mp_obj_type_t pyb_uart_type = {
664664
.name = MP_QSTR_UART,
665665
.print = pyb_uart_print,
666666
.make_new = pyb_uart_make_new,
667-
.getiter = mp_identity,
667+
.getiter = mp_identity_getiter,
668668
.iternext = mp_stream_unbuffered_iter,
669669
.protocol = &uart_stream_p,
670670
.locals_dict = (mp_obj_t)&pyb_uart_locals_dict,

‎esp8266/machine_uart.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ const mp_obj_type_t pyb_uart_type = {
285285
.name = MP_QSTR_UART,
286286
.print = pyb_uart_print,
287287
.make_new = pyb_uart_make_new,
288-
.getiter = mp_identity,
288+
.getiter = mp_identity_getiter,
289289
.iternext = mp_stream_unbuffered_iter,
290290
.protocol = &uart_stream_p,
291291
.locals_dict = (mp_obj_dict_t*)&pyb_uart_locals_dict,

‎extmod/modbtree.c‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,8 @@ STATIC mp_obj_t btree_items(size_t n_args, const mp_obj_t *args) {
184184
}
185185
STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_items_obj, 1, 4, btree_items);
186186

187-
STATIC mp_obj_t btree_getiter(mp_obj_t self_in) {
187+
STATIC mp_obj_t btree_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
188+
(void)iter_buf;
188189
mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
189190
if (self->next_flags != 0) {
190191
// If we're called immediately after keys(), values(), or items(),

‎extmod/vfs_fat_file.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,7 @@ const mp_obj_type_t mp_type_fileio = {
264264
.name = MP_QSTR_FileIO,
265265
.print = file_obj_print,
266266
.make_new = file_obj_make_new,
267-
.getiter = mp_identity,
267+
.getiter = mp_identity_getiter,
268268
.iternext = mp_stream_unbuffered_iter,
269269
.protocol = &fileio_stream_p,
270270
.locals_dict = (mp_obj_dict_t*)&rawfile_locals_dict,
@@ -283,7 +283,7 @@ const mp_obj_type_t mp_type_textio = {
283283
.name = MP_QSTR_TextIOWrapper,
284284
.print = file_obj_print,
285285
.make_new = file_obj_make_new,
286-
.getiter = mp_identity,
286+
.getiter = mp_identity_getiter,
287287
.iternext = mp_stream_unbuffered_iter,
288288
.protocol = &textio_stream_p,
289289
.locals_dict = (mp_obj_dict_t*)&rawfile_locals_dict,

‎py/emitnative.c‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1806,6 +1806,7 @@ STATIC void emit_native_get_iter(emit_t *emit) {
18061806
vtype_kind_t vtype;
18071807
emit_pre_pop_reg(emit, &vtype, REG_ARG_1);
18081808
assert(vtype == VTYPE_PYOBJ);
1809+
assert(0); // TODO allocate memory for iter_buf
18091810
emit_call(emit, MP_F_GETITER);
18101811
emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET);
18111812
}

‎py/modbuiltins.c‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,8 @@ STATIC mp_obj_t mp_builtin_abs(mp_obj_t o_in) {
117117
MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_abs_obj, mp_builtin_abs);
118118

119119
STATIC mp_obj_t mp_builtin_all(mp_obj_t o_in) {
120-
mp_obj_t iterable = mp_getiter(o_in);
120+
mp_obj_iter_buf_t iter_buf;
121+
mp_obj_t iterable = mp_getiter(o_in, &iter_buf);
121122
mp_obj_t item;
122123
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
123124
if (!mp_obj_is_true(item)) {
@@ -129,7 +130,8 @@ STATIC mp_obj_t mp_builtin_all(mp_obj_t o_in) {
129130
MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_all_obj, mp_builtin_all);
130131

131132
STATIC mp_obj_t mp_builtin_any(mp_obj_t o_in) {
132-
mp_obj_t iterable = mp_getiter(o_in);
133+
mp_obj_iter_buf_t iter_buf;
134+
mp_obj_t iterable = mp_getiter(o_in, &iter_buf);
133135
mp_obj_t item;
134136
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
135137
if (mp_obj_is_true(item)) {
@@ -258,7 +260,7 @@ STATIC mp_obj_t mp_builtin_hex(mp_obj_t o_in) {
258260
MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_hex_obj, mp_builtin_hex);
259261

260262
STATIC mp_obj_t mp_builtin_iter(mp_obj_t o_in) {
261-
return mp_getiter(o_in);
263+
return mp_getiter(o_in, NULL);
262264
}
263265
MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_iter_obj, mp_builtin_iter);
264266

@@ -270,7 +272,8 @@ STATIC mp_obj_t mp_builtin_min_max(size_t n_args, const mp_obj_t *args, mp_map_t
270272
mp_obj_t key_fn = key_elem == NULL ? MP_OBJ_NULL : key_elem->value;
271273
if (n_args == 1) {
272274
// given an iterable
273-
mp_obj_t iterable = mp_getiter(args[0]);
275+
mp_obj_iter_buf_t iter_buf;
276+
mp_obj_t iterable = mp_getiter(args[0], &iter_buf);
274277
mp_obj_t best_key = MP_OBJ_NULL;
275278
mp_obj_t best_obj = MP_OBJ_NULL;
276279
mp_obj_t item;
@@ -495,7 +498,8 @@ STATIC mp_obj_t mp_builtin_sum(size_t n_args, const mp_obj_t *args) {
495498
case 1: value = MP_OBJ_NEW_SMALL_INT(0); break;
496499
default: value = args[1]; break;
497500
}
498-
mp_obj_t iterable = mp_getiter(args[0]);
501+
mp_obj_iter_buf_t iter_buf;
502+
mp_obj_t iterable = mp_getiter(args[0], &iter_buf);
499503
mp_obj_t item;
500504
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
501505
value = mp_binary_op(MP_BINARY_OP_ADD, value, item);

‎py/obj.c‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,11 @@ mp_obj_t mp_identity(mp_obj_t self) {
481481
}
482482
MP_DEFINE_CONST_FUN_OBJ_1(mp_identity_obj, mp_identity);
483483

484+
mp_obj_t mp_identity_getiter(mp_obj_t self, mp_obj_iter_buf_t *iter_buf) {
485+
(void)iter_buf;
486+
return self;
487+
}
488+
484489
bool mp_get_buffer(mp_obj_t obj, mp_buffer_info_t *bufinfo, mp_uint_t flags) {
485490
mp_obj_type_t *type = mp_obj_get_type(obj);
486491
if (type->buffer_p.get_buffer == NULL) {

‎py/obj.h‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -417,13 +417,19 @@ typedef enum {
417417
PRINT_EXC_SUBCLASS = 0x80, // Internal flag for printing exception subclasses
418418
} mp_print_kind_t;
419419

420+
typedef struct _mp_obj_iter_buf_t {
421+
mp_obj_base_t base;
422+
mp_obj_t buf[3];
423+
} mp_obj_iter_buf_t;
424+
420425
typedef void (*mp_print_fun_t)(const mp_print_t *print, mp_obj_t o, mp_print_kind_t kind);
421426
typedef mp_obj_t (*mp_make_new_fun_t)(const mp_obj_type_t *type, size_t n_args, size_t n_kw, const mp_obj_t *args);
422427
typedef mp_obj_t (*mp_call_fun_t)(mp_obj_t fun, size_t n_args, size_t n_kw, const mp_obj_t *args);
423428
typedef mp_obj_t (*mp_unary_op_fun_t)(mp_uint_t op, mp_obj_t);
424429
typedef mp_obj_t (*mp_binary_op_fun_t)(mp_uint_t op, mp_obj_t, mp_obj_t);
425430
typedef void (*mp_attr_fun_t)(mp_obj_t self_in, qstr attr, mp_obj_t *dest);
426431
typedef mp_obj_t (*mp_subscr_fun_t)(mp_obj_t self_in, mp_obj_t index, mp_obj_t value);
432+
typedef mp_obj_t (*mp_getiter_fun_t)(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf);
427433

428434
// Buffer protocol
429435
typedef struct _mp_buffer_info_t {
@@ -486,7 +492,11 @@ struct _mp_obj_type_t {
486492
// value=MP_OBJ_NULL means delete, value=MP_OBJ_SENTINEL means load, else store
487493
// can return MP_OBJ_NULL if op not supported
488494

489-
mp_fun_1_t getiter; // corresponds to __iter__ special method
495+
// corresponds to __iter__ special method
496+
// can use given mp_obj_iter_buf_t to store iterator
497+
// otherwise can return a pointer to an object on the heap
498+
mp_getiter_fun_t getiter;
499+
490500
mp_fun_1_t iternext; // may return MP_OBJ_STOP_ITERATION as an optimisation instead of raising StopIteration() (with no args)
491501

492502
mp_buffer_p_t buffer_p;
@@ -637,7 +647,7 @@ mp_obj_t mp_obj_new_set(size_t n_args, mp_obj_t *items);
637647
mp_obj_t mp_obj_new_slice(mp_obj_t start, mp_obj_t stop, mp_obj_t step);
638648
mp_obj_t mp_obj_new_super(mp_obj_t type, mp_obj_t obj);
639649
mp_obj_t mp_obj_new_bound_meth(mp_obj_t meth, mp_obj_t self);
640-
mp_obj_t mp_obj_new_getitem_iter(mp_obj_t *args);
650+
mp_obj_t mp_obj_new_getitem_iter(mp_obj_t *args, mp_obj_iter_buf_t *iter_buf);
641651
mp_obj_t mp_obj_new_module(qstr module_name);
642652
mp_obj_t mp_obj_new_memoryview(byte typecode, size_t nitems, void *items);
643653

@@ -775,6 +785,7 @@ qstr mp_obj_code_get_name(const byte *code_info);
775785

776786
mp_obj_t mp_identity(mp_obj_t self);
777787
MP_DECLARE_CONST_FUN_OBJ_1(mp_identity_obj);
788+
mp_obj_t mp_identity_getiter(mp_obj_t self, mp_obj_iter_buf_t *iter_buf);
778789

779790
// module
780791
typedef struct _mp_obj_module_t {

‎py/objarray.c‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@
5959
#define TYPECODE_MASK (~(size_t)0)
6060
#endif
6161

62-
STATIC mp_obj_t array_iterator_new(mp_obj_t array_in);
62+
STATIC mp_obj_t array_iterator_new(mp_obj_t array_in, mp_obj_iter_buf_t *iter_buf);
6363
STATIC mp_obj_t array_append(mp_obj_t self_in, mp_obj_t arg);
6464
STATIC mp_obj_t array_extend(mp_obj_t self_in, mp_obj_t arg_in);
6565
STATIC mp_int_t array_get_buffer(mp_obj_t o_in, mp_buffer_info_t *bufinfo, mp_uint_t flags);
@@ -141,7 +141,8 @@ STATIC mp_obj_t array_construct(char typecode, mp_obj_t initializer) {
141141

142142
mp_obj_array_t *array = array_new(typecode, len);
143143

144-
mp_obj_t iterable = mp_getiter(initializer);
144+
mp_obj_iter_buf_t iter_buf;
145+
mp_obj_t iterable = mp_getiter(initializer, &iter_buf);
145146
mp_obj_t item;
146147
size_t i = 0;
147148
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
@@ -608,15 +609,18 @@ STATIC mp_obj_t array_it_iternext(mp_obj_t self_in) {
608609
STATIC const mp_obj_type_t array_it_type = {
609610
{ &mp_type_type },
610611
.name = MP_QSTR_iterator,
611-
.getiter = mp_identity,
612+
.getiter = mp_identity_getiter,
612613
.iternext = array_it_iternext,
613614
};
614615

615-
STATIC mp_obj_t array_iterator_new(mp_obj_t array_in) {
616+
STATIC mp_obj_t array_iterator_new(mp_obj_t array_in, mp_obj_iter_buf_t *iter_buf) {
617+
assert(sizeof(mp_obj_array_t) <= sizeof(mp_obj_iter_buf_t));
616618
mp_obj_array_t *array = MP_OBJ_TO_PTR(array_in);
617-
mp_obj_array_it_t *o = m_new0(mp_obj_array_it_t, 1);
619+
mp_obj_array_it_t *o = (mp_obj_array_it_t*)iter_buf;
618620
o->base.type = &array_it_type;
619621
o->array = array;
622+
o->offset = 0;
623+
o->cur = 0;
620624
#if MICROPY_PY_BUILTINS_MEMORYVIEW
621625
if (array->base.type == &mp_type_memoryview) {
622626
o->offset = array->free;

‎py/objdict.c‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -210,8 +210,9 @@ STATIC mp_obj_t dict_it_iternext(mp_obj_t self_in) {
210210
}
211211
}
212212

213-
STATIC mp_obj_t dict_getiter(mp_obj_t self_in) {
214-
mp_obj_dict_it_t *o = m_new_obj(mp_obj_dict_it_t);
213+
STATIC mp_obj_t dict_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
214+
assert(sizeof(mp_obj_dict_it_t) <= sizeof(mp_obj_iter_buf_t));
215+
mp_obj_dict_it_t *o = (mp_obj_dict_it_t*)iter_buf;
215216
o->base.type = &mp_type_polymorph_iter;
216217
o->iternext = dict_it_iternext;
217218
o->dict = self_in;
@@ -249,7 +250,8 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_1(dict_copy_obj, dict_copy);
249250

250251
// this is a classmethod
251252
STATIC mp_obj_t dict_fromkeys(size_t n_args, const mp_obj_t *args) {
252-
mp_obj_t iter = mp_getiter(args[1]);
253+
mp_obj_iter_buf_t iter_buf;
254+
mp_obj_t iter = mp_getiter(args[1], &iter_buf);
253255
mp_obj_t value = mp_const_none;
254256
mp_obj_t next = MP_OBJ_NULL;
255257

@@ -375,10 +377,12 @@ STATIC mp_obj_t dict_update(size_t n_args, const mp_obj_t *args, mp_map_t *kwarg
375377
}
376378
} else {
377379
// update from a generic iterable of pairs
378-
mp_obj_t iter = mp_getiter(args[1]);
380+
mp_obj_iter_buf_t iter_buf;
381+
mp_obj_t iter = mp_getiter(args[1], &iter_buf);
379382
mp_obj_t next = MP_OBJ_NULL;
380383
while ((next = mp_iternext(iter)) != MP_OBJ_STOP_ITERATION) {
381-
mp_obj_t inneriter = mp_getiter(next);
384+
mp_obj_iter_buf_t inner_iter_buf;
385+
mp_obj_t inneriter = mp_getiter(next, &inner_iter_buf);
382386
mp_obj_t key = mp_iternext(inneriter);
383387
mp_obj_t value = mp_iternext(inneriter);
384388
mp_obj_t stop = mp_iternext(inneriter);
@@ -457,14 +461,15 @@ STATIC mp_obj_t dict_view_it_iternext(mp_obj_t self_in) {
457461
STATIC const mp_obj_type_t dict_view_it_type = {
458462
{ &mp_type_type },
459463
.name = MP_QSTR_iterator,
460-
.getiter = mp_identity,
464+
.getiter = mp_identity_getiter,
461465
.iternext = dict_view_it_iternext,
462466
};
463467

464-
STATIC mp_obj_t dict_view_getiter(mp_obj_t view_in) {
468+
STATIC mp_obj_t dict_view_getiter(mp_obj_t view_in, mp_obj_iter_buf_t *iter_buf) {
469+
assert(sizeof(mp_obj_dict_view_it_t) <= sizeof(mp_obj_iter_buf_t));
465470
mp_check_self(MP_OBJ_IS_TYPE(view_in, &dict_view_type));
466471
mp_obj_dict_view_t *view = MP_OBJ_TO_PTR(view_in);
467-
mp_obj_dict_view_it_t *o = m_new_obj(mp_obj_dict_view_it_t);
472+
mp_obj_dict_view_it_t *o = (mp_obj_dict_view_it_t*)iter_buf;
468473
o->base.type = &dict_view_it_type;
469474
o->kind = view->kind;
470475
o->dict = view->dict;
@@ -479,7 +484,8 @@ STATIC void dict_view_print(const mp_print_t *print, mp_obj_t self_in, mp_print_
479484
bool first = true;
480485
mp_print_str(print, mp_dict_view_names[self->kind]);
481486
mp_print_str(print, "([");
482-
mp_obj_t self_iter = dict_view_getiter(self_in);
487+
mp_obj_iter_buf_t iter_buf;
488+
mp_obj_t self_iter = dict_view_getiter(self_in, &iter_buf);
483489
mp_obj_t next = MP_OBJ_NULL;
484490
while ((next = dict_view_it_iternext(self_iter)) != MP_OBJ_STOP_ITERATION) {
485491
if (!first) {

0 commit comments

Comments
 (0)