feat!: Rego -> RVM Compiler and extensive testsuite (#506)

# RVM compiler test cases

Coverage:
- arithmetic
- arrays
- chained lookups
- comparisons
- comprehensions
- default rules
- destructuring
- function rules
- loops/quantifiers
- multiple entrypoints
- objects/sets
- variables
- negative/edge scenarios such as data/rule conflicts
- virtual data lookups
- etc

 # Modify interpreter and compiled policy for RVM Compilation

- Interpreter::eval_default_rule_for_compiler:
   evaluates a named default rule in isolation - allows compiler to emit a constant value instead of instructions
   for the default value

#  feat: Rego Compiler Scaffolding

- Introduce the rego::compiler module surface and entry point wiring
- Add the core compiler concepts:
  - register allocator
  - scope tracking
  - literal/builtin tables
  - rule worklists
  - instruction emit helpers
  - compiler-specific error types
  - context structs for rules, comprehensions, and loops to support later lowering passes.

# feat: Compile Rules/Queries

- add compiler::compile_from_policy workflow plus rule worklist, entry-point wiring, and recursion checks
- implement query lowering:
  - scheduling-aware statement ordering
  - loop hoisting
  - “every/some” semantics
  - context yields
  -  literal assertions
- finalize Program construction

# feat: Expression Lowering

- add compile_rego_expr and helpers to translate every AST expression into RVM instructions,
- interop with binding plans, comprehensions, and membership checks.
- implement collection literal builders (ArrayCreate, SetCreate, ObjectCreate)
  - dedupe literal keys and handle mixed literal/dynamic fields via instruction data blocks.
- operations:
  - arithmetic/boolean/bin operators
  - membership
  - unary minus
  - set unions/intersections
  - etc
- user-defined and builtin function calls
- reference handling
  - analyse chained refs
  - distinguishe data/input/local roots
  - perform rule dispatch or virtual document lookups
  - emits optimized Index/ChainedIndex instructions.

# feat: Comprehensions & Loops

- shared comprehension emitter
 - wraps array/set/object comprehensions with ComprehensionBegin/End
 - context management
- loop lowering utilities
 - read hoisting metadata
 - emit LoopStart/LoopNext
 - some in lowering
 - every quantifiers
 - index iteration
 - propagate binding plans into stored registers so downstream statements see bound variables.

# feat: Destructuring Lowering

- destructuring planner integration
 - assignment/parameter/loop bindings use hoisted plans instead of re-walking ASTs.
- handle :=, =, wildcard matches, and equality
 - evaluate RHS
 - applying destructuring plans
 - emit assert condition as needed
- support nested array/object destructuring, dynamic keys, and some ... in forms

# test: Shared Testing + RVM Suites

- move YAML test helpers into test_utils.rs and re-export via common.rs for use by interpreter and vm test suites
- comprehensive compiler test suite
  - compiles policies with the new Rego→RVM compiler
  - runs them through RegoVM
  - compares against interpreter behavior
  - supports multiple entry points
  - provides assembly listings
  - filterable YAML suites.

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2025-11-24 12:08:37 -06:00
committed by GitHub
parent 688e6128d4
commit a3a20a1235
43 changed files with 6945 additions and 139 deletions
+406
View File
@@ -0,0 +1,406 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
use super::{CompilationContext, Compiler, CompilerError, ContextType, Register, Result};
use crate::ast::{self, ExprRef, LiteralStmt, Query};
use crate::compiler::destructuring_planner::plans::BindingPlan;
use crate::compiler::hoist::{HoistedLoop, LoopType};
use crate::lexer::Span;
use crate::rvm::instructions::{LoopMode, LoopStartParams};
use crate::rvm::Instruction;
use crate::Value;
use alloc::format;
use alloc::string::ToString;
use alloc::vec::Vec;
impl<'a> Compiler<'a> {
pub(super) fn get_statement_loops(&self, stmt: &LiteralStmt) -> Result<Vec<HoistedLoop>> {
self.policy
.inner
.loop_hoisting_table
.get_statement_loops(self.current_module_index, stmt.sidx)
.cloned()
.ok_or_else(|| CompilerError::General {
message: format!(
"missing loop hoisting data for statement at {}:{}",
stmt.span.line, stmt.span.col
),
})
}
pub(super) fn get_expr_loops(&self, expr: &ExprRef) -> Vec<HoistedLoop> {
let module_idx = self.current_module_index;
let expr_idx = expr.as_ref().eidx();
self.policy
.inner
.loop_hoisting_table
.get_expr_loops(module_idx, expr_idx)
.cloned()
.unwrap_or_default()
}
pub(super) fn compile_hoisted_loops(
&mut self,
stmts: &[&LiteralStmt],
loops: &[HoistedLoop],
) -> Result<()> {
if loops.is_empty() {
if !stmts.is_empty() {
self.compile_single_statement(stmts[0])?;
return self.hoist_loops_and_compile_statements(&stmts[1..]);
} else {
self.hoist_loops_and_emit_context_yield()?;
}
}
let current_loop = &loops[0];
let remaining_loops = &loops[1..];
match current_loop.loop_type {
LoopType::IndexIteration => {
self.compile_index_iteration_loop(
&current_loop.loop_expr,
&current_loop.key,
&current_loop.value,
&current_loop.collection,
stmts,
remaining_loops,
)?;
Ok(())
}
LoopType::Walk => Err(CompilerError::General {
message: "walk loops are not yet supported in the RVM compiler".to_string(),
}),
}
}
pub(super) fn compile_every_quantifier(
&mut self,
key: &Option<Span>,
value: &Span,
domain: &ExprRef,
query: &Query,
span: &Span,
) -> Result<()> {
let collection_reg = self.compile_rego_expr(domain)?;
let key_reg = self.alloc_register();
let value_reg = self.alloc_register();
let result_reg = self.alloc_register();
let value_var_name = value.text().to_string();
let key_var_name = key.as_ref().map(|k| k.text().to_string());
let actual_key_reg = if key_var_name.is_none() || key_var_name.as_deref() == Some("_") {
value_reg
} else {
key_reg
};
let loop_params_index = self.program.add_loop_params(LoopStartParams {
mode: LoopMode::Every,
collection: collection_reg,
key_reg: actual_key_reg,
value_reg,
result_reg,
body_start: 0,
loop_end: 0,
});
self.emit_instruction(
Instruction::LoopStart {
params_index: loop_params_index,
},
span,
);
let body_start = self.program.instructions.len() as u16;
self.push_scope();
let every_context = CompilationContext {
context_type: ContextType::Every,
dest_register: result_reg,
key_expr: None,
value_expr: None,
span: span.clone(),
key_value_loops_hoisted: false,
};
self.push_context(every_context);
self.add_variable(&value_var_name, value_reg);
if let Some(ref key_name) = key_var_name {
self.add_variable(key_name, key_reg);
}
self.compile_query(query)?;
self.pop_context();
self.pop_scope();
self.emit_instruction(
Instruction::LoopNext {
body_start,
loop_end: 0,
},
span,
);
let loop_end = self.program.instructions.len() as u16;
self.program
.update_loop_params(loop_params_index, |params| {
params.body_start = body_start;
params.loop_end = loop_end;
});
let loop_next_idx = self.program.instructions.len() - 1;
if let Instruction::LoopNext {
loop_end: ref mut end,
..
} = &mut self.program.instructions[loop_next_idx]
{
*end = loop_end;
}
Ok(())
}
fn compile_index_iteration_loop(
&mut self,
loop_expr: &Option<ExprRef>,
key_var: &Option<ExprRef>,
_value_var: &ExprRef,
collection: &ExprRef,
remaining_stmts: &[&LiteralStmt],
remaining_loops: &[HoistedLoop],
) -> Result<()> {
let collection_reg = self.compile_rego_expr(collection)?;
let key_reg = self.alloc_register();
let value_reg = self.alloc_register();
let result_reg = self.alloc_register();
if let Some(loop_expr) = loop_expr {
self.loop_expr_register_map
.insert(loop_expr.clone(), value_reg);
}
let mut key_binding_plan: Option<(BindingPlan, Span)> = None;
if let Some(key_var) = key_var {
if let Some(binding_plan) = self.get_binding_plan_for_expr(key_var) {
if let BindingPlan::LoopIndex { .. } = &binding_plan {
key_binding_plan = Some((binding_plan, key_var.span().clone()));
} else {
return Err(CompilerError::UnexpectedBindingPlan {
context: format!("loop index pattern {}", key_var.span().text()),
found: format!("{binding_plan:?}"),
});
}
} else {
match key_var.as_ref() {
ast::Expr::Var { value, .. } => {
let var_name = match value {
Value::String(s) => {
if s.as_ref() == "_" {
"".to_string()
} else {
s.to_string()
}
}
_ => value.to_string(),
};
if !var_name.is_empty() && var_name != "_" {
self.store_variable(var_name, key_reg);
}
}
_ => {
return Err(CompilerError::MissingBindingPlan {
context: format!("loop index pattern {}", key_var.span().text()),
});
}
}
}
self.loop_expr_register_map.insert(key_var.clone(), key_reg);
}
let loop_params_index = self.program.add_loop_params(LoopStartParams {
mode: LoopMode::ForEach,
collection: collection_reg,
key_reg,
value_reg,
result_reg,
body_start: 0,
loop_end: 0,
});
self.emit_instruction(
Instruction::LoopStart {
params_index: loop_params_index,
},
collection.span(),
);
let body_start = self.program.instructions.len() as u16;
if let Some((binding_plan, plan_span)) = key_binding_plan.as_ref() {
self.apply_binding_plan(binding_plan, key_reg, plan_span)
.map_err(CompilerError::from)?;
}
let body_stmts = &remaining_stmts[0..];
self.compile_hoisted_loops(body_stmts, remaining_loops)?;
self.emit_instruction(
Instruction::LoopNext {
body_start,
loop_end: 0,
},
collection.span(),
);
let loop_end = self.program.instructions.len() as u16;
self.program
.update_loop_params(loop_params_index, |params| {
params.body_start = body_start;
params.loop_end = loop_end;
});
let loop_next_idx = self.program.instructions.len() - 1;
if let Instruction::LoopNext {
loop_end: ref mut end,
..
} = &mut self.program.instructions[loop_next_idx]
{
*end = loop_end;
}
Ok(())
}
pub(super) fn compile_some_in_loop_with_remaining_statements(
&mut self,
key: &Option<ExprRef>,
value: &ExprRef,
collection: &ExprRef,
remaining_stmts: &[&LiteralStmt],
) -> Result<Register> {
let loop_body_stmts = &remaining_stmts[1..];
self.compile_some_in_loop_with_body(key, value, collection, loop_body_stmts)
}
fn compile_some_in_loop_with_body(
&mut self,
key: &Option<ExprRef>,
value: &ExprRef,
collection: &ExprRef,
loop_body_stmts: &[&LiteralStmt],
) -> Result<Register> {
let collection_reg = self.compile_rego_expr(collection)?;
let key_reg = self.alloc_register();
let value_reg = self.alloc_register();
let result_reg = self.alloc_register();
let loop_params_index = self.program.add_loop_params(LoopStartParams {
mode: LoopMode::ForEach,
collection: collection_reg,
key_reg,
value_reg,
result_reg,
body_start: 0,
loop_end: 0,
});
self.emit_instruction(
Instruction::LoopStart {
params_index: loop_params_index,
},
collection.span(),
);
let body_start = self.program.instructions.len() as u16;
if let Some(binding_plan) = self.get_binding_plan_for_expr(collection) {
if let BindingPlan::SomeIn {
key_plan,
value_plan,
..
} = &binding_plan
{
let key_register = key_plan.as_ref().map(|_| key_reg);
self.apply_some_in_binding_plan(
key_plan.as_ref(),
key_register,
value_plan,
value_reg,
collection.span(),
)
.map_err(CompilerError::from)?;
} else {
return Err(CompilerError::UnexpectedBindingPlan {
context: format!("some-in binding {}", collection.span().text()),
found: format!("{binding_plan:?}"),
});
}
} else {
if let Some(key_expr) = key {
match key_expr.as_ref() {
ast::Expr::Var {
value: var_name, ..
} => {
let var_name = var_name.as_string()?.to_string();
self.store_variable(var_name, key_reg);
}
_ => {
return Err(CompilerError::MissingBindingPlan {
context: format!("some-in key pattern {}", key_expr.span().text()),
});
}
}
}
match value.as_ref() {
ast::Expr::Var {
value: var_name, ..
} => {
let var_name = var_name.as_string()?.to_string();
self.store_variable(var_name, value_reg);
}
_ => {
return Err(CompilerError::MissingBindingPlan {
context: format!("some-in value pattern {}", value.span().text()),
});
}
}
}
self.hoist_loops_and_compile_statements(loop_body_stmts)?;
self.emit_instruction(
Instruction::LoopNext {
body_start,
loop_end: 0,
},
collection.span(),
);
let loop_end = self.program.instructions.len() as u16;
self.program
.update_loop_params(loop_params_index, |params| {
params.body_start = body_start;
params.loop_end = loop_end;
});
let loop_next_idx = self.program.instructions.len() - 1;
if let Instruction::LoopNext {
loop_end: ref mut end,
..
} = &mut self.program.instructions[loop_next_idx]
{
*end = loop_end;
}
Ok(result_reg)
}
}