use crate::{GasSpent, Module, Trait, BalanceOf, NegativeImbalanceOf};
use sp_std::convert::TryFrom;
use sp_runtime::traits::{
CheckedMul, Zero, SaturatedConversion, AtLeast32Bit, UniqueSaturatedInto,
};
use frame_support::{
traits::{Currency, ExistenceRequirement, Imbalance, OnUnbalanced, WithdrawReason}, StorageValue,
dispatch::DispatchError,
};
#[cfg(test)]
use std::{any::Any, fmt::Debug};
pub type Gas = u64;
#[must_use]
#[derive(Debug, PartialEq, Eq)]
pub enum GasMeterResult {
Proceed,
OutOfGas,
}
impl GasMeterResult {
pub fn is_out_of_gas(&self) -> bool {
match *self {
GasMeterResult::OutOfGas => true,
GasMeterResult::Proceed => false,
}
}
}
#[cfg(not(test))]
pub trait TestAuxiliaries {}
#[cfg(not(test))]
impl<T> TestAuxiliaries for T {}
#[cfg(test)]
pub trait TestAuxiliaries: Any + Debug + PartialEq + Eq {}
#[cfg(test)]
impl<T: Any + Debug + PartialEq + Eq> TestAuxiliaries for T {}
pub trait Token<T: Trait>: Copy + Clone + TestAuxiliaries {
type Metadata;
fn calculate_amount(&self, metadata: &Self::Metadata) -> Gas;
}
#[cfg(test)]
pub struct ErasedToken {
pub description: String,
pub token: Box<dyn Any>,
}
pub struct GasMeter<T: Trait> {
limit: Gas,
gas_left: Gas,
gas_price: BalanceOf<T>,
#[cfg(test)]
tokens: Vec<ErasedToken>,
}
impl<T: Trait> GasMeter<T> {
pub fn with_limit(gas_limit: Gas, gas_price: BalanceOf<T>) -> GasMeter<T> {
GasMeter {
limit: gas_limit,
gas_left: gas_limit,
gas_price,
#[cfg(test)]
tokens: Vec::new(),
}
}
#[inline]
pub fn charge<Tok: Token<T>>(
&mut self,
metadata: &Tok::Metadata,
token: Tok,
) -> GasMeterResult {
#[cfg(test)]
{
let erased_tok = ErasedToken {
description: format!("{:?}", token),
token: Box::new(token),
};
self.tokens.push(erased_tok);
}
let amount = token.calculate_amount(metadata);
let new_value = match self.gas_left.checked_sub(amount) {
None => None,
Some(val) => Some(val),
};
self.gas_left = new_value.unwrap_or_else(Zero::zero);
match new_value {
Some(_) => GasMeterResult::Proceed,
None => GasMeterResult::OutOfGas,
}
}
pub fn with_nested<R, F: FnOnce(Option<&mut GasMeter<T>>) -> R>(
&mut self,
amount: Gas,
f: F,
) -> R {
if self.gas_left < amount {
f(None)
} else {
self.gas_left = self.gas_left - amount;
let mut nested = GasMeter::with_limit(amount, self.gas_price);
let r = f(Some(&mut nested));
self.gas_left = self.gas_left + nested.gas_left;
r
}
}
pub fn gas_price(&self) -> BalanceOf<T> {
self.gas_price
}
pub fn gas_left(&self) -> Gas {
self.gas_left
}
fn spent(&self) -> Gas {
self.limit - self.gas_left
}
#[cfg(test)]
pub fn tokens(&self) -> &[ErasedToken] {
&self.tokens
}
}
pub fn buy_gas<T: Trait>(
transactor: &T::AccountId,
gas_limit: Gas,
) -> Result<(GasMeter<T>, NegativeImbalanceOf<T>), DispatchError> {
let gas_price = <Module<T>>::gas_price();
let cost = if gas_price.is_zero() {
<BalanceOf<T>>::zero()
} else {
<BalanceOf<T> as TryFrom<Gas>>::try_from(gas_limit).ok()
.and_then(|gas_limit| gas_price.checked_mul(&gas_limit))
.ok_or("overflow multiplying gas limit by price")?
};
let imbalance = T::Currency::withdraw(
transactor,
cost,
WithdrawReason::Fee.into(),
ExistenceRequirement::KeepAlive
)?;
Ok((GasMeter::with_limit(gas_limit, gas_price), imbalance))
}
pub fn refund_unused_gas<T: Trait>(
transactor: &T::AccountId,
gas_meter: GasMeter<T>,
imbalance: NegativeImbalanceOf<T>,
) {
let gas_spent = gas_meter.spent();
let gas_left = gas_meter.gas_left();
GasSpent::mutate(|block_gas_spent| *block_gas_spent += gas_spent);
let refund = gas_meter.gas_price * gas_left.unique_saturated_into();
let refund_imbalance = T::Currency::deposit_creating(transactor, refund);
if let Ok(imbalance) = imbalance.offset(refund_imbalance) {
T::GasPayment::on_unbalanced(imbalance);
}
}
pub fn approx_gas_for_balance<Balance>(gas_price: Balance, balance: Balance) -> Gas
where Balance: AtLeast32Bit
{
if gas_price.is_zero() {
Zero::zero()
} else {
(balance / gas_price).saturated_into::<Gas>()
}
}
#[macro_export]
macro_rules! match_tokens {
($tokens_iter:ident,) => {
};
($tokens_iter:ident, $x:expr, $($rest:tt)*) => {
{
let next = ($tokens_iter).next().unwrap();
let pattern = $x;
let mut _pattern_typed_next_ref = &pattern;
_pattern_typed_next_ref = match next.token.downcast_ref() {
Some(p) => {
assert_eq!(p, &pattern);
p
}
None => {
panic!("expected type {} got {}", stringify!($x), next.description);
}
};
}
match_tokens!($tokens_iter, $($rest)*);
};
}
#[cfg(test)]
mod tests {
use super::{GasMeter, Token};
use crate::{tests::Test, gas::approx_gas_for_balance};
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
struct SimpleToken(u64);
impl Token<Test> for SimpleToken {
type Metadata = ();
fn calculate_amount(&self, _metadata: &()) -> u64 { self.0 }
}
struct MultiplierTokenMetadata {
multiplier: u64,
}
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
struct MultiplierToken(u64);
impl Token<Test> for MultiplierToken {
type Metadata = MultiplierTokenMetadata;
fn calculate_amount(&self, metadata: &MultiplierTokenMetadata) -> u64 {
self.0 * metadata.multiplier
}
}
#[test]
fn it_works() {
let gas_meter = GasMeter::<Test>::with_limit(50000, 10);
assert_eq!(gas_meter.gas_left(), 50000);
}
#[test]
fn simple() {
let mut gas_meter = GasMeter::<Test>::with_limit(50000, 10);
let result = gas_meter
.charge(&MultiplierTokenMetadata { multiplier: 3 }, MultiplierToken(10));
assert!(!result.is_out_of_gas());
assert_eq!(gas_meter.gas_left(), 49_970);
assert_eq!(gas_meter.spent(), 30);
assert_eq!(gas_meter.gas_price(), 10);
}
#[test]
fn tracing() {
let mut gas_meter = GasMeter::<Test>::with_limit(50000, 10);
assert!(!gas_meter.charge(&(), SimpleToken(1)).is_out_of_gas());
assert!(!gas_meter
.charge(&MultiplierTokenMetadata { multiplier: 3 }, MultiplierToken(10))
.is_out_of_gas());
let mut tokens = gas_meter.tokens()[0..2].iter();
match_tokens!(tokens, SimpleToken(1), MultiplierToken(10),);
}
#[test]
fn refuse_to_execute_anything_if_zero() {
let mut gas_meter = GasMeter::<Test>::with_limit(0, 10);
assert!(gas_meter.charge(&(), SimpleToken(1)).is_out_of_gas());
}
#[test]
fn overcharge_is_unrecoverable() {
let mut gas_meter = GasMeter::<Test>::with_limit(200, 10);
assert!(gas_meter.charge(&(), SimpleToken(300)).is_out_of_gas());
assert!(gas_meter.charge(&(), SimpleToken(1)).is_out_of_gas());
}
#[test]
fn charge_exact_amount() {
let mut gas_meter = GasMeter::<Test>::with_limit(25, 10);
assert!(!gas_meter.charge(&(), SimpleToken(25)).is_out_of_gas());
}
#[test]
fn approx_gas_for_balance_works() {
let tests = vec![
(approx_gas_for_balance(0_u64, 123), 0),
(approx_gas_for_balance(0_u64, 456), 0),
(approx_gas_for_balance(1_u64, 123), 123),
(approx_gas_for_balance(1_u64, 456), 456),
(approx_gas_for_balance(100_u64, 900), 9),
(approx_gas_for_balance(123_u64, 900), 7),
];
for (lhs, rhs) in tests {
assert_eq!(lhs, rhs);
}
}
}