Add `board_deinit` for use with sleep

This changes lots of files to unify `board.h` across ports. It adds
`board_deinit` when CIRCUITPY_ALARM is set. `main.c` uses it to
deinit the board before deep sleeping (even when pretending.)

Deep sleep is now a two step process for the port. First, the
port should prepare to deep sleep based on the given alarms. It
should set alarms for both deep and pretend sleep. In particular,
the pretend versions should be set immediately so that we don't
miss an alarm as we shutdown. These alarms should also wake from
`port_idle_until_interrupt` which is used when pretending to deep
sleep.

Second, when real deep sleeping, `alarm_enter_deep_sleep` is called.
The port should set any alarms it didn't during prepare based on
data it saved internally during prepare.

ESP32-S2 sleep is a bit reorganized to locate more logic with
TimeAlarm. This will help it scale to more alarm types.

Fixes #3786

3 files changed, 121 insertions, 96 deletions

diff --git a/ports/esp32s2/common-hal/alarm/__init__.c b/ports/esp32s2/common-hal/alarm/__init__.c
index 11e173fe2..529179200 100644
--- a/ports/esp32s2/common-hal/alarm/__init__.c
+++ b/ports/esp32s2/common-hal/alarm/__init__.c
@@ -32,38 +32,41 @@
 #include "shared-bindings/alarm/pin/PinAlarm.h"
 #include "shared-bindings/alarm/time/TimeAlarm.h"
 #include "shared-bindings/microcontroller/__init__.h"
-#include "shared-bindings/time/__init__.h"
 #include "shared-bindings/wifi/__init__.h"
 
+#include "supervisor/port.h"
+#include "supervisor/shared/workflow.h"
+
 #include "common-hal/alarm/__init__.h"
 
-#include "esp_log.h"
 #include "esp_sleep.h"
 
-STATIC mp_obj_tuple_t *_deep_sleep_alarms;
-
 void alarm_reset(void) {
-    _deep_sleep_alarms = mp_const_empty_tuple;
+    alarm_time_timealarm_reset();
+    esp_sleep_disable_wakeup_source(ESP_SLEEP_WAKEUP_ALL);
 }
 
