[nrf fromtree] modified: nrf_wifi: Refactor heap calls

Instead if littering ifdef's across the code, use a single API and
initialize different pools only once.

Signed-off-by: Chaitanya Tata <[email protected]>
(cherry picked from commit ab9c531)

Author: krish2718

modules/nrf_wifi/os/shim.c

@@ -34,14 +34,23 @@
 
 LOG_MODULE_REGISTER(wifi_nrf, CONFIG_WIFI_NRF70_LOG_LEVEL);
 
-#if !defined(CONFIG_NRF_WIFI_GLOBAL_HEAP)
+/* Memory pool management - unified pool-based API */
+#if defined(CONFIG_NRF_WIFI_GLOBAL_HEAP)
+/* Use global system heap */
+extern struct sys_heap _system_heap;
+static struct k_heap * const wifi_ctrl_pool = &_system_heap;
+static struct k_heap * const wifi_data_pool = &_system_heap;
+#else
+/* Use dedicated heaps */
 #if defined(CONFIG_NOCACHE_MEMORY)
 K_HEAP_DEFINE_NOCACHE(wifi_drv_ctrl_mem_pool, CONFIG_NRF_WIFI_CTRL_HEAP_SIZE);
 K_HEAP_DEFINE_NOCACHE(wifi_drv_data_mem_pool, CONFIG_NRF_WIFI_DATA_HEAP_SIZE);
 #else
 K_HEAP_DEFINE(wifi_drv_ctrl_mem_pool, CONFIG_NRF_WIFI_CTRL_HEAP_SIZE);
 K_HEAP_DEFINE(wifi_drv_data_mem_pool, CONFIG_NRF_WIFI_DATA_HEAP_SIZE);
 #endif /* CONFIG_NOCACHE_MEMORY */
+static struct k_heap * const wifi_ctrl_pool = &wifi_drv_ctrl_mem_pool;
+static struct k_heap * const wifi_data_pool = &wifi_drv_data_mem_pool;
 #endif /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
 
 #define WORD_SIZE 4
@@ -52,22 +61,14 @@ static void *zep_shim_mem_alloc(size_t size)
 {
 	size_t size_aligned = ROUND_UP(size, 4);
 
-#if defined(CONFIG_NRF_WIFI_GLOBAL_HEAP)
-	return k_malloc(size_aligned);
-#else /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
-	return k_heap_aligned_alloc(&wifi_drv_ctrl_mem_pool, WORD_SIZE, size_aligned, K_FOREVER);
-#endif /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
+	return k_heap_aligned_alloc(wifi_ctrl_pool, WORD_SIZE, size_aligned, K_FOREVER);
 }
 
 static void *zep_shim_data_mem_alloc(size_t size)
 {
 	size_t size_aligned = ROUND_UP(size, 4);
 
-#if defined(CONFIG_NRF_WIFI_GLOBAL_HEAP)
-	return k_malloc(size_aligned);
-#else /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
-	return k_heap_aligned_alloc(&wifi_drv_data_mem_pool, WORD_SIZE, size_aligned, K_FOREVER);
-#endif /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
+	return k_heap_aligned_alloc(wifi_data_pool, WORD_SIZE, size_aligned, K_FOREVER);
 }
 
 static void *zep_shim_mem_zalloc(size_t size)
@@ -111,22 +112,14 @@ static void *zep_shim_data_mem_zalloc(size_t size)
 static void zep_shim_mem_free(void *buf)
 {
 	if (buf) {
-#if defined(CONFIG_NRF_WIFI_GLOBAL_HEAP)
-		k_free(buf);
-#else /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
-		k_heap_free(&wifi_drv_ctrl_mem_pool, buf);
-#endif /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
+		k_heap_free(wifi_ctrl_pool, buf);
 	}
 }
 
 static void zep_shim_data_mem_free(void *buf)
 {
 	if (buf) {
-#if defined(CONFIG_NRF_WIFI_GLOBAL_HEAP)
-		k_free(buf);
-#else /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
-		k_heap_free(&wifi_drv_data_mem_pool, buf);
-#endif /* CONFIG_NRF_WIFI_GLOBAL_HEAP */
+		k_heap_free(wifi_data_pool, buf);
 	}
 }
 
@@ -211,21 +204,21 @@ static void *zep_shim_spinlock_alloc(void)
 {
 	struct zep_shim_spinlock *slock = NULL;
 
-	slock = k_calloc(sizeof(*slock), sizeof(char));
+	slock = k_heap_aligned_alloc(wifi_ctrl_pool, WORD_SIZE, sizeof(*slock), K_FOREVER);
 	if (!slock) {
 		LOG_ERR("%s: Unable to allocate memory for spinlock", __func__);
+	} else {
+		memset(slock, 0, sizeof(*slock));
 	}
 
 	return slock;
 }
 
 static void zep_shim_spinlock_free(void *lock)
 {
-#if defined(CONFIG_NRF_WIFI_GLOBAL_HEAP)
-	k_free(lock);
-#else
-	k_heap_free(&wifi_drv_ctrl_mem_pool, lock);
-#endif
+	if (lock) {
+		k_heap_free(wifi_ctrl_pool, lock);
+	}
 }
 
 static void zep_shim_spinlock_init(void *lock)