evm 1.1.1

use alloc::{
	boxed::Box,
	collections::{BTreeMap, BTreeSet},
	vec::Vec,
};
use core::mem;

use crate::interpreter::{
	ExitError, ExitException, ExitFatal,
	runtime::{
		Log, RuntimeBackend, RuntimeBaseBackend, RuntimeConfig, RuntimeEnvironment, SetCodeOrigin,
		TouchKind,
	},
};
#[allow(unused_imports)]
use crate::uint::{H160, H256, U256, U256Ext};
use crate::{MergeStrategy, backend::TransactionalBackend};
use sha3::{Digest, Keccak256};

const RIPEMD: H160 = H160([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3]);

/// Changeset of an [OverlayedBackend].
#[derive(Clone, Debug)]
pub struct OverlayedChangeSet {
	/// Logs.
	pub logs: Vec<Log>,
	/// Balance changes.
	pub balances: BTreeMap<H160, U256>,
	/// Code changes.
	pub codes: BTreeMap<H160, Vec<u8>>,
	/// Nonce changes.
	pub nonces: BTreeMap<H160, U256>,
	/// Storage resets.
	pub storage_resets: BTreeSet<H160>,
	/// Storage changes.
	pub storages: BTreeMap<(H160, H256), H256>,
	/// Transient storage values.
	pub transient_storage: BTreeMap<(H160, H256), H256>,
	/// Accessed addresses or storage values.
	pub accessed: BTreeSet<(H160, Option<H256>)>,
	/// Touch those account and create empty if necessary.
	pub touched: BTreeSet<H160>,
	/// Deleted accounts.
	pub deletes: BTreeSet<H160>,
}

/// Overlayed backend.
pub struct OverlayedBackend<'config, B> {
	backend: B,
	substate: Box<Substate>,
	touched_ripemd: bool,
	config: &'config RuntimeConfig,
}

impl<'config, B> OverlayedBackend<'config, B> {
	/// Create a new overlayed backend, wrapping another backend.
	pub fn new(backend: B, config: &'config RuntimeConfig) -> Self {
		Self {
			backend,
			substate: Box::new(Substate::new()),
			touched_ripemd: false,
			config,
		}
	}
}

impl<B: RuntimeEnvironment + RuntimeBaseBackend> OverlayedBackend<'_, B> {
	/// Deconstruct the current overlayed backend. The change set can then be applied into the actual backend.
	pub fn deconstruct(mut self) -> (B, OverlayedChangeSet) {
		if self.config.eip161_empty_check && self.touched_ripemd {
			self.substate.touched.insert(RIPEMD);
		}

		let mut touched = BTreeSet::new();

		if self.config.eip161_empty_check {
			for address in self.substate.touched.clone() {
				if !self.exists(address) {
					self.substate.deletes.insert(address);
				}
			}
		} else {
			for address in self.substate.touched.clone() {
				if self.exists(address) {
					touched.insert(address);
				}
			}
			touched.insert(self.block_coinbase());
		}

		(
			self.backend,
			OverlayedChangeSet {
				logs: self.substate.logs,
				balances: self.substate.balances,
				codes: self.substate.codes,
				nonces: self.substate.nonces,
				storage_resets: self.substate.storage_resets,
				storages: self.substate.storages,
				transient_storage: self.substate.transient_storage,
				deletes: self.substate.deletes,
				accessed: self.substate.accessed,
				touched,
			},
		)
	}
}

impl<B: RuntimeEnvironment> RuntimeEnvironment for OverlayedBackend<'_, B> {
	fn block_hash(&self, number: U256) -> H256 {
		self.backend.block_hash(number)
	}

	fn block_number(&self) -> U256 {
		self.backend.block_number()
	}

	fn block_coinbase(&self) -> H160 {
		self.backend.block_coinbase()
	}

	fn block_timestamp(&self) -> U256 {
		self.backend.block_timestamp()
	}

	fn block_difficulty(&self) -> U256 {
		self.backend.block_difficulty()
	}

	fn block_randomness(&self) -> Option<H256> {
		self.backend.block_randomness()
	}

	fn block_gas_limit(&self) -> U256 {
		self.backend.block_gas_limit()
	}

	fn block_base_fee_per_gas(&self) -> U256 {
		self.backend.block_base_fee_per_gas()
	}

	fn blob_base_fee_per_gas(&self) -> U256 {
		self.backend.blob_base_fee_per_gas()
	}

	fn blob_versioned_hash(&self, index: U256) -> H256 {
		self.backend.blob_versioned_hash(index)
	}

	fn chain_id(&self) -> U256 {
		self.backend.chain_id()
	}
}

