1 files changed, 82 insertions, 0 deletions
diff --git a/pkg/corpus/fenwik_test.go b/pkg/corpus/fenwik_test.go
new file mode 100644
index 000000000..177e21aa2
--- /dev/null
+++ b/pkg/corpus/fenwik_test.go
@@ -0,0 +1,82 @@
+package corpus
+
+import (
+	"math/rand"
+	"testing"
+)
+
+func naivePrefixSum(a []int64, r int) int64 {
+	var s int64
+	for i := 1; i <= r; i++ {
+		s += a[i]
+	}
+	return s
+}
+
+func TestFenwickBasic(t *testing.T) {
+	var ft fenwikTree
+
+	values := []int64{0, 5, -3, 7, 10}
+	for _, v := range values[1:] {
+		ft.addNew(v)
+	}
+
+	if ft.size != len(values)-1 {
+		t.Fatalf("size mismatch: got %d want %d", ft.size, len(values)-1)
+	}
+
+	for i := 1; i < ft.size; i++ {
+		got := ft.prefixSumTo(i)
+		want := naivePrefixSum(values, i)
+		if got != want {
+			t.Fatalf("prefixSumTo(%d) = %d, want %d", i, got, want)
+		}
+	}
+}
+
+func TestFenwickExtend(t *testing.T) {
+	var ft fenwikTree
+
+	n := 1000
+	ref := make([]int64, 1)
+
+	for i := 1; i <= n; i++ {
+		ft.addNew(int64(i))
+		ref = append(ref, int64(i))
+
+		if ft.prefixSumTo(i) != naivePrefixSum(ref, i) {
+			t.Fatalf("extend failed at i=%d, fw=%d and naive=%d", i, ft.prefixSumTo(i), naivePrefixSum(ref, i))
+		}
+	}
+}
+
+func TestFenwickRandom(t *testing.T) {
+	var ft fenwikTree
+	ft.extendAndRebuild()
+
+	ref := make([]int64, 1)
+
+	rng := rand.New(rand.NewSource(1))
+
+	for step := range 10_000 {
+		if rng.Intn(2) == 0 || ft.size <= 2 {
+			v := int64(rng.Intn(1000) - 500)
+			ft.addNew(v)
+			ref = append(ref, v)
+		} else {
+			i := rng.Intn(ft.size-1) + 1
+			v := int64(rng.Intn(1000) - 500)
+			ft.add(i, v)
+			ref[i] += v
+		}
+
+		i := rng.Intn(ft.size)
+		if ft.prefixSumTo(i) != naivePrefixSum(ref, i) {
+			t.Fatalf("mismatch at step %d, i=%d", step, i)
+		}
+
+		if ft.prefixSumTo(ft.size) != ft.fullSum {
+			t.Fatalf("fullSum mismatch %d, sm=%d, fullSum=%d, naive=%d", step, ft.prefixSumTo(ft.size), ft.fullSum, naivePrefixSum(ref, ft.size))
+		}
+	}
+}