gear_common/gas_provider/
lockable.rs

1// This file is part of Gear.
2
3// Copyright (C) 2021-2025 Gear Technologies Inc.
4// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
5
6// This program is free software: you can redistribute it and/or modify
7// it under the terms of the GNU General Public License as published by
8// the Free Software Foundation, either version 3 of the License, or
9// (at your option) any later version.
10
11// This program is distributed in the hope that it will be useful,
12// but WITHOUT ANY WARRANTY; without even the implied warranty of
13// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14// GNU General Public License for more details.
15
16// You should have received a copy of the GNU General Public License
17// along with this program. If not, see <https://www.gnu.org/licenses/>.
18
19use super::{scheduler::StorageType, *};
20pub use gear_core::gas::LockId;
21
22/// An error indicating there is no corresponding enum variant to the one provided
23#[derive(Debug)]
24pub struct TryFromStorageTypeError;
25
26impl fmt::Display for TryFromStorageTypeError {
27    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
28        write!(formatter, "Corresponding enum variant not found")
29    }
30}
31
32impl TryFrom<StorageType> for LockId {
33    type Error = TryFromStorageTypeError;
34
35    fn try_from(storage: StorageType) -> Result<Self, Self::Error> {
36        match storage {
37            StorageType::Mailbox => Ok(Self::Mailbox),
38            StorageType::Waitlist => Ok(Self::Waitlist),
39            StorageType::Reservation => Ok(Self::Reservation),
40            StorageType::DispatchStash => Ok(Self::DispatchStash),
41            _ => Err(TryFromStorageTypeError),
42        }
43    }
44}
45
46pub trait LockableTree: Tree {
47    /// Locking some value from underlying node balance.
48    ///
49    /// If `key` does not identify any value or the `amount` exceeds what's
50    /// locked under that key, an error is returned.
51    ///
52    /// This can't create imbalance as no value is burned or created.
53    fn lock(
54        key: impl Into<Self::NodeId>,
55        id: LockId,
56        amount: Self::Balance,
57    ) -> Result<(), Self::Error>;
58
59    /// Unlocking some value from node's locked balance.
60    ///
61    /// If `key` does not identify any value or the `amount` exceeds what's
62    /// locked under that key, an error is returned.
63    ///
64    /// This can't create imbalance as no value is burned or created.
65    fn unlock(
66        key: impl Into<Self::NodeId>,
67        id: LockId,
68        amount: Self::Balance,
69    ) -> Result<(), Self::Error>;
70
71    /// Unlocking all value from node's locked balance. Returns the actual amount having been unlocked
72    /// (wrapped in a `Result`)
73    ///
74    /// See [`unlock`](Self::unlock) for details.
75    fn unlock_all(key: impl Into<Self::NodeId>, id: LockId) -> Result<Self::Balance, Self::Error> {
76        let key = key.into();
77        let amount = Self::get_lock(key.clone(), id)?;
78        Self::unlock(key, id, amount.clone()).map(|_| amount)
79    }
80
81    /// Get locked value associated with given id.
82    ///
83    /// Returns errors in cases of absence associated with given key node,
84    /// or if such functionality is forbidden for specific node type:
85    /// for example, for `GasNode::ReservedLocal`.
86    fn get_lock(key: impl Into<Self::NodeId>, id: LockId) -> Result<Self::Balance, Self::Error>;
87}
88
89#[test]
90fn lock_id_enum_discriminants_are_consistent() {
91    // Important for the [`gclient::api::GearApi`] implementation:
92    // the function `migrate_program()` relies on `LockId::Reservation` having discriminant 2
93    assert_eq!(0, LockId::Mailbox as usize);
94    assert_eq!(1, LockId::Waitlist as usize);
95    assert_eq!(2, LockId::Reservation as usize);
96    assert_eq!(3, LockId::DispatchStash as usize);
97}