miden_stdlib/handlers/
sorted_array.rs

1use alloc::{vec, vec::Vec};
2
3use miden_core::{EventName, Felt, FieldElement, LexicographicWord, Word};
4use miden_processor::{AdviceMutation, EventError, MemoryError, ProcessState};
5
6/// Event name for the lowerbound_array operation.
7pub const LOWERBOUND_ARRAY_EVENT_NAME: EventName =
8    EventName::new("stdlib::collections::sorted_array::lowerbound_array");
9
10/// Event name for the lowerbound_key_value operation.
11pub const LOWERBOUND_KEY_VALUE_EVENT_NAME: EventName =
12    EventName::new("stdlib::collections::sorted_array::lowerbound_key_value");
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15enum KeySize {
16    Full,
17    Half,
18}
19
20/// Pushes onto the advice stack the first pointer in [start_ptr, end_ptr) such that
21/// `mem[word_ptr] >= KEY` in lexicographic order of words. If all words are < KEY, returns end_ptr.
22/// The array must be sorted in non-decreasing order.
23///
24/// Inputs:
25///   Operand stack: [KEY, start_ptr, end_ptr, ...]
26///   Advice stack: [...]
27///
28/// Outputs:
29///   Operand stack: [KEY, start_ptr, end_ptr, ...]
30///   Advice stack: [maybe_key_ptr, was_key_found, ...]
31///
32/// # Errors
33/// Returns an error if the provided word array is not sorted in non-decreasing order.
34pub fn handle_lowerbound_array(process: &ProcessState) -> Result<Vec<AdviceMutation>, EventError> {
35    push_lowerbound_result(process, 4, KeySize::Full)
36}
37
38/// Pushes onto the advice stack the first pointer in [start_ptr, end_ptr) such that
39/// `mem[word_ptr] >= KEY` in lexicographic order of words. If all keys are < KEY, returns end_ptr.
40/// The array must be a list of key-value pairs (word tuples) where all keys are in non-decreasing
41/// order.
42///
43/// This event returns
44///
45/// Inputs:
46///   Operand stack: [event_id, KEY, start_ptr, end_ptr, use_full_key, ...]
47///   Advice stack: [...]
48///
49/// Outputs:
50///   Operand stack: [event_id, KEY, start_ptr, end_ptr, use_full_key, ...]
51///   Advice stack: [maybe_key_ptr, was_key_found, ...]
52///
53/// # Errors
54/// Returns an error if the keys are not sorted in non-decreasing order.
55pub fn handle_lowerbound_key_value(
56    process: &ProcessState,
57) -> Result<Vec<AdviceMutation>, EventError> {
58    let use_full_key = process.get_stack_item(7);
59
60    let key_size = match use_full_key.as_int() {
61        0 => KeySize::Half,
62        1 => KeySize::Full,
63        _ => {
64            return Err(EventError::from(alloc::format!(
65                "use_full_key must be 0 or 1, was {use_full_key}"
66            )));
67        },
68    };
69
70    push_lowerbound_result(process, 8, key_size)
71}
72
73/// Offsets for the push_lowerbound_result inputs from the top of the stack
74const KEY_OFFSET: usize = 1;
75const START_ADDR_OFFSET: usize = 5;
76const END_ADDR_OFFSET: usize = 6;
77
78fn push_lowerbound_result(
79    process: &ProcessState,
80    stride: u32,
81    key_size: KeySize,
82) -> Result<Vec<AdviceMutation>, EventError> {
83    // only support sorted arrays (stride = 4) and sorted key-value arrays (stride = 8)
84    assert!(stride == 4 || stride == 8);
85
86    // Read inputs from the stack
87    let key = LexicographicWord::new(process.get_stack_word_be(KEY_OFFSET));
88    let addr_range = process.get_mem_addr_range(START_ADDR_OFFSET, END_ADDR_OFFSET)?;
89
90    // Validate the start_addr is word-aligned (multiple of 4)
91    if addr_range.start % 4 != 0 {
92        return Err(MemoryError::unaligned_word_access(
93            addr_range.start,
94            process.ctx(),
95            Felt::from(process.clk()),
96            &(),
97        )
98        .into());
99    }
100
101    // Validate the end_addr is properly aligned (i.e. the entire array has size divisible by
102    // stride)
103    if (addr_range.end - addr_range.start) % stride != 0 {
104        if stride == 4 {
105            return Err(
106                SortedArrayError::InvalidArrayRange { size: addr_range.len() as u32 }.into()
107            );
108        } else {
109            return Err(
110                SortedArrayError::InvalidKeyValueRange { size: addr_range.len() as u32 }.into()
111            );
112        }
113    }
114
115    // If range is empty, result is end_ptr
116    if addr_range.is_empty() {
117        return Ok(vec![AdviceMutation::extend_stack(vec![
118            Felt::from(false),
119            Felt::from(addr_range.end),
120        ])]);
121    }
122
123    // Helper function to get a word from memory and convert it to a LexicographicWord
124    let get_word = {
125        |addr: u32| {
126            process
127                .get_mem_word(process.ctx(), addr)
128                .map(|word| word_to_search_key(word.unwrap_or_default(), key_size))
129        }
130    };
131
132    let mut was_key_found = false;
133    let mut result = None;
134
135    // Test the first element
136    let mut previous_word = get_word(addr_range.start)?;
137    if previous_word >= key {
138        was_key_found = previous_word == key;
139        result = Some(addr_range.start);
140    }
141
142    // Validate the entire array is non-decreasing and find the first element where `element >= key`
143    for addr in addr_range.clone().step_by(stride as usize).skip(1) {
144        let word = get_word(addr)?;
145        if word < previous_word {
146            return Err(SortedArrayError::NotAscendingOrder {
147                index: addr,
148                value: word.into(),
149                predecessor: previous_word.into(),
150            }
151            .into());
152        }
153        if word >= key && result.is_none() {
154            was_key_found = word == key;
155            result = Some(addr);
156        }
157        previous_word = word;
158    }
159
160    Ok(vec![AdviceMutation::extend_stack(vec![
161        Felt::from(was_key_found),
162        Felt::from(result.unwrap_or(addr_range.end)),
163    ])])
164}
165
166/// Selectively zeroizes the felts in a [`Word`] based on the provided [`KeySize`].
167///
168/// - If the `key_size` is [`KeySize::Full`], the word is returned untouched.
169/// - If the `key_size` is [`KeySize::Half`], the word is returned with the two least significant
170///   elements zeroized.
171fn word_to_search_key(mut word: Word, key_size: KeySize) -> LexicographicWord {
172    match key_size {
173        KeySize::Full => LexicographicWord::new(word),
174        KeySize::Half => {
175            word[0] = Felt::ZERO;
176            word[1] = Felt::ZERO;
177
178            LexicographicWord::new(word)
179        },
180    }
181}
182
183// ERROR TYPES
184// ================================================================================================
185
186/// Error types that can occur during LOWERBOUND event operations.
187#[derive(Debug, thiserror::Error)]
188pub enum SortedArrayError {
189    /// Elements are not sorted in non-decreasing order.
190    #[error("element at index {index} ({value}) is smaller than the predecessor ({predecessor})")]
191    NotAscendingOrder {
192        index: u32,
193        value: Word,
194        predecessor: Word,
195    },
196
197    /// Last element is an incomplete word.
198    #[error("array size must be divisible by 4, but was {size}")]
199    InvalidArrayRange { size: u32 },
200
201    /// Last key or value is an incomplete word.
202    #[error("key-value array must have size divisible by 4 or 8, but was {size}")]
203    InvalidKeyValueRange { size: u32 },
204}