From aacd61893952af24049fd64f3735551ff9fc14ea Mon Sep 17 00:00:00 2001 From: Damien George Date: Wed, 13 Nov 2019 21:05:34 +1100 Subject: [PATCH] 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. --- py/objtype.c | 8 ++++++-- py/objtype.h | 3 +++ py/runtime.c | 15 ++++++++++----- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/py/objtype.c b/py/objtype.c index bf089dc49..5d4dd7d2f 100644 --- a/py/objtype.c +++ b/py/objtype.c @@ -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 return mp_call_method_self_n_kw(member[0], member[1], n_args, n_kw, args); } -STATIC mp_obj_t instance_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) { +// Note that iter_buf may be NULL, and needs to be allocated if needed +mp_obj_t mp_obj_instance_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) { mp_obj_instance_t *self = MP_OBJ_TO_PTR(self_in); mp_obj_t member[2] = {MP_OBJ_NULL}; 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) return MP_OBJ_NULL; } else if (member[0] == MP_OBJ_SENTINEL) { mp_obj_type_t *type = mp_obj_get_type(self->subobj[0]); + if (iter_buf == NULL) { + iter_buf = m_new_obj(mp_obj_iter_buf_t); + } return type->getiter(self->subobj[0], iter_buf); } else { 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) o->binary_op = instance_binary_op; o->attr = mp_obj_instance_attr; o->subscr = instance_subscr; - o->getiter = instance_getiter; + o->getiter = mp_obj_instance_getiter; //o->iternext = ; not implemented o->buffer_p.get_buffer = instance_get_buffer; diff --git a/py/objtype.h b/py/objtype.h index 3fc8c6e1b..2c613b904 100644 --- a/py/objtype.h +++ b/py/objtype.h @@ -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 // this needs to be exposed for the above macros to work correctly 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); +// this needs to be exposed for mp_getiter +mp_obj_t mp_obj_instance_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf); + #endif // MICROPY_INCLUDED_PY_OBJTYPE_H diff --git a/py/runtime.c b/py/runtime.c index cf4fc5d38..d5511236d 100644 --- a/py/runtime.c +++ b/py/runtime.c @@ -34,6 +34,7 @@ #include "py/objstr.h" #include "py/objtuple.h" #include "py/objlist.h" +#include "py/objtype.h" #include "py/objmodule.h" #include "py/objgenerator.h" #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) { return o_in; } - // if caller did not provide a buffer then allocate one on the heap - if (iter_buf == NULL) { - iter_buf = m_new_obj(mp_obj_iter_buf_t); - } - // check for native getiter (corresponds to __iter__) if (type->getiter != NULL) { + if (iter_buf == NULL && type->getiter != mp_obj_instance_getiter) { + // if caller did not provide a buffer then allocate one on the heap + // mp_obj_instance_getiter is special, it will allocate only if needed + iter_buf = m_new_obj(mp_obj_iter_buf_t); + } mp_obj_t iter = type->getiter(o_in, iter_buf); if (iter != MP_OBJ_NULL) { return iter; @@ -1183,6 +1184,10 @@ mp_obj_t mp_getiter(mp_obj_t o_in, mp_obj_iter_buf_t *iter_buf) { mp_load_method_maybe(o_in, MP_QSTR___getitem__, dest); if (dest[0] != MP_OBJ_NULL) { // __getitem__ exists, create and return an iterator + if (iter_buf == NULL) { + // if caller did not provide a buffer then allocate one on the heap + iter_buf = m_new_obj(mp_obj_iter_buf_t); + } return mp_obj_new_getitem_iter(dest, iter_buf); }