2 files changed, 60 insertions, 50 deletions

diff --git a/src/Kaleidoscope-LayerHighlighter.cpp b/src/Kaleidoscope-LayerHighlighter.cpp
index d81c6a7..798aab6 100644
--- a/src/Kaleidoscope-LayerHighlighter.cpp
+++ b/src/Kaleidoscope-LayerHighlighter.cpp
@@ -1,50 +1,53 @@
 #include "Kaleidoscope-LayerHighlighter.h"
 
-#include "Kaleidoscope-LEDControl.h"
+#include "kaleidoscope/KeyAddr.h"                    // for KeyAddr
+#include "kaleidoscope/Runtime.h"                    // for Runtime, Runtime_
+#include "kaleidoscope/plugin/LEDControl.h"          // for LEDControl
+#include "kaleidoscope/plugin/LEDControl/LEDUtils.h" // for breathe_compute
 
-#include <Kaleidoscope.h>
+namespace kaleidoscope {
+  namespace plugin {
 
-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());
-  }
+    EventHandlerResult LayerHighlighter::onSetup(void) {
+      return EventHandlerResult::OK;
+    }
 
-  return kaleidoscope::EventHandlerResult::OK;
-}
+    EventHandlerResult LayerHighlighter::onKeyEvent(KeyEvent &event) {
+      if (event.key != LockLayer(layer)) {
+        return EventHandlerResult::OK;
+      }
 
-kaleidoscope::EventHandlerResult LayerHighlighter::afterEachCycle() {
-  if (!Layer.isOn(layer)) {
-    return kaleidoscope::EventHandlerResult::OK;
-  }
+      if (keyToggledOff(event.state)) {
+        ::LEDControl.set_mode(::LEDControl.get_mode_index());
+      }
 
-  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);
-      Key layer_key = Layer.getKey(layer, r, c);
+      return EventHandlerResult::OK;
+    }
 
-      if (k == LockLayer(layer)) {
-        cRGB lock_color = breath_compute(lockHue);
-        LEDControl.setCrgbAt(r, c, lock_color);
+    EventHandlerResult LayerHighlighter::afterEachCycle() {
+      if (!Layer.isActive(layer)) {
+        return EventHandlerResult::OK;
       }
 
-      if ((k != layer_key) || (k == Key_NoKey)) {
-        LEDControl.refreshAt(r, c);
-      } else {
-        LEDControl.setCrgbAt(r, c, color);
+      ::LEDControl.set_mode(::LEDControl.get_mode_index());
+      for (auto keyAddr : KeyAddr::all()) {
+        Key k = Layer.lookupOnActiveLayer(keyAddr);
+        Key layer_key = Layer.getKey(layer, keyAddr);
+
+        if (k == LockLayer(layer)) {
+          cRGB lock_color = breath_compute(this->lockHue);
+          ::LEDControl.setCrgbAt(keyAddr, lock_color);
+        }
+
+        if ((k != layer_key) || (k == Key_NoKey)) {
+          ::LEDControl.refreshAt(keyAddr);
+        } else {
+          ::LEDControl.setCrgbAt(keyAddr, this->color);
+        }
       }
+
+      return EventHandlerResult::OK;
     }
-  }
 
-  return kaleidoscope::EventHandlerResult::OK;
+  }
 }
diff --git a/src/Kaleidoscope-LayerHighlighter.h b/src/Kaleidoscope-LayerHighlighter.h
index ba67395..3807cd1 100644
--- a/src/Kaleidoscope-LayerHighlighter.h
+++ b/src/Kaleidoscope-LayerHighlighter.h
@@ -1,20 +1,27 @@
 #pragma once
 
-#include <Kaleidoscope.h>
+#include "kaleidoscope/KeyEvent.h"             // for KeyEvent
+#include "kaleidoscope/event_handler_result.h" // for EventHandlerResult
+#include "kaleidoscope/keyswitch_state.h"      // for keyToggledOn…
+#include "kaleidoscope/plugin.h"               // for Plugin
 
-class LayerHighlighter : public kaleidoscope::Plugin {
- public:
-  LayerHighlighter(uint8_t l): color(CRGB(160, 0, 0)), lockHue(170), layer(l) {}
+namespace kaleidoscope {
+  namespace plugin {
+    class LayerHighlighter : public Plugin {
+    public:
+      LayerHighlighter(uint8_t l): color(CRGB(160, 0, 0)), lockHue(170), layer(l) {}
 
-  cRGB color;
-  uint8_t lockHue;
+      cRGB color;
+      uint8_t lockHue;
 
-  kaleidoscope::EventHandlerResult onSetup(void);
-  kaleidoscope::EventHandlerResult onKeyswitchEvent(Key &mapped_key, byte row, byte col, uint8_t key_state);
-  kaleidoscope::EventHandlerResult afterEachCycle();
+      EventHandlerResult onSetup(void);
+      EventHandlerResult onKeyEvent(KeyEvent &event);
+      // EventHandlerResult onKeyswitchEvent(Key &mapped_key, KeyAddr key_addr, uint8_t key_state);
+      EventHandlerResult afterEachCycle();
 
- private:
-  const uint8_t layer;
-  static byte row, col;
-  uint8_t savedLEDMode;
-};
+    private:
+      const uint8_t layer;
+    };
+
+  }
+}