impl<B: RuntimeBaseBackend> RuntimeBaseBackend for OverlayedBackend<'_, B> {
	fn balance(&self, address: H160) -> U256 {
		if let Some(balance) = self.substate.known_balance(address) {
			balance
		} else {
			self.backend.balance(address)
		}
	}

	fn code(&self, address: H160) -> Vec<u8> {
		if let Some(code) = self.substate.known_code(address) {
			code
		} else {
			self.backend.code(address)
		}
	}

	fn code_hash(&self, address: H160) -> H256 {
		if self.exists(address) {
			H256::from_slice(&Keccak256::digest(&self.code(address)[..]))
		} else {
			H256::default()
		}
	}

	fn storage(&self, address: H160, index: H256) -> H256 {
		if let Some(value) = self.substate.known_storage(address, index) {
			value
		} else {
			self.backend.storage(address, index)
		}
	}

	fn transient_storage(&self, address: H160, index: H256) -> H256 {
		if let Some(value) = self.substate.known_transient_storage(address, index) {
			value
		} else {
			self.backend.transient_storage(address, index)
		}
	}

	fn exists(&self, address: H160) -> bool {
		if self.config.eip161_empty_check {
			let is_empty = self.balance(address) == U256::ZERO
				&& self.code_size(address) == U256::ZERO
				&& self.nonce(address) == U256::ZERO;
			!is_empty
		} else if let Some(exists) = self.substate.known_exists(address) {
			exists
		} else {
			self.backend.exists(address)
		}
	}

	fn nonce(&self, address: H160) -> U256 {
		if let Some(nonce) = self.substate.known_nonce(address) {
			nonce
		} else {
			self.backend.nonce(address)
		}
	}

	fn can_create(&self, address: H160) -> bool {
		if self.config.eip7610_create_check_storage {
			if let Some(nonce) = self.substate.known_nonce(address)
				&& nonce != U256::ZERO
			{
				return false;
			}
			if let Some(code) = self.substate.known_code(address)
				&& !code.is_empty()
			{
				return false;
			}
			if let Some(is_storage_empty) = self.substate.known_storage_empty(address)
				&& !is_storage_empty
			{
				return false;
			}

			self.backend.can_create(address)
		} else {
			self.nonce(address) == U256::ZERO && self.code_size(address) == U256::ZERO
		}
	}
}

