1 files changed, 47 insertions, 41 deletions

diff --git a/quantum/rgb_matrix/animations/typing_heatmap_anim.h b/quantum/rgb_matrix/animations/typing_heatmap_anim.h
index 4b17c4c3ed..a05c07760e 100644
--- a/quantum/rgb_matrix/animations/typing_heatmap_anim.h
+++ b/quantum/rgb_matrix/animations/typing_heatmap_anim.h
@@ -6,30 +6,41 @@ RGB_MATRIX_EFFECT(TYPING_HEATMAP)
 #            define RGB_MATRIX_TYPING_HEATMAP_DECREASE_DELAY_MS 25
 #        endif
 
+#        ifndef RGB_MATRIX_TYPING_HEATMAP_SPREAD
+#            define RGB_MATRIX_TYPING_HEATMAP_SPREAD 40
+#        endif
+
+#        ifndef RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT
+#            define RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT 16
+#        endif
 void process_rgb_matrix_typing_heatmap(uint8_t row, uint8_t col) {
 #        ifdef RGB_MATRIX_TYPING_HEATMAP_SLIM
     // Limit effect to pressed keys
     g_rgb_frame_buffer[row][col] = qadd8(g_rgb_frame_buffer[row][col], 32);
 #        else
-    uint8_t m_row = row - 1;
-    uint8_t p_row = row + 1;
-    uint8_t m_col = col - 1;
-    uint8_t p_col = col + 1;
-
-    if (m_col < col) g_rgb_frame_buffer[row][m_col] = qadd8(g_rgb_frame_buffer[row][m_col], 16);
-    g_rgb_frame_buffer[row][col] = qadd8(g_rgb_frame_buffer[row][col], 32);
-    if (p_col < MATRIX_COLS) g_rgb_frame_buffer[row][p_col] = qadd8(g_rgb_frame_buffer[row][p_col], 16);
-
-    if (p_row < MATRIX_ROWS) {
-        if (m_col < col) g_rgb_frame_buffer[p_row][m_col] = qadd8(g_rgb_frame_buffer[p_row][m_col], 13);
-        g_rgb_frame_buffer[p_row][col] = qadd8(g_rgb_frame_buffer[p_row][col], 16);
-        if (p_col < MATRIX_COLS) g_rgb_frame_buffer[p_row][p_col] = qadd8(g_rgb_frame_buffer[p_row][p_col], 13);
+    if (g_led_config.matrix_co[row][col] == NO_LED) { // skip as pressed key doesn't have an led position
+        return;
     }
-
-    if (m_row < row) {
-        if (m_col < col) g_rgb_frame_buffer[m_row][m_col] = qadd8(g_rgb_frame_buffer[m_row][m_col], 13);
-        g_rgb_frame_buffer[m_row][col] = qadd8(g_rgb_frame_buffer[m_row][col], 16);
-        if (p_col < MATRIX_COLS) g_rgb_frame_buffer[m_row][p_col] = qadd8(g_rgb_frame_buffer[m_row][p_col], 13);
+    for (uint8_t i_row = 0; i_row < MATRIX_ROWS; i_row++) {
+        for (uint8_t i_col = 0; i_col < MATRIX_COLS; i_col++) {
+            if (g_led_config.matrix_co[i_row][i_col] == NO_LED) { // skip as target key doesn't have an led position
+                continue;
+            }
+            if (i_row == row && i_col == col) {
+                g_rgb_frame_buffer[row][col] = qadd8(g_rgb_frame_buffer[row][col], 32);
+            } else {
+#            define LED_DISTANCE(led_a, led_b) sqrt16(((int8_t)(led_a.x - led_b.x) * (int8_t)(led_a.x - led_b.x)) + ((int8_t)(led_a.y - led_b.y) * (int8_t)(led_a.y - led_b.y)))
+                uint8_t distance = LED_DISTANCE(g_led_config.point[g_led_config.matrix_co[row][col]], g_led_config.point[g_led_config.matrix_co[i_row][i_col]]);
+#            undef LED_DISTANCE
+                if (distance <= RGB_MATRIX_TYPING_HEATMAP_SPREAD) {
+                    uint8_t amount = qsub8(RGB_MATRIX_TYPING_HEATMAP_SPREAD, distance);
+                    if (amount > RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT) {
+                        amount = RGB_MATRIX_TYPING_HEATMAP_AREA_LIMIT;
+                    }
+                    g_rgb_frame_buffer[i_row][i_col] = qadd8(g_rgb_frame_buffer[i_row][i_col], amount);
+                }
+            }
+        }
     }
 #        endif
 }
@@ -40,10 +51,7 @@ static uint16_t heatmap_decrease_timer;
 static bool decrease_heatmap_values;
 
 bool TYPING_HEATMAP(effect_params_t* params) {
-    // Modified version of RGB_MATRIX_USE_LIMITS to work off of matrix row / col size
-    uint8_t led_min = RGB_MATRIX_LED_PROCESS_LIMIT * params->iter;
-    uint8_t led_max = led_min + RGB_MATRIX_LED_PROCESS_LIMIT;
-    if (led_max > sizeof(g_rgb_frame_buffer)) led_max = sizeof(g_rgb_frame_buffer);
+    RGB_MATRIX_USE_LIMITS(led_min, led_max);
 
     if (params->init) {
         rgb_matrix_set_color_all(0, 0, 0);
@@ -63,28 +71,26 @@ bool TYPING_HEATMAP(effect_params_t* params) {
     }
 
     // Render heatmap & decrease
-    for (int i = led_min; i < led_max; i++) {
-        uint8_t row = i % MATRIX_ROWS;
-        uint8_t col = i / MATRIX_ROWS;
-        uint8_t val = g_rgb_frame_buffer[row][col];
-
-        // set the pixel colour
-        uint8_t led[LED_HITS_TO_REMEMBER];
-        uint8_t led_count = rgb_matrix_map_row_column_to_led(row, col, led);
-        for (uint8_t j = 0; j < led_count; ++j) {
-            if (!HAS_ANY_FLAGS(g_led_config.flags[led[j]], params->flags)) continue;
-
-            HSV hsv = {170 - qsub8(val, 85), rgb_matrix_config.hsv.s, scale8((qadd8(170, val) - 170) * 3, rgb_matrix_config.hsv.v)};
-            RGB rgb = rgb_matrix_hsv_to_rgb(hsv);
-            rgb_matrix_set_color(led[j], rgb.r, rgb.g, rgb.b);
-        }
-
-        if (decrease_heatmap_values) {
-            g_rgb_frame_buffer[row][col] = qsub8(val, 1);
+    uint8_t count = 0;
+    for (uint8_t row = 0; row < MATRIX_ROWS && count < RGB_MATRIX_LED_PROCESS_LIMIT; row++) {
+        for (uint8_t col = 0; col < MATRIX_COLS && RGB_MATRIX_LED_PROCESS_LIMIT; col++) {
+            if (g_led_config.matrix_co[row][col] >= led_min && g_led_config.matrix_co[row][col] < led_max) {
+                count++;
+                uint8_t val = g_rgb_frame_buffer[row][col];
+                if (!HAS_ANY_FLAGS(g_led_config.flags[g_led_config.matrix_co[row][col]], params->flags)) continue;
+
+                HSV hsv = {170 - qsub8(val, 85), rgb_matrix_config.hsv.s, scale8((qadd8(170, val) - 170) * 3, rgb_matrix_config.hsv.v)};
+                RGB rgb = rgb_matrix_hsv_to_rgb(hsv);
+                rgb_matrix_set_color(g_led_config.matrix_co[row][col], rgb.r, rgb.g, rgb.b);
+
+                if (decrease_heatmap_values) {
+                    g_rgb_frame_buffer[row][col] = qsub8(val, 1);
+                }
+            }
         }
     }
 
-    return led_max < sizeof(g_rgb_frame_buffer);
+    return rgb_matrix_check_finished_leds(led_max);
 }
 
 #    endif // RGB_MATRIX_CUSTOM_EFFECT_IMPLS