use crate::Config;
use tetcore_std::marker::PhantomData;
use tp_runtime::traits::Zero;
use fabric_support::dispatch::{
DispatchResultWithPostInfo, PostDispatchInfo, DispatchErrorWithPostInfo,
};
use noble_contracts_primitives::ExecError;
#[cfg(test)]
use std::{any::Any, fmt::Debug};
pub type Gas = fabric_support::weights::Weight;
#[must_use]
#[derive(Debug, PartialEq, Eq)]
pub enum GasMeterResult {
Proceed(ChargedAmount),
OutOfGas,
}
impl GasMeterResult {
pub fn is_out_of_gas(&self) -> bool {
match *self {
GasMeterResult::OutOfGas => true,
GasMeterResult::Proceed(_) => false,
}
}
}
#[derive(Debug, PartialEq, Eq)]
pub struct ChargedAmount(Gas);
impl ChargedAmount {
pub fn amount(&self) -> Gas {
self.0
}
}
#[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: Config>: 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: Config> {
gas_limit: Gas,
gas_left: Gas,
_phantom: PhantomData<T>,
#[cfg(test)]
tokens: Vec<ErasedToken>,
}
impl<T: Config> GasMeter<T> {
pub fn new(gas_limit: Gas) -> Self {
GasMeter {
gas_limit,
gas_left: gas_limit,
_phantom: PhantomData,
#[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(ChargedAmount(amount)),
None => GasMeterResult::OutOfGas,
}
}
pub fn refund(&mut self, amount: ChargedAmount) {
self.gas_left = self.gas_left.saturating_add(amount.0).min(self.gas_limit)
}
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::new(amount);
let r = f(Some(&mut nested));
self.gas_left = self.gas_left + nested.gas_left;
r
}
}
pub fn gas_spent(&self) -> Gas {
self.gas_limit - self.gas_left
}
pub fn gas_left(&self) -> Gas {
self.gas_left
}
pub fn into_dispatch_result<R, E>(self, result: Result<R, E>) -> DispatchResultWithPostInfo
where
E: Into<ExecError>,
{
let post_info = PostDispatchInfo {
actual_weight: Some(self.gas_spent()),
pays_fee: Default::default(),
};
result
.map(|_| post_info)
.map_err(|e| DispatchErrorWithPostInfo { post_info, error: e.into().error })
}
#[cfg(test)]
pub fn tokens(&self) -> &[ErasedToken] {
&self.tokens
}
}
#[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;
#[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>::new(50000);
assert_eq!(gas_meter.gas_left(), 50000);
}
#[test]
fn simple() {
let mut gas_meter = GasMeter::<Test>::new(50000);
let result = gas_meter
.charge(&MultiplierTokenMetadata { multiplier: 3 }, MultiplierToken(10));
assert!(!result.is_out_of_gas());
assert_eq!(gas_meter.gas_left(), 49_970);
}
#[test]
fn tracing() {
let mut gas_meter = GasMeter::<Test>::new(50000);
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>::new(0);
assert!(gas_meter.charge(&(), SimpleToken(1)).is_out_of_gas());
}
#[test]
fn overcharge_is_unrecoverable() {
let mut gas_meter = GasMeter::<Test>::new(200);
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>::new(25);
assert!(!gas_meter.charge(&(), SimpleToken(25)).is_out_of_gas());
}
}