feat(rvm): new instructions and loop semantics for Azure Policy support (microsoft#659)

The Rego VM was designed around Rego's semantics, but Azure Policy needs
a few things Rego doesn't: host-supplied context alongside input/data,
undefined-to-null coercion for missing fields, skip-undefined collection
behavior for wildcard aliases, and non-vacuous iteration over non-array
values.

This commit adds five new instructions to bridge those gaps:

  LoadContext / LoadMetadata — give programs access to host-supplied
  evaluation context and cached program metadata at runtime.

  ArrayPushDefined — like ArrayPush but silently drops undefined values,
  so wildcard alias collection (field[*].property) excludes absent
  nested properties instead of leaking undefined entries into the array.

  ReturnUndefinedIfNotTrue — early return with Undefined when a guard
  condition isn't satisfied, without tripping a VM assertion failure.
  This models "condition doesn't match" cleanly.

  CoalesceUndefinedToNull — turns Undefined into Null in-place so that
  downstream builtins see null rather than short-circuiting on undefined.

The loop engine also gains an Azure Policy mode: when the source language
is "azure_policy", an Every loop over a non-array value (scalars, null,
objects) iterates once over a virtual Null element instead of being
vacuously true.  This matches how field[*] behaves on non-array fields
in Azure Policy — the condition body runs once against Null, which
typically evaluates to false.

On the plumbing side: the VM gets a context field with set_context(),
metadata is cached as a Value on program load, and map_limit_error is
inlined into memory_check since it had only one call site.

Four new YAML test suites (~880 lines) cover the new instructions and
context/metadata loading, along with instruction parser, display, and
assembly listing support for everything added here.

Author: anakrish

src/rvm/instructions/display.rs

@@ -143,6 +143,8 @@ impl core::fmt::Display for Instruction {
             Instruction::LoadBool { dest, value } => format!("LOAD_BOOL R({}) {}", dest, value),
             Instruction::LoadData { dest } => format!("LOAD_DATA R({})", dest),
             Instruction::LoadInput { dest } => format!("LOAD_INPUT R({})", dest),
+            Instruction::LoadContext { dest } => format!("LOAD_CONTEXT R({})", dest),
+            Instruction::LoadMetadata { dest } => format!("LOAD_METADATA R({})", dest),
             Instruction::Move { dest, src } => format!("MOVE R({}) R({})", dest, src),
             Instruction::Add { dest, left, right } => {
                 format!("ADD R({}) R({}) R({})", dest, left, right)
@@ -220,6 +222,9 @@ impl core::fmt::Display for Instruction {
             }
             Instruction::ArrayNew { dest } => format!("ARRAY_NEW R({})", dest),
             Instruction::ArrayPush { arr, value } => format!("ARRAY_PUSH R({}) R({})", arr, value),
+            Instruction::ArrayPushDefined { arr, value } => {
+                format!("ARRAY_PUSH_DEFINED R({}) R({})", arr, value)
+            }
             Instruction::ArrayCreate { params_index } => {
                 format!("ARRAY_CREATE P({})", params_index)
             }
@@ -247,6 +252,12 @@ impl core::fmt::Display for Instruction {
                 };
                 format!("{} R({})", name, register)
             }
+            Instruction::ReturnUndefinedIfNotTrue { condition } => {
+                format!("RETURN_UNDEFINED_IF_NOT_TRUE R({})", condition)
+            }
+            Instruction::CoalesceUndefinedToNull { register } => {
+                format!("COALESCE_UNDEF_TO_NULL R({})", register)
+            }
             Instruction::LoopStart { params_index } => {
                 format!("LOOP_START P({})", params_index)
             }

src/rvm/instructions/mod.rs

@@ -54,6 +54,16 @@ pub enum Instruction {
         dest: u8,
     },
 
+    /// Load host-supplied context value into register
+    LoadContext {
+        dest: u8,
+    },
+
+    /// Load program metadata value into register
+    LoadMetadata {
+        dest: u8,
+    },
+
     /// Move value from one register to another
     Move {
         dest: u8,
@@ -206,6 +216,16 @@ pub enum Instruction {
         value: u8,
     },
 
+    /// Push element to array, but skip if the value is undefined.
+    ///
+    /// Used by Azure Policy's `field('alias[*].property')` wildcard collection
+    /// so that absent nested properties are excluded from the collected array
+    /// rather than producing undefined entries.
+    ArrayPushDefined {
+        arr: u8,
+        value: u8,
+    },
+
     /// Create array from registers - returns undefined if any element is undefined
     ArrayCreate {
         /// Index into program's instruction_data.array_create_params table
@@ -254,6 +274,25 @@ pub enum Instruction {
         mode: GuardMode,
     },
 
+    /// Return undefined immediately when the condition register is not exactly
+    /// `Bool(true)`.  Any other value — including `false`, `Undefined`, `Null`,
+    /// numbers, strings, etc. — causes an immediate return of `Undefined`.
+    ///
+    /// This is used by Azure Policy compilation to model "condition does not match"
+    /// without treating it as a VM assertion failure.
+    ReturnUndefinedIfNotTrue {
+        condition: u8,
+    },
+
+    /// Replace Undefined with Null in a register.
+    ///
+    /// Azure Policy treats missing fields as null rather than undefined.
+    /// This instruction prevents the RVM's undefined-propagation from
+    /// short-circuiting subsequent builtin calls.
+    CoalesceUndefinedToNull {
+        register: u8,
+    },
+
     /// Start a loop over a collection with specified semantics - uses parameter table
     LoopStart {
         /// Index into program's instruction_data.loop_params table

src/rvm/program/listing.rs

@@ -307,6 +307,14 @@ fn format_instruction_readable(
             let base = format!("{}LoadInput    r{} ← input", indent, dest);
             align_comment(&base, "Load global input document", config.comment_column)
         }
+        Instruction::LoadContext { dest } => {
+            let base = format!("{}LoadContext  r{} ← context", indent, dest);
+            align_comment(&base, "Load evaluation context", config.comment_column)
+        }
+        Instruction::LoadMetadata { dest } => {
+            let base = format!("{}LoadMetadata r{} ← metadata", indent, dest);
+            align_comment(&base, "Load program metadata", config.comment_column)
+        }
         Instruction::Move { dest, src } => {
             let base = format!("{}Move         r{} ← r{}", indent, dest, src);
             let comment = format!("Copy value from r{} to r{}", src, dest);
@@ -565,6 +573,11 @@ fn format_instruction_readable(
             let comment = format!("Append r{} to array r{}", value, arr);
             align_comment(&base, &comment, config.comment_column)
         }
+        Instruction::ArrayPushDefined { arr, value } => {
+            let base = format!("{}ArrayPushDef r{}.push(r{})", indent, arr, value);
+            let comment = format!("Append r{} to array r{} (skip if undefined)", value, arr);
+            align_comment(&base, &comment, config.comment_column)
+        }
         Instruction::ArrayCreate { params_index } => instruction_data
             .get_array_create_params(params_index)
             .map_or_else(
@@ -658,6 +671,25 @@ fn format_instruction_readable(
             };
             align_comment(&keyword, &comment, config.comment_column)
         }
+        Instruction::ReturnUndefinedIfNotTrue { condition } => {
+            let base = format!(
+                "{}ReturnUndefinedIfNotTrue if r{} != true return undefined",
+                indent, condition
+            );
+            let comment = format!(
+                "Return undefined unless r{} is exactly boolean true",
+                condition
+            );
+            align_comment(&base, &comment, config.comment_column)
+        }
+        Instruction::CoalesceUndefinedToNull { register } => {
+            let base = format!(
+                "{}CoalesceUndefinedToNull r{} = null if undefined",
+                indent, register
+            );
+            let comment = format!("Azure Policy: absent field → null (r{})", register);
+            align_comment(&base, &comment, config.comment_column)
+        }
         Instruction::LoopStart { params_index } => {
             instruction_data.get_loop_params(params_index).map_or_else(
                 || {
@@ -959,6 +991,8 @@ const fn get_instruction_name(instruction: &Instruction) -> &'static str {
         Instruction::LoadBool { .. } => "LOAD_BOOL",
         Instruction::LoadData { .. } => "LOAD_DATA",
         Instruction::LoadInput { .. } => "LOAD_INPUT",
+        Instruction::LoadContext { .. } => "LOAD_CONTEXT",
+        Instruction::LoadMetadata { .. } => "LOAD_METADATA",
         Instruction::Move { .. } => "MOVE",
         Instruction::Add { .. } => "ADD",
         Instruction::Sub { .. } => "SUB",
@@ -984,6 +1018,7 @@ const fn get_instruction_name(instruction: &Instruction) -> &'static str {
         Instruction::IndexLiteral { .. } => "INDEX_LIT",
         Instruction::ArrayNew { .. } => "ARRAY_NEW",
         Instruction::ArrayPush { .. } => "ARRAY_PUSH",
+        Instruction::ArrayPushDefined { .. } => "ARRAY_PUSH_DEF",
         Instruction::ArrayCreate { .. } => "ARRAY_CREATE",
         Instruction::SetNew { .. } => "SET_NEW",
         Instruction::SetAdd { .. } => "SET_ADD",
@@ -996,6 +1031,8 @@ const fn get_instruction_name(instruction: &Instruction) -> &'static str {
             crate::rvm::instructions::GuardMode::Condition => "ASSERT",
             crate::rvm::instructions::GuardMode::NotUndefined => "ASSERT_NOT_UNDEF",
         },
+        Instruction::ReturnUndefinedIfNotTrue { .. } => "RET_UNDEF_IF_NOT_TRUE",
+        Instruction::CoalesceUndefinedToNull { .. } => "COALESCE_UNDEF_TO_NULL",
         Instruction::LoopStart { .. } => "LOOP_START",
         Instruction::LoopNext { .. } => "LOOP_NEXT",
         Instruction::CallRule { .. } => "CALL_RULE",
@@ -1024,6 +1061,12 @@ fn format_operation_compact(
         Instruction::LoadInput { dest } => {
             format!("{}r{} ← input", indent, dest)
         }
+        Instruction::LoadContext { dest } => {
+            format!("{}r{} ← context", indent, dest)
+        }
+        Instruction::LoadMetadata { dest } => {
+            format!("{}r{} ← metadata", indent, dest)
+        }
         Instruction::LoadData { dest } => {
             format!("{}r{} ← data", indent, dest)
         }

src/rvm/tests/instruction_parser.rs

@@ -31,6 +31,8 @@ pub fn parse_instruction(text: &str) -> Result<Instruction> {
             "LoadBool" => parse_load_bool(params_text),
             "LoadData" => parse_load_data(params_text),
             "LoadInput" => parse_load_input(params_text),
+            "LoadContext" => parse_load_context(params_text),
+            "LoadMetadata" => parse_load_metadata(params_text),
             "Move" => parse_move(params_text),
             "Add" => parse_add(params_text),
             "Sub" => parse_sub(params_text),
@@ -59,6 +61,7 @@ pub fn parse_instruction(text: &str) -> Result<Instruction> {
             "ArrayCreate" => parse_array_create(params_text),
             "SetCreate" => parse_set_create(params_text),
             "ArrayPush" => parse_array_push(params_text),
+            "ArrayPushDefined" => parse_array_push_defined(params_text),
             "SetNew" => parse_set_new(params_text),
             "SetAdd" => parse_set_add(params_text),
             "Contains" => parse_contains(params_text),
@@ -78,6 +81,8 @@ pub fn parse_instruction(text: &str) -> Result<Instruction> {
             "ComprehensionAdd" => parse_comprehension_add(params_text),
             "ComprehensionBegin" => parse_comprehension_start(params_text),
             "ComprehensionYield" => parse_comprehension_add(params_text),
+            "ReturnUndefinedIfNotTrue" => parse_return_undefined_if_not_true(params_text),
+            "CoalesceUndefinedToNull" => parse_coalesce_undefined_to_null(params_text),
             _ => bail!("Unknown instruction: {}", name),
         }
     } else {
@@ -414,6 +419,16 @@ fn parse_array_push(params_text: &str) -> Result<Instruction> {
     })
 }
 
+fn parse_array_push_defined(params_text: &str) -> Result<Instruction> {
+    let params = parse_params(params_text)?;
+    let arr = get_param_u16(&params, "arr")?;
+    let value = get_param_u16(&params, "value")?;
+    Ok(Instruction::ArrayPushDefined {
+        arr: arr.try_into().unwrap(),
+        value: value.try_into().unwrap(),
+    })
+}
+
 fn parse_array_create(params_text: &str) -> Result<Instruction> {
     let params = parse_params(params_text)?;
     let params_index = get_param_u16(&params, "params_index")?;
@@ -555,6 +570,22 @@ fn parse_load_input(params_text: &str) -> Result<Instruction> {
     })
 }
 
+fn parse_load_context(params_text: &str) -> Result<Instruction> {
+    let params = parse_params(params_text)?;
+    let dest = get_param_u16(&params, "dest")?;
+    Ok(Instruction::LoadContext {
+        dest: dest.try_into().unwrap(),
+    })
+}
+
+fn parse_load_metadata(params_text: &str) -> Result<Instruction> {
+    let params = parse_params(params_text)?;
+    let dest = get_param_u16(&params, "dest")?;
+    Ok(Instruction::LoadMetadata {
+        dest: dest.try_into().unwrap(),
+    })
+}
+
 fn parse_mod(params_text: &str) -> Result<Instruction> {
     let params = parse_params(params_text)?;
     let dest = get_param_u16(&params, "dest")?;
@@ -660,3 +691,19 @@ fn parse_comprehension_add(params_text: &str) -> Result<Instruction> {
         key_reg,
     })
 }
+
+fn parse_return_undefined_if_not_true(params_text: &str) -> Result<Instruction> {
+    let params = parse_params(params_text)?;
+    let condition = get_param_u16(&params, "condition")?;
+    Ok(Instruction::ReturnUndefinedIfNotTrue {
+        condition: condition.try_into().unwrap(),
+    })
+}
+
+fn parse_coalesce_undefined_to_null(params_text: &str) -> Result<Instruction> {
+    let params = parse_params(params_text)?;
+    let register = get_param_u16(&params, "register")?;
+    Ok(Instruction::CoalesceUndefinedToNull {
+        register: register.try_into().unwrap(),
+    })
+}

src/rvm/tests/vm.rs

@@ -83,6 +83,12 @@ mod tests {
         data: Option<crate::Value>,
         #[serde(default)]
         input: Option<crate::Value>,
+        #[serde(default)]
+        context: Option<crate::Value>,
+        #[serde(default)]
+        metadata_language: Option<String>,
+        #[serde(default)]
+        metadata_annotations: Option<BTreeMap<String, crate::Value>>,
         literals: Vec<crate::Value>,
         #[serde(default)]
         rule_infos: Vec<RuleInfoSpec>,
@@ -266,6 +272,9 @@ mod tests {
         instruction_params: Option<InstructionParamsSpec>,
         data: Option<Value>,
         input: Option<Value>,
+        context: Option<Value>,
+        metadata_language: Option<String>,
+        metadata_annotations: Option<BTreeMap<String, Value>>,
         max_instructions: Option<usize>,
         host_await_responses: Option<Vec<HostAwaitResponseSpec>>,
         host_await_responses_run_to_completion: Option<Vec<HostAwaitResponseSpec>>,
@@ -285,6 +294,12 @@ mod tests {
             None
         };
 
+        let processed_context = if let Some(ref context_value) = context {
+            Some(process_value(context_value)?)
+        } else {
+            None
+        };
+
         let processed_rule_tree = if let Some(ref tree_value) = rule_tree {
             Some(process_value(tree_value)?)
         } else {
@@ -631,6 +646,21 @@ mod tests {
         program.max_rule_window_size = 255;
         program.dispatch_window_size = 50;
 
+        // Recompute derived flags since instructions were assigned directly
+        // (bypassing add_instruction which normally tracks has_host_await)
+        program.recompute_host_await_presence();
+
+        // Set metadata if provided
+        if let Some(lang) = metadata_language {
+            program.metadata.language = lang;
+        }
+        if let Some(annotations) = metadata_annotations {
+            program.metadata.annotations = annotations
+                .into_iter()
+                .map(|(key, value)| process_value(&value).map(|processed| (key, processed)))
+                .collect::<anyhow::Result<_>>()?;
+        }
+
         // Initialize resolved builtins if we have builtin info
         if !program.builtin_info_table.is_empty() {
             if let Err(e) = program.initialize_resolved_builtins() {
@@ -664,6 +694,10 @@ mod tests {
                 vm.set_input(input_value);
             }
 
+            if let Some(context_value) = processed_context.clone() {
+                vm.set_context(context_value);
+            }
+
             if let Some(limit) = max_instructions {
                 vm.set_max_instructions(limit);
             }
@@ -931,6 +965,9 @@ mod tests {
                     test_case.instruction_params.clone(),
                     test_case.data.clone(),
                     test_case.input.clone(),
+                    test_case.context.clone(),
+                    test_case.metadata_language.clone(),
+                    test_case.metadata_annotations.clone(),
                     test_case.max_instructions,
                     test_case.host_await_responses.clone(),
                     test_case.host_await_responses_run_to_completion.clone(),

src/rvm/vm/comprehension.rs

@@ -312,7 +312,7 @@ impl RegoVM {
                     *current_item =
                         Some(self.get_register(comprehension_context.value_reg)?.clone());
                 }
-                IterationState::Array { .. } => {}
+                IterationState::Array { .. } | IterationState::Single { .. } => {}
             }
 
             iter_state.advance();
@@ -468,7 +468,7 @@ impl RegoVM {
                         } => {
                             *current_item = Some(iteration_value.clone());
                         }
-                        IterationState::Array { .. } => {}
+                        IterationState::Array { .. } | IterationState::Single { .. } => {}
                     }
 
                     iter_state.advance();
@@ -599,7 +599,7 @@ impl RegoVM {
             } => {
                 *current_item = Some(self.get_register(value_reg)?.clone());
             }
-            IterationState::Array { .. } => {}
+            IterationState::Array { .. } | IterationState::Single { .. } => {}
         }
 
         Ok(())