Skip to main content

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) if lib.root.repr_hash() == library_hash => Ok(Some(lib.root)),
190            _ => 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) if lib.root.repr_hash() == library_hash => Ok(Some(lib.root)),
200            _ => 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) if lib.root.repr_hash() == library_hash => return Ok(Some(lib.root)),
210                _ => 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) if lib.root.repr_hash() == library_hash => return Ok(Some(lib.root)),
223                _ => 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)?.and_then(|descr| {
233            if descr.lib.repr_hash() == library_hash {
234                Some(descr.lib.clone())
235            } else {
236                None
237            }
238        }))
239    }
240
241    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
242        struct LibDescrRef<'tlb> {
243            lib: &'tlb DynCell,
244        }
245
246        impl<'a> Load<'a> for LibDescrRef<'a> {
247            fn load_from(slice: &mut CellSlice<'a>) -> Result<Self, Error> {
248                if slice.load_small_uint(2)? != 0 {
249                    return Err(Error::InvalidTag);
250                }
251                Ok(Self {
252                    lib: slice.load_reference()?,
253                })
254            }
255        }
256
257        impl EquivalentRepr<LibDescr> for LibDescrRef<'_> {}
258
259        Ok(self
260            .cast_ref::<HashBytes, LibDescrRef<'a>>()
261            .get(library_hash)?
262            .and_then(|descr| {
263                if descr.lib.repr_hash() == library_hash {
264                    Some(descr.lib)
265                } else {
266                    None
267                }
268            }))
269    }
270}
271
272impl<S: BuildHasher> LibraryProvider for std::collections::HashMap<HashBytes, SimpleLib, S> {
273    fn find(&self, library_hash: &HashBytes) -> Result<Option<Cell>, Error> {
274        Ok(self.get(library_hash).and_then(|lib| {
275            if lib.root.repr_hash() == library_hash {
276                Some(lib.root.clone())
277            } else {
278                None
279            }
280        }))
281    }
282
283    fn find_ref<'a>(&'a self, library_hash: &HashBytes) -> Result<Option<&'a DynCell>, Error> {
284        Ok(self.get(library_hash).and_then(|lib| {
285            if lib.root.repr_hash() == library_hash {
286                Some(lib.root.as_ref())
287            } else {
288                None
289            }
290        }))
291    }
292}
293
294struct SimpleLibRef<'tlb> {
295    root: &'tlb DynCell,
296}
297
298impl<'a> Load<'a> for SimpleLibRef<'a> {
299    fn load_from(slice: &mut CellSlice<'a>) -> Result<Self, Error> {
300        slice.load_bit()?;
301        Ok(Self {
302            root: slice.load_reference()?,
303        })
304    }
305}
306
307impl EquivalentRepr<SimpleLib> for SimpleLibRef<'_> {}
308
309/// Gas tracking context.
310pub struct GasConsumer<'l> {
311    /// Maximum possible value of the `limit`.
312    gas_max: u64,
313    /// Gas limit for the out-of-gas exception.
314    gas_limit: std::cell::Cell<u64>,
315    /// Free gas (e.g. for external messages without any balance).
316    gas_credit: std::cell::Cell<u64>,
317    /// Initial gas to compute the consumed amount.
318    gas_base: std::cell::Cell<u64>,
319    /// Remaining gas available.
320    gas_remaining: std::cell::Cell<i64>,
321    /// Gas price (fixed point with 16 bits for fractional part).
322    gas_price: NonZeroU64,
323
324    /// A set of visited cells.
325    loaded_cells: std::cell::UnsafeCell<HashSet<HashBytes>>,
326    /// Libraries provider.
327    libraries: &'l dyn LibraryProvider,
328
329    /// Number of signature checks.
330    chksign_counter: std::cell::Cell<usize>,
331    /// Free gas (can be paid when using isolated consumer).
332    free_gas_consumed: std::cell::Cell<u64>,
333    /// Number of balance calls with cheap gas consumer
334    get_extra_balance_counter: std::cell::Cell<usize>,
335
336    // Missing library in case of resolving error occured.
337    missing_library: std::cell::Cell<Option<HashBytes>>,
338}
339
340impl<'l> GasConsumer<'l> {
341    pub const BUILD_CELL_GAS: u64 = 500;
342    pub const NEW_CELL_GAS: u64 = 100;
343    pub const OLD_CELL_GAS: u64 = 25;
344
345    pub const FREE_STACK_DEPTH: usize = 32;
346    pub const FREE_SIGNATURE_CHECKS: usize = 10;
347    pub const CHEAP_GET_BALANCE_CALLS: usize = 5;
348    pub const CHEAP_GET_BALANCE_GAS_THRESHOLD: u64 = 200;
349    pub const FREE_NESTED_CONT_JUMP: usize = 8;
350
351    pub const STACK_VALUE_GAS_PRICE: u64 = 1;
352    pub const TUPLE_ENTRY_GAS_PRICE: u64 = 1;
353    pub const HASH_EXT_ENTRY_GAS_PRICE: u64 = 1;
354    pub const CHK_SGN_GAS_PRICE: u64 = 4000;
355    pub const IMPLICIT_JMPREF_GAS_PRICE: u64 = 10;
356    pub const IMPLICIT_RET_GAS_PRICE: u64 = 5;
357    pub const EXCEPTION_GAS_PRICE: u64 = 50;
358    pub const RUNVM_GAS_PRICE: u64 = 40;
359
360    pub fn new(params: GasParams) -> Self {
361        static NO_LIBRARIES: NoLibraries = NoLibraries;
362
363        Self::with_libraries(params, &NO_LIBRARIES)
364    }
365
366    pub fn with_libraries(params: GasParams, libraries: &'l dyn LibraryProvider) -> Self {
367        let gas_remaining = truncate_gas(params.limit.saturating_add(params.credit));
368
369        Self {
370            gas_max: truncate_gas(params.max),
371            gas_limit: std::cell::Cell::new(truncate_gas(params.limit)),
372            gas_credit: std::cell::Cell::new(truncate_gas(params.credit)),
373            gas_base: std::cell::Cell::new(gas_remaining),
374            gas_remaining: std::cell::Cell::new(gas_remaining as i64),
375            gas_price: NonZeroU64::new(params.price).unwrap_or(NonZeroU64::MIN),
376            loaded_cells: Default::default(),
377            libraries,
378            chksign_counter: std::cell::Cell::new(0),
379            free_gas_consumed: std::cell::Cell::new(0),
380            get_extra_balance_counter: std::cell::Cell::new(0),
381            missing_library: std::cell::Cell::new(None),
382        }
383    }
384
385    pub fn limited(&'l self, remaining: u64) -> LimitedGasConsumer<'l> {
386        LimitedGasConsumer::<'l> {
387            gas: self,
388            remaining: std::cell::Cell::new(remaining),
389        }
390    }
391
392    // TODO: Consume free gas as non-free when `isolate`.
393    pub fn derive(&mut self, params: GasConsumerDeriveParams) -> VmResult<ParentGasConsumer<'l>> {
394        use std::cell::Cell;
395
396        Ok(if params.isolate {
397            // Pay for the free gas to prevent abuse.
398            self.try_consume(self.free_gas_consumed.get())?;
399
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            // Reset current free gas counters.
406            self.chksign_counter.set(0);
407            self.free_gas_consumed.set(0);
408            self.get_extra_balance_counter.set(0);
409
410            // Create a new gas consumer.
411            let params = GasParams {
412                max: std::cmp::min(params.gas_max, gas_remaining as u64),
413                limit: std::cmp::min(params.gas_limit, gas_remaining as u64),
414                credit: 0,
415                price: self.price(),
416            };
417            let libraries = self.libraries;
418
419            ParentGasConsumer::Isolated(std::mem::replace(
420                self,
421                Self::with_libraries(params, libraries),
422            ))
423        } else {
424            // NOTE: Compute remaining gas only when all operations
425            //       with parent consumer are made.
426            let gas_remaining = self.remaining();
427            vm_ensure!(gas_remaining >= 0, OutOfGas);
428
429            // Create child gas consumer params.
430            let gas_max = std::cmp::min(params.gas_max, gas_remaining as u64);
431            let gas_limit = std::cmp::min(params.gas_limit, gas_remaining as u64);
432
433            // Move gas params from child to parent.
434            ParentGasConsumer::Shared(GasConsumer {
435                gas_max: std::mem::replace(&mut self.gas_max, gas_max),
436                gas_limit: std::mem::replace(&mut self.gas_limit, Cell::new(gas_limit)),
437                gas_credit: std::mem::replace(&mut self.gas_credit, Cell::new(0)),
438                gas_base: std::mem::replace(&mut self.gas_base, Cell::new(gas_limit)),
439                gas_remaining: std::mem::replace(
440                    &mut self.gas_remaining,
441                    Cell::new(gas_limit as i64),
442                ),
443                gas_price: self.gas_price,
444                loaded_cells: Default::default(),
445                libraries: self.libraries,
446                chksign_counter: self.chksign_counter.clone(),
447                free_gas_consumed: self.free_gas_consumed.clone(),
448                get_extra_balance_counter: self.get_extra_balance_counter.clone(),
449                missing_library: self.missing_library.clone(),
450            })
451        })
452    }
453
454    pub fn restore(&mut self, mut parent: ParentGasConsumer<'l>) -> RestoredGasConsumer {
455        let meta = RestoredGasConsumer {
456            gas_consumed: self.consumed(),
457            gas_limit: self.limit(),
458        };
459
460        if let ParentGasConsumer::Shared(parent) = &mut parent {
461            // Merge loaded cells.
462            parent.loaded_cells = std::mem::take(&mut self.loaded_cells);
463        }
464
465        match parent {
466            ParentGasConsumer::Isolated(mut parent) | ParentGasConsumer::Shared(mut parent) => {
467                // Merge missing library.
468                let missing_lib = self.missing_library.get_mut();
469                let parent_lib = parent.missing_library.get_mut();
470                if parent_lib.is_none() && missing_lib.is_some() {
471                    *parent_lib = *missing_lib;
472                }
473
474                // Merge free gas counters.
475                parent.chksign_counter = self.chksign_counter.clone();
476                parent.free_gas_consumed = self.free_gas_consumed.clone();
477                parent.get_extra_balance_counter = self.get_extra_balance_counter.clone();
478
479                *self = parent
480            }
481        }
482
483        meta
484    }
485
486    pub fn libraries(&self) -> &'l dyn LibraryProvider {
487        self.libraries
488    }
489
490    pub fn credit(&self) -> u64 {
491        self.gas_credit.get()
492    }
493
494    pub fn consumed(&self) -> u64 {
495        (self.gas_base.get() as i64).saturating_sub(self.gas_remaining.get()) as u64
496    }
497
498    pub fn free_gas_consumed(&self) -> u64 {
499        self.free_gas_consumed.get()
500    }
501
502    pub fn remaining(&self) -> i64 {
503        self.gas_remaining.get()
504    }
505
506    pub fn base(&self) -> u64 {
507        self.gas_base.get()
508    }
509
510    pub fn limit(&self) -> u64 {
511        self.gas_limit.get()
512    }
513
514    pub fn set_limit(&self, limit: u64) {
515        let limit = std::cmp::min(limit, self.gas_max);
516        vm_log_trace!("changing gas limit: new_limit={limit}");
517
518        self.gas_credit.set(0);
519        self.gas_limit.set(limit);
520        self.set_base(limit);
521    }
522
523    fn set_base(&self, mut base: u64) {
524        base = truncate_gas(base);
525        let diff = base as i64 - self.gas_base.get() as i64;
526        self.gas_remaining.set(self.gas_remaining.get() + diff);
527        self.gas_base.set(base);
528    }
529
530    pub fn price(&self) -> u64 {
531        self.gas_price.get()
532    }
533
534    pub fn try_get_extra_balance_consumer(&'l self) -> Option<LimitedGasConsumer<'l>> {
535        self.get_extra_balance_counter
536            .set(self.get_extra_balance_counter.get() + 1);
537        if self.get_extra_balance_counter.get() <= Self::CHEAP_GET_BALANCE_CALLS {
538            return Some(self.limited(Self::CHEAP_GET_BALANCE_GAS_THRESHOLD));
539        }
540
541        None
542    }
543
544    pub fn try_consume_exception_gas(&self) -> Result<(), Error> {
545        self.try_consume(Self::EXCEPTION_GAS_PRICE)
546    }
547
548    pub fn try_consume_implicit_jmpref_gas(&self) -> Result<(), Error> {
549        self.try_consume(Self::IMPLICIT_JMPREF_GAS_PRICE)
550    }
551
552    pub fn try_consume_implicit_ret_gas(&self) -> Result<(), Error> {
553        self.try_consume(Self::IMPLICIT_RET_GAS_PRICE)
554    }
555
556    pub fn try_consume_check_signature_gas(&self) -> Result<(), Error> {
557        self.chksign_counter.set(self.chksign_counter.get() + 1);
558        if self.chksign_counter.get() > Self::FREE_SIGNATURE_CHECKS {
559            self.try_consume(Self::CHK_SGN_GAS_PRICE)?;
560        } else {
561            self.consume_free_gas(Self::CHK_SGN_GAS_PRICE);
562        }
563        Ok(())
564    }
565
566    pub fn try_consume_stack_gas(&self, stack: Option<&SafeRc<Stack>>) -> Result<(), Error> {
567        if let Some(stack) = stack {
568            self.try_consume_stack_depth_gas(stack.depth())?;
569        }
570        Ok(())
571    }
572
573    pub fn try_consume_tuple_gas(&self, tuple_len: u64) -> Result<(), Error> {
574        self.try_consume(tuple_len * Self::TUPLE_ENTRY_GAS_PRICE)?;
575        Ok(())
576    }
577
578    pub fn try_consume_stack_depth_gas(&self, depth: usize) -> Result<(), Error> {
579        self.try_consume(
580            (std::cmp::max(depth, Self::FREE_STACK_DEPTH) - Self::FREE_STACK_DEPTH) as u64
581                * Self::STACK_VALUE_GAS_PRICE,
582        )
583    }
584
585    pub fn try_consume(&self, amount: u64) -> Result<(), Error> {
586        let remaining = self
587            .gas_remaining
588            .get()
589            .saturating_sub(truncate_gas(amount) as i64);
590        self.gas_remaining.set(remaining);
591
592        if remaining >= 0 {
593            Ok(())
594        } else {
595            Err(Error::Cancelled)
596        }
597    }
598
599    pub fn consume_free_gas(&self, amount: u64) {
600        let consumed = truncate_gas(self.free_gas_consumed.get().saturating_add(amount));
601        self.free_gas_consumed.set(consumed);
602    }
603
604    pub fn missing_library(&self) -> Option<HashBytes> {
605        self.missing_library.get()
606    }
607
608    pub fn set_missing_library(&self, hash: &HashBytes) {
609        self.missing_library.set(Some(*hash));
610    }
611
612    pub fn load_cell_as_slice(&self, cell: Cell, mode: LoadMode) -> Result<OwnedCellSlice, Error> {
613        let cell = ok!(self.load_cell_impl(cell, mode));
614        Ok(OwnedCellSlice::new_allow_exotic(cell))
615    }
616
617    fn load_cell_impl<'s: 'a, 'a, T: LoadLibrary<'a>>(
618        &'s self,
619        mut cell: T,
620        mode: LoadMode,
621    ) -> Result<T, Error> {
622        let mut library_loaded = false;
623        loop {
624            if mode.use_gas() {
625                // SAFETY: This is the only place where we borrow `loaded_cells` as mut.
626                let is_new =
627                    unsafe { (*self.loaded_cells.get()).insert(*cell.as_ref().repr_hash()) };
628
629                ok!(self.try_consume(if is_new {
630                    GasConsumer::NEW_CELL_GAS
631                } else {
632                    GasConsumer::OLD_CELL_GAS
633                }));
634            }
635
636            if !mode.resolve() {
637                return Ok(cell);
638            }
639
640            match cell.as_ref().cell_type() {
641                CellType::Ordinary => return Ok(cell),
642                CellType::LibraryReference if !library_loaded => {
643                    // Library data structure is enforced by `CellContext::finalize_cell`.
644                    debug_assert_eq!(cell.as_ref().bit_len(), 8 + 256);
645
646                    // Find library by hash.
647                    let mut library_hash = HashBytes::ZERO;
648                    ok!(cell
649                        .as_ref()
650                        .as_slice_allow_exotic()
651                        .get_raw(8, &mut library_hash.0, 256));
652
653                    let Some(library_cell) = ok!(T::load_library(self, &library_hash)) else {
654                        self.missing_library.set(Some(library_hash));
655                        return Err(Error::CellUnderflow);
656                    };
657
658                    cell = library_cell;
659                    library_loaded = true;
660                }
661                _ => return Err(Error::CellUnderflow),
662            }
663        }
664    }
665}
666
667impl CellContext for GasConsumer<'_> {
668    fn finalize_cell(&self, cell: CellParts<'_>) -> Result<Cell, Error> {
669        ok!(self.try_consume(GasConsumer::BUILD_CELL_GAS));
670        Cell::empty_context().finalize_cell(cell)
671    }
672
673    fn load_cell(&self, cell: Cell, mode: LoadMode) -> Result<Cell, Error> {
674        self.load_cell_impl(cell, mode)
675    }
676
677    fn load_dyn_cell<'s: 'a, 'a>(
678        &'s self,
679        cell: &'a DynCell,
680        mode: LoadMode,
681    ) -> Result<&'a DynCell, Error> {
682        self.load_cell_impl(cell, mode)
683    }
684}
685
686/// Consumes at most N gas units, everything else is treated as free gas.
687pub struct LimitedGasConsumer<'l> {
688    /// Main gas consumer
689    gas: &'l GasConsumer<'l>,
690    /// Remaining gas that can be consumed.
691    remaining: std::cell::Cell<u64>,
692}
693
694impl CellContext for LimitedGasConsumer<'_> {
695    fn finalize_cell(&self, _: CellParts<'_>) -> Result<Cell, Error> {
696        panic!("limited gas consumer must not be used for making new cells")
697    }
698
699    fn load_cell(&self, cell: Cell, mode: LoadMode) -> Result<Cell, Error> {
700        self.load_cell_impl(cell, mode)
701    }
702
703    fn load_dyn_cell<'s: 'a, 'a>(
704        &'s self,
705        cell: &'a DynCell,
706        mode: LoadMode,
707    ) -> Result<&'a DynCell, Error> {
708        self.load_cell_impl(cell, mode)
709    }
710}
711
712impl LimitedGasConsumer<'_> {
713    fn load_cell_impl<'s: 'a, 'a, T: LoadLibrary<'a>>(
714        &'s self,
715        mut cell: T,
716        mode: LoadMode,
717    ) -> Result<T, Error> {
718        let mut library_loaded = false;
719        loop {
720            if mode.use_gas() {
721                // SAFETY: This is the only place where we borrow `loaded_cells` as mut.
722                let is_new =
723                    unsafe { (*self.gas.loaded_cells.get()).insert(*cell.as_ref().repr_hash()) };
724
725                ok!(self.try_consume(if is_new {
726                    GasConsumer::NEW_CELL_GAS
727                } else {
728                    GasConsumer::OLD_CELL_GAS
729                }));
730            }
731
732            if !mode.resolve() {
733                return Ok(cell);
734            }
735
736            match cell.as_ref().cell_type() {
737                CellType::Ordinary => return Ok(cell),
738                CellType::LibraryReference if !library_loaded => {
739                    // Library data structure is enforced by `CellContext::finalize_cell`.
740                    debug_assert_eq!(cell.as_ref().bit_len(), 8 + 256);
741
742                    // Find library by hash.
743                    let mut library_hash = HashBytes::ZERO;
744                    ok!(cell
745                        .as_ref()
746                        .as_slice_allow_exotic()
747                        .get_raw(8, &mut library_hash.0, 256));
748
749                    let Some(library_cell) = ok!(T::load_library(self.gas, &library_hash)) else {
750                        self.gas.missing_library.set(Some(library_hash));
751                        return Err(Error::CellUnderflow);
752                    };
753
754                    cell = library_cell;
755                    library_loaded = true;
756                }
757                _ => return Err(Error::CellUnderflow),
758            }
759        }
760    }
761
762    pub fn try_consume(&self, mut amount: u64) -> Result<(), Error> {
763        amount = truncate_gas(amount);
764
765        let mut remaining = self.remaining.get();
766        let consumed = std::cmp::min(amount, remaining);
767
768        self.gas.try_consume(consumed)?;
769
770        remaining -= consumed;
771        self.remaining.set(remaining);
772
773        // threshold reached
774        if remaining == 0 {
775            self.gas.consume_free_gas(amount - consumed);
776        }
777        Ok(())
778    }
779}
780
781trait LoadLibrary<'a>: AsRef<DynCell> + 'a {
782    fn load_library(gas: &'a GasConsumer, library_hash: &HashBytes) -> Result<Option<Self>, Error>
783    where
784        Self: Sized;
785}
786
787impl<'a> LoadLibrary<'a> for &'a DynCell {
788    fn load_library(gas: &'a GasConsumer, library_hash: &HashBytes) -> Result<Option<Self>, Error>
789    where
790        Self: Sized,
791    {
792        gas.libraries.find_ref(library_hash)
793    }
794}
795
796impl LoadLibrary<'_> for Cell {
797    fn load_library(gas: &'_ GasConsumer, library_hash: &HashBytes) -> Result<Option<Self>, Error>
798    where
799        Self: Sized,
800    {
801        gas.libraries.find(library_hash)
802    }
803}
804
805/// Params to replace the current gas consumer.
806#[derive(Debug, Clone, Copy)]
807pub struct GasConsumerDeriveParams {
808    pub gas_max: u64,
809    pub gas_limit: u64,
810    pub isolate: bool,
811}
812
813/// Parent of the derived gas consumer.
814pub enum ParentGasConsumer<'l> {
815    Isolated(GasConsumer<'l>),
816    Shared(GasConsumer<'l>),
817}
818
819/// Info extracted when parent gas consumer is restored.
820#[derive(Debug, Clone, Copy)]
821pub struct RestoredGasConsumer {
822    pub gas_consumed: u64,
823    pub gas_limit: u64,
824}
825
826const fn truncate_gas(gas: u64) -> u64 {
827    if gas <= i64::MAX as u64 {
828        gas
829    } else {
830        i64::MAX as u64
831    }
832}
833
834#[cfg(test)]
835mod tests {
836    use super::*;
837
838    #[test]
839    fn find_lib_dict_ref() {
840        let lib1 = Boc::decode(tvmasm!("NOP")).unwrap();
841        let lib2 = Boc::decode(tvmasm!("NOP NOP")).unwrap();
842
843        // Dict with SimpleLib
844        let mut libraries = vec![
845            (*lib1.repr_hash(), SimpleLib {
846                public: true,
847                root: lib1.clone(),
848            }),
849            (*lib2.repr_hash(), SimpleLib {
850                public: true,
851                root: lib2.clone(),
852            }),
853        ];
854        libraries.sort_unstable_by(|(l, _), (r, _)| l.cmp(r));
855        let libraries = Dict::<HashBytes, SimpleLib>::try_from_sorted_slice(&libraries).unwrap();
856
857        assert!(libraries.find(&HashBytes::ZERO).unwrap().is_none());
858        assert!(libraries.find_ref(&HashBytes::ZERO).unwrap().is_none());
859
860        assert_eq!(
861            libraries.find(lib1.repr_hash()).unwrap().unwrap().as_ref(),
862            libraries.find_ref(lib1.repr_hash()).unwrap().unwrap()
863        );
864        assert_eq!(
865            libraries.find(lib2.repr_hash()).unwrap().unwrap().as_ref(),
866            libraries.find_ref(lib2.repr_hash()).unwrap().unwrap()
867        );
868
869        // Dict with LibDescr
870        let mut publishers = Dict::new();
871        publishers.add(HashBytes::ZERO, ()).unwrap();
872
873        {
874            let lib = LibDescr {
875                lib: lib1.clone(),
876                publishers: publishers.clone(),
877            };
878            let c = CellBuilder::build_from(&lib).unwrap();
879            let parsed = c.parse::<LibDescr>().unwrap();
880
881            assert_eq!(lib, parsed);
882        }
883
884        let mut libraries = vec![
885            (*lib1.repr_hash(), LibDescr {
886                lib: lib1.clone(),
887                publishers: publishers.clone(),
888            }),
889            (*lib2.repr_hash(), LibDescr {
890                lib: lib2.clone(),
891                publishers,
892            }),
893        ];
894        libraries.sort_unstable_by(|(l, _), (r, _)| l.cmp(r));
895
896        let libraries = Dict::<HashBytes, LibDescr>::try_from_sorted_slice(&libraries).unwrap();
897
898        assert!(libraries.find(&HashBytes::ZERO).unwrap().is_none());
899        assert!(libraries.find_ref(&HashBytes::ZERO).unwrap().is_none());
900
901        assert_eq!(
902            libraries.find(lib1.repr_hash()).unwrap().unwrap().as_ref(),
903            libraries.find_ref(lib1.repr_hash()).unwrap().unwrap()
904        );
905        assert_eq!(
906            libraries.find(lib2.repr_hash()).unwrap().unwrap().as_ref(),
907            libraries.find_ref(lib2.repr_hash()).unwrap().unwrap()
908        );
909    }
910}