6 files changed, 258 insertions, 45 deletions

diff --git a/keyboards/ergodox/infinity/config.h b/keyboards/ergodox/infinity/config.h
index 9e264083b..6cde193e1 100644
--- a/keyboards/ergodox/infinity/config.h
+++ b/keyboards/ergodox/infinity/config.h
@@ -54,6 +54,8 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 // The visualizer needs gfx thread priorities
 #define VISUALIZER_THREAD_PRIORITY (NORMAL_PRIORITY - 2)
 
+#define VISUALIZER_USER_DATA_SIZE 16
+
 /*
  * Feature disable options
  *  These options are also useful to firmware size reduction.
diff --git a/keyboards/ergodox/infinity/infinity.c b/keyboards/ergodox/infinity/infinity.c
index ecc072abb..62259ed3f 100644
--- a/keyboards/ergodox/infinity/infinity.c
+++ b/keyboards/ergodox/infinity/infinity.c
@@ -126,34 +126,48 @@ void matrix_scan_kb(void) {
         matrix_scan_user();
 }
 
+__attribute__ ((weak))
 void ergodox_board_led_on(void){
 }
 
+__attribute__ ((weak))
 void ergodox_right_led_1_on(void){
 }
 
+__attribute__ ((weak))
 void ergodox_right_led_2_on(void){
 }
 
+__attribute__ ((weak))
 void ergodox_right_led_3_on(void){
 }
 
-void ergodox_right_led_on(uint8_t led){
-}
-
+__attribute__ ((weak))
 void ergodox_board_led_off(void){
 }
 
+__attribute__ ((weak))
 void ergodox_right_led_1_off(void){
 }
 
+__attribute__ ((weak))
 void ergodox_right_led_2_off(void){
 }
 
+__attribute__ ((weak))
 void ergodox_right_led_3_off(void){
 }
 
-void ergodox_right_led_off(uint8_t led){
+__attribute__ ((weak))
+void ergodox_right_led_1_set(uint8_t n) {
+}
+
+__attribute__ ((weak))
+void ergodox_right_led_2_set(uint8_t n) {
+}
+
+__attribute__ ((weak))
+void ergodox_right_led_3_set(uint8_t n) {
 }
 
 #ifdef ONEHAND_ENABLE
diff --git a/keyboards/ergodox/infinity/infinity.h b/keyboards/ergodox/infinity/infinity.h
index fec9e565c..73a0f4bf7 100644
--- a/keyboards/ergodox/infinity/infinity.h
+++ b/keyboards/ergodox/infinity/infinity.h
@@ -7,13 +7,38 @@ void ergodox_board_led_on(void);
 void ergodox_right_led_1_on(void);
 void ergodox_right_led_2_on(void);
 void ergodox_right_led_3_on(void);
-void ergodox_right_led_on(uint8_t led);
+
+inline void ergodox_right_led_on(uint8_t led) {
+    switch (led) {
+    case 0:
+        ergodox_right_led_1_on();
+        break;
+    case 1:
+        ergodox_right_led_2_on();
+        break;
+    case 2:
+        ergodox_right_led_3_on();
+        break;
+    }
+}
 
 void ergodox_board_led_off(void);
 void ergodox_right_led_1_off(void);
 void ergodox_right_led_2_off(void);
 void ergodox_right_led_3_off(void);
-void ergodox_right_led_off(uint8_t led);
+inline void ergodox_right_led_off(uint8_t led) {
+    switch (led) {
+    case 0:
+        ergodox_right_led_1_off();
+        break;
+    case 1:
+        ergodox_right_led_2_off();
+        break;
+    case 2:
+        ergodox_right_led_3_off();
+        break;
+    }
+}
 
 inline void ergodox_led_all_on(void)
 {
@@ -31,36 +56,22 @@ inline void ergodox_led_all_off(void)
     ergodox_right_led_3_off();
 }
 
-inline void ergodox_right_led_1_set(uint8_t n){
-        if (n) {
-                ergodox_right_led_1_on();
-        } else {
-                ergodox_right_led_1_off();
-        }
-}
-
-inline void ergodox_right_led_2_set(uint8_t n){
-        if (n) {
-                ergodox_right_led_2_on();
-        } else {
-                ergodox_right_led_2_off();
-        }
-}
-
-inline void ergodox_right_led_3_set(uint8_t n){
-        if (n) {
-                ergodox_right_led_3_on();
-        } else {
-                ergodox_right_led_3_off();
-        }
-}
+void ergodox_right_led_1_set(uint8_t n);
+void ergodox_right_led_2_set(uint8_t n);
+void ergodox_right_led_3_set(uint8_t n);
 
 inline void ergodox_right_led_set(uint8_t led, uint8_t n){
-        if (n) {
-                ergodox_right_led_on(led);
-        } else {
-                ergodox_right_led_off(led);
-        }
+    switch (led) {
+    case 0:
+        ergodox_right_led_1_set(n);
+        break;
+    case 1:
+        ergodox_right_led_2_set(n);
+        break;
+    case 2:
+        ergodox_right_led_3_set(n);
+        break;
+    }
 }
 
 inline void ergodox_led_all_set(uint8_t n) {
diff --git a/keyboards/ergodox/infinity/visualizer.c b/keyboards/ergodox/infinity/visualizer.c
index 6af3a72ee..7bc42fbc7 100644
--- a/keyboards/ergodox/infinity/visualizer.c
+++ b/keyboards/ergodox/infinity/visualizer.c
@@ -27,6 +27,7 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #endif
 
 #include "visualizer.h"
+#include "system/serial_link.h"
 
 // To generate an image array like this
 // Ensure the image is 128 x 32 or smaller
@@ -73,6 +74,15 @@ static const uint8_t image_data_lcd_logo[512] = {
 static const uint32_t logo_background_color = LCD_COLOR(0x00, 0x00, 0xFF);
 static const uint32_t initial_color = LCD_COLOR(0, 0, 0);
 
+static const uint32_t led_emulation_colors[4] = {
+    LCD_COLOR(0, 0, 0),
+    LCD_COLOR(255, 255, 255),
+    LCD_COLOR(84, 255, 255),
+    LCD_COLOR(168, 255, 255),
+};
+
+static uint32_t next_led_target_color = 0;
+
 typedef enum {
     LCD_STATE_INITIAL,
     LCD_STATE_LAYER_BITMAP,
@@ -81,6 +91,19 @@ typedef enum {
 
 static lcd_state_t lcd_state = LCD_STATE_INITIAL;
 
+typedef struct {
+    uint8_t led_on;
+    uint8_t led1;
+    uint8_t led2;
+    uint8_t led3;
+} visualizer_user_data_t;
+
+// Don't access from visualization function, use the visualizer state instead
+static visualizer_user_data_t user_data_keyboard = {};
+
+_Static_assert(sizeof(visualizer_user_data_t) <= VISUALIZER_USER_DATA_SIZE,
+    "Please increase the VISUALIZER_USER_DATA_SIZE");
+
 bool display_logo(keyframe_animation_t* animation, visualizer_state_t* state) {
     (void)state;
     (void)animation;
@@ -117,16 +140,27 @@ static keyframe_animation_t startup_animation = {
 };
 
 // The color animation animates the LCD color when you change layers
-static keyframe_animation_t color_animation = {
-    .num_frames = 2,
+static keyframe_animation_t one_led_color = {
+    .num_frames = 1,
     .loop = false,
-    // Note that there's a 200 ms no-operation frame,
-    // this prevents the color from changing when activating the layer
-    // momentarily
-    .frame_lengths = {gfxMillisecondsToTicks(200), gfxMillisecondsToTicks(500)},
-    .frame_functions = {keyframe_no_operation, keyframe_animate_backlight_color},
+    .frame_lengths = {gfxMillisecondsToTicks(0)},
+    .frame_functions = {keyframe_set_backlight_color},
 };
 
+bool swap_led_target_color(keyframe_animation_t* animation, visualizer_state_t* state) {
+    uint32_t temp = next_led_target_color;
+    next_led_target_color = state->target_lcd_color;
+    state->target_lcd_color = temp;
+    return false;
+}
+
+// The color animation animates the LCD color when you change layers
+static keyframe_animation_t two_led_colors = {
+    .num_frames = 2,
+    .loop = true,
+    .frame_lengths = {gfxMillisecondsToTicks(1000), gfxMillisecondsToTicks(0)},
+    .frame_functions = {keyframe_set_backlight_color, swap_led_target_color},
+};
 
 // The LCD animation alternates between the layer name display and a
 // bitmap that displays all active layers
@@ -177,6 +211,45 @@ void initialize_user_visualizer(visualizer_state_t* state) {
     start_keyframe_animation(&startup_animation);
 }
 
+static const uint32_t red;
+static const uint32_t green;
+static const uint32_t blue;
+
+inline bool is_led_on(visualizer_user_data_t* user_data, uint8_t num) {
+    return user_data->led_on & (1u << num);
+}
+
+static uint8_t get_led_index_master(visualizer_user_data_t* user_data) {
+    for (int i=0; i < 3; i++) {
+        if (is_led_on(user_data, i)) {
+            return i + 1;
+        }
+    }
+    return 0;
+}
+
+static uint8_t get_led_index_slave(visualizer_user_data_t* user_data) {
+    uint8_t master_index = get_led_index_master(user_data);
+    if (master_index!=0) {
+        for (int i=master_index; i < 3; i++) {
+            if (is_led_on(user_data, i)) {
+                return i + 1;
+            }
+        }
+    }
+
+    return 0;
+}
+
+static uint8_t get_secondary_led_index(visualizer_user_data_t* user_data) {
+    if (is_led_on(user_data, 0) &&
+            is_led_on(user_data, 1) &&
+            is_led_on(user_data, 2)) {
+        return 3;
+    }
+    return 0;
+}
+
 void update_user_visualizer_state(visualizer_state_t* state, visualizer_keyboard_status_t prev_status) {
     // Check the status here to start and stop animations
     // You might have to save some state, like the current animation here so that you can start the right
@@ -186,9 +259,38 @@ void update_user_visualizer_state(visualizer_state_t* state, visualizer_keyboard
     // This is also important because the slave won't have access to the active layer for example outside the
     // status.
 
-    if (lcd_state == LCD_STATE_INITIAL) {
-        state->target_lcd_color = LCD_COLOR(0x40, 0xB0, 0xFF);
-        start_keyframe_animation(&color_animation);
+
+    visualizer_user_data_t* user_data_new = (visualizer_user_data_t*)state->status.user_data;
+    visualizer_user_data_t* user_data_old = (visualizer_user_data_t*)prev_status.user_data;
+
+    uint8_t new_index;
+    uint8_t old_index;
+
+    if (is_serial_link_master()) {
+        new_index = get_led_index_master(user_data_new);
+        old_index = get_led_index_master(user_data_old);
+    }
+    else {
+        new_index = get_led_index_slave(user_data_new);
+        old_index = get_led_index_slave(user_data_old);
+    }
+    uint8_t new_secondary_index = get_secondary_led_index(user_data_new);
+    uint8_t old_secondary_index = get_secondary_led_index(user_data_old);
+
+    if (lcd_state == LCD_STATE_INITIAL ||
+            new_index != old_index ||
+            new_secondary_index != old_secondary_index) {
+
+        if (new_secondary_index != 0) {
+            state->target_lcd_color = led_emulation_colors[new_index];
+            next_led_target_color = led_emulation_colors[new_secondary_index];
+            stop_keyframe_animation(&one_led_color);
+            start_keyframe_animation(&two_led_colors);
+        } else {
+            state->target_lcd_color = led_emulation_colors[new_index];
+            stop_keyframe_animation(&two_led_colors);
+            start_keyframe_animation(&one_led_color);
+        }
     }
 
     if (state->status.leds) {
@@ -233,3 +335,56 @@ void user_visualizer_resume(visualizer_state_t* state) {
     lcd_state = LCD_STATE_INITIAL;
     start_keyframe_animation(&resume_animation);
 }
+
+void ergodox_board_led_on(void){
+    // No board led support
+}
+
+void ergodox_right_led_1_on(void){
+    user_data_keyboard.led_on |= (1u << 0);
+    visualizer_set_user_data(&user_data_keyboard);
+}
+
+void ergodox_right_led_2_on(void){
+    user_data_keyboard.led_on |= (1u << 1);
+    visualizer_set_user_data(&user_data_keyboard);
+}
+
+void ergodox_right_led_3_on(void){
+    user_data_keyboard.led_on |= (1u << 2);
+    visualizer_set_user_data(&user_data_keyboard);
+}
+
+void ergodox_board_led_off(void){
+    // No board led support
+}
+
+void ergodox_right_led_1_off(void){
+    user_data_keyboard.led_on &= ~(1u << 0);
+    visualizer_set_user_data(&user_data_keyboard);
+}
+
+void ergodox_right_led_2_off(void){
+    user_data_keyboard.led_on &= ~(1u << 1);
+    visualizer_set_user_data(&user_data_keyboard);
+}
+
+void ergodox_right_led_3_off(void){
+    user_data_keyboard.led_on &= ~(1u << 2);
+    visualizer_set_user_data(&user_data_keyboard);
+}
+
+void ergodox_right_led_1_set(uint8_t n) {
+    user_data_keyboard.led1 = n;
+    visualizer_set_user_data(&user_data_keyboard);
+}
+
+void ergodox_right_led_2_set(uint8_t n) {
+    user_data_keyboard.led2 = n;
+    visualizer_set_user_data(&user_data_keyboard);
+}
+
+void ergodox_right_led_3_set(uint8_t n) {
+    user_data_keyboard.led3 = n;
+    visualizer_set_user_data(&user_data_keyboard);
+}
diff --git a/quantum/visualizer/visualizer.c b/quantum/visualizer/visualizer.c
index 62ebebcee..bff0df494 100644
--- a/quantum/visualizer/visualizer.c
+++ b/quantum/visualizer/visualizer.c
@@ -64,6 +64,9 @@ static visualizer_keyboard_status_t current_status = {
     .mods = 0xFF,
     .leds = 0xFFFFFFFF,
     .suspended = false,
+#ifdef VISUALIZER_USER_DATA_SIZE
+    .user_data = {0}
+#endif
 };
 
 static bool same_status(visualizer_keyboard_status_t* status1, visualizer_keyboard_status_t* status2) {
@@ -71,11 +74,19 @@ static bool same_status(visualizer_keyboard_status_t* status1, visualizer_keyboa
         status1->default_layer == status2->default_layer &&
         status1->mods == status2->mods &&
         status1->leds == status2->leds &&
-        status1->suspended == status2->suspended;
+        status1->suspended == status2->suspended
+#ifdef VISUALIZER_USER_DATA_SIZE
+        && memcmp(status1->user_data, status2->user_data, VISUALIZER_USER_DATA_SIZE) == 0
+#endif
+    ;
 }
 
 static bool visualizer_enabled = false;
 
+#ifdef VISUALIZER_USER_DATA_SIZE
+static uint8_t user_data[VISUALIZER_USER_DATA_SIZE];
+#endif
+
 #define MAX_SIMULTANEOUS_ANIMATIONS 4
 static keyframe_animation_t* animations[MAX_SIMULTANEOUS_ANIMATIONS] = {};
 
@@ -431,6 +442,9 @@ static DECLARE_THREAD_FUNCTION(visualizerThread, arg) {
         .mods = 0xFF,
         .leds = 0xFFFFFFFF,
         .suspended = false,
+#ifdef VISUALIZER_USER_DATA_SIZE
+        .user_data = {0},
+#endif
     };
 
     visualizer_state_t state = {
@@ -590,6 +604,12 @@ uint8_t visualizer_get_mods() {
   return mods;
 }
 
+#ifdef VISUALIZER_USER_DATA_SIZE
+void visualizer_set_user_data(void* u) {
+    memcpy(user_data, u, VISUALIZER_USER_DATA_SIZE);
+}
+#endif
+
 void visualizer_update(uint32_t default_state, uint32_t state, uint8_t mods, uint32_t leds) {
     // Note that there's a small race condition here, the thread could read
     // a state where one of these are set but not the other. But this should
@@ -618,6 +638,9 @@ void visualizer_update(uint32_t default_state, uint32_t state, uint8_t mods, uin
             .leds = leds,
             .suspended = current_status.suspended,
         };
+#ifdef VISUALIZER_USER_DATA_SIZE
+       memcpy(new_status.user_data, user_data, VISUALIZER_USER_DATA_SIZE);
+#endif
         if (!same_status(&current_status, &new_status)) {
             changed = true;
             current_status = new_status;
diff --git a/quantum/visualizer/visualizer.h b/quantum/visualizer/visualizer.h
index 2c81cb9f7..beb8f2ab4 100644
--- a/quantum/visualizer/visualizer.h
+++ b/quantum/visualizer/visualizer.h
@@ -68,6 +68,9 @@ typedef struct {
     uint8_t mods;
     uint32_t leds; // See led.h for available statuses
     bool suspended;
+#ifdef VISUALIZER_USER_DATA_SIZE
+    uint8_t user_data[VISUALIZER_USER_DATA_SIZE];
+#endif
 } visualizer_keyboard_status_t;
 
 // The state struct is used by the various keyframe functions
@@ -146,6 +149,11 @@ bool keyframe_enable_lcd_and_backlight(keyframe_animation_t* animation, visualiz
 // directly from the initalize_user_visualizer function (the animation can be null)
 bool enable_visualization(keyframe_animation_t* animation, visualizer_state_t* state);
 
+// The master can set userdata which will be transferred to the slave
+#ifdef VISUALIZER_USER_DATA_SIZE
+void visualizer_set_user_data(void* user_data);
+#endif
+
 // These functions have to be implemented by the user
 void initialize_user_visualizer(visualizer_state_t* state);
 void update_user_visualizer_state(visualizer_state_t* state, visualizer_keyboard_status_t prev_status);