Added support for locking One Shot modifiers.

5 files changed, 34 insertions, 14 deletions

diff --git a/docs/key_lock.md b/docs/key_lock.md
index 03cea2089..e424061a9 100644
--- a/docs/key_lock.md
+++ b/docs/key_lock.md
@@ -8,4 +8,4 @@ Here's how to use it:
 2. Enable key lock by including `KEY_LOCK_ENABLE = yes` in your Makefile.
 3. That's it!
 
-Important: switching layers does not cancel the key lock. Additionally, key lock is only able to hold standard action keys. This does not include any of the QMK special functions, or shifted versions of keys such as KC_LPRN. If it's in the [basic_keycodes](basic_keycodes.md) list, it can be held. If it's not, then it can't be.
+Important: switching layers does not cancel the key lock. Additionally, key lock is only able to hold standard action keys and One Shot modifier keys (for example, if you have your shift defined as `OSM(KC_LSFT)`; see [One Shot Keys](quantum_keycodes.md#one-shot-keys)). This does not include any of the QMK special functions (except One Shot modifiers), or shifted versions of keys such as KC_LPRN. If it's in the [basic_keycodes](basic_keycodes.md) list, it can be held. If it's not, then it can't be.
diff --git a/keyboards/nyquist/keymaps/333fred/Makefile b/keyboards/nyquist/keymaps/333fred/Makefile
index 576bb3c30..f85443280 100644
--- a/keyboards/nyquist/keymaps/333fred/Makefile
+++ b/keyboards/nyquist/keymaps/333fred/Makefile
@@ -1,5 +1,6 @@
 KEY_LOCK_ENABLE = yes
 NKRO_ENABLE = yes
+CONSOLE_ENABLE = yes
 
 ifndef QUANTUM_DIR
         include ../../../../Makefile
diff --git a/quantum/process_keycode/process_key_lock.c b/quantum/process_keycode/process_key_lock.c
index e3632b74f..b1ba397a0 100644
--- a/quantum/process_keycode/process_key_lock.c
+++ b/quantum/process_keycode/process_key_lock.c
@@ -50,7 +50,16 @@
 uint64_t key_state[4] = { 0x0, 0x0, 0x0, 0x0 };
 bool watching = false;
 
-bool process_key_lock(uint16_t keycode, keyrecord_t *record) {
+// Translate any OSM keycodes back to their unmasked versions.
+uint16_t inline translate_keycode(uint16_t keycode) {
+    if (keycode > QK_ONE_SHOT_MOD && keycode <= QK_ONE_SHOT_MOD_MAX) {
+        return keycode ^ QK_ONE_SHOT_MOD;
+    } else {
+        return keycode;
+    }
+}
+
+bool process_key_lock(uint16_t *keycode, keyrecord_t *record) {
     // We start by categorizing the keypress event. In the event of a down
     // event, there are several possibilities:
     // 1. The key is not being locked, and we are not watching for new keys.
@@ -76,44 +85,54 @@ bool process_key_lock(uint16_t keycode, keyrecord_t *record) {
     // 2. The key is being locked. In this case, we will mask the up event
     //    by returning false, so the OS never sees that the key was released
     //    until the user pressed the key again.
+
+    // We translate any OSM keycodes back to their original keycodes, so that if the key being
+    // one-shot modded is a standard keycode, we can handle it. This is the only set of special
+    // keys that we handle
+    uint16_t translated_keycode = translate_keycode(*keycode);
+
     if (record->event.pressed) {
         // Non-standard keycode, reset and return
-        if (!(IS_STANDARD_KEYCODE(keycode) || keycode == KC_LOCK)) {
+        if (!(IS_STANDARD_KEYCODE(translated_keycode) || translated_keycode == KC_LOCK)) {
             watching = false;
             return true;
         }
 
         // If we're already watching, turn off the watch.
-        if (keycode == KC_LOCK) {
+        if (translated_keycode == KC_LOCK) {
             watching = !watching;
             return false;
         }
-        
-        if (IS_STANDARD_KEYCODE(keycode)) {
+
+        if (IS_STANDARD_KEYCODE(translated_keycode)) {
             // We check watching first. This is so that in the following scenario, we continue to
             // hold the key: KC_LOCK, KC_F, KC_LOCK, KC_F
             // If we checked in reverse order, we'd end up holding the key pressed after the second
             // KC_F press is registered, when the user likely meant to hold F
             if (watching) {
                 watching = false;
-                SET_KEY_STATE(keycode);
+                SET_KEY_STATE(translated_keycode);
+                // We need to set the keycode passed in to be the translated keycode, in case we
+                // translated a OSM back to the original keycode.
+                *keycode = translated_keycode;
                 // Let the standard keymap send the keycode down event. The up event will be masked.
                 return true;
             }
-            
-            if (KEY_STATE(keycode)) {
-                UNSET_KEY_STATE(keycode);
+
+            if (KEY_STATE(translated_keycode)) {
+                UNSET_KEY_STATE(translated_keycode);
                 // The key is already held, stop this process. The up event will be sent when the user
                 // releases the key.
                 return false;
             }
         }
-        
+
         // Either the key isn't a standard key, or we need to send the down event. Continue standard
         // processing
         return true;
     } else {
         // Stop processing if it's a standard key and we're masking up.
-        return !(IS_STANDARD_KEYCODE(keycode) && KEY_STATE(keycode));
+        return !(IS_STANDARD_KEYCODE(translated_keycode) && KEY_STATE(translated_keycode));
     }
 }
+
diff --git a/quantum/process_keycode/process_key_lock.h b/quantum/process_keycode/process_key_lock.h
index 237e103bc..876db4a32 100644
--- a/quantum/process_keycode/process_key_lock.h
+++ b/quantum/process_keycode/process_key_lock.h
@@ -19,6 +19,6 @@
 
 #include "quantum.h"
 
-bool process_key_lock(uint16_t keycode, keyrecord_t *record);
+bool process_key_lock(uint16_t *keycode, keyrecord_t *record);
 
 #endif // PROCESS_KEY_LOCK_H
diff --git a/quantum/quantum.c b/quantum/quantum.c
index c71a97bf2..0243a7e01 100644
--- a/quantum/quantum.c
+++ b/quantum/quantum.c
@@ -195,7 +195,7 @@ bool process_record_quantum(keyrecord_t *record) {
   if (!(
   #if defined(KEY_LOCK_ENABLE)
     // Must run first to be able to mask key_up events.
-    process_key_lock(keycode, record) &&
+    process_key_lock(&keycode, record) &&
   #endif
     process_record_kb(keycode, record) &&
   #if defined(MIDI_ENABLE) && defined(MIDI_ADVANCED)