openethereum/ethcore/src/blockchain/helpers/generators.rs

// Copyright 2015, 2016 Ethcore (UK) Ltd.
// This file is part of Parity.

// Parity is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// Parity is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with Parity.  If not, see <http://www.gnu.org/licenses/>.

use util::rlp::*;
use util::hash::{H256, H2048};
use util::numbers::{U256};
use util::bytes::Bytes;
use header::{BlockNumber, Header};
use transaction::SignedTransaction;

pub trait Forkable {
	fn fork(self, fork_number: usize) -> Self where Self: Sized;
}

pub struct Fork<I> {
	iter: I,
	fork_number: usize,
}

impl<I> Iterator for Fork<I> where I: Iterator, <I as Iterator>::Item: Forkable {
	type Item = <I as Iterator>::Item;

	#[inline]
	fn next(&mut self) -> Option<Self::Item> {
		self.iter.next().map(|item| item.fork(self.fork_number))
	}
}

pub trait WithBloom {
	fn with_bloom(self, bloom: H2048) -> Self where Self: Sized;
}

pub struct Bloom<I> {
	iter: I,
	bloom: H2048,
}

impl<I> Iterator for Bloom<I> where I: Iterator, <I as Iterator>::Item: WithBloom {
	type Item = <I as Iterator>::Item;

	#[inline]
	fn next(&mut self) -> Option<Self::Item> {
		self.iter.next().map(|item| item.with_bloom(self.bloom.clone()))
	}
}

/// Chain iterator interface.
pub trait ChainIterator: Iterator {
	/// Should be called to create a fork of current iterator.
	/// Blocks generated by fork will have lower difficulty than current chain.
	fn fork(&mut self, fork_number: usize) -> Fork<Self> where Self: Sized;
	/// Should be called to make every consecutive block have given bloom.
	fn with_bloom(&mut self, bloom: H2048) -> Bloom<Self> where Self: Sized;
}

impl<I> ChainIterator for I where I: Iterator + Sized + Clone {
	fn fork(&mut self, fork_number: usize) -> Fork<Self> {
		Fork {
			iter: self.clone(),
			fork_number: fork_number
		}
	}

	fn with_bloom(&mut self, bloom: H2048) -> Bloom<Self> {
		Bloom {
			iter: self.clone(),
			bloom: bloom
		}
	}
}

/// Helper structure, used for encoding blocks.
#[derive(Default)]
pub struct Block {
	header: Header,
	transactions: Vec<SignedTransaction>,
	uncles: Vec<Header>
}

impl Block {
	pub fn rlp(&self) -> Bytes {
		encode(self).to_vec()
	}
}

impl Encodable for Block {
	fn rlp_append(&self, s: &mut RlpStream) {
		s.begin_list(3);
		s.append(&self.header);
		s.append(&self.transactions);
		s.append(&self.uncles);
	}
}

impl Forkable for Block {
	fn fork(mut self, fork_number: usize) -> Self where Self: Sized {
		self.header.difficulty = self.header.difficulty - U256::from(fork_number);
		self
	}
}

impl WithBloom for Block {
	fn with_bloom(mut self, bloom: H2048) -> Self where Self: Sized {
		self.header.log_bloom = bloom;
		self
	}
}

/// Blockchain generator.
#[derive(Clone)]
pub struct ChainGenerator {
	/// Next block number.
	number: BlockNumber,
	/// Next block parent hash.
	parent_hash: H256,
	/// Next block difficulty.
	difficulty: U256,
}

impl ChainGenerator {
	fn prepare_block(&self) -> Block {
		let mut block = Block::default();
		block.header.parent_hash = self.parent_hash.clone();
		block.header.number = self.number;
		block.header.difficulty = self.difficulty;
		block
	}
}

impl Default for ChainGenerator {
	fn default() -> Self {
		ChainGenerator {
			number: 0,
			parent_hash: H256::default(),
			difficulty: U256::from(1000),
		}
	}
}

impl Iterator for ChainGenerator {
	type Item = Block;

	fn next(&mut self) -> Option<Self::Item> {
		let block = self.prepare_block();
		self.number += 1;
		self.parent_hash = block.header.hash();
		Some(block)
	}
}


#[cfg(test)]
mod tests {
	use util::hash::H256;
	use util::sha3::Hashable;
	use views::BlockView;
	use super::{ChainIterator, ChainGenerator};

	#[test]
	fn canon_chain_generator() {
		let mut canon_chain = ChainGenerator::default();

		let genesis_rlp = canon_chain.next().unwrap().rlp();
		let genesis = BlockView::new(&genesis_rlp);

		assert_eq!(genesis.header_view().parent_hash(), H256::default());
		assert_eq!(genesis.header_view().number(), 0);

		let b1_rlp = canon_chain.next().unwrap().rlp();
		let b1 = BlockView::new(&b1_rlp);

		assert_eq!(b1.header_view().parent_hash(), genesis.header_view().sha3());
		assert_eq!(b1.header_view().number(), 1);

		let mut fork_chain = canon_chain.fork(1);

		let b2_rlp_fork = fork_chain.next().unwrap().rlp();
		let b2_fork = BlockView::new(&b2_rlp_fork);

		assert_eq!(b2_fork.header_view().parent_hash(), b1.header_view().sha3());
		assert_eq!(b2_fork.header_view().number(), 2);

		let b2_rlp = canon_chain.next().unwrap().rlp();
		let b2 = BlockView::new(&b2_rlp);

		assert_eq!(b2.header_view().parent_hash(), b1.header_view().sha3());
		assert_eq!(b2.header_view().number(), 2);
		assert!(b2.header_view().difficulty() > b2_fork.header_view().difficulty());
	}
}