|
1 | 1 | use std::collections::HashMap; |
2 | 2 |
|
| 3 | +use arbitrary::Arbitrary; |
3 | 4 | use itertools::Itertools; |
| 5 | +use num_traits::ConstOne; |
4 | 6 | use triton_vm::memory_layout::MemoryRegion; |
5 | 7 | use triton_vm::prelude::*; |
6 | 8 |
|
@@ -30,6 +32,34 @@ pub struct Library { |
30 | 32 | num_allocated_words: u32, |
31 | 33 | } |
32 | 34 |
|
| 35 | +/// Represents a [static memory allocation][kmalloc] within Triton VM. |
| 36 | +/// Both its location within Triton VM's memory and its size and are fix. |
| 37 | +/// |
| 38 | +/// [kmalloc]: Library::kmalloc |
| 39 | +#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Arbitrary)] |
| 40 | +pub struct StaticAllocation { |
| 41 | + write_address: BFieldElement, |
| 42 | + num_words: u32, |
| 43 | +} |
| 44 | + |
| 45 | +impl StaticAllocation { |
| 46 | + /// The address from which the allocated memory can be read. |
| 47 | + pub fn read_address(&self) -> BFieldElement { |
| 48 | + let offset = bfe!(self.num_words) - BFieldElement::ONE; |
| 49 | + self.write_address() + offset |
| 50 | + } |
| 51 | + |
| 52 | + /// The address to which the allocated memory can be written. |
| 53 | + pub fn write_address(&self) -> BFieldElement { |
| 54 | + self.write_address |
| 55 | + } |
| 56 | + |
| 57 | + /// The number of words allocated in this memory block. |
| 58 | + pub fn num_words(&self) -> u32 { |
| 59 | + self.num_words |
| 60 | + } |
| 61 | +} |
| 62 | + |
33 | 63 | impl Default for Library { |
34 | 64 | fn default() -> Self { |
35 | 65 | Self::new() |
@@ -119,19 +149,26 @@ impl Library { |
119 | 149 | self.all_external_dependencies().concat() |
120 | 150 | } |
121 | 151 |
|
122 | | - /// Statically allocate `num_words` words of memory. Panics if more static |
123 | | - /// memory is required than what the capacity allows for. |
124 | | - pub fn kmalloc(&mut self, num_words: u32) -> BFieldElement { |
| 152 | + /// Statically allocate `num_words` words of memory. |
| 153 | + /// |
| 154 | + /// # Panics |
| 155 | + /// |
| 156 | + /// Panics if |
| 157 | + /// - `num_words` is zero, |
| 158 | + /// - the total number of statically allocated words exceeds `u32::MAX`. |
| 159 | + pub fn kmalloc(&mut self, num_words: u32) -> StaticAllocation { |
125 | 160 | assert!(num_words > 0, "must allocate a positive number of words"); |
126 | | - let address = STATIC_MEMORY_FIRST_ADDRESS |
127 | | - - bfe!(self.num_allocated_words) |
128 | | - - BFieldElement::new(num_words as u64 - 1); |
| 161 | + let write_address = |
| 162 | + STATIC_MEMORY_FIRST_ADDRESS - bfe!(self.num_allocated_words) - bfe!(num_words - 1); |
129 | 163 | self.num_allocated_words = self |
130 | 164 | .num_allocated_words |
131 | 165 | .checked_add(num_words) |
132 | 166 | .expect("Cannot allocate more that u32::MAX words through `kmalloc`."); |
133 | 167 |
|
134 | | - address |
| 168 | + StaticAllocation { |
| 169 | + write_address, |
| 170 | + num_words, |
| 171 | + } |
135 | 172 | } |
136 | 173 | } |
137 | 174 |
|
@@ -466,13 +503,13 @@ mod tests { |
466 | 503 | const MINUS_TWO: BFieldElement = BFieldElement::new(BFieldElement::MAX - 1); |
467 | 504 | let mut lib = Library::new(); |
468 | 505 |
|
469 | | - let first_free_address = lib.kmalloc(1); |
470 | | - assert_eq!(MINUS_TWO, first_free_address); |
| 506 | + let first_chunk = lib.kmalloc(1); |
| 507 | + assert_eq!(MINUS_TWO, first_chunk.write_address()); |
471 | 508 |
|
472 | | - let second_free_address = lib.kmalloc(7); |
473 | | - assert_eq!(-BFieldElement::new(9), second_free_address,); |
| 509 | + let second_chunk = lib.kmalloc(7); |
| 510 | + assert_eq!(-bfe!(9), second_chunk.write_address()); |
474 | 511 |
|
475 | | - let third_free_address = lib.kmalloc(1000); |
476 | | - assert_eq!(-BFieldElement::new(1009), third_free_address); |
| 512 | + let third_chunk = lib.kmalloc(1000); |
| 513 | + assert_eq!(-bfe!(1009), third_chunk.write_address()); |
477 | 514 | } |
478 | 515 | } |
0 commit comments