From 2b597fb7802a95a39ea2107e289cc54da4294b17 Mon Sep 17 00:00:00 2001
From: Jakub Sztandera <kubuxu@protocol.ai>
Date: Tue, 23 Jun 2020 16:18:52 +0200
Subject: [PATCH] Add tests for lambda

Signed-off-by: Jakub Sztandera <kubuxu@protocol.ai>
---
 chain/types/electionproof.go                  |  7 +++--
 chain/types/electionproof_test.go             | 28 ++++++++++++++++++-
 .../testdata/TestLambdaFunction/10-100.golden |  1 +
 .../TestLambdaFunction/1024-2048.golden       |  1 +
 ...2000000000000000-100000000000000000.golden |  1 +
 5 files changed, 34 insertions(+), 4 deletions(-)
 create mode 100644 chain/types/testdata/TestLambdaFunction/10-100.golden
 create mode 100644 chain/types/testdata/TestLambdaFunction/1024-2048.golden
 create mode 100644 chain/types/testdata/TestLambdaFunction/2000000000000000-100000000000000000.golden

diff --git a/chain/types/electionproof.go b/chain/types/electionproof.go
index 96acd3ceed7..1e12d1182bf 100644
--- a/chain/types/electionproof.go
+++ b/chain/types/electionproof.go
@@ -112,7 +112,7 @@ type poiss struct {
 	pmf  *big.Int
 	icdf *big.Int
 
-	tmp *big.Int // temporary variable for optmization
+	tmp *big.Int // temporary variable for optimization
 
 	k uint64
 }
@@ -125,7 +125,8 @@ func newPoiss(lambda *big.Int) (*poiss, *big.Int) {
 
 	// pmf(k) = (lambda^k)*(e^lambda) / k!
 	// k = 0 here, so it simplifies to just e^-lambda
-	pmf := expneg(lambda) // Q.256
+	elam := expneg(lambda) // Q.256
+	pmf := new(big.Int).Set(elam)
 
 	// icdf(k) = 1 - ∑ᵏᵢ₌₀ pmf(i)
 	// icdf(0) = 1 - pmf(0)
@@ -139,7 +140,7 @@ func newPoiss(lambda *big.Int) (*poiss, *big.Int) {
 		lam: lambda,
 		pmf: pmf,
 
-		tmp:  new(big.Int),
+		tmp:  elam,
 		icdf: icdf,
 
 		k: k,
diff --git a/chain/types/electionproof_test.go b/chain/types/electionproof_test.go
index 4303a17832b..27ac634570c 100644
--- a/chain/types/electionproof_test.go
+++ b/chain/types/electionproof_test.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/xorcare/golden"
 )
 
@@ -27,7 +28,6 @@ func TestPoissonFunction(t *testing.T) {
 	for _, test := range tests {
 		test := test
 		t.Run(fmt.Sprintf("lam-%d-%d", test.lambdaBase, test.lambdaShift), func(t *testing.T) {
-
 			b := &bytes.Buffer{}
 			b.WriteString("icdf\n")
 
@@ -47,6 +47,32 @@ func TestPoissonFunction(t *testing.T) {
 	}
 }
 
+func TestLambdaFunction(t *testing.T) {
+	tests := []struct {
+		power      string
+		totalPower string
+		target     float64
+	}{
+		{"10", "100", .1 * 5.},
+		{"1024", "2048", 0.5 * 5.},
+		{"2000000000000000", "100000000000000000", 0.02 * 5.},
+	}
+
+	for _, test := range tests {
+		test := test
+		t.Run(fmt.Sprintf("%s-%s", test.power, test.totalPower), func(t *testing.T) {
+			pow, ok := new(big.Int).SetString(test.power, 10)
+			assert.True(t, ok)
+			total, ok := new(big.Int).SetString(test.totalPower, 10)
+			assert.True(t, ok)
+			lam := lambda(pow, total)
+			assert.Equal(t, test.target, q256ToF(lam))
+			golden.Assert(t, []byte(lam.String()))
+		})
+	}
+
+}
+
 func q256ToF(x *big.Int) float64 {
 	deno := big.NewInt(1)
 	deno = deno.Lsh(deno, 256)
diff --git a/chain/types/testdata/TestLambdaFunction/10-100.golden b/chain/types/testdata/TestLambdaFunction/10-100.golden
new file mode 100644
index 00000000000..adb7eaf5e9b
--- /dev/null
+++ b/chain/types/testdata/TestLambdaFunction/10-100.golden
@@ -0,0 +1 @@
+57896044618658097711785492504343953926634992332820282019728792003956564819968
\ No newline at end of file
diff --git a/chain/types/testdata/TestLambdaFunction/1024-2048.golden b/chain/types/testdata/TestLambdaFunction/1024-2048.golden
new file mode 100644
index 00000000000..d449e4eabaa
--- /dev/null
+++ b/chain/types/testdata/TestLambdaFunction/1024-2048.golden
@@ -0,0 +1 @@
+289480223093290488558927462521719769633174961664101410098643960019782824099840
\ No newline at end of file
diff --git a/chain/types/testdata/TestLambdaFunction/2000000000000000-100000000000000000.golden b/chain/types/testdata/TestLambdaFunction/2000000000000000-100000000000000000.golden
new file mode 100644
index 00000000000..76d1403ac2b
--- /dev/null
+++ b/chain/types/testdata/TestLambdaFunction/2000000000000000-100000000000000000.golden
@@ -0,0 +1 @@
+11579208923731619542357098500868790785326998466564056403945758400791312963993
\ No newline at end of file