gear_common/gas_provider/
internal.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::*;
20use crate::storage::MapStorage;
21
22/// Output of `TreeImpl::catch_value` call.
23#[derive(Debug, Clone, Copy)]
24enum CatchValueOutput<Balance> {
25    /// Catching value is impossible, therefore blocked.
26    Blocked,
27    /// Value was not caught, because was moved to the patron node.
28    ///
29    /// For more info about patron nodes see `TreeImpl::find_ancestor_patron`
30    Missed,
31    /// Value was caught and will be removed from the node
32    Caught(Balance),
33}
34
35impl<Balance: BalanceTrait> CatchValueOutput<Balance> {
36    fn into_consume_output<ExternalId, Funds>(
37        self,
38        origin: ExternalId,
39        multiplier: GasMultiplier<Funds, Balance>,
40    ) -> Option<(
41        NegativeImbalance<Balance>,
42        GasMultiplier<Funds, Balance>,
43        ExternalId,
44    )> {
45        match self {
46            CatchValueOutput::Caught(value) => {
47                Some((NegativeImbalance::new(value), multiplier, origin))
48            }
49            _ => None,
50        }
51    }
52
53    fn is_blocked(&self) -> bool {
54        matches!(self, CatchValueOutput::Blocked)
55    }
56
57    fn is_caught(&self) -> bool {
58        matches!(self, CatchValueOutput::Caught(_))
59    }
60}
61
62pub struct TreeImpl<TotalValue, InternalError, Error, ExternalId, NodeId, StorageMap>(
63    PhantomData<(
64        TotalValue,
65        InternalError,
66        Error,
67        ExternalId,
68        NodeId,
69        StorageMap,
70    )>,
71);
72
73impl<TotalValue, Balance, Funds, InternalError, Error, ExternalId, NodeId, StorageMap>
74    TreeImpl<TotalValue, InternalError, Error, ExternalId, NodeId, StorageMap>
75where
76    Balance: BalanceTrait,
77    Funds: Clone,
78    TotalValue: ValueStorage<Value = Balance>,
79    InternalError: super::Error,
80    Error: From<InternalError>,
81    ExternalId: Clone,
82    NodeId: Copy,
83    StorageMap: MapStorage<Key = NodeId, Value = GasNode<ExternalId, NodeId, Balance, Funds>>,
84{
85    pub(super) fn get_node(key: impl Into<NodeId>) -> Option<StorageMap::Value> {
86        StorageMap::get(&key.into())
87    }
88
89    /// Returns the first parent, that is able to hold a concrete value, but
90    /// doesn't necessarily have a non-zero value, along with it's id.
91    ///
92    /// Node itself is considered as a self-parent too. The gas tree holds
93    /// invariant, that all the nodes with unspecified value always have a
94    /// parent with a specified value.
95    ///
96    /// The id of the returned node is of `Option` type. If it's `None`, it
97    /// means, that the ancestor and `self` are the same.
98    pub(super) fn node_with_value(
99        node: StorageMap::Value,
100    ) -> Result<(StorageMap::Value, Option<NodeId>), Error> {
101        let mut ret_node = node;
102        let mut ret_id = None;
103        if let GasNode::UnspecifiedLocal { parent, .. } = ret_node {
104            ret_id = Some(parent);
105            ret_node = Self::get_node(parent).ok_or_else(InternalError::parent_is_lost)?;
106            if !(ret_node.is_external() || ret_node.is_specified_local() || ret_node.is_reserved())
107            {
108                return Err(InternalError::unexpected_node_type().into());
109            }
110        }
111
112        Ok((ret_node, ret_id))
113    }
114
115    pub(super) fn decrease_parents_ref(node: &StorageMap::Value) -> Result<(), Error> {
116        let id = match node.parent() {
117            Some(id) => id,
118            None => return Ok(()),
119        };
120
121        let mut parent = Self::get_node(id).ok_or_else(InternalError::parent_is_lost)?;
122        if parent.refs() == 0 {
123            return Err(InternalError::parent_has_no_children().into());
124        }
125
126        match node {
127            GasNode::SpecifiedLocal { .. } => {
128                parent.decrease_spec_refs();
129            }
130            GasNode::UnspecifiedLocal { .. } => {
131                parent.decrease_unspec_refs();
132            }
133            _ => return Err(InternalError::unexpected_node_type().into()),
134        }
135
136        // Update parent node
137        StorageMap::insert(id, parent);
138
139        Ok(())
140    }
141
142    /// Tries to __"catch"__ the value inside the node if possible.
143    ///
144    /// If the node is a patron or of unspecified type, value is blocked, i.e.
145    /// can't be removed or impossible to hold value to be removed.
146    ///
147    /// If the node is not a patron, but it has an ancestor patron, value is
148    /// moved to it. So the patron's balance is increased (mutated).
149    /// Otherwise the value is caught and removed from the tree. In both
150    /// cases the `self` node's balance is zeroed.
151    ///
152    /// # Note
153    /// Method doesn't mutate `self` in the storage, but only changes it's
154    /// balance in memory.
155    fn catch_value(node: &mut StorageMap::Value) -> Result<CatchValueOutput<Balance>, Error> {
156        if node.is_patron() {
157            return Ok(CatchValueOutput::Blocked);
158        }
159
160        if !node.is_unspecified_local() {
161            if let Some((mut patron, patron_id)) = Self::find_ancestor_patron(node)? {
162                let self_value = node
163                    .value_mut()
164                    .ok_or_else(InternalError::unexpected_node_type)?;
165                if self_value.is_zero() {
166                    // Early return to prevent redundant storage look-ups
167                    return Ok(CatchValueOutput::Missed);
168                }
169                let patron_value = patron
170                    .value_mut()
171                    .ok_or_else(InternalError::unexpected_node_type)?;
172                *patron_value = patron_value.saturating_add(*self_value);
173                *self_value = Zero::zero();
174                StorageMap::insert(patron_id, patron);
175
176                Ok(CatchValueOutput::Missed)
177            } else {
178                let self_value = node
179                    .value_mut()
180                    .ok_or_else(InternalError::unexpected_node_type)?;
181                let value_copy = *self_value;
182                *self_value = Zero::zero();
183
184                Ok(CatchValueOutput::Caught(value_copy))
185            }
186        } else {
187            Ok(CatchValueOutput::Blocked)
188        }
189    }
190
191    /// Looks for `self` node's patron ancestor.
192    ///
193    /// A patron node is the node, on which some other nodes in the tree rely.
194    /// More precisely, unspecified local nodes rely on nodes with value, so
195    /// specified nodes as `GasNode::External` and `GasNode::SpecifiedLocal`
196    /// are patron ones. The other criteria for a node to be marked as the
197    /// patron one is not being consumed - value of such nodes mustn't be
198    /// moved, because node itself rely on it.
199    #[allow(clippy::type_complexity)]
200    fn find_ancestor_patron(
201        node: &StorageMap::Value,
202    ) -> Result<Option<(StorageMap::Value, NodeId)>, Error> {
203        match node {
204            GasNode::External { .. } | GasNode::Cut { .. } | GasNode::Reserved { .. } => Ok(None),
205            GasNode::SpecifiedLocal { parent, .. } => {
206                let mut ret_id = *parent;
207                let mut ret_node =
208                    Self::get_node(*parent).ok_or_else(InternalError::parent_is_lost)?;
209                while !ret_node.is_patron() {
210                    match ret_node {
211                        GasNode::External { .. } | GasNode::Reserved { .. } => return Ok(None),
212                        GasNode::SpecifiedLocal { parent, .. } => {
213                            ret_id = parent;
214                            ret_node =
215                                Self::get_node(parent).ok_or_else(InternalError::parent_is_lost)?;
216                        }
217                        _ => return Err(InternalError::unexpected_node_type().into()),
218                    }
219                }
220                Ok(Some((ret_node, ret_id)))
221            }
222            // Although unspecified local type has a patron parent, it's considered
223            // an error to call the method from that type of gas node.
224            GasNode::UnspecifiedLocal { .. } => Err(InternalError::forbidden().into()),
225        }
226    }
227
228    /// Tries to remove consumed nodes on the same path from the `key` node to
229    /// the root (including it). While trying to remove nodes, also catches
230    /// value stored in them is performed.
231    ///
232    /// Value catch is performed for all the non-patron nodes on the path from
233    /// `key` to root, until some patron node is reached. By the invariant,
234    /// catching can't be blocked, because the node is not a patron.
235    ///
236    /// For node removal there are 2 main requirements:
237    /// 1. It's not a patron node
238    /// 2. It doesn't have any children nodes.
239    ///
240    /// Although the value in nodes is moved or returned to the origin, calling
241    /// `GasNode::catch_value` in this procedure can still result in catching
242    /// non-zero value. That's possible for example, when gasful parent is
243    /// consumed and has a gas-less child. When gas-less child is consumed
244    /// in `ValueTree::consume` call, the gasful parent's value is caught
245    /// in this function.
246    ///
247    /// # Invariants
248    /// Internal invariant of the procedure:
249    ///
250    /// 1. If `catch_value` call ended up with `CatchValueOutput::Missed` in
251    ///    `consume`, all the calls of catch_value on ancestor nodes will be
252    ///    `CatchValueOutput::Missed` as well.
253    ///
254    /// That's because if there is an existing ancestor patron on the path from
255    /// the `key` node to the root, catching value on all the nodes before that
256    /// patron on this same path will give the same `CatchValueOutput::Missed`
257    /// result due to the fact that they all have same ancestor patron, which
258    /// will receive their values.
259    ///
260    /// 2. Also in that case cascade ancestors consumption will last until
261    ///    either the patron node or the first ancestor with specified child found.
262    ///
263    /// 3. If `catch_value` call ended up with `CatchValueOutput::Caught(x)` in
264    ///    `consume`, all the calls of `catch_value` on ancestor nodes will be
265    ///    `CatchValueOutput::Caught(0)`.
266    ///
267    /// That's due to the 12-th invariant stated in [`super::property_tests`]
268    ///   module docs. When node becomes consumed without unspec refs (i.e.,
269    ///   stops being a patron) `consume` procedure call on such node either
270    ///   moves value upstream (if there is an ancestor patron) or returns
271    ///   value to the origin. So any repetitive `catch_value` call on such
272    ///   nodes results in `CatchValueOutput::Caught(0)` (if there is an
273    ///   ancestor patron).
274    ///
275    /// So if `consume` procedure on the node with `key` id resulted in value
276    ///    being caught, it means that there are no ancestor patrons, so none of
277    ///    `catch_value` calls on the node's ancestors will return
278    ///    `CatchValueOutput::Missed`, but will return
279    ///    `CatchValueOutput::Caught(0)`.
280    fn try_remove_consumed_ancestors(
281        key: NodeId,
282        descendant_catch_output: CatchValueOutput<Balance>,
283    ) -> ConsumeResultOf<Self> {
284        let mut node_id = key;
285
286        let mut node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
287        let mut consume_output = None;
288        let (external, multiplier, _) = Self::get_origin_node(key)?;
289
290        // Descendant's `catch_value` output is used for the sake of optimization.
291        // We could easily run `catch_value` in the below `while` loop each time
292        // we process the ancestor. But that would lead to quadratic complexity
293        // of the `consume` & `try_remove_consumed_ancestors` procedures.
294        //
295        // In order to optimize that we use internal properties of the `consume`
296        // procedure described in the function's docs. The general idea of the
297        // optimization is that in some situations there is no need in
298        // `catch_value` call, because results will be the same for
299        // all the ancestors.
300        let mut catch_output = if descendant_catch_output.is_caught() {
301            CatchValueOutput::Caught(Zero::zero())
302        } else {
303            descendant_catch_output
304        };
305        while !node.is_patron() {
306            if catch_output.is_blocked() {
307                catch_output = Self::catch_value(&mut node)?;
308            }
309
310            // The node is not a patron and can't be of unspecified type.
311            if catch_output.is_blocked() {
312                return Err(InternalError::value_is_blocked().into());
313            }
314
315            consume_output = consume_output
316                .or_else(|| catch_output.into_consume_output(external.clone(), multiplier.clone()));
317
318            if node.spec_refs() == 0 {
319                Self::decrease_parents_ref(&node)?;
320                StorageMap::remove(node_id);
321
322                match node {
323                    GasNode::External { .. } | GasNode::Reserved { .. } => {
324                        if !catch_output.is_caught() {
325                            return Err(InternalError::value_is_not_caught().into());
326                        }
327                        return Ok(consume_output);
328                    }
329                    GasNode::SpecifiedLocal { parent, .. } => {
330                        node_id = parent;
331                        node = Self::get_node(parent).ok_or_else(InternalError::parent_is_lost)?;
332                    }
333                    _ => return Err(InternalError::unexpected_node_type().into()),
334                }
335            } else {
336                StorageMap::insert(node_id, node);
337                return Ok(consume_output);
338            }
339        }
340
341        Ok(consume_output)
342    }
343
344    /// Create ValueNode from node key with value
345    fn create_from_with_value(
346        key: impl Into<NodeId>,
347        new_node_key: impl Into<NodeId>,
348        amount: Balance,
349        constructor: impl FnOnce(
350            NodeId,
351            Balance,
352            &mut GasNode<ExternalId, NodeId, Balance, Funds>,
353            NodeId,
354        ) -> Result<GasNode<ExternalId, NodeId, Balance, Funds>, Error>,
355    ) -> Result<(), Error> {
356        let key = key.into();
357        let new_node_key = new_node_key.into();
358
359        // Check if there is no node with such key yet first.
360        // This also checks if key == new_node_key.
361        if StorageMap::contains_key(&new_node_key) {
362            return Err(InternalError::node_already_exists().into());
363        }
364
365        let (mut node, node_id) =
366            Self::node_with_value(Self::get_node(key).ok_or_else(InternalError::node_not_found)?)?;
367        // Check if the parent node is cut
368        if node.is_cut() {
369            return Err(InternalError::forbidden().into());
370        }
371
372        // A `node` is guaranteed to have inner_value here, because
373        // it was queried after `Self::node_with_value` call.
374        if node
375            .value()
376            .ok_or_else(InternalError::unexpected_node_type)?
377            < amount
378        {
379            return Err(InternalError::insufficient_balance().into());
380        }
381
382        let node_id = node_id.unwrap_or(key);
383
384        let new_node = constructor(key, amount, &mut node, node_id)?;
385
386        // Save new node
387        StorageMap::insert(new_node_key, new_node);
388
389        let node_value = node
390            .value_mut()
391            .ok_or_else(InternalError::unexpected_node_type)?;
392
393        *node_value = node_value.saturating_sub(amount);
394
395        StorageMap::insert(node_id, node);
396
397        Ok(())
398    }
399
400    // Get limit node fn that may work with both: consumed and not, depending on `validate` argument.
401    fn get_limit_node_impl(
402        key: impl Into<NodeId>,
403        validate: impl FnOnce(&GasNode<ExternalId, NodeId, Balance, Funds>) -> Result<(), Error>,
404    ) -> Result<(Balance, NodeId), Error> {
405        let key = key.into();
406
407        let node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
408
409        validate(&node)?;
410
411        let (node_with_value, maybe_key) = Self::node_with_value(node)?;
412
413        // The node here is external, specified or reserved hence has the inner value
414        let v = node_with_value
415            .value()
416            .ok_or_else(InternalError::unexpected_node_type)?;
417
418        Ok((v, maybe_key.unwrap_or(key)))
419    }
420}
421
422impl<TotalValue, Balance, Funds, InternalError, Error, ExternalId, NodeId, StorageMap> Tree
423    for TreeImpl<TotalValue, InternalError, Error, ExternalId, NodeId, StorageMap>
424where
425    Balance: BalanceTrait,
426    Funds: Clone,
427    TotalValue: ValueStorage<Value = Balance>,
428    InternalError: super::Error,
429    Error: From<InternalError>,
430    ExternalId: Clone,
431    NodeId: Copy,
432    StorageMap: MapStorage<Key = NodeId, Value = GasNode<ExternalId, NodeId, Balance, Funds>>,
433{
434    type ExternalOrigin = ExternalId;
435    type NodeId = NodeId;
436    type Balance = Balance;
437    type Funds = Funds;
438
439    type PositiveImbalance = PositiveImbalance<Balance>;
440    type NegativeImbalance = NegativeImbalance<Balance>;
441
442    type InternalError = InternalError;
443    type Error = Error;
444
445    fn total_supply() -> Self::Balance {
446        TotalValue::get().unwrap_or_else(Zero::zero)
447    }
448
449    fn create(
450        origin: Self::ExternalOrigin,
451        multiplier: GasMultiplier<Self::Funds, Self::Balance>,
452        key: impl Into<Self::NodeId>,
453        amount: Self::Balance,
454    ) -> Result<Self::PositiveImbalance, Self::Error> {
455        let key = key.into();
456
457        if StorageMap::contains_key(&key) {
458            return Err(InternalError::node_already_exists().into());
459        }
460
461        let node = GasNode::new(origin, multiplier, amount, false);
462
463        // Save value node to storage
464        StorageMap::insert(key, node);
465
466        let positive_imbalance = PositiveImbalance::new(amount);
467
468        // Update Total in storage
469        TotalValue::mutate(|total| {
470            positive_imbalance.apply_to(total).map_err(|_| {
471                *total = None;
472                InternalError::total_value_is_overflowed()
473            })
474        })?;
475
476        Ok(positive_imbalance)
477    }
478
479    fn get_origin_node(
480        key: impl Into<Self::NodeId>,
481    ) -> Result<OriginNodeDataOf<Self>, Self::Error> {
482        let key = key.into();
483        let node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
484
485        if let Some((external_origin, multiplier)) = node.external_data() {
486            Ok((external_origin, multiplier, key))
487        } else {
488            let root_id = node
489                .root_id()
490                .unwrap_or_else(|| unreachable!("Guaranteed by GasNode::root_id() fn"));
491
492            let root_node = Self::get_node(root_id).ok_or_else(InternalError::node_not_found)?;
493
494            let (external_origin, multiplier) = root_node
495                .external_data()
496                .unwrap_or_else(|| unreachable!("Guaranteed by GasNode::root_id() fn"));
497            Ok((external_origin, multiplier, root_id))
498        }
499    }
500
501    fn get_limit_node(
502        key: impl Into<Self::NodeId>,
503    ) -> Result<(Self::Balance, Self::NodeId), Self::Error> {
504        let key = key.into();
505
506        Self::get_limit_node_impl(key, |node| {
507            if node.is_consumed() {
508                Err(InternalError::node_was_consumed().into())
509            } else {
510                Ok(())
511            }
512        })
513    }
514
515    fn get_limit_node_consumed(
516        key: impl Into<Self::NodeId>,
517    ) -> Result<(Self::Balance, Self::NodeId), Self::Error> {
518        let key = key.into();
519
520        Self::get_limit_node_impl(key, |node| {
521            if node.is_consumed() {
522                Ok(())
523            } else {
524                Err(InternalError::forbidden().into())
525            }
526        })
527    }
528
529    /// Marks a node with `key` as consumed, if possible, and tries to return
530    /// it's value and delete it. The function performs same procedure with all
531    /// the nodes on the path from it to the root, if possible.
532    ///
533    /// Marking a node as `consumed` is possible only for `GasNode::External`
534    /// and `GasNode::SpecifiedLocal` nodes. That is because these nodes can
535    /// be not deleted after the function call, because of, for instance,
536    /// having children refs. Such nodes as `GasNode::UnspecifiedLocal`
537    /// and `GasNode::ReservedLocal` are removed when the function is
538    /// called, so there is no need for marking them as consumed.
539    ///
540    /// When consuming the node, it's value is mutated by calling `catch_value`,
541    /// which tries to either return or move value upstream if possible.
542    /// Read the `catch_value` function's documentation for details.
543    ///
544    /// To delete node, here should be two requirements:
545    /// 1. `Self::consume` was called on the node.
546    /// 2. The node has no children, i.e. spec/unspec refs.
547    ///
548    /// So if it's impossible to delete a node, then it's impossible to delete
549    /// its parent in the current call. Also if it's possible to delete a node,
550    /// then it doesn't necessarily mean that its parent will be deleted. An
551    /// example here could be the case, when during async execution original
552    /// message went to wait list, so wasn't consumed but the one generated
553    /// during the execution of the original message went to message queue
554    /// and was successfully executed.
555    fn consume(key: impl Into<Self::NodeId>) -> ConsumeResultOf<Self> {
556        let key = key.into();
557        let mut node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
558
559        #[cfg(fuzz)]
560        {
561            let s = fail::FailScenario::setup();
562            // This is a fail point with name `fail_fuzzer`.
563            // It's supposed to return an error if `FAILPOINTS`
564            // env variable is set.
565            fail::fail_point!("fail_fuzzer", |_| {
566                // We intentionally return this error, as it has
567                // unique usage here and we won't confuse it with
568                // other real errors.
569                Err(InternalError::node_already_exists().into())
570            });
571            s.teardown();
572        }
573
574        if node.is_consumed() {
575            return Err(InternalError::node_was_consumed().into());
576        }
577
578        // Check if at least one lock has not been released
579        if !node.lock().is_zero() {
580            return Err(InternalError::consumed_with_lock().into());
581        }
582
583        if let Some(system_reserve) = node.system_reserve()
584            && !system_reserve.is_zero()
585        {
586            return Err(InternalError::consumed_with_system_reservation().into());
587        }
588
589        node.mark_consumed();
590        let catch_output = Self::catch_value(&mut node)?;
591        let (external, multiplier, _) = Self::get_origin_node(key)?;
592
593        let res = if node.refs() == 0 {
594            Self::decrease_parents_ref(&node)?;
595            StorageMap::remove(key);
596
597            match node {
598                GasNode::External { .. } | GasNode::Cut { .. } | GasNode::Reserved { .. } => {
599                    if !catch_output.is_caught() {
600                        return Err(InternalError::value_is_not_caught().into());
601                    }
602                    catch_output.into_consume_output(external, multiplier)
603                }
604                GasNode::UnspecifiedLocal { parent, .. } => {
605                    if !catch_output.is_blocked() {
606                        return Err(InternalError::value_is_not_blocked().into());
607                    }
608                    Self::try_remove_consumed_ancestors(parent, catch_output)?
609                }
610                GasNode::SpecifiedLocal { parent, .. } => {
611                    if catch_output.is_blocked() {
612                        return Err(InternalError::value_is_blocked().into());
613                    }
614                    let consume_output = catch_output.into_consume_output(external, multiplier);
615                    let consume_ancestors_output =
616                        Self::try_remove_consumed_ancestors(parent, catch_output)?;
617                    match (&consume_output, consume_ancestors_output) {
618                        // value can't be caught in both procedures
619                        (Some(_), Some((neg_imb, ..))) if neg_imb.peek().is_zero() => {
620                            consume_output
621                        }
622                        (None, None) => consume_output,
623                        _ => return Err(InternalError::unexpected_consume_output().into()),
624                    }
625                }
626            }
627        } else {
628            if node.is_cut() || node.is_unspecified_local() {
629                return Err(InternalError::unexpected_node_type().into());
630            }
631
632            StorageMap::insert(key, node);
633            catch_output.into_consume_output(external, multiplier)
634        };
635
636        // Update Total in storage
637        if let Some((negative_imbalance, ..)) = res.as_ref() {
638            TotalValue::mutate(|total| {
639                negative_imbalance.apply_to(total).map_err(|_| {
640                    *total = None;
641                    InternalError::total_value_is_underflowed()
642                })
643            })?;
644        }
645
646        Ok(res)
647    }
648
649    /// Spends `amount` of gas from the ancestor of node with `key` id.
650    ///
651    /// Calling the function is possible even if an ancestor is consumed.
652    ///
653    /// ### Note:
654    /// Node is considered as an ancestor of itself.
655    fn spend(
656        key: impl Into<Self::NodeId>,
657        amount: Self::Balance,
658    ) -> Result<Self::NegativeImbalance, Self::Error> {
659        let key = key.into();
660
661        // Upstream node with a concrete value exist for any node.
662        // If it doesn't, the tree is considered invalidated.
663        let (mut node, node_id) =
664            Self::node_with_value(Self::get_node(key).ok_or_else(InternalError::node_not_found)?)?;
665
666        // A `node` is guaranteed to have inner_value here, because it was
667        // queried after `Self::node_with_value` call.
668        let node_value = node
669            .value_mut()
670            .ok_or_else(InternalError::unexpected_node_type)?;
671
672        if *node_value < amount {
673            return Err(InternalError::insufficient_balance().into());
674        }
675
676        *node_value = node_value.saturating_sub(amount);
677        log::debug!("Spent {amount:?} of gas");
678
679        // Save node that delivers limit
680        StorageMap::insert(node_id.unwrap_or(key), node);
681
682        let negative_imbalance = NegativeImbalance::new(amount);
683
684        // Update Total in storage
685        TotalValue::mutate(|total| {
686            negative_imbalance.apply_to(total).map_err(|_| {
687                *total = None;
688                InternalError::total_value_is_underflowed()
689            })
690        })?;
691
692        Ok(negative_imbalance)
693    }
694
695    fn split_with_value(
696        key: impl Into<Self::NodeId>,
697        new_key: impl Into<Self::NodeId>,
698        amount: Self::Balance,
699    ) -> Result<(), Self::Error> {
700        Self::create_from_with_value(
701            key,
702            new_key,
703            amount,
704            |_key, value, parent_node, parent_id| {
705                parent_node.increase_spec_refs();
706
707                Ok(GasNode::SpecifiedLocal {
708                    root: parent_node.root_id().unwrap_or(parent_id),
709                    value,
710                    lock: Zero::zero(),
711                    system_reserve: Zero::zero(),
712                    parent: parent_id,
713                    refs: Default::default(),
714                    consumed: false,
715                })
716            },
717        )
718    }
719
720    fn split(
721        key: impl Into<Self::NodeId>,
722        new_key: impl Into<Self::NodeId>,
723    ) -> Result<(), Self::Error> {
724        let key = key.into();
725        let new_key = new_key.into();
726
727        let (mut node, node_id) =
728            Self::node_with_value(Self::get_node(key).ok_or_else(InternalError::node_not_found)?)?;
729        let node_id = node_id.unwrap_or(key);
730
731        // Check if the value node is cut
732        if node.is_cut() {
733            return Err(InternalError::forbidden().into());
734        }
735
736        // This also checks if key == new_node_key
737        if StorageMap::contains_key(&new_key) {
738            return Err(InternalError::node_already_exists().into());
739        }
740
741        node.increase_unspec_refs();
742
743        let new_node = GasNode::UnspecifiedLocal {
744            root: node.root_id().unwrap_or(node_id),
745            parent: node_id,
746            lock: Zero::zero(),
747            system_reserve: Zero::zero(),
748        };
749
750        // Save new node
751        StorageMap::insert(new_key, new_node);
752        // Update current node
753        StorageMap::insert(node_id, node);
754
755        Ok(())
756    }
757
758    fn cut(
759        key: impl Into<Self::NodeId>,
760        new_key: impl Into<Self::NodeId>,
761        amount: Self::Balance,
762    ) -> Result<(), Self::Error> {
763        Self::create_from_with_value(
764            key,
765            new_key,
766            amount,
767            |key, value, _parent_node, _parent_id| {
768                let (id, multiplier, _) = Self::get_origin_node(key)?;
769                Ok(GasNode::Cut {
770                    id,
771                    multiplier,
772                    value,
773                    lock: Zero::zero(),
774                })
775            },
776        )
777    }
778
779    fn create_deposit(
780        key: impl Into<Self::NodeId>,
781        new_key: impl Into<Self::NodeId>,
782        amount: Self::Balance,
783    ) -> Result<(), Self::Error> {
784        Self::create_from_with_value(
785            key,
786            new_key,
787            amount,
788            |key, value, _parent_node, _parent_id| {
789                let (id, multiplier, _) = Self::get_origin_node(key)?;
790                Ok(GasNode::new(id, multiplier, value, true))
791            },
792        )
793    }
794
795    fn exists(key: impl Into<Self::NodeId>) -> bool {
796        Self::get_node(key).is_some()
797    }
798
799    fn exists_and_deposit(key: impl Into<Self::NodeId>) -> bool {
800        Self::get_node(key)
801            .map(|node| matches!(node, GasNode::External { deposit: true, .. }))
802            .unwrap_or(false)
803    }
804
805    fn clear() {
806        TotalValue::kill();
807        StorageMap::clear();
808    }
809}
810
811impl<TotalValue, Balance, Funds, InternalError, Error, ExternalId, NodeId, StorageMap> LockableTree
812    for TreeImpl<TotalValue, InternalError, Error, ExternalId, NodeId, StorageMap>
813where
814    Balance: BalanceTrait,
815    Funds: Clone,
816    TotalValue: ValueStorage<Value = Balance>,
817    InternalError: super::Error,
818    Error: From<InternalError>,
819    ExternalId: Clone,
820    NodeId: Copy,
821    StorageMap: MapStorage<Key = NodeId, Value = GasNode<ExternalId, NodeId, Balance, Funds>>,
822{
823    fn lock(
824        key: impl Into<Self::NodeId>,
825        id: LockId,
826        amount: Self::Balance,
827    ) -> Result<(), Self::Error> {
828        let key = key.into();
829
830        // Taking node to lock into.
831        let node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
832
833        // Validating that node is not consumed.
834        if node.is_consumed() {
835            return Err(InternalError::node_was_consumed().into());
836        }
837
838        // Quick quit on queried zero lock.
839        if amount.is_zero() {
840            return Ok(());
841        }
842
843        // Taking value provider for this node.
844        let (mut ancestor_node, ancestor_id) = Self::node_with_value(node)?;
845
846        // Mutating value of provider.
847        let ancestor_node_value = ancestor_node
848            .value_mut()
849            .ok_or_else(InternalError::unexpected_node_type)?;
850
851        if *ancestor_node_value < amount {
852            return Err(InternalError::insufficient_balance().into());
853        }
854
855        *ancestor_node_value = ancestor_node_value.saturating_sub(amount);
856
857        // If provider is a parent, we save it to storage, otherwise mutating
858        // current node further, saving it afterward.
859        let mut node = if let Some(ancestor_id) = ancestor_id {
860            StorageMap::insert(ancestor_id, ancestor_node);
861
862            // Unreachable error: the same queried at the beginning of function.
863            Self::get_node(key).ok_or_else(InternalError::node_not_found)?
864        } else {
865            ancestor_node
866        };
867
868        let locked = node.lock()[id];
869        node.lock_mut()[id] = locked.saturating_add(amount);
870
871        StorageMap::insert(key, node);
872
873        Ok(())
874    }
875
876    // Such implementation of moving value upper works, because:
877    //
878    // - For value-holding types (`GasNode::External` and
879    // `GasNode::SpecifiedLocal`) locking and unlocking on consumed node is denied at the moment,
880    // so on lock and unlock they will update only themselves.
881    //
882    // - For non-value-holding type (`GasNode::UnspecifiedLocal`) locking and
883    // unlocking chained with value-holding parent, which cannot be freed
884    // (can't move its balance upstream), due to existence of this
885    // unspecified node, referring it.
886    //
887    // - For reservation type (`GasNode::ReservedLocal`) locking is denied.
888    fn unlock(
889        key: impl Into<Self::NodeId>,
890        id: LockId,
891        amount: Self::Balance,
892    ) -> Result<(), Self::Error> {
893        let key = key.into();
894
895        // Taking node to unlock from.
896        let mut node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
897
898        // Validating that node is not consumed.
899        if node.is_consumed() {
900            return Err(InternalError::node_was_consumed().into());
901        }
902
903        // Quick quit on queried zero unlock.
904        if amount.is_zero() {
905            return Ok(());
906        }
907
908        // Mutating locked value of queried node.
909        let node_lock = &mut node.lock_mut()[id];
910        if *node_lock < amount {
911            return Err(InternalError::insufficient_balance().into());
912        }
913
914        *node_lock = node_lock.saturating_sub(amount);
915
916        // Taking value provider for this node.
917        let (ancestor_node, ancestor_id) = Self::node_with_value(node.clone())?;
918
919        // Mutating value of provider.
920        // If provider is a current node, we save it to storage, otherwise mutating
921        // provider node further, saving it afterward.
922        let (mut ancestor_node, ancestor_id) = if let Some(ancestor_id) = ancestor_id {
923            StorageMap::insert(key, node);
924
925            (ancestor_node, ancestor_id)
926        } else {
927            (node, key)
928        };
929
930        let ancestor_value = ancestor_node
931            .value_mut()
932            .ok_or_else(InternalError::unexpected_node_type)?;
933
934        *ancestor_value = ancestor_value.saturating_add(amount);
935
936        StorageMap::insert(ancestor_id, ancestor_node);
937
938        Ok(())
939    }
940
941    fn get_lock(key: impl Into<Self::NodeId>, id: LockId) -> Result<Self::Balance, Self::Error> {
942        let key = key.into();
943        let node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
944
945        Ok(node.lock()[id])
946    }
947}
948
949impl<TotalValue, Balance, Funds, InternalError, Error, ExternalId, NodeId, StorageMap>
950    ReservableTree for TreeImpl<TotalValue, InternalError, Error, ExternalId, NodeId, StorageMap>
951where
952    Balance: BalanceTrait,
953    Funds: Clone,
954    TotalValue: ValueStorage<Value = Balance>,
955    InternalError: super::Error,
956    Error: From<InternalError>,
957    ExternalId: Clone,
958    NodeId: Copy,
959    StorageMap: MapStorage<Key = NodeId, Value = GasNode<ExternalId, NodeId, Balance, Funds>>,
960{
961    fn reserve(
962        key: impl Into<Self::NodeId>,
963        new_key: impl Into<Self::NodeId>,
964        amount: Self::Balance,
965    ) -> Result<(), Self::Error> {
966        Self::create_from_with_value(
967            key,
968            new_key,
969            amount,
970            |key, value, _parent_node, _parent_id| {
971                let (id, multiplier, _) = Self::get_origin_node(key)?;
972                Ok(GasNode::Reserved {
973                    id,
974                    multiplier,
975                    value,
976                    lock: Zero::zero(),
977                    refs: Default::default(),
978                    consumed: false,
979                })
980            },
981        )
982    }
983
984    fn system_reserve(
985        key: impl Into<Self::NodeId>,
986        amount: Self::Balance,
987    ) -> Result<(), Self::Error> {
988        let key = key.into();
989
990        // Taking node to lock into.
991        let node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
992
993        // Validating node type to be able to contain system reservation.
994        if !node.is_system_reservable() {
995            return Err(InternalError::forbidden().into());
996        }
997
998        // Validating that node is not consumed.
999        if node.is_consumed() {
1000            return Err(InternalError::node_was_consumed().into());
1001        }
1002
1003        // Quick quit on queried zero lock.
1004        if amount.is_zero() {
1005            return Ok(());
1006        }
1007
1008        // Taking value provider for this node.
1009        let (mut ancestor_node, ancestor_id) = Self::node_with_value(node)?;
1010
1011        // Mutating value of provider.
1012        let ancestor_node_value = ancestor_node
1013            .value_mut()
1014            .ok_or_else(InternalError::unexpected_node_type)?;
1015
1016        if *ancestor_node_value < amount {
1017            return Err(InternalError::insufficient_balance().into());
1018        }
1019
1020        *ancestor_node_value = ancestor_node_value.saturating_sub(amount);
1021
1022        // If provider is a parent, we save it to storage, otherwise mutating
1023        // current node further, saving it afterward.
1024        let mut node = if let Some(ancestor_id) = ancestor_id {
1025            StorageMap::insert(ancestor_id, ancestor_node);
1026
1027            // Unreachable error: the same queried at the beginning of function.
1028            Self::get_node(key).ok_or_else(InternalError::node_not_found)?
1029        } else {
1030            ancestor_node
1031        };
1032
1033        let system_reservation = node
1034            .system_reserve_mut()
1035            .ok_or_else(InternalError::unexpected_node_type)?;
1036
1037        *system_reservation = system_reservation.saturating_add(amount);
1038
1039        StorageMap::insert(key, node);
1040
1041        Ok(())
1042    }
1043
1044    fn system_unreserve(key: impl Into<Self::NodeId>) -> Result<Self::Balance, Self::Error> {
1045        let key = key.into();
1046
1047        // Taking node to unlock from.
1048        let mut node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
1049
1050        // Validating node type to be able to contain system reservation.
1051        if !node.is_system_reservable() {
1052            return Err(InternalError::forbidden().into());
1053        }
1054
1055        // Validating that node is not consumed.
1056        if node.is_consumed() {
1057            return Err(InternalError::node_was_consumed().into());
1058        }
1059
1060        let amount = node
1061            .system_reserve()
1062            .ok_or_else(InternalError::unexpected_node_type)?;
1063
1064        // Quick quit on queried zero unlock.
1065        if amount.is_zero() {
1066            return Ok(Zero::zero());
1067        }
1068
1069        // Mutating locked value of queried node.
1070        let system_reservation = node
1071            .system_reserve_mut()
1072            .ok_or_else(InternalError::unexpected_node_type)?;
1073
1074        *system_reservation = Zero::zero();
1075
1076        // Taking value provider for this node.
1077        let (ancestor_node, ancestor_id) = Self::node_with_value(node.clone())?;
1078
1079        // Mutating value of provider.
1080        // If provider is a current node, we save it to storage, otherwise mutating
1081        // provider node further, saving it afterward.
1082        let (mut ancestor_node, ancestor_id) = if let Some(ancestor_id) = ancestor_id {
1083            StorageMap::insert(key, node);
1084
1085            (ancestor_node, ancestor_id)
1086        } else {
1087            (node, key)
1088        };
1089
1090        let ancestor_value = ancestor_node
1091            .value_mut()
1092            .ok_or_else(InternalError::unexpected_node_type)?;
1093
1094        *ancestor_value = ancestor_value.saturating_add(amount);
1095
1096        StorageMap::insert(ancestor_id, ancestor_node);
1097
1098        Ok(amount)
1099    }
1100
1101    fn get_system_reserve(key: impl Into<Self::NodeId>) -> Result<Self::Balance, Self::Error> {
1102        let node = Self::get_node(key).ok_or_else(InternalError::node_not_found)?;
1103
1104        node.system_reserve()
1105            .ok_or_else(|| InternalError::forbidden().into())
1106    }
1107}