From f69dcc573564cee78eddeb97f3cebd93b1ba3c15 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Mon, 19 Oct 2026 21:18:32 +0000 Subject: [PATCH 01/13] fix(rvm): accumulate keyed partial sets without losing object keys Compile keyed contains rules as partial objects and accumulate each leaf with ObjectSetAdd. Preserve interpreter parity for head-key loops, missing contributions, and empty results. Keep existing instruction discriminants and parameter tables unchanged. Depends on the hoisted output-loop compiler fix. No explanation or partial-evaluation functionality is introduced. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 2 +- src/compiler/hoist.rs | 9 +- src/interpreter.rs | 30 ++- src/languages/rego/compiler/comprehensions.rs | 1 + src/languages/rego/compiler/loops.rs | 1 + src/languages/rego/compiler/mod.rs | 1 + src/languages/rego/compiler/queries.rs | 16 +- src/languages/rego/compiler/rules.rs | 26 +- src/rvm/instructions/display.rs | 3 + src/rvm/instructions/mod.rs | 8 + src/rvm/program/listing.rs | 12 + src/rvm/tests/instruction_parser.rs | 8 + src/rvm/vm/dispatch.rs | 29 +++ tests/rvm/rego/cases/keyed_partial_sets.yaml | 234 ++++++++++++++++++ tests/rvm/vm/suites/keyed_set_operations.yaml | 73 ++++++ 15 files changed, 440 insertions(+), 13 deletions(-) create mode 100644 tests/rvm/rego/cases/keyed_partial_sets.yaml create mode 100644 tests/rvm/vm/suites/keyed_set_operations.yaml diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index bb252a470..85088de23 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -68,6 +68,7 @@ comprehension contexts. Outside loops it aborts the current rule. | Mnemonic | Operands / Params | Behaviour | | :------------------------- | :---------------------------- | :---------------------------------------------------------- | | `ObjectSet` | `obj, key, value` | Mutates object in `obj` with key/value from registers. | +| `ObjectSetAdd` | `obj, key, value` | Adds a defined value to the set at `obj[key]`, creating it if absent. Undefined keys and values are omitted. | | `ObjectCreate` | `params_index=pN` | Builds object from literal template and register entries. | | `ArrayNew` | `dest` | Creates empty array. | | `ArrayPush` | `arr, value` | Appends to array. | @@ -304,4 +305,3 @@ artifacts. Some contain `Vec` fields; the compiler is responsible for ensuring indices remain valid and stable across serialization boundaries. --- - diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index 79eba4b74..d6d385a94 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -482,9 +482,12 @@ impl LoopHoister { }; (key_expr, output_expr) } - RuleHead::Set { key, .. } => { - // For set rules, no separate key_expr, output_expr is the key - (None, key.clone()) + RuleHead::Set { refr, key, .. } => { + let key_expr = match refr.as_ref() { + Expr::RefBrack { index, .. } => Some(index.clone()), + _ => None, + }; + (key_expr, key.clone()) } RuleHead::Func { assign, .. } => { // Function rules return the assignment value (if any) diff --git a/src/interpreter.rs b/src/interpreter.rs index e9e17c587..74c5b3de1 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -2005,7 +2005,9 @@ impl Interpreter { } if output == Value::Undefined || !comps_defined { - ctx_mut.rule_value = Value::Undefined; + if !is_set || key_expr.is_none() { + ctx_mut.rule_value = Value::Undefined; + } return Ok(false); } @@ -3618,11 +3620,24 @@ impl Interpreter { RuleHead::Set { refr, key, .. } => { Parser::get_path_ref_components_into(refr, &mut path)?; let is_old_style_set = key.is_none(); + let (rule_ref, key_expr) = match refr.as_ref() { + Expr::RefBrack { refr, index, .. } + if !is_old_style_set + && !matches!( + index.as_ref(), + Expr::String { .. } | Expr::RawString { .. } + ) => + { + (refr.clone(), Some(index.clone())) + } + _ => (refr.clone(), None), + }; Ok(( Context { + key_expr, output_expr: key.clone(), value: Value::new_set(), - rule_ref: Some(refr.clone()), + rule_ref: Some(rule_ref), is_set: true, is_old_style_set, ..Context::default() @@ -4136,8 +4151,17 @@ impl Interpreter { let is_set = ctx.is_set; let is_object = ctx.key_expr.is_some() && !is_set; - let value = self.eval_rule_bodies(ctx, span, rule_body)?; let package_components = self.eval_rule_ref(&module.package.refr)?; + if is_set && ctx.key_expr.is_some() { + let mut prefix = package_components.clone(); + prefix.extend( + Self::get_rule_path_components(refr)? + .into_iter() + .map(Value::String), + ); + self.update_rule_value(span, prefix, Value::Undefined, false)?; + } + let value = self.eval_rule_bodies(ctx, span, rule_body)?; if value != Value::Undefined { for (path, value_in_map) in value.as_object()? { diff --git a/src/languages/rego/compiler/comprehensions.rs b/src/languages/rego/compiler/comprehensions.rs index 5f9964736..00288dbd2 100644 --- a/src/languages/rego/compiler/comprehensions.rs +++ b/src/languages/rego/compiler/comprehensions.rs @@ -42,6 +42,7 @@ impl<'a> Compiler<'a> { context_type: ContextType::Comprehension(context_type), dest_register: result_reg, key_expr: key_expr.cloned(), + multi_value: false, value_expr: value_expr.cloned(), span: span.clone(), key_value_loops_hoisted: false, diff --git a/src/languages/rego/compiler/loops.rs b/src/languages/rego/compiler/loops.rs index 8c17f1596..926c22fd7 100644 --- a/src/languages/rego/compiler/loops.rs +++ b/src/languages/rego/compiler/loops.rs @@ -156,6 +156,7 @@ impl<'a> Compiler<'a> { context_type: ContextType::Every, dest_register: result_reg, key_expr: None, + multi_value: false, value_expr: None, span: span.clone(), key_value_loops_hoisted: false, diff --git a/src/languages/rego/compiler/mod.rs b/src/languages/rego/compiler/mod.rs index 60121a60f..fa08510c6 100644 --- a/src/languages/rego/compiler/mod.rs +++ b/src/languages/rego/compiler/mod.rs @@ -61,6 +61,7 @@ pub struct CompilationContext { pub(super) context_type: ContextType, pub(super) dest_register: Register, pub(super) key_expr: Option, + pub(super) multi_value: bool, pub(super) value_expr: Option, pub(super) span: Span, pub(super) key_value_loops_hoisted: bool, diff --git a/src/languages/rego/compiler/queries.rs b/src/languages/rego/compiler/queries.rs index 8ba78e2d8..8c0356c33 100644 --- a/src/languages/rego/compiler/queries.rs +++ b/src/languages/rego/compiler/queries.rs @@ -208,10 +208,18 @@ impl<'a> Compiler<'a> { } ContextType::Rule(RuleType::PartialObject) => { self.emit_instruction( - Instruction::ObjectSet { - obj: dest_register, - key: key_register, - value: value_register, + if context.multi_value { + Instruction::ObjectSetAdd { + obj: dest_register, + key: key_register, + value: value_register, + } + } else { + Instruction::ObjectSet { + obj: dest_register, + key: key_register, + value: value_register, + } }, span, ); diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index abd6b0641..dbb81bd2c 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -55,6 +55,13 @@ impl<'a> Compiler<'a> { .map(|def| { if let Rule::Spec { head, .. } = def.as_ref() { match head { + RuleHead::Set { + refr, key: Some(_), .. + } if matches!(refr.as_ref(), Expr::RefBrack { index, .. } + if !Self::is_string_literal(index)) => + { + RuleType::PartialObject + } RuleHead::Set { .. } => RuleType::PartialSet, RuleHead::Compr { refr, assign, .. } => match refr.as_ref() { crate::ast::Expr::RefBrack { .. } if assign.is_some() => { @@ -470,12 +477,26 @@ impl<'a> Compiler<'a> { }; (key_expr, output_expr) } - RuleHead::Set { key, .. } => { + RuleHead::Set { refr, key, .. } => { self.rule_definition_function_params[rule_index as usize].push(None); self.rule_definition_destructuring_patterns[rule_index as usize] .push(None); - (None, key.clone()) + let key_expr = match refr.as_ref() { + Expr::RefBrack { refr, index, .. } + if rule_type == RuleType::PartialObject => + { + if Self::has_unsupported_bracket_prefix(refr) { + return Err( + CompilerError::PartialObjectNestedKeyUnsupported + .at(refr.span()), + ); + } + Some(index.clone()) + } + _ => None, + }; + (key_expr, key.clone()) } RuleHead::Func { assign, args, .. } => { let mut param_names = Vec::new(); @@ -572,6 +593,7 @@ impl<'a> Compiler<'a> { dest_register: result_register, context_type: ContextType::Rule(rule_type.clone()), key_expr, + multi_value: matches!(head, RuleHead::Set { .. }), value_expr, span, key_value_loops_hoisted: false, diff --git a/src/rvm/instructions/display.rs b/src/rvm/instructions/display.rs index 1b8a00fd2..9131a27ca 100644 --- a/src/rvm/instructions/display.rs +++ b/src/rvm/instructions/display.rs @@ -201,6 +201,9 @@ impl core::fmt::Display for Instruction { Instruction::ObjectSet { obj, key, value } => { format!("OBJECT_SET R({}) R({}) R({})", obj, key, value) } + Instruction::ObjectSetAdd { obj, key, value } => { + format!("OBJECT_SET_ADD R({}) R({}) R({})", obj, key, value) + } Instruction::ObjectCreate { params_index } => { format!("OBJECT_CREATE P({})", params_index) } diff --git a/src/rvm/instructions/mod.rs b/src/rvm/instructions/mod.rs index 82fe0f70a..fb33a30af 100644 --- a/src/rvm/instructions/mod.rs +++ b/src/rvm/instructions/mod.rs @@ -416,6 +416,14 @@ pub enum Instruction { /// Register that accumulates the result. result: u8, }, + + /// Add a defined value to the set at `obj[key]`, creating the set if absent. + /// Used by keyed partial-set rules (`rule[key] contains value`). + ObjectSetAdd { + obj: u8, + key: u8, + value: u8, + }, } impl Instruction { diff --git a/src/rvm/program/listing.rs b/src/rvm/program/listing.rs index 9b825e819..27153c9a7 100644 --- a/src/rvm/program/listing.rs +++ b/src/rvm/program/listing.rs @@ -507,6 +507,17 @@ fn format_instruction_readable( let comment = format!("Set field r{}[r{}] = r{}", obj, key, value); align_comment(&base, &comment, config.comment_column) } + Instruction::ObjectSetAdd { obj, key, value } => { + let base = format!( + "{}ObjectSetAdd r{}[r{}].insert(r{})", + indent, obj, key, value + ); + align_comment( + &base, + "Accumulate a keyed partial set", + config.comment_column, + ) + } Instruction::ObjectCreate { params_index } => { let params = program .instruction_data @@ -1022,6 +1033,7 @@ const fn get_instruction_name(instruction: &Instruction) -> &'static str { Instruction::HostAwait { .. } => "HOST_AWAIT", Instruction::Return { .. } => "RETURN", Instruction::ObjectSet { .. } => "OBJ_SET", + Instruction::ObjectSetAdd { .. } => "OBJ_SET_ADD", Instruction::ObjectCreate { .. } => "OBJ_CREATE", Instruction::Index { .. } => "INDEX", Instruction::IndexLiteral { .. } => "INDEX_LIT", diff --git a/src/rvm/tests/instruction_parser.rs b/src/rvm/tests/instruction_parser.rs index 5d3237524..e956f13da 100644 --- a/src/rvm/tests/instruction_parser.rs +++ b/src/rvm/tests/instruction_parser.rs @@ -53,6 +53,14 @@ pub fn parse_instruction(text: &str) -> Result { "RuleReturn" => parse_rule_return(params_text), "DestructuringSuccess" => parse_destructuring_success(params_text), "ObjectSet" => parse_object_set(params_text), + "ObjectSetAdd" => { + let params = parse_params(params_text)?; + Ok(Instruction::ObjectSetAdd { + obj: get_param_u16(¶ms, "obj")?.try_into()?, + key: get_param_u16(¶ms, "key")?.try_into()?, + value: get_param_u16(¶ms, "value")?.try_into()?, + }) + } "ObjectCreate" => parse_object_create(params_text), "Index" => parse_index(params_text), "IndexLiteral" => parse_index_literal(params_text), diff --git a/src/rvm/vm/dispatch.rs b/src/rvm/vm/dispatch.rs index b1913fbb8..842b00db4 100644 --- a/src/rvm/vm/dispatch.rs +++ b/src/rvm/vm/dispatch.rs @@ -444,6 +444,35 @@ impl RegoVM { ) -> Result { use Instruction::*; match instruction { + ObjectSetAdd { obj, key, value } => { + let key_value = self.get_register(key)?.clone(); + let value_value = self.get_register(value)?.clone(); + let mut obj_value = self.take_register(obj)?; + let result = (|| { + let Value::Object(ref mut fields) = obj_value else { + return Err(VmError::RegisterNotObject { + register: obj, + value: obj_value.clone(), + pc: self.pc, + }); + }; + if key_value != Value::Undefined && value_value != Value::Undefined { + let set = crate::Rc::make_mut(fields) + .get_or_insert_with(key_value, Value::new_set); + let Value::Set(ref mut elements) = *set else { + return Err(VmError::RegisterNotSet { + register: obj, + value: set.clone(), + pc: self.pc, + }); + }; + crate::Rc::make_mut(elements).insert(value_value); + } + Ok(InstructionOutcome::Continue) + })(); + self.set_register(obj, obj_value)?; + result + } ObjectSet { obj, key, value } => { let key_value = self.get_register(key)?.clone(); let value_value = self.get_register(value)?.clone(); diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml new file mode 100644 index 000000000..e297a7a0b --- /dev/null +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -0,0 +1,234 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +cases: + - note: keyed_partial_set_accumulates_per_key + modules: + - | + package test + ports[host] contains port if { + some host, values in {"web": [80, 443, 80], "db": [5432]} + some port in values + } + query: data.test.ports + want_result: + web: + set!: [80, 443] + db: + set!: [5432] + + - note: keyed_partial_set_multiple_definitions + modules: + - | + package test + values[key] contains value if { + key := "a" + value := 1 + } + values[key] contains value if { + key := "a" + value := 2 + } + values[key] contains value if { + key := "b" + value := 1 + } + query: data.test.values + want_result: + a: + set!: [1, 2] + b: + set!: [1] + + - note: keyed_partial_set_preserves_defined_values + modules: + - | + package test + values[key] contains value if { + some key, value in {"false": false, "null": null, "zero": 0} + } + query: data.test.values + want_result: + "false": + set!: [false] + "null": + set!: [null] + zero: + set!: [0] + + - note: keyed_partial_set_undefined_key + input: {} + modules: + - | + package test + values[input.missing] contains 1 if { true } + query: data.test.values + want_result: {} + + - note: keyed_partial_set_undefined_value + input: {} + modules: + - | + package test + values[key] contains input.missing if { key := "a" } + query: data.test.values + want_result: {} + + - note: keyed_partial_set_empty_body + modules: + - | + package test + values["a"] contains 1 + values["a"] contains 2 + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: keyed_partial_set_unmatched_body + modules: + - | + package test + values[key] contains 1 if { + key := "a" + false + } + query: data.test.values + want_result: {} + + - note: keyed_partial_set_head_key_loop + input: + keys: ["a", "b", "a"] + modules: + - | + package test + values[input.keys[_]] contains 1 if { true } + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [1] + + - note: keyed_partial_set_static_prefix + modules: + - | + package test + values.fixed[key] contains value if { + some key, value in {"a": 1, "b": 2} + } + query: data.test.values.fixed + want_result: + a: + set!: [1] + b: + set!: [2] + + - note: keyed_partial_set_mixed_missing_values + input: + values: + a: 1 + b: false + modules: + - | + package test + values[key] contains input.values[key] if { + some key in ["a", "missing", "b"] + } + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [false] + + - note: keyed_partial_set_multiple_definitions_suspendable + execution_mode: suspendable + modules: + - | + package test + values[key] contains value if { + some key, value in {"a": 1, "b": 2} + } + values[key] contains 3 if { key := "a" } + query: data.test.values + want_result: + a: + set!: [1, 3] + b: + set!: [2] + + - note: keyed_partial_set_head_key_loop_suspendable + execution_mode: suspendable + input: + keys: ["a", "b", "a"] + modules: + - | + package test + values[input.keys[_]] contains 1 if { true } + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [1] + + - note: keyed_partial_set_mixed_missing_values_suspendable + execution_mode: suspendable + input: + values: + a: 1 + b: false + modules: + - | + package test + values[key] contains input.values[key] if { + some key in ["a", "missing", "b"] + } + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [false] + + - note: keyed_partial_set_undefined_value_suspendable + execution_mode: suspendable + input: {} + modules: + - | + package test + values[key] contains input.missing if { key := "a" } + query: data.test.values + want_result: {} + + - note: keyed_partial_set_unmatched_body_suspendable + execution_mode: suspendable + modules: + - | + package test + values[key] contains 1 if { + key := "a" + false + } + query: data.test.values + want_result: {} + + - note: keyed_partial_set_numeric_key + modules: + - | + package test + values[1] contains value if { some value in [false, null] } + query: data.test.values + want_result: + 1: + set!: [false, null] + + - note: keyed_partial_set_preserves_legacy_set_head + rego_v0: true + modules: + - | + package test + values[value] { value := [1, 2, 1][_] } + query: data.test.values + want_result: + set!: [1, 2] diff --git a/tests/rvm/vm/suites/keyed_set_operations.yaml b/tests/rvm/vm/suites/keyed_set_operations.yaml new file mode 100644 index 000000000..e7dfdd983 --- /dev/null +++ b/tests/rvm/vm/suites/keyed_set_operations.yaml @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +cases: + - note: keyed_set_accumulates_and_deduplicates + literals: + - {} + - "key" + instructions: + - "Load { dest: 0, literal_idx: 0 }" + - "Load { dest: 1, literal_idx: 1 }" + - "LoadFalse { dest: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "LoadNull { dest: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "Return { value: 0 }" + want_result: + key: + set!: [false, null] + + - note: keyed_set_skips_undefined_key_and_value + literals: + - {} + - "key" + - "#undefined" + instructions: + - "Load { dest: 0, literal_idx: 0 }" + - "Load { dest: 1, literal_idx: 1 }" + - "Load { dest: 2, literal_idx: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "ObjectSetAdd { obj: 0, key: 2, value: 1 }" + - "Return { value: 0 }" + want_result: {} + + - note: keyed_set_preserves_shared_object_and_set + literals: + - {} + - "key" + instructions: + - "Load { dest: 0, literal_idx: 0 }" + - "Load { dest: 1, literal_idx: 1 }" + - "LoadFalse { dest: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "Move { dest: 3, src: 0 }" + - "LoadNull { dest: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "Return { value: 3 }" + want_result: + key: + set!: [false] + + - note: keyed_set_rejects_non_object_even_for_undefined + literals: + - "#undefined" + instructions: + - "LoadNull { dest: 0 }" + - "Load { dest: 1, literal_idx: 0 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 1 }" + - "Return { value: 0 }" + want_error: "Register 0 does not contain an object" + + - note: keyed_set_rejects_non_set_leaf + literals: + - {"key": 1} + - "key" + instructions: + - "Load { dest: 0, literal_idx: 0 }" + - "Load { dest: 1, literal_idx: 1 }" + - "LoadNull { dest: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "Return { value: 0 }" + want_error: "Register 0 does not contain a set" From 3f96fbec0bd7c071d41ae2d83f5942b7aa9db587 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:23:51 +0000 Subject: [PATCH 02/13] test(rvm): constrain keyed-set planning and document split scope Leave static and legacy set-head planning unchanged and cover explicit rejection of multiple dynamic key levels. Document the narrow three-register accumulation instruction. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 5 +++++ src/compiler/hoist.rs | 10 +++++++++- tests/rvm/rego/cases/keyed_partial_sets.yaml | 13 +++++++++++++ 3 files changed, 27 insertions(+), 1 deletion(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index 85088de23..2f4d1771e 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -85,6 +85,11 @@ comprehension contexts. Outside loops it aborts the current rule. Parameter structures: +- `ObjectSetAdd` uses three register operands without a parameter table. It + accumulates `rule[key] contains value` into an object of sets, preserving + existing members and omitting undefined keys and values. The target must be + an object, and an existing leaf must be a set. Static string prefixes remain + part of the rule path; multiple dynamic key levels remain unsupported. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register pairs. Literal keys must be sorted to match template order. - `ArrayCreateParams` and `SetCreateParams` store register lists. The VM checks diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index d6d385a94..dce6db4f7 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -484,7 +484,15 @@ impl LoopHoister { } RuleHead::Set { refr, key, .. } => { let key_expr = match refr.as_ref() { - Expr::RefBrack { index, .. } => Some(index.clone()), + Expr::RefBrack { index, .. } + if key.is_some() + && !matches!( + index.as_ref(), + Expr::String { .. } | Expr::RawString { .. } + ) => + { + Some(index.clone()) + } _ => None, }; (key_expr, key.clone()) diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index e297a7a0b..4d7ae7d7c 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -232,3 +232,16 @@ cases: query: data.test.values want_result: set!: [1, 2] + + - note: keyed_partial_set_rejects_multiple_dynamic_key_levels + modules: + - | + package test + values[group][key] contains 1 if { + group := "g" + key := "k" + } + main := values + query: data.test.main + want_error: "partial object rules with nested bracket keys are not yet supported by the RVM compiler" + allow_interpreter_success: true From 875eacb0e1aa12e32f66792d1b14e49ae158b995 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:48:58 +0000 Subject: [PATCH 03/13] fix(rvm): handle keyed-set review edge cases Classify literal set heads relative to the compiled document, reject mixed rule kinds, and hoist bodyless and empty-query head loops. Treat missing interpreter output collections as empty iterations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 6 +- src/interpreter.rs | 1 + src/languages/rego/compiler/queries.rs | 4 +- src/languages/rego/compiler/rules.rs | 86 +++++++--- tests/rvm/rego/cases/keyed_partial_sets.yaml | 158 +++++++++++++++++++ 5 files changed, 230 insertions(+), 25 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index 2f4d1771e..23f01ef2d 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -89,7 +89,11 @@ Parameter structures: accumulates `rule[key] contains value` into an object of sets, preserving existing members and omitting undefined keys and values. The target must be an object, and an existing leaf must be a set. Static string prefixes remain - part of the rule path; multiple dynamic key levels remain unsupported. + part of the rule path; when a parent rule also has dynamic keys, its literal + final keys accumulate into the same object. Head-key loops are hoisted even + for bodyless rules, and undefined collections contribute nothing. The compiler + rejects conflicting `contains` and single-valued rule kinds. Multiple dynamic + key levels remain unsupported. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register pairs. Literal keys must be sorted to match template order. - `ArrayCreateParams` and `SetCreateParams` store register lists. The VM checks diff --git a/src/interpreter.rs b/src/interpreter.rs index 74c5b3de1..421359038 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -2149,6 +2149,7 @@ impl Interpreter { result = self.eval_output_expr_in_loop(loop_tail)? || result; } } + Value::Undefined => {} _ => { let span = Self::loop_span(loop_info); return Err(span diff --git a/src/languages/rego/compiler/queries.rs b/src/languages/rego/compiler/queries.rs index 8c0356c33..4faf001d3 100644 --- a/src/languages/rego/compiler/queries.rs +++ b/src/languages/rego/compiler/queries.rs @@ -128,7 +128,9 @@ impl<'a> Compiler<'a> { let mut key_value_loops = Vec::new(); if let Some(expr) = key_expr.as_ref() { - key_value_loops.extend(self.get_expr_loops(expr)?); + if !Self::is_simple_literal(expr) { + key_value_loops.extend(self.get_expr_loops(expr)?); + } } if let Some(expr) = value_expr.as_ref() { diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index dbb81bd2c..e8145a8f7 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -50,16 +50,18 @@ impl<'a> Compiler<'a> { .into()); }; + let mut has_keyed_set = false; + let mut has_single_value = false; let rule_types: BTreeSet = definitions .iter() .map(|def| { - if let Rule::Spec { head, .. } = def.as_ref() { + Ok(if let Rule::Spec { head, .. } = def.as_ref() { + has_single_value |= !matches!(head, RuleHead::Set { .. }); match head { RuleHead::Set { refr, key: Some(_), .. - } if matches!(refr.as_ref(), Expr::RefBrack { index, .. } - if !Self::is_string_literal(index)) => - { + } if self.keyed_set_key(rule_path, def, refr)?.is_some() => { + has_keyed_set = true; RuleType::PartialObject } RuleHead::Set { .. } => RuleType::PartialSet, @@ -74,9 +76,18 @@ impl<'a> Compiler<'a> { } } else { RuleType::Complete - } + }) }) - .collect(); + .collect::>()?; + + if has_keyed_set && has_single_value { + return Err(CompilerError::General { + message: format!( + "conflicting rule kinds for '{rule_path}': contains and single-valued heads" + ), + } + .into()); + } if rule_types.len() > 1 { return Err(CompilerError::General { @@ -96,6 +107,43 @@ impl<'a> Compiler<'a> { }) } + fn keyed_set_key( + &self, + rule_path: &str, + rule: &crate::ast::NodeRef, + refr: &ExprRef, + ) -> Result> { + let (prefix, key) = match refr.as_ref() { + Expr::RefBrack { refr, index, .. } => (refr, index.clone()), + Expr::RefDot { + refr, field, eidx, .. + } => ( + refr, + ExprRef::new(Expr::String { + span: field.0.clone(), + value: field.1.clone(), + eidx: *eidx, + }), + ), + _ => return Ok(None), + }; + if !Self::is_string_literal(&key) { + return Ok(Some(key)); + } + + let module_index = self.find_module_index_for_rule(rule)?; + let module = &self.policy.get_modules()[module_index as usize]; + let package = + get_path_string(&module.package.refr, Some("data")).map_err(CompilerError::from)?; + if get_path_string(refr, Some(&package)).map_err(CompilerError::from)? == rule_path { + return Ok(None); + } + if get_path_string(prefix, Some(&package)).map_err(CompilerError::from)? != rule_path { + return Err(CompilerError::PartialObjectNestedKeyUnsupported.at(refr.span())); + } + Ok(Some(key)) + } + fn validate_partial_object_shape(&self, refr: &ExprRef) -> Result<()> { let Expr::RefBrack { refr: prefix, @@ -131,7 +179,7 @@ impl<'a> Compiler<'a> { matches!(expr.as_ref(), Expr::String { .. } | Expr::RawString { .. }) } - fn is_simple_literal(expr: &ExprRef) -> bool { + pub(super) fn is_simple_literal(expr: &ExprRef) -> bool { match expr.as_ref() { Expr::String { .. } | Expr::RawString { .. } @@ -482,19 +530,18 @@ impl<'a> Compiler<'a> { self.rule_definition_destructuring_patterns[rule_index as usize] .push(None); - let key_expr = match refr.as_ref() { - Expr::RefBrack { refr, index, .. } - if rule_type == RuleType::PartialObject => - { + let key_expr = if rule_type == RuleType::PartialObject { + if let Expr::RefBrack { refr, .. } = refr.as_ref() { if Self::has_unsupported_bracket_prefix(refr) { return Err( CompilerError::PartialObjectNestedKeyUnsupported .at(refr.span()), ); } - Some(index.clone()) } - _ => None, + self.keyed_set_key(rule_path, rule_ref, refr)? + } else { + None }; (key_expr, key.clone()) } @@ -619,7 +666,7 @@ impl<'a> Compiler<'a> { value_expr.span(), ); - self.emit_context_yield()?; + self.hoist_loops_and_emit_context_yield()?; self.emit_instruction(Instruction::RuleReturn {}, value_expr.span()); self.pop_scope(); @@ -665,15 +712,8 @@ impl<'a> Compiler<'a> { if !body.query.stmts.is_empty() { self.compile_query(&body.query)?; - } else if let Some(value_expr) = body_value_expr.clone() { - let value_reg = self.compile_rego_expr(&value_expr)?; - self.emit_instruction( - Instruction::Move { - dest: result_register, - src: value_reg, - }, - value_expr.span(), - ); + } else { + self.hoist_loops_and_emit_context_yield()?; } self.emit_instruction(Instruction::RuleReturn {}, &body.span); diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index 4d7ae7d7c..4b0803f92 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,164 @@ # Licensed under the MIT License. cases: + - note: review_missing_head_key_collection + input: {} + modules: + - | + package test + values[input.keys[_]] contains 1 if { true } + query: data.test.values + want_result: {} + + - note: review_missing_head_key_collection_suspendable + execution_mode: suspendable + input: {} + modules: + - | + package test + values[input.keys[_]] contains 1 if { true } + query: data.test.values + want_result: {} + + - note: review_mixed_keyed_rule_kinds + modules: + - | + package test + values[k] contains 1 if { k := "a" } + values[k] := 2 if { k := "a" } + query: data.test.values + want_error: "conflicting rule kinds" + + - note: review_mixed_keyed_rule_kinds_reverse + execution_mode: suspendable + modules: + - | + package test + values[k] := 2 if { k := "a" } + values[k] contains 1 if { k := "a" } + query: data.test.values + want_error: "conflicting rule kinds" + + - note: review_mixed_keyed_rule_kinds_equal_set + modules: + - | + package test + values[k] contains 1 if { k := "a" } + values[k] := {1} if { k := "a" } + query: data.test.values + want_error: "conflicting rule kinds" + allow_interpreter_success: true + + - note: review_mixed_literal_dynamic_keys + modules: + - | + package test + values["a"] contains 1 + values[k] contains 2 if { k := "a" } + query: data.test.values + want_result: + a: + set!: [1, 2] + + - note: review_mixed_literal_dynamic_keys_suspendable + execution_mode: suspendable + modules: + - | + package test + values[k] contains 2 if { k := "a" } + values["a"] contains 1 + query: data.test.values + want_result: + a: + set!: [1, 2] + + - note: review_bodyless_head_key_loop + input: + keys: ["a", "b"] + modules: + - | + package test + values[input.keys[_]] contains 1 + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [1] + + - note: review_mixed_dot_dynamic_keys + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { k := "a" } + query: data.test.values + want_result: + a: + set!: [1, 2] + + - note: review_mixed_literal_dynamic_keys_across_modules + execution_mode: suspendable + modules: + - | + package test.nested + values["a"] contains 1 if { true } + - | + package test.nested + values[k] contains 2 if { k := "a" } + query: data.test.nested.values + want_result: + a: + set!: [1, 2] + + - note: review_bodyless_missing_head_key_collection + input: {} + modules: + - | + package test + values[input.keys[_]] contains 1 + query: data.test.values + want_result: {} + + - note: review_mixed_literal_dynamic_child_entrypoint + modules: + - | + package test + values["a"] contains 1 + values[k] contains 2 if { k := "b" } + query: data.test.values.a + want_result: + set!: [1] + + - note: review_empty_query_head_key_loop + input: + keys: ["a", "b"] + modules: + - | + package test + values[input.keys[_]] contains 1 if {} + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [1] + + - note: review_bodyless_head_key_loop_suspendable + execution_mode: suspendable + input: + keys: ["a", "b"] + modules: + - | + package test + values[input.keys[_]] contains 1 + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [1] + - note: keyed_partial_set_accumulates_per_key modules: - | From d80e66e539c4d213703e098f496fb34c0895e47c Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:57:14 +0000 Subject: [PATCH 04/13] fix(rvm): preserve keyed-set omission and target validation Do not create interpreter leaf sets when keyed rule bodies fail. Validate existing ObjectSetAdd leaves even when the incoming value is undefined, without creating absent leaves. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 8 ++-- src/interpreter.rs | 5 ++- src/rvm/vm/dispatch.rs | 11 +++++ tests/rvm/rego/cases/keyed_partial_sets.yaml | 40 +++++++++++++++++++ tests/rvm/vm/suites/keyed_set_operations.yaml | 31 ++++++++++++++ 5 files changed, 90 insertions(+), 5 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index 23f01ef2d..f81bbab5d 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -88,10 +88,12 @@ Parameter structures: - `ObjectSetAdd` uses three register operands without a parameter table. It accumulates `rule[key] contains value` into an object of sets, preserving existing members and omitting undefined keys and values. The target must be - an object, and an existing leaf must be a set. Static string prefixes remain - part of the rule path; when a parent rule also has dynamic keys, its literal + an object, and an existing leaf must be a set even when the incoming value is + undefined. An undefined contribution does not create a missing leaf. Static + string prefixes remain part of the rule path; when a parent rule also has dynamic keys, its literal final keys accumulate into the same object. Head-key loops are hoisted even - for bodyless rules, and undefined collections contribute nothing. The compiler + for bodyless rules, and undefined collections contribute nothing. Failed keyed + rule bodies do not create empty leaf sets. The compiler rejects conflicting `contains` and single-valued rule kinds. Multiple dynamic key levels remain unsupported. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register diff --git a/src/interpreter.rs b/src/interpreter.rs index 421359038..cf9b56edf 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -4150,10 +4150,11 @@ impl Interpreter { RuleHead::Compr { refr, .. } | RuleHead::Set { refr, .. } => { let (ctx, _) = self.make_rule_context(rule_head)?; let is_set = ctx.is_set; + let is_keyed_set = is_set && ctx.key_expr.is_some(); let is_object = ctx.key_expr.is_some() && !is_set; let package_components = self.eval_rule_ref(&module.package.refr)?; - if is_set && ctx.key_expr.is_some() { + if is_keyed_set { let mut prefix = package_components.clone(); prefix.extend( Self::get_rule_path_components(refr)? @@ -4176,7 +4177,7 @@ impl Interpreter { is_set, )?; } - } else if is_set { + } else if is_set && !is_keyed_set { if let Ok(mut comps) = self.eval_rule_ref(refr) { let mut full_path = package_components; full_path.append(&mut comps); diff --git a/src/rvm/vm/dispatch.rs b/src/rvm/vm/dispatch.rs index 842b00db4..5d5fac0f7 100644 --- a/src/rvm/vm/dispatch.rs +++ b/src/rvm/vm/dispatch.rs @@ -456,6 +456,17 @@ impl RegoVM { pc: self.pc, }); }; + if key_value != Value::Undefined && value_value == Value::Undefined { + if let Some(set) = fields.get(&key_value) { + if !matches!(set, Value::Set(_)) { + return Err(VmError::RegisterNotSet { + register: obj, + value: set.clone(), + pc: self.pc, + }); + } + } + } if key_value != Value::Undefined && value_value != Value::Undefined { let set = crate::Rc::make_mut(fields) .get_or_insert_with(key_value, Value::new_set); diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index 4b0803f92..f0573ffba 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,46 @@ # Licensed under the MIT License. cases: + - note: keyed_partial_set_resolvable_key_failing_body + input: {key: "a", enabled: false} + modules: + - | + package test + values[input.key] contains 1 if { input.enabled } + query: data.test.values + want_result: {} + + - note: keyed_partial_set_resolvable_key_failing_body_suspendable + execution_mode: suspendable + input: {key: "a", enabled: false} + modules: + - | + package test + values[input.key] contains 1 if { input.enabled } + query: data.test.values + want_result: {} + + - note: keyed_partial_set_failing_body_preserves_other_keys + input: {key: "a", enabled: false} + modules: + - | + package test + values[input.key] contains 1 if { input.enabled } + values[k] contains 2 if { k := "b" } + query: data.test.values + want_result: + b: + set!: [2] + + - note: keyed_partial_set_resolvable_key_undefined_body + input: {key: "a"} + modules: + - | + package test + values[input.key] contains 1 if { input.enabled } + query: data.test.values + want_result: {} + - note: review_missing_head_key_collection input: {} modules: diff --git a/tests/rvm/vm/suites/keyed_set_operations.yaml b/tests/rvm/vm/suites/keyed_set_operations.yaml index e7dfdd983..7b4d79792 100644 --- a/tests/rvm/vm/suites/keyed_set_operations.yaml +++ b/tests/rvm/vm/suites/keyed_set_operations.yaml @@ -71,3 +71,34 @@ cases: - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" - "Return { value: 0 }" want_error: "Register 0 does not contain a set" + + - note: keyed_set_rejects_non_set_leaf_even_for_undefined_value + literals: + - {"key": 1} + - "key" + - "#undefined" + instructions: + - "Load { dest: 0, literal_idx: 0 }" + - "Load { dest: 1, literal_idx: 1 }" + - "Load { dest: 2, literal_idx: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "Return { value: 0 }" + want_error: "Register 0 does not contain a set" + + - note: keyed_set_undefined_value_preserves_existing_set + literals: + - {} + - "key" + - "#undefined" + instructions: + - "Load { dest: 0, literal_idx: 0 }" + - "Load { dest: 1, literal_idx: 1 }" + - "LoadNull { dest: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "Move { dest: 3, src: 0 }" + - "Load { dest: 2, literal_idx: 2 }" + - "ObjectSetAdd { obj: 0, key: 1, value: 2 }" + - "Return { value: 0 }" + want_result: + key: + set!: [null] From 9fbc70e83fdcce85a173a1b428a2da1736283ce2 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:18:37 +0000 Subject: [PATCH 05/13] fix(rvm): validate keyed paths and isolate body loop planning Reject unsupported keyed-set prefixes rather than dropping path components. Preserve compiler errors during rule reference resolution, reset head-loop hoisting per body, and skip scalar output-loop collections consistently. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 7 +- src/interpreter.rs | 8 +- src/languages/rego/compiler/references.rs | 18 +++- src/languages/rego/compiler/rules.rs | 17 ++-- tests/rvm/rego/cases/keyed_partial_sets.yaml | 88 ++++++++++++++++++++ 5 files changed, 116 insertions(+), 22 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index f81bbab5d..caca064d9 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -92,10 +92,13 @@ Parameter structures: undefined. An undefined contribution does not create a missing leaf. Static string prefixes remain part of the rule path; when a parent rule also has dynamic keys, its literal final keys accumulate into the same object. Head-key loops are hoisted even - for bodyless rules, and undefined collections contribute nothing. Failed keyed + for bodyless rules and independently for each body. Undefined and scalar + collections contribute nothing. Failed keyed rule bodies do not create empty leaf sets. The compiler rejects conflicting `contains` and single-valued rule kinds. Multiple dynamic - key levels remain unsupported. + key levels remain unsupported. Dynamic prefixes followed by dotted keys and + ancestor evaluation requiring intervening static prefixes are also explicitly + rejected; a static prefix is supported when it is the compiled rule path. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register pairs. Literal keys must be sorted to match template order. - `ArrayCreateParams` and `SetCreateParams` store register lists. The VM checks diff --git a/src/interpreter.rs b/src/interpreter.rs index cf9b56edf..e4a5f0c0f 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -2149,13 +2149,7 @@ impl Interpreter { result = self.eval_output_expr_in_loop(loop_tail)? || result; } } - Value::Undefined => {} - _ => { - let span = Self::loop_span(loop_info); - return Err(span - .source - .error(span.line, span.col, "item cannot be indexed")); - } + _ => {} } self.remove_loop_var_value(loop_target_expr); Ok(result) diff --git a/src/languages/rego/compiler/references.rs b/src/languages/rego/compiler/references.rs index ca8d8c5bd..e3d0aee96 100644 --- a/src/languages/rego/compiler/references.rs +++ b/src/languages/rego/compiler/references.rs @@ -169,7 +169,14 @@ impl<'a> Compiler<'a> { let rule_candidate = format_string_path(&static_prefix[0..=i]).map_err(CompilerError::from)?; - if let Ok(rule_index) = self.get_or_assign_rule_index(&rule_candidate) { + if self.policy.inner.rules.contains_key(&rule_candidate) + || self + .policy + .inner + .default_rules + .contains_key(&rule_candidate) + { + let rule_index = self.get_or_assign_rule_index(&rule_candidate)?; // Found a rule match! Call the rule let rule_result_reg = self.alloc_register(); self.emit_instruction( @@ -309,7 +316,14 @@ impl<'a> Compiler<'a> { append_path_component(&rule_candidate, field).map_err(CompilerError::from)?; } - if let Ok(rule_index) = self.get_or_assign_rule_index(&rule_candidate) { + if self.policy.inner.rules.contains_key(&rule_candidate) + || self + .policy + .inner + .default_rules + .contains_key(&rule_candidate) + { + let rule_index = self.get_or_assign_rule_index(&rule_candidate)?; let rule_result_reg = self.alloc_register(); self.emit_instruction( Instruction::CallRule { diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index e8145a8f7..c15459d88 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -127,15 +127,17 @@ impl<'a> Compiler<'a> { ), _ => return Ok(None), }; - if !Self::is_string_literal(&key) { - return Ok(Some(key)); + if Self::has_unsupported_bracket_prefix(prefix) { + return Err(CompilerError::PartialObjectNestedKeyUnsupported.at(refr.span())); } let module_index = self.find_module_index_for_rule(rule)?; let module = &self.policy.get_modules()[module_index as usize]; let package = get_path_string(&module.package.refr, Some("data")).map_err(CompilerError::from)?; - if get_path_string(refr, Some(&package)).map_err(CompilerError::from)? == rule_path { + if Self::is_string_literal(&key) + && get_path_string(refr, Some(&package)).map_err(CompilerError::from)? == rule_path + { return Ok(None); } if get_path_string(prefix, Some(&package)).map_err(CompilerError::from)? != rule_path { @@ -531,14 +533,6 @@ impl<'a> Compiler<'a> { .push(None); let key_expr = if rule_type == RuleType::PartialObject { - if let Expr::RefBrack { refr, .. } = refr.as_ref() { - if Self::has_unsupported_bracket_prefix(refr) { - return Err( - CompilerError::PartialObjectNestedKeyUnsupported - .at(refr.span()), - ); - } - } self.keyed_set_key(rule_path, rule_ref, refr)? } else { None @@ -700,6 +694,7 @@ impl<'a> Compiler<'a> { if let Some(context) = self.context_stack.last_mut() { context.value_expr = body_value_expr.clone(); + context.key_value_loops_hoisted = false; } self.emit_instruction( diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index f0573ffba..cece8357c 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,94 @@ # Licensed under the MIT License. cases: + - note: keyed_partial_set_rejects_dynamic_dotted_prefix + modules: + - | + package test + values[group].fixed contains 1 if { group := "g" } + main := values + query: data.test.main + want_error: "partial object rules with nested bracket keys are not yet supported by the RVM compiler" + allow_interpreter_success: true + + - note: keyed_partial_set_rejects_dynamic_dotted_prefix_suspendable + execution_mode: suspendable + modules: + - | + package test + values[group].fixed contains 1 if { group := "g" } + main := data.test.values + query: data.test.main + want_error: "partial object rules with nested bracket keys are not yet supported by the RVM compiler" + allow_interpreter_success: true + + - note: keyed_partial_set_rejects_flattened_static_prefix + modules: + - | + package test + values.fixed[k] contains 1 if { k := "a" } + main := values + query: data.test.main + want_error: "partial object rules with nested bracket keys are not yet supported by the RVM compiler" + allow_interpreter_success: true + + - note: keyed_partial_set_rejects_flattened_static_prefix_suspendable + execution_mode: suspendable + modules: + - | + package test + values.fixed[k] contains 1 if { k := "a" } + main := data.test.values + query: data.test.main + want_error: "partial object rules with nested bracket keys are not yet supported by the RVM compiler" + allow_interpreter_success: true + + - note: keyed_partial_set_multiple_body_head_loops + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[_]] := 1 if { false } { true } + query: data.test.values + want_result: + a: 1 + b: 1 + + - note: keyed_partial_set_multiple_body_head_loops_suspendable + execution_mode: suspendable + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[_]] := 1 if { false } { true } + query: data.test.values + want_result: + a: 1 + b: 1 + + - note: keyed_partial_set_scalar_head_collections + input: {collections: [null, false, true, 42, "abc"]} + modules: + - | + package test + values[collection[_]] contains 1 if { + some collection in input.collections + } + query: data.test.values + want_result: {} + + - note: keyed_partial_set_scalar_head_collections_suspendable + execution_mode: suspendable + input: {collections: [null, false, true, 42, "abc"]} + modules: + - | + package test + values[collection[_]] contains 1 if { + some collection in input.collections + } + query: data.test.values + want_result: {} + - note: keyed_partial_set_resolvable_key_failing_body input: {key: "a", enabled: false} modules: From 166e4c40ddf3b3fc8686fc8a56d1544ce5a1c9f7 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:39:30 +0000 Subject: [PATCH 06/13] fix(rvm): share output index bindings and seed static sets Plan output keys and values against an evolving context. Apply interpreter loop-index binding plans in iteration-local scopes. Seed static set leaves once before evaluating rule bodies, preserving empty declarations and prior contributions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 4 + src/compiler/hoist.rs | 54 ++--- src/interpreter.rs | 129 ++++++---- src/languages/rego/compiler/rules.rs | 57 +++++ tests/rvm/rego/cases/keyed_partial_sets.yaml | 234 +++++++++++++++++++ 5 files changed, 402 insertions(+), 76 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index caca064d9..b431b661b 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -99,6 +99,10 @@ Parameter structures: key levels remain unsupported. Dynamic prefixes followed by dotted keys and ancestor evaluation requiring intervening static prefixes are also explicitly rejected; a static prefix is supported when it is the compiled rule path. + Static string leaf declarations are seeded as empty sets once before body + evaluation, without overwriting accumulated members. Key and value output + expressions share index bindings; named indices pair entries, while independent + wildcards retain Cartesian-product semantics. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register pairs. Literal keys must be sorted to match template order. - `ArrayCreateParams` and `SetCreateParams` store register lists. The VM checks diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index dce6db4f7..ba920c8b6 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -515,7 +515,7 @@ impl LoopHoister { context.current_scope_bound_vars.clone(); // Store the context for this query - let populated_body_context = + let mut populated_body_context = self.populate_query(module_idx, &body.query, &body_context)?; self.lookup .ensure_query_capacity(module_idx, body.query.qidx); @@ -527,21 +527,17 @@ impl LoopHoister { // Process the key expression if present if let Some(ref key) = key_expr { - self.populate_output_expr(module_idx, key, &populated_body_context)?; + self.populate_output_expr(module_idx, key, &mut populated_body_context)?; } // Process the head value expression (if present) - if let Some(ref head_value) = value_expr { - self.populate_output_expr(module_idx, head_value, &populated_body_context)?; - } - - // Process the value expression if present in the assign - if let Some(ref assign) = body.assign { - self.populate_output_expr( - module_idx, - &assign.value, - &populated_body_context, - )?; + if let Some(output) = body + .assign + .as_ref() + .map(|assign| &assign.value) + .or(value_expr.as_ref()) + { + self.populate_output_expr(module_idx, output, &mut populated_body_context)?; } } @@ -556,18 +552,18 @@ impl LoopHoister { context.current_scope_bound_vars.clone(); if let Some(ref key) = key_expr { - self.populate_output_expr(module_idx, key, &body_context)?; + self.populate_output_expr(module_idx, key, &mut body_context)?; } if let Some(ref value) = value_expr { - self.populate_output_expr(module_idx, value, &body_context)?; + self.populate_output_expr(module_idx, value, &mut body_context)?; } } } Rule::Default { value, .. } => { // For default rules, just process the value expression - let context = self.create_scope_context(module_idx)?; - self.populate_output_expr(module_idx, value, &context)?; + let mut context = self.create_scope_context(module_idx)?; + self.populate_output_expr(module_idx, value, &mut context)?; } } @@ -733,12 +729,12 @@ impl LoopHoister { None, Some(term.clone()), ); - let populated_context = + let mut populated_context = self.populate_query(module_idx, query.as_ref(), &compr_context)?; self.lookup.ensure_query_capacity(module_idx, query.qidx); self.lookup .set_query_context(module_idx, query.qidx, populated_context.clone())?; - self.populate_output_expr_with_context(module_idx, term, &populated_context)?; + self.populate_output_expr(module_idx, term, &mut populated_context)?; } E::ObjectCompr { key, value, query, .. @@ -748,13 +744,13 @@ impl LoopHoister { Some(key.clone()), Some(value.clone()), ); - let populated_context = + let mut populated_context = self.populate_query(module_idx, query.as_ref(), &compr_context)?; self.lookup.ensure_query_capacity(module_idx, query.qidx); self.lookup .set_query_context(module_idx, query.qidx, populated_context.clone())?; - self.populate_output_expr_with_context(module_idx, key, &populated_context)?; - self.populate_output_expr_with_context(module_idx, value, &populated_context)?; + self.populate_output_expr(module_idx, key, &mut populated_context)?; + self.populate_output_expr(module_idx, value, &mut populated_context)?; } E::Call { fcn, params, .. } => { self.analyze_expr(module_idx, fcn, context, loops)?; @@ -920,20 +916,10 @@ impl LoopHoister { &mut self, module_idx: u32, expr: &ExprRef, - context: &ScopeContext, - ) -> Result<()> { - self.populate_output_expr_with_context(module_idx, expr, context) - } - - fn populate_output_expr_with_context( - &mut self, - module_idx: u32, - expr: &ExprRef, - context: &ScopeContext, + context: &mut ScopeContext, ) -> Result<()> { let mut loops = Vec::new(); - let mut expr_context = context.clone(); - self.analyze_expr(module_idx, expr, &mut expr_context, &mut loops)?; + self.analyze_expr(module_idx, expr, context, &mut loops)?; let expr_idx = expr.as_ref().eidx(); self.lookup.ensure_expr_capacity(module_idx, expr_idx); diff --git a/src/interpreter.rs b/src/interpreter.rs index e4a5f0c0f..e86919c53 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -1709,29 +1709,7 @@ impl Interpreter { } } - let index_plan = if let Some(index) = index_expr { - let module_idx = self.current_module_index; - let expr_idx = index.as_ref().eidx(); - let plan = self - .compiled_policy - .loop_hoisting_table - .get_expr_binding_plan(module_idx, expr_idx) - .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; - - match plan.cloned() { - Some(BindingPlan::LoopIndex { - destructuring_plan, .. - }) => destructuring_plan, - Some(other_plan) => { - bail!("internal error: expected LoopIndex for loop index expression, got {:?}", other_plan); - } - None => { - bail!("internal error: no binding plan found for loop index expression"); - } - } - } else { - bail!("internal error: no binding plan found for loop index expression"); - }; + let index_plan = self.get_loop_index_plan(loop_info)?; // Create a new scope. self.scopes.push(Scope::default()); @@ -1950,6 +1928,47 @@ impl Interpreter { Ok(is_const && Self::is_simple_literal(output_expr)?) } + fn get_loop_index_plan(&self, loop_info: &HoistedLoop) -> Result { + let index = Self::loop_index_expr(loop_info) + .ok_or_else(|| anyhow!("internal error: no loop index expression"))?; + let plan = self + .compiled_policy + .loop_hoisting_table + .get_expr_binding_plan(self.current_module_index, index.as_ref().eidx()) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; + match plan { + Some(BindingPlan::LoopIndex { + destructuring_plan, .. + }) => Ok(destructuring_plan.clone()), + Some(other_plan) => { + bail!( + "internal error: expected LoopIndex for loop index expression, got {:?}", + other_plan + ) + } + None => bail!("internal error: no binding plan found for loop index expression"), + } + } + + fn eval_output_loop_iteration( + &mut self, + loop_info: &HoistedLoop, + index_plan: &DestructuringPlan, + index: &Value, + value: &Value, + loops: &[HoistedLoop], + ) -> Result { + self.memory_check()?; + self.set_loop_var_value(Self::loop_assignment_expr(loop_info), value.clone())?; + let result = if self.execute_destructuring_plan(index_plan, index)? == Value::Bool(true) { + self.eval_output_expr_in_loop(loops) + } else { + Ok(false) + }; + Self::clear_scope(self.current_scope_mut()?); + result + } + fn eval_output_expr_in_loop(&mut self, loops: &[HoistedLoop]) -> Result { self.check_execution_time()?; if loops.is_empty() { @@ -2128,31 +2147,57 @@ impl Interpreter { let (loop_info, loop_tail) = loops .split_first() .ok_or_else(|| anyhow!("internal error: expected loop info"))?; - let mut result = false; let loop_target_expr = Self::loop_assignment_expr(loop_info); - match self.eval_expr(Self::loop_collection_expr(loop_info))? { - Value::Array(items) => { - for v in items.iter() { - self.set_loop_var_value(loop_target_expr, v.clone())?; - result = self.eval_output_expr_in_loop(loop_tail)? || result; + let collection = self.eval_expr(Self::loop_collection_expr(loop_info))?; + let index_plan = if Self::loop_index_expr(loop_info).is_some() { + self.get_loop_index_plan(loop_info)? + } else { + DestructuringPlan::Ignore + }; + self.scopes.push(Scope::default()); + let result = (|| { + let mut result = false; + match collection { + Value::Array(items) => { + for (idx, value) in items.iter().enumerate() { + result = self.eval_output_loop_iteration( + loop_info, + &index_plan, + &Value::from(idx), + value, + loop_tail, + )? || result; + } } - } - Value::Set(items) => { - for v in items.iter() { - self.set_loop_var_value(loop_target_expr, v.clone())?; - result = self.eval_output_expr_in_loop(loop_tail)? || result; + Value::Set(items) => { + for value in items.iter() { + result = self.eval_output_loop_iteration( + loop_info, + &index_plan, + value, + value, + loop_tail, + )? || result; + } } - } - Value::Object(obj) => { - for (_, v) in obj.iter() { - self.set_loop_var_value(loop_target_expr, v.clone())?; - result = self.eval_output_expr_in_loop(loop_tail)? || result; + Value::Object(obj) => { + for (key, value) in obj.iter() { + result = self.eval_output_loop_iteration( + loop_info, + &index_plan, + key, + value, + loop_tail, + )? || result; + } } + _ => {} } - _ => {} - } + Ok(result) + })(); + self.scopes.pop(); self.remove_loop_var_value(loop_target_expr); - Ok(result) + result } fn get_current_context(&self) -> Result<&Context> { diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index c15459d88..08fc7b0f7 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -497,6 +497,41 @@ impl<'a> Compiler<'a> { self.rule_definition_static_values.push(Vec::new()); } + // Seed all static leaves once, before any body can fail or accumulate members. + let mut initial_object = Value::new_object(); + if rule_type == RuleType::PartialObject { + for rule in rule_definitions { + if let Rule::Spec { + head: + RuleHead::Set { + refr, key: Some(_), .. + }, + .. + } = rule.as_ref() + { + if let Some(key) = self.keyed_set_key(rule_path, rule, refr)? { + if let Expr::String { value, .. } | Expr::RawString { value, .. } = + key.as_ref() + { + initial_object + .as_object_mut() + .map_err(CompilerError::from)? + .insert(value.clone(), Value::new_set()); + } + } + } + } + } + let initial_literal = if initial_object + .as_object() + .map_err(CompilerError::from)? + .is_empty() + { + None + } else { + Some(self.add_literal(initial_object)) + }; + let mut num_registers_used = 0; let mut rule_param_count: Option = None; @@ -659,6 +694,17 @@ impl<'a> Compiler<'a> { }, value_expr.span(), ); + if def_idx == 0 { + if let Some(literal_idx) = initial_literal { + self.emit_instruction( + Instruction::Load { + dest: result_register, + literal_idx, + }, + value_expr.span(), + ); + } + } self.hoist_loops_and_emit_context_yield()?; @@ -704,6 +750,17 @@ impl<'a> Compiler<'a> { }, &body.span, ); + if def_idx == 0 && body_idx == 0 { + if let Some(literal_idx) = initial_literal { + self.emit_instruction( + Instruction::Load { + dest: result_register, + literal_idx, + }, + &body.span, + ); + } + } if !body.query.stmts.is_empty() { self.compile_query(&body.query)?; diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index cece8357c..d1374f7b5 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,240 @@ # Licensed under the MIT License. cases: + - note: keyed_partial_set_shared_head_index + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[i]] contains input.vals[i] if { true } + query: data.test.values + want_result: + a: + set!: [10] + b: + set!: [20] + + - note: keyed_partial_set_shared_head_index_bodyless + execution_mode: suspendable + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[i]] contains input.vals[i] + query: data.test.values + want_result: + a: + set!: [10] + b: + set!: [20] + + - note: keyed_partial_set_shared_head_index_suspendable + execution_mode: suspendable + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[i]] contains input.vals[i] if { true } + query: data.test.values + want_result: + a: + set!: [10] + b: + set!: [20] + + - note: keyed_partial_set_shared_head_index_bodyless_completion + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[i]] contains input.vals[i] + query: data.test.values + want_result: + a: + set!: [10] + b: + set!: [20] + + - note: keyed_partial_set_shared_index_with_nested_value_loop + execution_mode: suspendable + input: {keys: ["a", "b"], vals: [[10, 11], [20, 21]]} + modules: + - | + package test + values[input.keys[i]] contains input.vals[i][j] if { true } + query: data.test.values + want_result: + a: + set!: [10, 11] + b: + set!: [20, 21] + + - note: keyed_partial_set_body_bound_head_index + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[i]] contains input.vals[i] if { i := 1 } + query: data.test.values + want_result: + b: + set!: [20] + + - note: keyed_partial_set_shared_comprehension_head_index + execution_mode: suspendable + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + main := {input.keys[i]: input.vals[i] | true} + query: data.test.main + want_result: {a: 10, b: 20} + + - note: keyed_partial_set_head_index_value + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[i]] contains i if { true } + query: data.test.values + want_result: + a: + set!: [0] + b: + set!: [1] + + - note: keyed_partial_set_shared_object_head_index + execution_mode: suspendable + input: {keys: {x: "a", y: "b"}, vals: {x: 10, y: 20}} + modules: + - | + package test + values[input.keys[i]] contains input.vals[i] if { true } + query: data.test.values + want_result: + a: + set!: [10] + b: + set!: [20] + + - note: keyed_partial_set_shared_set_head_index + modules: + - | + package test + values[keys[i]] contains i if { keys := {"a", "b"} } + query: data.test.values + want_result: + a: + set!: ["a"] + b: + set!: ["b"] + + - note: keyed_partial_set_short_shared_value_collection + execution_mode: suspendable + input: {keys: ["a", "b", "c"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[i]] contains input.vals[i] if { true } + query: data.test.values + want_result: + a: + set!: [10] + b: + set!: [20] + + - note: keyed_partial_set_independent_head_wildcards + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[_]] contains input.vals[_] if { true } + query: data.test.values + want_result: + a: + set!: [10, 20] + b: + set!: [10, 20] + + - note: keyed_partial_set_static_empty_leaf + modules: + - | + package test + values.a contains 1 if { false } + main := values + query: data.test.main + want_result: + a: + set!: [] + + - note: keyed_partial_set_static_empty_leaf_suspendable + execution_mode: suspendable + modules: + - | + package test + values["a"] contains 1 if { false } + main := values + query: data.test.main + want_result: + a: + set!: [] + + - note: keyed_partial_set_static_empty_with_dynamic + modules: + - | + package test + values.a contains 1 if { false } + values[k] contains 2 if { k := "b" } + values[k] contains 3 if { k := "c"; false } + query: data.test.values + want_result: + a: + set!: [] + b: + set!: [2] + + - note: keyed_partial_set_static_empty_after_dynamic + execution_mode: suspendable + modules: + - | + package test + values[k] contains 2 if { k := "a" } + values.a contains 1 if { false } + query: data.test.values + want_result: + a: + set!: [2] + + - note: keyed_partial_set_static_empty_with_bodyless_first + execution_mode: suspendable + modules: + - | + package test + values.b contains 2 + values.a contains 1 if { false } + values.c contains 3 if { false } + main := values + query: data.test.main + want_result: + a: + set!: [] + b: + set!: [2] + c: + set!: [] + + - note: keyed_partial_set_static_empty_before_dynamic + modules: + - | + package test + values.a contains 1 if { false } + values[k] contains 2 if { k := "a" } + values.a contains 3 if { false } + query: data.test.values + want_result: + a: + set!: [2] + - note: keyed_partial_set_rejects_dynamic_dotted_prefix modules: - | From 21632f2fc846276711b932d3ef27a87ec7b4e324 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:25:15 +0000 Subject: [PATCH 07/13] fix(rvm): isolate output loop plans across rule bodies Store body-scoped output loops and binding plans and select them in both evaluators. Clear stale compiler loop-register mappings between bodies and align suspendable partial-rule accumulation with run-to-completion execution. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 6 +- src/compiler/hoist.rs | 140 +++++++++++++++--- src/interpreter.rs | 51 +++++-- src/languages/rego/compiler/comprehensions.rs | 1 + src/languages/rego/compiler/core.rs | 9 +- src/languages/rego/compiler/loops.rs | 9 +- src/languages/rego/compiler/mod.rs | 1 + src/languages/rego/compiler/rules.rs | 2 + src/rvm/vm/rules.rs | 25 +++- tests/rvm/rego/cases/keyed_partial_sets.yaml | 86 +++++++++++ 10 files changed, 290 insertions(+), 40 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index b431b661b..dc915f86e 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -102,7 +102,11 @@ Parameter structures: Static string leaf declarations are seeded as empty sets once before body evaluation, without overwriting accumulated members. Key and value output expressions share index bindings; named indices pair entries, while independent - wildcards retain Cartesian-product semantics. + wildcards retain Cartesian-product semantics. Output-loop and binding metadata + is selected by body query as well as module and expression, so differing body + bindings cannot overwrite each other's head plans. Both execution modes + accumulate all independent partial-rule bodies, while `else` bodies retain + first-success semantics. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register pairs. Literal keys must be sorted to match template order. - `ArrayCreateParams` and `SetCreateParams` store register lists. The VM checks diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index c4a2515a8..192b78134 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -24,7 +24,7 @@ use crate::scheduler::compute_module_globals; use crate::*; use anyhow::{anyhow, Result}; -use alloc::collections::BTreeSet; +use alloc::collections::{BTreeMap, BTreeSet}; use alloc::vec::Vec; /// Implementation of VariableBindingContext for ScopeContext @@ -120,6 +120,10 @@ pub struct HoistedLoopsLookup { /// Maps (module_index, query_index) -> ScopeContext /// Stores compilation contexts for queries (rules, comprehensions, every) query_contexts: Lookup, + + // Rule heads reuse expression IDs across bodies with different bindings. + output_loops: BTreeMap<(u32, u32, u32), Vec>, + output_binding_plans: BTreeMap<(u32, u32, u32), BindingPlan>, } impl HoistedLoopsLookup { @@ -195,6 +199,37 @@ impl HoistedLoopsLookup { self.expr_loops.get_checked(module_idx, expr_idx) } + pub fn get_output_loops( + &self, + module_idx: u32, + query_idx: Option, + expr_idx: u32, + ) -> LookupResult>> { + if let Some(query_idx) = query_idx { + if let Some(loops) = self.output_loops.get(&(module_idx, query_idx, expr_idx)) { + return Ok(Some(loops)); + } + } + self.get_expr_loops(module_idx, expr_idx) + } + + pub fn get_output_binding_plan( + &self, + module_idx: u32, + query_idx: Option, + expr_idx: u32, + ) -> LookupResult> { + if let Some(query_idx) = query_idx { + if let Some(plan) = self + .output_binding_plans + .get(&(module_idx, query_idx, expr_idx)) + { + return Ok(Some(plan)); + } + } + self.get_expr_binding_plan(module_idx, expr_idx) + } + /// Store the compilation context for a query pub fn set_query_context( &mut self, @@ -246,6 +281,7 @@ pub struct LoopHoister { lookup: HoistedLoopsLookup, schedule: Option>, module_globals: Lookup>>, + output_query: Option, } impl LoopHoister { @@ -255,6 +291,7 @@ impl LoopHoister { lookup: HoistedLoopsLookup::new(), schedule: None, module_globals: Lookup::new(), + output_query: None, } } @@ -264,6 +301,7 @@ impl LoopHoister { lookup: HoistedLoopsLookup::new(), schedule: Some(schedule), module_globals: Lookup::new(), + output_query: None, } } @@ -477,7 +515,12 @@ impl LoopHoister { // Process the key expression if present if let Some(ref key) = key_expr { - self.populate_output_expr(module_idx, key, &mut populated_body_context)?; + self.populate_body_output_expr( + module_idx, + body.query.qidx, + key, + &mut populated_body_context, + )?; } // Process the head value expression (if present) @@ -487,7 +530,12 @@ impl LoopHoister { .map(|assign| &assign.value) .or(value_expr.as_ref()) { - self.populate_output_expr(module_idx, output, &mut populated_body_context)?; + self.populate_body_output_expr( + module_idx, + body.query.qidx, + output, + &mut populated_body_context, + )?; } } @@ -526,6 +574,18 @@ impl LoopHoister { module_idx: u32, query: &Query, parent_context: &ScopeContext, + ) -> Result { + let output_query = self.output_query.take(); + let result = self.populate_query_impl(module_idx, query, parent_context); + self.output_query = output_query; + result + } + + fn populate_query_impl( + &mut self, + module_idx: u32, + query: &Query, + parent_context: &ScopeContext, ) -> Result { self.lookup.ensure_query_capacity(module_idx, query.qidx); let mut context = parent_context.clone(); @@ -684,7 +744,12 @@ impl LoopHoister { self.lookup.ensure_query_capacity(module_idx, query.qidx); self.lookup .set_query_context(module_idx, query.qidx, populated_context.clone())?; - self.populate_output_expr(module_idx, term, &mut populated_context)?; + self.populate_body_output_expr( + module_idx, + query.qidx, + term, + &mut populated_context, + )?; } E::ObjectCompr { key, value, query, .. @@ -699,8 +764,18 @@ impl LoopHoister { self.lookup.ensure_query_capacity(module_idx, query.qidx); self.lookup .set_query_context(module_idx, query.qidx, populated_context.clone())?; - self.populate_output_expr(module_idx, key, &mut populated_context)?; - self.populate_output_expr(module_idx, value, &mut populated_context)?; + self.populate_body_output_expr( + module_idx, + query.qidx, + key, + &mut populated_context, + )?; + self.populate_body_output_expr( + module_idx, + query.qidx, + value, + &mut populated_context, + )?; } E::Call { fcn, params, .. } => { self.analyze_expr(module_idx, fcn, context, loops)?; @@ -741,11 +816,7 @@ impl LoopHoister { // Immediately bind variables from the plan to context Self::bind_vars_from_plan_to_context(&binding_plan, context); - self.lookup.set_expr_binding_plan( - module_idx, - expr_idx, - binding_plan, - )?; + self.set_output_binding_plan(module_idx, expr_idx, binding_plan)?; } Err(err) => return Err(map_binding_error(err)), } @@ -769,11 +840,7 @@ impl LoopHoister { let expr_idx = index.as_ref().eidx(); self.lookup.ensure_expr_capacity(module_idx, expr_idx); Self::bind_vars_from_plan_to_context(&binding_plan, context); - self.lookup.set_expr_binding_plan( - module_idx, - expr_idx, - binding_plan, - )?; + self.set_output_binding_plan(module_idx, expr_idx, binding_plan)?; } Err(err) => return Err(map_binding_error(err)), } @@ -805,8 +872,7 @@ impl LoopHoister { let expr_idx = expr.as_ref().eidx(); self.lookup.ensure_expr_capacity(module_idx, expr_idx); Self::bind_vars_from_plan_to_context(&binding_plan, context); - self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan)?; + self.set_output_binding_plan(module_idx, expr_idx, binding_plan)?; self.analyze_expr(module_idx, lhs, context, loops)?; self.analyze_expr(module_idx, rhs, context, loops)?; @@ -862,6 +928,36 @@ impl LoopHoister { } /// Populate loop information for output expressions (rule values, comprehension terms) + fn populate_body_output_expr( + &mut self, + module_idx: u32, + query_idx: u32, + expr: &ExprRef, + context: &mut ScopeContext, + ) -> Result<()> { + let previous = self.output_query.replace(query_idx); + let result = self.populate_output_expr(module_idx, expr, context); + self.output_query = previous; + result + } + + fn set_output_binding_plan( + &mut self, + module_idx: u32, + expr_idx: u32, + plan: BindingPlan, + ) -> Result<()> { + if let Some(query_idx) = self.output_query { + self.lookup + .output_binding_plans + .insert((module_idx, query_idx, expr_idx), plan); + Ok(()) + } else { + self.lookup + .set_expr_binding_plan(module_idx, expr_idx, plan) + } + } + fn populate_output_expr( &mut self, module_idx: u32, @@ -873,7 +969,13 @@ impl LoopHoister { let expr_idx = expr.as_ref().eidx(); self.lookup.ensure_expr_capacity(module_idx, expr_idx); - self.lookup.set_expr_loops(module_idx, expr_idx, loops)?; + if let Some(query_idx) = self.output_query { + self.lookup + .output_loops + .insert((module_idx, query_idx, expr_idx), loops); + } else { + self.lookup.set_expr_loops(module_idx, expr_idx, loops)?; + } Ok(()) } diff --git a/src/interpreter.rs b/src/interpreter.rs index 7b31accb9..10998c6e3 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -190,6 +190,7 @@ struct Context { is_old_style_set: bool, output_constness_determined: bool, early_return: bool, + output_query: Option, } impl Default for Context { @@ -207,6 +208,7 @@ impl Default for Context { is_old_style_set: false, output_constness_determined: false, early_return: false, + output_query: None, } } } @@ -369,6 +371,20 @@ impl Interpreter { } } + fn expression_binding_plan( + &self, + expr_idx: u32, + ) -> crate::lookup::LookupResult> { + self.loop_hoisting_table(self.current_module_index) + .get_output_binding_plan( + self.current_module_index, + self.contexts + .last() + .and_then(|context| context.output_query), + expr_idx, + ) + } + pub const fn get_data_mut(&mut self) -> &mut Value { &mut self.data } @@ -500,11 +516,9 @@ impl Interpreter { if let Expr::Call { params, .. } = Self::loop_assignment_expr(loop_info).as_ref() { if let Some(last_param) = params.last() { - let module_idx = self.current_module_index; let expr_idx = last_param.as_ref().eidx(); let binding_plan = self - .loop_hoisting_table(module_idx) - .get_expr_binding_plan(module_idx, expr_idx) + .expression_binding_plan(expr_idx) .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; return match binding_plan.cloned() { @@ -1021,12 +1035,10 @@ impl Interpreter { let mut count: usize = 0; // Fetch the binding plan for this some..in expression - let module_idx = self.current_module_index; let expr_idx = collection.as_ref().eidx(); let binding_plan = self - .loop_hoisting_table(module_idx) - .get_expr_binding_plan(module_idx, expr_idx) + .expression_binding_plan(expr_idx) .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; let Some(BindingPlan::SomeIn { @@ -1940,8 +1952,7 @@ impl Interpreter { let index = Self::loop_index_expr(loop_info) .ok_or_else(|| anyhow!("internal error: no loop index expression"))?; let plan = self - .loop_hoisting_table(self.current_module_index) - .get_expr_binding_plan(self.current_module_index, index.as_ref().eidx()) + .expression_binding_plan(index.as_ref().eidx()) .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; match plan { Some(BindingPlan::LoopIndex { @@ -2237,7 +2248,13 @@ impl Interpreter { if let Some(ke) = &key_expr { match self .loop_hoisting_table(self.current_module_index) - .get_expr_loops(self.current_module_index, ke.as_ref().eidx()) + .get_output_loops( + self.current_module_index, + self.contexts + .last() + .and_then(|context| context.output_query), + ke.as_ref().eidx(), + ) .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? { Some(hoisted_loops) => { @@ -2253,7 +2270,13 @@ impl Interpreter { if let Some(oe) = &output_expr { match self .loop_hoisting_table(self.current_module_index) - .get_expr_loops(self.current_module_index, oe.as_ref().eidx()) + .get_output_loops( + self.current_module_index, + self.contexts + .last() + .and_then(|context| context.output_query), + oe.as_ref().eidx(), + ) .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? { Some(hoisted_loops) => { @@ -2356,6 +2379,7 @@ impl Interpreter { fn eval_query(&mut self, query: &Ref) -> Result { self.check_execution_time()?; + self.get_current_context_mut()?.output_query = Some(query.qidx); // Execute the query in a new scope self.scopes.push(Scope::new()); let order_indices = { @@ -3160,13 +3184,11 @@ impl Interpreter { let value = self.eval_call_impl(span, expr, fcn, arg_prefix)?; if allow_return_arg { - let module_idx = self.current_module_index; let expr_idx = last_param.as_ref().eidx(); if let Some(BindingPlan::Parameter { destructuring_plan, .. }) = self - .loop_hoisting_table(module_idx) - .get_expr_binding_plan(module_idx, expr_idx) + .expression_binding_plan(expr_idx) .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? .cloned() { @@ -3563,8 +3585,7 @@ impl Interpreter { let expr_idx = expr.as_ref().eidx(); let expr_text = expr.span().text().to_string(); let binding_plan = self - .loop_hoisting_table(module_idx) - .get_expr_binding_plan(module_idx, expr_idx) + .expression_binding_plan(expr_idx) .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? .cloned() .ok_or_else(|| { diff --git a/src/languages/rego/compiler/comprehensions.rs b/src/languages/rego/compiler/comprehensions.rs index 00288dbd2..6157088da 100644 --- a/src/languages/rego/compiler/comprehensions.rs +++ b/src/languages/rego/compiler/comprehensions.rs @@ -46,6 +46,7 @@ impl<'a> Compiler<'a> { value_expr: value_expr.cloned(), span: span.clone(), key_value_loops_hoisted: false, + output_query: Some(query.qidx), }; self.push_context(context); self.compile_query(query)?; diff --git a/src/languages/rego/compiler/core.rs b/src/languages/rego/compiler/core.rs index 314a01b2f..22cbffb3a 100644 --- a/src/languages/rego/compiler/core.rs +++ b/src/languages/rego/compiler/core.rs @@ -116,6 +116,7 @@ impl<'a> Compiler<'a> { pub fn reset_rule_definition_registers(&mut self) { self.current_input_register = None; self.current_data_register = None; + self.loop_expr_register_map.clear(); } /// Push a new compilation context onto the context stack @@ -198,7 +199,13 @@ impl<'a> Compiler<'a> { self.policy .inner .loop_hoisting_table - .get_expr_binding_plan(module_idx, expr_idx) + .get_output_binding_plan( + module_idx, + self.context_stack + .last() + .and_then(|context| context.output_query), + expr_idx, + ) .map_err(|err| { CompilerError::General { message: format!("loop hoisting table out of bounds: {err}"), diff --git a/src/languages/rego/compiler/loops.rs b/src/languages/rego/compiler/loops.rs index 926c22fd7..508b00c3a 100644 --- a/src/languages/rego/compiler/loops.rs +++ b/src/languages/rego/compiler/loops.rs @@ -52,7 +52,13 @@ impl<'a> Compiler<'a> { .policy .inner .loop_hoisting_table - .get_expr_loops(module_idx, expr_idx) + .get_output_loops( + module_idx, + self.context_stack + .last() + .and_then(|context| context.output_query), + expr_idx, + ) .map_err(|err| { CompilerError::General { message: format!("loop hoisting table out of bounds: {err}"), @@ -160,6 +166,7 @@ impl<'a> Compiler<'a> { value_expr: None, span: span.clone(), key_value_loops_hoisted: false, + output_query: Some(query.qidx), }; self.push_context(every_context); diff --git a/src/languages/rego/compiler/mod.rs b/src/languages/rego/compiler/mod.rs index fa08510c6..6b425d39b 100644 --- a/src/languages/rego/compiler/mod.rs +++ b/src/languages/rego/compiler/mod.rs @@ -65,6 +65,7 @@ pub struct CompilationContext { pub(super) value_expr: Option, pub(super) span: Span, pub(super) key_value_loops_hoisted: bool, + pub(super) output_query: Option, } /// Entry in the rule compilation worklist that tracks both rule path and full call stack for recursion detection diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index 08fc7b0f7..cec2996cd 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -673,6 +673,7 @@ impl<'a> Compiler<'a> { value_expr, span, key_value_loops_hoisted: false, + output_query: None, }; self.push_context(context); let mut body_entry_points = Vec::new(); @@ -741,6 +742,7 @@ impl<'a> Compiler<'a> { if let Some(context) = self.context_stack.last_mut() { context.value_expr = body_value_expr.clone(); context.key_value_loops_hoisted = false; + context.output_query = Some(body.query.qidx); } self.emit_instruction( diff --git a/src/rvm/vm/rules.rs b/src/rvm/vm/rules.rs index 7921ca38b..3b0f74b45 100644 --- a/src/rvm/vm/rules.rs +++ b/src/rvm/vm/rules.rs @@ -708,11 +708,30 @@ impl RegoVM { } } - if let Some(definition_bodies) = rule_info - .definitions + let next_is_else = rule_info + .else_bodies .get(frame_data.current_definition_index) + .and_then(|bodies| { + frame_data + .current_body_index + .checked_add(1) + .and_then(|next| bodies.get(next)) + }) + .copied() + .unwrap_or(false); + if !matches!( + frame_data.rule_type, + RuleType::PartialSet | RuleType::PartialObject + ) || next_is_else { - frame_data.current_body_index = definition_bodies.len(); + if let Some(definition_bodies) = rule_info + .definitions + .get(frame_data.current_definition_index) + { + frame_data.current_body_index = definition_bodies.len(); + } else { + self.increment_counter(&mut frame_data.current_body_index, "rule body index")?; + } } else { self.increment_counter(&mut frame_data.current_body_index, "rule body index")?; } diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index d1374f7b5..e8eaa572a 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,92 @@ # Licensed under the MIT License. cases: + - note: output_head_body_specific_unbound_then_bound + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[i]] := i if { true } { i := 1 } + query: data.test.values + want_result: {a: 0, b: 1} + + - note: output_head_body_specific_unbound_then_bound_suspendable + execution_mode: suspendable + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[i]] := i if { true } { i := 1 } + query: data.test.values + want_result: {a: 0, b: 1} + + - note: output_head_body_specific_bound_then_unbound + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[i]] := i if { i := 0 } { true } + query: data.test.values + want_result: {a: 0, b: 1} + + - note: output_head_body_specific_bound_then_unbound_suspendable + execution_mode: suspendable + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[i]] := i if { i := 0 } { true } + query: data.test.values + want_result: {a: 0, b: 1} + + - note: output_head_body_specific_bound_then_failed + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[i]] := i if { i := 0 } { false } + query: data.test.values + want_result: {a: 0} + + - note: output_head_body_specific_bound_then_failed_suspendable + execution_mode: suspendable + input: {keys: ["a", "b"]} + modules: + - | + package test + values[input.keys[i]] := i if { i := 0 } { false } + query: data.test.values + want_result: {a: 0} + + - note: output_head_body_specific_paired_values + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[i]] := input.vals[i] if { true } { i := 1 } + query: data.test.values + want_result: {a: 10, b: 20} + + - note: output_head_body_specific_paired_values_suspendable + execution_mode: suspendable + input: {keys: ["a", "b"], vals: [10, 20]} + modules: + - | + package test + values[input.keys[i]] := input.vals[i] if { i := 0 } { true } + query: data.test.values + want_result: {a: 10, b: 20} + + - note: output_head_body_specific_builtin_values + execution_mode: suspendable + input: {keys: ["a", "b"], vals: ["A", "B"]} + modules: + - | + package test + values[input.keys[i]] := lower(input.vals[i]) if { true } { i := 1 } + query: data.test.values + want_result: {a: "a", b: "b"} + - note: keyed_partial_set_shared_head_index input: {keys: ["a", "b"], vals: [10, 20]} modules: From 4f283f6c387da0bce368e95859cb5a4ded9ed7b5 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:13:30 +0000 Subject: [PATCH 08/13] fix(rvm): include dynamic keyed sets in static child lookups Share lookup-only child definition expansions between evaluators and filter projected dynamic contributions by the requested key. Preserve declaration-based override suppression and avoid whole-parent calls that introduce sibling recursion. Resolve each compiled definition in its declaring package. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 5 +- src/compiled_policy.rs | 10 + src/interpreter.rs | 69 ++++++- src/languages/rego/compiler/comprehensions.rs | 1 + src/languages/rego/compiler/loops.rs | 1 + src/languages/rego/compiler/mod.rs | 2 + src/languages/rego/compiler/queries.rs | 28 +++ src/languages/rego/compiler/rules.rs | 60 ++++-- tests/engine/mod.rs | 26 +++ tests/rvm/rego/cases/keyed_partial_sets.yaml | 183 ++++++++++++++++++ 10 files changed, 368 insertions(+), 17 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index dc915f86e..e495e957f 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -106,7 +106,10 @@ Parameter structures: is selected by body query as well as module and expression, so differing body bindings cannot overwrite each other's head plans. Both execution modes accumulate all independent partial-rule bodies, while `else` bodies retain - first-success semantics. + first-success semantics. Static child references and entrypoints include + dynamic keyed-set definitions from their immediate parent. The compiler + filters those definitions by the requested key before set insertion, rather + than calling the entire parent and introducing cycles through sibling rules. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register pairs. Literal keys must be sorted to match template order. - `ArrayCreateParams` and `SetCreateParams` store register lists. The VM checks diff --git a/src/compiled_policy.rs b/src/compiled_policy.rs index 468426e4f..eaafa7709 100644 --- a/src/compiled_policy.rs +++ b/src/compiled_policy.rs @@ -216,6 +216,8 @@ pub(crate) struct CompiledPolicyData { pub(crate) parser_contexts: Rc>, pub(crate) schedule: Option>, pub(crate) rules: Map>>, + // Lookup-only expansions must not suppress dynamic siblings under a `with` override. + pub(crate) keyed_set_child_rules: Map>>, pub(crate) default_rules: Map>, pub(crate) rule_path_components: Map>, pub(crate) imports: BTreeMap>, @@ -241,3 +243,11 @@ pub(crate) struct CompiledPolicyData { // Pre-computed loop hoisting information pub(crate) loop_hoisting_table: HoistedLoopsLookup, } + +impl CompiledPolicyData { + pub(crate) fn rules_for_lookup(&self, path: &str) -> Option<&Vec>> { + self.keyed_set_child_rules + .get(path) + .or_else(|| self.rules.get(path)) + } +} diff --git a/src/interpreter.rs b/src/interpreter.rs index 10998c6e3..bc297aa4c 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -3303,7 +3303,7 @@ impl Interpreter { fn ensure_rule_evaluated(&mut self, path: String) -> Result<()> { self.check_execution_time()?; let mut matched = false; - if let Some(rules) = self.compiled_policy.rules.get(&path) { + if let Some(rules) = self.compiled_policy.rules_for_lookup(&path) { matched = true; for r in rules.clone() { if !self.processed.contains(&r) { @@ -4642,6 +4642,7 @@ impl Interpreter { fn clear_rule_registrations(&mut self) { let compiled_policy = self.compiled_policy_mut(); compiled_policy.rules.clear(); + compiled_policy.keyed_set_child_rules.clear(); compiled_policy.default_rules.clear(); compiled_policy.rule_paths.clear(); compiled_policy.rule_path_components.clear(); @@ -4742,7 +4743,7 @@ impl Interpreter { self.set_current_module(prev_module)?; registration_result?; } - Ok(()) + self.record_keyed_set_children() })(); if result.is_err() { self.clear_rule_registrations(); @@ -4750,6 +4751,70 @@ impl Interpreter { result } + // Include dynamic contributions in child lookups without evaluating unrelated siblings. + fn record_keyed_set_children(&mut self) -> Result<()> { + let mut dynamic_sets = Map::>>::new(); + for module in self.compiled_policy.modules.clone().iter() { + let package = get_path_string(&module.package.refr, Some("data"))?; + for rule in &module.policy { + let Rule::Spec { + head: + RuleHead::Set { + refr, key: Some(_), .. + }, + .. + } = rule.as_ref() + else { + continue; + }; + let Expr::RefBrack { + refr: prefix, + index, + .. + } = refr.as_ref() + else { + continue; + }; + if matches!(index.as_ref(), Expr::String { .. } | Expr::RawString { .. }) { + continue; + } + if get_rule_path_components(prefix)? + .iter() + .any(|component| !matches!(component, PathComponent::String(_))) + { + continue; + } + let parent = get_path_string(prefix, Some(&package))?; + dynamic_sets.entry(parent).or_default().push(rule.clone()); + self.memory_check()?; + } + } + if dynamic_sets.is_empty() { + return Ok(()); + } + + let mut child_rules = Map::new(); + let policy = self.compiled_policy.clone(); + for (path, definitions) in &policy.rules { + let components = policy + .rule_path_components + .get(path) + .ok_or_else(|| anyhow!("missing components for registered rule path {path}"))?; + if let Some((_, parent)) = components.split_last() { + let parent: Vec<&str> = parent.iter().map(String::as_str).collect(); + if let Some(rules) = dynamic_sets.get(&format_string_path(&parent)?) { + let mut definitions = definitions.clone(); + definitions.extend(rules.iter().cloned()); + child_rules.insert(path.clone(), definitions); + self.memory_check()?; + } + } + } + drop(policy); + self.compiled_policy_mut().keyed_set_child_rules = child_rules; + Ok(()) + } + pub fn add_extension( &mut self, path: String, diff --git a/src/languages/rego/compiler/comprehensions.rs b/src/languages/rego/compiler/comprehensions.rs index 6157088da..e485d8335 100644 --- a/src/languages/rego/compiler/comprehensions.rs +++ b/src/languages/rego/compiler/comprehensions.rs @@ -42,6 +42,7 @@ impl<'a> Compiler<'a> { context_type: ContextType::Comprehension(context_type), dest_register: result_reg, key_expr: key_expr.cloned(), + key_filter: None, multi_value: false, value_expr: value_expr.cloned(), span: span.clone(), diff --git a/src/languages/rego/compiler/loops.rs b/src/languages/rego/compiler/loops.rs index 508b00c3a..bebd4a48a 100644 --- a/src/languages/rego/compiler/loops.rs +++ b/src/languages/rego/compiler/loops.rs @@ -162,6 +162,7 @@ impl<'a> Compiler<'a> { context_type: ContextType::Every, dest_register: result_reg, key_expr: None, + key_filter: None, multi_value: false, value_expr: None, span: span.clone(), diff --git a/src/languages/rego/compiler/mod.rs b/src/languages/rego/compiler/mod.rs index 6b425d39b..e03e06831 100644 --- a/src/languages/rego/compiler/mod.rs +++ b/src/languages/rego/compiler/mod.rs @@ -61,6 +61,8 @@ pub struct CompilationContext { pub(super) context_type: ContextType, pub(super) dest_register: Register, pub(super) key_expr: Option, + /// Requested static child when projecting a dynamic keyed-set definition. + pub(super) key_filter: Option, pub(super) multi_value: bool, pub(super) value_expr: Option, pub(super) span: Span, diff --git a/src/languages/rego/compiler/queries.rs b/src/languages/rego/compiler/queries.rs index 4faf001d3..8a58e549a 100644 --- a/src/languages/rego/compiler/queries.rs +++ b/src/languages/rego/compiler/queries.rs @@ -171,6 +171,34 @@ impl<'a> Compiler<'a> { .map(|key_expr| self.compile_rego_expr(&key_expr)) .unwrap_or(Ok(value_register))?; + if let Some(key) = context.key_filter { + let literal_idx = self.add_literal(key); + let expected = self.alloc_register(); + self.emit_instruction( + Instruction::Load { + dest: expected, + literal_idx, + }, + span, + ); + let matches = self.alloc_register(); + self.emit_instruction( + Instruction::Eq { + dest: matches, + left: key_register, + right: expected, + }, + span, + ); + self.emit_instruction( + Instruction::Guard { + register: matches, + mode: GuardMode::Condition, + }, + span, + ); + } + match context.context_type { ContextType::Comprehension(ComprehensionType::Array) => { self.emit_instruction( diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index cec2996cd..61289f2a3 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -38,7 +38,7 @@ impl<'a> Compiler<'a> { } pub(super) fn compute_rule_type(&self, rule_path: &str) -> Result { - let Some(definitions) = self.policy.inner.rules.get(rule_path) else { + let Some(definitions) = self.policy.inner.rules_for_lookup(rule_path) else { // Default-only rules (e.g., `default deny := true`) have no regular definitions // in the `rules` map — they only exist in `default_rules`. Treat them as Complete. if self.policy.inner.default_rules.contains_key(rule_path) { @@ -60,9 +60,14 @@ impl<'a> Compiler<'a> { match head { RuleHead::Set { refr, key: Some(_), .. - } if self.keyed_set_key(rule_path, def, refr)?.is_some() => { - has_keyed_set = true; - RuleType::PartialObject + } => { + let (key, filter) = self.keyed_set_key(rule_path, def, refr)?; + has_keyed_set |= key.is_some(); + if key.is_some() && filter.is_none() { + RuleType::PartialObject + } else { + RuleType::PartialSet + } } RuleHead::Set { .. } => RuleType::PartialSet, RuleHead::Compr { refr, assign, .. } => match refr.as_ref() { @@ -112,7 +117,7 @@ impl<'a> Compiler<'a> { rule_path: &str, rule: &crate::ast::NodeRef, refr: &ExprRef, - ) -> Result> { + ) -> Result<(Option, Option)> { let (prefix, key) = match refr.as_ref() { Expr::RefBrack { refr, index, .. } => (refr, index.clone()), Expr::RefDot { @@ -125,7 +130,7 @@ impl<'a> Compiler<'a> { eidx: *eidx, }), ), - _ => return Ok(None), + _ => return Ok((None, None)), }; if Self::has_unsupported_bracket_prefix(prefix) { return Err(CompilerError::PartialObjectNestedKeyUnsupported.at(refr.span())); @@ -138,12 +143,27 @@ impl<'a> Compiler<'a> { if Self::is_string_literal(&key) && get_path_string(refr, Some(&package)).map_err(CompilerError::from)? == rule_path { - return Ok(None); + return Ok((None, None)); } - if get_path_string(prefix, Some(&package)).map_err(CompilerError::from)? != rule_path { + let prefix_path = get_path_string(prefix, Some(&package)).map_err(CompilerError::from)?; + if prefix_path != rule_path { + if !Self::is_string_literal(&key) { + if let Some((leaf, parent)) = self + .policy + .inner + .rule_path_components + .get(rule_path) + .and_then(|components| components.split_last()) + { + let parent: Vec<&str> = parent.iter().map(String::as_str).collect(); + if format_string_path(&parent).map_err(CompilerError::from)? == prefix_path { + return Ok((Some(key), Some(Value::from(leaf.clone())))); + } + } + } return Err(CompilerError::PartialObjectNestedKeyUnsupported.at(refr.span())); } - Ok(Some(key)) + Ok((Some(key), None)) } fn validate_partial_object_shape(&self, refr: &ExprRef) -> Result<()> { @@ -458,7 +478,7 @@ impl<'a> Compiler<'a> { self.current_module_index = module_index; let saved_register_counter = self.register_counter; - if let Some(rule_definitions) = rules.get(rule_path) { + if let Some(rule_definitions) = self.policy.inner.rules_for_lookup(rule_path) { let Some(rule_index) = self.rule_index_map.get(rule_path).copied() else { return Err(CompilerError::General { message: format!( @@ -509,7 +529,7 @@ impl<'a> Compiler<'a> { .. } = rule.as_ref() { - if let Some(key) = self.keyed_set_key(rule_path, rule, refr)? { + if let (Some(key), _) = self.keyed_set_key(rule_path, rule, refr)? { if let Expr::String { value, .. } | Expr::RawString { value, .. } = key.as_ref() { @@ -544,7 +564,15 @@ impl<'a> Compiler<'a> { let result_register = self.alloc_register(); self.current_module_index = self.find_module_index_for_rule(rule_ref)?; - + self.current_package = get_path_string( + &self.policy.get_modules()[self.current_module_index as usize] + .package + .refr, + Some("data"), + ) + .map_err(CompilerError::from)?; + + let mut key_filter = None; let (key_expr, value_expr) = match head { RuleHead::Compr { refr, assign, .. } => { if rule_type == RuleType::PartialObject { @@ -567,8 +595,11 @@ impl<'a> Compiler<'a> { self.rule_definition_destructuring_patterns[rule_index as usize] .push(None); - let key_expr = if rule_type == RuleType::PartialObject { - self.keyed_set_key(rule_path, rule_ref, refr)? + let key_expr = if key.is_some() { + let (key_expr, filter) = + self.keyed_set_key(rule_path, rule_ref, refr)?; + key_filter = filter; + key_expr } else { None }; @@ -669,6 +700,7 @@ impl<'a> Compiler<'a> { dest_register: result_register, context_type: ContextType::Rule(rule_type.clone()), key_expr, + key_filter, multi_value: matches!(head, RuleHead::Set { .. }), value_expr, span, diff --git a/tests/engine/mod.rs b/tests/engine/mod.rs index 9c71e982f..22035ff6c 100644 --- a/tests/engine/mod.rs +++ b/tests/engine/mod.rs @@ -3,6 +3,32 @@ use anyhow::{bail, Result}; use regorus::*; +use std::collections::BTreeSet; + +#[test] +fn keyed_set_child_override_preserves_dynamic_sibling_contributions() -> Result<()> { + let mut engine = Engine::new(); + engine.add_policy( + "test.rego".into(), + r#"package test + values.a contains 1 + values[k] contains 2 if { k := "b" } + main := result if { + result := values.b with data.test.values.a as {9} + }"# + .into(), + )?; + + assert_eq!( + engine.eval_rule("data.test.main".into())?, + Value::from(BTreeSet::from([Value::from(2)])) + ); + assert_eq!( + engine.eval_rule("data.test.values.a".into())?, + Value::from(BTreeSet::from([Value::from(1)])) + ); + Ok(()) +} #[test] fn package_components_distinguish_literal_dots_from_nested_namespaces() -> Result<()> { diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index e8eaa572a..eca119ef8 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,189 @@ # Licensed under the MIT License. cases: + - note: overlapping_keyed_set_local_child + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { k := "a" } + main := values.a + query: data.test.main + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_local_child_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { k := "a" } + main := values.a + query: data.test.main + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_child_entrypoint + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { k := "a" } + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_child_entrypoint_suspendable + execution_mode: suspendable + modules: + - | + package test + values[k] contains 2 if { k := "a" } + values["a"] contains 1 + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_external_child + modules: + - | + package test + main := data.other.values.a + - | + package other + values.a contains 1 + values[k] contains 2 if { k := "a" } + query: data.test.main + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_imported_child_suspendable + execution_mode: suspendable + modules: + - | + package test + import data.other.values.a as child + main := child + - | + package other + values.a contains 1 + - | + package other + values[k] contains 2 if { k := "a" } + query: data.test.main + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_parent_then_child + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { k := "a" } + main := [values, values.a] + query: data.test.main + want_result: + - a: + set!: [1, 2] + - set!: [1, 2] + + - note: overlapping_keyed_set_child_nested_package_suspendable + execution_mode: suspendable + modules: + - | + package test + extra := 2 + values[k] contains extra if { k := "a" } + - | + package test.values + a contains 1 + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_child_then_parent_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { k := "a" } + main := [values.a, values] + query: data.test.main + want_result: + - set!: [1, 2] + - a: + set!: [1, 2] + + - note: overlapping_keyed_set_child_head_loop + input: {keys: ["a", "b", "a"], vals: [2, 99, 3]} + modules: + - | + package test + values.a contains 1 + values[input.keys[i]] contains input.vals[i] + query: data.test.values.a + want_result: + set!: [1, 2, 3] + + - note: overlapping_keyed_set_child_failed_dynamic_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 if { false } + values[k] contains 2 if { k := "a"; false } + query: data.test.values.a + want_result: + set!: [] + + - note: overlapping_keyed_set_child_head_loop_suspendable + execution_mode: suspendable + input: {keys: ["a", "b", "a"], vals: [false, 99, null]} + modules: + - | + package test + values.a contains 1 + values[input.keys[i]] contains input.vals[i] if { true } + query: data.test.values.a + want_result: + set!: [1, false, null] + + - note: overlapping_keyed_set_child_sibling_dependency + modules: + - | + package test + values.a contains count(values.b) + values.b contains 4 + values[k] contains 2 if { k := "a" } + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_child_sibling_dependency_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains count(values.b) + values.b contains 4 + values[k] contains 2 if { k := "a" } + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_escaped_child_suspendable + execution_mode: suspendable + modules: + - | + package test + values["a.b"] contains 1 + values[k] contains 2 if { k := "a.b" } + query: data.test.values["a.b"] + want_result: + set!: [1, 2] + - note: output_head_body_specific_unbound_then_bound input: {keys: ["a", "b"]} modules: From cbe0d9579cb31a4dbb4f961d2b24bb448a16058c Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:52:38 +0000 Subject: [PATCH 09/13] fix(rvm): filter projected keyed-set definitions early Index dynamic definitions once per keyed parent and merge them lazily for static child lookups, avoiding per-child copies. Reject nonmatching projected keys before dependent body expressions when a constant key binding or some-in binding is available. Add regressions for nonmatching child dependencies in both execution modes, including reversed statement order and some-in bindings. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 9 +- src/compiled_policy.rs | 62 +++++++++-- src/compiler/hoist.rs | 99 ++++++++++++++++- src/interpreter.rs | 104 +++++++++++++++--- src/languages/rego/compiler/comprehensions.rs | 1 + src/languages/rego/compiler/loops.rs | 3 + src/languages/rego/compiler/mod.rs | 1 + src/languages/rego/compiler/queries.rs | 79 ++++++++++++- src/languages/rego/compiler/rules.rs | 3 +- tests/rvm/rego/cases/keyed_partial_sets.yaml | 81 ++++++++++++++ 10 files changed, 411 insertions(+), 31 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index e495e957f..0485ff29f 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -107,9 +107,12 @@ Parameter structures: bindings cannot overwrite each other's head plans. Both execution modes accumulate all independent partial-rule bodies, while `else` bodies retain first-success semantics. Static child references and entrypoints include - dynamic keyed-set definitions from their immediate parent. The compiler - filters those definitions by the requested key before set insertion, rather - than calling the entire parent and introducing cycles through sibling rules. + dynamic keyed-set definitions indexed once per parent and merged only when a + child is looked up. A known constant child key is bound and checked before + unrelated body statements, avoiding false recursion through nonmatching + dynamic definitions. The compiler also filters those definitions by the + requested key before set insertion, rather than calling the entire parent + and introducing cycles through sibling rules. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register pairs. Literal keys must be sorted to match template order. - `ArrayCreateParams` and `SetCreateParams` store register lists. The VM checks diff --git a/src/compiled_policy.rs b/src/compiled_policy.rs index eaafa7709..e10d5c698 100644 --- a/src/compiled_policy.rs +++ b/src/compiled_policy.rs @@ -15,7 +15,8 @@ use crate::scheduler::*; use crate::utils::*; use crate::*; -use alloc::collections::BTreeMap; +use alloc::borrow::Cow; +use alloc::collections::{BTreeMap, BTreeSet}; use anyhow::Result; #[cfg(feature = "azure_policy")] @@ -216,8 +217,9 @@ pub(crate) struct CompiledPolicyData { pub(crate) parser_contexts: Rc>, pub(crate) schedule: Option>, pub(crate) rules: Map>>, - // Lookup-only expansions must not suppress dynamic siblings under a `with` override. - pub(crate) keyed_set_child_rules: Map>>, + // Dynamic keyed-set definitions are indexed once per parent for child lookups. + pub(crate) keyed_set_parent_rules: Map, + pub(crate) keyed_set_child_parents: Map, pub(crate) default_rules: Map>, pub(crate) rule_path_components: Map>, pub(crate) imports: BTreeMap>, @@ -244,10 +246,56 @@ pub(crate) struct CompiledPolicyData { pub(crate) loop_hoisting_table: HoistedLoopsLookup, } +#[derive(Debug, Clone, Default)] +pub(crate) struct KeyedSetParentRules { + pub(crate) definitions: Vec>, + rule_set: BTreeSet>, +} + +impl KeyedSetParentRules { + pub(crate) fn insert(&mut self, rule: Ref) { + if self.rule_set.insert(rule.clone()) { + self.definitions.push(rule); + } + } + + pub(crate) fn contains(&self, rule: &Ref) -> bool { + self.rule_set.contains(rule) + } +} + impl CompiledPolicyData { - pub(crate) fn rules_for_lookup(&self, path: &str) -> Option<&Vec>> { - self.keyed_set_child_rules - .get(path) - .or_else(|| self.rules.get(path)) + pub(crate) fn rules_for_lookup(&self, path: &str) -> Option]>> { + let static_rules = self.rules.get(path)?; + let Some(parent) = self.keyed_set_child_parents.get(path) else { + return Some(Cow::Borrowed(static_rules)); + }; + let Some(dynamic_rules) = self.keyed_set_parent_rules.get(parent) else { + return Some(Cow::Borrowed(static_rules)); + }; + + let mut rules = static_rules.clone(); + rules.extend(dynamic_rules.definitions.iter().cloned()); + Some(Cow::Owned(rules)) + } + + pub(crate) fn keyed_set_child_filter(&self, path: &str, rule: &Ref) -> Option { + let parent = self.keyed_set_child_parents.get(path)?; + if !self.keyed_set_parent_rules.get(parent)?.contains(rule) { + return None; + } + self.rule_path_components + .get(path)? + .last() + .map(|key| Value::from(key.clone())) + } + + pub(crate) fn is_projected_keyed_set_rule(&self, path: &str, rule: &Ref) -> bool { + let Some(parent) = self.keyed_set_child_parents.get(path) else { + return false; + }; + self.keyed_set_parent_rules + .get(parent) + .is_some_and(|rules| rules.contains(rule)) } } diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index 192b78134..7635c2a72 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -16,7 +16,9 @@ use super::destructuring_planner::{ map_binding_error, BindingPlan, ScopingMode, VariableBindingContext, }; -use crate::ast::{Expr, ExprRef, Literal, LiteralStmt, Module, Query, Ref, Rule, RuleHead}; +use crate::ast::{ + AssignOp, Expr, ExprRef, Literal, LiteralStmt, Module, Query, Ref, Rule, RuleHead, +}; use crate::compiler::context::{ContextType, ScopeContext}; use crate::lookup::Lookup; use crate::lookup::LookupResult; @@ -27,6 +29,101 @@ use anyhow::{anyhow, Result}; use alloc::collections::{BTreeMap, BTreeSet}; use alloc::vec::Vec; +pub fn prioritize_key_binding( + statements: &mut Vec<&LiteralStmt>, + key_expr: Option<&ExprRef>, + key_filter: Option<&Value>, +) -> Result<()> { + if key_filter.is_none() { + return Ok(()); + } + let Some(Expr::Var { span, .. }) = key_expr.map(|expr| expr.as_ref()) else { + return Ok(()); + }; + let key_name = span.text(); + let binding_idx = statements.iter().position(|stmt| { + if !stmt.with_mods.is_empty() { + return false; + } + let Literal::Expr { expr, .. } = &stmt.literal else { + return false; + }; + let Expr::AssignExpr { + op: AssignOp::ColEq, + lhs, + rhs, + .. + } = expr.as_ref() + else { + return false; + }; + let Expr::Var { span: lhs_span, .. } = lhs.as_ref() else { + return false; + }; + lhs_span.text() == key_name && is_scalar_literal(rhs.as_ref()) + }); + if let Some(idx) = binding_idx { + if idx != 0 { + for stmt in statements.iter().take(idx) { + if stmt_references_var(stmt, key_name)? { + return Ok(()); + } + } + if let Some(stmt) = statements.get(idx).copied() { + statements.remove(idx); + statements.insert(0, stmt); + } + } + } + Ok(()) +} + +fn stmt_references_var(stmt: &LiteralStmt, name: &str) -> Result { + let mut found = false; + let mut check_expr = |expr: &ExprRef| -> Result { + if matches!(expr.as_ref(), Expr::Var { span, .. } if span.text() == name) { + found = true; + Ok(false) + } else { + Ok(true) + } + }; + match &stmt.literal { + Literal::Expr { expr, .. } | Literal::NotExpr { expr, .. } => { + crate::query::traversal::traverse(expr, &mut check_expr)?; + } + Literal::SomeIn { + key, + value, + collection, + .. + } => { + if let Some(key) = key { + crate::query::traversal::traverse(key, &mut check_expr)?; + } + crate::query::traversal::traverse(value, &mut check_expr)?; + crate::query::traversal::traverse(collection, &mut check_expr)?; + } + Literal::SomeVars { vars, .. } => { + found = vars.iter().any(|var| var.text() == name); + } + Literal::Every { .. } => return Ok(true), + } + Ok(found) +} + +fn is_scalar_literal(expr: &Expr) -> bool { + match expr { + Expr::String { .. } + | Expr::RawString { .. } + | Expr::Number { .. } + | Expr::Bool { .. } + | Expr::Null { .. } => true, + Expr::UnaryExpr { expr, .. } => matches!(expr.as_ref(), Expr::Number { .. }), + _ => false, + } +} + /// Implementation of VariableBindingContext for ScopeContext impl VariableBindingContext for ScopeContext { fn is_var_unbound(&self, var_name: &str, scoping: ScopingMode) -> bool { diff --git a/src/interpreter.rs b/src/interpreter.rs index bc297aa4c..fdaebd5fc 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -191,6 +191,7 @@ struct Context { output_constness_determined: bool, early_return: bool, output_query: Option, + key_filter: Option, } impl Default for Context { @@ -209,6 +210,7 @@ impl Default for Context { output_constness_determined: false, early_return: false, output_query: None, + key_filter: None, } } } @@ -1072,6 +1074,11 @@ impl Interpreter { continue; } + if self.reject_nonmatching_child_key()? { + *self.current_scope_mut()? = scope_saved.clone(); + continue; + } + let mut should_break = false; if self.eval_stmts(stmts)? { count = count.saturating_add(1); @@ -1108,6 +1115,11 @@ impl Interpreter { continue; } + if self.reject_nonmatching_child_key()? { + *self.current_scope_mut()? = scope_saved.clone(); + continue; + } + let mut should_break = false; if self.eval_stmts(stmts)? { count = count.saturating_add(1); @@ -1145,6 +1157,11 @@ impl Interpreter { continue; } + if self.reject_nonmatching_child_key()? { + *self.current_scope_mut()? = scope_saved.clone(); + continue; + } + let mut should_break = false; if self.eval_stmts(stmts)? { count = count.saturating_add(1); @@ -2336,6 +2353,9 @@ impl Interpreter { .unwrap_or(&[]); eval_success = self.eval_stmt(stmt, tail)?; + if eval_success && self.reject_nonmatching_child_key()? { + eval_success = false; + } if matches!(&stmt.literal, Literal::SomeIn { .. }) { return Ok(eval_success); @@ -2377,6 +2397,24 @@ impl Interpreter { Ok(eval_success) } + fn reject_nonmatching_child_key(&mut self) -> Result { + let Some((key_expr, filter)) = self.contexts.last().and_then(|ctx| { + ctx.key_filter + .as_ref() + .zip(ctx.key_expr.as_ref()) + .map(|(filter, expr)| (expr.clone(), filter.clone())) + }) else { + return Ok(false); + }; + let Expr::Var { span, .. } = key_expr.as_ref() else { + return Ok(false); + }; + if self.lookup_local_var(&span.source_str()).is_none() { + return Ok(false); + } + Ok(self.eval_expr(&key_expr)? != filter) + } + fn eval_query(&mut self, query: &Ref) -> Result { self.check_execution_time()?; self.get_current_context_mut()?.output_query = Some(query.qidx); @@ -2429,7 +2467,7 @@ impl Interpreter { } }; - let ordered_stmts: Vec<&LiteralStmt> = match order_indices { + let mut ordered_stmts: Vec<&LiteralStmt> = match order_indices { Some(order) => { let stmts_len = query.stmts.len(); if order.len() != stmts_len { @@ -2456,6 +2494,13 @@ impl Interpreter { } None => query.stmts.iter().collect(), }; + if let Some(context) = self.contexts.last() { + crate::compiler::hoist::prioritize_key_binding( + &mut ordered_stmts, + context.key_expr.as_ref(), + context.key_filter.as_ref(), + )?; + } let r = self.eval_stmts(&ordered_stmts); self.scopes.pop(); @@ -3303,12 +3348,20 @@ impl Interpreter { fn ensure_rule_evaluated(&mut self, path: String) -> Result<()> { self.check_execution_time()?; let mut matched = false; - if let Some(rules) = self.compiled_policy.rules_for_lookup(&path) { + let lookup_rules = self + .compiled_policy + .rules_for_lookup(&path) + .map(|rules| rules.into_owned()); + if let Some(rules) = lookup_rules { matched = true; - for r in rules.clone() { + for r in rules { if !self.processed.contains(&r) { let module = self.get_rule_module(&r)?; - self.eval_rule(&module, &r)?; + let key_filter = self.compiled_policy.keyed_set_child_filter(&path, &r); + self.eval_rule_with_key_filter(&module, &r, key_filter)?; + if self.compiled_policy.is_projected_keyed_set_rule(&path, &r) { + self.processed.remove(&r); + } } } } @@ -4199,7 +4252,12 @@ impl Interpreter { Ok(()) } - fn eval_rule_impl(&mut self, module: &Ref, rule: &Ref) -> Result<()> { + fn eval_rule_impl( + &mut self, + module: &Ref, + rule: &Ref, + key_filter: Option, + ) -> Result<()> { self.check_execution_time()?; match rule.as_ref() { Rule::Spec { @@ -4209,7 +4267,8 @@ impl Interpreter { } => { match rule_head { RuleHead::Compr { refr, .. } | RuleHead::Set { refr, .. } => { - let (ctx, _) = self.make_rule_context(rule_head)?; + let (mut ctx, _) = self.make_rule_context(rule_head)?; + ctx.key_filter = key_filter; let is_set = ctx.is_set; let is_keyed_set = is_set && ctx.key_expr.is_some(); let is_object = ctx.key_expr.is_some() && !is_set; @@ -4309,6 +4368,15 @@ impl Interpreter { } pub fn eval_rule(&mut self, module: &Ref, rule: &Ref) -> Result<()> { + self.eval_rule_with_key_filter(module, rule, None) + } + + fn eval_rule_with_key_filter( + &mut self, + module: &Ref, + rule: &Ref, + key_filter: Option, + ) -> Result<()> { self.check_execution_time()?; // Set current module index self.current_module_index = self.find_module_index(module); @@ -4351,7 +4419,7 @@ impl Interpreter { let scopes = core::mem::take(&mut self.scopes); let prev_module = self.set_current_module(Some(module.clone()))?; - let res = self.eval_rule_impl(module, rule); + let res = self.eval_rule_impl(module, rule, key_filter); self.set_current_module(prev_module)?; self.scopes = scopes; @@ -4642,7 +4710,8 @@ impl Interpreter { fn clear_rule_registrations(&mut self) { let compiled_policy = self.compiled_policy_mut(); compiled_policy.rules.clear(); - compiled_policy.keyed_set_child_rules.clear(); + compiled_policy.keyed_set_parent_rules.clear(); + compiled_policy.keyed_set_child_parents.clear(); compiled_policy.default_rules.clear(); compiled_policy.rule_paths.clear(); compiled_policy.rule_path_components.clear(); @@ -4753,7 +4822,7 @@ impl Interpreter { // Include dynamic contributions in child lookups without evaluating unrelated siblings. fn record_keyed_set_children(&mut self) -> Result<()> { - let mut dynamic_sets = Map::>>::new(); + let mut dynamic_sets = Map::::new(); for module in self.compiled_policy.modules.clone().iter() { let package = get_path_string(&module.package.refr, Some("data"))?; for rule in &module.policy { @@ -4785,7 +4854,7 @@ impl Interpreter { continue; } let parent = get_path_string(prefix, Some(&package))?; - dynamic_sets.entry(parent).or_default().push(rule.clone()); + dynamic_sets.entry(parent).or_default().insert(rule.clone()); self.memory_check()?; } } @@ -4793,25 +4862,26 @@ impl Interpreter { return Ok(()); } - let mut child_rules = Map::new(); + let mut child_parents = Map::new(); let policy = self.compiled_policy.clone(); - for (path, definitions) in &policy.rules { + for path in policy.rules.keys() { let components = policy .rule_path_components .get(path) .ok_or_else(|| anyhow!("missing components for registered rule path {path}"))?; if let Some((_, parent)) = components.split_last() { let parent: Vec<&str> = parent.iter().map(String::as_str).collect(); - if let Some(rules) = dynamic_sets.get(&format_string_path(&parent)?) { - let mut definitions = definitions.clone(); - definitions.extend(rules.iter().cloned()); - child_rules.insert(path.clone(), definitions); + let parent = format_string_path(&parent)?; + if dynamic_sets.contains_key(&parent) { + child_parents.insert(path.clone(), parent); self.memory_check()?; } } } drop(policy); - self.compiled_policy_mut().keyed_set_child_rules = child_rules; + let compiled_policy = self.compiled_policy_mut(); + compiled_policy.keyed_set_parent_rules = dynamic_sets; + compiled_policy.keyed_set_child_parents = child_parents; Ok(()) } diff --git a/src/languages/rego/compiler/comprehensions.rs b/src/languages/rego/compiler/comprehensions.rs index e485d8335..277149808 100644 --- a/src/languages/rego/compiler/comprehensions.rs +++ b/src/languages/rego/compiler/comprehensions.rs @@ -43,6 +43,7 @@ impl<'a> Compiler<'a> { dest_register: result_reg, key_expr: key_expr.cloned(), key_filter: None, + key_filter_guarded: false, multi_value: false, value_expr: value_expr.cloned(), span: span.clone(), diff --git a/src/languages/rego/compiler/loops.rs b/src/languages/rego/compiler/loops.rs index bebd4a48a..737edb43f 100644 --- a/src/languages/rego/compiler/loops.rs +++ b/src/languages/rego/compiler/loops.rs @@ -163,6 +163,7 @@ impl<'a> Compiler<'a> { dest_register: result_reg, key_expr: None, key_filter: None, + key_filter_guarded: false, multi_value: false, value_expr: None, span: span.clone(), @@ -305,6 +306,7 @@ impl<'a> Compiler<'a> { .apply_binding_plan(binding_plan, key_reg, plan_span) .map_err(|e| CompilerError::from(e).at(plan_span))?; } + self.emit_key_filter_guard(collection.span())?; let body_stmts = &remaining_stmts[0..]; self.compile_hoisted_loops(body_stmts, remaining_loops)?; @@ -475,6 +477,7 @@ impl<'a> Compiler<'a> { } } + self.emit_key_filter_guard(collection.span())?; self.hoist_loops_and_compile_statements(loop_body_stmts)?; self.emit_instruction( diff --git a/src/languages/rego/compiler/mod.rs b/src/languages/rego/compiler/mod.rs index e03e06831..9550057c2 100644 --- a/src/languages/rego/compiler/mod.rs +++ b/src/languages/rego/compiler/mod.rs @@ -63,6 +63,7 @@ pub struct CompilationContext { pub(super) key_expr: Option, /// Requested static child when projecting a dynamic keyed-set definition. pub(super) key_filter: Option, + pub(super) key_filter_guarded: bool, pub(super) multi_value: bool, pub(super) value_expr: Option, pub(super) span: Span, diff --git a/src/languages/rego/compiler/queries.rs b/src/languages/rego/compiler/queries.rs index 8a58e549a..b3c862c68 100644 --- a/src/languages/rego/compiler/queries.rs +++ b/src/languages/rego/compiler/queries.rs @@ -8,10 +8,12 @@ use super::{Compiler, CompilerError, ComprehensionType, ContextType, Result}; use crate::ast::{self, LiteralStmt, Query}; +use crate::compiler::hoist::prioritize_key_binding; use crate::rvm::instructions::GuardMode; use crate::rvm::program::RuleType; use crate::rvm::Instruction; use alloc::format; +use alloc::string::ToString as _; use alloc::vec::Vec; impl<'a> Compiler<'a> { @@ -32,7 +34,7 @@ impl<'a> Compiler<'a> { None => None, }; - let ordered_stmts: Vec<&LiteralStmt> = match schedule { + let mut ordered_stmts: Vec<&LiteralStmt> = match schedule { Some(schedule) => schedule .order .iter() @@ -40,6 +42,19 @@ impl<'a> Compiler<'a> { .collect(), None => query.stmts.iter().collect(), }; + if let Some(context) = self.context_stack.last() { + prioritize_key_binding( + &mut ordered_stmts, + context.key_expr.as_ref(), + context.key_filter.as_ref(), + ) + .map_err(|err| { + CompilerError::General { + message: err.to_string(), + } + .at(&query.span) + })?; + } self.hoist_loops_and_compile_statements(&ordered_stmts) }; @@ -59,7 +74,9 @@ impl<'a> Compiler<'a> { let loop_exprs = self.get_statement_loops(stmt)?; if !loop_exprs.is_empty() { - return self.compile_hoisted_loops(&stmts[idx..], &loop_exprs); + self.compile_hoisted_loops(&stmts[idx..], &loop_exprs)?; + self.emit_key_filter_guard(&stmt.span)?; + return Ok(()); } if matches!(&stmt.literal, ast::Literal::SomeIn { .. }) { @@ -95,16 +112,74 @@ impl<'a> Compiler<'a> { &stmt.span, ); } + self.emit_key_filter_guard(&stmt.span)?; return Ok(()); } } self.compile_single_statement(stmt)?; + self.emit_key_filter_guard(&stmt.span)?; } self.hoist_loops_and_emit_context_yield() } + pub(super) fn emit_key_filter_guard(&mut self, span: &crate::lexer::Span) -> Result<()> { + let Some((var_name, filter)) = self.context_stack.last().and_then(|context| { + if context.key_filter_guarded { + return None; + } + let (ContextType::Rule(_), Some(filter), Some(expr)) = ( + &context.context_type, + context.key_filter.as_ref(), + context.key_expr.as_ref(), + ) else { + return None; + }; + let crate::ast::Expr::Var { span: key_span, .. } = expr.as_ref() else { + return None; + }; + Some((key_span.text().to_string(), filter.clone())) + }) else { + return Ok(()); + }; + if !self.is_var_bound_in_current_scope(&var_name) { + return Ok(()); + } + let Some(key_register) = self.lookup_variable(&var_name) else { + return Ok(()); + }; + let literal_idx = self.add_literal(filter); + let expected = self.alloc_register(); + self.emit_instruction( + Instruction::Load { + dest: expected, + literal_idx, + }, + span, + ); + let matches = self.alloc_register(); + self.emit_instruction( + Instruction::Eq { + dest: matches, + left: key_register, + right: expected, + }, + span, + ); + self.emit_instruction( + Instruction::Guard { + register: matches, + mode: GuardMode::Condition, + }, + span, + ); + if let Some(context) = self.context_stack.last_mut() { + context.key_filter_guarded = true; + } + Ok(()) + } + pub(super) fn hoist_loops_and_emit_context_yield(&mut self) -> Result<()> { if let Some(context) = self.context_stack.last() { match &context.context_type { diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index 61289f2a3..44a96383e 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -520,7 +520,7 @@ impl<'a> Compiler<'a> { // Seed all static leaves once, before any body can fail or accumulate members. let mut initial_object = Value::new_object(); if rule_type == RuleType::PartialObject { - for rule in rule_definitions { + for rule in rule_definitions.iter() { if let Rule::Spec { head: RuleHead::Set { @@ -701,6 +701,7 @@ impl<'a> Compiler<'a> { context_type: ContextType::Rule(rule_type.clone()), key_expr, key_filter, + key_filter_guarded: false, multi_value: matches!(head, RuleHead::Set { .. }), value_expr, span, diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index eca119ef8..fbdbb6641 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,87 @@ # Licensed under the MIT License. cases: + - note: overlapping_keyed_set_nonmatching_dependent_child + modules: + - | + package test + values.a contains 1 + values[k] contains n if { + k := "b" + n := count(values.a) + } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_nonmatching_dependent_child_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { + k := "b" + n := count(values.a) + } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_nonmatching_dependency_before_key + modules: + - | + package test + values.a contains 1 + values[k] contains n if { + n := count(values.a) + k := "b" + } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_nonmatching_dependency_before_key_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { + n := count(values.a) + k := "b" + } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_nonmatching_some_in_dependency + modules: + - | + package test + values.a contains 1 + values[k] contains n if { + some k in ["b"] + n := count(values.a) + } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_nonmatching_some_in_dependency_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { + some k in ["b"] + n := count(values.a) + } + query: data.test.values.a + want_result: + set!: [1] + - note: overlapping_keyed_set_local_child modules: - | From 76c502b349bd8e504bc66ae7892a5c984ce67393 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Tue, 20 Oct 2026 07:39:08 +0000 Subject: [PATCH 10/13] fix(rvm): preserve keyed projection bindings and recursion checks Defer child filters until declared keys are defined and check hoisted bindings before dependent statements. Track interpreter recursion per projection, reset compiler guard state per body, and propagate genuine projected cycles as typed runtime errors in both VM modes using the existing recursion-check flag. Cover unification, parent and sibling dependencies, indexed assignments, lookup order, genuine cycles, and the unsupported multi-body set parser boundary. Review the patch and cumulative implementation locally before committing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 7 + src/interpreter.rs | 31 ++- src/languages/rego/compiler/loops.rs | 1 + src/languages/rego/compiler/rules.rs | 4 + src/rvm/vm/errors.rs | 3 + src/rvm/vm/rules.rs | 22 +- tests/engine/mod.rs | 58 +++++ tests/rvm/rego/cases/keyed_partial_sets.yaml | 211 +++++++++++++++++++ 8 files changed, 327 insertions(+), 10 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index 0485ff29f..2988b4a97 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -113,6 +113,13 @@ Parameter structures: dynamic definitions. The compiler also filters those definitions by the requested key before set insertion, rather than calling the entire parent and introducing cycles through sibling rules. + Early filters defer declared-but-undefined keys until binding and run after + hoisted assignments as well as ordinary statements. Interpreter recursion + tracking distinguishes parent and child projections of the same definition, + while repeated evaluation of the same projection remains a cycle. + Programs containing child projections enable runtime recursion checking in + both VM modes; genuine cycles propagate an error rather than becoming + undefined results. - `ObjectCreateParams` reuses arrays of literal key/value pairs and register pairs. Literal keys must be sorted to match template order. - `ArrayCreateParams` and `SetCreateParams` store register lists. The VM checks diff --git a/src/interpreter.rs b/src/interpreter.rs index fdaebd5fc..15dbc3245 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -119,7 +119,8 @@ pub struct Interpreter { processed: BTreeSet>, processed_paths: Value, rule_values: RuleValues, - active_rules: Vec>, + // A parent invocation and its filtered child projections have distinct identities. + active_rules: Vec<(Ref, Option)>, builtins_cache: BTreeMap<(&'static str, Vec), Value>, no_rules_lookup: bool, execution_timer: ExecutionTimer, @@ -1630,9 +1631,14 @@ impl Interpreter { self.memory_check()?; self.check_execution_time()?; if loops.is_empty() { + if self.reject_nonmatching_child_key()? { + return Ok(false); + } if let Some((first_stmt, tail_stmts)) = stmts.split_first() { // Evaluate the current statement whose loop expressions have been hoisted. - if self.eval_stmt(first_stmt, tail_stmts)? { + if self.eval_stmt(first_stmt, tail_stmts)? + && !self.reject_nonmatching_child_key()? + { if !matches!(&first_stmt.literal, Literal::SomeIn { .. }) { self.eval_stmts(tail_stmts) } else { @@ -2409,10 +2415,10 @@ impl Interpreter { let Expr::Var { span, .. } = key_expr.as_ref() else { return Ok(false); }; - if self.lookup_local_var(&span.source_str()).is_none() { + let Some(key) = self.lookup_local_var(&span.source_str()) else { return Ok(false); - } - Ok(self.eval_expr(&key_expr)? != filter) + }; + Ok(key != Value::Undefined && key != filter) } fn eval_query(&mut self, query: &Ref) -> Result { @@ -4391,10 +4397,17 @@ impl Interpreter { return Ok(()); } - self.active_rules.push(rule.clone()); - if self.active_rules.iter().filter(|&r| r == rule).count() == 2 { + let invocation = (rule.clone(), key_filter.clone()); + self.active_rules.push(invocation.clone()); + if self + .active_rules + .iter() + .filter(|&r| r == &invocation) + .count() + == 2 + { let mut msg = String::default(); - for r in &self.active_rules { + for (r, _) in &self.active_rules { let refr = Self::get_rule_refr(r); let span = refr.span(); msg.push_str( @@ -4424,7 +4437,7 @@ impl Interpreter { self.set_current_module(prev_module)?; self.scopes = scopes; match self.active_rules.pop() { - Some(ref r) if r == rule => res, + Some(ref r) if r == &invocation => res, _ => bail!("internal error: current rule not active"), } } diff --git a/src/languages/rego/compiler/loops.rs b/src/languages/rego/compiler/loops.rs index 737edb43f..1d6f4f2b5 100644 --- a/src/languages/rego/compiler/loops.rs +++ b/src/languages/rego/compiler/loops.rs @@ -86,6 +86,7 @@ impl<'a> Compiler<'a> { if loops.is_empty() { if !stmts.is_empty() { self.compile_single_statement(stmts[0])?; + self.emit_key_filter_guard(&stmts[0].span)?; return self.hoist_loops_and_compile_statements(&stmts[1..]); } else { return self.hoist_loops_and_emit_context_yield(); diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index 44a96383e..6260a06ee 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -599,6 +599,9 @@ impl<'a> Compiler<'a> { let (key_expr, filter) = self.keyed_set_key(rule_path, rule_ref, refr)?; key_filter = filter; + if key_filter.is_some() { + self.program.needs_runtime_recursion_check = true; + } key_expr } else { None @@ -775,6 +778,7 @@ impl<'a> Compiler<'a> { if let Some(context) = self.context_stack.last_mut() { context.value_expr = body_value_expr.clone(); context.key_value_loops_hoisted = false; + context.key_filter_guarded = false; context.output_query = Some(body.query.qidx); } diff --git a/src/rvm/vm/errors.rs b/src/rvm/vm/errors.rs index cdbd753f7..9a5355645 100644 --- a/src/rvm/vm/errors.rs +++ b/src/rvm/vm/errors.rs @@ -87,6 +87,9 @@ pub enum VmError { available: usize, }, + #[error("Runtime recursion detected in rule {rule_index} (pc={pc})")] + RuleRecursionDetected { rule_index: u16, pc: usize }, + #[error("Invalid object create params index: {index} (pc={pc}, available={available})")] InvalidObjectCreateParams { index: u16, diff --git a/src/rvm/vm/rules.rs b/src/rvm/vm/rules.rs index 3b0f74b45..d3a4cb8ff 100644 --- a/src/rvm/vm/rules.rs +++ b/src/rvm/vm/rules.rs @@ -17,7 +17,7 @@ use super::execution_model::{ use super::machine::RegoVM; impl RegoVM { - /// Returns true if the error represents a resource-limit violation that + /// Returns true if the error represents a fatal evaluation failure that /// must never be silently absorbed by rule evaluation. pub(super) const fn is_fatal_vm_error(err: &VmError) -> bool { matches!( @@ -27,9 +27,25 @@ impl RegoVM { | VmError::MemoryBudgetExceeded { .. } | VmError::RegexSizeLimitExceeded { .. } | VmError::InstructionLimitExceeded { .. } + | VmError::RuleRecursionDetected { .. } ) } + fn check_rule_recursion(&self, rule_index: u16) -> Result<()> { + if self.program.needs_runtime_recursion_check + && self + .call_rule_stack + .iter() + .any(|context| context.rule_index == rule_index) + { + return Err(VmError::RuleRecursionDetected { + rule_index, + pc: self.pc, + }); + } + Ok(()) + } + /// Restore VM state that was swapped out for rule execution. /// Must be called before returning an error from `execute_rule_definitions_common` /// to avoid leaving the VM in an inconsistent state. @@ -249,6 +265,8 @@ impl RegoVM { } } + self.check_rule_recursion(rule_index)?; + let rule_type = rule_info.rule_type.clone(); let rule_definitions = rule_info.definitions.clone(); @@ -416,6 +434,8 @@ impl RegoVM { } } + self.check_rule_recursion(rule_index)?; + if rule_info.definitions.is_empty() { // No compiled definitions — check for a default value before returning Undefined. let result = rule_info diff --git a/tests/engine/mod.rs b/tests/engine/mod.rs index 22035ff6c..c158842d5 100644 --- a/tests/engine/mod.rs +++ b/tests/engine/mod.rs @@ -5,6 +5,64 @@ use anyhow::{bail, Result}; use regorus::*; use std::collections::BTreeSet; +#[test] +fn keyed_set_projection_preserves_cycle_detection() -> Result<()> { + for (policy, entrypoint) in [ + ( + r#"package test + values.a contains 1 + values[k] contains n if { k := "a"; n := count(values.a) }"#, + "data.test.values.a", + ), + ( + r#"package test + values.a contains 1 + values[k] contains n if { k := "a"; n := count(values.a) }"#, + "data.test.values", + ), + ( + r#"package test + values.a contains 1 + values.b contains 2 + values[k] contains n if { k := "a"; n := count(values.b) } + values[k] contains n if { k := "b"; n := count(values.a) }"#, + "data.test.values.a", + ), + ] { + let mut engine = Engine::new(); + engine.add_policy("test.rego".into(), policy.into())?; + let error = engine + .eval_rule(entrypoint.into()) + .expect_err("a genuine keyed-set cycle must be rejected"); + assert!( + error.to_string().contains("recursion detected"), + "{entrypoint}: {error}" + ); + } + Ok(()) +} + +#[test] +fn keyed_set_multiple_bodies_remain_rejected() { + let mut engine = Engine::new(); + let error = engine + .add_policy( + "test.rego".into(), + r#"package test + values.a contains 1 + values[k] contains n if { k := "b"; n := 2 } + { k := "b"; n := count(values.a) }"# + .into(), + ) + .expect_err("multiple set bodies are unsupported"); + assert!( + error + .to_string() + .contains("else cannot be used with set rules"), + "{error}" + ); +} + #[test] fn keyed_set_child_override_preserves_dynamic_sibling_contributions() -> Result<()> { let mut engine = Engine::new(); diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index fbdbb6641..fc9100995 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,217 @@ # Licensed under the MIT License. cases: + - note: overlapping_keyed_set_matching_cycle + modules: + - | + package test + values.a contains 1 + values[k] contains n if { k := "a"; n := count(values.a) } + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: overlapping_keyed_set_matching_cycle_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { k := "a"; n := count(values.a) } + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: overlapping_keyed_set_sibling_cycle + modules: + - | + package test + values.a contains 1 + values.b contains 2 + values[k] contains n if { k := "a"; n := count(values.b) } + values[k] contains n if { k := "b"; n := count(values.a) } + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: overlapping_keyed_set_sibling_cycle_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values.b contains 2 + values[k] contains n if { k := "a"; n := count(values.b) } + values[k] contains n if { k := "b"; n := count(values.a) } + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: overlapping_keyed_set_index_bound_before_dependency + modules: + - | + package test + values.a contains 1 + values[k] contains n if { n := {"b": 2}[k] + count(values.a) } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_index_bound_before_dependency_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { n := {"b": 2}[k] + count(values.a) } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_dependent_lookup_order + modules: + - | + package test + values.a contains 0 + values.b contains 4 + values[k] contains n if { k := "a"; n := count(values.b) } + main := [values.a, values, values.b] + query: data.test.main + want_result: + - set!: [0, 1] + - a: + set!: [0, 1] + b: + set!: [4] + - set!: [4] + + - note: overlapping_keyed_set_dependent_lookup_order_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 0 + values.b contains 4 + values[k] contains n if { k := "a"; n := count(values.b) } + main := [values, values.b, values.a] + query: data.test.main + want_result: + - a: + set!: [0, 1] + b: + set!: [4] + - set!: [4] + - set!: [0, 1] + + - note: overlapping_keyed_set_declared_key_unification + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { some k; k = "a" } + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_declared_key_unification_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { some k; k = "a" } + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_parent_dependent_child + modules: + - | + package test + values.a contains 1 + values[k] contains n if { k := "b"; n := count(values.a) } + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [1] + + - note: overlapping_keyed_set_parent_dependent_child_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { k := "b"; n := count(values.a) } + query: data.test.values + want_result: + a: + set!: [1] + b: + set!: [1] + + - note: overlapping_keyed_set_sibling_projection + modules: + - | + package test + values.a contains 0 + values.b contains 4 + values[k] contains n if { k := "a"; n := count(values.b) } + query: data.test.values.a + want_result: + set!: [0, 1] + + - note: overlapping_keyed_set_sibling_projection_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 0 + values.b contains 4 + values[k] contains n if { k := "a"; n := count(values.b) } + query: data.test.values.a + want_result: + set!: [0, 1] + + - note: overlapping_keyed_set_indexed_assignment + modules: + - | + package test + values.a contains 1 + values[k] contains n if { k := ["b"][_]; n := count(values.a) } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_indexed_assignment_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { k := ["b"][_]; n := count(values.a) } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_indexed_matching_assignment + modules: + - | + package test + values.a contains 1 + values[k] contains n if { k := ["b", "a"][_]; n := 2 } + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_indexed_matching_assignment_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { k := ["b", "a"][_]; n := 2 } + query: data.test.values.a + want_result: + set!: [1, 2] + - note: overlapping_keyed_set_nonmatching_dependent_child modules: - | From 6c9551d728997f17ff6c37d5683e44f174712f79 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:11:32 +0000 Subject: [PATCH 11/13] fix: guard expression-key projections before dependencies Share conservative key readiness across evaluators, filter each output-loop level, and fence actual projected output keys before values and accumulation. Prioritize scalar unification bindings in either operand order. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/rvm/instruction-set.md | 11 +- src/compiler/hoist.rs | 40 +++- src/interpreter.rs | 36 +++- src/languages/rego/compiler/queries.rs | 97 +++++---- tests/engine/mod.rs | 31 +++ tests/rvm/rego/cases/keyed_partial_sets.yaml | 203 +++++++++++++++++++ 6 files changed, 359 insertions(+), 59 deletions(-) diff --git a/docs/rvm/instruction-set.md b/docs/rvm/instruction-set.md index 2988b4a97..1fe967838 100644 --- a/docs/rvm/instruction-set.md +++ b/docs/rvm/instruction-set.md @@ -113,8 +113,15 @@ Parameter structures: dynamic definitions. The compiler also filters those definitions by the requested key before set insertion, rather than calling the entire parent and introducing cycles through sibling rules. - Early filters defer declared-but-undefined keys until binding and run after - hoisted assignments as well as ordinary statements. Interpreter recursion + Both evaluators check ready input/reference keys before body evaluation, + between hoisted loops, and before output evaluation. Early filters defer + declared-but-undefined locals until binding and run after hoisted assignments + as well as ordinary statements. Scalar key bindings using `:=` or `=` (in + either operand order for `=`) can precede unrelated dependent statements. + Calls, comprehensions, and unresolved references conservatively defer early + filtering. Both evaluators always check the actual output key before + evaluating its value or accumulating a projected contribution, preventing + sibling writes even when early filtering is deferred. Interpreter recursion tracking distinguishes parent and child projections of the same definition, while repeated evaluation of the same projection remains a cycle. Programs containing child projections enable runtime recursion checking in diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index 7635c2a72..68becc272 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -29,6 +29,35 @@ use anyhow::{anyhow, Result}; use alloc::collections::{BTreeMap, BTreeSet}; use alloc::vec::Vec; +pub fn key_expr_is_ready( + expr: &ExprRef, + mut is_bound: impl FnMut(&ExprRef) -> Result, + mut is_var_bound: impl FnMut(&crate::lexer::Span) -> bool, +) -> Result { + let mut ready = true; + crate::query::traversal::traverse(expr, &mut |node| { + if is_bound(node)? { + return Ok(false); + } + match node.as_ref() { + Expr::Var { span, .. } => { + ready &= span.text() == "input" || is_var_bound(span); + Ok(false) + } + Expr::ArrayCompr { .. } + | Expr::SetCompr { .. } + | Expr::ObjectCompr { .. } + | Expr::Call { .. } + | Expr::AssignExpr { .. } => { + ready = false; + Ok(false) + } + _ => Ok(true), + } + })?; + Ok(ready) +} + pub fn prioritize_key_binding( statements: &mut Vec<&LiteralStmt>, key_expr: Option<&ExprRef>, @@ -49,7 +78,7 @@ pub fn prioritize_key_binding( return false; }; let Expr::AssignExpr { - op: AssignOp::ColEq, + op, lhs, rhs, .. @@ -57,10 +86,11 @@ pub fn prioritize_key_binding( else { return false; }; - let Expr::Var { span: lhs_span, .. } = lhs.as_ref() else { - return false; - }; - lhs_span.text() == key_name && is_scalar_literal(rhs.as_ref()) + matches!(lhs.as_ref(), Expr::Var { span: lhs_span, .. } if lhs_span.text() == key_name) + && is_scalar_literal(rhs.as_ref()) + || *op == AssignOp::Eq + && matches!(rhs.as_ref(), Expr::Var { span: rhs_span, .. } if rhs_span.text() == key_name) + && is_scalar_literal(lhs.as_ref()) }); if let Some(idx) = binding_idx { if idx != 0 { diff --git a/src/interpreter.rs b/src/interpreter.rs index 15dbc3245..c5eb844b9 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -2012,6 +2012,9 @@ impl Interpreter { fn eval_output_expr_in_loop(&mut self, loops: &[HoistedLoop]) -> Result { self.check_execution_time()?; + if self.reject_nonmatching_child_key()? { + return Ok(false); + } if loops.is_empty() { let (key_expr, output_expr) = self.get_exprs_from_context()?; @@ -2035,7 +2038,16 @@ impl Interpreter { let mut comps = self.eval_rule_ref(&rule_ref)?; if let Some(ke) = &key_expr { is_const_rule = is_const_rule && Self::is_simple_literal(ke)?; - comps.push(self.eval_expr(ke)?); + let key = self.eval_expr(ke)?; + if self + .get_current_context()? + .key_filter + .as_ref() + .is_some_and(|filter| filter != &key) + { + return Ok(false); + } + comps.push(key); } let output = if let Some(oe) = &output_expr { // Rule is constant only if its ref, key and output are constant. @@ -2262,6 +2274,9 @@ impl Interpreter { fn eval_output_expr(&mut self) -> Result { self.check_execution_time()?; + if self.reject_nonmatching_child_key()? { + return Ok(false); + } // Evaluate output expression after all the statements have been executed. let (key_expr, output_expr) = self.get_exprs_from_context()?; @@ -2320,6 +2335,9 @@ impl Interpreter { } fn eval_stmts(&mut self, stmts: &[&LiteralStmt]) -> Result { + if self.reject_nonmatching_child_key()? { + return Ok(false); + } let mut eval_success = true; for (idx, stmt) in stmts.iter().enumerate() { @@ -2412,13 +2430,17 @@ impl Interpreter { }) else { return Ok(false); }; - let Expr::Var { span, .. } = key_expr.as_ref() else { - return Ok(false); - }; - let Some(key) = self.lookup_local_var(&span.source_str()) else { + if !crate::compiler::hoist::key_expr_is_ready( + &key_expr, + |expr| self.get_loop_var_value(expr).map(|value| value.is_some()), + |span| { + self.lookup_local_var(&span.source_str()) + .is_some_and(|value| value != Value::Undefined) + }, + )? { return Ok(false); - }; - Ok(key != Value::Undefined && key != filter) + } + Ok(self.eval_expr(&key_expr)? != filter) } fn eval_query(&mut self, query: &Ref) -> Result { diff --git a/src/languages/rego/compiler/queries.rs b/src/languages/rego/compiler/queries.rs index b3c862c68..661e90331 100644 --- a/src/languages/rego/compiler/queries.rs +++ b/src/languages/rego/compiler/queries.rs @@ -67,6 +67,10 @@ impl<'a> Compiler<'a> { &mut self, stmts: &[&LiteralStmt], ) -> Result<()> { + if let Some(context) = self.context_stack.last() { + let span = context.span.clone(); + self.emit_key_filter_guard(&span)?; + } for (idx, stmt) in stmts.iter().enumerate() { if !stmt.with_mods.is_empty() { return Err(CompilerError::WithKeywordUnsupported.at(&stmt.span)); @@ -125,7 +129,7 @@ impl<'a> Compiler<'a> { } pub(super) fn emit_key_filter_guard(&mut self, span: &crate::lexer::Span) -> Result<()> { - let Some((var_name, filter)) = self.context_stack.last().and_then(|context| { + let Some((expr, filter)) = self.context_stack.last().and_then(|context| { if context.key_filter_guarded { return None; } @@ -136,19 +140,33 @@ impl<'a> Compiler<'a> { ) else { return None; }; - let crate::ast::Expr::Var { span: key_span, .. } = expr.as_ref() else { - return None; - }; - Some((key_span.text().to_string(), filter.clone())) + Some((expr.clone(), filter.clone())) }) else { return Ok(()); }; - if !self.is_var_bound_in_current_scope(&var_name) { + if !crate::compiler::hoist::key_expr_is_ready( + &expr, + |node| Ok(self.loop_expr_register_map.contains_key(node)), + |key_span| self.lookup_variable(key_span.text()).is_some(), + ) + .map_err(CompilerError::from)? + { return Ok(()); } - let Some(key_register) = self.lookup_variable(&var_name) else { - return Ok(()); - }; + let key_register = self.compile_rego_expr(&expr)?; + self.emit_key_match_guard(key_register, filter, span); + if let Some(context) = self.context_stack.last_mut() { + context.key_filter_guarded = true; + } + Ok(()) + } + + fn emit_key_match_guard( + &mut self, + key_register: super::Register, + filter: crate::Value, + span: &crate::lexer::Span, + ) { let literal_idx = self.add_literal(filter); let expected = self.alloc_register(); self.emit_instruction( @@ -174,13 +192,13 @@ impl<'a> Compiler<'a> { }, span, ); - if let Some(context) = self.context_stack.last_mut() { - context.key_filter_guarded = true; - } - Ok(()) } pub(super) fn hoist_loops_and_emit_context_yield(&mut self) -> Result<()> { + if let Some(context) = self.context_stack.last() { + let span = context.span.clone(); + self.emit_key_filter_guard(&span)?; + } if let Some(context) = self.context_stack.last() { match &context.context_type { ContextType::Every => { @@ -226,6 +244,20 @@ impl<'a> Compiler<'a> { if let Some(context) = self.context_stack.last().cloned() { let dest_register = context.dest_register; let span = &context.span; + let projected_key_register = if context.key_filter.is_some() { + context + .key_expr + .as_ref() + .map(|expr| self.compile_rego_expr(expr)) + .transpose()? + } else { + None + }; + if let (Some(key_register), Some(key)) = + (projected_key_register, context.key_filter.clone()) + { + self.emit_key_match_guard(key_register, key, span); + } let value_register = match context.value_expr { Some(expr) => self.compile_rego_expr(&expr)?, None => { @@ -241,38 +273,13 @@ impl<'a> Compiler<'a> { } }; - let key_register = context - .key_expr - .map(|key_expr| self.compile_rego_expr(&key_expr)) - .unwrap_or(Ok(value_register))?; - - if let Some(key) = context.key_filter { - let literal_idx = self.add_literal(key); - let expected = self.alloc_register(); - self.emit_instruction( - Instruction::Load { - dest: expected, - literal_idx, - }, - span, - ); - let matches = self.alloc_register(); - self.emit_instruction( - Instruction::Eq { - dest: matches, - left: key_register, - right: expected, - }, - span, - ); - self.emit_instruction( - Instruction::Guard { - register: matches, - mode: GuardMode::Condition, - }, - span, - ); - } + let key_register = match projected_key_register { + Some(register) => register, + None => context + .key_expr + .map(|key_expr| self.compile_rego_expr(&key_expr)) + .unwrap_or(Ok(value_register))?, + }; match context.context_type { ContextType::Comprehension(ComprehensionType::Array) => { diff --git a/tests/engine/mod.rs b/tests/engine/mod.rs index c158842d5..878628a6d 100644 --- a/tests/engine/mod.rs +++ b/tests/engine/mod.rs @@ -5,6 +5,37 @@ use anyhow::{bail, Result}; use regorus::*; use std::collections::BTreeSet; +#[test] +fn keyed_set_expression_projection_preserves_sibling_override() -> Result<()> { + for head in [ + "values[input.key] contains 2", + "values[input.keys[_]] contains 2", + "values[input.keys[_]] contains 2 if { true }", + ] { + let mut engine = Engine::new(); + engine.add_policy( + "test.rego".into(), + format!("package test\nvalues.a contains 1\n{head}"), + )?; + engine.set_input(Value::from_json_str(r#"{"key":"b","keys":["a","b"]}"#)?); + let result = engine.eval_query( + "[data.test.values.a, data.test.values.b] with data.test.values.b as {9}".into(), + false, + )?; + let first = if head.contains("keys") { + vec![1, 2] + } else { + vec![1] + }; + let expected = Value::from(vec![ + Value::from(first.into_iter().map(Value::from).collect::>()), + Value::from(BTreeSet::from([Value::from(9)])), + ]); + assert_eq!(result.result[0].expressions[0].value, expected, "{head}"); + } + Ok(()) +} + #[test] fn keyed_set_projection_preserves_cycle_detection() -> Result<()> { for (policy, entrypoint) in [ diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index fc9100995..045aa68f1 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,209 @@ # Licensed under the MIT License. cases: + - note: overlapping_keyed_set_head_value_loop_dependency + input: {keys: [b, c]} + modules: + - | + package test + values.a contains 1 + values[input.keys[_]] contains [count(values.a)][_] + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_head_value_loop_dependency_suspendable + execution_mode: suspendable + input: {keys: [b, c]} + modules: + - | + package test + values.a contains 1 + values[input.keys[_]] contains [count(values.a)][_] + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_input_key_dependency + input: {key: b} + modules: + - | + package test + values.a contains 1 + values[input.key] contains n if { n := count(values.a) } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_input_key_dependency_suspendable + execution_mode: suspendable + input: {key: b} + modules: + - | + package test + values.a contains 1 + values[input.key] contains n if { n := count(values.a) } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_bodyless_input_key_dependency + input: {key: b} + modules: + - | + package test + values.a contains 1 + values[input.key] contains count(values.a) + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_bodyless_input_key_dependency_suspendable + execution_mode: suspendable + input: {key: b} + modules: + - | + package test + values.a contains 1 + values[input.key] contains count(values.a) + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_missing_input_key_dependency + input: {} + modules: + - | + package test + values.a contains 1 + values[input.key] contains count(values.a) + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_missing_input_key_dependency_suspendable + execution_mode: suspendable + input: {} + modules: + - | + package test + values.a contains 1 + values[input.key] contains count(values.a) + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_unification_before_dependency + modules: + - | + package test + values.a contains 1 + values[z] contains n if { n := count(values.a); z = "b" } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_unification_before_dependency_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[z] contains n if { n := count(values.a); z = "b" } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_reversed_unification_before_dependency + modules: + - | + package test + values.a contains 1 + values[z] contains n if { n := count(values.a); "b" = z } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_reversed_unification_before_dependency_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[z] contains n if { n := count(values.a); "b" = z } + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_input_head_loop_dependency + input: {keys: [b, c]} + modules: + - | + package test + values.a contains 1 + values[input.keys[_]] contains count(values.a) + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_input_head_loop_dependency_suspendable + execution_mode: suspendable + input: {keys: [b, c]} + modules: + - | + package test + values.a contains 1 + values[input.keys[_]] contains count(values.a) + query: data.test.values.a + want_result: + set!: [1] + + - note: overlapping_keyed_set_matching_input_key + input: {key: a} + modules: + - | + package test + values.a contains 1 + values[input.key] contains 2 + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_matching_input_key_suspendable + execution_mode: suspendable + input: {key: a} + modules: + - | + package test + values.a contains 1 + values[input.key] contains 2 + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_later_bound_expression_key + input: {keys: {chosen: a}} + modules: + - | + package test + values.a contains 1 + values[input.keys[k]] contains 2 if { k := "chosen" } + query: data.test.values.a + want_result: + set!: [1, 2] + + - note: overlapping_keyed_set_later_bound_expression_key_suspendable + execution_mode: suspendable + input: {keys: {chosen: a}} + modules: + - | + package test + values.a contains 1 + values[input.keys[k]] contains 2 if { k := "chosen" } + query: data.test.values.a + want_result: + set!: [1, 2] + - note: overlapping_keyed_set_matching_cycle modules: - | From e757d007f5e2d4619ebb3c33021965222e4a5c77 Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:56:44 +0000 Subject: [PATCH 12/13] fix(rvm): address keyed projection review findings Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/compiler/hoist.rs | 110 +++++--- src/interpreter.rs | 205 +++++++++++++- src/languages/rego/compiler/comprehensions.rs | 1 + src/languages/rego/compiler/function_calls.rs | 2 +- src/languages/rego/compiler/loops.rs | 1 + src/languages/rego/compiler/mod.rs | 1 + src/languages/rego/compiler/queries.rs | 83 +++++- src/languages/rego/compiler/rules.rs | 2 + src/rvm/instructions/mod.rs | 1 + src/rvm/vm/errors.rs | 1 + tests/engine/mod.rs | 113 ++++++++ tests/rvm/rego/cases/keyed_partial_sets.yaml | 251 ++++++++++++++++++ 12 files changed, 734 insertions(+), 37 deletions(-) diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index 68becc272..0854a8cda 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -33,6 +33,7 @@ pub fn key_expr_is_ready( expr: &ExprRef, mut is_bound: impl FnMut(&ExprRef) -> Result, mut is_var_bound: impl FnMut(&crate::lexer::Span) -> bool, + mut is_safe_call: impl FnMut(&ExprRef) -> Result, ) -> Result { let mut ready = true; crate::query::traversal::traverse(expr, &mut |node| { @@ -47,61 +48,108 @@ pub fn key_expr_is_ready( Expr::ArrayCompr { .. } | Expr::SetCompr { .. } | Expr::ObjectCompr { .. } - | Expr::Call { .. } + | Expr::BinExpr { .. } + | Expr::UnaryExpr { .. } + | Expr::ArithExpr { .. } | Expr::AssignExpr { .. } => { ready = false; Ok(false) } + Expr::Call { fcn, params, .. } => { + ready &= is_safe_literal_key_call(fcn, params)? && is_safe_call(node)?; + Ok(false) + } _ => Ok(true), } })?; Ok(ready) } +pub fn is_safe_literal_key_call(fcn: &ExprRef, params: &[ExprRef]) -> Result { + let name = crate::utils::get_path_string(fcn, None)?; + let is_string = + |expr: &ExprRef| matches!(expr.as_ref(), Expr::String { .. } | Expr::RawString { .. }); + Ok(match (name.as_str(), params) { + ("lower", [value]) => is_string(value), + ("concat", [separator, values]) => { + is_string(separator) + && matches!( + values.as_ref(), + Expr::Array { items, .. } if items.iter().all(is_string) + ) + } + _ => false, + }) +} + +fn is_safe_key_binding_value( + expr: &ExprRef, + is_safe_call: &mut impl FnMut(&ExprRef) -> Result, +) -> Result { + match expr.as_ref() { + Expr::Call { fcn, params, .. } => { + Ok(is_safe_literal_key_call(fcn, params)? && is_safe_call(expr)?) + } + other => Ok(is_scalar_literal(other)), + } +} + pub fn prioritize_key_binding( statements: &mut Vec<&LiteralStmt>, key_expr: Option<&ExprRef>, key_filter: Option<&Value>, + mut is_safe_call: impl FnMut(&ExprRef) -> Result, ) -> Result<()> { if key_filter.is_none() { return Ok(()); } - let Some(Expr::Var { span, .. }) = key_expr.map(|expr| expr.as_ref()) else { + let Some(key_expr) = key_expr else { return Ok(()); }; - let key_name = span.text(); - let binding_idx = statements.iter().position(|stmt| { - if !stmt.with_mods.is_empty() { - return false; + let mut key_names = BTreeSet::new(); + crate::query::traversal::traverse(key_expr, &mut |node| { + if let Expr::Var { span, .. } = node.as_ref() { + let name = span.text(); + if name != "_" && name != "input" { + key_names.insert(name.to_string()); + } } - let Literal::Expr { expr, .. } = &stmt.literal else { - return false; - }; - let Expr::AssignExpr { - op, - lhs, - rhs, - .. - } = expr.as_ref() - else { - return false; - }; - matches!(lhs.as_ref(), Expr::Var { span: lhs_span, .. } if lhs_span.text() == key_name) - && is_scalar_literal(rhs.as_ref()) - || *op == AssignOp::Eq - && matches!(rhs.as_ref(), Expr::Var { span: rhs_span, .. } if rhs_span.text() == key_name) - && is_scalar_literal(lhs.as_ref()) - }); - if let Some(idx) = binding_idx { - if idx != 0 { + Ok(true) + })?; + for key_name in key_names { + let mut binding_idx = None; + for (idx, stmt) in statements.iter().enumerate() { + if !stmt.with_mods.is_empty() { + continue; + } + let Literal::Expr { expr, .. } = &stmt.literal else { + continue; + }; + let Expr::AssignExpr { op, lhs, rhs, .. } = expr.as_ref() else { + continue; + }; + if matches!(lhs.as_ref(), Expr::Var { span, .. } if span.text() == key_name) + && is_safe_key_binding_value(rhs, &mut is_safe_call)? + || *op == AssignOp::Eq + && matches!(rhs.as_ref(), Expr::Var { span, .. } if span.text() == key_name) + && is_safe_key_binding_value(lhs, &mut is_safe_call)? + { + binding_idx = Some(idx); + } + } + if let Some(idx) = binding_idx { + let mut referenced_earlier = false; for stmt in statements.iter().take(idx) { - if stmt_references_var(stmt, key_name)? { - return Ok(()); + if stmt_references_var(stmt, &key_name)? { + referenced_earlier = true; + break; } } - if let Some(stmt) = statements.get(idx).copied() { - statements.remove(idx); - statements.insert(0, stmt); + if idx != 0 && !referenced_earlier { + if let Some(stmt) = statements.get(idx).copied() { + statements.remove(idx); + statements.insert(0, stmt); + } } } } diff --git a/src/interpreter.rs b/src/interpreter.rs index c5eb844b9..1a8a872ce 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -556,6 +556,15 @@ impl Interpreter { } } + fn ensure_loop_var_values_capacity_for_module( + &mut self, + module_idx: u32, + num_expressions: u32, + ) { + self.loop_var_values + .ensure_capacity(module_idx, num_expressions); + } + fn current_module(&self) -> Result> { self.module .clone() @@ -2421,6 +2430,25 @@ impl Interpreter { Ok(eval_success) } + fn is_safe_projected_key_call(&self, expr: &ExprRef) -> Result { + let Expr::Call { fcn, params, .. } = expr.as_ref() else { + return Ok(false); + }; + if !crate::compiler::hoist::is_safe_literal_key_call(fcn, params)? { + return Ok(false); + } + let original_path = crate::utils::get_path_string(fcn, None)?; + let full_path = self + .resolve_fcn_path_through_imports(&original_path) + .unwrap_or_else(|| format!("{}.{}", self.current_module_path, original_path)); + Ok(!self.extensions.contains_key(&original_path) + && !self.with_functions.contains_key(&original_path) + && !self.with_functions.contains_key(&full_path) + && self.lookup_function_by_name(&original_path).is_none() + && self.lookup_function_by_name(&full_path).is_none() + && Self::is_builtin(fcn.span(), &original_path)) + } + fn reject_nonmatching_child_key(&mut self) -> Result { let Some((key_expr, filter)) = self.contexts.last().and_then(|ctx| { ctx.key_filter @@ -2437,12 +2465,170 @@ impl Interpreter { self.lookup_local_var(&span.source_str()) .is_some_and(|value| value != Value::Undefined) }, + |call| self.is_safe_projected_key_call(call), )? { return Ok(false); } Ok(self.eval_expr(&key_expr)? != filter) } + fn projected_head_key_loops_match(&mut self) -> Result> { + let Some((key_expr, filter, query_idx)) = self.contexts.last().and_then(|context| { + context + .key_filter + .as_ref() + .zip(context.key_expr.as_ref()) + .zip(context.output_query) + .map(|((filter, expr), query_idx)| (expr.clone(), filter.clone(), query_idx)) + }) else { + return Ok(None); + }; + let loops = self + .loop_hoisting_table(self.current_module_index) + .get_output_loops(self.current_module_index, Some(query_idx), key_expr.eidx()) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? + .cloned() + .unwrap_or_default(); + if loops.is_empty() + || loops + .iter() + .any(|loop_info| loop_info.loop_type != LoopType::IndexIteration) + { + return Ok(None); + } + self.scopes.push(Scope::default()); + let result = self.key_matches_in_hoisted_loops(&loops, &key_expr, &filter); + self.scopes.pop(); + result + } + + fn key_matches_in_hoisted_loops( + &mut self, + loops: &[HoistedLoop], + key_expr: &ExprRef, + filter: &Value, + ) -> Result> { + let Some((loop_info, remaining)) = loops.split_first() else { + if !crate::compiler::hoist::key_expr_is_ready( + key_expr, + |expr| self.get_loop_var_value(expr).map(|value| value.is_some()), + |span| { + self.lookup_local_var(&span.source_str()) + .is_some_and(|value| value != Value::Undefined) + }, + |call| self.is_safe_projected_key_call(call), + )? { + return Ok(None); + } + return Ok(Some(self.eval_expr(key_expr)? == *filter)); + }; + if !crate::compiler::hoist::key_expr_is_ready( + &loop_info.collection, + |expr| self.get_loop_var_value(expr).map(|value| value.is_some()), + |span| { + self.lookup_local_var(&span.source_str()) + .is_some_and(|value| value != Value::Undefined) + }, + |call| self.is_safe_projected_key_call(call), + )? { + return Ok(None); + } + let collection = self.eval_expr(Self::loop_collection_expr(loop_info))?; + let index_plan = if Self::loop_index_expr(loop_info).is_some() { + self.get_loop_index_plan(loop_info)? + } else { + DestructuringPlan::Ignore + }; + let loop_target = Self::loop_assignment_expr(loop_info).clone(); + self.scopes.push(Scope::default()); + let result = (|| { + let mut unresolved = false; + match collection { + Value::Array(values) => { + for (index, value) in values.iter().enumerate() { + self.memory_check()?; + match self.key_matches_for_hoisted_loop_item( + &loop_target, + &index_plan, + &Value::from(index), + value, + remaining, + key_expr, + filter, + )? { + Some(true) => return Ok(Some(true)), + Some(false) => {} + None => unresolved = true, + } + } + } + Value::Set(values) => { + for value in values.iter() { + self.memory_check()?; + match self.key_matches_for_hoisted_loop_item( + &loop_target, + &index_plan, + value, + value, + remaining, + key_expr, + filter, + )? { + Some(true) => return Ok(Some(true)), + Some(false) => {} + None => unresolved = true, + } + } + } + Value::Object(values) => { + for (key, value) in values.iter() { + self.memory_check()?; + match self.key_matches_for_hoisted_loop_item( + &loop_target, + &index_plan, + key, + value, + remaining, + key_expr, + filter, + )? { + Some(true) => return Ok(Some(true)), + Some(false) => {} + None => unresolved = true, + } + } + } + Value::Undefined => {} + _ => {} + } + Ok(if unresolved { None } else { Some(false) }) + })(); + self.scopes.pop(); + self.remove_loop_var_value(&loop_target); + result + } + + #[allow(clippy::too_many_arguments)] + fn key_matches_for_hoisted_loop_item( + &mut self, + loop_target: &ExprRef, + index_plan: &DestructuringPlan, + index: &Value, + value: &Value, + remaining: &[HoistedLoop], + key_expr: &ExprRef, + filter: &Value, + ) -> Result> { + self.set_loop_var_value(loop_target, value.clone())?; + let matches = if self.execute_destructuring_plan(index_plan, index)? == Value::Bool(true) { + self.key_matches_in_hoisted_loops(remaining, key_expr, filter)? + } else { + Some(false) + }; + Self::clear_scope(self.current_scope_mut()?); + Ok(matches) + } + fn eval_query(&mut self, query: &Ref) -> Result { self.check_execution_time()?; self.get_current_context_mut()?.output_query = Some(query.qidx); @@ -2527,6 +2713,7 @@ impl Interpreter { &mut ordered_stmts, context.key_expr.as_ref(), context.key_filter.as_ref(), + |call| self.is_safe_projected_key_call(call), )?; } @@ -3854,7 +4041,14 @@ impl Interpreter { } self.contexts.push(next_ctx); } - result = self.eval_query(&body.query); + if let Some(context) = self.contexts.last_mut() { + context.output_query = Some(body.query.qidx); + } + result = match self.projected_head_key_loops_match() { + Ok(Some(false)) => Ok(false), + Ok(_) => self.eval_query(&body.query), + Err(error) => Err(error), + }; match &result { Ok(true) => { any_success = true; @@ -4453,11 +4647,18 @@ impl Interpreter { // the local variables of callee function. let scopes = core::mem::take(&mut self.scopes); let prev_module = self.set_current_module(Some(module.clone()))?; + let loop_var_values = core::mem::take(&mut self.loop_var_values); + self.ensure_loop_var_values_capacity_for_module( + self.current_module_index, + module.num_expressions, + ); let res = self.eval_rule_impl(module, rule, key_filter); - self.set_current_module(prev_module)?; + let restore_module = self.set_current_module(prev_module); self.scopes = scopes; + self.loop_var_values = loop_var_values; + restore_module?; match self.active_rules.pop() { Some(ref r) if r == &invocation => res, _ => bail!("internal error: current rule not active"), diff --git a/src/languages/rego/compiler/comprehensions.rs b/src/languages/rego/compiler/comprehensions.rs index 277149808..89d5e063f 100644 --- a/src/languages/rego/compiler/comprehensions.rs +++ b/src/languages/rego/compiler/comprehensions.rs @@ -47,6 +47,7 @@ impl<'a> Compiler<'a> { multi_value: false, value_expr: value_expr.cloned(), span: span.clone(), + key_loops_hoisted: false, key_value_loops_hoisted: false, output_query: Some(query.qidx), }; diff --git a/src/languages/rego/compiler/function_calls.rs b/src/languages/rego/compiler/function_calls.rs index ef93804a7..b210bbb6e 100644 --- a/src/languages/rego/compiler/function_calls.rs +++ b/src/languages/rego/compiler/function_calls.rs @@ -232,7 +232,7 @@ impl<'a> Compiler<'a> { /// keep their prior meaning (e.g. a builtin call). OPA instead rewrites /// aliases unconditionally and rejects calls to a missing target at /// compile time. - fn resolve_fcn_path_through_imports(&self, path: &str) -> Option { + pub(super) fn resolve_fcn_path_through_imports(&self, path: &str) -> Option { if self.policy.inner.imports.is_empty() { return None; } diff --git a/src/languages/rego/compiler/loops.rs b/src/languages/rego/compiler/loops.rs index 1d6f4f2b5..cefc29908 100644 --- a/src/languages/rego/compiler/loops.rs +++ b/src/languages/rego/compiler/loops.rs @@ -168,6 +168,7 @@ impl<'a> Compiler<'a> { multi_value: false, value_expr: None, span: span.clone(), + key_loops_hoisted: false, key_value_loops_hoisted: false, output_query: Some(query.qidx), }; diff --git a/src/languages/rego/compiler/mod.rs b/src/languages/rego/compiler/mod.rs index 9550057c2..8ca47ff88 100644 --- a/src/languages/rego/compiler/mod.rs +++ b/src/languages/rego/compiler/mod.rs @@ -67,6 +67,7 @@ pub struct CompilationContext { pub(super) multi_value: bool, pub(super) value_expr: Option, pub(super) span: Span, + pub(super) key_loops_hoisted: bool, pub(super) key_value_loops_hoisted: bool, pub(super) output_query: Option, } diff --git a/src/languages/rego/compiler/queries.rs b/src/languages/rego/compiler/queries.rs index 661e90331..43bbaa636 100644 --- a/src/languages/rego/compiler/queries.rs +++ b/src/languages/rego/compiler/queries.rs @@ -7,8 +7,8 @@ )] use super::{Compiler, CompilerError, ComprehensionType, ContextType, Result}; -use crate::ast::{self, LiteralStmt, Query}; -use crate::compiler::hoist::prioritize_key_binding; +use crate::ast::{self, Expr, ExprRef, LiteralStmt, Query}; +use crate::compiler::hoist::{is_safe_literal_key_call, prioritize_key_binding}; use crate::rvm::instructions::GuardMode; use crate::rvm::program::RuleType; use crate::rvm::Instruction; @@ -17,6 +17,28 @@ use alloc::string::ToString as _; use alloc::vec::Vec; impl<'a> Compiler<'a> { + fn is_safe_projected_key_call(&self, expr: &ExprRef) -> Result { + let Expr::Call { fcn, params, .. } = expr.as_ref() else { + return Ok(false); + }; + if !is_safe_literal_key_call(fcn, params).map_err(CompilerError::from)? { + return Ok(false); + } + let original_path = crate::utils::get_path_string(fcn, None) + .map_err(|_| CompilerError::InvalidFunctionExpression)?; + let full_path = if self.policy.inner.rules.contains_key(&original_path) { + original_path.clone() + } else if let Some(resolved) = self.resolve_fcn_path_through_imports(&original_path) { + resolved + } else { + crate::utils::get_path_string(fcn, Some(&self.current_package)) + .map_err(|_| CompilerError::InvalidFunctionExpressionWithPackage)? + }; + Ok(self.is_builtin(&original_path) + && !self.is_user_defined_function(&full_path) + && !self.host_await_builtins.contains_key(&original_path)) + } + pub(super) fn compile_query(&mut self, query: &Query) -> Result<()> { self.push_scope(); @@ -47,6 +69,10 @@ impl<'a> Compiler<'a> { &mut ordered_stmts, context.key_expr.as_ref(), context.key_filter.as_ref(), + |expr| { + self.is_safe_projected_key_call(expr) + .map_err(|err| anyhow::anyhow!("{err}")) + }, ) .map_err(|err| { CompilerError::General { @@ -67,6 +93,9 @@ impl<'a> Compiler<'a> { &mut self, stmts: &[&LiteralStmt], ) -> Result<()> { + if self.hoist_projected_key_loops_before_body(stmts)? { + return Ok(()); + } if let Some(context) = self.context_stack.last() { let span = context.span.clone(); self.emit_key_filter_guard(&span)?; @@ -128,6 +157,44 @@ impl<'a> Compiler<'a> { self.hoist_loops_and_emit_context_yield() } + fn hoist_projected_key_loops_before_body(&mut self, stmts: &[&LiteralStmt]) -> Result { + let key_expr = match self.context_stack.last() { + Some(context) + if matches!(context.context_type, ContextType::Rule(_)) + && context.key_filter.is_some() + && !context.key_loops_hoisted => + { + context.key_expr.clone() + } + _ => None, + }; + let Some(key_expr) = key_expr else { + return Ok(false); + }; + let loops = self.get_expr_loops(&key_expr)?; + if loops.is_empty() { + return Ok(false); + } + let Some(first_loop) = loops.first() else { + return Ok(false); + }; + if !crate::compiler::hoist::key_expr_is_ready( + &first_loop.collection, + |node| Ok(self.loop_expr_register_map.contains_key(node)), + |span| self.lookup_variable(span.text()).is_some(), + |call| self.is_safe_projected_key_call(call).map_err(Into::into), + ) + .map_err(CompilerError::from)? + { + return Ok(false); + } + if let Some(context) = self.context_stack.last_mut() { + context.key_loops_hoisted = true; + } + self.compile_hoisted_loops(stmts, &loops)?; + Ok(true) + } + pub(super) fn emit_key_filter_guard(&mut self, span: &crate::lexer::Span) -> Result<()> { let Some((expr, filter)) = self.context_stack.last().and_then(|context| { if context.key_filter_guarded { @@ -148,6 +215,7 @@ impl<'a> Compiler<'a> { &expr, |node| Ok(self.loop_expr_register_map.contains_key(node)), |key_span| self.lookup_variable(key_span.text()).is_some(), + |call| self.is_safe_projected_key_call(call).map_err(Into::into), ) .map_err(CompilerError::from)? { @@ -221,7 +289,11 @@ impl<'a> Compiler<'a> { let mut key_value_loops = Vec::new(); if let Some(expr) = key_expr.as_ref() { - if !Self::is_simple_literal(expr) { + let key_loops_hoisted = self + .context_stack + .last() + .is_some_and(|context| context.key_loops_hoisted); + if !key_loops_hoisted && !Self::is_simple_literal(expr) { key_value_loops.extend(self.get_expr_loops(expr)?); } } @@ -230,6 +302,11 @@ impl<'a> Compiler<'a> { key_value_loops.extend(self.get_expr_loops(expr)?); } + if let Some(context) = self.context_stack.last() { + let span = context.span.clone(); + self.emit_key_filter_guard(&span)?; + } + if !key_value_loops.is_empty() { if let Some(context) = self.context_stack.last_mut() { context.key_value_loops_hoisted = true; diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs index 6260a06ee..a193c3159 100644 --- a/src/languages/rego/compiler/rules.rs +++ b/src/languages/rego/compiler/rules.rs @@ -708,6 +708,7 @@ impl<'a> Compiler<'a> { multi_value: matches!(head, RuleHead::Set { .. }), value_expr, span, + key_loops_hoisted: false, key_value_loops_hoisted: false, output_query: None, }; @@ -777,6 +778,7 @@ impl<'a> Compiler<'a> { if let Some(context) = self.context_stack.last_mut() { context.value_expr = body_value_expr.clone(); + context.key_loops_hoisted = false; context.key_value_loops_hoisted = false; context.key_filter_guarded = false; context.output_query = Some(body.query.qidx); diff --git a/src/rvm/instructions/mod.rs b/src/rvm/instructions/mod.rs index fb33a30af..2d46632b8 100644 --- a/src/rvm/instructions/mod.rs +++ b/src/rvm/instructions/mod.rs @@ -17,6 +17,7 @@ pub use types::{ use serde::{Deserialize, Serialize}; /// RVM Instructions - simplified enum-based design +#[non_exhaustive] #[derive(Debug, Clone, Copy, Serialize, Deserialize)] pub enum Instruction { /// Load literal value from literal table into register diff --git a/src/rvm/vm/errors.rs b/src/rvm/vm/errors.rs index 9a5355645..7b45ae9fa 100644 --- a/src/rvm/vm/errors.rs +++ b/src/rvm/vm/errors.rs @@ -9,6 +9,7 @@ use core::time::Duration; use thiserror::Error; /// VM execution errors +#[non_exhaustive] #[derive(Error, Debug, Clone, PartialEq)] pub enum VmError { #[error("Execution stopped: exceeded maximum instruction limit of {limit} after {executed} instructions (pc={pc})")] diff --git a/tests/engine/mod.rs b/tests/engine/mod.rs index 878628a6d..b5e5633b7 100644 --- a/tests/engine/mod.rs +++ b/tests/engine/mod.rs @@ -73,6 +73,119 @@ fn keyed_set_projection_preserves_cycle_detection() -> Result<()> { Ok(()) } +#[test] +fn keyed_set_parent_loop_projection_state_is_invocation_local() -> Result<()> { + let mut engine = Engine::new(); + engine.add_policy( + "test.rego".into(), + r#"package test + values.a contains 1 + values[input.keys[_]] contains count(data.test.values.a)"# + .into(), + )?; + + for keys in [r#"["b","a"]"#, r#"["a","b"]"#] { + engine.set_input(Value::from_json_str(&format!(r#"{{"keys":{keys}}}"#))?); + for path in ["data.test.values", "data.test.values.a"] { + let error = engine + .eval_rule(path.into()) + .expect_err("matching projected recursion must be rejected"); + assert!( + error.to_string().contains("recursion detected"), + "{path}: {error}" + ); + } + + engine.set_input(Value::from_json_str(r#"{"keys":["b"]}"#)?); + assert_eq!( + engine.eval_rule("data.test.values.a".into())?, + Value::from(BTreeSet::from([Value::from(1)])) + ); + } + Ok(()) +} + +#[test] +fn keyed_set_guarded_head_arithmetic_runs_after_body_success() -> Result<()> { + let mut engine = Engine::new(); + engine.add_policy( + "test.rego".into(), + r#"package test + values.a contains 1 + values[k] contains 2 if { is_number(input.port); k := input.port + 1 }"# + .into(), + )?; + + engine.set_input(Value::from_json_str(r#"{"port":"invalid"}"#)?); + assert_eq!( + engine.eval_rule("data.test.values.a".into())?, + Value::from(BTreeSet::from([Value::from(1)])) + ); + + engine.set_input(Value::from_json_str(r#"{"port":80}"#)?); + let result = engine.eval_rule("data.test.values".into())?; + assert_eq!( + result, + Value::from(std::collections::BTreeMap::from([ + ( + Value::from("a"), + Value::from(BTreeSet::from([Value::from(1)])), + ), + ( + Value::from(81), + Value::from(BTreeSet::from([Value::from(2)])), + ), + ])) + ); + + let mut guarded_engine = Engine::new(); + guarded_engine.add_policy( + "guard.rego".into(), + r#"package guard + values[k] contains 2 if { input.enabled; k := input.port + 1 }"# + .into(), + )?; + guarded_engine.set_input(Value::from_json_str( + r#"{"port":"invalid","enabled":true}"#, + )?); + let error = guarded_engine + .eval_rule("data.guard.values".into()) + .expect_err("a successful body must still report a failing head expression"); + assert!(!error.to_string().contains("recursion detected"), "{error}"); + guarded_engine.set_input(Value::from_json_str(r#"{"port":80,"enabled":true}"#)?); + assert_eq!( + guarded_engine.eval_rule("data.guard.values".into())?, + Value::from(std::collections::BTreeMap::from([( + Value::from(81), + Value::from(BTreeSet::from([Value::from(2)])), + )])) + ); + Ok(()) +} + +#[test] +fn keyed_set_does_not_speculate_overridden_builtin_key_calls() -> Result<()> { + let mut engine = Engine::new(); + engine.add_extension( + "lower".to_string(), + 1, + Box::new(|_| Err(anyhow::anyhow!("overridden lower must not run"))), + )?; + engine.add_policy( + "test.rego".into(), + r#"package test + values.a contains 1 + values[k] contains 2 if { false; k := lower("B") }"# + .into(), + )?; + + assert_eq!( + engine.eval_rule("data.test.values.a".into())?, + Value::from(BTreeSet::from([Value::from(1)])) + ); + Ok(()) +} + #[test] fn keyed_set_multiple_bodies_remain_rejected() { let mut engine = Engine::new(); diff --git a/tests/rvm/rego/cases/keyed_partial_sets.yaml b/tests/rvm/rego/cases/keyed_partial_sets.yaml index 045aa68f1..b1349f9f9 100644 --- a/tests/rvm/rego/cases/keyed_partial_sets.yaml +++ b/tests/rvm/rego/cases/keyed_partial_sets.yaml @@ -2,6 +2,243 @@ # Licensed under the MIT License. cases: + - note: keyed_set_compound_key_binding_before_dependency + input: {keys: {chosen: b}} + modules: + - | + package test + values.a contains 1 + values[input.keys[z]] contains n if { n := count(values.a); z := "chosen" } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_compound_key_binding_before_dependency_suspendable + execution_mode: suspendable + input: {keys: {chosen: b}} + modules: + - | + package test + values.a contains 1 + values[input.keys[z]] contains n if { n := count(values.a); z := "chosen" } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_multiple_compound_key_bindings_before_dependency + input: {map: {foo: {bar: b}}} + modules: + - | + package test + values.a contains 1 + values[input.map[x][y]] contains n if { + n := count(values.a) + x := "foo" + y := "bar" + } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_call_key_body_dependency + modules: + - | + package test + values.a contains 1 + values[k] contains n if { n := count(values.a); k := lower("B") } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_call_key_body_dependency_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { n := count(values.a); k := lower("B") } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_call_key_body_matching_cycle + modules: + - | + package test + values.a contains 1 + values[k] contains n if { n := count(values.a); k := lower("A") } + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: keyed_set_call_key_body_matching_cycle_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains n if { n := count(values.a); k := lower("A") } + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: keyed_set_call_key_value_loop + modules: + - | + package test + values.a contains 1 + values[k] contains [count(data.test.values.a)][_] if { k := concat("", ["b"]) } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_call_key_value_loop_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains [count(data.test.values.a)][_] if { k := concat("", ["b"]) } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_call_key_value_loop_matching_cycle + modules: + - | + package test + values.a contains 1 + values[k] contains [count(data.test.values.a)][_] if { k := concat("", ["a"]) } + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: keyed_set_call_key_value_loop_matching_cycle_suspendable + execution_mode: suspendable + modules: + - | + package test + values.a contains 1 + values[k] contains [count(data.test.values.a)][_] if { k := concat("", ["a"]) } + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: keyed_set_head_loop_before_dependent_body + input: {keys: [b]} + modules: + - | + package test + values.a contains 1 + values[input.keys[_]] contains count(data.test.values.a) + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_head_loop_before_dependent_body_suspendable + execution_mode: suspendable + input: {keys: [b]} + modules: + - | + package test + values.a contains 1 + values[input.keys[_]] contains count(data.test.values.a) + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_head_loop_matching_cycle + input: {keys: [a]} + modules: + - | + package test + values.a contains 1 + values[input.keys[_]] contains count(data.test.values.a) + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: keyed_set_head_loop_matching_cycle_suspendable + execution_mode: suspendable + input: {keys: [a]} + modules: + - | + package test + values.a contains 1 + values[input.keys[_]] contains count(data.test.values.a) + query: data.test.values.a + want_error: "Runtime recursion detected" + + - note: keyed_set_guarded_arithmetic_invalid + input: {port: "invalid"} + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { is_number(input.port); k := input.port + 1 } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_guarded_arithmetic_invalid_suspendable + execution_mode: suspendable + input: {port: "invalid"} + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { is_number(input.port); k := input.port + 1 } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_guarded_arithmetic_missing_input + input: {} + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { is_number(input.port); k := input.port + 1 } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_guarded_arithmetic_missing_input_suspendable + execution_mode: suspendable + input: {} + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { is_number(input.port); k := input.port + 1 } + query: data.test.values.a + want_result: + set!: [1] + + - note: keyed_set_guarded_arithmetic_valid + input: {port: 80} + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { is_number(input.port); k := input.port + 1 } + query: data.test.values + want_result: + a: + set!: [1] + 81: + set!: [2] + + - note: keyed_set_guarded_arithmetic_valid_suspendable + execution_mode: suspendable + input: {port: 80} + modules: + - | + package test + values.a contains 1 + values[k] contains 2 if { is_number(input.port); k := input.port + 1 } + query: data.test.values + want_result: + a: + set!: [1] + 81: + set!: [2] + - note: overlapping_keyed_set_head_value_loop_dependency input: {keys: [b, c]} modules: @@ -834,6 +1071,20 @@ cases: b: set!: [20, 21] + - note: keyed_partial_set_nested_head_key_loops + input: + groups: + - keys: ["a", "b"] + vals: [10, 20] + modules: + - | + package test + values.b contains 99 + values[input.groups[i].keys[j]] contains input.groups[i].vals[j] if { true } + query: data.test.values.b + want_result: + set!: [20, 99] + - note: keyed_partial_set_body_bound_head_index input: {keys: ["a", "b"], vals: [10, 20]} modules: From 1c705bdf2f62a4b98faefd0a6c532cc5ea19e11d Mon Sep 17 00:00:00 2001 From: snehabandla <62899142+snehabandla@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:01:52 +0000 Subject: [PATCH 13/13] fix(interpreter): size isolated loop cache for all modules Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/interpreter.rs | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/src/interpreter.rs b/src/interpreter.rs index 1a8a872ce..269001423 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -556,15 +556,6 @@ impl Interpreter { } } - fn ensure_loop_var_values_capacity_for_module( - &mut self, - module_idx: u32, - num_expressions: u32, - ) { - self.loop_var_values - .ensure_capacity(module_idx, num_expressions); - } - fn current_module(&self) -> Result> { self.module .clone() @@ -4648,10 +4639,7 @@ impl Interpreter { let scopes = core::mem::take(&mut self.scopes); let prev_module = self.set_current_module(Some(module.clone()))?; let loop_var_values = core::mem::take(&mut self.loop_var_values); - self.ensure_loop_var_values_capacity_for_module( - self.current_module_index, - module.num_expressions, - ); + self.ensure_loop_var_values_capacity(); let res = self.eval_rule_impl(module, rule, key_filter);