1 files changed, 23 insertions, 25 deletions

diff --git a/src/core/ao_task.c b/src/core/ao_task.c
index a11979f0..0411fbdd 100644
--- a/src/core/ao_task.c
+++ b/src/core/ao_task.c
@@ -114,6 +114,8 @@ ao_task_validate_alarm_queue(void)
 #define ao_task_validate_alarm_queue()
 #endif
 
+uint16_t	ao_task_alarm_tick;
+
 static void
 ao_task_to_alarm_queue(struct ao_task *task)
 {
@@ -126,6 +128,7 @@ ao_task_to_alarm_queue(struct ao_task *task)
 		}
 	}
 	ao_list_append(&task->alarm_queue, &alarm_queue);
+	ao_task_alarm_tick = ao_list_first_entry(&alarm_queue, struct ao_task, alarm_queue)->alarm;
 	ao_task_validate_alarm_queue();
 }
 
@@ -133,6 +136,10 @@ static void
 ao_task_from_alarm_queue(struct ao_task *task)
 {
 	ao_list_del(&task->alarm_queue);
+	if (ao_list_is_empty(&alarm_queue))
+		ao_task_alarm_tick = 0;
+	else
+		ao_task_alarm_tick = ao_list_first_entry(&alarm_queue, struct ao_task, alarm_queue)->alarm;
 	ao_task_validate_alarm_queue();
 }
 
@@ -154,10 +161,7 @@ void
 ao_task_check_alarm(uint16_t tick)
 {
 	struct ao_task	*alarm, *next;
-	int		i;
 
-	if (ao_num_tasks == 0)
-		return;
 	ao_list_for_each_entry_safe(alarm, next, &alarm_queue, struct ao_task, alarm_queue) {
 		if ((int16_t) (tick - alarm->alarm) < 0)
 			break;
@@ -173,6 +177,7 @@ ao_task_init(void)
 	uint8_t	i;
 	ao_list_init(&run_queue);
 	ao_list_init(&alarm_queue);
+	ao_task_alarm_tick = 0;
 	for (i = 0; i < SLEEP_HASH_SIZE; i++)
 		ao_list_init(&sleep_queue[i]);
 }
@@ -264,14 +269,9 @@ ao_task_validate(void)
 		}
 	}
 }
-#else
-#define ao_task_validate()
-#endif
+#endif /* DEBUG */
 
-#else
-#define ao_task_to_run_queue(task)
-#define ao_task_to_alarm_queue(task)
-#endif
+#endif /* HAS_TASK_QUEUE */
 
 void
 ao_add_task(__xdata struct ao_task * task, void (*start)(void), __code char *name) __reentrant
@@ -331,6 +331,7 @@ ao_yield(void) ao_arch_naked_define
 	}
 
 	ao_arch_isr_stack();
+	ao_arch_block_interrupts();
 
 #if AO_CHECK_STACK
 	in_yield = 1;
@@ -339,21 +340,19 @@ ao_yield(void) ao_arch_naked_define
 	 * this loop will run forever, which is just fine
 	 */
 #if HAS_TASK_QUEUE
-	if (ao_cur_task->wchan == NULL) {
-		uint32_t flags;
-		flags = ao_arch_irqsave();
+	/* If the current task is running, move it to the
+	 * end of the queue to allow other tasks a chance
+	 */
+	if (ao_cur_task->wchan == NULL)
 		ao_task_to_run_queue(ao_cur_task);
-		ao_arch_irqrestore(flags);
-	}
 	ao_cur_task = NULL;
-
 	for (;;) {
 		ao_arch_memory_barrier();
 		if (!ao_list_is_empty(&run_queue))
 			break;
-		ao_arch_cpu_idle();
+		/* Wait for interrupts when there's nothing ready */
+		ao_arch_wait_interrupt();
 	}
-		
 	ao_cur_task = ao_list_first_entry(&run_queue, struct ao_task, queue);
 #else
 	{
@@ -374,20 +373,19 @@ ao_yield(void) ao_arch_naked_define
 			    (int16_t) (ao_time() - ao_cur_task->alarm) >= 0)
 				break;
 
-			/* Enter lower power mode when there isn't anything to do */
+			/* Wait for interrupts when there's nothing ready */
 			if (ao_cur_task_index == ao_last_task_index)
-				ao_arch_cpu_idle();
+				ao_arch_wait_interrupt();
 		}
-#if HAS_SAMPLE_PROFILE
-		ao_cur_task->start = ao_sample_profile_timer_value();
-#endif
 	}
 #endif
+#if HAS_SAMPLE_PROFILE
+	ao_cur_task->start = ao_sample_profile_timer_value();
+#endif
 #if HAS_STACK_GUARD
 	ao_mpu_stack_guard(ao_cur_task->stack);
 #endif
 #if AO_CHECK_STACK
-	ao_arch_block_interrupts();
 	in_yield = 0;
 #endif
 	ao_arch_restore_stack();
@@ -519,7 +517,7 @@ ao_task_info(void)
 		       task->name,
 		       (int) task->wchan);
 	}
-#if HAS_TASK_QUEUE
+#if HAS_TASK_QUEUE && DEBUG
 	ao_task_validate();
 #endif
 }