tycho_vm/
gas.rs

1use std::hash::BuildHasher;
2use std::num::NonZeroU64;
3use std::rc::Rc;
4use std::sync::Arc;
5
6use ahash::HashSet;
7use tycho_types::cell::{CellParts, LoadMode};
8use tycho_types::error::Error;
9use tycho_types::models::{LibDescr, SimpleLib};
10use tycho_types::prelude::*;
11
12use crate::error::VmResult;
13use crate::saferc::SafeRc;
14use crate::stack::Stack;
15use crate::util::OwnedCellSlice;
16
17/// Initialization params for [`GasConsumer`].
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub struct GasParams {
20    /// Maximum possible value of the `limit`.
21    pub max: u64,
22    /// Gas limit for the out-of-gas exception.
23    pub limit: u64,
24    /// Free gas (e.g. for external messages without any balance).
25    pub credit: u64,
26    /// Gas price (fixed point with 16 bits for fractional part).
27    pub price: u64,
28}
29
30impl GasParams {
31    pub const MAX_GAS: u64 = i64::MAX as u64;
32
33    const STUB_GAS_PRICE: u64 = 1000 << 16;
34
35    pub const fn unlimited() -> Self {
36        Self {
37            max: Self::MAX_GAS,
38            limit: Self::MAX_GAS,
39            credit: 0,
40            price: Self::STUB_GAS_PRICE,
41        }
42    }
43
44    pub const fn getter() -> Self {
45        Self {
46            max: 1000000,
47            limit: 1000000,
48            credit: 0,
49            price: Self::STUB_GAS_PRICE,
50        }
51    }
52}
53
54impl Default for GasParams {
55    #[inline]
56    fn default() -> Self {
57        Self::unlimited()
58    }
59}
60
61/// Library cells resolver.
62pub trait LibraryProvider {
63    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error>;
64
65    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error>;
66}
67
68impl<T: LibraryProvider> LibraryProvider for &'_ T {
69    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
70        T::find(self, library_hash)
71    }
72
73    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
74        T::find_ref(self, library_hash)
75    }
76}
77
78impl<T: LibraryProvider> LibraryProvider for Option<T> {
79    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
80        match self {
81            Some(this) => T::find(this, library_hash),
82            None => Ok(None),
83        }
84    }
85
86    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
87        match self {
88            Some(this) => T::find_ref(this, library_hash),
89            None => Ok(None),
90        }
91    }
92}
93
94impl<T1, T2> LibraryProvider for (T1, T2)
95where
96    T1: LibraryProvider,
97    T2: LibraryProvider,
98{
99    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
100        if let res @ Some(_) = ok!(T1::find(&self.0, library_hash)) {
101            return Ok(res);
102        }
103        T2::find(&self.1, library_hash)
104    }
105
106    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
107        if let res @ Some(_) = ok!(T1::find_ref(&self.0, library_hash)) {
108            return Ok(res);
109        }
110        T2::find_ref(&self.1, library_hash)
111    }
112}
113
114impl<T1, T2, T3> LibraryProvider for (T1, T2, T3)
115where
116    T1: LibraryProvider,
117    T2: LibraryProvider,
118    T3: LibraryProvider,
119{
120    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
121        if let res @ Some(_) = ok!(T1::find(&self.0, library_hash)) {
122            return Ok(res);
123        }
124        if let res @ Some(_) = ok!(T2::find(&self.1, library_hash)) {
125            return Ok(res);
126        }
127        T3::find(&self.2, library_hash)
128    }
129
130    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
131        if let res @ Some(_) = ok!(T1::find_ref(&self.0, library_hash)) {
132            return Ok(res);
133        }
134        if let res @ Some(_) = ok!(T2::find_ref(&self.1, library_hash)) {
135            return Ok(res);
136        }
137        T3::find_ref(&self.2, library_hash)
138    }
139}
140
141impl<T: LibraryProvider> LibraryProvider for Box<T> {
142    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
143        T::find(self, library_hash)
144    }
145
146    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
147        T::find_ref(self, library_hash)
148    }
149}
150
151impl<T: LibraryProvider> LibraryProvider for Rc<T> {
152    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
153        T::find(self, library_hash)
154    }
155
156    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
157        T::find_ref(self, library_hash)
158    }
159}
160
161impl<T: LibraryProvider> LibraryProvider for Arc<T> {
162    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
163        T::find(self, library_hash)
164    }
165
166    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
167        T::find_ref(self, library_hash)
168    }
169}
170
171/// Empty libraries provider.
172#[derive(Default, Debug, Clone, Copy)]
173pub struct NoLibraries;
174
175impl LibraryProvider for NoLibraries {
176    #[inline]
177    fn find(&self, _library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
178        Ok(None)
179    }
180
181    fn find_ref<'a>(&'a self, _library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
182        Ok(None)
183    }
184}
185
186impl LibraryProvider for Dict<HashBytes, SimpleLib> {
187    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
188        match self.get(library_hash)? {
189            Some(lib) => Ok(Some(lib.root)),
190            None => Ok(None),
191        }
192    }
193
194    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
195        match self
196            .cast_ref::<HashBytes, SimpleLibRef<'_>>()
197            .get(library_hash)?
198        {
199            Some(lib) => Ok(Some(lib.root)),
200            None => Ok(None),
201        }
202    }
203}
204
205impl LibraryProvider for Vec<Dict<HashBytes, SimpleLib>> {
206    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
207        for lib in self {
208            match lib.get(library_hash)? {
209                Some(lib) => return Ok(Some(lib.root)),
210                None => continue,
211            }
212        }
213        Ok(None)
214    }
215
216    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
217        for lib in self {
218            match lib
219                .cast_ref::<HashBytes, SimpleLibRef<'_>>()
220                .get(library_hash)?
221            {
222                Some(lib) => return Ok(Some(lib.root)),
223                None => continue,
224            }
225        }
226        Ok(None)
227    }
228}
229
230impl LibraryProvider for Dict<HashBytes, LibDescr> {
231    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
232        Ok(self.get(library_hash)?.map(|lib| lib.lib.clone()))
233    }
234
235    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
236        struct LibDescrRef<'tlb> {
237            lib: &'tlb DynCell,
238        }
239
240        impl<'a> Load<'a> for LibDescrRef<'a> {
241            fn load_from(slice: &mut CellSlice<'a>) -> Result<Self, Error> {
242                if slice.load_small_uint(2)? != 0 {
243                    return Err(Error::InvalidTag);
244                }
245                Ok(Self {
246                    lib: slice.load_reference()?,
247                })
248            }
249        }
250
251        impl EquivalentRepr<LibDescr> for LibDescrRef<'_> {}
252
253        Ok(self
254            .cast_ref::<HashBytes, LibDescrRef<'a>>()
255            .get(library_hash)?
256            .map(|lib| lib.lib))
257    }
258}
259
260impl<S: BuildHasher> LibraryProvider for std::collections::HashMap<HashBytes, SimpleLib, S> {
261    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
262        Ok(self.get(library_hash).map(|lib| lib.root.clone()))
263    }
264
265    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
266        Ok(self.get(library_hash).map(|lib| lib.root.as_ref()))
267    }
268}
269
270struct SimpleLibRef<'tlb> {
271    root: &'tlb DynCell,
272}
273
274impl<'a> Load<'a> for SimpleLibRef<'a> {
275    fn load_from(slice: &mut CellSlice<'a>) -> Result<Self, Error> {
276        slice.load_bit()?;
277        Ok(Self {
278            root: slice.load_reference()?,
279        })
280    }
281}
282
283impl EquivalentRepr<SimpleLib> for SimpleLibRef<'_> {}
284
285/// Gas tracking context.
286pub struct GasConsumer<'l> {
287    /// Maximum possible value of the `limit`.
288    gas_max: u64,
289    /// Gas limit for the out-of-gas exception.
290    gas_limit: std::cell::Cell<u64>,
291    /// Free gas (e.g. for external messages without any balance).
292    gas_credit: std::cell::Cell<u64>,
293    /// Initial gas to compute the consumed amount.
294    gas_base: std::cell::Cell<u64>,
295    /// Remaining gas available.
296    gas_remaining: std::cell::Cell<i64>,
297    /// Gas price (fixed point with 16 bits for fractional part).
298    gas_price: NonZeroU64,
299
300    /// A set of visited cells.
301    loaded_cells: std::cell::UnsafeCell<HashSet<HashBytes>>,
302    /// Libraries provider.
303    libraries: &'l dyn LibraryProvider,
304
305    /// Number of signature checks.
306    chksign_counter: std::cell::Cell<usize>,
307    /// Free gas (can be paid when using isolated consumer).
308    free_gas_consumed: std::cell::Cell<u64>,
309    /// Number of balance calls with cheap gas consumer
310    get_extra_balance_counter: std::cell::Cell<usize>,
311
312    // Missing library in case of resolving error occured.
313    missing_library: std::cell::Cell<Option<HashBytes>>,
314}
315
316impl<'l> GasConsumer<'l> {
317    pub const BUILD_CELL_GAS: u64 = 500;
318    pub const NEW_CELL_GAS: u64 = 100;
319    pub const OLD_CELL_GAS: u64 = 25;
320
321    pub const FREE_STACK_DEPTH: usize = 32;
322    pub const FREE_SIGNATURE_CHECKS: usize = 10;
323    pub const CHEAP_GET_BALANCE_CALLS: usize = 5;
324    pub const CHEAP_GET_BALANCE_GAS_THRESHOLD: u64 = 200;
325    pub const FREE_NESTED_CONT_JUMP: usize = 8;
326
327    pub const STACK_VALUE_GAS_PRICE: u64 = 1;
328    pub const TUPLE_ENTRY_GAS_PRICE: u64 = 1;
329    pub const HASH_EXT_ENTRY_GAS_PRICE: u64 = 1;
330    pub const CHK_SGN_GAS_PRICE: u64 = 4000;
331    pub const IMPLICIT_JMPREF_GAS_PRICE: u64 = 10;
332    pub const IMPLICIT_RET_GAS_PRICE: u64 = 5;
333    pub const EXCEPTION_GAS_PRICE: u64 = 50;
334    pub const RUNVM_GAS_PRICE: u64 = 40;
335
336    pub fn new(params: GasParams) -> Self {
337        static NO_LIBRARIES: NoLibraries = NoLibraries;
338
339        Self::with_libraries(params, &NO_LIBRARIES)
340    }
341
342    pub fn with_libraries(params: GasParams, libraries: &'l dyn LibraryProvider) -> Self {
343        let gas_remaining = truncate_gas(params.limit.saturating_add(params.credit));
344
345        Self {
346            gas_max: truncate_gas(params.max),
347            gas_limit: std::cell::Cell::new(truncate_gas(params.limit)),
348            gas_credit: std::cell::Cell::new(truncate_gas(params.credit)),
349            gas_base: std::cell::Cell::new(gas_remaining),
350            gas_remaining: std::cell::Cell::new(gas_remaining as i64),
351            gas_price: NonZeroU64::new(params.price).unwrap_or(NonZeroU64::MIN),
352            loaded_cells: Default::default(),
353            libraries,
354            chksign_counter: std::cell::Cell::new(0),
355            free_gas_consumed: std::cell::Cell::new(0),
356            get_extra_balance_counter: std::cell::Cell::new(0),
357            missing_library: std::cell::Cell::new(None),
358        }
359    }
360
361    pub fn limited(&'l self, remaining: u64) -> LimitedGasConsumer<'l> {
362        LimitedGasConsumer::<'l> {
363            gas: self,
364            remaining: std::cell::Cell::new(remaining),
365        }
366    }
367
368    // TODO: Consume free gas as non-free when `isolate`.
369    pub fn derive(&mut self, params: GasConsumerDeriveParams) -> VmResult<ParentGasConsumer<'l>> {
370        use std::cell::Cell;
371
372        Ok(if params.isolate {
373            // Pay for the free gas to prevent abuse.
374            self.try_consume(self.free_gas_consumed.get())?;
375
376            // NOTE: Compute remaining gas only when all operations
377            //       with parent consumer are made.
378            let gas_remaining = self.remaining();
379            vm_ensure!(gas_remaining >= 0, OutOfGas);
380
381            // Reset current free gas counters.
382            self.chksign_counter.set(0);
383            self.free_gas_consumed.set(0);
384            self.get_extra_balance_counter.set(0);
385
386            // Create a new gas consumer.
387            let params = GasParams {
388                max: std::cmp::min(params.gas_max, gas_remaining as u64),
389                limit: std::cmp::min(params.gas_limit, gas_remaining as u64),
390                credit: 0,
391                price: self.price(),
392            };
393            let libraries = self.libraries;
394
395            ParentGasConsumer::Isolated(std::mem::replace(
396                self,
397                Self::with_libraries(params, libraries),
398            ))
399        } else {
400            // NOTE: Compute remaining gas only when all operations
401            //       with parent consumer are made.
402            let gas_remaining = self.remaining();
403            vm_ensure!(gas_remaining >= 0, OutOfGas);
404
405            // Create child gas consumer params.
406            let gas_max = std::cmp::min(params.gas_max, gas_remaining as u64);
407            let gas_limit = std::cmp::min(params.gas_limit, gas_remaining as u64);
408
409            // Move gas params from child to parent.
410            ParentGasConsumer::Shared(GasConsumer {
411                gas_max: std::mem::replace(&mut self.gas_max, gas_max),
412                gas_limit: std::mem::replace(&mut self.gas_limit, Cell::new(gas_limit)),
413                gas_credit: std::mem::replace(&mut self.gas_credit, Cell::new(0)),
414                gas_base: std::mem::replace(&mut self.gas_base, Cell::new(gas_limit)),
415                gas_remaining: std::mem::replace(
416                    &mut self.gas_remaining,
417                    Cell::new(gas_limit as i64),
418                ),
419                gas_price: self.gas_price,
420                loaded_cells: Default::default(),
421                libraries: self.libraries,
422                chksign_counter: self.chksign_counter.clone(),
423                free_gas_consumed: self.free_gas_consumed.clone(),
424                get_extra_balance_counter: self.get_extra_balance_counter.clone(),
425                missing_library: self.missing_library.clone(),
426            })
427        })
428    }
429
430    pub fn restore(&mut self, mut parent: ParentGasConsumer<'l>) -> RestoredGasConsumer {
431        let meta = RestoredGasConsumer {
432            gas_consumed: self.consumed(),
433            gas_limit: self.limit(),
434        };
435
436        if let ParentGasConsumer::Shared(parent) = &mut parent {
437            // Merge loaded cells.
438            parent.loaded_cells = std::mem::take(&mut self.loaded_cells);
439        }
440
441        match parent {
442            ParentGasConsumer::Isolated(mut parent) | ParentGasConsumer::Shared(mut parent) => {
443                // Merge missing library.
444                let missing_lib = self.missing_library.get_mut();
445                let parent_lib = parent.missing_library.get_mut();
446                if parent_lib.is_none() && missing_lib.is_some() {
447                    *parent_lib = *missing_lib;
448                }
449
450                // Merge free gas counters.
451                parent.chksign_counter = self.chksign_counter.clone();
452                parent.free_gas_consumed = self.free_gas_consumed.clone();
453                parent.get_extra_balance_counter = self.get_extra_balance_counter.clone();
454
455                *self = parent
456            }
457        }
458
459        meta
460    }
461
462    pub fn libraries(&self) -> &'l dyn LibraryProvider {
463        self.libraries
464    }
465
466    pub fn credit(&self) -> u64 {
467        self.gas_credit.get()
468    }
469
470    pub fn consumed(&self) -> u64 {
471        (self.gas_base.get() as i64).saturating_sub(self.gas_remaining.get()) as u64
472    }
473
474    pub fn free_gas_consumed(&self) -> u64 {
475        self.free_gas_consumed.get()
476    }
477
478    pub fn remaining(&self) -> i64 {
479        self.gas_remaining.get()
480    }
481
482    pub fn base(&self) -> u64 {
483        self.gas_base.get()
484    }
485
486    pub fn limit(&self) -> u64 {
487        self.gas_limit.get()
488    }
489
490    pub fn set_limit(&self, limit: u64) {
491        let limit = std::cmp::min(limit, self.gas_max);
492        vm_log_trace!("changing gas limit: new_limit={limit}");
493
494        self.gas_credit.set(0);
495        self.gas_limit.set(limit);
496        self.set_base(limit);
497    }
498
499    fn set_base(&self, mut base: u64) {
500        base = truncate_gas(base);
501        let diff = base as i64 - self.gas_base.get() as i64;
502        self.gas_remaining.set(self.gas_remaining.get() + diff);
503        self.gas_base.set(base);
504    }
505
506    pub fn price(&self) -> u64 {
507        self.gas_price.get()
508    }
509
510    pub fn try_get_extra_balance_consumer(&'l self) -> Option<LimitedGasConsumer<'l>> {
511        self.get_extra_balance_counter
512            .set(self.get_extra_balance_counter.get() + 1);
513        if self.get_extra_balance_counter.get() <= Self::CHEAP_GET_BALANCE_CALLS {
514            return Some(self.limited(Self::CHEAP_GET_BALANCE_GAS_THRESHOLD));
515        }
516
517        None
518    }
519
520    pub fn try_consume_exception_gas(&self) -> Result<(), Error> {
521        self.try_consume(Self::EXCEPTION_GAS_PRICE)
522    }
523
524    pub fn try_consume_implicit_jmpref_gas(&self) -> Result<(), Error> {
525        self.try_consume(Self::IMPLICIT_JMPREF_GAS_PRICE)
526    }
527
528    pub fn try_consume_implicit_ret_gas(&self) -> Result<(), Error> {
529        self.try_consume(Self::IMPLICIT_RET_GAS_PRICE)
530    }
531
532    pub fn try_consume_check_signature_gas(&self) -> Result<(), Error> {
533        self.chksign_counter.set(self.chksign_counter.get() + 1);
534        if self.chksign_counter.get() > Self::FREE_SIGNATURE_CHECKS {
535            self.try_consume(Self::CHK_SGN_GAS_PRICE)?;
536        } else {
537            self.consume_free_gas(Self::CHK_SGN_GAS_PRICE);
538        }
539        Ok(())
540    }
541
542    pub fn try_consume_stack_gas(&self, stack: Option<&SafeRc<Stack>>) -> Result<(), Error> {
543        if let Some(stack) = stack {
544            self.try_consume_stack_depth_gas(stack.depth())?;
545        }
546        Ok(())
547    }
548
549    pub fn try_consume_tuple_gas(&self, tuple_len: u64) -> Result<(), Error> {
550        self.try_consume(tuple_len * Self::TUPLE_ENTRY_GAS_PRICE)?;
551        Ok(())
552    }
553
554    pub fn try_consume_stack_depth_gas(&self, depth: usize) -> Result<(), Error> {
555        self.try_consume(
556            (std::cmp::max(depth, Self::FREE_STACK_DEPTH) - Self::FREE_STACK_DEPTH) as u64
557                * Self::STACK_VALUE_GAS_PRICE,
558        )
559    }
560
561    pub fn try_consume(&self, amount: u64) -> Result<(), Error> {
562        let remaining = self
563            .gas_remaining
564            .get()
565            .saturating_sub(truncate_gas(amount) as i64);
566        self.gas_remaining.set(remaining);
567
568        if remaining >= 0 {
569            Ok(())
570        } else {
571            Err(Error::Cancelled)
572        }
573    }
574
575    pub fn consume_free_gas(&self, amount: u64) {
576        let consumed = truncate_gas(self.free_gas_consumed.get().saturating_add(amount));
577        self.free_gas_consumed.set(consumed);
578    }
579
580    pub fn missing_library(&self) -> Option<HashBytes> {
581        self.missing_library.get()
582    }
583
584    pub fn set_missing_library(&self, hash: &HashBytes) {
585        self.missing_library.set(Some(*hash));
586    }
587
588    pub fn load_cell_as_slice(&self, cell: Cell, mode: LoadMode) -> Result<OwnedCellSlice, Error> {
589        let cell = ok!(self.load_cell_impl(cell, mode));
590        Ok(OwnedCellSlice::new_allow_exotic(cell))
591    }
592
593    fn load_cell_impl<'s: 'a, 'a, T: LoadLibrary<'a>>(
594        &'s self,
595        mut cell: T,
596        mode: LoadMode,
597    ) -> Result<T, Error> {
598        let mut library_loaded = false;
599        loop {
600            if mode.use_gas() {
601                // SAFETY: This is the only place where we borrow `loaded_cells` as mut.
602                let is_new =
603                    unsafe { (*self.loaded_cells.get()).insert(*cell.as_ref().repr_hash()) };
604
605                ok!(self.try_consume(if is_new {
606                    GasConsumer::NEW_CELL_GAS
607                } else {
608                    GasConsumer::OLD_CELL_GAS
609                }));
610            }
611
612            if !mode.resolve() {
613                return Ok(cell);
614            }
615
616            match cell.as_ref().cell_type() {
617                CellType::Ordinary => return Ok(cell),
618                CellType::LibraryReference if !library_loaded => {
619                    // Library data structure is enforced by `CellContext::finalize_cell`.
620                    debug_assert_eq!(cell.as_ref().bit_len(), 8 + 256);
621
622                    // Find library by hash.
623                    let mut library_hash = HashBytes::ZERO;
624                    ok!(cell
625                        .as_ref()
626                        .as_slice_allow_exotic()
627                        .get_raw(8, &mut library_hash.0, 256));
628
629                    let Some(library_cell) = ok!(T::load_library(self, &library_hash)) else {
630                        self.missing_library.set(Some(library_hash));
631                        return Err(Error::CellUnderflow);
632                    };
633
634                    cell = library_cell;
635                    library_loaded = true;
636                }
637                _ => return Err(Error::CellUnderflow),
638            }
639        }
640    }
641}
642
643impl CellContext for GasConsumer<'_> {
644    fn finalize_cell(&self, cell: CellParts<'_>) -> Result<Cell, Error> {
645        ok!(self.try_consume(GasConsumer::BUILD_CELL_GAS));
646        Cell::empty_context().finalize_cell(cell)
647    }
648
649    fn load_cell(&self, cell: Cell, mode: LoadMode) -> Result<Cell, Error> {
650        self.load_cell_impl(cell, mode)
651    }
652
653    fn load_dyn_cell<'s: 'a, 'a>(
654        &'s self,
655        cell: &'a DynCell,
656        mode: LoadMode,
657    ) -> Result<&'a DynCell, Error> {
658        self.load_cell_impl(cell, mode)
659    }
660}
661
662/// Consumes at most N gas units, everything else is treated as free gas.
663pub struct LimitedGasConsumer<'l> {
664    /// Main gas consumer
665    gas: &'l GasConsumer<'l>,
666    /// Remaining gas that can be consumed.
667    remaining: std::cell::Cell<u64>,
668}
669
670impl CellContext for LimitedGasConsumer<'_> {
671    fn finalize_cell(&self, _: CellParts<'_>) -> Result<Cell, Error> {
672        panic!("limited gas consumer must not be used for making new cells")
673    }
674
675    fn load_cell(&self, cell: Cell, mode: LoadMode) -> Result<Cell, Error> {
676        self.load_cell_impl(cell, mode)
677    }
678
679    fn load_dyn_cell<'s: 'a, 'a>(
680        &'s self,
681        cell: &'a DynCell,
682        mode: LoadMode,
683    ) -> Result<&'a DynCell, Error> {
684        self.load_cell_impl(cell, mode)
685    }
686}
687
688impl LimitedGasConsumer<'_> {
689    fn load_cell_impl<'s: 'a, 'a, T: LoadLibrary<'a>>(
690        &'s self,
691        mut cell: T,
692        mode: LoadMode,
693    ) -> Result<T, Error> {
694        let mut library_loaded = false;
695        loop {
696            if mode.use_gas() {
697                // SAFETY: This is the only place where we borrow `loaded_cells` as mut.
698                let is_new =
699                    unsafe { (*self.gas.loaded_cells.get()).insert(*cell.as_ref().repr_hash()) };
700
701                ok!(self.try_consume(if is_new {
702                    GasConsumer::NEW_CELL_GAS
703                } else {
704                    GasConsumer::OLD_CELL_GAS
705                }));
706            }
707
708            if !mode.resolve() {
709                return Ok(cell);
710            }
711
712            match cell.as_ref().cell_type() {
713                CellType::Ordinary => return Ok(cell),
714                CellType::LibraryReference if !library_loaded => {
715                    // Library data structure is enforced by `CellContext::finalize_cell`.
716                    debug_assert_eq!(cell.as_ref().bit_len(), 8 + 256);
717
718                    // Find library by hash.
719                    let mut library_hash = HashBytes::ZERO;
720                    ok!(cell
721                        .as_ref()
722                        .as_slice_allow_exotic()
723                        .get_raw(8, &mut library_hash.0, 256));
724
725                    let Some(library_cell) = ok!(T::load_library(self.gas, &library_hash)) else {
726                        self.gas.missing_library.set(Some(library_hash));
727                        return Err(Error::CellUnderflow);
728                    };
729
730                    cell = library_cell;
731                    library_loaded = true;
732                }
733                _ => return Err(Error::CellUnderflow),
734            }
735        }
736    }
737
738    pub fn try_consume(&self, mut amount: u64) -> Result<(), Error> {
739        amount = truncate_gas(amount);
740
741        let mut remaining = self.remaining.get();
742        let consumed = std::cmp::min(amount, remaining);
743
744        self.gas.try_consume(consumed)?;
745
746        remaining -= consumed;
747        self.remaining.set(remaining);
748
749        // threshold reached
750        if remaining == 0 {
751            self.gas.consume_free_gas(amount - consumed);
752        }
753        Ok(())
754    }
755}
756
757trait LoadLibrary<'a>: AsRef<DynCell> + 'a {
758    fn load_library(gas: &'a GasConsumer, library_hash: &HashBytes) -> Result<Option<Self>, Error>
759    where
760        Self: Sized;
761}
762
763impl<'a> LoadLibrary<'a> for &'a DynCell {
764    fn load_library(gas: &'a GasConsumer, library_hash: &HashBytes) -> Result<Option<Self>, Error>
765    where
766        Self: Sized,
767    {
768        gas.libraries.find_ref(library_hash)
769    }
770}
771
772impl LoadLibrary<'_> for Cell {
773    fn load_library(gas: &'_ GasConsumer, library_hash: &HashBytes) -> Result<Option<Self>, Error>
774    where
775        Self: Sized,
776    {
777        gas.libraries.find(library_hash)
778    }
779}
780
781/// Params to replace the current gas consumer.
782#[derive(Debug, Clone, Copy)]
783pub struct GasConsumerDeriveParams {
784    pub gas_max: u64,
785    pub gas_limit: u64,
786    pub isolate: bool,
787}
788
789/// Parent of the derived gas consumer.
790pub enum ParentGasConsumer<'l> {
791    Isolated(GasConsumer<'l>),
792    Shared(GasConsumer<'l>),
793}
794
795/// Info extracted when parent gas consumer is restored.
796#[derive(Debug, Clone, Copy)]
797pub struct RestoredGasConsumer {
798    pub gas_consumed: u64,
799    pub gas_limit: u64,
800}
801
802const fn truncate_gas(gas: u64) -> u64 {
803    if gas <= i64::MAX as u64 {
804        gas
805    } else {
806        i64::MAX as u64
807    }
808}
809
810#[cfg(test)]
811mod tests {
812    use super::*;
813
814    #[test]
815    fn find_lib_dict_ref() {
816        let lib1 = Boc::decode(tvmasm!("NOP")).unwrap();
817        let lib2 = Boc::decode(tvmasm!("NOP NOP")).unwrap();
818
819        // Dict with SimpleLib
820        let mut libraries = vec![
821            (*lib1.repr_hash(), SimpleLib {
822                public: true,
823                root: lib1.clone(),
824            }),
825            (*lib2.repr_hash(), SimpleLib {
826                public: true,
827                root: lib2.clone(),
828            }),
829        ];
830        libraries.sort_unstable_by(|(l, _), (r, _)| l.cmp(r));
831        let libraries = Dict::<HashBytes, SimpleLib>::try_from_sorted_slice(&libraries).unwrap();
832
833        assert!(libraries.find(&HashBytes::ZERO).unwrap().is_none());
834        assert!(libraries.find_ref(&HashBytes::ZERO).unwrap().is_none());
835
836        assert_eq!(
837            libraries.find(lib1.repr_hash()).unwrap().unwrap().as_ref(),
838            libraries.find_ref(lib1.repr_hash()).unwrap().unwrap()
839        );
840        assert_eq!(
841            libraries.find(lib2.repr_hash()).unwrap().unwrap().as_ref(),
842            libraries.find_ref(lib2.repr_hash()).unwrap().unwrap()
843        );
844
845        // Dict with LibDescr
846        let mut publishers = Dict::new();
847        publishers.add(HashBytes::ZERO, ()).unwrap();
848
849        {
850            let lib = LibDescr {
851                lib: lib1.clone(),
852                publishers: publishers.clone(),
853            };
854            let c = CellBuilder::build_from(&lib).unwrap();
855            let parsed = c.parse::<LibDescr>().unwrap();
856
857            assert_eq!(lib, parsed);
858        }
859
860        let mut libraries = vec![
861            (*lib1.repr_hash(), LibDescr {
862                lib: lib1.clone(),
863                publishers: publishers.clone(),
864            }),
865            (*lib2.repr_hash(), LibDescr {
866                lib: lib2.clone(),
867                publishers,
868            }),
869        ];
870        libraries.sort_unstable_by(|(l, _), (r, _)| l.cmp(r));
871
872        let libraries = Dict::<HashBytes, LibDescr>::try_from_sorted_slice(&libraries).unwrap();
873
874        assert!(libraries.find(&HashBytes::ZERO).unwrap().is_none());
875        assert!(libraries.find_ref(&HashBytes::ZERO).unwrap().is_none());
876
877        assert_eq!(
878            libraries.find(lib1.repr_hash()).unwrap().unwrap().as_ref(),
879            libraries.find_ref(lib1.repr_hash()).unwrap().unwrap()
880        );
881        assert_eq!(
882            libraries.find(lib2.repr_hash()).unwrap().unwrap().as_ref(),
883            libraries.find_ref(lib2.repr_hash()).unwrap().unwrap()
884        );
885    }
886}