Implemented more efficient and memory savvy RSI

4 files changed, 65 insertions, 25 deletions

diff --git a/example/ndpiReader.c b/example/ndpiReader.c
index 1e127d1cd..5ece4e414 100644
--- a/example/ndpiReader.c
+++ b/example/ndpiReader.c
@@ -3699,12 +3699,12 @@ void rsiUnitTest() {
   };
   u_int i, n = sizeof(v) / sizeof(unsigned int);
 
-  ndpi_init_rsi(&s, 8);
+  assert(ndpi_alloc_rsi(&s, 8) == 0);
   
   for(i=0; i<n; i++) {
     float rsi = ndpi_rsi_add_value(&s, v[i]);
 
-#if DEBUG
+#if 0
     printf("%2d) RSI = %f\n", i, rsi);
 #endif
   }
diff --git a/src/include/ndpi_api.h.in b/src/include/ndpi_api.h.in
index d90bcde19..9055546fc 100644
--- a/src/include/ndpi_api.h.in
+++ b/src/include/ndpi_api.h.in
@@ -1391,7 +1391,7 @@ extern "C" {
 
   /* ******************************* */
 
-  void  ndpi_init_rsi(struct ndpi_rsi_struct *s, u_int16_t num_learning_values);
+  int ndpi_alloc_rsi(struct ndpi_rsi_struct *s, u_int16_t num_learning_values);
   void  ndpi_free_rsi(struct ndpi_rsi_struct *s);
   float ndpi_rsi_add_value(struct ndpi_rsi_struct *s, const u_int32_t value);
   
diff --git a/src/include/ndpi_typedefs.h b/src/include/ndpi_typedefs.h
index 859a7f0e4..94556f6a3 100644
--- a/src/include/ndpi_typedefs.h
+++ b/src/include/ndpi_typedefs.h
@@ -1545,8 +1545,10 @@ struct ndpi_analyze_struct {
 /* **************************************** */
 
 struct ndpi_rsi_struct {
-  struct ndpi_analyze_struct gains, losses;
-  u_int32_t last_value;
+  u_int8_t empty:1, rsi_ready:1, _notused:6;
+  u_int16_t num_values, next_index;
+  u_int32_t *gains, *losses;
+  u_int32_t last_value, total_gains, total_losses;
 };
 
 /* **************************************** */
diff --git a/src/lib/ndpi_analyze.c b/src/lib/ndpi_analyze.c
index 32b1a4920..1ee6e99b2 100644
--- a/src/lib/ndpi_analyze.c
+++ b/src/lib/ndpi_analyze.c
@@ -821,42 +821,80 @@ int ndpi_cluster_bins(struct ndpi_bin *bins, u_int16_t num_bins,
 /* 
    RSI (Relative Strength Index)
 
-   RSI = 100 − [ 100/ (1 + (Average loss/Average gain)) ]
+   RSI = 100 − [ 100/ (1 + (Average gain/Average loss)) ]
 
    https://www.investopedia.com/terms/r/rsi.asp
 */
 
-void ndpi_init_rsi(struct ndpi_rsi_struct *s, u_int16_t num_learning_values) {
-  ndpi_init_data_analysis(&s->gains, num_learning_values);
-  ndpi_init_data_analysis(&s->losses, num_learning_values);
-  s->last_value = 0;
+int ndpi_alloc_rsi(struct ndpi_rsi_struct *s, u_int16_t num_learning_values) {
+  memset(s, 0, sizeof(struct ndpi_rsi_struct));
+
+  s->empty  = 1, s->num_values = num_learning_values;
+  s->gains  = (u_int32_t*)ndpi_calloc(num_learning_values, sizeof(u_int32_t));
+  s->losses = (u_int32_t*)ndpi_calloc(num_learning_values, sizeof(u_int32_t));
+
+  if(s->gains && s->losses) {
+    s->last_value = 0;
+    return(0);
+  } else {
+    if(s->gains)  free(s->gains);
+    if(s->losses) free(s->losses);
+    return(-1);
+  }
 }
 
 void ndpi_free_rsi(struct ndpi_rsi_struct *s) {
-  ndpi_free_data_analysis(&s->gains, 0), ndpi_free_data_analysis(&s->losses, 0);
+  ndpi_free(s->gains), ndpi_free(s->losses);
 }
 
 /*
   This function adds a new value and returns the computed RSI, or -1
   if there are too many points (< num_learning_values)
+
+  RSI < 30 (too many losses)
+  RSI > 70 (too many gains)
 */
 float ndpi_rsi_add_value(struct ndpi_rsi_struct *s, const u_int32_t value) {
-  if(s->gains.num_data_entries == 0)
-    ndpi_data_add_value(&s->gains, 0), ndpi_data_add_value(&s->losses, 0);
-  else {
-    if(value > s->last_value)    
-      ndpi_data_add_value(&s->gains, value - s->last_value), ndpi_data_add_value(&s->losses, 0);
-    else
-      ndpi_data_add_value(&s->losses, s->last_value - value), ndpi_data_add_value(&s->gains, 0);
-  }
+  float relative_strength;
   
-  s->last_value = value;
+  if(!s->empty) {
+    u_int32_t val;
 
-  if(s->gains.num_data_entries >= s->gains.num_values_array_len) {
-    float relative_strength = ndpi_data_window_average(&s->gains) / ndpi_data_window_average(&s->losses);
+    s->total_gains -= s->gains[s->next_index], s->total_losses -= s->losses[s->next_index];
+    
+    if(value > s->last_value) {
+      val = value - s->last_value;
+      s->gains[s->next_index] = val, s->losses[s->next_index] = 0;
+      s->total_gains += val;
+#ifdef DEBUG_RSI
+      printf("Gain: %u\n", val);
+#endif
+    } else {
+      val = s->last_value - value;
+      s->losses[s->next_index] = val, s->gains[s->next_index] = 0;
+      s->total_losses += val;
+#ifdef DEBUG_RSI
+      printf("Loss: %u\n", val);
+#endif
+    }
 
-    /* printf("RSI: %f\n", relative_strength); */
+#ifdef DEBUG_RSI
+    printf("[value: %u][total_gains: %u][total_losses: %u][cur_idx: %u]\n", value, s->total_gains, s->total_losses, s->next_index);
+#endif
+  }
+  
+  s->last_value = value, s->next_index = (s->next_index + 1) % s->num_values, s->empty = 0;
+  if(s->next_index == 0) s->rsi_ready = 1; /* We have completed one round */
+  
+  if(!s->rsi_ready)
+    return(0);
+  else if(s->total_losses == 0) /* Avoid division by zero (**) */
+    return(100.);
+  else {
+    relative_strength = (float)s->total_gains / (float)s->total_losses; /* (**) */
+#ifdef DEBUG_RSI
+    printf("RSI: %f\n", relative_strength);
+#endif
     return(100. - (100. / (1. + relative_strength)));
-  } else
-    return(-1); /* Too early */
+  }
 }