Fix OSMs getting stuck (#20034)

4 files changed, 170 insertions, 19 deletions

diff --git a/quantum/action.c b/quantum/action.c
index 21fa98dc54..41686b5a5e 100644
--- a/quantum/action.c
+++ b/quantum/action.c
@@ -435,39 +435,32 @@ void process_action(keyrecord_t *record, action_t action) {
                     } else {
                         if (event.pressed) {
                             if (tap_count == 0) {
+                                // Not a tap, but a hold: register the held mod
                                 ac_dprintf("MODS_TAP: Oneshot: 0\n");
-                                register_mods(mods | get_oneshot_mods());
+                                register_mods(mods);
                             } else if (tap_count == 1) {
                                 ac_dprintf("MODS_TAP: Oneshot: start\n");
-                                set_oneshot_mods(mods | get_oneshot_mods());
+                                add_oneshot_mods(mods);
 #        if defined(ONESHOT_TAP_TOGGLE) && ONESHOT_TAP_TOGGLE > 1
                             } else if (tap_count == ONESHOT_TAP_TOGGLE) {
                                 ac_dprintf("MODS_TAP: Toggling oneshot");
                                 register_mods(mods);
-                                clear_oneshot_mods();
-                                set_oneshot_locked_mods(mods | get_oneshot_locked_mods());
+                                del_oneshot_mods(mods);
+                                add_oneshot_locked_mods(mods);
 #        endif
-                            } else {
-                                register_mods(mods | get_oneshot_mods());
                             }
                         } else {
                             if (tap_count == 0) {
-                                clear_oneshot_mods();
+                                // Release hold: unregister the held mod and its variants
                                 unregister_mods(mods);
-                            } else if (tap_count == 1) {
-                                // Retain Oneshot mods
+                                del_oneshot_mods(mods);
+                                del_oneshot_locked_mods(mods);
 #        if defined(ONESHOT_TAP_TOGGLE) && ONESHOT_TAP_TOGGLE > 1
-                                if (mods & get_mods()) {
-                                    unregister_mods(mods);
-                                    clear_oneshot_mods();
-                                    set_oneshot_locked_mods(~mods & get_oneshot_locked_mods());
-                                }
-                            } else if (tap_count == ONESHOT_TAP_TOGGLE) {
-                                // Toggle Oneshot Layer
-#        endif
-                            } else {
+                            } else if (tap_count == 1 && (mods & get_mods())) {
                                 unregister_mods(mods);
-                                clear_oneshot_mods();
+                                del_oneshot_mods(mods);
+                                del_oneshot_locked_mods(mods);
+#        endif
                             }
                         }
                     }
diff --git a/quantum/action_util.c b/quantum/action_util.c
index 7f7d32887b..034c2ecaae 100644
--- a/quantum/action_util.c
+++ b/quantum/action_util.c
@@ -46,6 +46,12 @@ static uint8_t oneshot_locked_mods = 0;
 uint8_t        get_oneshot_locked_mods(void) {
     return oneshot_locked_mods;
 }
+void add_oneshot_locked_mods(uint8_t mods) {
+    if ((oneshot_locked_mods & mods) != mods) {
+        oneshot_locked_mods |= mods;
+        oneshot_locked_mods_changed_kb(oneshot_locked_mods);
+    }
+}
 void set_oneshot_locked_mods(uint8_t mods) {
     if (mods != oneshot_locked_mods) {
         oneshot_locked_mods = mods;
@@ -58,6 +64,12 @@ void clear_oneshot_locked_mods(void) {
         oneshot_locked_mods_changed_kb(oneshot_locked_mods);
     }
 }
+void del_oneshot_locked_mods(uint8_t mods) {
+    if (oneshot_locked_mods & mods) {
+        oneshot_locked_mods &= ~mods;
+        oneshot_locked_mods_changed_kb(oneshot_locked_mods);
+    }
+}
 #    if (defined(ONESHOT_TIMEOUT) && (ONESHOT_TIMEOUT > 0))
 static uint16_t oneshot_time = 0;
 bool            has_oneshot_mods_timed_out(void) {
diff --git a/quantum/action_util.h b/quantum/action_util.h
index 0ecf15ae4b..02f6e9e6df 100644
--- a/quantum/action_util.h
+++ b/quantum/action_util.h
@@ -65,8 +65,10 @@ void    clear_oneshot_mods(void);
 bool    has_oneshot_mods_timed_out(void);
 
 uint8_t get_oneshot_locked_mods(void);
+void    add_oneshot_locked_mods(uint8_t mods);
 void    set_oneshot_locked_mods(uint8_t mods);
 void    clear_oneshot_locked_mods(void);
+void    del_oneshot_locked_mods(uint8_t mods);
 
 typedef enum { ONESHOT_PRESSED = 0b01, ONESHOT_OTHER_KEY_PRESSED = 0b10, ONESHOT_START = 0b11, ONESHOT_TOGGLED = 0b100 } oneshot_fullfillment_t;
 void    set_oneshot_layer(uint8_t layer, uint8_t state);
diff --git a/tests/basic/test_one_shot_keys.cpp b/tests/basic/test_one_shot_keys.cpp
index 2401c2c837..2a3434bf16 100644
--- a/tests/basic/test_one_shot_keys.cpp
+++ b/tests/basic/test_one_shot_keys.cpp
@@ -160,6 +160,150 @@ INSTANTIATE_TEST_CASE_P(
         ));
 // clang-format on
 
+TEST_F(OneShot, OSMChainingTwoOSMs) {
+    TestDriver driver;
+    InSequence s;
+    KeymapKey  osm_key1    = KeymapKey{0, 0, 0, OSM(MOD_LSFT), KC_LSFT};
+    KeymapKey  osm_key2    = KeymapKey{0, 0, 1, OSM(MOD_LCTL), KC_LCTL};
+    KeymapKey  regular_key = KeymapKey{0, 1, 0, KC_A};
+
+    set_keymap({osm_key1, osm_key2, regular_key});
+
+    /* Press and release OSM1 */
+    EXPECT_NO_REPORT(driver);
+    osm_key1.press();
+    run_one_scan_loop();
+    osm_key1.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Press and relesea OSM2 */
+    EXPECT_NO_REPORT(driver);
+    osm_key2.press();
+    run_one_scan_loop();
+    osm_key2.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Press regular key */
+    EXPECT_REPORT(driver, (osm_key1.report_code, osm_key2.report_code, regular_key.report_code)).Times(1);
+    regular_key.press();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Release regular key */
+    EXPECT_EMPTY_REPORT(driver);
+    regular_key.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+}
+
+TEST_F(OneShot, OSMDoubleTapNotLockingOSMs) {
+    TestDriver driver;
+    InSequence s;
+    KeymapKey  osm_key1    = KeymapKey{0, 0, 0, OSM(MOD_LSFT), KC_LSFT};
+    KeymapKey  osm_key2    = KeymapKey{0, 0, 1, OSM(MOD_LCTL), KC_LCTL};
+    KeymapKey  regular_key = KeymapKey{0, 1, 0, KC_A};
+
+    set_keymap({osm_key1, osm_key2, regular_key});
+
+    /* Press and release OSM1 */
+    EXPECT_NO_REPORT(driver);
+    osm_key1.press();
+    run_one_scan_loop();
+    osm_key1.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Press and release OSM2 twice */
+    EXPECT_NO_REPORT(driver);
+    osm_key2.press();
+    run_one_scan_loop();
+    osm_key2.release();
+    run_one_scan_loop();
+    osm_key2.press();
+    run_one_scan_loop();
+    osm_key2.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Press regular key */
+    EXPECT_REPORT(driver, (osm_key1.report_code, osm_key2.report_code, regular_key.report_code)).Times(1);
+    regular_key.press();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Release regular key */
+    EXPECT_EMPTY_REPORT(driver);
+    regular_key.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Press regular key */
+    EXPECT_REPORT(driver, (regular_key.report_code)).Times(1);
+    regular_key.press();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Release regular key */
+    EXPECT_EMPTY_REPORT(driver);
+    regular_key.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+}
+
+TEST_F(OneShot, OSMHoldNotLockingOSMs) {
+    TestDriver driver;
+    InSequence s;
+    KeymapKey  osm_key1    = KeymapKey{0, 0, 0, OSM(MOD_LSFT), KC_LSFT};
+    KeymapKey  osm_key2    = KeymapKey{0, 0, 1, OSM(MOD_LCTL), KC_LCTL};
+    KeymapKey  regular_key = KeymapKey{0, 1, 0, KC_A};
+
+    set_keymap({osm_key1, osm_key2, regular_key});
+
+    /* Press and release OSM1 */
+    EXPECT_NO_REPORT(driver);
+    osm_key1.press();
+    run_one_scan_loop();
+    osm_key1.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Press and hold OSM2 */
+    EXPECT_REPORT(driver, (osm_key1.report_code, osm_key2.report_code)).Times(1);
+    osm_key2.press();
+    run_one_scan_loop();
+    idle_for(TAPPING_TERM);
+    VERIFY_AND_CLEAR(driver);
+
+    /* Press and release regular key */
+    EXPECT_REPORT(driver, (osm_key1.report_code, osm_key2.report_code, regular_key.report_code)).Times(1);
+    EXPECT_REPORT(driver, (osm_key2.report_code)).Times(1);
+    regular_key.press();
+    run_one_scan_loop();
+    regular_key.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Release OSM2 */
+    EXPECT_EMPTY_REPORT(driver);
+    osm_key2.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Press regular key */
+    EXPECT_REPORT(driver, (regular_key.report_code)).Times(1);
+    regular_key.press();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+
+    /* Release regular key */
+    EXPECT_EMPTY_REPORT(driver);
+    regular_key.release();
+    run_one_scan_loop();
+    VERIFY_AND_CLEAR(driver);
+}
+
 TEST_F(OneShot, OSLWithAdditionalKeypress) {
     TestDriver driver;
     InSequence s;