update wake-alarm implementation

4 files changed, 6 insertions, 10 deletions

diff --git a/main.c b/main.c
index 9c4055a28..91b6a062d 100755
--- a/main.c
+++ b/main.c
@@ -157,7 +157,7 @@ STATIC void start_mp(supervisor_allocation* heap) {
 
     #if CIRCUITPY_ALARM
     // Record which alarm woke us up, if any. An object may be created so the heap must be functional.
-    alarm_save_wakeup_alarm();
+    alarm_save_wake_alarm();
     // Reset alarm module only after we retrieved the wakeup alarm.
     alarm_reset();
     #endif
diff --git a/ports/esp32s2/common-hal/alarm/__init__.c b/ports/esp32s2/common-hal/alarm/__init__.c
index 0b1a7f049..47764a03a 100644
--- a/ports/esp32s2/common-hal/alarm/__init__.c
+++ b/ports/esp32s2/common-hal/alarm/__init__.c
@@ -121,9 +121,9 @@ STATIC void _idle_until_alarm(void) {
         RUN_BACKGROUND_TASKS;
         // Allow ctrl-C interrupt.
         if (alarm_woken_from_sleep()) {
+            alarm_save_wake_alarm();
             return;
         }
-
         port_idle_until_interrupt();
     }
 }
diff --git a/shared-bindings/alarm/__init__.c b/shared-bindings/alarm/__init__.c
index e8230a4ca..7023c70e5 100644
--- a/shared-bindings/alarm/__init__.c
+++ b/shared-bindings/alarm/__init__.c
@@ -227,21 +227,17 @@ mp_obj_t alarm_get_wake_alarm(void) {
     }
 }
 
-STATIC void alarm_set_wake_alarm(mp_obj_t alarm) {
+// Initialize .wake_alarm value.
+void alarm_save_wake_alarm(void) {
     // Equivalent of:
     // alarm.wake_alarm = alarm
     mp_map_elem_t *elem =
         mp_map_lookup(&alarm_module_globals.map, MP_ROM_QSTR(MP_QSTR_wake_alarm), MP_MAP_LOOKUP);
     if (elem) {
-        elem->value = alarm;
+        elem->value = common_hal_alarm_get_wake_alarm();
     }
 }
 
-// Initialize .wake_alarm value.
-void alarm_save_wakeup_alarm(void) {
-    alarm_set_wake_alarm(common_hal_alarm_get_wake_alarm());
-}
-
 const mp_obj_module_t alarm_module = {
     .base = { &mp_type_module },
     .globals = (mp_obj_dict_t*)&alarm_module_globals,
diff --git a/shared-bindings/alarm/__init__.h b/shared-bindings/alarm/__init__.h
index 4c12e0e70..154f91e26 100644
--- a/shared-bindings/alarm/__init__.h
+++ b/shared-bindings/alarm/__init__.h
@@ -50,7 +50,7 @@ extern void common_hal_alarm_gc_collect(void);
 extern mp_obj_t common_hal_alarm_get_wake_alarm(void);
 
 // Used by wake-up code.
-void alarm_save_wakeup_alarm(void);
+void alarm_save_wake_alarm(void);
 
 
 // True if an alarm is alerting. This is most useful for pretend deep sleep.