diff --git a/src/builtins/azure_policy/helpers.rs b/src/builtins/azure_policy/helpers.rs index f9eadad..cd07c66 100644 --- a/src/builtins/azure_policy/helpers.rs +++ b/src/builtins/azure_policy/helpers.rs @@ -319,7 +319,7 @@ pub fn resolve_path(root: &Value, path: &str) -> Value { match ¤t { Value::Object(map) => { let mut next = None; - for (key, value) in map.iter() { + for (key, value) in map.iter_sorted() { if let Value::String(ref key_str) = *key { if strings::keys::eq(key_str, &segment) { next = Some(value.clone()); diff --git a/src/builtins/azure_policy/template_functions_collection.rs b/src/builtins/azure_policy/template_functions_collection.rs index 88cd6c8..1f774c5 100644 --- a/src/builtins/azure_policy/template_functions_collection.rs +++ b/src/builtins/azure_policy/template_functions_collection.rs @@ -8,10 +8,10 @@ use crate::ast::{Expr, Ref}; use crate::builtins; use crate::lexer::Span; +use crate::value::Object; use crate::value::Value; use crate::Rc; -use alloc::collections::BTreeMap; use alloc::vec::Vec; use anyhow::Result; @@ -72,7 +72,7 @@ fn fn_intersection( // Intersection of objects: keep key-value pairs from the first // object only when the key exists in every other object AND // the value is equal across all of them. - let mut result: BTreeMap = first.as_ref().clone(); + let mut result: Object = first.as_ref().clone(); for arg in rest { let Value::Object(ref other) = *arg else { return Ok(Value::Undefined); @@ -114,7 +114,7 @@ fn fn_union(_span: &Span, _params: &[Ref], args: &[Value], _strict: bool) Value::Object(_) => { // Union of objects: recursive merge. Nested objects are merged // recursively; all other types (including arrays) use last-writer-wins. - let mut result = BTreeMap::::new(); + let mut result = Object::new(); for arg in args { let Value::Object(ref obj) = *arg else { return Ok(Value::Undefined); @@ -264,7 +264,7 @@ fn fn_create_object( ); } - let mut map = BTreeMap::::new(); + let mut map = Object::new(); for pair in args.chunks(2) { #[allow(clippy::pattern_type_mismatch)] @@ -280,9 +280,9 @@ fn fn_create_object( /// Recursively merge two objects. Nested objects are merged; everything /// else (including arrays) uses the value from `incoming`. -fn merge_objects(base: &BTreeMap, overlay: &BTreeMap) -> Value { +fn merge_objects(base: &Object, overlay: &Object) -> Value { let mut result = base.clone(); - for (k, v) in overlay { + for (k, v) in overlay.iter() { #[allow(clippy::needless_borrowed_reference)] let merged = match (result.get(k), v) { (Some(&Value::Object(ref prev)), &Value::Object(ref next)) => merge_objects(prev, next), diff --git a/src/builtins/azure_policy/template_functions_misc.rs b/src/builtins/azure_policy/template_functions_misc.rs index 6f97861..1fbcacb 100644 --- a/src/builtins/azure_policy/template_functions_misc.rs +++ b/src/builtins/azure_policy/template_functions_misc.rs @@ -8,10 +8,10 @@ use crate::ast::{Expr, Ref}; use crate::builtins; use crate::lexer::Span; +use crate::value::Object; use crate::value::Value; use crate::Rc; -use alloc::collections::BTreeMap; use alloc::string::{String, ToString as _}; use alloc::vec::Vec; use anyhow::Result; @@ -84,8 +84,8 @@ fn fn_items(_span: &Span, _params: &[Ref], args: &[Value], _strict: bool) return Ok(Value::Undefined); }; let mut result = Vec::with_capacity(obj.len()); - for (k, v) in obj.as_ref() { - let mut entry = BTreeMap::::new(); + for (k, v) in obj.iter_sorted() { + let mut entry = Object::new(); entry.insert(Value::from("key"), k.clone()); entry.insert(Value::from("value"), v.clone()); result.push(Value::Object(Rc::new(entry))); diff --git a/src/builtins/encoding.rs b/src/builtins/encoding.rs index 4e732a0..588f99e 100644 --- a/src/builtins/encoding.rs +++ b/src/builtins/encoding.rs @@ -308,7 +308,7 @@ fn urlquery_encode_object( { let mut pairs = url.query_pairs_mut(); - for (key, value) in obj.iter() { + for (key, value) in obj.iter_sorted() { let key = ensure_string(name, ¶ms[0], key)?; match value { Value::String(v) => { diff --git a/src/builtins/graph.rs b/src/builtins/graph.rs index 09f5ce8..c294299 100644 --- a/src/builtins/graph.rs +++ b/src/builtins/graph.rs @@ -7,10 +7,11 @@ use crate::ast::{Expr, Ref}; use crate::builtins; use crate::builtins::utils::{enforce_limit, ensure_args_count, ensure_object}; use crate::lexer::Span; +use crate::value::Object; use crate::value::Value; use crate::*; -use alloc::collections::{BTreeMap, BTreeSet}; +use alloc::collections::BTreeSet; use anyhow::{bail, Result}; @@ -80,7 +81,7 @@ fn reachable(span: &Span, params: &[Ref], args: &[Value], strict: bool) -> } fn visit( - graph: &BTreeMap, + graph: &Object, visited: &mut BTreeSet, node: &Value, path: &mut Vec, @@ -211,7 +212,7 @@ fn walk_visit(path: &mut Vec, value: &Value, paths: &mut Vec) -> R } } Value::Object(obj) => { - for (key, value) in obj.iter() { + for (key, value) in obj.iter_sorted() { path.push(key.clone()); // Guard path stack growth while traversing object entries. enforce_limit()?; diff --git a/src/builtins/objects.rs b/src/builtins/objects.rs index 49daa8f..b5fac7b 100644 --- a/src/builtins/objects.rs +++ b/src/builtins/objects.rs @@ -205,7 +205,7 @@ fn merge_filters( let vref = match f { Value::Object(obj) => { let obj = Rc::make_mut(obj); - let entry = obj.entry(p.clone()).or_insert_with(Value::new_object); + let entry = obj.get_or_insert_with(p.clone(), Value::new_object); // Guard filter map growth when creating nested objects. enforce_limit()?; entry diff --git a/src/builtins/strings.rs b/src/builtins/strings.rs index fad7309..fd3a2ee 100644 --- a/src/builtins/strings.rs +++ b/src/builtins/strings.rs @@ -207,7 +207,7 @@ fn to_string(v: &Value, unescape: bool) -> String { } Value::Object(o) => { "{".to_owned() - + &o.iter() + + &o.iter_sorted() .map(|(k, v)| to_string(k, true) + ": " + &to_string(v, true)) .collect::>() .join(", ") @@ -568,7 +568,7 @@ fn replace_n(span: &Span, params: &[Ref], args: &[Value], _strict: bool) - let mut s = ensure_string(name, ¶ms[1], &args[1])?; let span = params[0].span(); - for item in obj.as_ref().iter() { + for item in obj.as_ref().iter_sorted() { match item { (Value::String(k), Value::String(v)) => { s = s.replace(k.as_ref(), v.as_ref()).into(); diff --git a/src/builtins/utils.rs b/src/builtins/utils.rs index 65de2c9..77f3b85 100644 --- a/src/builtins/utils.rs +++ b/src/builtins/utils.rs @@ -5,11 +5,12 @@ use crate::ast::{Expr, Ref}; use crate::lexer::Span; use crate::number::Number; +use crate::value::Object; use crate::Rc; use crate::Value; use crate::*; -use alloc::collections::{BTreeMap, BTreeSet}; +use alloc::collections::BTreeSet; use anyhow::{bail, Result}; @@ -168,7 +169,7 @@ pub fn ensure_set(fcn: &str, arg: &Expr, v: Value) -> Result> }) } -pub fn ensure_object(fcn: &str, arg: &Expr, v: Value) -> Result>> { +pub fn ensure_object(fcn: &str, arg: &Expr, v: Value) -> Result> { Ok(match v { Value::Object(o) => o, _ => { diff --git a/src/interpreter.rs b/src/interpreter.rs index 7f7e154..2028eb9 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -28,7 +28,6 @@ use crate::{Expression, Extension, Location, QueryResult, QueryResults}; use crate::query::traversal::traverse; use crate::Rc; -use alloc::collections::btree_map::Entry as BTreeMapEntry; use alloc::collections::{BTreeMap, BTreeSet}; use anyhow::{anyhow, bail, Result}; use core::ops::Bound::*; @@ -1312,10 +1311,10 @@ impl Interpreter { *obj = Value::new_object(); } - obj = obj - .as_object_mut()? - .entry(Value::String(p.to_string().into())) - .or_insert(Value::new_object()); + obj = obj.as_object_mut()?.get_or_insert_with( + Value::String(p.to_string().into()), + Value::new_object, + ); } *obj = value; // Mark modified rules as processed. @@ -1682,8 +1681,7 @@ impl Interpreter { let set = obj .as_object_mut() .map_err(|_| anyhow!(span.error("previous value is not an object")))? - .entry(p) - .or_insert(Value::new_set()) + .get_or_insert_with(p, Value::new_set) .as_set_mut() .map_err(|_| anyhow!(span.error("previous value is not a set")))?; set.append(value.as_set_mut()?); @@ -1691,20 +1689,13 @@ impl Interpreter { let obj = obj .as_object_mut() .map_err(|_| anyhow!(span.error("previous value is not an object")))?; - match obj.entry(p) { - BTreeMapEntry::Vacant(v) => { - if value != Value::Undefined { - v.insert(value); - } else { - // TODO: clean this assumption between Undefined vs Object. - v.insert(Value::new_object()); - } - } - BTreeMapEntry::Occupied(o) => { - if o.get() != &value && value != Value::Undefined { - bail!(span - .error("complete rules should not produce multiple outputs")) - } + if value == Value::Undefined { + // TODO: clean this assumption between Undefined vs Object. + obj.get_or_insert_with(p, Value::new_object); + } else { + let existing = obj.get_or_insert_with(p, || value.clone()); + if *existing != value { + bail!(span.error("complete rules should not produce multiple outputs")) } } } @@ -1713,8 +1704,7 @@ impl Interpreter { obj = obj .as_object_mut() .map_err(|_| anyhow!(span.error("previous value is not an object")))? - .entry(p) - .or_insert(Value::new_object()); + .get_or_insert_with(p, Value::new_object); } } Ok(()) @@ -1822,8 +1812,7 @@ impl Interpreter { let set = ctx_mut .rule_value .as_object_mut()? - .entry(Value::from_array(comps)) - .or_insert(Value::new_set()); + .get_or_insert_with(Value::from_array(comps), Value::new_set); if output != Value::Undefined { set.as_set_mut()?.insert(output); return Ok(true); @@ -1832,20 +1821,13 @@ impl Interpreter { } // Non-set rule. - match ctx_mut - .rule_value - .as_object_mut()? - .entry(Value::from_array(comps)) - { - BTreeMapEntry::Vacant(v) => { - v.insert(output); - } - BTreeMapEntry::Occupied(o) if o.get() != &output => bail!(rule_ref + let key = Value::from_array(comps); + let obj_mut = ctx_mut.rule_value.as_object_mut()?; + let existing = obj_mut.get_or_insert_with(key, || output.clone()); + if *existing != output { + bail!(rule_ref .span() - .error("rules must not produce multiple outputs")), - _ => { - // Rule produced same value. - } + .error("rules must not produce multiple outputs")); } return Ok(true); @@ -2471,7 +2453,7 @@ impl Interpreter { } Value::Object(map) => { s.push('{'); - for (idx, (k, entry_value)) in map.iter().enumerate() { + for (idx, (k, entry_value)) in map.iter_sorted().enumerate() { if idx > 0 { s.push_str(", "); } diff --git a/src/languages/azure_policy/aliases/denormalizer/mod.rs b/src/languages/azure_policy/aliases/denormalizer/mod.rs index 67f9639..8021e53 100644 --- a/src/languages/azure_policy/aliases/denormalizer/mod.rs +++ b/src/languages/azure_policy/aliases/denormalizer/mod.rs @@ -213,10 +213,10 @@ pub fn denormalize_with_aliases( // Phase 4: Attach properties to result. if !properties.is_empty() { if let Some(Value::Object(existing_rc)) = result.get_mut("properties") { - // Merge directly into the BTreeMap, avoiding full ObjMap round-trip. + // Merge directly into the Object, avoiding full ObjMap round-trip. let existing = Rc::make_mut(existing_rc); for (k, v) in properties { - existing.entry(Value::String(k)).or_insert(v); + existing.get_or_insert_with(Value::String(k), || v); } } else { obj_insert(&mut result, "properties", make_value(properties)); diff --git a/src/languages/azure_policy/aliases/denormalizer/sub_resource.rs b/src/languages/azure_policy/aliases/denormalizer/sub_resource.rs index 6ab9897..69f6719 100644 --- a/src/languages/azure_policy/aliases/denormalizer/sub_resource.rs +++ b/src/languages/azure_policy/aliases/denormalizer/sub_resource.rs @@ -7,6 +7,7 @@ use alloc::collections::{BTreeMap, BTreeSet}; use alloc::string::String; use alloc::vec::Vec; +use crate::value::Object; use crate::Value; use super::super::obj_map::{make_value, new_map, obj_insert, val_str, ObjMap}; @@ -141,7 +142,7 @@ fn rewrap_nested_array( /// BTreeMap-native recursion for nested sub-resource array re-wrapping, /// avoiding ObjMap round-trips on each array element. fn rewrap_nested_array_in_btree( - btree: &mut alloc::collections::BTreeMap, + btree: &mut Object, parent_parts: &[&str], array_name: &str, envelope_fields: &BTreeSet, @@ -187,10 +188,7 @@ fn rewrap_nested_array_in_btree( } /// Find a key in a BTreeMap using case-insensitive comparison. -fn find_key_ci_btree( - btree: &alloc::collections::BTreeMap, - key: &str, -) -> Option { +fn find_key_ci_btree(btree: &Object, key: &str) -> Option { btree .keys() .find(|k| val_str(k).is_some_and(|s| s.eq_ignore_ascii_case(key))) diff --git a/src/languages/azure_policy/aliases/normalizer/element_remap.rs b/src/languages/azure_policy/aliases/normalizer/element_remap.rs index c1b413b..2782e28 100644 --- a/src/languages/azure_policy/aliases/normalizer/element_remap.rs +++ b/src/languages/azure_policy/aliases/normalizer/element_remap.rs @@ -6,11 +6,12 @@ use alloc::string::String; use alloc::vec::Vec; +use crate::value::Object; use crate::Value; use super::super::obj_map::{ - obj_get, obj_get_mut, obj_insert, set_nested_in_btree, set_nested_lowercased, - set_nested_verbatim, ObjMap, + obj_get, obj_get_mut, obj_insert, set_nested, set_nested_lowercased, set_nested_verbatim, + ObjMap, }; use super::super::types::PrecomputedRemap; @@ -118,7 +119,7 @@ fn apply_remap_at_depth( /// BTreeMap-native recursion for element-level remap, avoiding ObjMap /// round-trips on each array element. fn remap_at_depth_in_btree( - btree: &mut alloc::collections::BTreeMap, + btree: &mut Object, array_chain: &[Vec], depth: usize, source_field: &str, @@ -177,12 +178,7 @@ fn remap_at_depth_in_btree( } /// Remap a value between dotted paths directly in a BTreeMap. -fn remap_deep_field_in_btree( - btree: &mut alloc::collections::BTreeMap, - source: &str, - target: &str, - lowercase: bool, -) { +fn remap_deep_field_in_btree(btree: &mut Object, source: &str, target: &str, lowercase: bool) { let val = match read_dotted_path_btree(btree, source) { Some(v) => v, None => return, @@ -198,14 +194,11 @@ fn remap_deep_field_in_btree( } return; } - set_nested_in_btree(btree, &segments, val, lowercase); + set_nested(btree, &segments, val, lowercase); } /// Read a value at a dotted path from a BTreeMap. -fn read_dotted_path_btree( - btree: &alloc::collections::BTreeMap, - path: &str, -) -> Option { +fn read_dotted_path_btree(btree: &Object, path: &str) -> Option { let segments: Vec<&str> = path.split('.').collect(); let first = segments.first()?; let mut cur: &Value = btree.get(&Value::from(*first))?; diff --git a/src/languages/azure_policy/aliases/normalizer/mod.rs b/src/languages/azure_policy/aliases/normalizer/mod.rs index c577964..09b366b 100644 --- a/src/languages/azure_policy/aliases/normalizer/mod.rs +++ b/src/languages/azure_policy/aliases/normalizer/mod.rs @@ -13,6 +13,7 @@ mod flatten; // Re-export items used by the denormalizer. pub(crate) use element_remap::{apply_element_remap, ElementRemap}; +use crate::value::Object; use crate::Value; use super::obj_map::{ @@ -109,7 +110,7 @@ pub fn normalize_with_aliases( /// Merge `properties` fields into the result map, skipping keys that already /// exist. fn merge_properties( - obj: &alloc::collections::BTreeMap, + obj: &Object, result: &mut ObjMap, sub_arrays: Option<&alloc::collections::BTreeSet>, ) { diff --git a/src/languages/azure_policy/aliases/obj_map.rs b/src/languages/azure_policy/aliases/obj_map.rs index 65c5bb0..3598cd2 100644 --- a/src/languages/azure_policy/aliases/obj_map.rs +++ b/src/languages/azure_policy/aliases/obj_map.rs @@ -4,7 +4,7 @@ //! Lightweight string-keyed map used during normalization/denormalization. //! //! Internally uses `hashbrown::HashMap, Value>` for O(1) lookups, -//! then converts to `Value::Object` (a `BTreeMap`) only at +//! then converts to `Value::Object` (an `Object`) only at //! the output boundary via [`make_value`]. use alloc::string::String; @@ -12,6 +12,7 @@ use alloc::vec::Vec; use hashbrown::HashMap; +use crate::value::Object; use crate::Rc; use crate::Value; @@ -81,14 +82,13 @@ pub fn obj_remove(map: &mut ObjMap, key: &str) -> Option { /// Convert an [`ObjMap`] into a [`Value::Object`]. /// /// Keys are converted from `Rc` to `Value::String` and inserted into -/// a `BTreeMap` to match the `Value::Object` representation. +/// an `Object` to match the `Value::Object` representation. pub fn make_value(map: ObjMap) -> Value { - use alloc::collections::BTreeMap; - let mut btree = BTreeMap::new(); - for (k, v) in map { - btree.insert(Value::String(k), v); - } - Value::Object(Rc::new(btree)) + let obj: Object = map + .into_iter() + .map(|(k, v)| (Value::String(k), v)) + .collect(); + Value::Object(Rc::new(obj)) } /// Convert a `Vec` into a `Value::Array`. @@ -115,14 +115,14 @@ pub fn extract_type_field(resource: &Value) -> Option<&str> { }) } -/// Convert a `Value::Object` (BTreeMap) into an [`ObjMap`]. +/// Convert a `Value::Object` (Object) into an [`ObjMap`]. /// /// Non-string keys are silently skipped. #[allow(dead_code)] pub fn value_to_obj_map(value: &Value) -> Option { - let btree = value.as_object().ok()?; - let mut map = ObjMap::with_capacity(btree.len()); - for (k, v) in btree.iter() { + let obj = value.as_object().ok()?; + let mut map = ObjMap::with_capacity(obj.len()); + for (k, v) in obj.iter() { if let Value::String(s) = k { map.insert(Rc::clone(s), v.clone()); } @@ -194,7 +194,7 @@ fn set_nested_inner(obj: &mut ObjMap, segments: &[&str], value: Value, lowercase // Descend directly into the BTreeMap, avoiding ObjMap round-trip. if let Some(Value::Object(inner_rc)) = obj.get_mut(&*seg) { let inner_btree = Rc::make_mut(inner_rc); - set_nested_in_btree( + set_nested( inner_btree, segments.get(1..).unwrap_or_default(), value, @@ -203,17 +203,12 @@ fn set_nested_inner(obj: &mut ObjMap, segments: &[&str], value: Value, lowercase } } -/// Set a value at a path directly in a `BTreeMap`, creating +/// Set a value at a path directly in an `Object`, creating /// intermediate `Value::Object` nodes as needed. /// /// This avoids the `btree_to_obj_map` / `obj_map_to_btree` round-trip that /// would clone every sibling entry at each nesting level. -pub fn set_nested_in_btree( - btree: &mut alloc::collections::BTreeMap, - segments: &[&str], - value: Value, - lowercase: bool, -) { +pub fn set_nested(obj: &mut Object, segments: &[&str], value: Value, lowercase: bool) { let Some(&first) = segments.first() else { return; }; @@ -226,18 +221,18 @@ pub fn set_nested_in_btree( let key_val = Value::String(Rc::clone(&key_rc)); if segments.len() == 1 { - btree.insert(key_val, value); + obj.insert(key_val, value); return; } // Ensure an intermediate object exists. - if !btree.contains_key(&key_val) { - btree.insert(key_val.clone(), make_value(new_map())); + if !obj.contains_key(&key_val) { + obj.insert(key_val.clone(), make_value(new_map())); } - if let Some(Value::Object(inner_rc)) = btree.get_mut(&key_val) { + if let Some(Value::Object(inner_rc)) = obj.get_mut(&key_val) { let inner = Rc::make_mut(inner_rc); - set_nested_in_btree( + set_nested( inner, segments.get(1..).unwrap_or_default(), value, @@ -352,20 +347,15 @@ fn remove_field_at_depth(obj: &mut ObjMap, array_chain: &[Vec], depth: u for elem in inner.iter_mut() { if let Value::Object(obj_rc) = elem { let inner_btree = Rc::make_mut(obj_rc); - remove_field_at_depth_in_btree( - inner_btree, - array_chain, - depth.saturating_add(1), - field, - ); + remove_field_at_depth_obj(inner_btree, array_chain, depth.saturating_add(1), field); } } } } -/// BTreeMap-native recursion for element-level field removal. -fn remove_field_at_depth_in_btree( - btree: &mut alloc::collections::BTreeMap, +/// Object-native recursion for element-level field removal. +fn remove_field_at_depth_obj( + obj: &mut Object, array_chain: &[Vec], depth: usize, field: &str, @@ -374,10 +364,10 @@ fn remove_field_at_depth_in_btree( let segments: Vec<&str> = field.split('.').collect(); if segments.len() == 1 { if let Some(&seg) = segments.first() { - btree.remove(&Value::from(seg)); + obj.remove(&Value::from(seg)); } } else if segments.len() > 1 { - remove_at_dotted_path_in_btree(btree, &segments); + remove_at_dotted_path_obj(obj, &segments); } return; }; @@ -389,12 +379,12 @@ fn remove_field_at_depth_in_btree( let key_val = Value::from(first); let arr_val = if nav.len() == 1 { - match btree.get_mut(&key_val) { + match obj.get_mut(&key_val) { Some(v) => v, None => return, } } else { - let mut cur: &mut Value = match btree.get_mut(&key_val) { + let mut cur: &mut Value = match obj.get_mut(&key_val) { Some(v) => v, None => return, }; @@ -415,27 +405,19 @@ fn remove_field_at_depth_in_btree( for elem in inner.iter_mut() { if let Value::Object(obj_rc) = elem { let inner_btree = Rc::make_mut(obj_rc); - remove_field_at_depth_in_btree( - inner_btree, - array_chain, - depth.saturating_add(1), - field, - ); + remove_field_at_depth_obj(inner_btree, array_chain, depth.saturating_add(1), field); } } } } -/// Remove the leaf segment at a dotted path directly in a BTreeMap. -fn remove_at_dotted_path_in_btree( - btree: &mut alloc::collections::BTreeMap, - segments: &[&str], -) { +/// Remove the leaf segment at a dotted path directly in an Object. +fn remove_at_dotted_path_obj(obj: &mut Object, segments: &[&str]) { let Some((&leaf, parent_segs)) = segments.split_last() else { return; }; if parent_segs.is_empty() { - btree.remove(&Value::from(leaf)); + obj.remove(&Value::from(leaf)); return; } @@ -443,7 +425,7 @@ fn remove_at_dotted_path_in_btree( return; }; let first_key = Value::from(first); - let parent_val = match btree.get_mut(&first_key) { + let parent_val = match obj.get_mut(&first_key) { Some(v) => v, None => return, }; diff --git a/src/languages/azure_policy/compiler/effects.rs b/src/languages/azure_policy/compiler/effects.rs index afad39d..9987897 100644 --- a/src/languages/azure_policy/compiler/effects.rs +++ b/src/languages/azure_policy/compiler/effects.rs @@ -11,7 +11,7 @@ //! to fetch a related resource and an optional `existenceCondition` evaluated //! inline. -use alloc::collections::BTreeMap; +use crate::value::Object; use alloc::format; use alloc::string::ToString as _; use alloc::vec::Vec; @@ -814,7 +814,7 @@ pub(super) fn build_object_from_keys( span: &crate::lexer::Span, ) -> Result { // Build template: object with all keys set to Undefined. - let mut template = BTreeMap::new(); + let mut template = Object::new(); for &(key_idx, _) in &keys { // key_idx was returned by `add_literal_u16` in the calling code, // so it is always in bounds. We use `.get()` + `?` instead of diff --git a/src/languages/azure_policy/compiler/metadata.rs b/src/languages/azure_policy/compiler/metadata.rs index 6ba3066..5226539 100644 --- a/src/languages/azure_policy/compiler/metadata.rs +++ b/src/languages/azure_policy/compiler/metadata.rs @@ -272,7 +272,7 @@ impl Compiler { fn insert_string_set_annotation( annot: &mut alloc::collections::BTreeMap, key: &str, - observed: &BTreeSet, + observed: &alloc::collections::BTreeSet, ) { if !observed.is_empty() { let set: BTreeSet = observed diff --git a/src/languages/rego/compiler/expressions/collection_literals.rs b/src/languages/rego/compiler/expressions/collection_literals.rs index ac9ea91..2f3b39f 100644 --- a/src/languages/rego/compiler/expressions/collection_literals.rs +++ b/src/languages/rego/compiler/expressions/collection_literals.rs @@ -11,8 +11,9 @@ use crate::ast::{Expr, ExprRef}; use crate::lexer::Span; use crate::rvm::instructions::{ArrayCreateParams, ObjectCreateParams, SetCreateParams}; use crate::rvm::Instruction; +use crate::value::Object; use crate::{Rc, Value}; -use alloc::collections::{BTreeMap, BTreeSet}; +use alloc::collections::BTreeSet; use alloc::vec::Vec; /// Try to evaluate an expression as a compile-time constant. @@ -43,7 +44,7 @@ pub(in crate::languages::rego::compiler) fn try_eval_const(expr: &Expr) -> Optio Expr::Object { fields, .. } => fields .iter() .map(|(_, k, v)| Some((try_eval_const(k.as_ref())?, try_eval_const(v.as_ref())?))) - .collect::>>() + .collect::>() .map(|m| Value::Object(Rc::new(m))), _ => None, } @@ -117,7 +118,7 @@ impl<'a> Compiler<'a> { fields: &[(crate::lexer::Span, ExprRef, ExprRef)], span: &Span, ) -> Result { - let all_const: Option> = fields + let all_const: Option = fields .iter() .map(|(_, k, v)| Some((try_eval_const(k.as_ref())?, try_eval_const(v.as_ref())?))) .collect(); @@ -166,7 +167,7 @@ impl<'a> Compiler<'a> { let mut template_keys = literal_keys.clone(); template_keys.sort(); - let mut template_obj = BTreeMap::new(); + let mut template_obj = Object::new(); for key in &template_keys { template_obj.insert(key.clone(), Value::Undefined); } diff --git a/src/rvm/program/metadata.rs b/src/rvm/program/metadata.rs index 539d5f8..ec783e3 100644 --- a/src/rvm/program/metadata.rs +++ b/src/rvm/program/metadata.rs @@ -11,9 +11,9 @@ //! values are converted through [`MetadataValue`] — a postcard/bincode-safe //! enum that avoids `deserialize_any`. +use crate::value::Object; use crate::Rc; -use alloc::collections::BTreeMap; -use alloc::collections::BTreeSet; +use alloc::collections::{BTreeMap, BTreeSet}; use alloc::string::String; use alloc::vec::Vec; use serde::{Deserialize, Serialize}; @@ -52,7 +52,7 @@ impl ProgramMetadata { pub fn to_value(&self) -> crate::value::Value { use crate::value::Value; - let mut obj = BTreeMap::new(); + let mut obj = Object::new(); obj.insert( Value::String("compiler_version".into()), Value::String(self.compiler_version.as_str().into()), @@ -75,7 +75,7 @@ impl ProgramMetadata { ); if !self.annotations.is_empty() { - let mut annotations_obj = BTreeMap::new(); + let mut annotations_obj = Object::new(); for (k, v) in &self.annotations { annotations_obj.insert(Value::String(k.as_str().into()), v.clone()); } @@ -198,7 +198,7 @@ impl MetadataValue { match *self { MetadataValue::String(ref s) => Value::String(s.as_str().into()), MetadataValue::StringSet(ref set) => { - let mut bset = alloc::collections::BTreeSet::new(); + let mut bset = BTreeSet::new(); for s in set { bset.insert(Value::String(s.as_str().into())); } @@ -211,7 +211,7 @@ impl MetadataValue { Value::Array(Rc::new(values)) } MetadataValue::Map(ref map) => { - let mut obj = BTreeMap::new(); + let mut obj = Object::new(); for (k, v) in map { obj.insert(Value::String(k.as_str().into()), v.to_value()); } @@ -257,7 +257,6 @@ mod metadata_serde { mod tests { use super::*; use crate::value::Value; - use alloc::collections::BTreeSet; /// Round-trip: Value → MetadataValue → Value must be equivalent for /// all lossless variants (strings, bools, integers, arrays, objects). diff --git a/src/rvm/program/serialization/value.rs b/src/rvm/program/serialization/value.rs index d0bacee..0e98a17 100644 --- a/src/rvm/program/serialization/value.rs +++ b/src/rvm/program/serialization/value.rs @@ -1,6 +1,6 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -use alloc::collections::{BTreeMap, BTreeSet}; +use alloc::collections::BTreeSet; use alloc::format; use alloc::string::String; use alloc::vec::Vec; @@ -11,6 +11,7 @@ use serde::ser::{SerializeSeq as _, SerializeTuple as _}; use serde::{Deserialize, Serialize}; use crate::number::Number; +use crate::value::Object; use crate::value::Value; const VARIANT_NULL: u32 = 0; @@ -132,7 +133,7 @@ impl<'a> Serialize for BinarySetRef<'a> { } } -struct BinaryObjectRef<'a>(&'a BTreeMap); +struct BinaryObjectRef<'a>(&'a Object); impl<'a> Serialize for BinaryObjectRef<'a> { fn serialize(&self, serializer: S) -> Result @@ -140,7 +141,7 @@ impl<'a> Serialize for BinaryObjectRef<'a> { S: serde::Serializer, { let mut seq = serializer.serialize_seq(Some(self.0.len()))?; - for (key, value) in self.0.iter() { + for (key, value) in self.0.iter_sorted() { seq.serialize_element(&BinaryEntryRef(key, value))?; } seq.end() @@ -261,11 +262,11 @@ impl<'de> Visitor<'de> for BinaryValueVisitor { } (BinaryVariant::Object, variant) => { let entries: Vec<(BinaryValue, BinaryValue)> = variant.newtype_variant()?; - let mut map = BTreeMap::new(); + let mut map = Object::new(); for (key, value) in entries { map.insert(key.into_value(), value.into_value()); } - Ok(BinaryValue(Value::from(map))) + Ok(BinaryValue(Value::Object(crate::Rc::new(map)))) } (BinaryVariant::Undefined, variant) => { variant.unit_variant()?; diff --git a/src/rvm/vm/arithmetic.rs b/src/rvm/vm/arithmetic.rs index 3d40488..89362d1 100644 --- a/src/rvm/vm/arithmetic.rs +++ b/src/rvm/vm/arithmetic.rs @@ -6,8 +6,6 @@ // Disable both to keep patterns consistent within this file. #![allow(clippy::pattern_type_mismatch, clippy::needless_borrowed_reference)] -use alloc::collections::BTreeSet; - use crate::number::Number; use crate::value::Value; @@ -32,8 +30,9 @@ impl RegoVM { match (a, b) { (&Value::Number(ref x), &Value::Number(ref y)) => Ok(Value::from(x.sub(y)?)), (&Value::Set(ref left), &Value::Set(ref right)) => { - let diff: BTreeSet = left.difference(right).cloned().collect(); - Ok(Value::from_set(diff)) + let diff: alloc::collections::BTreeSet = + left.difference(right).cloned().collect(); + Ok(Value::from(diff)) } _ => Err(VmError::InvalidSubtraction { left: a.clone(), diff --git a/src/rvm/vm/comprehension.rs b/src/rvm/vm/comprehension.rs index b4ed70a..aafa52b 100644 --- a/src/rvm/vm/comprehension.rs +++ b/src/rvm/vm/comprehension.rs @@ -2,9 +2,9 @@ // Licensed under the MIT License. use crate::rvm::instructions::{ComprehensionBeginParams, ComprehensionMode}; +use crate::value::Object; use crate::value::Value; use crate::Rc; -use alloc::collections::BTreeMap; use alloc::format; use alloc::sync::Arc; use alloc::vec::Vec; @@ -34,12 +34,12 @@ impl RegoVM { let initial_result = match params.mode { ComprehensionMode::Set => Value::new_set(), ComprehensionMode::Array => Value::new_array(), - ComprehensionMode::Object => Value::Object(Rc::new(BTreeMap::new())), + ComprehensionMode::Object => Value::Object(Rc::new(Object::new())), }; self.set_register(params.result_reg, initial_result.clone())?; let auto_iterate = params.collection_reg != params.result_reg; - let iteration_state = if auto_iterate { + let mut iteration_state = if auto_iterate { let source_value = self.get_register(params.collection_reg)?.clone(); match source_value { Value::Array(items) => { @@ -53,11 +53,9 @@ impl RegoVM { if obj.is_empty() { None } else { - Some(IterationState::Object { - obj, - current_key: None, - first_iteration: true, - }) + // O(1) cursor over shared Rc. + let cursor = obj.cursor(); + Some(IterationState::Object { obj, cursor }) } } Value::Set(set) => { @@ -79,7 +77,7 @@ impl RegoVM { None }; - let has_iteration = if let Some(state) = iteration_state.as_ref() { + let has_iteration = if let Some(state) = iteration_state.as_mut() { self.setup_next_iteration(state, params.key_reg, params.value_reg)? } else { false @@ -123,12 +121,12 @@ impl RegoVM { let initial_result = match params.mode { ComprehensionMode::Set => Value::new_set(), ComprehensionMode::Array => Value::new_array(), - ComprehensionMode::Object => Value::Object(Rc::new(BTreeMap::new())), + ComprehensionMode::Object => Value::Object(Rc::new(Object::new())), }; self.set_register(params.result_reg, initial_result.clone())?; let auto_iterate = params.collection_reg != params.result_reg; - let iteration_state = if auto_iterate { + let mut iteration_state = if auto_iterate { let source_value = self.get_register(params.collection_reg)?.clone(); match source_value { Value::Array(items) => { @@ -142,11 +140,8 @@ impl RegoVM { if obj.is_empty() { None } else { - Some(IterationState::Object { - obj, - current_key: None, - first_iteration: true, - }) + let cursor = obj.cursor(); + Some(IterationState::Object { obj, cursor }) } } Value::Set(set) => { @@ -168,7 +163,7 @@ impl RegoVM { None }; - let has_iteration = if let Some(state) = iteration_state.as_ref() { + let has_iteration = if let Some(state) = iteration_state.as_mut() { self.setup_next_iteration(state, params.key_reg, params.value_reg)? } else { false @@ -255,6 +250,21 @@ impl RegoVM { }; let result_reg = comprehension_context.result_reg; + // Snapshot the iteration value register BEFORE taking the result + // register: if the comprehension compiler ever allocates + // `result_reg == context.value_reg`, the writeback at the bottom + // of this function would clobber the value register, and a + // post-writeback read here would feed the wrong value into + // `IterationState::Set::current_item`. Only Set needs the snapshot + // (Object uses a self-advancing cursor; Array advances by index). + let set_resume_snapshot = if matches!( + comprehension_context.iteration_state, + Some(IterationState::Set { .. }) + ) { + Some(self.get_register(comprehension_context.value_reg)?.clone()) + } else { + None + }; // Take ownership of the result register so Rc refcount stays at 1, // allowing Rc::make_mut to mutate in-place instead of deep-cloning. let mut current_result = self.take_register(result_reg)?; @@ -292,29 +302,16 @@ impl RegoVM { self.set_register(result_reg, current_result)?; if let Some(iter_state) = comprehension_context.iteration_state.as_mut() { - match *iter_state { - IterationState::Object { - ref mut current_key, - .. - } => { - let tracked_key = - if comprehension_context.key_reg != comprehension_context.value_reg { - self.get_register(comprehension_context.key_reg)?.clone() - } else { - self.get_register(comprehension_context.value_reg)?.clone() - }; - *current_key = Some(tracked_key); - } - IterationState::Set { - ref mut current_item, - .. - } => { - *current_item = - Some(self.get_register(comprehension_context.value_reg)?.clone()); - } - IterationState::Array { .. } | IterationState::Single { .. } => {} + // Set's `Bound::Excluded(current_item)` resume scheme needs the + // pre-mutation snapshot taken at the top of this function. + // Object uses a self-advancing cursor and needs no snapshot. + if let IterationState::Set { + ref mut current_item, + .. + } = *iter_state + { + *current_item = set_resume_snapshot; } - iter_state.advance(); let has_next = self.setup_next_iteration( iter_state, @@ -359,8 +356,7 @@ impl RegoVM { result_reg_idx, key_reg_idx, value_reg_idx, - iteration_key, - iteration_value, + iter_is_set, ) = { let frame = self.execution_stack @@ -382,8 +378,8 @@ impl RegoVM { let result_reg_idx = context.result_reg; let mode = context.mode.clone(); - let iteration_key = self.get_register(context.key_reg)?.clone(); - let iteration_value = self.get_register(context.value_reg)?.clone(); + let iter_is_set = + matches!(context.iteration_state, Some(IterationState::Set { .. })); ( value_to_add, @@ -392,8 +388,7 @@ impl RegoVM { result_reg_idx, context.key_reg, context.value_reg, - iteration_key, - iteration_value, + iter_is_set, ) } else { return Err(VmError::InvalidIteration { @@ -403,6 +398,18 @@ impl RegoVM { } }; + // Snapshot the iteration value register BEFORE the result writeback: + // if the compiler ever allocates `result_reg == value_reg_idx`, a + // post-writeback read would feed the result accumulator into + // `IterationState::Set::current_item`, breaking the next iteration. + // Only Set needs this (Object cursor self-advances; Array advances + // by index). + let set_resume_snapshot = if iter_is_set { + Some(self.get_register(value_reg_idx)?.clone()) + } else { + None + }; + // Take ownership of the result register so Rc refcount stays at 1, // allowing Rc::make_mut to mutate in-place instead of deep-cloning. let mut current_result = self.take_register(result_reg_idx)?; @@ -450,27 +457,13 @@ impl RegoVM { } = &mut frame.kind { if let Some(iter_state) = context.iteration_state.as_mut() { - match *iter_state { - IterationState::Object { - ref mut current_key, - .. - } => { - let tracked_key = if context.key_reg != context.value_reg { - iteration_key.clone() - } else { - iteration_value.clone() - }; - *current_key = Some(tracked_key); - } - IterationState::Set { - ref mut current_item, - .. - } => { - *current_item = Some(iteration_value.clone()); - } - IterationState::Array { .. } | IterationState::Single { .. } => {} + if let IterationState::Set { + ref mut current_item, + .. + } = *iter_state + { + *current_item = set_resume_snapshot; } - iter_state.advance(); } @@ -487,8 +480,21 @@ impl RegoVM { } }; - if let Some(state) = iteration_state_snapshot.as_ref() { - let has_next = self.setup_next_iteration(state, key_reg_idx, value_reg_idx)?; + if let Some(mut state) = iteration_state_snapshot { + let has_next = self.setup_next_iteration(&mut state, key_reg_idx, value_reg_idx)?; + + // `setup_next_iteration` advances Object's internal cursor; the + // owning frame holds the iteration_state, so we must write the + // updated state back. (The Array/Set variants are also unchanged + // by copy, so the writeback is uniform.) + if let Some(frame) = self.execution_stack.get_mut(comprehension_index) { + if let FrameKind::Comprehension { + ref mut context, .. + } = frame.kind + { + context.iteration_state = Some(state); + } + } if has_next { if let Some(frame) = self.execution_stack.get_mut(comprehension_index) { @@ -553,11 +559,16 @@ impl RegoVM { context: &mut ComprehensionContext, ) -> Result<()> { if let Some(iter_state) = context.iteration_state.as_mut() { - self.capture_comprehension_iteration_position( - iter_state, - context.key_reg, - context.value_reg, - )?; + // Snapshot the current value into Set's `current_item` so the + // next iteration can resume from `Bound::Excluded(current)`. + // Object uses a self-advancing cursor and needs no snapshot here. + if let IterationState::Set { + ref mut current_item, + .. + } = *iter_state + { + *current_item = Some(self.get_register(context.value_reg)?.clone()); + } iter_state.advance(); let has_next = self.setup_next_iteration(iter_state, context.key_reg, context.value_reg)?; @@ -574,36 +585,6 @@ impl RegoVM { Ok(()) } - fn capture_comprehension_iteration_position( - &mut self, - iter_state: &mut IterationState, - key_reg: u8, - value_reg: u8, - ) -> Result<()> { - match *iter_state { - IterationState::Object { - ref mut current_key, - .. - } => { - let tracked_key = if key_reg != value_reg { - self.get_register(key_reg)?.clone() - } else { - self.get_register(value_reg)?.clone() - }; - *current_key = Some(tracked_key); - } - IterationState::Set { - ref mut current_item, - .. - } => { - *current_item = Some(self.get_register(value_reg)?.clone()); - } - IterationState::Array { .. } | IterationState::Single { .. } => {} - } - - Ok(()) - } - fn execute_comprehension_end_run_to_completion(&mut self) -> Result<()> { // `ComprehensionEnd` is reached from a loaded program; an empty stack // here means malformed user-supplied bytecode, which must still surface diff --git a/src/rvm/vm/context.rs b/src/rvm/vm/context.rs index 7b5a714..56439b9 100644 --- a/src/rvm/vm/context.rs +++ b/src/rvm/vm/context.rs @@ -3,8 +3,9 @@ use crate::rvm::instructions::{ComprehensionMode, LoopMode}; use crate::value::Value; +use crate::value::{Object, ObjectCursor}; use crate::Rc; -use alloc::collections::{BTreeMap, BTreeSet}; +use alloc::collections::BTreeSet; use alloc::vec::Vec; /// Loop execution context for managing iteration state @@ -24,7 +25,18 @@ pub struct LoopContext { pub current_iteration_failed: bool, // Track if current iteration had condition failures } -/// Iterator state for different collection types +/// Iterator state for different collection types. +/// +/// Snapshot independence for `Object` is provided by the shared +/// `Rc` — `Rc::make_mut` on an aliased Rc allocates a new +/// collection, leaving the iterator's Rc pointing at the original +/// pre-mutation state. The `ObjectCursor` is opaque and resumes in +/// O(log n) for the BTree backend. +/// +/// `Set` continues to use the pre-existing snapshot-by-cloned-key +/// approach (`current_item` + `first_iteration`); migration of `Set` +/// to a cursor-based iterator ships with the `Set` storage abstraction +/// in a follow-up PR. #[derive(Debug, Clone)] pub enum IterationState { Array { @@ -32,9 +44,8 @@ pub enum IterationState { index: usize, }, Object { - obj: Rc>, - current_key: Option, - first_iteration: bool, + obj: Rc, + cursor: ObjectCursor, }, Set { items: Rc>, @@ -64,11 +75,11 @@ impl IterationState { ); *index = index.saturating_add(1); } - Self::Object { - ref mut first_iteration, - .. - } - | Self::Set { + // For Object the cursor advances inside `setup_next_iteration` + // when it pulls the next item via `Object::next`, so `advance` + // is a no-op for the cursor-backed Object variant. + Self::Object { .. } => {} + Self::Set { ref mut first_iteration, .. } => { @@ -121,3 +132,71 @@ pub(super) struct ComprehensionContext { /// Resume location for the parent frame once this comprehension completes pub(super) resume_pc: usize, } + +#[cfg(test)] +#[allow( + clippy::expect_used, + clippy::unwrap_used, + clippy::unreachable, + clippy::pattern_type_mismatch, + clippy::shadow_unrelated, + clippy::panic +)] +mod tests { + use super::*; + use crate::value::Object; + + /// IterationState::Object holds an `Rc` plus an opaque cursor. + /// Mutating an aliased Rc via `Rc::make_mut` allocates a new collection + /// (CoW) so the in-flight iterator's source is unaffected. + #[test] + fn iteration_state_object_is_snapshot_independent_of_source() { + let mut obj = Object::new(); + obj.insert(Value::from("a"), Value::from(1)); + obj.insert(Value::from("b"), Value::from(2)); + obj.insert(Value::from("c"), Value::from(3)); + + let source = Value::Object(Rc::new(obj)); + + let snapshot_obj = match &source { + Value::Object(o) => Rc::clone(o), + _ => unreachable!(), + }; + let state = IterationState::Object { + obj: Rc::clone(&snapshot_obj), + cursor: snapshot_obj.cursor(), + }; + + // Mutate a clone of the source mid-iteration. + let mut alias = source.clone(); + let inner = alias.as_object_mut().expect("object"); + inner.insert(Value::from("a"), Value::from(999)); + inner.insert(Value::from("d"), Value::from(4)); + inner.remove(&Value::from("b")); + + // Drain the snapshot via the cursor — must still report the original + // 3 entries with original values. + let mut collected: Vec<(Value, Value)> = Vec::new(); + if let IterationState::Object { + ref obj, + mut cursor, + } = state + { + while let Some((k, v)) = obj.next(&mut cursor) { + collected.push((k.clone(), v.clone())); + } + } else { + unreachable!(); + } + assert_eq!(collected.len(), 3); + assert!(collected.contains(&(Value::from("a"), Value::from(1)))); + assert!(collected.contains(&(Value::from("b"), Value::from(2)))); + assert!(collected.contains(&(Value::from("c"), Value::from(3)))); + assert!(!collected.iter().any(|kv| kv.0 == Value::from("d"))); + + // The original source Value (untouched) is also unchanged. + let src_obj = source.as_object().expect("object"); + assert_eq!(src_obj.len(), 3); + assert_eq!(src_obj.get(&Value::from("a")), Some(&Value::from(1))); + } +} diff --git a/src/rvm/vm/dispatch.rs b/src/rvm/vm/dispatch.rs index d0cb044..3a9379c 100644 --- a/src/rvm/vm/dispatch.rs +++ b/src/rvm/vm/dispatch.rs @@ -4,7 +4,6 @@ use crate::rvm::instructions::{GuardMode, Instruction, LiteralOrRegister}; use crate::rvm::program::Program; use crate::value::Value; -use alloc::collections::BTreeSet; use alloc::vec::Vec; use core::mem; @@ -670,7 +669,7 @@ impl RegoVM { } } SetNew { dest } => { - let empty_set = Value::Set(crate::Rc::new(BTreeSet::new())); + let empty_set = Value::new_set(); self.set_register(dest, empty_set)?; Ok(InstructionOutcome::Continue) } @@ -707,7 +706,7 @@ impl RegoVM { if any_undefined { self.set_register(params.dest, Value::Undefined)?; } else { - let mut set = BTreeSet::new(); + let mut set = alloc::collections::BTreeSet::new(); for ® in params.element_registers() { set.insert(self.get_register(reg)?.clone()); } diff --git a/src/rvm/vm/loops.rs b/src/rvm/vm/loops.rs index 57458a7..f6ecf11 100644 --- a/src/rvm/vm/loops.rs +++ b/src/rvm/vm/loops.rs @@ -3,6 +3,7 @@ use crate::rvm::instructions::LoopMode; use crate::value::Value; +use crate::Rc; use super::context::{IterationState, LoopContext}; use super::errors::{Result, VmError}; @@ -89,13 +90,13 @@ impl RegoVM { ) -> Result<()> { self.set_register(params.result_reg, Value::Bool(false))?; - let iteration_state = match self.resolve_iteration_state(mode, ¶ms)? { + let mut iteration_state = match self.resolve_iteration_state(mode, ¶ms)? { Some(state) => state, None => return Ok(()), }; let has_next = - self.setup_next_iteration(&iteration_state, params.key_reg, params.value_reg)?; + self.setup_next_iteration(&mut iteration_state, params.key_reg, params.value_reg)?; if !has_next { self.pc = usize::from(params.loop_end); return Ok(()); @@ -155,15 +156,10 @@ impl RegoVM { LoopAction::Continue => {} } - if let &mut IterationState::Object { - ref mut current_key, - .. - } = &mut loop_ctx.iteration_state - { - if loop_ctx.key_reg != loop_ctx.value_reg { - *current_key = Some(self.get_register(loop_ctx.key_reg)?.clone()); - } - } else if let &mut IterationState::Set { + // Snapshot the current value for Set so its next iteration can resume + // from `Bound::Excluded(current)`. Object uses a cursor and advances + // inside `setup_next_iteration` itself. + if let &mut IterationState::Set { ref mut current_item, .. } = &mut loop_ctx.iteration_state @@ -173,7 +169,7 @@ impl RegoVM { loop_ctx.iteration_state.advance(); let has_next = self.setup_next_iteration( - &loop_ctx.iteration_state, + &mut loop_ctx.iteration_state, loop_ctx.key_reg, loop_ctx.value_reg, )?; @@ -211,13 +207,13 @@ impl RegoVM { ) -> Result<()> { self.set_register(params.result_reg, Value::Bool(false))?; - let iteration_state = match self.resolve_iteration_state(mode, ¶ms)? { + let mut iteration_state = match self.resolve_iteration_state(mode, ¶ms)? { Some(state) => state, None => return Ok(()), }; let has_next = - self.setup_next_iteration(&iteration_state, params.key_reg, params.value_reg)?; + self.setup_next_iteration(&mut iteration_state, params.key_reg, params.value_reg)?; if !has_next { self.pc = usize::from(params.loop_end); return Ok(()); @@ -316,7 +312,14 @@ impl RegoVM { Ok(()) } LoopAction::Continue => { - let (mode, success_count, total_iterations, key_reg, value_reg, iteration_state) = { + let ( + mode, + success_count, + total_iterations, + key_reg, + value_reg, + mut iteration_state, + ) = { let (mode, success_count, total_iterations, key_reg, value_reg) = { let frame = self .execution_stack @@ -334,11 +337,6 @@ impl RegoVM { } }; - let key_value = if key_reg != value_reg { - Some(self.get_register(key_reg)?.clone()) - } else { - None - }; let value_value = self.get_register(value_reg)?.clone(); let frame = self @@ -349,20 +347,16 @@ impl RegoVM { &mut FrameKind::Loop { ref mut context, .. } => { - if let &mut IterationState::Object { - ref mut current_key, - .. - } = &mut context.iteration_state - { - if context.key_reg != context.value_reg { - *current_key = key_value; - } - } else if let &mut IterationState::Set { + // Snapshot the current value for Set so its next + // iteration can resume from `Bound::Excluded(current)`. + // Object uses a cursor and advances inside + // `setup_next_iteration` itself. + if let &mut IterationState::Set { ref mut current_item, .. } = &mut context.iteration_state { - *current_item = Some(value_value.clone()); + *current_item = Some(value_value); } context.iteration_state.advance(); @@ -381,7 +375,21 @@ impl RegoVM { } }; - let has_next = self.setup_next_iteration(&iteration_state, key_reg, value_reg)?; + let has_next = + self.setup_next_iteration(&mut iteration_state, key_reg, value_reg)?; + + // `setup_next_iteration` advances Object's internal cursor; + // the owning frame holds the iteration_state, so we must + // write the updated state back. (Array/Set are unchanged by + // the call, so the writeback is uniform.) + if let Some(frame) = self.execution_stack.last_mut() { + if let FrameKind::Loop { + ref mut context, .. + } = frame.kind + { + context.iteration_state = iteration_state; + } + } if has_next { if let Some(frame) = self.execution_stack.last_mut() { @@ -459,10 +467,14 @@ impl RegoVM { self.handle_empty_collection(mode, params.result_reg, params.loop_end)?; return Ok(None); } + // O(1) resumable cursor over the shared Rc. + // No eager pair snapshot: avoids O(N) setup, O(N) memory + // floor, and O(N) memory-limit checks. Snapshot + // independence is via the shared Rc (CoW). + let cursor = obj.cursor(); Ok(Some(IterationState::Object { - obj: obj.clone(), - current_key: None, - first_iteration: true, + obj: Rc::clone(obj), + cursor, })) } } @@ -512,7 +524,7 @@ impl RegoVM { pub(super) fn setup_next_iteration( &mut self, - state: &IterationState, + state: &mut IterationState, key_reg: u8, value_reg: u8, ) -> Result { @@ -538,33 +550,19 @@ impl RegoVM { } IterationState::Object { ref obj, - ref current_key, - ref first_iteration, + ref mut cursor, } => { - if *first_iteration { - if let Some((key, value)) = obj.iter().next() { - if key_reg != value_reg { - self.set_register(key_reg, key.clone())?; - } - self.set_register(value_reg, value.clone())?; - Ok(true) - } else { - Ok(false) - } - } else if let Some(ref current) = *current_key { - let mut range_iter = obj.range(( - core::ops::Bound::Excluded(current), - core::ops::Bound::Unbounded, - )); - if let Some((key, value)) = range_iter.next() { - if key_reg != value_reg { - self.set_register(key_reg, key.clone())?; - } - self.set_register(value_reg, value.clone())?; - Ok(true) - } else { - Ok(false) + // Object iterates via a resumable cursor on the shared + // `Rc`; `next` both yields the current entry and + // advances the cursor. No explicit `current_key` snapshot is + // needed — see the doc on `IterationState`. + if let Some((key, value)) = obj.next(cursor) { + let value = value.clone(); + if key_reg != value_reg { + self.set_register(key_reg, key.clone())?; } + self.set_register(value_reg, value)?; + Ok(true) } else { Ok(false) } diff --git a/src/schema/tests/suite.rs b/src/schema/tests/suite.rs index 40cc551..4ba2267 100644 --- a/src/schema/tests/suite.rs +++ b/src/schema/tests/suite.rs @@ -1498,7 +1498,7 @@ fn test_deserialize_object_default_empty_object() { let s = Schema::from_serde_json_value(schema).unwrap(); match s.as_type() { Type::Object { default, .. } => { - assert_eq!(default, &Some(Value::Object(Rc::new(BTreeMap::new())))); + assert_eq!(default, &Some(Value::new_object())); } _ => panic!("Expected Type::Object"), } @@ -1778,8 +1778,11 @@ fn test_deserialize_enum_values_with_object_non_string_keys() { match s.as_type() { Type::Enum { values, .. } => match &values[0] { Value::Object(obj) => { - assert_eq!(obj[&Value::from("1")], Value::from("one")); - assert_eq!(obj[&Value::from("true")], Value::from("bool")); + assert_eq!(*obj.get(&Value::from("1")).expect("1"), Value::from("one")); + assert_eq!( + *obj.get(&Value::from("true")).expect("true"), + Value::from("bool") + ); } _ => panic!("Expected object in enum values"), }, @@ -1802,18 +1805,21 @@ fn test_deserialize_enum_values_with_deeply_nested_structures() { match s.as_type() { Type::Enum { values, .. } => match &values[0] { Value::Object(obj) => { - let a = &obj[&Value::from("a")]; + let a = obj.get(&Value::from("a")).expect("a"); match a { Value::Array(arr) => match &arr[0] { Value::Object(inner) => { - let b = &inner[&Value::from("b")]; + let b = inner.get(&Value::from("b")).expect("b"); match b { Value::Array(barr) => { assert_eq!(barr[0], Value::from(1)); assert_eq!(barr[1], Value::from(2)); match &barr[2] { Value::Object(cobj) => { - assert_eq!(cobj[&Value::from("c")], Value::Null); + assert_eq!( + *cobj.get(&Value::from("c")).expect("c"), + Value::Null + ); } _ => panic!("Expected object for 'c'"), } @@ -1886,8 +1892,11 @@ fn test_deserialize_const_value_object() { match s.as_type() { Type::Const { value, .. } => match value { Value::Object(ref obj) => { - assert_eq!(obj[&Value::from("foo")], Value::from("bar")); - assert_eq!(obj[&Value::from("baz")], Value::from(1)); + assert_eq!( + *obj.get(&Value::from("foo")).expect("foo"), + Value::from("bar") + ); + assert_eq!(*obj.get(&Value::from("baz")).expect("baz"), Value::from(1)); } _ => panic!("Expected object for const value"), }, @@ -1940,13 +1949,13 @@ fn test_deserialize_const_value_deeply_nested() { match s.as_type() { Type::Const { value, .. } => match value { Value::Object(ref obj) => { - let a = &obj[&Value::from("a")]; + let a = obj.get(&Value::from("a")).expect("a"); match a { Value::Array(arr) => { assert_eq!(arr[0], Value::from(1)); match &arr[1] { Value::Object(inner) => { - let b = &inner[&Value::from("b")]; + let b = inner.get(&Value::from("b")).expect("b"); match b { Value::Array(barr) => { assert_eq!(barr[0], Value::Null); diff --git a/src/schema/validate.rs b/src/schema/validate.rs index 9f3391c..2e99319 100644 --- a/src/schema/validate.rs +++ b/src/schema/validate.rs @@ -7,6 +7,7 @@ use crate::{ schema::{error::ValidationError, Schema, Type}, + value::Object, *, }; use alloc::collections::BTreeMap; @@ -537,7 +538,7 @@ impl SchemaValidator { } fn validate_discriminated_subobject_with_base( - object_value: &BTreeMap, + object_value: &Object, discriminated_subobject: &crate::schema::DiscriminatedSubobject, base_properties: &BTreeMap, base_additional_properties: Option<&Schema>, @@ -653,7 +654,7 @@ impl SchemaValidator { } fn validate_subobject( - object_value: &BTreeMap, + object_value: &Object, subobject: &crate::schema::Subobject, path: &str, ) -> Result<(), ValidationError> { diff --git a/src/value/mod.rs b/src/value/mod.rs index 678309c..7886ddc 100644 --- a/src/value/mod.rs +++ b/src/value/mod.rs @@ -35,7 +35,7 @@ use core::str::FromStr; use anyhow::{anyhow, bail, Result}; use serde::de::{self, Deserializer, Error as DeError, MapAccess, SeqAccess, Visitor}; -use serde::ser::{SerializeMap, Serializer}; +use serde::ser::Serializer; use serde::{Deserialize, Serialize}; use crate::*; @@ -75,7 +75,7 @@ pub enum Value { /// An object. /// Unlike JSON, keys can be any value, not just string. - Object(Rc>), + Object(Rc), /// Undefined value. /// Used to indicate the absence of a value. @@ -98,26 +98,15 @@ impl Serialize for Value { where S: Serializer, { - use serde::ser::Error; match self { Value::Null => serializer.serialize_unit(), Value::Bool(b) => serializer.serialize_bool(*b), Value::String(s) => serializer.serialize_str(s.as_ref()), Value::Number(n) => n.serialize(serializer), Value::Array(a) => a.serialize(serializer), - Value::Object(fields) => { - let mut map = serializer.serialize_map(Some(fields.len()))?; - for (k, v) in fields.iter() { - match k { - Value::String(_) => map.serialize_entry(k, v)?, - _ => { - let key_str = serde_json::to_string(k).map_err(Error::custom)?; - map.serialize_entry(&key_str, v)? - } - } - } - map.end() - } + // Delegate to the Object/Set serializers — single canonical path, + // handles non-string-key stringification internally. + Value::Object(fields) => fields.serialize(serializer), // display set as an array Value::Set(s) => s.serialize(serializer), @@ -357,7 +346,7 @@ impl Value { /// assert_eq!(array[4], Value::from(12345u64)); /// let obj = array[5].as_object().expect("not an object"); /// assert_eq!(obj.len(), 1); - /// assert_eq!(obj[&Value::from("name")], Value::from("regorus")); + /// assert_eq!(obj.get(&Value::from("name")).expect("missing name"), &Value::from("regorus")); /// # Ok(()) /// # } /// ``` @@ -812,7 +801,7 @@ impl From> for Value { /// # Ok(()) /// # } fn from(s: BTreeMap) -> Self { - Value::Object(Rc::new(s)) + Value::Object(Rc::new(Object::from(s))) } } @@ -1291,16 +1280,16 @@ impl Value { } } - /// Cast value to [`& BTreeMap`] if [`Value::Object`]. + /// Cast value to [`&Object`] if [`Value::Object`]. /// ``` /// # use regorus::*; - /// # use std::collections::BTreeMap; + /// # use regorus::value::Object; /// # fn main() -> anyhow::Result<()> { /// let v = Value::from( /// [(Value::from("Hello"), Value::from("World"))] /// .iter() /// .cloned() - /// .collect::>(), + /// .collect::(), /// ); /// assert_eq!( /// v.as_object()?.iter().next(), @@ -1308,28 +1297,28 @@ impl Value { /// ); /// # Ok(()) /// # } - pub fn as_object(&self) -> Result<&BTreeMap> { + pub fn as_object(&self) -> Result<&Object> { match self { Value::Object(m) => Ok(m), _ => Err(anyhow!("not an object")), } } - /// Cast value to [`&mut BTreeMap`] if [`Value::Object`]. + /// Cast value to [`&mut Object`] if [`Value::Object`]. /// ``` /// # use regorus::*; - /// # use std::collections::BTreeMap; + /// # use regorus::value::Object; /// # fn main() -> anyhow::Result<()> { /// let mut v = Value::from( /// [(Value::from("Hello"), Value::from("World"))] /// .iter() /// .cloned() - /// .collect::>(), + /// .collect::(), /// ); /// v.as_object_mut()?.insert(Value::from("Good"), Value::from("Bye")); /// # Ok(()) /// # } - pub fn as_object_mut(&mut self) -> Result<&mut BTreeMap> { + pub fn as_object_mut(&mut self) -> Result<&mut Object> { match self { Value::Object(m) => Ok(Rc::make_mut(m)), _ => Err(anyhow!("not an object")), diff --git a/src/value/object/mod.rs b/src/value/object/mod.rs index 0592124..04078eb 100644 --- a/src/value/object/mod.rs +++ b/src/value/object/mod.rs @@ -153,6 +153,12 @@ impl Object { self.inner.entry(key).or_insert_with(default) } + /// Wrap into a `Value::Object`. + #[inline] + pub fn into_value(self) -> Value { + Value::Object(crate::Rc::new(self)) + } + /// Create a resumable cursor over entries in implementation-defined /// order. Stable for the lifetime of `&self`. O(1). /// @@ -250,3 +256,10 @@ impl From> for Object { Self { inner: map } } } + +impl From for Value { + #[inline] + fn from(o: Object) -> Self { + o.into_value() + } +} diff --git a/tests/azure_policy/mod.rs b/tests/azure_policy/mod.rs index 5b62746..c7141fe 100644 --- a/tests/azure_policy/mod.rs +++ b/tests/azure_policy/mod.rs @@ -544,16 +544,14 @@ fn make_context(case: &TestCase) -> Result { let map = ctx.as_object_mut()?; // Only inject if the caller didn't already provide requestContext // in the context object, to avoid clobbering custom test setups. - map.entry(Value::from("requestContext")).or_insert(rc_val); + map.get_or_insert_with(Value::from("requestContext"), || rc_val); } else if let Some(ref api_ver) = case.api_version { let map = ctx.as_object_mut()?; - if let std::collections::btree_map::Entry::Vacant(e) = - map.entry(Value::from("requestContext")) - { + if !map.contains_key(&Value::from("requestContext")) { let mut req_ctx = Value::new_object(); let rc_map = req_ctx.as_object_mut()?; rc_map.insert(Value::from("apiVersion"), Value::from(api_ver.clone())); - e.insert(req_ctx); + map.insert(Value::from("requestContext"), req_ctx); } } diff --git a/tests/rvm/vm/suites/loops/set_comprehensions.yaml b/tests/rvm/vm/suites/loops/set_comprehensions.yaml index 88eef4e..3392fa2 100644 --- a/tests/rvm/vm/suites/loops/set_comprehensions.yaml +++ b/tests/rvm/vm/suites/loops/set_comprehensions.yaml @@ -137,3 +137,45 @@ cases: - "Return { value: 7 }" # Return result set want_result: set!: [1, null, 2] + + # Set iteration resumes from `Bound::Excluded(current_item)`, so the + # ComprehensionAdd path must snapshot the current value into + # `IterationState::Set.current_item` before advancing — otherwise + # iteration stops after the first element. This case exercises + # Set-source comprehension end-to-end to lock that requirement in. + - note: set_source_visits_all_elements + description: Comprehension over a Set source must yield every element, not just the first + example_rego: | + src := {1, 2, 3, 4} + {x | x := src[_]} # {1, 2, 3, 4} + literals: + - 1 + - 2 + - 3 + - 4 + instruction_params: + comprehension_start_params: + - mode: "Set" + collection_reg: 0 + key_reg: 4 + value_reg: 5 + result_reg: 7 + body_start: 11 + comprehension_end: 13 + instructions: + - "SetNew { dest: 0 }" # Build source set {1,2,3,4} in register 0 + - "Load { dest: 1, literal_idx: 0 }" + - "SetAdd { set: 0, value: 1 }" + - "Load { dest: 2, literal_idx: 1 }" + - "SetAdd { set: 0, value: 2 }" + - "Load { dest: 3, literal_idx: 2 }" + - "SetAdd { set: 0, value: 3 }" + - "Load { dest: 6, literal_idx: 3 }" + - "SetAdd { set: 0, value: 6 }" + - "SetNew { dest: 7 }" # Initialize result set + - "ComprehensionStart { params_index: 0 }" # Iterate over Set source + - "ComprehensionAdd { value_reg: 5 }" # Add current value to result + - "Halt" + - "Return { value: 7 }" + want_result: + set!: [1, 2, 3, 4]