2 files changed, 16 insertions, 3 deletions
diff --git a/src/Kaleidoscope-LayerHighlighter.cpp b/src/Kaleidoscope-LayerHighlighter.cpp
index 3ed56cb..7c30761 100644
--- a/src/Kaleidoscope-LayerHighlighter.cpp
+++ b/src/Kaleidoscope-LayerHighlighter.cpp
@@ -8,14 +8,25 @@ kaleidoscope::EventHandlerResult LayerHighlighter::onSetup(void) {
   return kaleidoscope::EventHandlerResult::OK;
 }
 
+kaleidoscope::EventHandlerResult LayerHighlighter::onKeyswitchEvent(Key &mappedKey, byte row, byte col, uint8_t keyState) {
+  if (mappedKey != LockLayer(layer)) {
+    return kaleidoscope::EventHandlerResult::OK;
+  }
+
+  if (keyToggledOn(keyState)) {
+    savedLEDMode = LEDControl.get_mode_index();
+  } else if (keyToggledOff(keyState)) {
+    LEDControl.set_mode(LEDControl.get_mode_index());
+  }
+
+  return kaleidoscope::EventHandlerResult::OK;
+}
+
 kaleidoscope::EventHandlerResult LayerHighlighter::afterEachCycle() {
   if (!Layer.isOn(layer)) {
-    LEDControl.set_mode(LEDControl.get_mode_index());
     return kaleidoscope::EventHandlerResult::OK;
   }
 
-  LEDControl.set_mode(LEDControl.get_mode_index());
-
   for (uint8_t r = 0; r < ROWS; r++) {
     for (uint8_t c = 0; c < COLS; c++) {
       Key k = Layer.lookupOnActiveLayer(r, c);
diff --git a/src/Kaleidoscope-LayerHighlighter.h b/src/Kaleidoscope-LayerHighlighter.h
index 65ae483..ba67395 100644
--- a/src/Kaleidoscope-LayerHighlighter.h
+++ b/src/Kaleidoscope-LayerHighlighter.h
@@ -10,9 +10,11 @@ class LayerHighlighter : public kaleidoscope::Plugin {
   uint8_t lockHue;
 
   kaleidoscope::EventHandlerResult onSetup(void);
+  kaleidoscope::EventHandlerResult onKeyswitchEvent(Key &mapped_key, byte row, byte col, uint8_t key_state);
   kaleidoscope::EventHandlerResult afterEachCycle();
 
  private:
   const uint8_t layer;
   static byte row, col;
+  uint8_t savedLEDMode;
 };