Fix partial object iteration dropping bindings and guard unsupported RVM patterns

src/interpreter.rs

@@ -1782,6 +1782,7 @@ 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 output = if let Some(oe) = &output_expr {

src/languages/rego/compiler/error.rs

@@ -63,6 +63,12 @@ pub enum CompilerError {
     #[error("Invalid function expression with package")]
     InvalidFunctionExpressionWithPackage,
 
+    #[error("partial object rules with constant keys are not yet supported by the RVM compiler")]
+    PartialObjectConstantKeyUnsupported,
+
+    #[error("partial object rules with nested bracket keys are not yet supported by the RVM compiler")]
+    PartialObjectNestedKeyUnsupported,
+
     #[error("Compilation error: {message}")]
     General { message: String },
 }

src/languages/rego/compiler/rules.rs

@@ -59,7 +59,7 @@ impl<'a> Compiler<'a> {
                             crate::ast::Expr::RefBrack { .. } if assign.is_some() => {
                                 RuleType::PartialObject
                             }
-                            crate::ast::Expr::RefBrack { .. } => RuleType::PartialSet,
+                            crate::ast::Expr::RefBrack { .. } => RuleType::PartialObject,
                             _ => RuleType::Complete,
                         },
                         _ => RuleType::Complete,
@@ -88,6 +88,34 @@ impl<'a> Compiler<'a> {
         })
     }
 
