micropython · andrewleech · Mar 21, 2019
diff --git a/py/mpstate.h b/py/mpstate.h
@@ -209,7 +209,8 @@ typedef struct _mp_state_vm_t {
 
     #if MICROPY_ENABLE_SCHEDULER
     volatile int16_t sched_state;
-    uint16_t sched_sp;
+    uint8_t sched_len;
+    uint8_t sched_idx;
     #endif
 
     #if MICROPY_PY_THREAD_GIL

diff --git a/py/runtime.c b/py/runtime.c
@@ -63,7 +63,8 @@ void mp_init(void) {
     MP_STATE_VM(mp_pending_exception) = MP_OBJ_NULL;
     #if MICROPY_ENABLE_SCHEDULER
     MP_STATE_VM(sched_state) = MP_SCHED_IDLE;
-    MP_STATE_VM(sched_sp) = 0;
+    MP_STATE_VM(sched_idx) = 0;
+    MP_STATE_VM(sched_len) = 0;
     #endif
 
 #if MICROPY_ENABLE_EMERGENCY_EXCEPTION_BUF

diff --git a/py/runtime.h b/py/runtime.h
@@ -70,7 +70,7 @@ void mp_handle_pending_tail(mp_uint_t atomic_state);
 #if MICROPY_ENABLE_SCHEDULER
 void mp_sched_lock(void);
 void mp_sched_unlock(void);
-static inline unsigned int mp_sched_num_pending(void) { return MP_STATE_VM(sched_sp); }
+static inline unsigned int mp_sched_num_pending(void) { return MP_STATE_VM(sched_len); }
 bool mp_sched_schedule(mp_obj_t function, mp_obj_t arg);
 #endif
 

diff --git a/py/scheduler.c b/py/scheduler.c
@@ -30,6 +30,18 @@
 
 #if MICROPY_ENABLE_SCHEDULER
 
+#define IDX_MASK(i) ((i) & (MICROPY_SCHEDULER_DEPTH - 1))
+
+static inline bool mp_sched_full(void) {
+    MP_STATIC_ASSERT((IDX_MASK(MICROPY_SCHEDULER_DEPTH) == 0)); // MICROPY_SCHEDULER_DEPTH must be a power of 2
+
+    return mp_sched_num_pending() == MICROPY_SCHEDULER_DEPTH;
+}
+
+static inline bool mp_sched_empty(void) {
+    return mp_sched_num_pending() == 0;
+}
+
 // A variant of this is inlined in the VM at the pending exception check
 void mp_handle_pending(void) {
     if (MP_STATE_VM(sched_state) == MP_SCHED_PENDING) {
@@ -51,8 +63,10 @@ void mp_handle_pending(void) {
 // or by the VM's inlined version of that function.
 void mp_handle_pending_tail(mp_uint_t atomic_state) {
     MP_STATE_VM(sched_state) = MP_SCHED_LOCKED;
-    if (MP_STATE_VM(sched_sp) > 0) {
-        mp_sched_item_t item = MP_STATE_VM(sched_stack)[--MP_STATE_VM(sched_sp)];
+    if (!mp_sched_empty()) {
+        mp_sched_item_t item = MP_STATE_VM(sched_stack)[MP_STATE_VM(sched_idx)];
+        MP_STATE_VM(sched_idx) = IDX_MASK(MP_STATE_VM(sched_idx)+1);
+        MP_STATE_VM(sched_len)--;
         MICROPY_END_ATOMIC_SECTION(atomic_state);
         mp_call_function_1_protected(item.func, item.arg);
     } else {
@@ -87,13 +101,13 @@ void mp_sched_unlock(void) {
 bool mp_sched_schedule(mp_obj_t function, mp_obj_t arg) {
     mp_uint_t atomic_state = MICROPY_BEGIN_ATOMIC_SECTION();
     bool ret;
-    if (MP_STATE_VM(sched_sp) < MICROPY_SCHEDULER_DEPTH) {
+    if (!mp_sched_full()) {
         if (MP_STATE_VM(sched_state) == MP_SCHED_IDLE) {
             MP_STATE_VM(sched_state) = MP_SCHED_PENDING;
         }
-        MP_STATE_VM(sched_stack)[MP_STATE_VM(sched_sp)].func = function;
-        MP_STATE_VM(sched_stack)[MP_STATE_VM(sched_sp)].arg = arg;
-        ++MP_STATE_VM(sched_sp);
+        uint8_t iput = IDX_MASK(MP_STATE_VM(sched_idx) + MP_STATE_VM(sched_len)++);
+        MP_STATE_VM(sched_stack)[iput].func = function;
+        MP_STATE_VM(sched_stack)[iput].arg = arg;
         ret = true;
     } else {
         // schedule stack is full

diff --git a/tests/unix/extra_coverage.py.exp b/tests/unix/extra_coverage.py.exp
@@ -70,10 +70,10 @@ sched(2)=1
 sched(3)=1
 sched(4)=0
 unlocked
-3
-2
-1
 0
+1
+2
+3
 0123456789 b'0123456789'
 7300
 7300