From 64bcdf713644b168d28d6232170af938d9727a38 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Fri, 25 Sep 2026 10:35:51 +0200 Subject: [PATCH 1/3] fix(rego): schedule comprehensions in rule outputs Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/scheduler.rs | 41 +++- tests/engine/mod.rs | 40 ++++ tests/rvm/rego/cases/comprehensions.yaml | 273 +++++++++++++++++++++++ 3 files changed, 351 insertions(+), 3 deletions(-) diff --git a/src/scheduler.rs b/src/scheduler.rs index a54774d60..026a6167f 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -371,11 +371,27 @@ impl Analyzer { // Args are maintained in a separate scope so that they aren't used for // scheduling. self.scopes.push(scope); - for b in bodies { - self.analyze_query(key.clone(), value.clone(), &b.query, Scope::default())?; + let is_partial = match head { + RuleHead::Set { .. } => true, + RuleHead::Compr { refr, .. } => { + matches!(refr.as_ref(), Expr::RefBrack { .. }) + } + RuleHead::Func { .. } => false, + }; + for (idx, b) in bodies.iter().enumerate() { + let use_head_output = idx == 0 || (is_partial && b.assign.is_none()); + let (body_key, body_value) = if use_head_output { + (key.clone(), value.clone()) + } else { + (None, b.assign.as_ref().map(|assign| assign.value.clone())) + }; + self.analyze_query(body_key, body_value, &b.query, Scope::default())?; } if bodies.is_empty() { + if let Some(key) = key { + self.analyze_value_expr(&key)?; + } if let Some(value) = value { self.analyze_value_expr(&value)?; } @@ -416,6 +432,17 @@ impl Analyzer { Ok(()) } + fn analyze_output_expr(&mut self, expr: &Ref, scope: &Scope) -> Result<()> { + let mut output_scope = scope.clone(); + output_scope + .unscoped + .extend(output_scope.locals.keys().cloned()); + self.scopes.push(output_scope); + let result = self.analyze_value_expr(expr); + self.scopes.pop(); + result + } + fn analyze_rule_head( &mut self, head: &RuleHead, @@ -869,7 +896,7 @@ impl Analyzer { mut scope: Scope, ) -> Result<()> { let empty_str = query.span.source_str().clone_empty(); - self.gather_local_vars(key, value, query, &mut scope)?; + self.gather_local_vars(key.clone(), value.clone(), query, &mut scope)?; let mut infos = vec![]; let mut first_use = BTreeMap::new(); @@ -1119,6 +1146,14 @@ impl Analyzer { .set_checked(self.current_module_index, query.qidx, query_schedule) .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))?; + // Output expressions are evaluated after the query body. + if let Some(key) = &key { + self.analyze_output_expr(key, &scope)?; + } + if let Some(value) = &value { + self.analyze_output_expr(value, &scope)?; + } + // Propagate input usage to parent scopes if scope.uses_input && !self.scopes.is_empty() { if let Some(parent_scope) = self.scopes.last_mut() { diff --git a/tests/engine/mod.rs b/tests/engine/mod.rs index 292bcb9b9..4c6927939 100644 --- a/tests/engine/mod.rs +++ b/tests/engine/mod.rs @@ -102,6 +102,46 @@ fn extension_with_state() -> Result<()> { Ok(()) } +#[test] +fn fresh_engine_evaluates_inline_comprehension_outputs_without_entrypoint_compile() -> Result<()> { + let policy = r#" + package test + deny := {"result": true, "reasons": [v | some v in input.values]} if { + count(input.values) > 0 + } + "#; + let expected = Value::from_json_str(r#"{"result":true,"reasons":[false,"x",false]}"#)?; + + let mut nonempty = Engine::new(); + nonempty.set_input(Value::from_json_str(r#"{"values":[false,"x",false]}"#)?); + nonempty.add_policy("test.rego".to_string(), policy.to_string())?; + assert_eq!(nonempty.eval_rule("data.test.deny".to_string())?, expected); + let query = nonempty.eval_query("data.test.deny".to_string(), false)?; + assert_eq!(query.result.len(), 1); + assert_eq!(query.result[0].expressions[0].value, expected); + + let mut empty = Engine::new(); + empty.set_input(Value::from_json_str(r#"{"values":[]}"#)?); + empty.add_policy("test.rego".to_string(), policy.to_string())?; + assert_eq!( + empty.eval_rule("data.test.deny".to_string())?, + Value::Undefined + ); + let query = empty.eval_query("data.test.deny".to_string(), false)?; + assert!(query.result.is_empty()); + + let mut missing = Engine::new(); + missing.add_policy("test.rego".to_string(), policy.to_string())?; + assert_eq!( + missing.eval_rule("data.test.deny".to_string())?, + Value::Undefined + ); + let query = missing.eval_query("data.test.deny".to_string(), false)?; + assert!(query.result.is_empty()); + + Ok(()) +} + #[test] #[cfg(feature = "azure_policy")] #[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))] diff --git a/tests/rvm/rego/cases/comprehensions.yaml b/tests/rvm/rego/cases/comprehensions.yaml index 4152c379b..a2d629042 100644 --- a/tests/rvm/rego/cases/comprehensions.yaml +++ b/tests/rvm/rego/cases/comprehensions.yaml @@ -23,3 +23,276 @@ cases: z := [x | x := 1; x == 2] query: data.test.z want_result: [] + + - note: inline_comprehension_module_scope_array_empty + input: + violations: [] + modules: + - | + package test + violations := input.violations + deny := {"result": true, "reasons": [v | some v in violations]} if { count(violations) > 0 } + query: data.test.deny + want_result: "#undefined" + + - note: inline_comprehension_module_scope_array_nonempty + input: + violations: ["one", "two"] + modules: + - | + package test + violations := input.violations + deny := {"result": true, "reasons": [v | some v in violations]} if { count(violations) > 0 } + query: data.test.deny + want_result: + result: true + reasons: ["one", "two"] + + - note: inline_comprehension_module_scope_set_empty + modules: + - | + package test + violations := set() + deny := {"result": true, "reasons": [v | some v in violations]} if { count(violations) > 0 } + query: data.test.deny + want_result: "#undefined" + + - note: inline_comprehension_module_scope_set_nonempty + modules: + - | + package test + violations := {"one", "two"} + deny := {"result": true, "reasons": [v | some v in violations]} if { count(violations) > 0 } + query: data.test.deny + want_result: + result: true + reasons: ["one", "two"] + + - note: inline_comprehension_module_scope_missing + modules: + - | + package test + violations := input.violations + deny := {"result": true, "reasons": [v | some v in violations]} if { count(violations) > 0 } + query: data.test.deny + want_result: "#undefined" + + - note: inline_comprehension_separate_module_array_nonempty + input: + violations: ["one", "two"] + modules: + - | + package test + violations := input.violations + - | + package test + reasons := [v | some v in violations] + deny := {"result": true, "reasons": reasons} if {count(violations)>0} + query: data.test.deny + want_result: + result: true + reasons: ["one", "two"] + + - note: extracted_comprehension_module_scope_array_empty + input: + violations: [] + modules: + - | + package test + violations := input.violations + reasons := [v | some v in violations] + deny := {"result": true, "reasons": reasons} if { count(violations) > 0 } + query: data.test.deny + want_result: "#undefined" + + - note: extracted_comprehension_module_scope_set_nonempty + modules: + - | + package test + violations := {"one", "two"} + reasons := [v | some v in violations] + deny := {"result": true, "reasons": reasons} if { count(violations) > 0 } + query: data.test.deny + want_result: + result: true + reasons: ["one", "two"] + + - note: extracted_comprehension_module_scope_set_empty + modules: + - | + package test + violations := set() + reasons := [v | some v in violations] + deny := {"result": true, "reasons": reasons} if { count(violations) > 0 } + query: data.test.deny + want_result: "#undefined" + + - note: extracted_comprehension_module_scope_missing + modules: + - | + package test + violations := input.violations + reasons := [v | some v in violations] + deny := {"result": true, "reasons": reasons} if { count(violations) > 0 } + query: data.test.deny + want_result: "#undefined" + + - note: inline_comprehension_function_result + input: + violations: ["one", "two"] + modules: + - | + package test + violations := input.violations + deny(flag) := {"result": flag, "reasons": [v | some v in violations]} if { count(violations) > 0 } + main := deny(true) + query: data.test.main + want_result: + result: true + reasons: ["one", "two"] + + - note: inline_comprehension_function_argument_result + modules: + - | + package test + deny(values) := [v | some v in values] if { count(values) > 0 } + main := deny(["one", "two"]) + query: data.test.main + want_result: ["one", "two"] + + - note: inline_array_comprehension_contains_nested_term_comprehension + input: + violations: ["one", "two"] + modules: + - | + package test + result := [[v | some v in xs] | some xs in [[1], [2]]] + query: data.test.result + want_result: + - [1] + - [2] + + - note: inline_object_comprehension_nested_key_and_value_comprehensions + modules: + - | + package test + result := { + [k | some k in ["one", "two"]][0]: { + [v | some v in ["a", "b"]][0]: [upper(v) | some v in ["a", "b"]] + | some marker in [true] + } + } + query: data.test.result + want_result: + one: + a: ["A", "B"] + + - note: inline_object_comprehension_key_value_output + input: + violations: ["one", "two"] + modules: + - | + package test + violations := input.violations + deny := {v: upper(v) | some v in violations} if { count(violations) > 0 } + query: data.test.deny + want_result: + one: "ONE" + two: "TWO" + + - note: inline_comprehension_else_output + input: + violations: [] + modules: + - | + package test + violations := input.violations + deny := {"reasons": [v | some v in violations]} if { count(violations) > 0 } + else := {"reasons": [v | some v in violations]} if { count(violations) == 0 } + query: data.test.deny + want_result: + reasons: [] + + - note: inline_comprehension_local_scope_array_nonempty + input: + violations: ["one", "two"] + modules: + - | + package test + deny := {"result": true, "reasons": [v | some v in violations]} if { + violations := input.violations + count(violations) > 0 + } + query: data.test.deny + want_result: + result: true + reasons: ["one", "two"] + + - note: inline_comprehension_else_skips_unsafe_head_output + modules: + - | + package test + deny := [v | some v in xs] if { + xs := [1] + false + } else := [] if { + true + } + query: data.test.deny + want_result: [] + + - note: partial_object_independent_body_inherits_head_comprehension_output + rego_v0: true + modules: + - | + package test + import future.keywords.in + result[k] = [v | some v in values] { + k := "first" + values := [1] + } + { + k := "second" + values := [2] + } + query: data.test.result + want_result: + first: [1] + second: [2] + + - note: inline_comprehension_embedded_in_object_key_and_value + input: + values: ["one", "two"] + modules: + - | + package test + result := { + [v | some v in input.values][0]: [upper(v) | some v in input.values] + } + query: data.test.result + want_result: + one: ["ONE", "TWO"] + + - note: bodyless_partial_set_key_with_comprehension + modules: + - | + package test + values := [1, 2] + result contains [v | some v in values] + query: data.test.result + want_result: + set!: + - [1, 2] + + - note: inline_set_comprehension_output + input: + values: [false, "x", false] + modules: + - | + package test + result := {v | some v in input.values} + query: data.test.result + want_result: + set!: + - false + - "x" From 6bfae01c4305ea16d1ad3b39930a80cdbe16f78f Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Fri, 25 Sep 2026 11:58:12 +0200 Subject: [PATCH 2/3] refactor(rego): borrow rule outputs during variable analysis Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/scheduler.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/scheduler.rs b/src/scheduler.rs index 026a6167f..7d190f4db 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -469,8 +469,8 @@ impl Analyzer { fn gather_local_vars( &mut self, - key: Option>, - value: Option>, + key: Option<&Ref>, + value: Option<&Ref>, query: &Query, scope: &mut Scope, ) -> Result<()> { @@ -522,10 +522,10 @@ impl Analyzer { } } - if let Some(key) = &key { + if let Some(key) = key { gather_vars(key, false, &self.scopes, scope)?; } - if let Some(value) = &value { + if let Some(value) = value { gather_vars(value, false, &self.scopes, scope)?; } @@ -896,7 +896,7 @@ impl Analyzer { mut scope: Scope, ) -> Result<()> { let empty_str = query.span.source_str().clone_empty(); - self.gather_local_vars(key.clone(), value.clone(), query, &mut scope)?; + self.gather_local_vars(key.as_ref(), value.as_ref(), query, &mut scope)?; let mut infos = vec![]; let mut first_use = BTreeMap::new(); From 9051686179b40ed40572c3b6f67907e301e8d6d1 Mon Sep 17 00:00:00 2001 From: Maksym Mishchenko Date: Fri, 25 Sep 2026 13:26:29 +0200 Subject: [PATCH 3/3] fix(rego): preserve nested output capture dependencies Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/scheduler.rs | 72 ++++++++++++++++-------- tests/engine/mod.rs | 46 +++++++++++++++ tests/rvm/rego/cases/comprehensions.yaml | 12 ++++ 3 files changed, 107 insertions(+), 23 deletions(-) diff --git a/src/scheduler.rs b/src/scheduler.rs index 7d190f4db..7a0f0d38f 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -389,22 +389,23 @@ impl Analyzer { } if bodies.is_empty() { + let mut scope = Scope::default(); if let Some(key) = key { - self.analyze_value_expr(&key)?; + self.analyze_value_expr(&key, &mut scope)?; } if let Some(value) = value { - self.analyze_value_expr(&value)?; + self.analyze_value_expr(&value, &mut scope)?; } } self.scopes.pop(); Ok(()) } - Rule::Default { value, .. } => self.analyze_value_expr(value), + Rule::Default { value, .. } => self.analyze_value_expr(value, &mut Scope::default()), } } - fn analyze_value_expr(&mut self, expr: &Ref) -> Result<()> { + fn analyze_value_expr(&mut self, expr: &Ref, scope: &mut Scope) -> Result<()> { let mut comprs = vec![]; traverse(expr, &mut |e| match e.as_ref() { ArrayCompr { .. } | SetCompr { .. } | ObjectCompr { .. } => { @@ -414,32 +415,57 @@ impl Analyzer { _ => Ok(true), })?; for compr in comprs { - match compr.as_ref() { + let qidx = match compr.as_ref() { Expr::ArrayCompr { query, term, .. } | Expr::SetCompr { query, term, .. } => { self.analyze_query(None, Some(term.clone()), query, Scope::default())?; + query.qidx } Expr::ObjectCompr { query, key, value, .. - } => self.analyze_query( - Some(key.clone()), - Some(value.clone()), - query, - Scope::default(), - )?, - _ => (), + } => { + self.analyze_query( + Some(key.clone()), + Some(value.clone()), + query, + Scope::default(), + )?; + query.qidx + } + _ => continue, + }; + + if let Some(compr_scope) = self + .schedule_table + .get_checked(self.current_module_index, qidx) + .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))? + .map(|qs| &qs.scope) + { + Self::propagate_nested_scope(compr_scope, scope); } } Ok(()) } - fn analyze_output_expr(&mut self, expr: &Ref, scope: &Scope) -> Result<()> { + fn propagate_nested_scope(compr_scope: &Scope, scope: &mut Scope) { + if compr_scope.uses_input { + scope.uses_input = true; + } + for iv in &compr_scope.inputs { + if !scope.locals.contains_key(iv) && !scope.unscoped.contains(iv) { + scope.inputs.insert(iv.clone()); + } + } + } + + fn analyze_output_expr(&mut self, expr: &Ref, scope: &mut Scope) -> Result<()> { let mut output_scope = scope.clone(); output_scope .unscoped .extend(output_scope.locals.keys().cloned()); - self.scopes.push(output_scope); - let result = self.analyze_value_expr(expr); + self.scopes.push(output_scope.clone()); + let result = self.analyze_value_expr(expr, &mut output_scope); self.scopes.pop(); + Self::propagate_nested_scope(&output_scope, scope); result } @@ -1138,6 +1164,14 @@ impl Analyzer { _ => Vec::new(), }; + // Output expressions are evaluated after the query body. + if let Some(key) = &key { + self.analyze_output_expr(key, &mut scope)?; + } + if let Some(value) = &value { + self.analyze_output_expr(value, &mut scope)?; + } + let query_schedule = QuerySchedule { scope: scope.clone(), order, @@ -1146,14 +1180,6 @@ impl Analyzer { .set_checked(self.current_module_index, query.qidx, query_schedule) .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))?; - // Output expressions are evaluated after the query body. - if let Some(key) = &key { - self.analyze_output_expr(key, &scope)?; - } - if let Some(value) = &value { - self.analyze_output_expr(value, &scope)?; - } - // Propagate input usage to parent scopes if scope.uses_input && !self.scopes.is_empty() { if let Some(parent_scope) = self.scopes.last_mut() { diff --git a/tests/engine/mod.rs b/tests/engine/mod.rs index 4c6927939..18fe68664 100644 --- a/tests/engine/mod.rs +++ b/tests/engine/mod.rs @@ -142,6 +142,52 @@ fn fresh_engine_evaluates_inline_comprehension_outputs_without_entrypoint_compil Ok(()) } +#[test] +fn eval_rule_schedules_unification_nested_output_capture() -> Result<()> { + let mut engine = Engine::new(); + engine.add_policy( + "test.rego".to_string(), + r#" + package test + result := a if { + a := [[v | some v in vals] | true] + vals = [1, 2] + } + "# + .to_string(), + )?; + + assert_eq!( + engine.eval_rule("data.test.result".to_string())?, + Value::from_json_str("[[1,2]]")? + ); + Ok(()) +} + +#[test] +fn eval_rule_rejects_forward_assignment_nested_output_capture() -> Result<()> { + let mut engine = Engine::new(); + engine.add_policy( + "test.rego".to_string(), + r#" + package test + result := a if { + a := [[v | some v in vals] | true] + vals := [1, 2] + } + "# + .to_string(), + )?; + + let error = engine + .eval_rule("data.test.result".to_string()) + .expect_err("forward assignment should remain unsafe"); + assert!(error + .to_string() + .contains("use of undefined variable `vals`")); + Ok(()) +} + #[test] #[cfg(feature = "azure_policy")] #[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))] diff --git a/tests/rvm/rego/cases/comprehensions.yaml b/tests/rvm/rego/cases/comprehensions.yaml index a2d629042..e0e86be6c 100644 --- a/tests/rvm/rego/cases/comprehensions.yaml +++ b/tests/rvm/rego/cases/comprehensions.yaml @@ -296,3 +296,15 @@ cases: set!: - false - "x" + + - note: output_comprehension_captures_unification_body_local + modules: + - | + package test + result := a if { + a := [[v | some v in vals] | true] + vals = [1, 2] + } + query: data.test.result + want_result: + - [1, 2]