+    fn validate_partial_object_shape(&self, refr: &ExprRef) -> Result<()> {
+        let Expr::RefBrack {
+            refr: prefix,
+            index,
+            ..
+        } = refr.as_ref()
+        else {
+            return Ok(());
+        };
+
+        if matches!(prefix.as_ref(), Expr::RefBrack { .. }) {
+            return Err(CompilerError::PartialObjectNestedKeyUnsupported.at(refr.span()));
+        }
+
+        if matches!(
+            index.as_ref(),
+            Expr::String { .. }
+                | Expr::RawString { .. }
+                | Expr::Number { .. }
+                | Expr::Bool { .. }
+                | Expr::Null { .. }
+        ) {
+            return Err(CompilerError::PartialObjectConstantKeyUnsupported.at(index.span()));
+        }
+
+        Ok(())
+    }
+
     pub(super) fn get_or_assign_rule_index(&mut self, rule_path: &str) -> Result<u16> {
         if let Some(&index) = self.rule_index_map.get(rule_path) {
             return Ok(index);
@@ -345,6 +373,10 @@ impl<'a> Compiler<'a> {
 
                     let (key_expr, value_expr) = match head {
                         RuleHead::Compr { refr, assign, .. } => {
+                            if rule_type == RuleType::PartialObject {
+                                self.validate_partial_object_shape(refr)?;
+                            }
+
                             self.rule_definition_function_params[rule_index as usize].push(None);
                             self.rule_definition_destructuring_patterns[rule_index as usize]
                                 .push(None);

tests/interpreter/cases/rule/partial_object_v1.yaml

@@ -0,0 +1,47 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: constant_key_partial_object_v1
+    data: {}
+    input:
+      enabled: true
+    modules:
+      - |
+        package test
+        import rego.v1
+
+        p["fixed"] if {
+          input.enabled
+        }
+    query: data.test
+    want_result:
+      p:
+        fixed: true
+
+  - note: multilevel_partial_object_v1
+    data: {}
+    input:
+      nested:
+        app:
+          read: 1
+          write: 2
+        ops:
+          deploy: 3
+    modules:
+      - |
+        package test
+        import rego.v1
+
+        p[a][b] if {
+          some a, obj in input.nested
+          some b, _ in obj
+        }
+    query: data.test
+    want_result:
+      p:
+        app:
+          read: true
+          write: true
+        ops:
+          deploy: true

tests/rvm/rego/cases/partial_object_rules.yaml

@@ -0,0 +1,218 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: partial_object_variable_key_collects_all_bindings
+    data: {}
+    input:
+      items:
+        FOO: 1
+        BAR: 2
+        BAZ: 3
+    modules:
+      - |
+        package test
+
+        p[k] if {
+          some k, _ in input.items
+        }
+    query: data.test.p
+    want_result:
+      BAR: true
+      BAZ: true
+      FOO: true
+
+  - note: partial_object_explicit_value_collects_all_bindings
+    data: {}
+    input:
+      items:
+        FOO: 1
+        BAR: 2
+        BAZ: 3
+    modules:
+      - |
+        package test
+
+        p[k] := v if {
+          some k, v in input.items
+        }
+    query: data.test.p
+    want_result:
+      BAR: 2
+      BAZ: 3
+      FOO: 1
+
+  - note: partial_object_dynamic_expression_key_collects_all_bindings
+    data: {}
+    input:
+      items:
+        FOO: 1
+        BAR: 2
+        BAZ: 3
+      aliases:
+        FOO: alias-foo
+        BAR: alias-bar
+        BAZ: alias-baz
+    modules:
+      - |
+        package test
+
+        p[input.aliases[k]] := v if {
+          some k, v in input.items
+        }
+    query: data.test.p
+    want_result:
+      alias-bar: 2
+      alias-baz: 3
+      alias-foo: 1
+
+  - note: partial_object_constant_key_unsupported_in_rvm
+    data: {}
+    input:
+      enabled: true
+    modules:
+      - |
+        package test
+
+        p["fixed"] if {
+          input.enabled
+        }
+    query: data.test.p.fixed
+    want_error: "partial object rules with constant keys are not yet supported by the RVM compiler"
+    allow_interpreter_success: true
+
+  - note: partial_object_constant_key_explicit_value_unsupported_in_rvm
+    data: {}
+    input:
+      enabled: true
+    modules:
+      - |
+        package test
+
+        p["fixed"] := 7 if {
+          input.enabled
+        }
+    query: data.test.p.fixed
+    want_error: "partial object rules with constant keys are not yet supported by the RVM compiler"
+    allow_interpreter_success: true
+
+  - note: partial_object_multiple_bodies_collects_all_bindings
+    data: {}
+    input:
+      items:
+        FOO: 1
+        BAR: 2
+        BAZ: 3
+    modules:
+      - |
+        package test
+
+        p[k] if {
+          some k, _ in input.items
+          k in {"FOO", "BAR"}
+        }
+
+        p[k] if {
+          some k, _ in input.items
+          k == "BAZ"
+        }
+    query: data.test.p
+    want_result:
+      BAR: true
+      BAZ: true
+      FOO: true
+
+  - note: partial_set_contains_collects_all_bindings
+    data: {}
+    input:
+      items:
+        FOO: 1
+        BAR: 2
+        BAZ: 3
+    modules:
+      - |
+        package test
+
+        p contains k if {
+          some k, _ in input.items
+        }
+    query: data.test.p
+    want_result:
+      set!: ["BAR", "BAZ", "FOO"]
+
+  - note: partial_object_multilevel_key_unsupported_in_rvm
+    data: {}
+    input:
+      nested:
+        app:
+          read: 1
+          write: 2
+        ops:
+          deploy: 3
+    modules:
+      - |
+        package test
+
+        p[a][b] if {
+          some a, obj in input.nested
+          some b, _ in obj
+        }
+
+        main := p
+    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: partial_object_multilevel_key_explicit_value_unsupported_in_rvm
+    data: {}
+    input:
+      nested:
+        app:
+          read: 1
+          write: 2
+        ops:
+          deploy: 3
+    modules:
+      - |
+        package test
+
+        p[a][b] := v if {
+          some a, obj in input.nested
+          some b, v in obj
+        }
+
+        main := p
+    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: partial_object_array_iteration_collects_all_bindings
+    data: {}
+    input:
+      items: ["FOO", "BAR", "BAZ"]
+    modules:
+      - |
+        package test
+
+        p[v] if {
+          some _, v in input.items
+        }
+    query: data.test.p
+    want_result:
+      BAR: true
+      BAZ: true
+      FOO: true
+
+  - note: partial_object_empty_input_is_empty_object
+    data: {}
+    input:
+      items: {}
+    modules:
+      - |
+        package test
+
+        p[k] if {
+          some k, _ in input.items
+        }
+    query: data.test.p
+    want_result: {}