From 4443164e3fa0358cbb7e26342795a21b235f862a Mon Sep 17 00:00:00 2001
From: zhiayang <gpg@tsumaran.ai>
Date: Fri, 9 Feb 2024 00:49:56 -0500
Subject: [PATCH] fix initial display of battery status on displays, and also
 fix a null deref

---
 app/boards/arm/corneish_zen/widgets/battery_status.c | 12 ++++++++++--
 app/boards/shields/nice_view/widgets/status.c        | 12 ++++++++++--
 app/src/display/widgets/battery_status.c             | 12 ++++++++++--
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/app/boards/arm/corneish_zen/widgets/battery_status.c b/app/boards/arm/corneish_zen/widgets/battery_status.c
index 011319c4..499e22a4 100644
--- a/app/boards/arm/corneish_zen/widgets/battery_status.c
+++ b/app/boards/arm/corneish_zen/widgets/battery_status.c
@@ -66,8 +66,16 @@ void battery_status_update_cb(struct battery_status_state state) {
 }
 
 static struct battery_status_state battery_status_get_state(const zmk_event_t *eh) {
-    return (struct battery_status_state) {
-        .level = zmk_battery_state_of_charge(),
+    const struct zmk_battery_state_changed *ev = as_zmk_battery_state_changed(eh);
+    uint8_t soc;
+    if (ev != NULL) {
+        soc = ev->state_of_charge;
+    } else {
+        soc = zmk_battery_state_of_charge();
+    }
+
+    return (struct battery_status_state){
+        .level = soc,
 #if IS_ENABLED(CONFIG_USB_DEVICE_STACK)
         .usb_present = zmk_usb_is_powered(),
 #endif /* IS_ENABLED(CONFIG_USB_DEVICE_STACK) */
diff --git a/app/boards/shields/nice_view/widgets/status.c b/app/boards/shields/nice_view/widgets/status.c
index 3346546e..592c0551 100644
--- a/app/boards/shields/nice_view/widgets/status.c
+++ b/app/boards/shields/nice_view/widgets/status.c
@@ -210,8 +210,16 @@ static void battery_status_update_cb(struct battery_status_state state) {
 }
 
 static struct battery_status_state battery_status_get_state(const zmk_event_t *eh) {
-    return (struct battery_status_state) {
-        .level = zmk_battery_state_of_charge(),
+    const struct zmk_battery_state_changed *ev = as_zmk_battery_state_changed(eh);
+    uint8_t soc;
+    if (ev != NULL) {
+        soc = ev->state_of_charge;
+    } else {
+        soc = zmk_battery_state_of_charge();
+    }
+
+    return (struct battery_status_state){
+        .level = soc,
 #if IS_ENABLED(CONFIG_USB_DEVICE_STACK)
         .usb_present = zmk_usb_is_powered(),
 #endif /* IS_ENABLED(CONFIG_USB_DEVICE_STACK) */
diff --git a/app/src/display/widgets/battery_status.c b/app/src/display/widgets/battery_status.c
index 06f55f3f..a3f7c9f2 100644
--- a/app/src/display/widgets/battery_status.c
+++ b/app/src/display/widgets/battery_status.c
@@ -9,6 +9,7 @@
 #include <zephyr/logging/log.h>
 LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 
+#include <zmk/battery.h>
 #include <zmk/display.h>
 #include <zmk/display/widgets/battery_status.h>
 #include <zmk/usb.h>
@@ -63,8 +64,15 @@ void battery_status_update_cb(struct battery_status_state state) {
 
 static struct battery_status_state battery_status_get_state(const zmk_event_t *eh) {
     const struct zmk_battery_state_changed *ev = as_zmk_battery_state_changed(eh);
-    return (struct battery_status_state) {
-        .level = ev->state_of_charge,
+    uint8_t soc;
+    if (ev != NULL) {
+        soc = ev->state_of_charge;
+    } else {
+        soc = zmk_battery_state_of_charge();
+    }
+
+    return (struct battery_status_state){
+        .level = soc,
 #if IS_ENABLED(CONFIG_USB_DEVICE_STACK)
         .usb_present = zmk_usb_is_powered(),
 #endif /* IS_ENABLED(CONFIG_USB_DEVICE_STACK) */