fix: Optimize layout changes by doing runtime mapping

* To avoid tons of migration, extra flash writes, etc, we keep
  the keymaps and settings using a key position index that's tied
  to the stock layout, and at runtime mapping key positions as
  needed.

4 files changed, 115 insertions, 57 deletions

diff --git a/app/include/zmk/physical_layouts.h b/app/include/zmk/physical_layouts.h
index 33004af738..4760b48dfe 100644
--- a/app/include/zmk/physical_layouts.h
+++ b/app/include/zmk/physical_layouts.h
@@ -49,3 +49,12 @@ int zmk_physical_layouts_revert_selected(void);
 
 int zmk_physical_layouts_get_position_map(uint8_t source, uint8_t dest, size_t map_size,
                                           uint32_t map[map_size]);
+
+/**
+ * @brief Get a pointer to a position map array for mapping a key position in the selected
+ *        physical layout to the stock/chosen physical layout
+ *
+ * @retval a negative errno value in the case of errors
+ * @retval a positive length of the position map array that map is updated to point to.
+ */
+int zmk_physical_layouts_get_selected_to_stock_position_map(uint32_t const **map);
\ No newline at end of file
diff --git a/app/src/keymap.c b/app/src/keymap.c
index 0ea64b340f..6a6928dd2b 100644
--- a/app/src/keymap.c
+++ b/app/src/keymap.c
@@ -13,6 +13,7 @@ LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 #include <zmk/stdlib.h>
 #include <zmk/behavior.h>
 #include <zmk/keymap.h>
+#include <zmk/physical_layouts.h>
 #include <zmk/matrix.h>
 #include <zmk/sensors.h>
 #include <zmk/virtual_key_position.h>
@@ -230,7 +231,26 @@ zmk_keymap_get_layer_binding_at_idx(zmk_keymap_layer_id_t layer_id, uint8_t bind
 
     ASSERT_LAYER_VAL(layer_id, NULL)
 
-    return &zmk_keymap[layer_id][binding_idx];
+    const uint32_t *pos_map;
+    int ret = zmk_physical_layouts_get_selected_to_stock_position_map(&pos_map);
+    if (ret < 0) {
+        LOG_WRN("Failed to get the position map, can't find the right binding to return (%d)", ret);
+        return NULL;
+    }
+
+    if (binding_idx >= ret) {
+        LOG_WRN("Can't return binding for unmapped binding index %d", binding_idx);
+        return NULL;
+    }
+
+    uint32_t mapped_idx = pos_map[binding_idx];
+
+    if (mapped_idx >= ZMK_KEYMAP_LEN) {
+        LOG_WRN("Binding index %d mapped to an invalid key position %d", binding_idx, mapped_idx);
+        return NULL;
+    }
+
+    return &zmk_keymap[layer_id][mapped_idx];
 }
 
 #if IS_ENABLED(CONFIG_ZMK_KEYMAP_SETTINGS_STORAGE)
@@ -247,12 +267,37 @@ int zmk_keymap_set_layer_binding_at_idx(zmk_keymap_layer_id_t layer_id, uint8_t
 
     ASSERT_LAYER_VAL(layer_id, -EINVAL)
 
+    const uint32_t *pos_map;
+    int ret = zmk_physical_layouts_get_selected_to_stock_position_map(&pos_map);
+    if (ret < 0) {
+        LOG_WRN("Failed to get the mapping to determine where to set the binding (%d)", ret);
+        return ret;
+    }
+
+    if (binding_idx >= ret) {
+        LOG_WRN("Unable to set binding at index %d which isn't mapped", binding_idx);
+        return -EINVAL;
+    }
+
+    uint32_t storage_binding_idx = pos_map[binding_idx];
+
+    if (storage_binding_idx >= ZMK_KEYMAP_LEN) {
+        LOG_WRN("Can't set layer binding at unmapped/invalid index %d", binding_idx);
+        return -EINVAL;
+    }
+
+    if (memcmp(&zmk_keymap[layer_id][storage_binding_idx], &binding, sizeof(binding)) == 0) {
+        LOG_DBG("Not setting, no change to layer %d at index %d (%d)", layer_id, binding_idx,
+                storage_binding_idx);
+        return 0;
+    }
+
     uint8_t *pending = zmk_keymap_layer_pending_changes[layer_id];
 
     WRITE_BIT(pending[binding_idx / 8], binding_idx % 8, 1);
 
     // TODO: Need a mutex to protect access to the keymap data?
-    memcpy(&zmk_keymap[layer_id][binding_idx], &binding, sizeof(binding));
+    memcpy(&zmk_keymap[layer_id][storage_binding_idx], &binding, sizeof(binding));
 
     return 0;
 }
