Fixes potential wpm sampling overflow, along with code comment fixes (#15277)

Co-authored-by: Trevor Powell <trevor@vectorstorm.com.au>

1 files changed, 43 insertions, 32 deletions

diff --git a/quantum/wpm.c b/quantum/wpm.c
index 925e2c416e..62d4128b8e 100644
--- a/quantum/wpm.c
+++ b/quantum/wpm.c
@@ -22,33 +22,37 @@
 // WPM Stuff
 static uint8_t  current_wpm = 0;
 static uint32_t wpm_timer   = 0;
-#ifndef WPM_UNFILTERED
-static uint32_t smoothing_timer = 0;
-#endif
 
 /* The WPM calculation works by specifying a certain number of 'periods' inside
  * a ring buffer, and we count the number of keypresses which occur in each of
  * those periods.  Then to calculate WPM, we add up all of the keypresses in
  * the whole ring buffer, divide by the number of keypresses in a 'word', and
- * then adjust for how much time is captured by our ring buffer.  Right now
- * the ring buffer is hardcoded below to be six half-second periods, accounting
- * for a total WPM sampling period of up to three seconds of typing.
+ * then adjust for how much time is captured by our ring buffer.  The size
+ * of the ring buffer can be configured using the keymap configuration
+ * value `WPM_SAMPLE_PERIODS`.
  *
- * Whenever our WPM drops to absolute zero due to no typing occurring within
- * any contiguous three seconds, we reset and start measuring fresh,
- * which lets our WPM immediately reach the correct value even before a full
- * three second sampling buffer has been filled.
  */
 #define MAX_PERIODS (WPM_SAMPLE_PERIODS)
 #define PERIOD_DURATION (1000 * WPM_SAMPLE_SECONDS / MAX_PERIODS)
-#define LATENCY (100)
-static int8_t  period_presses[MAX_PERIODS] = {0};
+
+static int16_t period_presses[MAX_PERIODS] = {0};
 static uint8_t current_period              = 0;
 static uint8_t periods                     = 1;
 
 #if !defined(WPM_UNFILTERED)
-static uint8_t prev_wpm = 0;
-static uint8_t next_wpm = 0;
+/* LATENCY is used as part of filtering, and controls how quickly the reported
+ * WPM trails behind our actual instantaneous measured WPM value, and is
+ * defined in milliseconds.  So for LATENCY == 100, the displayed WPM is
+ * smoothed out over periods of 0.1 seconds.  This results in a nice,
+ * smoothly-moving reported WPM value which nevertheless is never more than
+ * 0.1 seconds behind the typist's actual current WPM.
+ *
+ * LATENCY is not used if WPM_UNFILTERED is defined.
+ */
+#    define LATENCY (100)
+static uint32_t smoothing_timer = 0;
+static uint8_t  prev_wpm        = 0;
+static uint8_t  next_wpm        = 0;
 #endif
 
 void    set_current_wpm(uint8_t new_wpm) { current_wpm = new_wpm; }
@@ -71,7 +75,7 @@ __attribute__((weak)) bool wpm_keycode_user(uint16_t keycode) {
     return false;
 }
 
-#ifdef WPM_ALLOW_COUNT_REGRESSION
+#if defined(WPM_ALLOW_COUNT_REGRESSION)
 __attribute__((weak)) uint8_t wpm_regress_count(uint16_t keycode) {
     bool weak_modded = (keycode >= QK_LCTL && keycode < QK_LSFT) || (keycode >= QK_RCTL && keycode < QK_RSFT);
 
@@ -95,12 +99,12 @@ __attribute__((weak)) uint8_t wpm_regress_count(uint16_t keycode) {
 // Outside 'raw' mode we smooth results over time.
 
 void update_wpm(uint16_t keycode) {
-    if (wpm_keycode(keycode)) {
+    if (wpm_keycode(keycode) && period_presses[current_period] < INT16_MAX) {
         period_presses[current_period]++;
     }
-#ifdef WPM_ALLOW_COUNT_REGRESSION
+#if defined(WPM_ALLOW_COUNT_REGRESSION)
     uint8_t regress = wpm_regress_count(keycode);
-    if (regress) {
+    if (regress && period_presses[current_period] > INT16_MIN) {
         period_presses[current_period]--;
     }
 #endif
@@ -116,32 +120,41 @@ void decay_wpm(void) {
     }
     int32_t  elapsed  = timer_elapsed32(wpm_timer);
     uint32_t duration = (((periods)*PERIOD_DURATION) + elapsed);
-    uint32_t wpm_now  = (60000 * presses) / (duration * WPM_ESTIMATED_WORD_SIZE);
-    wpm_now           = (wpm_now > 240) ? 240 : wpm_now;
+    int32_t  wpm_now  = (60000 * presses) / (duration * WPM_ESTIMATED_WORD_SIZE);
+
+    if (wpm_now < 0)  // set some reasonable WPM measurement limits
+        wpm_now = 0;
+    if (wpm_now > 240) wpm_now = 240;
 
     if (elapsed > PERIOD_DURATION) {
         current_period                 = (current_period + 1) % MAX_PERIODS;
         period_presses[current_period] = 0;
         periods                        = (periods < MAX_PERIODS - 1) ? periods + 1 : MAX_PERIODS - 1;
         elapsed                        = 0;
-        /* if (wpm_timer == 0) { */
-        wpm_timer = timer_read32();
-        /* } else { */
-        /*     wpm_timer += PERIOD_DURATION; */
-        /* } */
+        wpm_timer                      = timer_read32();
     }
     if (presses < 2)  // don't guess high WPM based on a single keypress.
         wpm_now = 0;
 
-#if defined WPM_LAUNCH_CONTROL
+#if defined(WPM_LAUNCH_CONTROL)
+    /*
+     * If the `WPM_LAUNCH_CONTROL` option is enabled, then whenever our WPM
+     * drops to absolute zero due to no typing occurring within our sample
+     * ring buffer, we reset and start measuring fresh, which lets our WPM
+     * immediately reach the correct value even before a full sampling buffer
+     * has been filled.
+     */
     if (presses == 0) {
-        current_period = 0;
-        periods        = 0;
-        wpm_now        = 0;
+        current_period    = 0;
+        periods           = 0;
+        wpm_now           = 0;
+        period_presses[0] = 0;
     }
 #endif  // WPM_LAUNCH_CONTROL
 
-#ifndef WPM_UNFILTERED
+#if defined(WPM_UNFILTERED)
+    current_wpm = wpm_now;
+#else
     int32_t latency = timer_elapsed32(smoothing_timer);
     if (latency > LATENCY) {
         smoothing_timer = timer_read32();
@@ -150,7 +163,5 @@ void decay_wpm(void) {
     }
 
     current_wpm = prev_wpm + (latency * ((int)next_wpm - (int)prev_wpm) / LATENCY);
-#else
-    current_wpm = wpm_now;
 #endif
 }