feat(fc-traits-gas-tank): use the newly -defined `SelectNonFungibleIt…

…em` trait to allow selecting (or discarding) items based on a custom rule.
virto-network · Dec 19, 2024 · f79db1a · f79db1a
1 parent 71d38dc
commit f79db1a
Show file tree

Hide file tree

Showing 8 changed files with 108 additions and 11 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -21,6 +21,7 @@ quote = { version = "1" }
 fc-traits-authn = { path = "./traits/authn", default-features = false }
 fc-traits-authn-proc = { path = "./traits/authn/proc", default-features = false }
 fc-traits-gas-tank = { path = "./traits/gas-tank", default-features = false }
+fc-traits-nonfungibles-helpers = { path = "./traits/nonfungibles-helpers", default-features = false }
 fc-traits-tracks = { path = "./traits/tracks", default-features = false }
 
 frame-benchmarking = { git = "https://github.com/virto-network/polkadot-sdk", branch = "release-virto-stable2409", default-features = false }
@@ -53,5 +54,6 @@ members = [
   "traits/authn/proc",
   "traits/gas-tank",
   "traits/memberships",
+    "traits/nonfungibles-helpers",
   "traits/tracks",
 ]
diff --git a/traits/gas-tank/Cargo.toml b/traits/gas-tank/Cargo.toml
@@ -10,6 +10,7 @@ version = "0.1.0"
 codec.workspace = true
 frame-support.workspace = true
 frame-system.workspace = true
+fc-traits-nonfungibles-helpers.workspace = true
 sp-runtime.workspace = true
 
 [dev-dependencies]

diff --git a/traits/gas-tank/src/impl_nonfungibles.rs b/traits/gas-tank/src/impl_nonfungibles.rs
@@ -1,5 +1,6 @@
 use core::marker::PhantomData;
 
+use frame_support::traits::Get;
 use frame_support::{
     pallet_prelude::{Decode, Encode},
     traits::nonfungibles_v2,
@@ -10,6 +11,8 @@ use sp_runtime::traits::{Bounded, CheckedAdd, CheckedSub};
 
 use crate::*;
 
+pub use fc_traits_nonfungibles_helpers::SelectNonFungibleItem;
+
 pub const ATTR_MEMBERSHIP_GAS: &[u8] = b"membership_gas";
 pub const ATTR_GAS_TX_PAY_WITH_MEMBERSHIP: &[u8] = b"mbmshp_pays_gas";
 
@@ -64,22 +67,34 @@ where
     }
 }
 
-pub struct NonFungibleGasTank<T, F, I>(PhantomData<(T, F, I)>);
+pub struct Noop;
+impl Get<Box<()>> for Noop {
+    fn get() -> Box<()> {
+        Box::new(())
+    }
+}
 
-impl<T, F, ItemConfig> GasBurner for NonFungibleGasTank<T, F, ItemConfig>
+pub struct NonFungibleGasTank<T, F, I, S = Noop>(PhantomData<(T, F, I, S)>);
+
+impl<T, F, I, S> GasBurner for NonFungibleGasTank<T, F, I, S>
 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,
+        + nonfungibles_v2::Mutate<T::AccountId, I>,
+    I: Default,
+    S: Get<Box<dyn SelectNonFungibleItem<F::CollectionId, F::ItemId>>>,
 {
     type AccountId = T::AccountId;
     type Gas = Weight;
 
     fn check_available_gas(who: &Self::AccountId, estimated: &Self::Gas) -> Option<Self::Gas> {
         F::owned(who).find_map(|(collection, item)| {
+            if !S::get().select(collection.clone(), item.clone()) {
+                return None;
+            }
+
             let mut tank = WeightTank::<T>::get::<F>(&collection, &item)?;
 
             let block_number = frame_system::Pallet::<T>::block_number();
@@ -92,7 +107,7 @@ where
             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).ok()?;
+                tank.put::<F, I>(&collection, &item).ok()?;
             };
 
             let remaining = capacity.checked_sub(&tank.used.checked_add(estimated)?)?;
@@ -127,7 +142,7 @@ where
                     tank.used = tank.used.checked_add(used)?;
                 }
 
-                tank.put::<F, ItemConfig>(&collection, &item).ok()?;
+                tank.put::<F, I>(&collection, &item).ok()?;
 
                 let max_weight = tank.capacity_per_period?;
                 Some(max_weight.saturating_sub(tank.used))
@@ -136,7 +151,7 @@ where
     }
 }
 
-impl<T, F, ItemConfig> GasFueler for NonFungibleGasTank<T, F, ItemConfig>
+impl<T, F, ItemConfig, S> GasFueler for NonFungibleGasTank<T, F, ItemConfig, S>
 where
     T: frame_system::Config,
     BlockNumberFor<T>: Bounded,
@@ -146,11 +161,15 @@ where
     ItemConfig: Default,
     F::CollectionId: 'static,
     F::ItemId: 'static,
