From 91e023971d06b441c5b086ffc49e148d9b0c601f Mon Sep 17 00:00:00 2001
From: Yuxuan Shui <yshuiv7@gmail.com>
Date: Mon, 12 Dec 2022 08:51:42 +0000
Subject: [PATCH] utils: add rolling_max

For tracking rolling max of a stream of integers.

Signed-off-by: Yuxuan Shui <yshuiv7@gmail.com>
---
 src/utils.c | 120 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 src/utils.h |   8 ++++
 2 files changed, 126 insertions(+), 2 deletions(-)

diff --git a/src/utils.c b/src/utils.c
index 8a27f393..8190226c 100644
--- a/src/utils.c
+++ b/src/utils.c
@@ -4,6 +4,7 @@
 
 #include "compiler.h"
 #include "string_utils.h"
+#include "test.h"
 #include "utils.h"
 
 /// Report allocation failure without allocating memory
@@ -36,8 +37,7 @@ void report_allocation_failure(const char *func, const char *file, unsigned int
 /// Calculates next closest power of two of 32bit integer n
 /// ref: https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
 ///
-int next_power_of_two(int n)
-{
+int next_power_of_two(int n) {
 	n--;
 	n |= n >> 1;
 	n |= n >> 2;
@@ -48,4 +48,120 @@ int next_power_of_two(int n)
 	return n;
 }
 
+/// Track the rolling maximum of a stream of integers.
+struct rolling_max {
+	/// A priority queue holding the indices of the maximum element candidates.
+	/// The head of the queue is the index of the maximum element.
+	/// The indices in the queue are the "original" indices.
+	///
+	/// There are only `capacity` elements in `elem`, all previous elements are
+	/// discarded. But the discarded elements' indices are not forgotten, that's why
+	/// it's called the "original" indices.
+	int *p;
+	int p_head, np;
+
+	/// The elemets
+	int *elem;
+	int elem_head, nelem;
+
+	int window_size;
+};
+
+void rolling_max_destroy(struct rolling_max *rm) {
+	free(rm->elem);
+	free(rm->p);
+	free(rm);
+}
+
+struct rolling_max *rolling_max_new(int size) {
+	auto rm = ccalloc(1, struct rolling_max);
+	if (!rm) {
+		return NULL;
+	}
+
+	rm->p = ccalloc(size, int);
+	rm->elem = ccalloc(size, int);
+	rm->window_size = size;
+	if (!rm->p || !rm->elem) {
+		goto err;
+	}
+
+	return rm;
+
+err:
+	rolling_max_destroy(rm);
+	return NULL;
+}
+
+void rolling_max_reset(struct rolling_max *rm) {
+	rm->p_head = 0;
+	rm->np = 0;
+	rm->nelem = 0;
+	rm->elem_head = 0;
+}
+
+void rolling_max_push(struct rolling_max *rm, int val) {
+#define IDX(n) ((n) % rm->window_size)
+	if (rm->nelem == rm->window_size) {
+		auto old_head = rm->elem_head;
+		// Discard the oldest element.
+		// rm->elem.pop_front();
+		rm->nelem--;
+		rm->elem_head = IDX(rm->elem_head + 1);
+
+		// Remove discarded element from the priority queue too.
+		assert(rm->np);
+		if (rm->p[rm->p_head] == old_head) {
+			// rm->p.pop_front()
+			rm->p_head = IDX(rm->p_head + 1);
+			rm->np--;
+		}
+	}
+
+	// Add the new element to the queue.
+	// rm->elem.push_back(val)
+	rm->elem[IDX(rm->elem_head + rm->nelem)] = val;
+	rm->nelem++;
+
+	// Update the prority queue.
+	// Remove all elements smaller than the new element from the queue. Because
+	// the new element will become the maximum element before them, and since they
+	// come b1efore the new element, they will have been popped before the new
+	// element, so they will never become the maximum element.
+	while (rm->np) {
+		int p_tail = IDX(rm->p_head + rm->np - 1);
+		if (rm->elem[rm->p[p_tail]] > val) {
+			break;
+		}
+		// rm->p.pop_back()
+		rm->np--;
+	}
+	// Add the new element to the end of the queue.
+	// rm->p.push_back(rm->start_index + rm->nelem - 1)
+	rm->p[IDX(rm->p_head + rm->np)] = IDX(rm->elem_head + rm->nelem - 1);
+	rm->np++;
+#undef IDX
+}
+
+int rolling_max_get_max(struct rolling_max *rm) {
+	if (rm->np == 0) {
+		return INT_MIN;
+	}
+	return rm->elem[rm->p[rm->p_head]];
+}
+
+TEST_CASE(rolling_max_test) {
+#define NELEM 15
+	auto rm = rolling_max_new(3);
+	const int data[NELEM] = {1, 2, 3, 1, 4, 5, 2, 3, 6, 5, 4, 3, 2, 0, 0};
+	const int expected_max[NELEM] = {1, 2, 3, 3, 4, 5, 5, 5, 6, 6, 6, 5, 4, 3, 2};
+	int max[NELEM] = {0};
+	for (int i = 0; i < NELEM; i++) {
+		rolling_max_push(rm, data[i]);
+		max[i] = rolling_max_get_max(rm);
+	}
+	TEST_TRUE(memcmp(max, expected_max, sizeof(max)) == 0);
+#undef NELEM
+}
+
 // vim: set noet sw=8 ts=8 :
diff --git a/src/utils.h b/src/utils.h
index a35bfa84..a1a3d34f 100644
--- a/src/utils.h
+++ b/src/utils.h
@@ -289,6 +289,14 @@ static inline void free_charpp(char **str) {
 ///
 int next_power_of_two(int n);
 
+struct rolling_max;
+
+struct rolling_max *rolling_max_new(int window_size);
+void rolling_max_free(struct rolling_max *rm);
+void rolling_max_reset(struct rolling_max *rm);
+void rolling_max_push(struct rolling_max *rm, int val);
+int rolling_max_get_max(struct rolling_max *rm);
+
 // Some versions of the Android libc do not have timespec_get(), use
 // clock_gettime() instead.
 #ifdef __ANDROID__