-void common_hal_alarm_disable_all(void) {
-    esp_sleep_disable_wakeup_source(ESP_SLEEP_WAKEUP_ALL);
+STATIC esp_sleep_wakeup_cause_t _get_wakeup_cause(void) {
+    if (alarm_time_timealarm_woke_us_up()) {
+        return ESP_SLEEP_WAKEUP_TIMER;
+    }
+
+    return esp_sleep_get_wakeup_cause();
 }
 
-mp_obj_t common_hal_alarm_get_wake_alarm(void) {
-    switch (esp_sleep_get_wakeup_cause()) {
+bool alarm_woken_from_sleep(void) {
+    return _get_wakeup_cause() != ESP_SLEEP_WAKEUP_UNDEFINED;
+}
+
+STATIC mp_obj_t _get_wake_alarm(size_t n_alarms, const mp_obj_t *alarms) {
+    switch (_get_wakeup_cause()) {
         case ESP_SLEEP_WAKEUP_TIMER: {
-            // Wake up from timer.
-            alarm_time_time_alarm_obj_t *timer = m_new_obj(alarm_time_time_alarm_obj_t);
-            timer->base.type = &alarm_time_time_alarm_type;
-            return timer;
+            return alarm_time_timealarm_get_wakeup_alarm(n_alarms, alarms);
         }
 
         case ESP_SLEEP_WAKEUP_EXT0: {
-            // Wake up from GPIO
-            alarm_pin_pin_alarm_obj_t *ext0 = m_new_obj(alarm_pin_pin_alarm_obj_t);
-            ext0->base.type = &alarm_pin_pin_alarm_type;
-            return ext0;
+            // TODO: implement pin alarm wake.
+            break;
         }
 
         case ESP_SLEEP_WAKEUP_TOUCHPAD:
@@ -79,16 +82,19 @@ mp_obj_t common_hal_alarm_get_wake_alarm(void) {
     return mp_const_none;
 }
 
+mp_obj_t common_hal_alarm_get_wake_alarm(void) {
+    return _get_wake_alarm(0, NULL);
+}
+
 // Set up light sleep or deep sleep alarms.
-STATIC void setup_sleep_alarms(size_t n_alarms, const mp_obj_t *alarms) {
+STATIC void _setup_sleep_alarms(bool deep_sleep, size_t n_alarms, const mp_obj_t *alarms) {
     bool time_alarm_set = false;
     alarm_time_time_alarm_obj_t *time_alarm = MP_OBJ_NULL;
 
     for (size_t i = 0; i < n_alarms; i++) {
         if (MP_OBJ_IS_TYPE(alarms[i], &alarm_pin_pin_alarm_type)) {
             mp_raise_NotImplementedError(translate("PinAlarm not yet implemented"));
-        }
-        else if (MP_OBJ_IS_TYPE(alarms[i], &alarm_time_time_alarm_type)) {
+        } else if (MP_OBJ_IS_TYPE(alarms[i], &alarm_time_time_alarm_type)) {
             if (time_alarm_set) {
                 mp_raise_ValueError(translate("Only one alarm.time alarm can be set."));
             }
@@ -97,102 +103,50 @@ STATIC void setup_sleep_alarms(size_t n_alarms, const mp_obj_t *alarms) {
         }
     }
 
-    if (time_alarm != MP_OBJ_NULL) {
-        // Compute how long to actually sleep, considering the time now.
-        mp_float_t now_secs = uint64_to_float(common_hal_time_monotonic_ms()) / 1000.0f;
-        mp_float_t wakeup_in_secs = MAX(0.0f, time_alarm->monotonic_time - now_secs);
-        const uint64_t sleep_for_us = (uint64_t) (wakeup_in_secs * 1000000);
-        ESP_LOGI("ALARM", "will sleep for us: %lld", sleep_for_us);
-        esp_sleep_enable_timer_wakeup(sleep_for_us);
+    if (time_alarm_set) {
+        alarm_time_timealarm_set_alarm(time_alarm);
     }
 }
 
-mp_obj_t common_hal_alarm_wait_until_alarms(size_t n_alarms, const mp_obj_t *alarms) {
-    if (n_alarms == 0) {
-        return mp_const_none;
-    }
-
-    bool time_alarm_set = false;
-    alarm_time_time_alarm_obj_t *time_alarm = MP_OBJ_NULL;
-
-    for (size_t i = 0; i < n_alarms; i++) {
-        if (MP_OBJ_IS_TYPE(alarms[i], &alarm_pin_pin_alarm_type)) {
-            mp_raise_NotImplementedError(translate("PinAlarm not yet implemented"));
+STATIC void _idle_until_alarm(void) {
+    // Poll for alarms.
+    while (!mp_hal_is_interrupted()) {
+        RUN_BACKGROUND_TASKS;
+        // Allow ctrl-C interrupt.
+        if (alarm_woken_from_sleep()) {
+            return;
         }
-        else if (MP_OBJ_IS_TYPE(alarms[i], &alarm_time_time_alarm_type)) {
-            if (time_alarm_set) {
-                mp_raise_ValueError(translate("Only one alarm.time alarm can be set."));
-            }
-            time_alarm  = MP_OBJ_TO_PTR(alarms[i]);
-            time_alarm_set = true;
-        }
-    }
-
-    ESP_LOGI("ALARM", "waiting for alarms");
 
-    if (time_alarm_set && n_alarms == 1) {
-        // If we're only checking time, avoid a polling loop, so maybe we can save some power.
-        const mp_float_t now_secs = uint64_to_float(common_hal_time_monotonic_ms()) / 1000.0f;
-        const mp_float_t wakeup_in_secs = MAX(0.0f, time_alarm->monotonic_time - now_secs);
-        const uint32_t delay_ms = (uint32_t) (wakeup_in_secs * 1000.0f);
-        ESP_LOGI("ALARM", "Delay for ms: %d", delay_ms);
-        common_hal_time_delay_ms((uint32_t) delay_ms);
-    } else {
-        // Poll for alarms.
-        while (true) {
-            RUN_BACKGROUND_TASKS;
-            // Allow ctrl-C interrupt.
-            if (mp_hal_is_interrupted()) {
-                return mp_const_none;
-            }
-
-            // TODO: Check PinAlarms.
-
-            if (time_alarm != MP_OBJ_NULL &&
-                common_hal_time_monotonic_ms() * 1000.f >= time_alarm->monotonic_time) {
-                return time_alarm;
-            }
-        }
+        port_idle_until_interrupt();
     }
-
-    return mp_const_none;
 }
 
 // Is it safe to do a light sleep? Check whether WiFi is on or there are
 // other ongoing tasks that should not be shut down.
-static bool light_sleep_ok(void) {
-    return !common_hal_wifi_radio_get_enabled(&common_hal_wifi_radio_obj);
+STATIC bool _light_sleep_ok(void) {
+    return !common_hal_wifi_radio_get_enabled(&common_hal_wifi_radio_obj) && !supervisor_workflow_active();
 }
 
 mp_obj_t common_hal_alarm_light_sleep_until_alarms(size_t n_alarms, const mp_obj_t *alarms) {
-    if (n_alarms == 0) {
-        return mp_const_none;
-    }
+    _setup_sleep_alarms(false, n_alarms, alarms);
 
-    if (light_sleep_ok()) {
-        ESP_LOGI("ALARM", "start light sleep");
-        setup_sleep_alarms(n_alarms, alarms);
+    // Light sleep can break some functionality so only do it when possible. Otherwise we idle.
+    if (_light_sleep_ok()) {
         esp_light_sleep_start();
-        return common_hal_alarm_get_wake_alarm();
     } else {
-        // Don't do an ESP32 light sleep.
-        return common_hal_alarm_wait_until_alarms(n_alarms, alarms);
+        _idle_until_alarm();
     }
+    mp_obj_t wake_alarm = _get_wake_alarm(n_alarms, alarms);
+    alarm_reset();
+    return wake_alarm;
 }
 
-void common_hal_alarm_exit_and_deep_sleep_until_alarms(size_t n_alarms, const mp_obj_t *alarms) {
-    setup_sleep_alarms(n_alarms, alarms);
-
-    // Raise an exception, which will be processed in main.c.
-    mp_raise_arg1(&mp_type_DeepSleepRequest, NULL);
-}
-
-void common_hal_alarm_prepare_for_deep_sleep(void) {
-    // Turn off WiFi and anything else that should be shut down cleanly.
-    common_hal_wifi_radio_set_enabled(&common_hal_wifi_radio_obj, false);
+void common_hal_alarm_set_deep_sleep_alarms(size_t n_alarms, const mp_obj_t *alarms) {
+    _setup_sleep_alarms(true, n_alarms, alarms);
 }
 
-void NORETURN common_hal_alarm_enter_deep_sleep(void) {
-    ESP_LOGI("ALARM", "start deep sleep");
+void NORETURN alarm_enter_deep_sleep(void) {
+    // The ESP-IDF caches the deep sleep settings and applies them before sleep.
+    // We don't need to worry about resetting them in the interim.
     esp_deep_sleep_start();
 }
diff --git a/ports/esp32s2/common-hal/alarm/time/TimeAlarm.c b/ports/esp32s2/common-hal/alarm/time/TimeAlarm.c
index 7af369463..34c8b0052 100644
--- a/ports/esp32s2/common-hal/alarm/time/TimeAlarm.c
+++ b/ports/esp32s2/common-hal/alarm/time/TimeAlarm.c
@@ -27,8 +27,12 @@
 #include "esp_sleep.h"
 
 #include "py/runtime.h"
+#include "supervisor/esp_port.h"
+
+#include "components/esp_timer/include/esp_timer.h"
 
 #include "shared-bindings/alarm/time/TimeAlarm.h"
+#include "shared-bindings/time/__init__.h"
 
 void common_hal_alarm_time_time_alarm_construct(alarm_time_time_alarm_obj_t *self, mp_float_t monotonic_time) {
     self->monotonic_time = monotonic_time;
@@ -37,3 +41,63 @@ void common_hal_alarm_time_time_alarm_construct(alarm_time_time_alarm_obj_t *sel
 mp_float_t common_hal_alarm_time_time_alarm_get_monotonic_time(alarm_time_time_alarm_obj_t *self) {
     return self->monotonic_time;
 }
+
+mp_obj_t alarm_time_timealarm_get_wakeup_alarm(size_t n_alarms, const mp_obj_t *alarms) {
+    // First, check to see if we match
+    for (size_t i = 0; i < n_alarms; i++) {
+        if (MP_OBJ_IS_TYPE(alarms[i], &alarm_time_time_alarm_type)) {
+            return alarms[i];
+        }
+    }
+    alarm_time_time_alarm_obj_t *timer = m_new_obj(alarm_time_time_alarm_obj_t);
+    timer->base.type = &alarm_time_time_alarm_type;
+    return timer;
+}
+
+esp_timer_handle_t pretend_sleep_timer;
+STATIC bool woke_up = false;
+
+// This is run in the timer task. We use it to wake the main CircuitPython task.
+void timer_callback(void *arg) {
+    (void) arg;
+    woke_up = true;
+    if (sleeping_circuitpython_task) {
+        xTaskNotifyGive(sleeping_circuitpython_task);
+    }
+}
+
+bool alarm_time_timealarm_woke_us_up(void) {
+    return woke_up;
+}
+
+void alarm_time_timealarm_reset(void) {
+    esp_timer_stop(pretend_sleep_timer);
+    esp_timer_delete(pretend_sleep_timer);
+    pretend_sleep_timer = NULL;
+    woke_up = false;
+}
+
+void alarm_time_timealarm_set_alarm(alarm_time_time_alarm_obj_t *self) {
+    if (pretend_sleep_timer != NULL) {
+        esp_timer_stop(pretend_sleep_timer);
+    } else {
+        // Configure the timer to use during pretend sleep.
+        esp_timer_create_args_t args;
+        args.callback = timer_callback;
+        args.arg = NULL;
+        args.dispatch_method = ESP_TIMER_TASK;
+        args.name = "Pretend deep sleep";
+        esp_timer_create(&args, &pretend_sleep_timer);
+    }
+
+    // Compute how long to actually sleep, considering the time now.
+    mp_float_t now_secs = uint64_to_float(common_hal_time_monotonic_ms()) / 1000.0f;
+    mp_float_t wakeup_in_secs = MAX(0.0f, self->monotonic_time - now_secs);
+    const uint64_t sleep_for_us = (uint64_t) (wakeup_in_secs * 1000000);
+    esp_sleep_enable_timer_wakeup(sleep_for_us);
+
+    // Also set the RTC interrupt so it can wake our task. This will be wiped out
+    // if we actually deep sleep.
+    woke_up = false;
+    esp_timer_start_once(pretend_sleep_timer, sleep_for_us);
+}
diff --git a/ports/esp32s2/common-hal/alarm/time/TimeAlarm.h b/ports/esp32s2/common-hal/alarm/time/TimeAlarm.h
index b540541f4..04c553009 100644
--- a/ports/esp32s2/common-hal/alarm/time/TimeAlarm.h
+++ b/ports/esp32s2/common-hal/alarm/time/TimeAlarm.h
@@ -31,3 +31,10 @@ typedef struct {
     mp_obj_base_t base;
     mp_float_t monotonic_time;      // values compatible with time.monotonic_time()
 } alarm_time_time_alarm_obj_t;
+
+// Find the alarm object that caused us to wake up or create an equivalent one.
+mp_obj_t alarm_time_timealarm_get_wakeup_alarm(size_t n_alarms, const mp_obj_t *alarms);
+// Check for the wake up alarm from pretend deep sleep.
+bool alarm_time_timealarm_woke_us_up(void);
+void alarm_time_timealarm_set_alarm(alarm_time_time_alarm_obj_t *self);
+void alarm_time_timealarm_reset(void);