oled_ext_pwr: Apply ddudek’s OLED re-init patch

This commit allows zmk to properly reinitialize the OLED display after ext power is re-enabled.

More info here:
zmkfirmware/zmk#674

This code was written by ddudek and not me. I am just adding it as a commit.

drivers/display/ssd1306.c

@@ -438,7 +438,94 @@ static const struct ssd1306_config ssd1306_config = {
 
 static struct ssd1306_data ssd1306_driver;
 
+static int ssd1306_re_init(const struct device *dev)
+{
+	LOG_DBG("ssd1306 oled re-init now");
+
+	// Note: uncoment below for full re-init
+	return ssd1306_init(dev);
+
+	// Note: Below list is minimal working set experimentally:
+
+// 	if (ssd1306_set_charge_pump(dev)) {
+// 		return -EIO;
+// 	}
+
+// 	uint8_t cmd_buf[] = {
+// 		SSD1306_SET_ENTIRE_DISPLAY_OFF,
+// #ifdef CONFIG_SSD1306_REVERSE_MODE
+// 		SSD1306_SET_REVERSE_DISPLAY,
+// #else
+// 		SSD1306_SET_NORMAL_DISPLAY,
+// #endif
+// 	};
+
+// 	if (ssd1306_write_bus(dev, cmd_buf, sizeof(cmd_buf), true)) {
+// 		LOG_ERR("ssd1306_write_bus");
+// 		return -EIO;
+// 	}
+
+// 	if (ssd1306_set_contrast(dev, CONFIG_SSD1306_DEFAULT_CONTRAST)) {
+// 		return -EIO;
+// 	}
+
+// 	ssd1306_resume(dev);
+
+// 	return 0;
+}
+
+static int ext_power_status = true;
+
+static int ssd1306_update_ext_power(const struct device *dev, bool ext_power_status_new_value) {
+
+	LOG_DBG("Inside update_ext_pwr");
+
+	// minimum sleep needed when waking up
+	if (ext_power_status_new_value == true) {
+		k_sleep(K_MSEC(30));
+	}
+
+	LOG_DBG("dev-config: %p", dev->config);
+
+	// first update to I2C
+	if (dev->config != NULL) {
+		struct ssd1306_config *config = dev->config;
+
+		LOG_DBG("Before i2c_update_ext_power");
+
+		if(i2c_update_ext_power(&config->bus, ext_power_status_new_value)) {
+			LOG_ERR("Failed i2c_update_ext_power!");
+			return -EIO;
+		}
+		LOG_DBG("After i2c_update_ext_power");
+
+	} else {
+		LOG_ERR("display config is NULL");
+	}
+
+	if (ext_power_status != ext_power_status_new_value) {
+		if (ext_power_status_new_value == true)
+		{
+			// sleep after I2C reset
+			k_sleep(K_MSEC(30));
+
+			LOG_DBG("Before ssd1306_re_init");
+			// re-init oled, sends commands through i2c bus
+			ssd1306_re_init(dev);
+			LOG_DBG("After ssd1306_re_init");
+		} else {
+			// no-op for now
+		}
+
+		ext_power_status = ext_power_status_new_value;
+	}
+	return 0;
+}
+
+
+
 static struct display_driver_api ssd1306_driver_api = {
+	.update_ext_power = ssd1306_update_ext_power,
 	.blanking_on = ssd1306_suspend,
 	.blanking_off = ssd1306_resume,
 	.write = ssd1306_write,
@@ -451,7 +538,9 @@ static struct display_driver_api ssd1306_driver_api = {
 	.set_orientation = ssd1306_set_orientation,
 };
 
+define DISPLAY_INIT_PRIORITY 95
+
 DEVICE_DT_INST_DEFINE(0, ssd1306_init, NULL,
 		      &ssd1306_driver, &ssd1306_config,
-		      POST_KERNEL, CONFIG_DISPLAY_INIT_PRIORITY,
+		      POST_KERNEL, DISPLAY_INIT_PRIORITY,
 		      &ssd1306_driver_api);

drivers/i2c/i2c_nrfx_twi.c

@@ -197,14 +197,11 @@ static int i2c_nrfx_twi_recover_bus(const struct device *dev)
 	return (err == NRFX_SUCCESS ? 0 : -EBUSY);
 }
 
-static const struct i2c_driver_api i2c_nrfx_twi_driver_api = {
-	.configure   = i2c_nrfx_twi_configure,
-	.transfer    = i2c_nrfx_twi_transfer,
-	.recover_bus = i2c_nrfx_twi_recover_bus,
-};
+static int initialized = false;
 
 static int init_twi(const struct device *dev)
 {
+	LOG_WRN("init_twi");
 	const struct i2c_nrfx_twi_config *config = dev->config;
 	struct i2c_nrfx_twi_data *dev_data = dev->data;
 	nrfx_err_t result = nrfx_twi_init(&config->twi, &config->config,
@@ -214,28 +211,82 @@ static int init_twi(const struct device *dev)
 			    dev->name);
 		return -EBUSY;
 	}
+    initialized = true;
+	return 0;
+}
 
