Skip to content

Commit aacd618

Browse files
dpgeorgedpgeorge
committed
py/runtime: Don't allocate iter buf for user-defined types.
A user-defined type that defines __iter__ doesn't need any memory to be pre-allocated for its iterator (because it can't use such memory). So optimise for this case by not allocating the iter-buf.
1 parent 11b4524 commit aacd618

3 files changed

Lines changed: 19 additions & 7 deletions

File tree

‎py/objtype.c‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -882,7 +882,8 @@ mp_obj_t mp_obj_instance_call(mp_obj_t self_in, size_t n_args, size_t n_kw, cons
882882
return mp_call_method_self_n_kw(member[0], member[1], n_args, n_kw, args);
883883
}
884884

885-
STATIC mp_obj_t instance_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
885+
// Note that iter_buf may be NULL, and needs to be allocated if needed
886+
mp_obj_t mp_obj_instance_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
886887
mp_obj_instance_t *self = MP_OBJ_TO_PTR(self_in);
887888
mp_obj_t member[2] = {MP_OBJ_NULL};
888889
struct class_lookup_data lookup = {
@@ -897,6 +898,9 @@ STATIC mp_obj_t instance_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf)
897898
return MP_OBJ_NULL;
898899
} else if (member[0] == MP_OBJ_SENTINEL) {
899900
mp_obj_type_t *type = mp_obj_get_type(self->subobj[0]);
901+
if (iter_buf == NULL) {
902+
iter_buf = m_new_obj(mp_obj_iter_buf_t);
903+
}
900904
return type->getiter(self->subobj[0], iter_buf);
901905
} else {
902906
return mp_call_method_n_kw(0, 0, member);
@@ -1136,7 +1140,7 @@ mp_obj_t mp_obj_new_type(qstr name, mp_obj_t bases_tuple, mp_obj_t locals_dict)
11361140
o->binary_op = instance_binary_op;
11371141
o->attr = mp_obj_instance_attr;
11381142
o->subscr = instance_subscr;
1139-
o->getiter = instance_getiter;
1143+
o->getiter = mp_obj_instance_getiter;
11401144
//o->iternext = ; not implemented
11411145
o->buffer_p.get_buffer = instance_get_buffer;
11421146

‎py/objtype.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,4 +51,7 @@ mp_obj_t mp_obj_instance_call(mp_obj_t self_in, size_t n_args, size_t n_kw, cons
5151
// this needs to be exposed for the above macros to work correctly
5252
mp_obj_t mp_obj_instance_make_new(const mp_obj_type_t *self_in, size_t n_args, size_t n_kw, const mp_obj_t *args);
5353

54+
// this needs to be exposed for mp_getiter
55+
mp_obj_t mp_obj_instance_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf);
56+
5457
#endif // MICROPY_INCLUDED_PY_OBJTYPE_H

‎py/runtime.c‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
#include "py/objstr.h"
3535
#include "py/objtuple.h"
3636
#include "py/objlist.h"
37+
#include "py/objtype.h"
3738
#include "py/objmodule.h"
3839
#include "py/objgenerator.h"
3940
#include "py/smallint.h"
@@ -1165,13 +1166,13 @@ mp_obj_t mp_getiter(mp_obj_t o_in, mp_obj_iter_buf_t *iter_buf) {
11651166
return o_in;
11661167
}
11671168

1168-
// if caller did not provide a buffer then allocate one on the heap
1169-
if (iter_buf == NULL) {
1170-
iter_buf = m_new_obj(mp_obj_iter_buf_t);
1171-
}
1172-
11731169
// check for native getiter (corresponds to __iter__)
11741170
if (type->getiter != NULL) {
1171+
if (iter_buf == NULL && type->getiter != mp_obj_instance_getiter) {
1172+
// if caller did not provide a buffer then allocate one on the heap
1173+
// mp_obj_instance_getiter is special, it will allocate only if needed
1174+
iter_buf = m_new_obj(mp_obj_iter_buf_t);
1175+
}
11751176
mp_obj_t iter = type->getiter(o_in, iter_buf);
11761177
if (iter != MP_OBJ_NULL) {
11771178
return iter;
@@ -1183,6 +1184,10 @@ mp_obj_t mp_getiter(mp_obj_t o_in, mp_obj_iter_buf_t *iter_buf) {
11831184
mp_load_method_maybe(o_in, MP_QSTR___getitem__, dest);
11841185
if (dest[0] != MP_OBJ_NULL) {
11851186
// __getitem__ exists, create and return an iterator
1187+
if (iter_buf == NULL) {
1188+
// if caller did not provide a buffer then allocate one on the heap
1189+
iter_buf = m_new_obj(mp_obj_iter_buf_t);
1190+
}
11861191
return mp_obj_new_getitem_iter(dest, iter_buf);
11871192
}
11881193

0 commit comments

Comments
 (0)