consensus/bor: handle blockalloc balance changes (maticnetwork#1074)

* fix: set balance in blockalloc

* chg: don't update balance if not zero

* fix: lint

* fix logic, add test cases

consensus/bor/bor.go

@@ -883,6 +883,10 @@ func (c *Bor) changeContractCodeIfNeeded(headerNumber uint64, state *state.State
 			for addr, account := range allocs {
 				log.Info("change contract code", "address", addr)
 				state.SetCode(addr, account.Code)
+
+				if state.GetBalance(addr).Cmp(big.NewInt(0)) == 0 {
+					state.SetBalance(addr, account.Balance)
+				}
 			}
 		}
 	}

consensus/bor/bor_test.go

@@ -40,6 +40,12 @@ func TestGenesisContractChange(t *testing.T) {
 						"balance": "0x1000",
 					},
 				},
+				"6": map[string]interface{}{
+					addr0.Hex(): map[string]interface{}{
+						"code":    hexutil.Bytes{0x1, 0x4},
+						"balance": "0x2000",
+					},
+				},
 			},
 		},
 	}
@@ -85,24 +91,35 @@ func TestGenesisContractChange(t *testing.T) {
 
 	root := genesis.Root()
 
-	// code does not change
+	// code does not change, balance remains 0
 	root, statedb = addBlock(root, 1)
 	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x1})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(0))
 
-	// code changes 1st time
+	// code changes 1st time, balance remains 0
 	root, statedb = addBlock(root, 2)
 	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x2})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(0))
 
-	// code same as 1st change
+	// code same as 1st change, balance remains 0
 	root, statedb = addBlock(root, 3)
 	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x2})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(0))
 
-	// code changes 2nd time
-	_, statedb = addBlock(root, 4)
+	// code changes 2nd time, balance updates to 4096
+	root, statedb = addBlock(root, 4)
 	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x3})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(4096))
 
-	// make sure balance change DOES NOT take effect
-	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(0))
+	// code same as 2nd change, balance remains 4096
+	root, statedb = addBlock(root, 5)
+	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x3})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(4096))
+
+	// code changes 3rd time, balance remains 4096
+	_, statedb = addBlock(root, 6)
+	require.Equal(t, statedb.GetCode(addr0), []byte{0x1, 0x4})
+	require.Equal(t, statedb.GetBalance(addr0), big.NewInt(4096))
 }
 
 func TestEncodeSigHeaderJaipur(t *testing.T) {