1 files changed, 58 insertions, 1 deletions

diff --git a/tmk_core/common/action.c b/tmk_core/common/action.c
index 2ccc0e0b9..20e1fc614 100644
--- a/tmk_core/common/action.c
+++ b/tmk_core/common/action.c
@@ -53,6 +53,63 @@ void action_exec(keyevent_t event)
 #endif
 }
 
+#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)
+{
+    disable_action_cache = true;
+    process_action(record);
+    disable_action_cache = false;
+}
+#else
+void process_action_nocache(keyrecord_t *record)
+{
+    process_action(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) {}
 
@@ -67,7 +124,7 @@ void process_action(keyrecord_t *record)
 
     process_action_kb(record);
 
-    action_t action = layer_switch_get_action(event.key);
+    action_t action = store_or_get_action(event.pressed, event.key);
     dprint("ACTION: "); debug_action(action);
 #ifndef NO_ACTION_LAYER
     dprint(" layer_state: "); layer_debug();