+
+static int i2c_nrfx_twi_update_ext_power(const struct device *dev, bool ext_power_enabled) {
+	LOG_WRN("I2C update_ext_power now");
+	if(ext_power_enabled) {
+		LOG_WRN("New state power on, re-init");
+		if (!initialized) {
+			struct i2c_nrfx_twi_config *config = dev->config;
+			LOG_DBG("Before nrfx_twi_uninit");
+			nrfx_twi_uninit(&config->twi);
+			LOG_DBG("After nrfx_twi_uninit");
+
+			LOG_DBG("Before init_twi");
+			init_twi(dev);
+			LOG_DBG("After init_twi");
+
+			// if (get_dev_data(dev)->dev_config) {
+			// 	i2c_nrfx_twi_configure(
+			// 		dev,
+			// 		get_dev_data(dev)->dev_config);
+			// }
+		}
+	} else {
+		if (initialized) {
+			initialized = false;
+		}
+	}
 	return 0;
 }
 
+
+static const struct i2c_driver_api i2c_nrfx_twi_driver_api = {
+	.update_ext_power = i2c_nrfx_twi_update_ext_power,
+	.configure   = i2c_nrfx_twi_configure,
+	.transfer    = i2c_nrfx_twi_transfer,
+	.recover_bus = i2c_nrfx_twi_recover_bus,
+};
+
 #ifdef CONFIG_PM_DEVICE
 static int twi_nrfx_pm_action(const struct device *dev,
 			      enum pm_device_action action)
 {
 	const struct i2c_nrfx_twi_config *config = dev->config;
 	struct i2c_nrfx_twi_data *data = dev->data;
 	int ret = 0;
-
+	LOG_DBG("In twi_nrfx_pm_action");
 	switch (action) {
 	case PM_DEVICE_ACTION_RESUME:
-		init_twi(dev);
-		if (data->dev_config) {
-			i2c_nrfx_twi_configure(dev, data->dev_config);
+        LOG_DBG("In PM_DEVICE_ACTION_RESUME");
+
+		if (!initialized) {
+			LOG_DBG("Before init_twi");
+			init_twi(dev);
+			LOG_DBG("After init_twi");
+
+			if (data->dev_config) {
+				LOG_DBG("Before i2c_nrfx_twi_configure");
+
+				i2c_nrfx_twi_configure(dev, data->dev_config);
+				LOG_DBG("After i2c_nrfx_twi_configure");
+
+			}
 		}
 		break;
 
 	case PM_DEVICE_ACTION_SUSPEND:
-		nrfx_twi_uninit(&config->twi);
+	    LOG_DBG("In PM_DEVICE_ACTION_SUSPEND");
+		if (initialized) {
+			LOG_DBG("Before nrfx_twi_uninit");
+			nrfx_twi_uninit(&config->twi);
+			LOG_DBG("After nrfx_twi_uninit");
+			initialized = false;
+		}
 		break;
 
 	default:

include/drivers/display.h

@@ -108,7 +108,11 @@ struct display_buffer_descriptor {
 	/** Number of pixels between consecutive rows in the data buffer */
 	uint16_t pitch;
 };
-
+/**
+ * @typedef display_update_ext_power_api
+ * @brief Callback API to manually update display power state
+ */
+typedef int (*display_update_ext_power_api)(const struct device *dev, bool enabled);
 /**
  * @typedef display_blanking_on_api
  * @brief Callback API to turn on display blanking
@@ -198,6 +202,7 @@ typedef int (*display_set_orientation_api)(const struct device *dev,
  * API which a display driver should expose
  */
 struct display_driver_api {
+	display_update_ext_power_api update_ext_power;
 	display_blanking_on_api blanking_on;
 	display_blanking_off_api blanking_off;
 	display_write_api write;
@@ -209,6 +214,14 @@ struct display_driver_api {
 	display_set_pixel_format_api set_pixel_format;
 	display_set_orientation_api set_orientation;
 };
+static inline int display_update_ext_power(const struct device *dev, bool enabled)
+{
+	struct display_driver_api *api =
+		(struct display_driver_api *)dev->api;
+
+	return api->update_ext_power(dev, enabled);
+}
+
 
 /**
  * @brief Write data to display

include/drivers/i2c.h

@@ -161,6 +161,7 @@ struct i2c_msg {
  */
 struct i2c_slave_config;
 
+typedef int (*i2c_api_update_ext_power_t)(const struct device *dev, bool enabled);
 typedef int (*i2c_api_configure_t)(const struct device *dev,
 				   uint32_t dev_config);
 typedef int (*i2c_api_get_config_t)(const struct device *dev,
@@ -176,6 +177,7 @@ typedef int (*i2c_api_slave_unregister_t)(const struct device *dev,
 typedef int (*i2c_api_recover_bus_t)(const struct device *dev);
 
 __subsystem struct i2c_driver_api {
+	i2c_api_update_ext_power_t update_ext_power;
 	i2c_api_configure_t configure;
 	i2c_api_get_config_t get_config;
 	i2c_api_full_io_t transfer;
@@ -612,6 +614,19 @@ static inline int i2c_transfer_dt(const struct i2c_dt_spec *spec,
 	return i2c_transfer(spec->bus, msgs, num_msgs, spec->addr);
 }
 
+static inline int i2c_update_ext_power(const struct device *dev, bool enabled)
+{
+	const struct i2c_driver_api *api =
+		(const struct i2c_driver_api *)dev->api;
+
+	if (api->update_ext_power == NULL) {
+		return -ENOTSUP;
+	}
+
+	return api->update_ext_power(dev, enabled);
+}
+
+
 /**
  * @brief Recover the I2C bus
  *