refactor: optimization of total power calculation

execution/execution_test.go

@@ -142,7 +142,7 @@ func TestLockTime(t *testing.T) {
 		val := sb.MakeNewValidator(pub)
 		sb.UpdateValidator(val)
 
-		sb.AcceptTestSortition = true
+		sb.TestAcceptSortition = true
 		pld := &payload.SortitionPayload{
 			Address: pub.Address(),
 			Proof:   sortition.GenerateRandomProof(),

execution/executor/bond.go

@@ -77,6 +77,7 @@ func (e *BondExecutor) Execute(trx *tx.Tx, sb sandbox.Sandbox) error {
 	receiverVal.AddToStake(pld.Stake)
 	receiverVal.UpdateLastBondingHeight(sb.CurrentHeight())
 
+	sb.UpdatePowerDelta(pld.Stake)
 	sb.UpdateAccount(pld.Sender, senderAcc)
 	sb.UpdateValidator(receiverVal)
 

execution/executor/bond_test.go

@@ -17,11 +17,12 @@ func TestExecuteBondTx(t *testing.T) {
 	senderAddr, senderAcc := tSandbox.TestStore.RandomTestAcc()
 	senderBalance := senderAcc.Balance()
 	pub, _ := bls.GenerateTestKeyPair()
+	receiverAddr := pub.Address()
 	fee, amt := randomAmountAndFee(senderBalance / 2)
 
 	t.Run("Should fail, invalid sender", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, 1, crypto.GenerateTestAddress(),
-			pub.Address(), pub, amt, fee, "invalid sender")
+			receiverAddr, pub, amt, fee, "invalid sender")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInvalidAddress)
@@ -37,15 +38,15 @@ func TestExecuteBondTx(t *testing.T) {
 
 	t.Run("Should fail, invalid sequence", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+2, senderAddr,
-			pub.Address(), pub, amt, fee, "invalid sequence")
+			receiverAddr, pub, amt, fee, "invalid sequence")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInvalidSequence)
 	})
 
 	t.Run("Should fail, insufficient balance", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+1, senderAddr,
-			pub.Address(), pub, senderBalance+1, 0, "insufficient balance")
+			receiverAddr, pub, senderBalance+1, 0, "insufficient balance")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInsufficientFunds)
@@ -74,35 +75,33 @@ func TestExecuteBondTx(t *testing.T) {
 
 	t.Run("Should fail, public key is not set", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+1, senderAddr,
-			pub.Address(), nil, amt, fee, "no public key")
+			receiverAddr, nil, amt, fee, "no public key")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInvalidPublicKey)
 	})
 
 	t.Run("Ok", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+1, senderAddr,
-			pub.Address(), pub, amt, fee, "ok")
+			receiverAddr, pub, amt, fee, "ok")
 
 		assert.NoError(t, exe.Execute(trx, tSandbox), "Ok")
 		assert.Error(t, exe.Execute(trx, tSandbox), "Execute again, should fail")
 	})
 
 	t.Run("Should fail, public key should not set for existing validators", func(t *testing.T) {
 		trx := tx.NewBondTx(tStamp500000, senderAcc.Sequence()+2, senderAddr,
-			pub.Address(), pub, amt, fee, "with public key")
+			receiverAddr, pub, amt, fee, "with public key")
 
 		err := exe.Execute(trx, tSandbox)
 		assert.Equal(t, errors.Code(err), errors.ErrInvalidPublicKey)
 	})
 
-	assert.Equal(t, tSandbox.Account(senderAddr).Balance(),
-		senderBalance-(amt+fee))
-	assert.Equal(t, tSandbox.Validator(pub.Address()).Stake(), amt)
-	assert.Equal(t, tSandbox.Validator(pub.Address()).LastBondingHeight(),
-		tSandbox.CurrentHeight())
+	assert.Equal(t, tSandbox.Account(senderAddr).Balance(), senderBalance-(amt+fee))
+	assert.Equal(t, tSandbox.Validator(receiverAddr).Stake(), amt)
+	assert.Equal(t, tSandbox.Validator(receiverAddr).LastBondingHeight(), tSandbox.CurrentHeight())
+	assert.Equal(t, tSandbox.PowerDelta(), amt)
 	assert.Equal(t, exe.Fee(), fee)
-
 	checkTotalCoin(t, fee)
 }
 

execution/executor/sortition_test.go

