Skip to main content

gear_common/gas_provider/
internal.rs

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