feat(Gas Tank): Improve GasFueler and create MakeTank

traits/gas-tank/src/impl_nonfungibles.rs

@@ -14,14 +14,42 @@ pub const ATTR_MEMBERSHIP_GAS: &[u8] = b"membership_gas";
 pub const ATTR_GAS_TX_PAY_WITH_MEMBERSHIP: &[u8] = b"mbmshp_pays_gas";
 
 #[derive(Encode, Decode, Debug)]
-pub struct MembershipWeightTank<T: frame_system::Config> {
-    pub since: BlockNumberFor<T>,
-    pub used: Weight,
-    pub period: Option<BlockNumberFor<T>>,
-    pub max_per_period: Option<Weight>,
+pub struct WeightTank<T: frame_system::Config> {
+    pub(crate) since: BlockNumberFor<T>,
+    pub(crate) used: Weight,
+    pub(crate) period: Option<BlockNumberFor<T>>,
+    pub(crate) capacity_per_period: Option<Weight>,
 }
 
-impl<T> Default for MembershipWeightTank<T>
+impl<T> WeightTank<T>
+where
+    T: frame_system::Config,
+{
+    fn new(capacity_per_period: Option<Weight>, period: Option<BlockNumberFor<T>>) -> Self {
+        Self {
+            since: frame_system::Pallet::<T>::block_number(),
+            used: Weight::zero(),
+            period,
+            capacity_per_period,
+        }
+    }
+
+    pub(crate) fn get<F>(collection_id: &F::CollectionId, item_id: &F::ItemId) -> Option<Self>
+    where
+        F: nonfungibles_v2::Inspect<T::AccountId>,
+    {
+        F::typed_system_attribute(collection_id, Some(item_id), &ATTR_MEMBERSHIP_GAS)
+    }
+
+    fn put<F, I>(&self, collection_id: &F::CollectionId, item_id: &F::ItemId) -> Option<()>
+    where
+        F: nonfungibles_v2::Inspect<T::AccountId> + nonfungibles_v2::Mutate<T::AccountId, I>,
+    {
+        F::set_typed_attribute(collection_id, item_id, &ATTR_MEMBERSHIP_GAS, self).ok()
+    }
+}
+
+impl<T> Default for WeightTank<T>
 where
     T: frame_system::Config,
     BlockNumberFor<T>: Default,
@@ -31,14 +59,14 @@ where
             since: Default::default(),
             used: Default::default(),
             period: Default::default(),
-            max_per_period: Default::default(),
+            capacity_per_period: Default::default(),
         }
     }
 }
 
-pub struct NonFungibleGasBurner<T, F, I>(PhantomData<(T, F, I)>);
+pub struct NonFungibleGasTank<T, F, I>(PhantomData<(T, F, I)>);
 
-impl<T, F, ItemConfig> GasBurner for NonFungibleGasBurner<T, F, ItemConfig>
+impl<T, F, ItemConfig> GasBurner for NonFungibleGasTank<T, F, ItemConfig>
 where
     T: frame_system::Config,
     BlockNumberFor<T>: Bounded,