@@ -33,41 +33,41 @@ func TestExecuteSortitionTx(t *testing.T) {
 
 	t.Run("Should fail, Invalid address", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, 1, crypto.GenerateTestAddress(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidAddress)
 	})
 
 	newVal.UpdateLastBondingHeight(tSandbox.CurrentHeight() - tSandbox.Params().BondInterval + 1)
 	tSandbox.UpdateValidator(newVal)
 	t.Run("Should fail, Bonding period", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, newVal.Sequence()+1, newVal.Address(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidHeight)
 	})
 
 	tSandbox.TestStore.AddTestBlock(500001)
 
 	t.Run("Should fail, Invalid sequence", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, newVal.Sequence()+2, newVal.Address(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidSequence)
 	})
 
 	t.Run("Should fail, Invalid proof", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, newVal.Sequence()+1, newVal.Address(), proof)
-		tSandbox.AcceptTestSortition = false
+		tSandbox.TestAcceptSortition = false
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidProof)
 	})
 
 	t.Run("Should fail, Committee has free seats and validator is in the committee", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, existingVal.Sequence()+1, existingVal.Address(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidTx)
 	})
 
 	t.Run("Should be ok", func(t *testing.T) {
 		trx := tx.NewSortitionTx(tStamp500000, newVal.Sequence()+1, newVal.Address(), proof)
-		tSandbox.AcceptTestSortition = true
+		tSandbox.TestAcceptSortition = true
 		assert.NoError(t, exe.Execute(trx, tSandbox))
 
 		// Execute again, should fail
@@ -88,7 +88,7 @@ func TestSortitionNonStrictMode(t *testing.T) {
 	val := tSandbox.TestStore.RandomTestVal()
 	proof := sortition.GenerateRandomProof()
 
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 	trx := tx.NewSortitionTx(tStamp500000, val.Sequence(), val.Address(), proof)
 	assert.Error(t, exe1.Execute(trx, tSandbox))
 	assert.NoError(t, exe2.Execute(trx, tSandbox))
@@ -119,7 +119,7 @@ func TestChangePower1(t *testing.T) {
 	proof3 := sortition.GenerateRandomProof()
 
 	tSandbox.TestParams.CommitteeSize = 4
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 	trx1 := tx.NewSortitionTx(tStamp500000, val1.Sequence()+1, val1.Address(), proof1)
 	assert.NoError(t, exe.Execute(trx1, tSandbox))
 
@@ -162,7 +162,7 @@ func TestChangePower2(t *testing.T) {
 	proof4 := sortition.GenerateRandomProof()
 
 	tSandbox.TestParams.CommitteeSize = 7
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 	trx1 := tx.NewSortitionTx(tStamp500000, val1.Sequence()+1, val1.Address(), proof1)
 	assert.NoError(t, exe.Execute(trx1, tSandbox))
 
@@ -196,7 +196,7 @@ func TestOldestDidNotPropose(t *testing.T) {
 	}
 
 	tSandbox.TestParams.CommitteeSize = 7
-	tSandbox.AcceptTestSortition = true
+	tSandbox.TestAcceptSortition = true
 
 	stamp := tStamp500000
 	for i := 0; i < 8; i = i + 2 {

execution/executor/unbond.go

@@ -53,6 +53,11 @@ func (e *UnbondExecutor) Execute(trx *tx.Tx, sb sandbox.Sandbox) error {
 
 	val.IncSequence()
 	val.UpdateUnbondingHeight(sb.CurrentHeight())
+
+	// At this point, the validator's power is zero.
+	// However, we know the validator's stake.
+	// So, we can update the power delta with the negative of the validator's stake.
+	sb.UpdatePowerDelta(-1 * val.Power())
 	sb.UpdateValidator(val)
 
 	return nil

execution/executor/unbond_test.go

@@ -15,6 +15,7 @@ func TestExecuteUnbondTx(t *testing.T) {
 	exe := NewUnbondExecutor(true)
 
 	pub, _ := bls.GenerateTestKeyPair()
+	valAddr := pub.Address()
 	val := tSandbox.MakeNewValidator(pub)
 	tSandbox.UpdateValidator(val)
 
@@ -24,7 +25,7 @@ func TestExecuteUnbondTx(t *testing.T) {
 	})
 
 	t.Run("Should fail, Invalid sequence", func(t *testing.T) {
-		trx := tx.NewUnbondTx(tStamp500000, val.Sequence()+2, pub.Address(), "invalid sequence")
+		trx := tx.NewUnbondTx(tStamp500000, val.Sequence()+2, valAddr, "invalid sequence")
 		assert.Equal(t, errors.Code(exe.Execute(trx, tSandbox)), errors.ErrInvalidSequence)
 	})
 
@@ -45,17 +46,18 @@ func TestExecuteUnbondTx(t *testing.T) {
 	})
 
 	t.Run("Ok", func(t *testing.T) {
-		trx := tx.NewUnbondTx(tStamp500000, val.Sequence()+1, pub.Address(), "Ok")
+		trx := tx.NewUnbondTx(tStamp500000, val.Sequence()+1, valAddr, "Ok")
 
 		assert.NoError(t, exe.Execute(trx, tSandbox))
 
 		// Execute again, should fail
 		assert.Error(t, exe.Execute(trx, tSandbox))
 	})
 
-	assert.Zero(t, tSandbox.Validator(pub.Address()).Stake())
-	assert.Zero(t, tSandbox.Validator(pub.Address()).Power())
-	assert.Equal(t, tSandbox.Validator(pub.Address()).UnbondingHeight(), tSandbox.CurrentHeight())
+	assert.Zero(t, tSandbox.Validator(valAddr).Stake())
+	assert.Zero(t, tSandbox.Validator(valAddr).Power())
+	assert.Equal(t, tSandbox.Validator(valAddr).UnbondingHeight(), tSandbox.CurrentHeight())
+	assert.Equal(t, tSandbox.PowerDelta(), -1*val.Stake())
 	assert.Zero(t, exe.Fee())
 
 	checkTotalCoin(t, 0)

sandbox/interface.go

@@ -20,6 +20,8 @@ type Sandbox interface {
 	Validator(crypto.Address) *validator.Validator
 	MakeNewValidator(*bls.PublicKey) *validator.Validator
 	UpdateValidator(*validator.Validator)
+	UpdatePowerDelta(delta int64)
+	PowerDelta() int64
 
 	VerifyProof(hash.Stamp, sortition.Proof, *validator.Validator) bool
 	Committee() committee.Reader

sandbox/mock.go

@@ -21,7 +21,8 @@ type MockSandbox struct {
 	TestStore            *store.MockStore
 	TestCommittee        committee.Committee
 	TestCommitteeSigners []crypto.Signer
-	AcceptTestSortition  bool
+	TestAcceptSortition  bool
+	TestPowerDelta       int64
 }
 
 func MockingSandbox() *MockSandbox {
@@ -97,7 +98,12 @@ func (m *MockSandbox) IterateValidators(consumer func(*validator.Validator, bool
 func (m *MockSandbox) Committee() committee.Reader {
 	return m.TestCommittee
 }
-
+func (m *MockSandbox) UpdatePowerDelta(delta int64) {
+	m.TestPowerDelta += delta
+}
+func (m *MockSandbox) PowerDelta() int64 {
+	return m.TestPowerDelta
+}
 func (m *MockSandbox) VerifyProof(hash.Stamp, sortition.Proof, *validator.Validator) bool {
-	return m.AcceptTestSortition
+	return m.TestAcceptSortition
 }

sandbox/sandbox.go

@@ -28,6 +28,8 @@ type sandbox struct {
 	params          param.Params
 	totalAccounts   int32
 	totalValidators int32
+	totalPower      int64
+	powerDelta      int64
 }
 
 type sandboxValidator struct {
@@ -40,11 +42,13 @@ type sandboxAccount struct {
 	updated bool
 }
 
-func NewSandbox(store store.Reader, params param.Params, committee committee.Reader) Sandbox {
+func NewSandbox(store store.Reader, params param.Params,
+	committee committee.Reader, totalPower int64) Sandbox {
 	sb := &sandbox{
-		store:     store,
-		committee: committee,
-		params:    params,
+		store:      store,
+		committee:  committee,
+		totalPower: totalPower,
+		params:     params,
 	}
 
 	sb.accounts = make(map[crypto.Address]*sandboxAccount)
@@ -223,18 +227,28 @@ func (sb *sandbox) Committee() committee.Reader {
 	return sb.committee
 }
 
-// TODO: write test for me.
+// UpdatePowerDelta updates the change in the total power of the blockchain.
+// The delta is the amount of change in the total power and can be either positive or negative.
+func (sb *sandbox) UpdatePowerDelta(delta int64) {
+	sb.lk.Lock()
+	defer sb.lk.Unlock()
+
+	sb.powerDelta += delta
+}
+
+func (sb *sandbox) PowerDelta() int64 {
+	sb.lk.RLock()
+	defer sb.lk.RUnlock()
+
+	return sb.powerDelta
+}
+
 // VerifyProof verifies proof of a sortition transaction.
 func (sb *sandbox) VerifyProof(stamp hash.Stamp, proof sortition.Proof, val *validator.Validator) bool {
 	_, b := sb.store.RecentBlockByStamp(stamp)
 	if b == nil {
 		return false
 	}
 	seed := b.Header().SortitionSeed()
-	total := int64(0) // TODO: we can get it from state
-	sb.store.IterateValidators(func(val *validator.Validator) bool {
-		total += val.Power()
-		return false
-	})
-	return sortition.VerifyProof(seed, proof, val.PublicKey(), total, val.Power())
+	return sortition.VerifyProof(seed, proof, val.PublicKey(), sb.totalPower, val.Power())
 }