3 files changed, 71 insertions, 41 deletions

diff --git a/tmk_core/common/action.c b/tmk_core/common/action.c
index 20e1fc6149..6aa6dc2601 100644
--- a/tmk_core/common/action.c
+++ b/tmk_core/common/action.c
@@ -55,7 +55,6 @@ void action_exec(keyevent_t event)
 
 #if !defined(NO_ACTION_LAYER) && defined(PREVENT_STUCK_MODIFIERS)
 bool disable_action_cache = false;
-uint8_t source_layers_cache[5][(MATRIX_ROWS * MATRIX_COLS + 7) / 8] = {0};
 
 void process_action_nocache(keyrecord_t *record)
 {
@@ -70,46 +69,6 @@ void process_action_nocache(keyrecord_t *record)
 }
 #endif
 
-/*
- * Make sure the action triggered when the key is released is the same
- * one as the one triggered on press. It's important for the mod keys
- * when the layer is switched after the down event but before the up
- * event as they may get stuck otherwise.
- */
-action_t store_or_get_action(bool pressed, keypos_t key)
-{
-#if !defined(NO_ACTION_LAYER) && defined(PREVENT_STUCK_MODIFIERS)
-    if (disable_action_cache) {
-        return layer_switch_get_action(key);
-    }
-    const uint8_t key_number = key.col + (key.row * MATRIX_COLS);
-    const uint8_t storage_row = key_number / 8;
-    const uint8_t storage_bit = key_number % 8;
-    uint8_t layer;
-    if (pressed) {
-        layer = layer_switch_get_layer(key);
-        for (uint8_t bit_number = 0; bit_number < 5; bit_number++) {
-            source_layers_cache[bit_number][storage_row] ^=
-                (-((layer & (1U << bit_number)) != 0)
-                 ^ source_layers_cache[bit_number][storage_row])
-                & (1U << storage_bit);
-        }
-    }
-    else {
-        layer = 0;
-        for (uint8_t bit_number = 0; bit_number < 5; bit_number++) {
-            layer |=
-                ((source_layers_cache[bit_number][storage_row]
-                  & (1U << storage_bit)) != 0)
-                << bit_number;
-        }
-    }
-    return action_for_key(layer, key);
-#else
-    return layer_switch_get_action(key);
-#endif
-}
-
 __attribute__ ((weak))
 void process_action_kb(keyrecord_t *record) {}
 
diff --git a/tmk_core/common/action_layer.c b/tmk_core/common/action_layer.c
index 76164adb5d..fc721a7323 100644
--- a/tmk_core/common/action_layer.c
+++ b/tmk_core/common/action_layer.c
@@ -110,6 +110,69 @@ void layer_debug(void)
 }
 #endif
 
+#if !defined(NO_ACTION_LAYER) && defined(PREVENT_STUCK_MODIFIERS)
+uint8_t source_layers_cache[MAX_LAYER_BITS][(MATRIX_ROWS * MATRIX_COLS + 7) / 8] = {0};
+
+void update_source_layers_cache(keypos_t key, uint8_t layer)
+{
+    const uint8_t key_number = key.col + (key.row * MATRIX_COLS);
+    const uint8_t storage_row = key_number / 8;
+    const uint8_t storage_bit = key_number % 8;
+
+    for (uint8_t bit_number = 0; bit_number < MAX_LAYER_BITS; bit_number++) {
+        source_layers_cache[bit_number][storage_row] ^=
+            (-((layer & (1U << bit_number)) != 0)
+             ^ source_layers_cache[bit_number][storage_row])
+            & (1U << storage_bit);
+    }
+}
+
+uint8_t read_source_layers_cache(keypos_t key)
+{
+    const uint8_t key_number = key.col + (key.row * MATRIX_COLS);
+    const uint8_t storage_row = key_number / 8;
+    const uint8_t storage_bit = key_number % 8;
+    uint8_t layer = 0;
+
+    for (uint8_t bit_number = 0; bit_number < MAX_LAYER_BITS; bit_number++) {
+        layer |=
+            ((source_layers_cache[bit_number][storage_row]
+              & (1U << storage_bit)) != 0)
+            << bit_number;
+    }
+
+    return layer;
+}
+#endif
+
+/*
+ * Make sure the action triggered when the key is released is the same
+ * one as the one triggered on press. It's important for the mod keys
+ * when the layer is switched after the down event but before the up
+ * event as they may get stuck otherwise.
+ */
+action_t store_or_get_action(bool pressed, keypos_t key)
+{
+#if !defined(NO_ACTION_LAYER) && defined(PREVENT_STUCK_MODIFIERS)
+    if (disable_action_cache) {
+        return layer_switch_get_action(key);
+    }
+
+    uint8_t layer;
+
+    if (pressed) {
+        layer = layer_switch_get_layer(key);
+        update_source_layers_cache(key, layer);
+    }
+    else {
+        layer = read_source_layers_cache(key);
+    }
+    return action_for_key(layer, key);
+#else
+    return layer_switch_get_action(key);
+#endif
+}
+
 
 int8_t layer_switch_get_layer(keypos_t key)
 {
diff --git a/tmk_core/common/action_layer.h b/tmk_core/common/action_layer.h
index 1a313a2590..3a4b1e3349 100644
--- a/tmk_core/common/action_layer.h
+++ b/tmk_core/common/action_layer.h
@@ -70,6 +70,14 @@ void layer_xor(uint32_t state);
 #define layer_debug()
 #endif
 
+/* pressed actions cache */
+#if !defined(NO_ACTION_LAYER) && defined(PREVENT_STUCK_MODIFIERS)
+/* The number of bits needed to represent the layer number: log2(32). */
+#define MAX_LAYER_BITS 5
+void update_source_layers_cache(keypos_t key, uint8_t layer);
+uint8_t read_source_layers_cache(keypos_t key);
+#endif
+action_t store_or_get_action(bool pressed, keypos_t key);
 
 /* return the topmost non-transparent layer currently associated with key */
 int8_t layer_switch_get_layer(keypos_t key);