Add tests for lambda

Signed-off-by: Jakub Sztandera <kubuxu@protocol.ai>

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,

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)

chain/types/testdata/TestLambdaFunction/10-100.golden

@@ -0,0 +1 @@
+57896044618658097711785492504343953926634992332820282019728792003956564819968

chain/types/testdata/TestLambdaFunction/1024-2048.golden

@@ -0,0 +1 @@
+289480223093290488558927462521719769633174961664101410098643960019782824099840

chain/types/testdata/TestLambdaFunction/2000000000000000-100000000000000000.golden

@@ -0,0 +1 @@
+11579208923731619542357098500868790785326998466564056403945758400791312963993