impl<B: RuntimeBaseBackend> RuntimeBackend for OverlayedBackend<'_, B> {
	fn original_storage(&self, address: H160, index: H256) -> H256 {
		if let Some(value) = self.substate.known_original_storage(address, index) {
			value
		} else {
			self.backend.storage(address, index)
		}
	}

	fn created(&self, address: H160) -> bool {
		self.substate.created(address)
	}

	fn deleted(&self, address: H160) -> bool {
		self.substate.deleted(address)
	}

	fn is_cold(&self, address: H160, index: Option<H256>) -> bool {
		let accessed = self
			.substate
			.known_accessed((address, index))
			.unwrap_or(false);
		!accessed
	}

	fn mark_hot(&mut self, address: H160, kind: TouchKind) {
		match kind {
			TouchKind::Access => {
				self.substate.accessed.insert((address, None));
			}
			TouchKind::StateChange => {
				if address == RIPEMD {
					self.touched_ripemd = true;
				}
				self.substate.touched.insert(address);
			}
			TouchKind::Coinbase => {
				self.substate.touched.insert(address);
				if self.config.eip3651_warm_coinbase_address {
					self.substate.accessed.insert((address, None));
				}
			}
		}
	}

	fn mark_storage_hot(&mut self, address: H160, index: H256) {
		self.substate.accessed.insert((address, Some(index)));
	}

	fn set_storage(&mut self, address: H160, index: H256, value: H256) -> Result<(), ExitError> {
		self.substate.storages.insert((address, index), value);
		Ok(())
	}

	fn set_transient_storage(
		&mut self,
		address: H160,
		index: H256,
		value: H256,
	) -> Result<(), ExitError> {
		self.substate
			.transient_storage
			.insert((address, index), value);
		Ok(())
	}

	fn log(&mut self, log: Log) -> Result<(), ExitError> {
		self.substate.logs.push(log);
		Ok(())
	}

	fn mark_delete_reset(&mut self, address: H160) {
		if self.config.eip6780_suicide_only_in_same_tx {
			if self.created(address) {
				self.substate.deletes.insert(address);
				self.substate.balances.insert(address, U256::ZERO);
			}
		} else {
			self.substate.deletes.insert(address);
			self.substate.balances.insert(address, U256::ZERO);
		}
	}

	fn mark_create(&mut self, address: H160) {
		self.substate.creates.insert(address);
	}

	fn reset_storage(&mut self, address: H160) {
		self.substate.storage_resets.insert(address);
	}

	fn set_code(
		&mut self,
		address: H160,
		code: Vec<u8>,
		_origin: SetCodeOrigin,
	) -> Result<(), ExitError> {
		self.substate.codes.insert(address, code);
		Ok(())
	}

	fn deposit(&mut self, target: H160, value: U256) {
		if value == U256::ZERO {
			return;
		}

		let current_balance = self.balance(target);
		self.substate
			.balances
			.insert(target, current_balance.saturating_add(value));
	}

	fn withdrawal(&mut self, source: H160, value: U256) -> Result<(), ExitError> {
		if value == U256::ZERO {
			return Ok(());
		}

		let current_balance = self.balance(source);
		if current_balance < value {
			return Err(ExitException::OutOfFund.into());
		}
		let new_balance = current_balance - value;
		self.substate.balances.insert(source, new_balance);
		Ok(())
	}

	fn inc_nonce(&mut self, address: H160) -> Result<(), ExitError> {
		let old_nonce = self.nonce(address);
		if old_nonce >= U256::from(u64::MAX) {
			return Err(ExitException::MaxNonce.into());
		}
		let new_nonce = old_nonce.saturating_add(U256::ONE);
		self.substate.nonces.insert(address, new_nonce);
		Ok(())
	}
}

impl<'config, B: RuntimeBaseBackend> TransactionalBackend for OverlayedBackend<'config, B> {
	fn push_substate(&mut self) {
		let mut parent = Box::new(Substate::new());
		mem::swap(&mut parent, &mut self.substate);
		self.substate.parent = Some(parent);
	}

	fn pop_substate(&mut self, strategy: MergeStrategy) -> Result<(), ExitError> {
		let mut child = self
			.substate
			.parent
			.take()
			.ok_or(ExitError::Fatal(ExitFatal::UnevenSubstate))?;
		mem::swap(&mut child, &mut self.substate);
		let child = child;

		match strategy {
			MergeStrategy::Commit => {
				for log in child.logs {
					self.substate.logs.push(log);
				}
				for (address, balance) in child.balances {
					self.substate.balances.insert(address, balance);
				}
				for (address, code) in child.codes {
					self.substate.codes.insert(address, code);
				}
				for (address, nonce) in child.nonces {
					self.substate.nonces.insert(address, nonce);
				}
				for address in child.storage_resets {
					self.substate.storage_resets.insert(address);
				}
				for ((address, key), value) in child.storages {
					self.substate.storages.insert((address, key), value);
				}
				for ((address, key), value) in child.transient_storage {
					self.substate
						.transient_storage
						.insert((address, key), value);
				}
				for address in child.deletes {
					self.substate.deletes.insert(address);
				}
				for address in child.creates {
					self.substate.creates.insert(address);
				}
				for address in child.touched {
					self.substate.touched.insert(address);
				}
				for item in child.accessed {
					self.substate.accessed.insert(item);
				}
			}
			MergeStrategy::Revert | MergeStrategy::Discard => {}
		}

		Ok(())
	}
}

