Skip to main content

gear_core/
memory.rs

1// Copyright (C) Gear Technologies Inc.
2// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
3
4//! Module for memory and allocations context.
5
6use crate::{
7    gas::ChargeError,
8    pages::{GearPage, WasmPage, WasmPagesAmount},
9};
10use alloc::{boxed::Box, format};
11use byteorder::{ByteOrder, LittleEndian};
12use core::{
13    fmt::{self, Debug},
14    ops::{Deref, DerefMut},
15};
16use numerated::{
17    interval::{Interval, TryFromRangeError},
18    tree::IntervalsTree,
19};
20use scale_decode::DecodeAsType;
21use scale_encode::EncodeAsType;
22use scale_info::{
23    TypeInfo,
24    scale::{Decode, Encode},
25};
26
27/// Interval in wasm program memory.
28#[derive(Clone, Copy, Eq, PartialEq, Encode, EncodeAsType, Decode, DecodeAsType)]
29pub struct MemoryInterval {
30    /// Interval offset in bytes.
31    pub offset: u32,
32    /// Interval size in bytes.
33    pub size: u32,
34}
35
36impl MemoryInterval {
37    /// Convert `MemoryInterval` to `[u8; 8]` bytes.
38    /// `0..4` - `offset`
39    /// `4..8` - `size`
40    #[inline]
41    pub fn to_bytes(&self) -> [u8; 8] {
42        let mut bytes = [0u8; 8];
43        LittleEndian::write_u32(&mut bytes[0..4], self.offset);
44        LittleEndian::write_u32(&mut bytes[4..8], self.size);
45        bytes
46    }
47
48    /// Convert `[u8; 8]` bytes to `MemoryInterval`.
49    /// `0..4` - `offset`
50    /// `4..8` - `size`
51    #[inline]
52    pub fn try_from_bytes(bytes: &[u8]) -> Result<Self, &'static str> {
53        if bytes.len() != 8 {
54            return Err("bytes size != 8");
55        }
56        let offset = LittleEndian::read_u32(&bytes[0..4]);
57        let size = LittleEndian::read_u32(&bytes[4..8]);
58        Ok(MemoryInterval { offset, size })
59    }
60}
61
62impl From<(u32, u32)> for MemoryInterval {
63    fn from(val: (u32, u32)) -> Self {
64        MemoryInterval {
65            offset: val.0,
66            size: val.1,
67        }
68    }
69}
70
71impl From<MemoryInterval> for (u32, u32) {
72    fn from(val: MemoryInterval) -> Self {
73        (val.offset, val.size)
74    }
75}
76
77impl Debug for MemoryInterval {
78    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
79        f.debug_struct("MemoryInterval")
80            .field("offset", &format_args!("{:#x}", self.offset))
81            .field("size", &format_args!("{:#x}", self.size))
82            .finish()
83    }
84}
85
86/// Error in attempt to make wrong size page buffer.
87#[derive(Debug, Default, PartialEq, Eq, PartialOrd, Ord, Clone, TypeInfo, derive_more::Display)]
88#[display("Trying to make wrong size page buffer, must be {:#x}", GearPage::SIZE)]
89pub struct IntoPageBufError;
90
91// Inner page buffer is stored as `[u64; _]` instead of `[u8; _]`
92// to mitigate `polkadot-js` limitation of array length,
93// which requires it not to be more than 2048.
94const _: () = {
95    assert!((GearPage::SIZE as usize).is_multiple_of(size_of::<u64>()));
96    assert!(GearPage::SIZE as usize / size_of::<u64>() <= 2048);
97};
98
99/// Alias for inner type of page buffer.
100pub type PageBufInner = Box<[u64; GearPage::SIZE as usize / size_of::<u64>()]>;
101
102/// Buffer for gear page data.
103#[derive(
104    Clone,
105    PartialEq,
106    Eq,
107    PartialOrd,
108    Ord,
109    Hash,
110    Encode,
111    EncodeAsType,
112    Decode,
113    DecodeAsType,
114    TypeInfo,
115)]
116pub struct PageBuf(PageBufInner);
117
118impl Debug for PageBuf {
119    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
120        write!(
121            f,
122            "PageBuf({:?}..{:?})",
123            &self.as_ref()[0..10],
124            &self.as_ref()[GearPage::SIZE as usize - 10..GearPage::SIZE as usize]
125        )
126    }
127}
128
129impl Deref for PageBuf {
130    type Target = [u8];
131    fn deref(&self) -> &Self::Target {
132        bytemuck::must_cast_slice(&*self.0)
133    }
134}
135
136impl DerefMut for PageBuf {
137    fn deref_mut(&mut self) -> &mut Self::Target {
138        bytemuck::must_cast_slice_mut(&mut *self.0)
139    }
140}
141
142impl AsRef<[u8; GearPage::SIZE as usize]> for PageBuf {
143    fn as_ref(&self) -> &[u8; GearPage::SIZE as usize] {
144        bytemuck::must_cast_ref(&*self.0)
145    }
146}
147
148impl AsMut<[u8; GearPage::SIZE as usize]> for PageBuf {
149    fn as_mut(&mut self) -> &mut [u8; GearPage::SIZE as usize] {
150        bytemuck::must_cast_mut(&mut *self.0)
151    }
152}
153
154impl PageBuf {
155    /// Returns new page buffer with zeroed data.
156    pub fn new_zeroed() -> PageBuf {
157        Self::filled_with(0)
158    }
159
160    /// Returns new page buffer filled with given byte.
161    pub fn filled_with(byte: u8) -> PageBuf {
162        let chunk = u64::from_ne_bytes([byte; 8]);
163        Self([chunk; GearPage::SIZE as usize / size_of::<u64>()].into())
164    }
165
166    /// Creates PageBuf from inner buffer. If the buffer has
167    /// the size of [`GearPage`] then no reallocations occur.
168    /// In other case it will be extended with zeros.
169    ///
170    /// The method is implemented intentionally instead of trait From to
171    /// highlight conversion cases in the source code.
172    pub fn from_inner(inner: PageBufInner) -> Self {
173        Self(inner)
174    }
175}
176
177/// Host pointer type.
178/// Host pointer can be 64bit or less, to support both we use u64.
179pub type HostPointer = u64;
180
181const _: () = assert!(size_of::<HostPointer>() >= size_of::<usize>());
182
183/// Core memory error.
184#[derive(Debug, Copy, Clone, Eq, PartialEq, derive_more::Display)]
185pub enum MemoryError {
186    /// The error occurs in attempt to access memory outside wasm program memory.
187    #[display("Trying to access memory outside wasm program memory")]
188    AccessOutOfBounds,
189}
190
191/// Backend wasm memory interface.
192pub trait Memory<Context> {
193    /// Memory grow error.
194    type GrowError: Debug;
195
196    /// Grow memory by number of pages.
197    fn grow(&self, ctx: &mut Context, pages: WasmPagesAmount) -> Result<(), Self::GrowError>;
198
199    /// Return current size of the memory.
200    fn size(&self, ctx: &Context) -> WasmPagesAmount;
201
202    /// Set memory region at specific pointer.
203    fn write(&self, ctx: &mut Context, offset: u32, buffer: &[u8]) -> Result<(), MemoryError>;
204
205    /// Reads memory contents at the given offset into a buffer.
206    fn read(&self, ctx: &Context, offset: u32, buffer: &mut [u8]) -> Result<(), MemoryError>;
207
208    /// Returns native addr of wasm memory buffer in wasm executor
209    fn get_buffer_host_addr(&self, ctx: &Context) -> Option<HostPointer> {
210        if self.size(ctx) == WasmPagesAmount::from(0) {
211            None
212        } else {
213            // We call this method only in case memory size is not zero,
214            // so memory buffer exists and has addr in host memory.
215            unsafe { Some(self.get_buffer_host_addr_unsafe(ctx)) }
216        }
217    }
218
219    /// Get buffer addr unsafe.
220    ///
221    /// # Safety
222    /// If memory size is 0 then buffer addr can be garbage
223    unsafe fn get_buffer_host_addr_unsafe(&self, ctx: &Context) -> HostPointer;
224}
225
226/// Pages allocations context for the running program.
227#[derive(Debug)]
228pub struct AllocationsContext {
229    /// Pages which has been in storage before execution
230    allocations: IntervalsTree<WasmPage>,
231    /// Shows that `allocations` was modified at least once per execution
232    allocations_changed: bool,
233    heap: Option<Interval<WasmPage>>,
234    static_pages: WasmPagesAmount,
235}
236
237/// Before and after memory grow actions.
238#[must_use]
239pub trait GrowHandler<Context> {
240    /// Before grow action
241    fn before_grow_action(ctx: &mut Context, mem: &mut impl Memory<Context>) -> Self;
242    /// After grow action
243    fn after_grow_action(self, ctx: &mut Context, mem: &mut impl Memory<Context>);
244}
245
246/// Grow handler do nothing implementation
247pub struct NoopGrowHandler;
248
249impl<Context> GrowHandler<Context> for NoopGrowHandler {
250    fn before_grow_action(_ctx: &mut Context, _mem: &mut impl Memory<Context>) -> Self {
251        NoopGrowHandler
252    }
253    fn after_grow_action(self, _ctx: &mut Context, _mem: &mut impl Memory<Context>) {}
254}
255
256/// Inconsistency in memory parameters provided for wasm execution.
257#[derive(Debug, Clone, PartialEq, Eq, derive_more::Display)]
258pub enum MemorySetupError {
259    /// Memory size exceeds max pages
260    #[display("Memory size {memory_size:?} must be less than or equal to {max_pages:?}")]
261    MemorySizeExceedsMaxPages {
262        /// Memory size
263        memory_size: WasmPagesAmount,
264        /// Max allowed memory size
265        max_pages: WasmPagesAmount,
266    },
267    /// Insufficient memory size
268    #[display("Memory size {memory_size:?} must be at least {static_pages:?}")]
269    InsufficientMemorySize {
270        /// Memory size
271        memory_size: WasmPagesAmount,
272        /// Static memory size
273        static_pages: WasmPagesAmount,
274    },
275    /// Stack end is out of static memory
276    #[display("Stack end {stack_end:?} is out of static memory 0..{static_pages:?}")]
277    StackEndOutOfStaticMemory {
278        /// Stack end
279        stack_end: WasmPage,
280        /// Static memory size
281        static_pages: WasmPagesAmount,
282    },
283    /// Allocated page is out of allowed memory interval
284    #[display(
285        "Allocated page {page:?} is out of allowed memory interval {static_pages:?}..{memory_size:?}"
286    )]
287    AllocatedPageOutOfAllowedInterval {
288        /// Allocated page
289        page: WasmPage,
290        /// Static memory size
291        static_pages: WasmPagesAmount,
292        /// Memory size
293        memory_size: WasmPagesAmount,
294    },
295}
296
297/// Allocation error
298#[derive(Debug, Clone, Eq, PartialEq, derive_more::Display, derive_more::From)]
299pub enum AllocError {
300    /// The error occurs when a program tries to allocate more memory than
301    /// allowed.
302    #[display("Trying to allocate more wasm program memory than allowed")]
303    ProgramAllocOutOfBounds,
304    /// The error occurs in attempt to free-up a memory page from static area or
305    /// outside additionally allocated for this program.
306    #[display("{_0:?} cannot be freed by the current program")]
307    InvalidFree(WasmPage),
308    /// Invalid range for free_range
309    #[display("Invalid range {_0:?}..={_1:?} for free_range")]
310    InvalidFreeRange(WasmPage, WasmPage),
311    /// Gas charge error
312    GasCharge(ChargeError),
313}
314
315impl AllocationsContext {
316    /// New allocations context.
317    ///
318    /// Provide currently running `program_id`, boxed memory abstraction
319    /// and allocation manager. Also configurable `static_pages` and `max_pages`
320    /// are set.
321    ///
322    /// Returns `MemorySetupError` on incorrect memory params.
323    pub fn try_new(
324        memory_size: WasmPagesAmount,
325        allocations: IntervalsTree<WasmPage>,
326        static_pages: WasmPagesAmount,
327        stack_end: Option<WasmPage>,
328        max_pages: WasmPagesAmount,
329    ) -> Result<Self, MemorySetupError> {
330        Self::validate_memory_params(
331            memory_size,
332            &allocations,
333            static_pages,
334            stack_end,
335            max_pages,
336        )?;
337
338        let heap = match Interval::try_from(static_pages..max_pages) {
339            Ok(interval) => Some(interval),
340            Err(TryFromRangeError::EmptyRange) => None,
341            // Branch is unreachable due to the check `static_pages <= max_pages`` in `validate_memory_params`.
342            _ => unreachable!(),
343        };
344
345        Ok(Self {
346            allocations,
347            allocations_changed: false,
348            heap,
349            static_pages,
350        })
351    }
352
353    /// Checks memory parameters, that are provided for wasm execution.
354    /// NOTE: this params partially checked in `Code::try_new` in `gear-core`.
355    fn validate_memory_params(
356        memory_size: WasmPagesAmount,
357        allocations: &IntervalsTree<WasmPage>,
358        static_pages: WasmPagesAmount,
359        stack_end: Option<WasmPage>,
360        max_pages: WasmPagesAmount,
361    ) -> Result<(), MemorySetupError> {
362        if memory_size > max_pages {
363            return Err(MemorySetupError::MemorySizeExceedsMaxPages {
364                memory_size,
365                max_pages,
366            });
367        }
368
369        if static_pages > memory_size {
370            return Err(MemorySetupError::InsufficientMemorySize {
371                memory_size,
372                static_pages,
373            });
374        }
375
376        if let Some(stack_end) = stack_end
377            && stack_end > static_pages
378        {
379            return Err(MemorySetupError::StackEndOutOfStaticMemory {
380                stack_end,
381                static_pages,
382            });
383        }
384
385        if let Some(page) = allocations.end()
386            && page >= memory_size
387        {
388            return Err(MemorySetupError::AllocatedPageOutOfAllowedInterval {
389                page,
390                static_pages,
391                memory_size,
392            });
393        }
394        if let Some(page) = allocations.start()
395            && page < static_pages
396        {
397            return Err(MemorySetupError::AllocatedPageOutOfAllowedInterval {
398                page,
399                static_pages,
400                memory_size,
401            });
402        }
403
404        Ok(())
405    }
406
407    /// Allocates specified number of continuously going pages
408    /// and returns zero-based number of the first one.
409    pub fn alloc<Context, G: GrowHandler<Context>>(
410        &mut self,
411        ctx: &mut Context,
412        mem: &mut impl Memory<Context>,
413        pages: WasmPagesAmount,
414        charge_gas_for_grow: impl FnOnce(WasmPagesAmount) -> Result<(), ChargeError>,
415    ) -> Result<WasmPage, AllocError> {
416        // Empty heap means that all memory is static, then no pages can be allocated.
417        // NOTE: returns an error even if `pages` == 0.
418        let heap = self.heap.ok_or(AllocError::ProgramAllocOutOfBounds)?;
419
420        // If trying to allocate zero pages, then returns heap start page (legacy).
421        if pages == WasmPage::from(0) {
422            return Ok(heap.start());
423        }
424
425        let interval = self
426            .allocations
427            .voids(heap)
428            .find_map(|void| {
429                Interval::<WasmPage>::with_len(void.start(), u32::from(pages))
430                    .ok()
431                    .and_then(|interval| (interval.end() <= void.end()).then_some(interval))
432            })
433            .ok_or(AllocError::ProgramAllocOutOfBounds)?;
434
435        if let Ok(grow) = Interval::<WasmPage>::try_from(mem.size(ctx)..interval.end().inc()) {
436            charge_gas_for_grow(grow.len())?;
437            let grow_handler = G::before_grow_action(ctx, mem);
438            mem.grow(ctx, grow.len()).unwrap_or_else(|err| {
439                let err_msg = format!(
440                    "AllocationContext:alloc: Failed to grow memory. \
441                        Got error - {err:?}",
442                );
443
444                log::error!("{err_msg}");
445                unreachable!("{err_msg}")
446            });
447            grow_handler.after_grow_action(ctx, mem);
448        }
449
450        self.allocations.insert(interval);
451        self.allocations_changed = true;
452
453        Ok(interval.start())
454    }
455
456    /// Free specific memory page.
457    pub fn free(&mut self, page: WasmPage) -> Result<(), AllocError> {
458        if let Some(heap) = self.heap
459            && page >= heap.start()
460            && page <= heap.end()
461            && self.allocations.remove(page)
462        {
463            self.allocations_changed = true;
464            return Ok(());
465        }
466
467        Err(AllocError::InvalidFree(page))
468    }
469
470    /// Try to free pages in range. Will only return error if range is invalid.
471    ///
472    /// Currently running program should own this pages.
473    pub fn free_range(&mut self, interval: Interval<WasmPage>) -> Result<(), AllocError> {
474        if let Some(heap) = self.heap {
475            // `free_range` allows do not modify the allocations so we do not check the `remove` result here
476            if interval.start() >= heap.start() && interval.end() <= heap.end() {
477                if self.allocations.remove(interval) {
478                    self.allocations_changed = true;
479                }
480
481                return Ok(());
482            }
483        }
484
485        Err(AllocError::InvalidFreeRange(
486            interval.start(),
487            interval.end(),
488        ))
489    }
490
491    /// Decomposes this instance and returns `static_pages`, `allocations` and `allocations_changed` params.
492    pub fn into_parts(self) -> (WasmPagesAmount, IntervalsTree<WasmPage>, bool) {
493        (
494            self.static_pages,
495            self.allocations,
496            self.allocations_changed,
497        )
498    }
499}
500
501/// This module contains tests of `GearPage` and `AllocationContext`
502#[cfg(test)]
503mod tests {
504    use super::*;
505    use alloc::{vec, vec::Vec};
506    use core::{cell::Cell, iter};
507
508    struct TestMemory(Cell<WasmPagesAmount>);
509
510    impl TestMemory {
511        fn new(amount: WasmPagesAmount) -> Self {
512            Self(Cell::new(amount))
513        }
514    }
515
516    impl Memory<()> for TestMemory {
517        type GrowError = ();
518
519        fn grow(&self, _ctx: &mut (), pages: WasmPagesAmount) -> Result<(), Self::GrowError> {
520            let new_pages_amount = self.0.get().add(pages).ok_or(())?;
521            self.0.set(new_pages_amount);
522            Ok(())
523        }
524
525        fn size(&self, _ctx: &()) -> WasmPagesAmount {
526            self.0.get()
527        }
528
529        fn write(&self, _ctx: &mut (), _offset: u32, _buffer: &[u8]) -> Result<(), MemoryError> {
530            unimplemented!()
531        }
532
533        fn read(&self, _ctx: &(), _offset: u32, _buffer: &mut [u8]) -> Result<(), MemoryError> {
534            unimplemented!()
535        }
536
537        unsafe fn get_buffer_host_addr_unsafe(&self, _ctx: &()) -> HostPointer {
538            unimplemented!()
539        }
540    }
541
542    #[test]
543    fn page_buf() {
544        let _ = tracing_subscriber::fmt::try_init();
545
546        let mut page_buf = PageBuf::filled_with(199);
547        page_buf[1] = 2;
548        log::debug!("page buff = {page_buf:?}");
549    }
550
551    #[test]
552    fn page_buf_encode() {
553        let page_buf = PageBuf::filled_with(199);
554
555        assert_eq!(page_buf.encode(), vec![199u8; GearPage::SIZE as usize])
556    }
557
558    #[test]
559    fn free_fails() {
560        let mut ctx =
561            AllocationsContext::try_new(0.into(), Default::default(), 0.into(), None, 0.into())
562                .unwrap();
563        assert_eq!(ctx.free(1.into()), Err(AllocError::InvalidFree(1.into())));
564
565        let mut ctx = AllocationsContext::try_new(
566            1.into(),
567            [WasmPage::from(0)].into_iter().collect(),
568            0.into(),
569            None,
570            1.into(),
571        )
572        .unwrap();
573        assert_eq!(ctx.free(1.into()), Err(AllocError::InvalidFree(1.into())));
574
575        let mut ctx = AllocationsContext::try_new(
576            4.into(),
577            [WasmPage::from(1), WasmPage::from(3)].into_iter().collect(),
578            1.into(),
579            None,
580            4.into(),
581        )
582        .unwrap();
583        let interval = Interval::<WasmPage>::try_from(1u16..4).unwrap();
584        assert_eq!(ctx.free_range(interval), Ok(()));
585    }
586
587    #[track_caller]
588    fn alloc_ok(ctx: &mut AllocationsContext, mem: &mut TestMemory, pages: u16, expected: u16) {
589        let res = ctx.alloc::<(), NoopGrowHandler>(&mut (), mem, pages.into(), |_| Ok(()));
590        assert_eq!(res, Ok(expected.into()));
591    }
592
593    #[track_caller]
594    fn alloc_err(ctx: &mut AllocationsContext, mem: &mut TestMemory, pages: u16, err: AllocError) {
595        let res = ctx.alloc::<(), NoopGrowHandler>(&mut (), mem, pages.into(), |_| Ok(()));
596        assert_eq!(res, Err(err));
597    }
598
599    #[test]
600    fn alloc() {
601        let _ = tracing_subscriber::fmt::try_init();
602
603        let mut ctx = AllocationsContext::try_new(
604            256.into(),
605            Default::default(),
606            16.into(),
607            None,
608            256.into(),
609        )
610        .unwrap();
611        let mut mem = TestMemory::new(16.into());
612        alloc_ok(&mut ctx, &mut mem, 16, 16);
613        alloc_ok(&mut ctx, &mut mem, 0, 16);
614
615        // there is a space for 14 more
616        (2..16).for_each(|i| alloc_ok(&mut ctx, &mut mem, 16, i * 16));
617
618        // no more mem!
619        alloc_err(&mut ctx, &mut mem, 16, AllocError::ProgramAllocOutOfBounds);
620
621        // but we free some and then can allocate page that was freed
622        ctx.free(137.into()).unwrap();
623        alloc_ok(&mut ctx, &mut mem, 1, 137);
624
625        // if we free 2 in a row we can allocate even 2
626        ctx.free(117.into()).unwrap();
627        ctx.free(118.into()).unwrap();
628        alloc_ok(&mut ctx, &mut mem, 2, 117);
629
630        // same as above, if we free_range 2 in a row we can allocate 2
631        let interval = Interval::<WasmPage>::try_from(117..119).unwrap();
632        ctx.free_range(interval).unwrap();
633        alloc_ok(&mut ctx, &mut mem, 2, 117);
634
635        // but if 2 are not in a row, bad luck
636        ctx.free(117.into()).unwrap();
637        ctx.free(158.into()).unwrap();
638        alloc_err(&mut ctx, &mut mem, 2, AllocError::ProgramAllocOutOfBounds);
639    }
640
641    #[test]
642    fn memory_params_validation() {
643        assert_eq!(
644            AllocationsContext::validate_memory_params(
645                4.into(),
646                &iter::once(WasmPage::from(2)).collect(),
647                2.into(),
648                Some(2.into()),
649                4.into(),
650            ),
651            Ok(())
652        );
653
654        assert_eq!(
655            AllocationsContext::validate_memory_params(
656                4.into(),
657                &Default::default(),
658                2.into(),
659                Some(2.into()),
660                3.into(),
661            ),
662            Err(MemorySetupError::MemorySizeExceedsMaxPages {
663                memory_size: 4.into(),
664                max_pages: 3.into()
665            })
666        );
667
668        assert_eq!(
669            AllocationsContext::validate_memory_params(
670                1.into(),
671                &Default::default(),
672                2.into(),
673                Some(1.into()),
674                4.into(),
675            ),
676            Err(MemorySetupError::InsufficientMemorySize {
677                memory_size: 1.into(),
678                static_pages: 2.into()
679            })
680        );
681
682        assert_eq!(
683            AllocationsContext::validate_memory_params(
684                4.into(),
685                &Default::default(),
686                2.into(),
687                Some(3.into()),
688                4.into(),
689            ),
690            Err(MemorySetupError::StackEndOutOfStaticMemory {
691                stack_end: 3.into(),
692                static_pages: 2.into()
693            })
694        );
695
696        assert_eq!(
697            AllocationsContext::validate_memory_params(
698                4.into(),
699                &[WasmPage::from(1), WasmPage::from(3)].into_iter().collect(),
700                2.into(),
701                Some(2.into()),
702                4.into(),
703            ),
704            Err(MemorySetupError::AllocatedPageOutOfAllowedInterval {
705                page: 1.into(),
706                static_pages: 2.into(),
707                memory_size: 4.into()
708            })
709        );
710
711        assert_eq!(
712            AllocationsContext::validate_memory_params(
713                4.into(),
714                &[WasmPage::from(2), WasmPage::from(4)].into_iter().collect(),
715                2.into(),
716                Some(2.into()),
717                4.into(),
718            ),
719            Err(MemorySetupError::AllocatedPageOutOfAllowedInterval {
720                page: 4.into(),
721                static_pages: 2.into(),
722                memory_size: 4.into()
723            })
724        );
725
726        assert_eq!(
727            AllocationsContext::validate_memory_params(
728                13.into(),
729                &iter::once(WasmPage::from(1)).collect(),
730                10.into(),
731                None,
732                13.into()
733            ),
734            Err(MemorySetupError::AllocatedPageOutOfAllowedInterval {
735                page: 1.into(),
736                static_pages: 10.into(),
737                memory_size: 13.into()
738            })
739        );
740
741        assert_eq!(
742            AllocationsContext::validate_memory_params(
743                13.into(),
744                &iter::once(WasmPage::from(1)).collect(),
745                WasmPagesAmount::UPPER,
746                None,
747                13.into()
748            ),
749            Err(MemorySetupError::InsufficientMemorySize {
750                memory_size: 13.into(),
751                static_pages: WasmPagesAmount::UPPER
752            })
753        );
754
755        assert_eq!(
756            AllocationsContext::validate_memory_params(
757                WasmPagesAmount::UPPER,
758                &iter::once(WasmPage::from(1)).collect(),
759                10.into(),
760                None,
761                WasmPagesAmount::UPPER,
762            ),
763            Err(MemorySetupError::AllocatedPageOutOfAllowedInterval {
764                page: 1.into(),
765                static_pages: 10.into(),
766                memory_size: WasmPagesAmount::UPPER
767            })
768        );
769    }
770
771    #[test]
772    fn allocations_changed_correctness() {
773        let new_ctx = |allocations| {
774            AllocationsContext::try_new(16.into(), allocations, 0.into(), None, 16.into()).unwrap()
775        };
776
777        // correct `alloc`
778        let mut ctx = new_ctx(Default::default());
779        assert!(
780            !ctx.allocations_changed,
781            "Expecting no changes after creation"
782        );
783        let mut mem = TestMemory::new(16.into());
784        alloc_ok(&mut ctx, &mut mem, 16, 0);
785        assert!(ctx.allocations_changed);
786
787        let (_, allocations, allocations_changed) = ctx.into_parts();
788        assert!(allocations_changed);
789
790        // fail `alloc`
791        let mut ctx = new_ctx(allocations);
792        alloc_err(&mut ctx, &mut mem, 16, AllocError::ProgramAllocOutOfBounds);
793        assert!(
794            !ctx.allocations_changed,
795            "Expecting allocations don't change because of error"
796        );
797
798        // fail `free`
799        assert!(ctx.free(16.into()).is_err());
800        assert!(!ctx.allocations_changed);
801
802        // correct `free`
803        assert!(ctx.free(10.into()).is_ok());
804        assert!(ctx.allocations_changed);
805
806        let (_, allocations, allocations_changed) = ctx.into_parts();
807        assert!(allocations_changed);
808
809        // correct `free_range`
810        // allocations: [0..9] ∪ [11..15]
811        let mut ctx = new_ctx(allocations);
812        let interval = Interval::<WasmPage>::try_from(10u16..12).unwrap();
813        assert!(ctx.free_range(interval).is_ok());
814        assert!(
815            ctx.allocations_changed,
816            "Expected value is `true` because the 11th page was freed from allocations."
817        );
818
819        let (_, allocations, allocations_changed) = ctx.into_parts();
820        assert!(allocations_changed);
821
822        // fail `free_range`
823        let mut ctx = new_ctx(allocations);
824        let interval = Interval::<WasmPage>::try_from(0u16..17).unwrap();
825        assert!(ctx.free_range(interval).is_err());
826        assert!(!ctx.allocations_changed);
827        assert!(!ctx.into_parts().2);
828    }
829
830    mod property_tests {
831        use super::*;
832        use proptest::{
833            arbitrary::any,
834            collection::size_range,
835            prop_oneof, proptest,
836            strategy::{Just, Strategy},
837            test_runner::Config as ProptestConfig,
838        };
839
840        #[derive(Debug, Clone)]
841        enum Action {
842            Alloc { pages: WasmPagesAmount },
843            Free { page: WasmPage },
844            FreeRange { page: WasmPage, size: u8 },
845        }
846
847        fn actions() -> impl Strategy<Value = Vec<Action>> {
848            let action = prop_oneof![
849                // Allocate smaller number (0..32) of pages due to `BTree::extend` significantly slows down prop-test.
850                wasm_pages_amount_with_range(0, 32).prop_map(|pages| Action::Alloc { pages }),
851                wasm_page().prop_map(|page| Action::Free { page }),
852                (wasm_page(), any::<u8>())
853                    .prop_map(|(page, size)| Action::FreeRange { page, size }),
854            ];
855            proptest::collection::vec(action, 0..1024)
856        }
857
858        fn allocations(start: u16, end: u16) -> impl Strategy<Value = IntervalsTree<WasmPage>> {
859            proptest::collection::btree_set(wasm_page_with_range(start, end), size_range(0..1024))
860                .prop_map(|pages| pages.into_iter().collect::<IntervalsTree<WasmPage>>())
861        }
862
863        fn wasm_page_with_range(start: u16, end: u16) -> impl Strategy<Value = WasmPage> {
864            (start..=end).prop_map(WasmPage::from)
865        }
866
867        fn wasm_page() -> impl Strategy<Value = WasmPage> {
868            wasm_page_with_range(0, u16::MAX)
869        }
870
871        fn wasm_pages_amount_with_range(
872            start: u32,
873            end: u32,
874        ) -> impl Strategy<Value = WasmPagesAmount> {
875            (start..=end).prop_map(|x| {
876                if x == u16::MAX as u32 + 1 {
877                    WasmPagesAmount::UPPER
878                } else {
879                    WasmPagesAmount::from(x as u16)
880                }
881            })
882        }
883
884        fn wasm_pages_amount() -> impl Strategy<Value = WasmPagesAmount> {
885            wasm_pages_amount_with_range(0, u16::MAX as u32 + 1)
886        }
887
888        #[derive(Debug)]
889        struct MemoryParams {
890            max_pages: WasmPagesAmount,
891            mem_size: WasmPagesAmount,
892            static_pages: WasmPagesAmount,
893            allocations: IntervalsTree<WasmPage>,
894        }
895
896        // This high-order strategy generates valid memory parameters in a specific way that allows passing `AllocationContext::validate_memory_params` checks.
897        fn combined_memory_params() -> impl Strategy<Value = MemoryParams> {
898            wasm_pages_amount()
899                .prop_flat_map(|max_pages| {
900                    let mem_size = wasm_pages_amount_with_range(0, u32::from(max_pages));
901                    (Just(max_pages), mem_size)
902                })
903                .prop_flat_map(|(max_pages, mem_size)| {
904                    let static_pages = wasm_pages_amount_with_range(0, u32::from(mem_size));
905                    (Just(max_pages), Just(mem_size), static_pages)
906                })
907                .prop_filter(
908                    "filter out cases where allocation region has zero size",
909                    |(_max_pages, mem_size, static_pages)| static_pages < mem_size,
910                )
911                .prop_flat_map(|(max_pages, mem_size, static_pages)| {
912                    // Last allocated page should be < `mem_size`.
913                    let end_exclusive = u32::from(mem_size) - 1;
914                    (
915                        Just(max_pages),
916                        Just(mem_size),
917                        Just(static_pages),
918                        allocations(u32::from(static_pages) as u16, end_exclusive as u16),
919                    )
920                })
921                .prop_map(
922                    |(max_pages, mem_size, static_pages, allocations)| MemoryParams {
923                        max_pages,
924                        mem_size,
925                        static_pages,
926                        allocations,
927                    },
928                )
929        }
930
931        fn proptest_config() -> ProptestConfig {
932            ProptestConfig {
933                cases: 1024,
934                ..Default::default()
935            }
936        }
937
938        #[track_caller]
939        fn assert_free_error(err: AllocError) {
940            match err {
941                AllocError::InvalidFree(_) => {}
942                AllocError::InvalidFreeRange(_, _) => {}
943                err => panic!("{err:?}"),
944            }
945        }
946
947        proptest! {
948            #![proptest_config(proptest_config())]
949            #[test]
950            fn alloc(
951                mem_params in combined_memory_params(),
952                actions in actions(),
953            ) {
954                let _ = tracing_subscriber::fmt::try_init();
955
956                let MemoryParams{max_pages, mem_size, static_pages, allocations} = mem_params;
957                let mut ctx = AllocationsContext::try_new(mem_size, allocations, static_pages, None, max_pages).unwrap();
958
959                let mut mem = TestMemory::new(mem_size);
960
961                for action in actions {
962                    match action {
963                        Action::Alloc { pages } => {
964                            match ctx.alloc::<_, NoopGrowHandler>(&mut (), &mut mem, pages, |_| Ok(())) {
965                                Err(AllocError::ProgramAllocOutOfBounds) => {
966                                    let x = mem.size(&()).add(pages);
967                                    assert!(x.is_none() || x.unwrap() > max_pages);
968                                }
969                                Ok(page) => {
970                                    assert!(pages == WasmPagesAmount::from(0) || (page >= static_pages && page < max_pages));
971                                    assert!(mem.size(&()) <= max_pages);
972                                    assert!(WasmPagesAmount::from(page).add(pages).unwrap() <= mem.size(&()));
973                                }
974                                Err(err) => panic!("{err:?}"),
975                            }
976                        }
977                        Action::Free { page } => {
978                            if let Err(err) = ctx.free(page) {
979                                assert_free_error(err);
980                            }
981                        }
982                        Action::FreeRange { page, size } => {
983                            if let Ok(interval) = Interval::<WasmPage>::with_len(page, size as u32) {
984                                let _ = ctx.free_range(interval).map_err(assert_free_error);
985                            }
986                        }
987                    }
988                }
989            }
990        }
991    }
992}