From c3dea92cec2d09b4d646ae0fa88ef15914c2751c Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 15 Oct 2026 20:27:26 +0000
Subject: [PATCH] fix(rvm): omit undefined values from set accumulation

Co-authored-by: anakrish <35780660+anakrish@users.noreply.github.com>
---
 src/rvm/vm/comprehension.rs                   |  8 +-
 src/rvm/vm/dispatch.rs                        |  4 +-
 .../rvm/rego/cases/undefined_set_values.yaml  | 74 +++++++++++++++++++
 tests/rvm/vm/suites/set_operations.yaml       | 38 ++++++++++
 4 files changed, 121 insertions(+), 3 deletions(-)
 create mode 100644 tests/rvm/rego/cases/undefined_set_values.yaml

diff --git a/src/rvm/vm/comprehension.rs b/src/rvm/vm/comprehension.rs
index b5c4c19b9..84ccf2495 100644
--- a/src/rvm/vm/comprehension.rs
+++ b/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);
diff --git a/src/rvm/vm/dispatch.rs b/src/rvm/vm/dispatch.rs
index d24cb9860..b1913fbb8 100644
--- a/src/rvm/vm/dispatch.rs
+++ b/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();
diff --git a/tests/rvm/rego/cases/undefined_set_values.yaml b/tests/rvm/rego/cases/undefined_set_values.yaml
new file mode 100644
index 000000000..7a8584012
--- /dev/null
+++ b/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!: []
diff --git a/tests/rvm/vm/suites/set_operations.yaml b/tests/rvm/vm/suites/set_operations.yaml
index a0ae8112e..cc8246968 100644
--- a/tests/rvm/vm/suites/set_operations.yaml
+++ b/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}"