struct Substate {
	parent: Option<Box<Substate>>,
	logs: Vec<Log>,
	balances: BTreeMap<H160, U256>,
	codes: BTreeMap<H160, Vec<u8>>,
	nonces: BTreeMap<H160, U256>,
	storage_resets: BTreeSet<H160>,
	storages: BTreeMap<(H160, H256), H256>,
	transient_storage: BTreeMap<(H160, H256), H256>,
	deletes: BTreeSet<H160>,
	creates: BTreeSet<H160>,
	touched: BTreeSet<H160>,
	accessed: BTreeSet<(H160, Option<H256>)>,
}

impl Substate {
	pub fn new() -> Self {
		Self {
			parent: None,
			logs: Vec::new(),
			balances: Default::default(),
			codes: Default::default(),
			nonces: Default::default(),
			storage_resets: Default::default(),
			storages: Default::default(),
			transient_storage: Default::default(),
			deletes: Default::default(),
			creates: Default::default(),
			touched: Default::default(),
			accessed: Default::default(),
		}
	}

	pub fn known_balance(&self, address: H160) -> Option<U256> {
		if let Some(balance) = self.balances.get(&address) {
			Some(*balance)
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_balance(address)
		} else {
			None
		}
	}

	pub fn known_code(&self, address: H160) -> Option<Vec<u8>> {
		if let Some(code) = self.codes.get(&address) {
			Some(code.clone())
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_code(address)
		} else {
			None
		}
	}

	pub fn known_nonce(&self, address: H160) -> Option<U256> {
		if let Some(nonce) = self.nonces.get(&address) {
			Some(*nonce)
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_nonce(address)
		} else {
			None
		}
	}

	pub fn known_accessed(&self, item: (H160, Option<H256>)) -> Option<bool> {
		if self.accessed.contains(&item) {
			Some(true)
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_accessed(item)
		} else {
			None
		}
	}

	pub fn known_storage(&self, address: H160, key: H256) -> Option<H256> {
		if let Some(value) = self.storages.get(&(address, key)) {
			Some(*value)
		} else if self.deletes.contains(&address) {
			None
		} else if self.storage_resets.contains(&address) {
			Some(H256::default())
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_storage(address, key)
		} else {
			None
		}
	}

	pub fn known_original_storage(&self, address: H160, _key: H256) -> Option<H256> {
		if self.deletes.contains(&address) {
			None
		} else if self.storage_resets.contains(&address) {
			Some(H256::default())
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_original_storage(address, _key)
		} else {
			None
		}
	}

	pub fn known_transient_storage(&self, address: H160, key: H256) -> Option<H256> {
		if let Some(value) = self.transient_storage.get(&(address, key)) {
			Some(*value)
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_transient_storage(address, key)
		} else {
			None
		}
	}

	// Note that this is always pre-EIP161 rule. For post-EIP161, the handler would simply
	// check nonce, balance, and code.
	pub fn known_exists(&self, address: H160) -> Option<bool> {
		if self.balances.contains_key(&address)
			|| self.nonces.contains_key(&address)
			|| self.codes.contains_key(&address)
			|| self.touched.contains(&address)
		{
			Some(true)
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_exists(address)
		} else {
			None
		}
	}

	pub fn known_storage_empty(&self, address: H160) -> Option<bool> {
		if self.storages.iter().any(|((a, _), _)| *a == address) {
			Some(false)
		} else if self.deletes.contains(&address) || self.storage_resets.contains(&address) {
			Some(true)
		} else if let Some(parent) = self.parent.as_ref() {
			parent.known_storage_empty(address)
		} else {
			None
		}
	}

	pub fn deleted(&self, address: H160) -> bool {
		if self.deletes.contains(&address) {
			true
		} else if let Some(parent) = self.parent.as_ref() {
			parent.deleted(address)
		} else {
			false
		}
	}

	pub fn created(&self, address: H160) -> bool {
		if self.creates.contains(&address) {
			true
		} else if let Some(parent) = self.parent.as_ref() {
			parent.created(address)
		} else {
			false
		}
	}
}