miden_tx/host/
link_map.rs

1use alloc::vec::Vec;
2use core::cmp::Ordering;
3
4use miden_processor::fast::ExecutionOutput;
5use miden_processor::{AdviceMutation, ContextId, ProcessState};
6use miden_protocol::{Felt, LexicographicWord, Word, ZERO};
7
8// LINK MAP
9// ================================================================================================
10
11/// A map based on a sorted linked list.
12///
13/// This type enables access to the list in kernel memory.
14///
15/// See link_map.masm for docs.
16///
17/// # Warning
18///
19/// The functions on this type assume that the provided map_ptr points to a valid map in the
20/// provided memory viewer. If those assumptions are violated, the functions may panic.
21#[derive(Clone, Copy)]
22pub struct LinkMap<'process> {
23    map_ptr: u32,
24    mem: &'process MemoryViewer<'process>,
25}
26
27impl<'process> LinkMap<'process> {
28    // CONSTRUCTOR
29    // --------------------------------------------------------------------------------------------
30
31    /// Creates a new link map from the provided map_ptr in the provided process.
32    pub fn new(map_ptr: Felt, mem: &'process MemoryViewer<'process>) -> Self {
33        let map_ptr: u32 = map_ptr.try_into().expect("map_ptr must be a valid u32");
34
35        Self { map_ptr, mem }
36    }
37
38    // PUBLIC METHODS
39    // --------------------------------------------------------------------------------------------
40
41    /// Handles a `LINK_MAP_SET_EVENT` emitted from a VM.
42    ///
43    /// Expected operand stack state before: [map_ptr, KEY, NEW_VALUE]
44    /// Advice stack state after: [set_operation, entry_ptr]
45    pub fn handle_set_event(process: &ProcessState<'_>) -> Vec<AdviceMutation> {
46        let map_ptr = process.get_stack_item(1);
47        let map_key = process.get_stack_word_be(2);
48
49        let mem_viewer = MemoryViewer::ProcessState(process);
50        let link_map = LinkMap::new(map_ptr, &mem_viewer);
51
52        let (set_op, entry_ptr) = link_map.compute_set_operation(LexicographicWord::from(map_key));
53
54        vec![AdviceMutation::extend_stack([Felt::from(set_op as u8), Felt::from(entry_ptr)])]
55    }
56
57    /// Handles a `LINK_MAP_GET_EVENT` emitted from a VM.
58    ///
59    /// Expected operand stack state before: [map_ptr, KEY]
60    /// Advice stack state after: [get_operation, entry_ptr]
61    pub fn handle_get_event(process: &ProcessState<'_>) -> Vec<AdviceMutation> {
62        let map_ptr = process.get_stack_item(1);
63        let map_key = process.get_stack_word_be(2);
64
65        let mem_viewer = MemoryViewer::ProcessState(process);
66        let link_map = LinkMap::new(map_ptr, &mem_viewer);
67        let (get_op, entry_ptr) = link_map.compute_get_operation(LexicographicWord::from(map_key));
68
69        vec![AdviceMutation::extend_stack([Felt::from(get_op as u8), Felt::from(entry_ptr)])]
70    }
71
72    /// Returns `true` if the map is empty, `false` otherwise.
73    pub fn is_empty(&self) -> bool {
74        self.head().is_none()
75    }
76
77    /// Returns an iterator over the link map entries.
78    pub fn iter(&self) -> impl Iterator<Item = Entry> {
79        LinkMapIter {
80            current_entry_ptr: self.head().unwrap_or(0),
81            map: *self,
82        }
83    }
84
85    // PRIVATE METHODS
86    // --------------------------------------------------------------------------------------------
87
88    /// Returns the entry pointer at the head of the map or `None` if the map is empty.
89    fn head(&self) -> Option<u32> {
90        // Returns None if the value was either not yet initialized or points to 0.
91        // It can point to 0 for example if a get operation is executed before a set operation,
92        // which initializes the value in memory to 0 but does not change it.
93        self.mem.get_kernel_mem_element(self.map_ptr).and_then(|head_ptr| {
94            if head_ptr == ZERO {
95                None
96            } else {
97                Some(u32::try_from(head_ptr).expect("head ptr should be a valid ptr"))
98            }
99        })
100    }
101
102    /// Returns the [`Entry`] at the given pointer.
103    fn entry(&self, entry_ptr: u32) -> Entry {
104        let key = self.key(entry_ptr);
105        let (value0, value1) = self.value(entry_ptr);
106        let metadata = self.metadata(entry_ptr);
107
108        Entry {
109            ptr: entry_ptr,
110            metadata,
111            key,
112            value0,
113            value1,
114        }
115    }
116
117    /// Returns the key of the entry at the given pointer.
118    fn key(&self, entry_ptr: u32) -> LexicographicWord {
119        LexicographicWord::from(
120            self.mem
121                .get_kernel_mem_word(entry_ptr + 4)
122                .expect("entry pointer should be valid"),
123        )
124    }
125
126    /// Returns the values of the entry at the given pointer.
127    fn value(&self, entry_ptr: u32) -> (Word, Word) {
128        let value0 = self
129            .mem
130            .get_kernel_mem_word(entry_ptr + 8)
131            .expect("entry pointer should be valid");
132        let value1 = self
133            .mem
134            .get_kernel_mem_word(entry_ptr + 12)
135            .expect("entry pointer should be valid");
136        (value0, value1)
137    }
138
139    /// Returns the metadata of the entry at the given pointer.
140    fn metadata(&self, entry_ptr: u32) -> EntryMetadata {
141        let entry_metadata =
142            self.mem.get_kernel_mem_word(entry_ptr).expect("entry pointer should be valid");
143
144        let map_ptr = entry_metadata[0];
145        let map_ptr = map_ptr.try_into().expect("entry_ptr should point to a u32 map_ptr");
146
147        let prev_entry_ptr = entry_metadata[1];
148        let prev_entry_ptr = prev_entry_ptr
149            .try_into()
150            .expect("entry_ptr should point to a u32 prev_entry_ptr");
151
152        let next_entry_ptr = entry_metadata[2];
153        let next_entry_ptr = next_entry_ptr
154            .try_into()
155            .expect("entry_ptr should point to a u32 next_entry_ptr");
156
157        EntryMetadata { map_ptr, prev_entry_ptr, next_entry_ptr }
158    }
159
160    /// Computes what needs to be done to insert the given key into the link map.
161    ///
162    /// If the key already exists in the map, then its value must be updated and
163    /// [`SetOperation::Update`] and the pointer to the existing entry are returned.
164    ///
165    /// If the key does not exist in the map, find the place where it has to be inserted. This can
166    /// be at the head of the list ([`SetOperation::InsertAtHead`]) if the key is smaller than all
167    /// existing keys or if the map is empty. Otherwise it is after an existing entry
168    /// ([`SetOperation::InsertAfterEntry`]) in which case the key must be greater than the entry's
169    /// key after which it is inserted and smaller than the entry before which it is inserted
170    /// (unless it is the end of the map).
171    fn compute_set_operation(&self, key: LexicographicWord) -> (SetOperation, u32) {
172        let Some(current_head) = self.head() else {
173            return (SetOperation::InsertAtHead, 0);
174        };
175
176        let mut last_entry_ptr: u32 = current_head;
177
178        for entry in self.iter() {
179            match key.cmp(&entry.key) {
180                Ordering::Equal => {
181                    return (SetOperation::Update, entry.ptr);
182                },
183                Ordering::Less => {
184                    if entry.ptr == current_head {
185                        return (SetOperation::InsertAtHead, entry.ptr);
186                    }
187
188                    break;
189                },
190                Ordering::Greater => {
191                    last_entry_ptr = entry.ptr;
192                },
193            }
194        }
195
196        (SetOperation::InsertAfterEntry, last_entry_ptr)
197    }
198
199    /// Computes a get operation for a key in a link map.
200    ///
201    /// If the key exists, then [`GetOperation::Found`] is returned and the pointer to it.
202    ///
203    /// If it does not exist, its absence must be proven, otherwise the host could lie. To do that,
204    /// the in-kernel link map validates that the key is not in the list, so this function returns
205    /// information pointing to the entry where the key would be if it existed.
206    ///
207    /// The way to compute this is the same as a set operation, so this function simply remaps its
208    /// output.
209    fn compute_get_operation(&self, key: LexicographicWord) -> (GetOperation, u32) {
210        let (set_op, entry_ptr) = self.compute_set_operation(key);
211        let get_op = match set_op {
212            SetOperation::Update => GetOperation::Found,
213            SetOperation::InsertAtHead => GetOperation::AbsentAtHead,
214            SetOperation::InsertAfterEntry => GetOperation::AbsentAfterEntry,
215        };
216        (get_op, entry_ptr)
217    }
218}
219
220// LINK MAP ITER
221// ================================================================================================
222
223/// An iterator over a [`LinkMap`].
224struct LinkMapIter<'process> {
225    current_entry_ptr: u32,
226    map: LinkMap<'process>,
227}
228
229impl<'process> Iterator for LinkMapIter<'process> {
230    type Item = Entry;
231
232    fn next(&mut self) -> Option<Self::Item> {
233        if self.current_entry_ptr == 0 {
234            return None;
235        }
236
237        let current_entry = self.map.entry(self.current_entry_ptr);
238
239        self.current_entry_ptr = current_entry.metadata.next_entry_ptr;
240
241        Some(current_entry)
242    }
243}
244
245// LINK MAP TYPES
246// ================================================================================================
247
248/// An entry in a [`LinkMap`].
249///
250/// Exposed for testing purposes only.
251#[derive(Debug, Clone, Copy)]
252pub struct Entry {
253    pub ptr: u32,
254    pub metadata: EntryMetadata,
255    pub key: LexicographicWord,
256    pub value0: Word,
257    pub value1: Word,
258}
259
260/// An entry's metadata in a [`LinkMap`].
261///
262/// Exposed for testing purposes only.
263#[derive(Debug, Clone, Copy)]
264pub struct EntryMetadata {
265    pub map_ptr: u32,
266    pub prev_entry_ptr: u32,
267    pub next_entry_ptr: u32,
268}
269
270// HELPER TYPES
271// ================================================================================================
272
273/// The operation needed to get a key or prove its absence.
274#[derive(Debug, Clone, Copy)]
275#[repr(u8)]
276enum GetOperation {
277    Found = 0,
278    AbsentAtHead = 1,
279    AbsentAfterEntry = 2,
280}
281
282/// The operation needed to set a key.
283#[derive(Debug, Clone, Copy)]
284#[repr(u8)]
285enum SetOperation {
286    Update = 0,
287    InsertAtHead = 1,
288    InsertAfterEntry = 2,
289}
290
291// MEMORY VIEWER
292// ================================================================================================
293
294/// A abstraction over ways to view a process' memory.
295///
296/// More specifically, it allows using a [`LinkMap`] both with a [`ProcessState`], i.e. a process
297/// that is actively executing and also an [`ExecutionOutput`], i.e. a process that has finished
298/// execution.
299///
300/// This should all go away again once we change a LinkMap's implementation to be based on an actual
301/// map type instead of viewing a process' memory directly.
302pub enum MemoryViewer<'mem> {
303    ProcessState(&'mem ProcessState<'mem>),
304    ExecutionOutputs(&'mem ExecutionOutput),
305}
306
307impl<'mem> MemoryViewer<'mem> {
308    /// Reads an element from transaction kernel memory.
309    fn get_kernel_mem_element(&self, addr: u32) -> Option<Felt> {
310        match self {
311            MemoryViewer::ProcessState(process_state) => {
312                process_state.get_mem_value(ContextId::root(), addr)
313            },
314            MemoryViewer::ExecutionOutputs(_execution_output) => {
315                // TODO: Use Memory::read_element once it no longer requires &mut self.
316                // https://github.com/0xMiden/miden-vm/issues/2237
317
318                // Copy of how Memory::read_element is implemented in Miden VM.
319                let idx = addr % miden_protocol::WORD_SIZE as u32;
320                let word_addr = addr - idx;
321
322                Some(self.get_kernel_mem_word(word_addr)?[idx as usize])
323            },
324        }
325    }
326
327    /// Reads a word from transaction kernel memory.
328    fn get_kernel_mem_word(&self, addr: u32) -> Option<Word> {
329        match self {
330            MemoryViewer::ProcessState(process_state) => process_state
331                .get_mem_word(ContextId::root(), addr)
332                .expect("address should be word-aligned"),
333            MemoryViewer::ExecutionOutputs(execution_output) => {
334                let tx_kernel_context = ContextId::root();
335                let clk = 0u32;
336                let err_ctx = ();
337
338                // Note that this never returns None even if the location is uninitialized, but the
339                // link map does not rely on this.
340                Some(
341                    execution_output
342                        .memory
343                        .read_word(tx_kernel_context, Felt::from(addr), clk.into(), &err_ctx)
344                        .expect("expected address to be word-aligned"),
345                )
346            },
347        }
348    }
349}