Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix H01 #141

Merged
merged 1 commit into from
Sep 19, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 25 additions & 0 deletions contracts/src/token/psp22/extensions/capped.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ pub use crate::{
};
pub use capped::Internal as _;
use openbrush::traits::{
AccountId,
Balance,
Storage,
String,
Expand Down Expand Up @@ -81,3 +82,27 @@ pub trait InternalImpl: Storage<Data> + Internal + PSP22 {
self.data().cap.get_or_default()
}
}

pub trait PSP22TransferImpl: Internal {
fn _before_token_transfer(
&mut self,
_from: Option<&AccountId>,
_to: Option<&AccountId>,
_amount: &Balance,
) -> Result<(), PSP22Error> {
if _from.is_none() && _to.is_some() && Internal::_is_cap_exceeded(self, _amount) {
return Err(PSP22Error::Custom(String::from("Cap exceeded")))
}

Ok(())
}

fn _after_token_transfer(
&mut self,
_from: Option<&AccountId>,
_to: Option<&AccountId>,
_amount: &Balance,
) -> Result<(), PSP22Error> {
Ok(())
}
}
46 changes: 25 additions & 21 deletions contracts/src/token/psp22/psp22.rs
Original file line number Diff line number Diff line change
Expand Up @@ -142,23 +142,9 @@ pub trait Internal {
fn _mint_to(&mut self, account: AccountId, amount: Balance) -> Result<(), PSP22Error>;

fn _burn_from(&mut self, account: AccountId, amount: Balance) -> Result<(), PSP22Error>;

fn _before_token_transfer(
&mut self,
_from: Option<&AccountId>,
_to: Option<&AccountId>,
_amount: &Balance,
) -> Result<(), PSP22Error>;

fn _after_token_transfer(
&mut self,
_from: Option<&AccountId>,
_to: Option<&AccountId>,
_amount: &Balance,
) -> Result<(), PSP22Error>;
}

