core: add GenerateChain, GenesisBlockForTesting

5 files changed, 230 insertions, 80 deletions

diff --git a/core/block_processor_test.go b/core/block_processor_test.go
index bd77a2fc2..5931a5f5e 100644
--- a/core/block_processor_test.go
+++ b/core/block_processor_test.go
@@ -27,16 +27,17 @@ func proc() (*BlockProcessor, *ChainManager) {
 
 func TestNumber(t *testing.T) {
 	pow := ezp.New()
+	_, chain := proc()
 
-	bp, chain := proc()
-	header := makeHeader(chain.Genesis(), 0, bp.db, 0)
+	statedb := state.New(chain.Genesis().Root(), chain.stateDb)
+	header := makeHeader(chain.Genesis(), statedb)
 	header.Number = big.NewInt(3)
 	err := ValidateHeader(pow, header, chain.Genesis().Header(), false)
 	if err != BlockNumberErr {
 		t.Errorf("expected block number error, got %q", err)
 	}
 
-	header = makeHeader(chain.Genesis(), 0, bp.db, 0)
+	header = makeHeader(chain.Genesis(), statedb)
 	err = ValidateHeader(pow, header, chain.Genesis().Header(), false)
 	if err == BlockNumberErr {
 		t.Errorf("didn't expect block number error")
diff --git a/core/chain_makers.go b/core/chain_makers.go
index e704fd088..72ae7970e 100644
--- a/core/chain_makers.go
+++ b/core/chain_makers.go
@@ -1,7 +1,6 @@
 package core
 
 import (
-	"fmt"
 	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -11,7 +10,8 @@ import (
 	"github.com/ethereum/go-ethereum/pow"
 )
 
-// So we can generate blocks easily
+// FakePow is a non-validating proof of work implementation.
+// It returns true from Verify for any block.
 type FakePow struct{}
 
 func (f FakePow) Search(block pow.Block, stop <-chan struct{}) (uint64, []byte) {
@@ -23,69 +23,125 @@ func (f FakePow) Turbo(bool)                  {}
 
 // So we can deterministically seed different blockchains
 var (
-	CanonicalSeed = 1
-	ForkSeed      = 2
+	canonicalSeed = 1
+	forkSeed      = 2
 )
 
-// Utility functions for making chains on the fly
-// Exposed for sake of testing from other packages (eg. go-ethash)
-func MakeBlock(bman *BlockProcessor, parent *types.Block, i int, db common.Database, seed int) *types.Block {
-	return types.NewBlock(makeHeader(parent, i, db, seed), nil, nil, nil)
+// BlockGen creates blocks for testing.
+// See GenerateChain for a detailed explanation.
+type BlockGen struct {
+	i       int
+	parent  *types.Block
+	chain   []*types.Block
+	header  *types.Header
+	statedb *state.StateDB
+
+	coinbase *state.StateObject
+	txs      []*types.Transaction
+	receipts []*types.Receipt
+	uncles   []*types.Header
 }
 
-func MakeChain(bman *BlockProcessor, parent *types.Block, max int, db common.Database, seed int) types.Blocks {
-	return makeChain(bman, parent, max, db, seed)
+// SetCoinbase sets the coinbase of the generated block.
+// It can be called at most once.
+func (b *BlockGen) SetCoinbase(addr common.Address) {
+	if b.coinbase != nil {
+		if len(b.txs) > 0 {
+			panic("coinbase must be set before adding transactions")
+		}
+		panic("coinbase can only be set once")
+	}
+	b.header.Coinbase = addr
+	b.coinbase = b.statedb.GetOrNewStateObject(addr)
+	b.coinbase.SetGasLimit(b.header.GasLimit)
 }
 
-func NewChainMan(block *types.Block, eventMux *event.TypeMux, db common.Database) *ChainManager {
-	return newChainManager(block, eventMux, db)
+// SetExtra sets the extra data field of the generated block.
+func (b *BlockGen) SetExtra(data []byte) {
+	b.header.Extra = data
 }
 
-func NewBlockProc(db common.Database, cman *ChainManager, eventMux *event.TypeMux) *BlockProcessor {
-	return newBlockProcessor(db, cman, eventMux)
+// AddTx adds a transaction to the generated block. If no coinbase has
+// been set, the block's coinbase is set to the zero address.
+//
+// AddTx panics if the transaction cannot be executed. In addition to
+// the protocol-imposed limitations (gas limit, etc.), there are some
+// further limitations on the content of transactions that can be
+// added. Notably, contract code relying on the BLOCKHASH instruction
+// will panic during execution.
+func (b *BlockGen) AddTx(tx *types.Transaction) {
+	if b.coinbase == nil {
+		b.SetCoinbase(common.Address{})
+	}
+	_, gas, err := ApplyMessage(NewEnv(b.statedb, nil, tx, b.header), tx, b.coinbase)
+	if err != nil {
+		panic(err)
+	}
+	b.statedb.Update()
+	b.header.GasUsed.Add(b.header.GasUsed, gas)
+	receipt := types.NewReceipt(b.statedb.Root().Bytes(), b.header.GasUsed)
+	logs := b.statedb.GetLogs(tx.Hash())
+	receipt.SetLogs(logs)
+	receipt.Bloom = types.CreateBloom(types.Receipts{receipt})
+	b.txs = append(b.txs, tx)
+	b.receipts = append(b.receipts, receipt)
 }
 
-func NewCanonical(n int, db common.Database) (*BlockProcessor, error) {
-	return newCanonical(n, db)
+// TxNonce returns the next valid transaction nonce for the
+// account at addr. It panics if the account does not exist.
+func (b *BlockGen) TxNonce(addr common.Address) uint64 {
+	if !b.statedb.HasAccount(addr) {
+		panic("account does not exist")
+	}
+	return b.statedb.GetNonce(addr)
 }
 
-// makeHeader creates the header for a new empty block, simulating
-// what miner would do. We seed chains by the first byte of the coinbase.
-func makeHeader(parent *types.Block, i int, db common.Database, seed int) *types.Header {
-	var addr common.Address
-	addr[0], addr[19] = byte(seed), byte(i) // 'random' coinbase
-	time := parent.Time() + 10              // block time is fixed at 10 seconds
-
-	// ensure that the block's coinbase has the block reward in the state.
-	state := state.New(parent.Root(), db)
-	cbase := state.GetOrNewStateObject(addr)
-	cbase.SetGasLimit(CalcGasLimit(parent))
-	cbase.AddBalance(BlockReward)
-	state.Update()
+// AddUncle adds an uncle header to the generated block.
+func (b *BlockGen) AddUncle(h *types.Header) {
+	b.uncles = append(b.uncles, h)
+}
 
-	return &types.Header{
-		Root:       state.Root(),
-		ParentHash: parent.Hash(),
-		Coinbase:   addr,
-		Difficulty: CalcDifficulty(time, parent.Time(), parent.Difficulty()),
-		Number:     new(big.Int).Add(parent.Number(), common.Big1),
-		Time:       uint64(time),
-		GasLimit:   CalcGasLimit(parent),
+// PrevBlock returns a previously generated block by number. It panics if
+// num is greater or equal to the number of the block being generated.
+// For index -1, PrevBlock returns the parent block given to GenerateChain.
+func (b *BlockGen) PrevBlock(index int) *types.Block {
+	if index >= b.i {
+		panic("block index out of range")
 	}
+	if index == -1 {
+		return b.parent
+	}
+	return b.chain[index]
 }
 
-// makeChain creates a valid chain of empty blocks.
-func makeChain(bman *BlockProcessor, parent *types.Block, max int, db common.Database, seed int) types.Blocks {
-	bman.bc.currentBlock = parent
-	blocks := make(types.Blocks, max)
-	for i := 0; i < max; i++ {
-		block := types.NewBlock(makeHeader(parent, i, db, seed), nil, nil, nil)
-		// Use ProcessWithParent to verify that we have produced a valid block.
-		_, err := bman.processWithParent(block, parent)
-		if err != nil {
-			fmt.Println("process with parent failed", err)
-			panic(err)
+// GenerateChain creates a chain of n blocks. The first block's
+// parent will be the provided parent. db is used to store
+// intermediate states and should contain the parent's state trie.
+//
+// The generator function is called with a new block generator for
+// every block. Any transactions and uncles added to the generator
+// become part of the block. If gen is nil, the blocks will be empty
+// and their coinbase will be the zero address.
+//
+// Blocks created by GenerateChain do not contain valid proof of work
+// values. Inserting them into ChainManager requires use of FakePow or
+// a similar non-validating proof of work implementation.
+func GenerateChain(parent *types.Block, db common.Database, n int, gen func(int, *BlockGen)) []*types.Block {
+	statedb := state.New(parent.Root(), db)
+	blocks := make(types.Blocks, n)
+	genblock := func(i int, h *types.Header) *types.Block {
+		b := &BlockGen{parent: parent, i: i, chain: blocks, header: h, statedb: statedb}
+		if gen != nil {
+			gen(i, b)
 		}
+		AccumulateRewards(statedb, h, b.uncles)
+		statedb.Update()
+		h.Root = statedb.Root()
+		return types.NewBlock(h, b.txs, b.uncles, b.receipts)
+	}
+	for i := 0; i < n; i++ {
+		header := makeHeader(parent, statedb)
+		block := genblock(i, header)
 		block.Td = CalcTD(block, parent)
 		blocks[i] = block
 		parent = block
@@ -93,41 +149,38 @@ func makeChain(bman *BlockProcessor, parent *types.Block, max int, db common.Dat
 	return blocks
 }
 
-// Create a new chain manager starting from given block
-// Effectively a fork factory
-func newChainManager(block *types.Block, eventMux *event.TypeMux, db common.Database) *ChainManager {
-	genesis := GenesisBlock(0, db)
-	bc := &ChainManager{blockDb: db, stateDb: db, genesisBlock: genesis, eventMux: eventMux, pow: FakePow{}}
-	bc.txState = state.ManageState(state.New(genesis.Root(), db))
-	bc.futureBlocks = NewBlockCache(1000)
-	if block == nil {
-		bc.Reset()
-	} else {
-		bc.currentBlock = block
-		bc.td = block.Td
+func makeHeader(parent *types.Block, state *state.StateDB) *types.Header {
+	time := parent.Time() + 10 // block time is fixed at 10 seconds
+	return &types.Header{
+		Root:       state.Root(),
+		ParentHash: parent.Hash(),
+		Coinbase:   parent.Coinbase(),
+		Difficulty: CalcDifficulty(time, parent.Time(), parent.Difficulty()),
+		GasLimit:   CalcGasLimit(parent),
+		GasUsed:    new(big.Int),
+		Number:     new(big.Int).Add(parent.Number(), common.Big1),
+		Time:       uint64(time),
 	}
-	return bc
-}
-
-// block processor with fake pow
-func newBlockProcessor(db common.Database, cman *ChainManager, eventMux *event.TypeMux) *BlockProcessor {
-	chainMan := newChainManager(nil, eventMux, db)
-	bman := NewBlockProcessor(db, db, FakePow{}, chainMan, eventMux)
-	return bman
 }
 
-// Make a new, deterministic canonical chain by running InsertChain
-// on result of makeChain.
+// newCanonical creates a new deterministic canonical chain by running
+// InsertChain on the result of makeChain.
 func newCanonical(n int, db common.Database) (*BlockProcessor, error) {
-	eventMux := &event.TypeMux{}
-
-	bman := newBlockProcessor(db, newChainManager(nil, eventMux, db), eventMux)
+	evmux := &event.TypeMux{}
+	chainman, _ := NewChainManager(GenesisBlock(0, db), db, db, FakePow{}, evmux)
+	bman := NewBlockProcessor(db, db, FakePow{}, chainman, evmux)
 	bman.bc.SetProcessor(bman)
 	parent := bman.bc.CurrentBlock()
 	if n == 0 {
 		return bman, nil
 	}
-	lchain := makeChain(bman, parent, n, db, CanonicalSeed)
+	lchain := makeChain(parent, n, db, canonicalSeed)
 	_, err := bman.bc.InsertChain(lchain)
 	return bman, err
 }
+
+func makeChain(parent *types.Block, n int, db common.Database, seed int) []*types.Block {
+	return GenerateChain(parent, db, n, func(i int, b *BlockGen) {
+		b.SetCoinbase(common.Address{0: byte(seed), 19: byte(i)})
+	})
+}
diff --git a/core/chain_makers_test.go b/core/chain_makers_test.go
new file mode 100644
index 000000000..d5125e1c3
--- /dev/null
+++ b/core/chain_makers_test.go
@@ -0,0 +1,78 @@
+package core
+
+import (
+	"fmt"
+	"math/big"
+
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
+	"github.com/ethereum/go-ethereum/ethdb"
+	"github.com/ethereum/go-ethereum/event"
+	"github.com/ethereum/go-ethereum/params"
+)
+
+func ExampleGenerateChain() {
+	var (
+		key1, _ = crypto.HexToECDSA("b71c71a67e1177ad4e901695e1b4b9ee17ae16c6668d313eac2f96dbcda3f291")
+		key2, _ = crypto.HexToECDSA("8a1f9a8f95be41cd7ccb6168179afb4504aefe388d1e14474d32c45c72ce7b7a")
+		key3, _ = crypto.HexToECDSA("49a7b37aa6f6645917e7b807e9d1c00d4fa71f18343b0d4122a4d2df64dd6fee")
+		addr1   = crypto.PubkeyToAddress(key1.PublicKey)
+		addr2   = crypto.PubkeyToAddress(key2.PublicKey)
+		addr3   = crypto.PubkeyToAddress(key3.PublicKey)
+		db, _   = ethdb.NewMemDatabase()
+	)
+
+	// Ensure that key1 has some funds in the genesis block.
+	genesis := GenesisBlockForTesting(db, addr1, big.NewInt(1000000))
+
+	// This call generates a chain of 5 blocks. The function runs for
+	// each block and adds different features to gen based on the
+	// block index.
+	chain := GenerateChain(genesis, db, 5, func(i int, gen *BlockGen) {
+		switch i {
+		case 0:
+			// In block 1, addr1 sends addr2 some ether.
+			tx, _ := types.NewTransaction(gen.TxNonce(addr1), addr2, big.NewInt(10000), params.TxGas, nil, nil).SignECDSA(key1)
+			gen.AddTx(tx)
+		case 1:
+			// In block 2, addr1 sends some more ether to addr2.
+			// addr2 passes it on to addr3.
+			tx1, _ := types.NewTransaction(gen.TxNonce(addr1), addr2, big.NewInt(1000), params.TxGas, nil, nil).SignECDSA(key1)
+			tx2, _ := types.NewTransaction(gen.TxNonce(addr2), addr3, big.NewInt(1000), params.TxGas, nil, nil).SignECDSA(key2)
+			gen.AddTx(tx1)
+			gen.AddTx(tx2)
+		case 2:
+			// Block 3 is empty but was mined by addr3.
+			gen.SetCoinbase(addr3)
+			gen.SetExtra([]byte("yeehaw"))
+		case 3:
+			// Block 4 includes blocks 2 and 3 as uncle headers (with modified extra data).
+			b2 := gen.PrevBlock(1).Header()
+			b2.Extra = []byte("foo")
+			gen.AddUncle(b2)
+			b3 := gen.PrevBlock(2).Header()
+			b3.Extra = []byte("foo")
+			gen.AddUncle(b3)
+		}
+	})
+
+	// Import the chain. This runs all block validation rules.
+	evmux := &event.TypeMux{}
+	chainman, _ := NewChainManager(genesis, db, db, FakePow{}, evmux)
+	chainman.SetProcessor(NewBlockProcessor(db, db, FakePow{}, chainman, evmux))
+	if i, err := chainman.InsertChain(chain); err != nil {
+		fmt.Printf("insert error (block %d): %v\n", i, err)
+		return
+	}
+
+	state := chainman.State()
+	fmt.Printf("last block: #%d\n", chainman.CurrentBlock().Number())
+	fmt.Println("balance of addr1:", state.GetBalance(addr1))
+	fmt.Println("balance of addr2:", state.GetBalance(addr2))
+	fmt.Println("balance of addr3:", state.GetBalance(addr3))
+	// Output:
+	// last block: #5
+	// balance of addr1: 989000
+	// balance of addr2: 10000
+	// balance of addr3: 5906250000000001000
+}
diff --git a/core/chain_manager_test.go b/core/chain_manager_test.go
index afec1d0f1..c37cd1a5e 100644
--- a/core/chain_manager_test.go
+++ b/core/chain_manager_test.go
@@ -67,7 +67,7 @@ func testFork(t *testing.T, bman *BlockProcessor, i, N int, f func(td1, td2 *big
 
 	// extend the fork
 	parent := bman2.bc.CurrentBlock()
-	chainB := makeChain(bman2, parent, N, db, ForkSeed)
+	chainB := makeChain(parent, N, db, forkSeed)
 	_, err = bman2.bc.InsertChain(chainB)
 	if err != nil {
 		t.Fatal("Insert chain error for fork:", err)
@@ -256,7 +256,7 @@ func TestBrokenChain(t *testing.T) {
 	}
 	bman2.bc.SetProcessor(bman2)
 	parent := bman2.bc.CurrentBlock()
-	chainB := makeChain(bman2, parent, 5, db2, ForkSeed)
+	chainB := makeChain(parent, 5, db2, forkSeed)
 	chainB = chainB[1:]
 	_, err = testChain(chainB, bman)
 	if err == nil {
@@ -444,7 +444,7 @@ func TestInsertNonceError(t *testing.T) {
 		genesis := GenesisBlock(0, db)
 		bc := chm(genesis, db)
 		bc.processor = NewBlockProcessor(db, db, bc.pow, bc, bc.eventMux)
-		blocks := makeChain(bc.processor.(*BlockProcessor), bc.currentBlock, i, db, 0)
+		blocks := makeChain(bc.currentBlock, i, db, 0)
 
 		fail := rand.Int() % len(blocks)
 		failblock := blocks[fail]
diff --git a/core/genesis.go b/core/genesis.go
index de2eee9db..df13466ec 100644
--- a/core/genesis.go
+++ b/core/genesis.go
@@ -3,6 +3,7 @@ package core
 import (
 	"encoding/json"
 	"fmt"
+	"math/big"
 	"os"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -56,3 +57,20 @@ var GenesisAccounts = []byte(`{
 	"e6716f9544a56c530d868e4bfbacb172315bdead": {"balance": "1606938044258990275541962092341162602522202993782792835301376"},
 	"1a26338f0d905e295fccb71fa9ea849ffa12aaf4": {"balance": "1606938044258990275541962092341162602522202993782792835301376"}
 }`)
+
+// GenesisBlockForTesting creates a block in which addr has the given wei balance.
+// The state trie of the block is written to db.
+func GenesisBlockForTesting(db common.Database, addr common.Address, balance *big.Int) *types.Block {
+	statedb := state.New(common.Hash{}, db)
+	obj := statedb.GetOrNewStateObject(addr)
+	obj.SetBalance(balance)
+	statedb.Update()
+	statedb.Sync()
+	block := types.NewBlock(&types.Header{
+		Difficulty: params.GenesisDifficulty,
+		GasLimit:   params.GenesisGasLimit,
+		Root:       statedb.Root(),
+	}, nil, nil, nil)
+	block.Td = params.GenesisDifficulty
+	return block
+}