@@ -440,7 +485,8 @@ static int save_bindings(void) {
             if (pending[kp / 8] & BIT(kp % 8)) {
                 LOG_DBG("Pending save for layer %d at key position %d", l, kp);
 
-                struct zmk_behavior_binding *binding = &zmk_keymap[l][kp];
+                const struct zmk_behavior_binding *binding =
+                    zmk_keymap_get_layer_binding_at_idx(l, kp);
                 struct zmk_behavior_binding_setting binding_setting = {
                     .behavior_local_id = zmk_behavior_get_local_id(binding->behavior_dev),
                     .param1 = binding->param1,
@@ -564,6 +610,7 @@ int zmk_keymap_discard_changes(void) {
 
 int zmk_keymap_reset_settings(void) {
     settings_delete(LAYER_ORDER_SETTINGS_KEY);
+
     for (int l = 0; l < ZMK_KEYMAP_LAYERS_LEN; l++) {
         char layer_name_setting_name[14];
         sprintf(layer_name_setting_name, LAYER_NAME_SETTINGS_KEY, l);
@@ -600,7 +647,8 @@ int zmk_keymap_reset_settings(void) { return -ENOTSUP; }
 
 int zmk_keymap_apply_position_state(uint8_t source, zmk_keymap_layer_id_t layer_id,
                                     uint32_t position, bool pressed, int64_t timestamp) {
-    const struct zmk_behavior_binding *binding = &zmk_keymap[layer_id][position];
+    const struct zmk_behavior_binding *binding =
+        zmk_keymap_get_layer_binding_at_idx(layer_id, position);
     struct zmk_behavior_binding_event event = {
         .layer = layer_id,
         .position = position,
diff --git a/app/src/physical_layouts.c b/app/src/physical_layouts.c
index 6e719d29a2..ceb3b15181 100644
--- a/app/src/physical_layouts.c
+++ b/app/src/physical_layouts.c
@@ -18,6 +18,7 @@
 
 LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 
+#include <zmk/matrix.h>
 #include <zmk/physical_layouts.h>
 #include <zmk/event_manager.h>
 #include <zmk/events/position_state_changed.h>
@@ -213,6 +214,35 @@ static void zmk_physical_layouts_kscan_process_msgq(struct k_work *item) {
     }
 }
 
+static const struct zmk_physical_layout *get_default_layout(void) {
+    const struct zmk_physical_layout *initial;
+
+#if USE_PHY_LAYOUTS && DT_HAS_CHOSEN(zmk_physical_layout)
+    initial = &_CONCAT(_zmk_physical_layout_, DT_CHOSEN(zmk_physical_layout));
+#else
+    initial = layouts[0];
+#endif
+
+    return initial;
+}
+
+static int get_index_of_layout(const struct zmk_physical_layout *layout) {
+    for (int i = 0; i < ARRAY_SIZE(layouts); i++) {
+        if (layouts[i] == layout) {
+            return i;
+        }
+    }
+
+    return -ENODEV;
+}
+
+static uint32_t selected_to_stock_map[ZMK_KEYMAP_LEN];
+
+int zmk_physical_layouts_get_selected_to_stock_position_map(uint32_t const **map) {
+    *map = selected_to_stock_map;
+    return ZMK_KEYMAP_LEN;
+}
+
 int zmk_physical_layouts_select_layout(const struct zmk_physical_layout *dest_layout) {
     if (!dest_layout) {
         return -ENODEV;
@@ -233,6 +263,15 @@ int zmk_physical_layouts_select_layout(const struct zmk_physical_layout *dest_la
         }
     }
 
+    int new_idx = get_index_of_layout(dest_layout);
+    int stock_idx = get_index_of_layout(get_default_layout());
+    int ret = zmk_physical_layouts_get_position_map(stock_idx, new_idx, ZMK_KEYMAP_LEN,
+                                                    selected_to_stock_map);
+    if (ret < 0) {
+        LOG_ERR("Failed to refresh the selected to stock mapping (%d)", ret);
+        return ret;
+    }
+
     active = dest_layout;
 
     if (active->kscan) {
@@ -284,15 +323,7 @@ static int8_t saved_selected_index = -1;
 #endif
 
 int zmk_physical_layouts_select_initial(void) {
-    const struct zmk_physical_layout *initial;
-
-#if USE_PHY_LAYOUTS && DT_HAS_CHOSEN(zmk_physical_layout)
-    initial = &_CONCAT(_zmk_physical_layout_, DT_CHOSEN(zmk_physical_layout));
-#else
-    initial = layouts[0];
-#endif
-
-    int ret = zmk_physical_layouts_select_layout(initial);
+    int ret = zmk_physical_layouts_select_layout(get_default_layout());
 
     return ret;
 }
@@ -326,6 +357,14 @@ int zmk_physical_layouts_get_position_map(uint8_t source, uint8_t dest, size_t m
         return -EINVAL;
     }
 
+    if (source == dest) {
+        for (int i = 0; i < map_size; i++) {
+            map[i] = i;
+        }
+
+        return 0;
+    }
+
     const struct zmk_physical_layout *src_layout = layouts[source];
     const struct zmk_physical_layout *dest_layout = layouts[dest];
     int max_kp = dest_layout->keys_len;
@@ -435,6 +474,11 @@ static int zmk_physical_layouts_init(void) {
     }
 #endif // IS_ENABLED(CONFIG_PM_DEVICE)
 
+    // Initialize a sane mapping
+    for (int i = 0; i < ZMK_KEYMAP_LEN; i++) {
+        selected_to_stock_map[i] = i;
+    }
+
     return zmk_physical_layouts_select_initial();
 }
 
diff --git a/app/src/studio/keymap_subsystem.c b/app/src/studio/keymap_subsystem.c
index c3a78be01e..4d659def98 100644
--- a/app/src/studio/keymap_subsystem.c
+++ b/app/src/studio/keymap_subsystem.c
@@ -323,47 +323,6 @@ zmk_studio_Response get_physical_layouts(const zmk_studio_Request *req) {
     return KEYMAP_RESPONSE(get_physical_layouts, resp);
 }
 
-static void migrate_keymap(const uint8_t old) {
-    int new = zmk_physical_layouts_get_selected();
-
-    uint32_t new_to_old_map[ZMK_KEYMAP_LEN];
-    int layout_size =
-        zmk_physical_layouts_get_position_map(old, new, ZMK_KEYMAP_LEN, new_to_old_map);
-
-    if (layout_size < 0) {
-        return;
-    }
-
-    for (int l = 0; l < ZMK_KEYMAP_LAYERS_LEN; l++) {
-        struct zmk_behavior_binding new_layer[ZMK_KEYMAP_LEN];
-
-        for (int b = 0; b < layout_size; b++) {
-            uint32_t old_b = new_to_old_map[b];
-
-            if (old_b == UINT32_MAX) {
-                memset(&new_layer[b], 0, sizeof(struct zmk_behavior_binding));
-                continue;
-            }
-
-            const struct zmk_behavior_binding *binding =
-                zmk_keymap_get_layer_binding_at_idx(l, old_b);
-
-            if (!binding) {
-                memset(&new_layer[b], 0, sizeof(struct zmk_behavior_binding));
-                continue;
-            }
-
-            memcpy(&new_layer[b], binding, sizeof(struct zmk_behavior_binding));
-        }
-
-        for (int b = 0; b < layout_size; b++) {
-            zmk_keymap_set_layer_binding_at_idx(l, b, new_layer[b]);
-        }
-    }
-
-    // TODO: Migrate combos?
-}
-
 zmk_studio_Response set_active_physical_layout(const zmk_studio_Request *req) {
     LOG_DBG("");
     uint8_t index = (uint8_t)req->subsystem.keymap.request_type.set_active_physical_layout;
@@ -379,9 +338,7 @@ zmk_studio_Response set_active_physical_layout(const zmk_studio_Request *req) {
     }
 
     int ret = zmk_physical_layouts_select(index);
-    if (ret >= 0) {
-        migrate_keymap(old);
-    } else {
+    if (ret < 0) {
         resp.which_result = zmk_keymap_SetActivePhysicalLayoutResponse_err_tag;
         resp.result.err =
             zmk_keymap_SetActivePhysicalLayoutErrorCode_SET_ACTIVE_PHYSICAL_LAYOUT_ERR_GENERIC;