pub trait InternalImpl: Storage<Data> + Internal {
pub trait InternalImpl: Storage<Data> + Internal + PSP22Transfer {
fn _emit_transfer_event(&self, _from: Option<AccountId>, _to: Option<AccountId>, _amount: Balance) {}

fn _emit_approval_event(&self, _owner: AccountId, _spender: AccountId, _amount: Balance) {}
Expand Down Expand Up @@ -188,14 +174,14 @@ pub trait InternalImpl: Storage<Data> + Internal {
return Err(PSP22Error::InsufficientBalance)
}

Internal::_before_token_transfer(self, Some(&from), Some(&to), &amount)?;
PSP22Transfer::_before_token_transfer(self, Some(&from), Some(&to), &amount)?;

self.data().balances.insert(&from, &(from_balance - amount));

let to_balance = Internal::_balance_of(self, &to);
self.data().balances.insert(&to, &(to_balance + amount));

Internal::_after_token_transfer(self, Some(&from), Some(&to), &amount)?;
PSP22Transfer::_after_token_transfer(self, Some(&from), Some(&to), &amount)?;
Internal::_emit_transfer_event(self, Some(from), Some(to), amount);

Ok(())
Expand All @@ -208,15 +194,15 @@ pub trait InternalImpl: Storage<Data> + Internal {
}

fn _mint_to(&mut self, account: AccountId, amount: Balance) -> Result<(), PSP22Error> {
Internal::_before_token_transfer(self, None, Some(&account), &amount)?;
PSP22Transfer::_before_token_transfer(self, None, Some(&account), &amount)?;
let mut new_balance = Internal::_balance_of(self, &account);
new_balance += amount;
self.data().balances.insert(&account, &new_balance);

let new_supply = self.data().supply.get_or_default() + amount;
self.data().supply.set(&new_supply);

Internal::_after_token_transfer(self, None, Some(&account), &amount)?;
PSP22Transfer::_after_token_transfer(self, None, Some(&account), &amount)?;
Internal::_emit_transfer_event(self, None, Some(account), amount);

Ok(())
Expand All @@ -229,20 +215,38 @@ pub trait InternalImpl: Storage<Data> + Internal {
return Err(PSP22Error::InsufficientBalance)
}

Internal::_before_token_transfer(self, Some(&account), None, &amount)?;
PSP22Transfer::_before_token_transfer(self, Some(&account), None, &amount)?;

from_balance -= amount;
self.data().balances.insert(&account, &from_balance);

let new_supply = self.data().supply.get_or_default() - amount;
self.data().supply.set(&new_supply);

Internal::_after_token_transfer(self, Some(&account), None, &amount)?;
PSP22Transfer::_after_token_transfer(self, Some(&account), None, &amount)?;
Internal::_emit_transfer_event(self, Some(account), None, amount);

Ok(())
}
}

pub trait PSP22Transfer {
fn _before_token_transfer(
&mut self,
_from: Option<&AccountId>,
_to: Option<&AccountId>,
_amount: &Balance,
) -> Result<(), PSP22Error>;

fn _after_token_transfer(
&mut self,
_from: Option<&AccountId>,
_to: Option<&AccountId>,
_amount: &Balance,
) -> Result<(), PSP22Error>;
}

pub trait PSP22TransferImpl {
fn _before_token_transfer(
&mut self,
_from: Option<&AccountId>,
Expand Down
1 change: 0 additions & 1 deletion examples/psp22/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ pub struct HatedStorage {
#[openbrush::contract]
pub mod my_psp22 {
use crate::*;
use openbrush::traits::String;

#[ink(storage)]
#[derive(Storage)]
Expand Down
5 changes: 1 addition & 4 deletions examples/psp22_extensions/capped/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,7 @@
#[openbrush::implementation(PSP22, PSP22Capped, PSP22Mintable)]
#[openbrush::contract]
pub mod my_psp22_capped {
use openbrush::traits::{
Storage,
String,
};
use openbrush::traits::Storage;

#[ink(storage)]
#[derive(Default, Storage)]
Expand Down
6 changes: 5 additions & 1 deletion lang/codegen/src/implementation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ pub fn generate(attrs: TokenStream, ink_module: TokenStream) -> TokenStream {

let mut impl_args = ImplArgs::new(&map, &mut items, &mut imports, &mut overriden_traits, ident);

for to_implement in args {
for to_implement in args.clone() {
match to_implement.as_str() {
"PSP22" => impl_psp22(&mut impl_args),
"PSP22Mintable" => impl_psp22_mintable(&mut impl_args),
Expand Down Expand Up @@ -115,6 +115,10 @@ pub fn generate(attrs: TokenStream, ink_module: TokenStream) -> TokenStream {
}
}

if args.contains(&String::from("PSP22")) {
impl_psp22_transfer(&mut impl_args, args.contains(&String::from("PSP22Capped")));
}

cleanup_imports(impl_args.imports);

// add the imports
Expand Down
78 changes: 42 additions & 36 deletions lang/codegen/src/implementations.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,24 +101,6 @@ pub(crate) fn impl_psp22(impl_args: &mut ImplArgs) {
fn _burn_from(&mut self, account: AccountId, amount: Balance) -> Result<(), PSP22Error> {
psp22::InternalImpl::_burn_from(self, account, amount)
}

fn _before_token_transfer(
&mut self,
from: Option<&AccountId>,
to: Option<&AccountId>,
amount: &Balance,
) -> Result<(), PSP22Error> {
psp22::InternalImpl::_before_token_transfer(self, from, to, amount)
}

fn _after_token_transfer(
&mut self,
from: Option<&AccountId>,
to: Option<&AccountId>,
amount: &Balance,
) -> Result<(), PSP22Error> {
psp22::InternalImpl::_after_token_transfer(self, from, to, amount)
}
}
))
.expect("Should parse");
Expand Down Expand Up @@ -295,6 +277,48 @@ pub(crate) fn impl_psp22_metadata(impl_args: &mut ImplArgs) {
impl_args.items.push(syn::Item::Impl(metadata));
}

pub(crate) fn impl_psp22_transfer(impl_args: &mut ImplArgs, capped: bool) {
let storage_struct_name = impl_args.contract_name();

let implementation = if capped {
syn::parse2::<syn::ItemImpl>(quote!(
impl capped::PSP22TransferImpl for #storage_struct_name {}
))
.expect("Should parse")
} else {
syn::parse2::<syn::ItemImpl>(quote!(
impl psp22::PSP22TransferImpl for #storage_struct_name {}
))
.expect("Should parse")
};

let transfer = syn::parse2::<syn::ItemImpl>(quote!(
impl psp22::PSP22Transfer for #storage_struct_name {
fn _before_token_transfer(
&mut self,
_from: Option<&AccountId>,
_to: Option<&AccountId>,
_amount: &Balance,
) -> Result<(), PSP22Error> {
PSP22TransferImpl::_before_token_transfer(self, _from, _to, _amount)
}

fn _after_token_transfer(
&mut self,
_from: Option<&AccountId>,
_to: Option<&AccountId>,
_amount: &Balance,
) -> Result<(), PSP22Error> {
PSP22TransferImpl::_after_token_transfer(self, _from, _to, _amount)
}
}
))
.expect("Should parse");

impl_args.items.push(syn::Item::Impl(implementation));
impl_args.items.push(syn::Item::Impl(transfer));
}

pub(crate) fn impl_psp22_capped(impl_args: &mut ImplArgs) {
let storage_struct_name = impl_args.contract_name();
let internal_impl = syn::parse2::<syn::ItemImpl>(quote!(
Expand Down Expand Up @@ -835,24 +859,6 @@ pub(crate) fn impl_psp34(impl_args: &mut ImplArgs) {
fn _check_token_exists(&self, id: &Id) -> Result<AccountId, PSP34Error> {
psp34::InternalImpl::_check_token_exists(self, id)
}

fn _before_token_transfer(
&mut self,
from: Option<&AccountId>,
to: Option<&AccountId>,
id: &Id,
) -> Result<(), PSP34Error> {
psp34::InternalImpl::_before_token_transfer(self, from, to, id)
}

fn _after_token_transfer(
&mut self,
from: Option<&AccountId>,
to: Option<&AccountId>,
id: &Id,
) -> Result<(), PSP34Error> {
psp34::InternalImpl::_after_token_transfer(self, from, to, id)
}
}
))
.expect("Should parse");
Expand Down