+    S: Get<Box<dyn SelectNonFungibleItem<F::CollectionId, F::ItemId>>>,
 {
     type TankId = (F::CollectionId, F::ItemId);
     type Gas = Weight;
 
     fn refuel_gas((collection_id, item_id): &Self::TankId, gas: &Self::Gas) -> Self::Gas {
+        if !S::get().select(collection_id.clone(), item_id.clone()) {
+            return Self::Gas::zero();
+        }
         let Some(mut tank) = WeightTank::<T>::get::<F>(collection_id, item_id) else {
             return Self::Gas::zero();
         };
@@ -171,7 +190,7 @@ where
     }
 }
 
-impl<T, F, ItemConfig> MakeTank for NonFungibleGasTank<T, F, ItemConfig>
+impl<T, F, ItemConfig, S> MakeTank for NonFungibleGasTank<T, F, ItemConfig, S>
 where
     T: frame_system::Config,
     BlockNumberFor<T>: Bounded,

diff --git a/traits/gas-tank/src/lib.rs b/traits/gas-tank/src/lib.rs
@@ -11,7 +11,7 @@ mod impl_nonfungibles;
 
 pub trait GasTank: GasBurner + GasFueler {}
 
-pub use impl_nonfungibles::NonFungibleGasTank;
+pub use impl_nonfungibles::{NonFungibleGasTank, SelectNonFungibleItem};
 
 /// Handles burning _"gas"_ from a tank to be spendable in transactions
 pub trait GasBurner {

diff --git a/traits/gas-tank/src/tests.rs b/traits/gas-tank/src/tests.rs
@@ -1,5 +1,6 @@
 use super::*;
 
+use fc_traits_nonfungibles_helpers::SelectNonFungibleItem;
 use frame_support::{
     assert_ok, derive_impl, parameter_types,
     traits::{ConstU128, ConstU32},
@@ -89,7 +90,15 @@ impl pallet_nfts::Config for Test {
     type Helper = ();
 }
 
-pub type MembershipsGas = NonFungibleGasTank<Test, Memberships, pallet_nfts::ItemConfig>;
+#[frame_support::storage_alias]
+type Toggle = StorageValue<Prefix, bool, frame_support::pallet_prelude::ValueQuery>;
+
+parameter_types! {
+     pub ToggleBasedSelector: Box<dyn SelectNonFungibleItem<u16, u32>> = Box::new(|_, _| Toggle::get());
+}
+
+pub type MembershipsGas =
+    NonFungibleGasTank<Test, Memberships, pallet_nfts::ItemConfig, ToggleBasedSelector>;
 
 parameter_types! {
     const CollectionOwner: AccountId = AccountId::new([0u8;32]);
@@ -111,6 +120,8 @@ pub(crate) fn new_test_ext() -> sp_io::TestExternalities {
     let collection_id = 1;
     let mut ext = sp_io::TestExternalities::default();
     ext.execute_with(|| {
+        Toggle::put(true);
+
         assert_ok!(Memberships::create_collection_with_id(
             collection_id,
             &CollectionOwner::get(),
@@ -167,6 +178,18 @@ mod gas_burner {
 
     use super::*;
 
+    #[test]
+    fn fail_if_selector_discards_a_membership() {
+        new_test_ext().execute_with(|| {
+            Toggle::put(false);
+            assert!(MembershipsGas::check_available_gas(
+                &SmallMember::get(),
+                &<() as frame_system::WeightInfo>::remark(100),
+            )
+            .is_none());
+        })
+    }
+
     #[test]
     fn fail_if_gas_is_larger_than_membership_capacity() {
         new_test_ext().execute_with(|| {

diff --git a/traits/nonfungibles-helpers/Cargo.toml b/traits/nonfungibles-helpers/Cargo.toml
@@ -0,0 +1,7 @@
+[package]
+authors.workspace = true
+edition.workspace = true
+license.workspace = true
+name = "fc-traits-nonfungibles-helpers"
+repository.workspace = true
+version = "0.1.0"
diff --git a/traits/nonfungibles-helpers/src/lib.rs b/traits/nonfungibles-helpers/src/lib.rs
@@ -0,0 +1,40 @@
+/// Defines a predicate to determine whether to select an item or not within a
+/// filter operation.
+pub trait SelectNonFungibleItem<CollectionId, ItemId> {
+    /// Returns `true` if the item should be selected in a filter operation
+    fn select(&self, collection_id: CollectionId, item_id: ItemId) -> bool;
+}
+
+impl<C, I> SelectNonFungibleItem<C, I> for () {
+    fn select(&self, _: C, _: I) -> bool {
+        true
+    }
+}
+
+impl<CollectionId, ItemId, T> SelectNonFungibleItem<CollectionId, ItemId> for T
+where
+    T: Fn(CollectionId, ItemId) -> bool,
+    CollectionId: Clone,
+    ItemId: Clone,
+{
+    fn select(&self, collection_id: CollectionId, item_id: ItemId) -> bool {
+        self(collection_id, item_id)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn it_works() {
+        // No-op works
+        assert!(SelectNonFungibleItem::<u32, u32>::select(&(), 1, 1));
+        assert!(SelectNonFungibleItem::<u32, u32>::select(&(), 1, 2));
+
+        // Specific function works
+        let select: Box<dyn SelectNonFungibleItem<u32, u32>> = Box::new(|a, b| (a + b) < 3u32);
+        assert!(select.select(1, 1));
+        assert!(!select.select(1, 2));
+    }
+}