Skip to content

Commit e7843f3

Browse files
anakrishCopilot
anakrish
and
Copilot
committed
refactor(rvm): use cursor for IterationState; iter_sorted at serialization
Switches RVM IterationState::Object and IterationState::Set from an eager Rc<[(Value, Value)]> / Rc<[Value]> pair snapshot built at loop / comprehension entry to a shared Rc<Object> / Rc<Set> plus an opaque cursor. This restores pre-v7 setup cost (O(1)) and per-element memory behaviour (no upfront snapshot, no per-element memory-limit checks). Snapshot independence from mid-iteration mutations is preserved via the Rc + copy-on-write: if a holder of an aliased Rc mutates via Rc::make_mut, a new collection is allocated and the iterator's Rc keeps pointing at the pre-mutation state. setup_next_iteration now takes &mut IterationState so the cursor can advance in-place; the suspendable yield/continue paths write the cursor-advanced state back to the frame's IterationState. Also rewires the RVM binary serializer (BinarySetRef, BinaryObjectRef) to use iter_sorted(), so canonical / content- addressable binary output is independent of the storage iteration order. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent 644fdcd commit e7843f3

4 files changed

Lines changed: 119 additions & 107 deletions

File tree

‎src/rvm/program/serialization/value.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ impl<'a> Serialize for BinarySetRef<'a> {
125125
S: serde::Serializer,
126126
{
127127
let mut seq = serializer.serialize_seq(Some(self.0.len()))?;
128-
for value in self.0.iter() {
128+
for value in self.0.iter_sorted() {
129129
seq.serialize_element(&BinaryValueRef(value))?;
130130
}
131131
seq.end()
@@ -140,7 +140,7 @@ impl<'a> Serialize for BinaryObjectRef<'a> {
140140
S: serde::Serializer,
141141
{
142142
let mut seq = serializer.serialize_seq(Some(self.0.len()))?;
143-
for (key, value) in self.0.iter() {
143+
for (key, value) in self.0.iter_sorted() {
144144
seq.serialize_element(&BinaryEntryRef(key, value))?;
145145
}
146146
seq.end()

‎src/rvm/vm/comprehension.rs‎

Lines changed: 25 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ impl RegoVM {
3939
self.set_register(params.result_reg, initial_result.clone())?;
4040

4141
let auto_iterate = params.collection_reg != params.result_reg;
42-
let iteration_state = if auto_iterate {
42+
let mut iteration_state = if auto_iterate {
4343
let source_value = self.get_register(params.collection_reg)?.clone();
4444
match source_value {
4545
Value::Array(items) => {
@@ -53,30 +53,17 @@ impl RegoVM {
5353
if obj.is_empty() {
5454
None
5555
} else {
56-
// Snapshot (key, value) pairs for iteration.
57-
let mut buf: Vec<(Value, Value)> = Vec::new();
58-
for (k, v) in obj.iter() {
59-
crate::utils::limits::check_memory_limit_if_needed()
60-
.map_err(anyhow::Error::msg)?;
61-
buf.push((k.clone(), v.clone()));
62-
}
63-
let pairs: Rc<[(Value, Value)]> = buf.into();
64-
Some(IterationState::Object { pairs, pos: 0 })
56+
// O(1) cursor over shared Rc<Object>.
57+
let cursor = obj.cursor();
58+
Some(IterationState::Object { obj, cursor })
6559
}
6660
}
6761
Value::Set(set) => {
6862
if set.is_empty() {
6963
None
7064
} else {
71-
// Build snapshot of values for iteration
72-
let mut buf: Vec<Value> = Vec::new();
73-
for v in set.iter() {
74-
crate::utils::limits::check_memory_limit_if_needed()
75-
.map_err(anyhow::Error::msg)?;
76-
buf.push(v.clone());
77-
}
78-
let values: Rc<[Value]> = buf.into();
79-
Some(IterationState::Set { values, pos: 0 })
65+
let cursor = set.cursor();
66+
Some(IterationState::Set { set, cursor })
8067
}
8168
}
8269
Value::Undefined => None,
@@ -87,7 +74,7 @@ impl RegoVM {
8774
None
8875
};
8976

90-
let has_iteration = if let Some(state) = iteration_state.as_ref() {
77+
let has_iteration = if let Some(state) = iteration_state.as_mut() {
9178
self.setup_next_iteration(state, params.key_reg, params.value_reg)?
9279
} else {
9380
false
@@ -136,7 +123,7 @@ impl RegoVM {
136123
self.set_register(params.result_reg, initial_result.clone())?;
137124

138125
let auto_iterate = params.collection_reg != params.result_reg;
139-
let iteration_state = if auto_iterate {
126+
let mut iteration_state = if auto_iterate {
140127
let source_value = self.get_register(params.collection_reg)?.clone();
141128
match source_value {
142129
Value::Array(items) => {
@@ -150,30 +137,16 @@ impl RegoVM {
150137
if obj.is_empty() {
151138
None
152139
} else {
153-
// Snapshot (key, value) pairs for iteration.
154-
let mut buf: Vec<(Value, Value)> = Vec::new();
155-
for (k, v) in obj.iter() {
156-
crate::utils::limits::check_memory_limit_if_needed()
157-
.map_err(anyhow::Error::msg)?;
158-
buf.push((k.clone(), v.clone()));
159-
}
160-
let pairs: Rc<[(Value, Value)]> = buf.into();
161-
Some(IterationState::Object { pairs, pos: 0 })
140+
let cursor = obj.cursor();
141+
Some(IterationState::Object { obj, cursor })
162142
}
163143
}
164144
Value::Set(set) => {
165145
if set.is_empty() {
166146
None
167147
} else {
168-
// Build snapshot of values for iteration
169-
let mut buf: Vec<Value> = Vec::new();
170-
for v in set.iter() {
171-
crate::utils::limits::check_memory_limit_if_needed()
172-
.map_err(anyhow::Error::msg)?;
173-
buf.push(v.clone());
174-
}
175-
let values: Rc<[Value]> = buf.into();
176-
Some(IterationState::Set { values, pos: 0 })
148+
let cursor = set.cursor();
149+
Some(IterationState::Set { set, cursor })
177150
}
178151
}
179152
Value::Undefined => None,
@@ -184,7 +157,7 @@ impl RegoVM {
184157
None
185158
};
186159

187-
let has_iteration = if let Some(state) = iteration_state.as_ref() {
160+
let has_iteration = if let Some(state) = iteration_state.as_mut() {
188161
self.setup_next_iteration(state, params.key_reg, params.value_reg)?
189162
} else {
190163
false
@@ -446,8 +419,18 @@ impl RegoVM {
446419
}
447420
};
448421

449-
if let Some(state) = iteration_state_snapshot.as_ref() {
450-
let has_next = self.setup_next_iteration(state, key_reg_idx, value_reg_idx)?;
422+
if let Some(mut state) = iteration_state_snapshot {
423+
let has_next = self.setup_next_iteration(&mut state, key_reg_idx, value_reg_idx)?;
424+
425+
// Write the cursor-advanced state back to the frame.
426+
if let Some(frame) = self.execution_stack.get_mut(comprehension_index) {
427+
if let FrameKind::Comprehension {
428+
ref mut context, ..
429+
} = frame.kind
430+
{
431+
context.iteration_state = Some(state);
432+
}
433+
}
451434

452435
if has_next {
453436
if let Some(frame) = self.execution_stack.get_mut(comprehension_index) {

‎src/rvm/vm/context.rs‎

Lines changed: 37 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
// Copyright (c) Microsoft Corporation.
22
// Licensed under the MIT License.
33

4+
use crate::collections::{Object, ObjectCursor, Set, SetCursor};
45
use crate::rvm::instructions::{ComprehensionMode, LoopMode};
56
use crate::value::Value;
67
use crate::Rc;
@@ -23,20 +24,26 @@ pub struct LoopContext {
2324
pub current_iteration_failed: bool, // Track if current iteration had condition failures
2425
}
2526

26-
/// Iterator state for different collection types
27+
/// Iterator state for different collection types.
28+
///
29+
/// Snapshot independence is provided by the shared `Rc<Object>` /
30+
/// `Rc<Set>` / `Rc<Vec<Value>>` — `Rc::make_mut` on an aliased Rc allocates
31+
/// a new collection, leaving the iterator's Rc pointing at the original
32+
/// pre-mutation state. The cursor for `Object` / `Set` is opaque and resumes
33+
/// in O(log n) for the BTree backend.
2734
#[derive(Debug, Clone)]
2835
pub enum IterationState {
2936
Array {
3037
items: Rc<Vec<Value>>,
3138
index: usize,
3239
},
3340
Object {
34-
pairs: Rc<[(Value, Value)]>,
35-
pos: usize,
41+
obj: Rc<Object>,
42+
cursor: ObjectCursor,
3643
},
3744
Set {
38-
values: Rc<[Value]>,
39-
pos: usize,
45+
set: Rc<Set>,
46+
cursor: SetCursor,
4047
},
4148
/// Virtual single-element iteration for non-collection values.
4249
/// Used by Azure Policy's `[*]` on scalar/null fields: presents a single
@@ -53,9 +60,10 @@ impl IterationState {
5360
Self::Array { ref mut index, .. } => {
5461
*index = index.saturating_add(1);
5562
}
56-
Self::Object { ref mut pos, .. } | Self::Set { ref mut pos, .. } => {
57-
*pos = pos.saturating_add(1);
58-
}
63+
// For Object/Set the cursor advances inside `setup_next_iteration`
64+
// when it pulls the next item via `Object::next` / `Set::next`,
65+
// so `advance` is a no-op for cursor-backed variants.
66+
Self::Object { .. } | Self::Set { .. } => {}
5967
Self::Single {
6068
ref mut consumed, ..
6169
} => {
@@ -111,9 +119,9 @@ mod tests {
111119
use super::*;
112120
use crate::collections::Object;
113121

114-
/// IterationState::Object snapshots (key, value) pairs at construction
115-
/// time. Mutating the source Value::Object (via Rc::make_mut on a clone)
116-
/// while the iteration is in flight must not affect the snapshot.
122+
/// IterationState::Object holds an `Rc<Object>` plus an opaque cursor.
123+
/// Mutating an aliased Rc via `Rc::make_mut` allocates a new collection
124+
/// (CoW) so the in-flight iterator's source is unaffected.
117125
#[test]
118126
fn iteration_state_object_is_snapshot_independent_of_source() {
119127
let mut obj = Object::new();
@@ -123,18 +131,13 @@ mod tests {
123131

124132
let source = Value::Object(Rc::new(obj));
125133

126-
// Build the snapshot exactly as loops.rs / comprehension.rs do.
127-
let snapshot_pairs: Rc<[(Value, Value)]> = match &source {
128-
Value::Object(o) => o
129-
.iter()
130-
.map(|(k, v)| (k.clone(), v.clone()))
131-
.collect::<Vec<_>>()
132-
.into(),
134+
let snapshot_obj = match &source {
135+
Value::Object(o) => Rc::clone(o),
133136
_ => unreachable!(),
134137
};
135138
let state = IterationState::Object {
136-
pairs: Rc::clone(&snapshot_pairs),
137-
pos: 0,
139+
obj: Rc::clone(&snapshot_obj),
140+
cursor: snapshot_obj.cursor(),
138141
};
139142

140143
// Mutate a clone of the source mid-iteration.
@@ -144,11 +147,20 @@ mod tests {
144147
inner.insert(Value::from("d"), Value::from(4));
145148
inner.remove(&Value::from("b"));
146149

147-
// Snapshot must still report the original 3 entries with original values.
148-
let collected: Vec<(Value, Value)> = match &state {
149-
IterationState::Object { pairs, .. } => pairs.to_vec(),
150-
_ => unreachable!(),
151-
};
150+
// Drain the snapshot via the cursor — must still report the original
151+
// 3 entries with original values.
152+
let mut collected: Vec<(Value, Value)> = Vec::new();
153+
if let IterationState::Object {
154+
ref obj,
155+
mut cursor,
156+
} = state
157+
{
158+
while let Some((k, v)) = obj.next(&mut cursor) {
159+
collected.push((k.clone(), v.clone()));
160+
}
161+
} else {
162+
unreachable!();
163+
}
152164
assert_eq!(collected.len(), 3);
153165
assert!(collected.contains(&(Value::from("a"), Value::from(1))));
154166
assert!(collected.contains(&(Value::from("b"), Value::from(2))));

0 commit comments

Comments
 (0)