@@ -52,24 +80,22 @@ where
 
     fn check_available_gas(who: &Self::AccountId, estimated: &Self::Gas) -> Option<Self::Gas> {
         F::owned(who).find_map(|(collection, item)| {
-            let mut gas_tank: MembershipWeightTank<T> =
-                F::typed_system_attribute(&collection, Some(&item), &ATTR_MEMBERSHIP_GAS)?;
+            let mut tank = WeightTank::<T>::get::<F>(&collection, &item)?;
 
             let block_number = frame_system::Pallet::<T>::block_number();
-            let period = gas_tank.period.unwrap_or(BlockNumberFor::<T>::max_value());
+            let period = tank.period.unwrap_or(BlockNumberFor::<T>::max_value());
 
-            let Some(max_weight) = gas_tank.max_per_period else {
+            let Some(capacity) = tank.capacity_per_period else {
                 return Some(Weight::MAX);
             };
 
-            if block_number.checked_sub(&gas_tank.since)? > period {
-                gas_tank.since = block_number.checked_add(&period)?;
-                gas_tank.used = Weight::zero();
-
-                F::set_typed_attribute(&collection, &item, &ATTR_MEMBERSHIP_GAS, &gas_tank).ok()?;
+            if block_number.checked_sub(&tank.since)? > period {
+                tank.since = block_number.checked_add(&period)?;
+                tank.used = Weight::zero();
+                tank.put::<F, ItemConfig>(&collection, &item)?;
             };
 
-            let remaining = max_weight.checked_sub(&gas_tank.used.checked_add(estimated)?)?;
+            let remaining = capacity.checked_sub(&tank.used.checked_add(estimated)?)?;
             F::set_typed_attribute(
                 &collection,
                 &item,
@@ -95,15 +121,76 @@ where
                 F::clear_typed_attribute(&collection, &item, &ATTR_GAS_TX_PAY_WITH_MEMBERSHIP)
                     .ok()?;
 
-                let mut gas_tank: MembershipWeightTank<T> =
-                    F::typed_system_attribute(&collection, Some(&item), &ATTR_MEMBERSHIP_GAS)?;
+                let mut tank = WeightTank::<T>::get::<F>(&collection, &item)?;
+
+                if tank.capacity_per_period.is_some() {
+                    tank.used = tank.used.checked_add(used)?;
+                }
 
-                gas_tank.used = gas_tank.used.checked_add(used)?;
+                tank.put::<F, ItemConfig>(&collection, &item)?;
 
-                F::set_typed_attribute(&collection, &item, &ATTR_MEMBERSHIP_GAS, &gas_tank).ok()?;
-                let max_weight = gas_tank.max_per_period?;
-                Some(max_weight.saturating_sub(gas_tank.used))
+                let max_weight = tank.capacity_per_period?;
+                Some(max_weight.saturating_sub(tank.used))
             })
             .unwrap_or_default()
     }
 }
+
+impl<T, F, ItemConfig> GasFueler for NonFungibleGasTank<T, F, ItemConfig>
+where
+    T: frame_system::Config,
+    BlockNumberFor<T>: Bounded,
+    F: nonfungibles_v2::Inspect<T::AccountId>
+        + nonfungibles_v2::InspectEnumerable<T::AccountId>
+        + nonfungibles_v2::Mutate<T::AccountId, ItemConfig>,
+    ItemConfig: Default,
+    F::CollectionId: 'static,
+    F::ItemId: 'static,
+{
+    type TankId = (F::CollectionId, F::ItemId);
+    type Gas = Weight;
+
+    fn refuel_gas((collection_id, item_id): &Self::TankId, gas: &Self::Gas) -> Self::Gas {
+        let Some(mut tank) = WeightTank::<T>::get::<F>(collection_id, item_id) else {
+            return Self::Gas::zero();
+        };
+
+        if tank.capacity_per_period.is_none() {
+            return Self::Gas::MAX;
+        }
+
+        tank.used = tank.used.saturating_sub(*gas);
+
+        // Should infallibly save the tank, given that it already got a tank
+        tank.put::<F, ItemConfig>(collection_id, item_id)
+            .unwrap_or_default();
+
+        tank.capacity_per_period
+            .unwrap_or_default()
+            .saturating_sub(tank.used)
+    }
+}
+
+impl<T, F, ItemConfig> MakeTank for NonFungibleGasTank<T, F, ItemConfig>
+where
+    T: frame_system::Config,
+    BlockNumberFor<T>: Bounded,
+    F: nonfungibles_v2::Inspect<T::AccountId>
+        + nonfungibles_v2::InspectEnumerable<T::AccountId>
+        + nonfungibles_v2::Mutate<T::AccountId, ItemConfig>,
+    ItemConfig: Default,
+    F::CollectionId: 'static,
+    F::ItemId: 'static,
+{
+    type TankId = (F::CollectionId, F::ItemId);
+    type Gas = Weight;
+    type BlockNumber = BlockNumberFor<T>;
+
+    fn make_tank(
+        (collection_id, item_id): &Self::TankId,
+        capacity: Option<Self::Gas>,
+        periodicity: Option<Self::BlockNumber>,
+    ) -> Option<()> {
+        WeightTank::<T>::new(capacity, periodicity).put::<F, ItemConfig>(collection_id, item_id)
+    }
+}

traits/gas-tank/src/lib.rs

@@ -1,6 +1,7 @@
 #![cfg_attr(not(feature = "std"), no_std)]
 
 use frame_support::Parameter;
+use sp_runtime::traits::BlockNumber;
 
 #[cfg(test)]
 mod tests;
@@ -9,7 +10,7 @@ mod impl_nonfungibles;
 
 pub trait GasTank: GasBurner + GasFueler {}
 
-pub use impl_nonfungibles::NonFungibleGasBurner;
+pub use impl_nonfungibles::NonFungibleGasTank;
 
 /// Handles burning _"gas"_ from a tank to be spendable in transactions
 pub trait GasBurner {
@@ -29,11 +30,27 @@ pub trait GasBurner {
 
 /// Handles fueling _"gas"_ on a tank to spend in future transactions
 pub trait GasFueler {
-    type AccountId: Parameter;
+    type TankId: Parameter;
     type Gas: Parameter;
 
     /// Refills as much `gas` as possible returning what the updated amount of gas in the tank.
     ///
     /// This method is expected not to fail.
-    fn refuel_gas(who: &Self::AccountId, gas: &Self::Gas) -> Self::Gas;
+    fn refuel_gas(id: &Self::TankId, gas: &Self::Gas) -> Self::Gas;
+}
+
+pub trait MakeTank {
+    type TankId: Parameter;
+    type Gas: Parameter;
+    type BlockNumber: BlockNumber;
+
+    /// Creates a new tank, allowing to specify a max gas `capacity` and a `periodicity` after
+    /// which the tank gets renewed.
+    ///
+    /// Returns `Some(())` if the creation was successful, or `None` otherwise.
+    fn make_tank(
+        id: &Self::TankId,
+        capacity: Option<Self::Gas>,
+        periodicity: Option<Self::BlockNumber>,
+    ) -> Option<()>;
 }

traits/gas-tank/src/tests.rs

@@ -6,7 +6,7 @@ use frame_support::{
     weights::Weight,
 };
 use frame_system::EnsureNever;
-use impl_nonfungibles::{MembershipWeightTank, NonFungibleGasBurner, ATTR_MEMBERSHIP_GAS};
+use impl_nonfungibles::{NonFungibleGasTank, WeightTank, ATTR_MEMBERSHIP_GAS};
 use sp_runtime::{
     traits::{IdentifyAccount, IdentityLookup, Verify},
     MultiSignature,
@@ -89,18 +89,20 @@ impl pallet_nfts::Config for Test {
     type Helper = ();
 }
 
-pub type MembershipsGas = NonFungibleGasBurner<Test, Memberships, pallet_nfts::ItemConfig>;
+pub type MembershipsGas = NonFungibleGasTank<Test, Memberships, pallet_nfts::ItemConfig>;
 
 parameter_types! {
     const CollectionOwner: AccountId = AccountId::new([0u8;32]);
 
     const SmallMember: AccountId = AccountId::new([1u8;32]);
     const MediumMember: AccountId = AccountId::new([2u8;32]);
     const LargeMember: AccountId = AccountId::new([3u8;32]);
+    const ExtraLargeMember: AccountId = AccountId::new([4u8;32]);
 
     SmallTank: Weight = <() as frame_system::WeightInfo>::remark(100);
     MediumTank: Weight = <() as frame_system::WeightInfo>::remark(1000);
     LargeTank: Weight = <() as frame_system::WeightInfo>::remark(10000);
+    ExtraLargeTank: Weight = <() as frame_system::WeightInfo>::remark(100000);
 }
 
 pub(crate) fn new_test_ext() -> sp_io::TestExternalities {
@@ -120,24 +122,24 @@ pub(crate) fn new_test_ext() -> sp_io::TestExternalities {
             (
                 1,
                 SmallMember::get(),
-                MembershipWeightTank::<Test> {
-                    max_per_period: Some(SmallTank::get()),
+                WeightTank::<Test> {
+                    capacity_per_period: Some(SmallTank::get()),
                     ..Default::default()
                 },
             ),
             (
                 2,
                 MediumMember::get(),
-                MembershipWeightTank::<Test> {
-                    max_per_period: Some(MediumTank::get()),
+                WeightTank::<Test> {
+                    capacity_per_period: Some(MediumTank::get()),
                     ..Default::default()
                 },
             ),
             (
                 3,
                 LargeMember::get(),
-                MembershipWeightTank::<Test> {
-                    max_per_period: Some(LargeTank::get()),
+                WeightTank::<Test> {
+                    capacity_per_period: Some(LargeTank::get()),
                     ..Default::default()
                 },
             ),
@@ -239,3 +241,84 @@ mod gas_burner {
         });
     }
 }
+
+mod gas_fueler {
+    use super::*;
+
+    #[test]
+    fn it_works() {
+        new_test_ext().execute_with(|| {
+            // Burn gas on large tank
+            let remaining = MembershipsGas::check_available_gas(
+                &LargeMember::get(),
+                &<() as frame_system::WeightInfo>::remark(1000),
+            )
+            .expect("gas to burn equals tank capacity; qed");
+
+            assert_eq!(
+                MembershipsGas::burn_gas(
+                    &LargeMember::get(),
+                    &remaining,
+                    &<() as frame_system::WeightInfo>::remark(5000)
+                ),
+                LargeTank::get().saturating_sub(<() as frame_system::WeightInfo>::remark(5000))
+            );
+
+            // Refuels gas
+            assert_eq!(
+                MembershipsGas::refuel_gas(
+                    &(1, 3),
+                    &<() as frame_system::WeightInfo>::remark(5000)
+                ),
+                LargeTank::get()
+            );
+        })
+    }
+}
+
+mod make_tank {
+    use super::*;
+
+    #[test]
+    fn it_works() {
+        use frame_support::traits::nonfungibles_v2::Mutate;
+
+        new_test_ext().execute_with(|| {
+            assert_ok!(Memberships::mint_into(
+                &1,
+                &4,
+                &ExtraLargeMember::get(),
+                &Default::default(),
+                true,
+            ));
+
+            MembershipsGas::make_tank(&(1, 4), Some(ExtraLargeTank::get()), None)
+                .expect("failed to register the tank");
+
+            // Burn gas on large tank
+            let remaining = MembershipsGas::check_available_gas(
+                &ExtraLargeMember::get(),
+                &ExtraLargeTank::get(),
+            )
+            .expect("gas to burn equals tank capacity; qed");
+
+            assert_eq!(
+                MembershipsGas::burn_gas(
+                    &ExtraLargeMember::get(),
+                    &remaining,
+                    &ExtraLargeTank::get(),
+                ),
+                Weight::zero()
+            );
+
+            // Refuels gas
+            assert_eq!(
+                MembershipsGas::refuel_gas(
+                    &(1, 4),
+                    &<() as frame_system::WeightInfo>::remark(100000)
+                ),
+                ExtraLargeTank::get()
+            );
+        })
+    }
+}