1 files changed, 14 insertions, 7 deletions

diff --git a/py/objrange.c b/py/objrange.c
index 79459316b1..8c4e14f49c 100644
--- a/py/objrange.c
+++ b/py/objrange.c
@@ -55,12 +55,13 @@ STATIC mp_obj_t range_it_iternext(mp_obj_t o_in) {
 STATIC const mp_obj_type_t range_it_type = {
     { &mp_type_type },
     .name = MP_QSTR_iterator,
-    .getiter = mp_identity,
+    .getiter = mp_identity_getiter,
     .iternext = range_it_iternext,
 };
 
-STATIC mp_obj_t mp_obj_new_range_iterator(mp_int_t cur, mp_int_t stop, mp_int_t step) {
-    mp_obj_range_it_t *o = m_new_obj(mp_obj_range_it_t);
+STATIC mp_obj_t mp_obj_new_range_iterator(mp_int_t cur, mp_int_t stop, mp_int_t step, mp_obj_iter_buf_t *iter_buf) {
+    assert(sizeof(mp_obj_range_it_t) <= sizeof(mp_obj_iter_buf_t));
+    mp_obj_range_it_t *o = (mp_obj_range_it_t*)iter_buf;
     o->base.type = &range_it_type;
     o->cur = cur;
     o->stop = stop;
@@ -104,8 +105,10 @@ STATIC mp_obj_t range_make_new(const mp_obj_type_t *type, size_t n_args, size_t
         o->start = mp_obj_get_int(args[0]);
         o->stop = mp_obj_get_int(args[1]);
         if (n_args == 3) {
-            // TODO check step is non-zero
             o->step = mp_obj_get_int(args[2]);
+            if (o->step == 0) {
+                mp_raise_ValueError("zero step");
+            }
         }
     }
 
@@ -151,19 +154,23 @@ STATIC mp_obj_t range_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
             o->start = self->start + slice.start * self->step;
             o->stop = self->start + slice.stop * self->step;
             o->step = slice.step * self->step;
+            if (slice.step < 0) {
+                // Negative slice steps have inclusive stop, so adjust for exclusive
+                o->stop -= self->step;
+            }
             return MP_OBJ_FROM_PTR(o);
         }
 #endif
-        uint index_val = mp_get_index(self->base.type, len, index, false);
+        size_t index_val = mp_get_index(self->base.type, len, index, false);
         return MP_OBJ_NEW_SMALL_INT(self->start + index_val * self->step);
     } else {
         return MP_OBJ_NULL; // op not supported
     }
 }
 
-STATIC mp_obj_t range_getiter(mp_obj_t o_in) {
+STATIC mp_obj_t range_getiter(mp_obj_t o_in, mp_obj_iter_buf_t *iter_buf) {
     mp_obj_range_t *o = MP_OBJ_TO_PTR(o_in);
-    return mp_obj_new_range_iterator(o->start, o->stop, o->step);
+    return mp_obj_new_range_iterator(o->start, o->stop, o->step, iter_buf);
 }