From 9fa8036ce4be550595ea5ecc0590668e03546cf8 Mon Sep 17 00:00:00 2001
From: Anand Krishnamoorthi <anakrish@microsoft.com>
Date: Fri, 5 Dec 2025 12:26:26 -0600
Subject: [PATCH] feat: Implement Rego else block compilation

- teach the Rego compiler to compile else chains correctly
- test suite

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
---
 src/languages/rego/compiler/rules.rs |  40 ++++++---
 tests/opa.rs                         |   3 +-
 tests/rvm/rego/cases/else_rules.yaml | 126 +++++++++++++++++++++++++++
 3 files changed, 154 insertions(+), 15 deletions(-)
 create mode 100644 tests/rvm/rego/cases/else_rules.yaml

diff --git a/src/languages/rego/compiler/rules.rs b/src/languages/rego/compiler/rules.rs
index 01979cb..78b296a 100644
--- a/src/languages/rego/compiler/rules.rs
+++ b/src/languages/rego/compiler/rules.rs
@@ -471,6 +471,20 @@ impl<'a> Compiler<'a> {
 
                             ::core::convert::identity(body_idx);
 
+                            let previous_value_expr = self
+                                .context_stack
+                                .last()
+                                .and_then(|ctx| ctx.value_expr.clone());
+                            let mut body_value_expr =
+                                body.assign.as_ref().map(|assign| assign.value.clone());
+                            if body_value_expr.is_none() && body_idx == 0 {
+                                body_value_expr = previous_value_expr.clone();
+                            }
+
+                            if let Some(context) = self.context_stack.last_mut() {
+                                context.value_expr = body_value_expr.clone();
+                            }
+
                             self.emit_instruction(
                                 Instruction::RuleInit {
                                     result_reg: result_register,
@@ -481,23 +495,23 @@ impl<'a> Compiler<'a> {
 
                             if !body.query.stmts.is_empty() {
                                 self.compile_query(&body.query)?;
-                            } else {
-                                let value_expr_opt =
-                                    self.context_stack.last().unwrap().value_expr.clone();
-                                if let Some(value_expr) = value_expr_opt {
-                                    let value_reg = self.compile_rego_expr(&value_expr)?;
-                                    self.emit_instruction(
-                                        Instruction::Move {
-                                            dest: result_register,
-                                            src: value_reg,
-                                        },
-                                        value_expr.span(),
-                                    );
-                                }
+                            } 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(),
+                                );
                             }
 
                             self.emit_instruction(Instruction::RuleReturn {}, &body.span);
 
+                            if let Some(context) = self.context_stack.last_mut() {
+                                context.value_expr = previous_value_expr;
+                            }
+
                             self.pop_scope();
                         }
                     }
diff --git a/tests/opa.rs b/tests/opa.rs
index bba7001..9909381 100644
--- a/tests/opa.rs
+++ b/tests/opa.rs
@@ -26,7 +26,6 @@ const OPA_TODO_FOLDERS: &[&str] = &[
     "baseandvirtualdocs",
     "dataderef",
     "defaultkeyword",
-    "elsekeyword",
     "every",
     "fix1863",
     "functions",
@@ -34,7 +33,7 @@ const OPA_TODO_FOLDERS: &[&str] = &[
     "partialobjectdoc",
     "planner-ir",
     "refheads",
-    "virtualdocs",
+    "type",
     "walkbuiltin",
     // RVM Compiler does not support 'with' keyword yet.
     "withkeyword",
diff --git a/tests/rvm/rego/cases/else_rules.yaml b/tests/rvm/rego/cases/else_rules.yaml
new file mode 100644
index 0000000..62cbe61
--- /dev/null
+++ b/tests/rvm/rego/cases/else_rules.yaml
@@ -0,0 +1,126 @@
+# Rego Else Rules Test Suite
+# Exercises compiler support for else bodies, including assignment overrides and boolean fallback logic.
+
+cases:
+  - note: else_rule_short_circuit
+    modules:
+      - |
+        package test
+        decision := 1 if {
+          1 == 1
+        }
+        else := 2 if {
+          1 == 1
+        }
+    query: data.test.decision
+    want_result: 1
+
+  - note: else_rule_fallback
+    modules:
+      - |
+        package test
+        decision := 1 if {
+          1 == 2
+        }
+        else := 2 if {
+          1 == 2
+        }
+        else := 3 if {
+          1 == 1
+        }
+    query: data.test.decision
+    want_result: 3
+
+  - note: else_rule_assignment_only
+    modules:
+      - |
+        package test
+        decision := 1 if {
+          1 == 2
+        }
+        else := 99
+    query: data.test.decision
+    want_result: 99
+
+  - note: else_rule_no_assignment_boolean
+    input:
+      method: "POST"
+    modules:
+      - |
+        package test
+        allow if {
+          input.method == "GET"
+        }
+        else if {
+          input.method == "POST"
+        }
+    query: data.test.allow
+    want_result: true
+
+  - note: else_rule_multiple_definitions
+    modules:
+      - |
+        package test
+        decision := "first" if {
+          false
+        }
+        else := "first-else" if {
+          false
+        }
+        decision := "second" if {
+          false
+        }
+        else := "second-else" if {
+          true
+        }
+    query: data.test.decision
+    want_result: "second-else"
+
+  - note: else_rule_function_fallback
+    modules:
+      - |
+        package test
+        f(x) := "small" if {
+          x < 5
+        }
+        else := "medium" if {
+          x < 10
+        }
+        f(x) := "large" if {
+          x >= 10
+        }
+        result := f(8)
+    query: data.test.result
+    want_result: "medium"
+
+  - note: else_rule_multiple_defined_single
+    modules:
+      - |
+        package ex
+
+        multiple_defined := false if {
+          false
+        }
+        else if {
+          true
+        }
+        else := false
+    query: data.ex.multiple_defined
+    want_result: true
+
+  - note: else_rule_boolean_middle_then_assignment
+    modules:
+      - |
+        package corner
+
+        corner_case := 7 if {
+          false
+        }
+        else := 6 if {
+          false
+        } else if {
+          true
+        }
+        else := 99
+    query: data.corner.corner_case
+    want_result: true