sui_display/v2/
meter.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::v2::error::FormatError;
5
6/// Limits that the parser enforces while parsing potentially multiple Display formats.
7#[derive(Clone)]
8pub struct Limits {
9    /// Maximum number of times the parser can recurse into nested structures. Depth does not
10    /// account for all nodes, only nodes that can be contained within themselves.
11    pub max_depth: usize,
12
13    /// Maximum number of AST nodes that can be allocated during parsing. This counts all values
14    /// that are instances of AST types (but not, for example, `Vec<T>`).
15    pub max_nodes: usize,
16
17    /// Maximum number of times the format can try to load an object.
18    pub max_loads: usize,
19}
20
21/// The available budget left for limits that are tracked across all invocations to the parser for
22/// a single Display.
23pub(crate) struct Budget {
24    pub nodes: usize,
25    pub loads: usize,
26}
27
28pub(crate) struct Meter<'b> {
29    depth_budget: usize,
30    budget: &'b mut Budget,
31}
32
33impl Limits {
34    pub(crate) fn budget(&self) -> Budget {
35        Budget {
36            nodes: self.max_nodes,
37            loads: self.max_loads,
38        }
39    }
40}
41
42impl<'b> Meter<'b> {
43    pub fn new(max_depth: usize, budget: &'b mut Budget) -> Self {
44        Meter {
45            depth_budget: max_depth,
46            budget,
47        }
48    }
49
50    /// Create a nested meter, with a reduced depth budget.
51    pub fn nest(&mut self) -> Result<Meter<'_>, FormatError> {
52        if self.depth_budget == 0 {
53            return Err(FormatError::TooDeep);
54        }
55
56        Ok(Meter {
57            depth_budget: self.depth_budget - 1,
58            budget: self.budget,
59        })
60    }
61
62    /// Signal that a node has been allocated.
63    pub fn alloc(&mut self) -> Result<(), FormatError> {
64        if self.budget.nodes == 0 {
65            return Err(FormatError::TooBig);
66        }
67
68        self.budget.nodes -= 1;
69        Ok(())
70    }
71
72    /// Signal that a load could be performed.
73    pub fn load(&mut self) -> Result<(), FormatError> {
74        if self.budget.loads == 0 {
75            return Err(FormatError::TooManyLoads);
76        }
77
78        self.budget.loads -= 1;
79        Ok(())
80    }
81}
82
83impl Default for Limits {
84    fn default() -> Self {
85        Self {
86            max_depth: 32,
87            max_nodes: 32768,
88            max_loads: 8,
89        }
90    }
91}