Prevent undefined values from entering RVM sets

src/rvm/vm/comprehension.rs

@@ -250,7 +250,9 @@ impl RegoVM {
 
         match (&comprehension_context.mode, &mut current_result) {
             (&ComprehensionMode::Set, &mut Value::Set(ref mut set)) => {
-                crate::Rc::make_mut(set).insert(value_to_add);
+                if !matches!(value_to_add, Value::Undefined) {
+                    crate::Rc::make_mut(set).insert(value_to_add);
+                }
             }
             (&ComprehensionMode::Array, &mut Value::Array(ref mut arr)) => {
                 crate::Rc::make_mut(arr).push(value_to_add);
@@ -362,7 +364,9 @@ impl RegoVM {
 
         match (&mode, &mut current_result) {
             (&ComprehensionMode::Set, &mut Value::Set(ref mut set)) => {
-                crate::Rc::make_mut(set).insert(value_to_add);
+                if !matches!(value_to_add, Value::Undefined) {
+                    crate::Rc::make_mut(set).insert(value_to_add);
+                }
             }
             (&ComprehensionMode::Array, &mut Value::Array(ref mut arr)) => {
                 crate::Rc::make_mut(arr).push(value_to_add);

src/rvm/vm/dispatch.rs

@@ -679,7 +679,9 @@ impl RegoVM {
                 let mut set_value = self.take_register(set)?;
 
                 if let Ok(set_mut) = set_value.as_set_mut() {
-                    set_mut.insert(value_to_add);
+                    if !matches!(value_to_add, Value::Undefined) {
+                        set_mut.insert(value_to_add);
+                    }
                     self.set_register(set, set_value)?;
                 } else {
                     let offending = set_value.clone();

tests/rvm/rego/cases/undefined_set_values.yaml

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: partial_set_omits_undefined_head
+    data: {}
+    input: {}
+    modules:
+      - |
+        package test
+        values contains input.missing if { true }
+    query: data.test.values
+    want_result:
+      set!: []
+
+  - note: partial_set_omits_undefined_head_suspendable
+    execution_mode: suspendable
+    data: {}
+    input: {}
+    modules:
+      - |
+        package test
+        values contains input.missing if { true }
+    query: data.test.values
+    want_result:
+      set!: []
+
+  - note: partial_set_preserves_defined_values
+    data: {}
+    input:
+      items: [null, false, 0, "", false]
+    modules:
+      - |
+        package test
+        values contains item if { some item in input.items }
+    query: data.test.values
+    want_result:
+      set!: [null, false, 0, ""]
+
+  - note: partial_set_preserves_defined_values_suspendable
+    execution_mode: suspendable
+    data: {}
+    input:
+      items: [null, false, 0, "", false]
+    modules:
+      - |
+        package test
+        values contains item if { some item in input.items }
+    query: data.test.values
+    want_result:
+      set!: [null, false, 0, ""]
+
+  - note: set_comprehension_omits_undefined_values
+    data: {}
+    input: {}
+    modules:
+      - |
+        package test
+        values := {input.missing | true}
+    query: data.test.values
+    want_result:
+      set!: []
+
+  - note: set_comprehension_omits_undefined_values_suspendable
+    execution_mode: suspendable
+    data: {}
+    input: {}
+    modules:
+      - |
+        package test
+        values := {input.missing | true}
+    query: data.test.values
+    want_result:
+      set!: []

tests/rvm/vm/suites/set_operations.yaml

@@ -6,6 +6,44 @@
 # Covers complex value deduplication, nested sets, and undefined handling
 
 cases:
+  - note: set_add_undefined_rejects_non_set_target
+    literals:
+      - "#undefined"
+    instructions:
+      - "LoadNull { dest: 0 }"
+      - "Load { dest: 1, literal_idx: 0 }"
+      - "SetAdd { set: 0, value: 1 }"
+      - "Return { value: 0 }"
+    want_error: "Register 0 does not contain a set"
+
+  - note: set_add_undefined_to_empty_set
+    literals:
+      - "#undefined"
+    instructions:
+      - "SetNew { dest: 0 }"
+      - "Load { dest: 1, literal_idx: 0 }"
+      - "SetAdd { set: 0, value: 1 }"
+      - "Return { value: 0 }"
+    want_result:
+      set!: []
+
+  - note: set_add_undefined_preserves_shared_set
+    literals:
+      - "#undefined"
+    instructions:
+      - "SetNew { dest: 0 }"
+      - "LoadFalse { dest: 1 }"
+      - "SetAdd { set: 0, value: 1 }"
+      - "LoadNull { dest: 1 }"
+      - "SetAdd { set: 0, value: 1 }"
+      - "Move { dest: 2, src: 0 }"
+      - "Load { dest: 1, literal_idx: 0 }"
+      - "SetAdd { set: 0, value: 1 }"
+      - "AssertEq { left: 0, right: 2 }"
+      - "Return { value: 0 }"
+    want_result:
+      set!: [false, null]
+
   - note: set_deduplication_simple
     description: Set automatically deduplicates simple values
     example_rego: "{1, 2, 1, 3, 2}"