tasm_lib/memory/
dyn_malloc.rs

1use std::collections::HashMap;
2
3use num_traits::Zero;
4use rand::prelude::*;
5use triton_vm::memory_layout::MemoryRegion;
6use triton_vm::prelude::*;
7
8use crate::empty_stack;
9use crate::prelude::*;
10use crate::snippet_bencher::BenchmarkCase;
11use crate::traits::function::Function;
12use crate::traits::function::FunctionInitialState;
13use crate::traits::rust_shadow::RustShadowError;
14
15/// The location of the dynamic allocator state in memory.
16///
17/// See the [memory convention][super] for details.
18pub const DYN_MALLOC_ADDRESS: BFieldElement = BFieldElement::new(BFieldElement::MAX);
19
20/// The index of the first page that can be dynamically allocated.
21pub const DYN_MALLOC_FIRST_PAGE: u64 = 1;
22
23/// The number of pages that can be dynamically allocated.
24pub const NUM_ALLOCATABLE_PAGES: u64 = (1 << 31) - 1;
25
26/// The size of one dynamically allocated page.
27pub const DYN_MALLOC_PAGE_SIZE: u64 = 1 << 32;
28
29pub const DYN_MALLOC_FIRST_ADDRESS: BFieldElement =
30    BFieldElement::new(DYN_MALLOC_FIRST_PAGE * DYN_MALLOC_PAGE_SIZE);
31
32/// Return a pointer to the next free page of memory. Updates the dyn malloc state
33/// accordingly
34#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
35pub struct DynMalloc;
36
37impl DynMalloc {
38    pub fn memory_region() -> MemoryRegion {
39        MemoryRegion::new(
40            DYN_MALLOC_FIRST_ADDRESS,
41            NUM_ALLOCATABLE_PAGES * DYN_MALLOC_PAGE_SIZE,
42        )
43    }
44}
45
46impl BasicSnippet for DynMalloc {
47    fn parameters(&self) -> Vec<(DataType, String)> {
48        vec![]
49    }
50
51    fn return_values(&self) -> Vec<(DataType, String)> {
52        vec![(DataType::Bfe, "*addr".to_string())]
53    }
54
55    fn entrypoint(&self) -> String {
56        "tasmlib_memory_dyn_malloc".to_string()
57    }
58
59    fn code(&self, _: &mut Library) -> Vec<LabelledInstruction> {
60        let entrypoint = self.entrypoint();
61        let dyn_malloc_init = format!("{entrypoint}_initialize");
62
63        triton_asm! {
64        // BEFORE: _
65        // AFTER:  _ *addr
66        {entrypoint}:
67            push {DYN_MALLOC_ADDRESS}       // _ *dyn_malloc_state
68            read_mem 1 pop 1                // _ page_idx
69
70            dup 0 push 0 eq                 // _ page_idx (page_idx == 0)
71            skiz call {dyn_malloc_init}     // _ page_idx
72
73            // Verify that we are mapping inside allowed region
74            push {NUM_ALLOCATABLE_PAGES}
75            dup 1
76            lt                              // _ page_idx (page_idx < NUM_ALLOCATABLE_PAGES)
77            assert error_id 70
78
79            // update dynamic allocator state
80            dup 0                           // _ page_idx page_idx
81            push 1                          // _ page_idx page_idx 1
82            add                             // _ page_idx next_page_idx
83            push {DYN_MALLOC_ADDRESS}       // _ page_idx next_page_idx *dyn_malloc_state
84            write_mem 1 pop 1               // _ page_idx
85
86            // translate page number to address
87            push {DYN_MALLOC_PAGE_SIZE}     // _ page_idx page_size
88            mul                             // _ *free_page
89            return
90
91        // BEFORE: _ 0
92        // AFTER:  _ DYN_MALLOC_FIRST_PAGE
93        {dyn_malloc_init}:
94            pop 1
95            push {DYN_MALLOC_FIRST_PAGE}
96            return
97        }
98    }
99}
100
101impl Function for DynMalloc {
102    fn rust_shadow(
103        &self,
104        stack: &mut Vec<BFieldElement>,
105        memory: &mut HashMap<BFieldElement, BFieldElement>,
106    ) -> Result<(), RustShadowError> {
107        let mut page_idx = memory.get(&DYN_MALLOC_ADDRESS).copied().unwrap_or_default();
108        if page_idx.is_zero() {
109            page_idx = DYN_MALLOC_FIRST_PAGE.into();
110        }
111        let page_idx = page_idx;
112
113        if page_idx.value() >= NUM_ALLOCATABLE_PAGES {
114            return Err(RustShadowError::Other);
115        }
116
117        let next_page_idx = page_idx + bfe!(1);
118
119        memory.insert(DYN_MALLOC_ADDRESS, next_page_idx);
120
121        let page_address = page_idx * BFieldElement::new(DYN_MALLOC_PAGE_SIZE);
122        stack.push(page_address);
123        Ok(())
124    }
125
126    fn pseudorandom_initial_state(
127        &self,
128        seed: [u8; 32],
129        _: Option<BenchmarkCase>,
130    ) -> FunctionInitialState {
131        let mut rng = StdRng::from_seed(seed);
132
133        let stack = empty_stack();
134
135        let mut memory = HashMap::new();
136        let page_number = rng.random_range(0..1_u64 << 31);
137        memory.insert(DYN_MALLOC_ADDRESS, page_number.into());
138
139        FunctionInitialState { stack, memory }
140    }
141
142    fn corner_case_initial_states(&self) -> Vec<FunctionInitialState> {
143        let empty_vm_state = {
144            let stack = self.init_stack_for_isolated_run();
145            let memory = HashMap::new();
146
147            FunctionInitialState { stack, memory }
148        };
149
150        let one_page_has_been_allocated = {
151            let stack = self.init_stack_for_isolated_run();
152            let memory: HashMap<_, _> = [(DYN_MALLOC_ADDRESS, bfe!(1))].into_iter().collect();
153
154            FunctionInitialState { stack, memory }
155        };
156
157        let second_to_last_page_has_been_allocated = {
158            let stack = self.init_stack_for_isolated_run();
159            let memory: HashMap<_, _> = [(DYN_MALLOC_ADDRESS, bfe!(NUM_ALLOCATABLE_PAGES - 1))]
160                .into_iter()
161                .collect();
162
163            FunctionInitialState { stack, memory }
164        };
165
166        let third_to_last_page_has_been_allocated = {
167            let stack = self.init_stack_for_isolated_run();
168            let memory: HashMap<_, _> = [(DYN_MALLOC_ADDRESS, bfe!(NUM_ALLOCATABLE_PAGES - 2))]
169                .into_iter()
170                .collect();
171
172            FunctionInitialState { stack, memory }
173        };
174
175        vec![
176            empty_vm_state,
177            one_page_has_been_allocated,
178            second_to_last_page_has_been_allocated,
179            third_to_last_page_has_been_allocated,
180        ]
181    }
182}
183
184#[cfg(test)]
185pub mod tests {
186    use std::collections::HashMap;
187
188    use super::*;
189    use crate::test_helpers::negative_test;
190    use crate::test_prelude::*;
191
192    #[macro_rules_attr::apply(test)]
193    fn expected_address_chosen_for_dyn_malloc() {
194        assert_eq!(bfe!(-1), DYN_MALLOC_ADDRESS);
195    }
196
197    #[macro_rules_attr::apply(test)]
198    fn rust_shadow() {
199        ShadowedFunction::new(DynMalloc).test();
200    }
201
202    #[macro_rules_attr::apply(test)]
203    fn disallow_allocation_outside_of_region_unit_test() {
204        fn negative_prop_disallow_allocation_outside_of_region(memory_page_index: BFieldElement) {
205            let snippet = DynMalloc;
206            let init_stack = snippet.init_stack_for_isolated_run();
207            let shadowed_snippet = ShadowedFunction::new(snippet);
208
209            let memory_filled_dyn_malloc: HashMap<_, _> = [(DYN_MALLOC_ADDRESS, memory_page_index)]
210                .into_iter()
211                .collect();
212            let init_state =
213                InitVmState::with_stack_and_memory(init_stack, memory_filled_dyn_malloc);
214
215            test_assertion_failure(&shadowed_snippet, init_state, &[70]);
216        }
217
218        // Last page has been allocated, next call must fail
219        negative_prop_disallow_allocation_outside_of_region(bfe!(NUM_ALLOCATABLE_PAGES));
220        negative_prop_disallow_allocation_outside_of_region(bfe!(NUM_ALLOCATABLE_PAGES + 1));
221        negative_prop_disallow_allocation_outside_of_region(bfe!(NUM_ALLOCATABLE_PAGES + 2));
222        negative_prop_disallow_allocation_outside_of_region(bfe!(u32::MAX - 1));
223        negative_prop_disallow_allocation_outside_of_region(bfe!(u32::MAX));
224    }
225
226    #[macro_rules_attr::apply(proptest)]
227    fn disallow_allocation_if_page_counter_is_not_a_u32(
228        #[strategy(arb())]
229        #[filter(#address.value() > u64::from(u32::MAX))]
230        address: BFieldElement,
231    ) {
232        fn negative_prop_disallow_allocation_with_non_u32_page_counter(
233            mem_page_index: BFieldElement,
234        ) {
235            let snippet = DynMalloc;
236            let init_stack = snippet.init_stack_for_isolated_run();
237            let shadowed_snippet = ShadowedFunction::new(snippet);
238
239            let memory_filled_dyn_malloc: HashMap<_, _> =
240                [(DYN_MALLOC_ADDRESS, mem_page_index)].into_iter().collect();
241            let init_state =
242                InitVmState::with_stack_and_memory(init_stack, memory_filled_dyn_malloc);
243            let expected_err =
244                InstructionError::OpStackError(OpStackError::FailedU32Conversion(mem_page_index));
245            negative_test(&shadowed_snippet, init_state, &[expected_err]);
246        }
247
248        negative_prop_disallow_allocation_with_non_u32_page_counter(address);
249    }
250
251    #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
252    struct MultipleDynMallocCalls {
253        num_calls: usize,
254    }
255
256    impl BasicSnippet for MultipleDynMallocCalls {
257        fn parameters(&self) -> Vec<(DataType, String)> {
258            vec![]
259        }
260
261        fn return_values(&self) -> Vec<(DataType, String)> {
262            vec![(DataType::Bfe, "*addr".to_string()); self.num_calls]
263        }
264
265        fn entrypoint(&self) -> String {
266            "tasmlib_memory_dyn_malloc_multiple_calls".to_string()
267        }
268
269        fn code(&self, lib: &mut Library) -> Vec<LabelledInstruction> {
270            let dyn_malloc = lib.import(Box::new(DynMalloc));
271
272            let single_call = triton_asm!(call { dyn_malloc });
273            let multiple_calls = vec![single_call; self.num_calls].concat();
274
275            triton_asm!( {self.entrypoint()}: {&multiple_calls} return )
276        }
277    }
278
279    impl Function for MultipleDynMallocCalls {
280        fn rust_shadow(
281            &self,
282            stack: &mut Vec<BFieldElement>,
283            memory: &mut HashMap<BFieldElement, BFieldElement>,
284        ) -> Result<(), RustShadowError> {
285            for _ in 0..self.num_calls {
286                DynMalloc.rust_shadow(stack, memory)?;
287            }
288            Ok(())
289        }
290
291        fn pseudorandom_initial_state(
292            &self,
293            seed: [u8; 32],
294            _: Option<BenchmarkCase>,
295        ) -> FunctionInitialState {
296            DynMalloc.pseudorandom_initial_state(seed, None)
297        }
298    }
299
300    #[macro_rules_attr::apply(proptest(cases = 20))]
301    fn dynamic_allocator_can_be_called_multiple_times(
302        #[strategy(0_usize..1_000)] num_calls: usize,
303    ) {
304        let multiple_calls = MultipleDynMallocCalls { num_calls };
305        ShadowedFunction::new(multiple_calls).test();
306    }
307}
308
309#[cfg(test)]
310mod benches {
311    use super::*;
312    use crate::test_prelude::*;
313
314    #[macro_rules_attr::apply(test)]
315    fn benchmark() {
316        ShadowedFunction::new(DynMalloc).bench();
317    }
318}
tasm_lib/memory/dyn_malloc.rs

tasm_lib/memory/
dyn_malloc.rs
