kata-containers/regorus
1
0
Fork 0
mirror of https://github.com/microsoft/regorus.git synced 2026-08-05 02:16:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(hoist): pre-compute loop hoisting metadata at compilation time (#483)

Browse Source 
Introduce a compiler pass that analyzes and pre-computes loop hoisting information
during policy compilation. This hoisted metadata is stored in lookup tables and made
available to downstream consumers:

- interpreter: use HoistedLoop entries during evaluation (replaces runtime scanning)
- type inference: can leverage pre-computed loop structure for type propagation
- RVM compiler: will consume hoisting metadata for optimized bytecode generation

Changes:
- populate loop hoisting tables during engine preparation and query snippet execution
- refactor eval_stmts_in_loop and eval_output_expr_in_loop to consume HoistedLoop directly
- add helper methods for accessing loop expressions, collections, and indices from HoistedLoop
- extend Lookup with get_checked and into_slots for safe query context access and merging

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2025-10-08 11:11:30 -05:00
committed by GitHub
parent 9604fe86f1
commit 5d8387f4d9
14 changed files with 1365 additions and 267 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
benches/schema_validation_benchmark.rs
View File

@@ -269,7 +269,7 @@ fn bench_mixed_type_array(c: &mut Criterion) {
}
});
let schema = Schema::from_serde_json_value(schema_json).unwrap();
let value = Value::from(json!(["hello", 42, true, "world", 3.14, false]));
let value = Value::from(json!(["hello", 42, true, "world", 99.5, false]));
c.bench_function("validate_mixed_type_array", |b| {
b.iter(|| {

6
src/compiled_policy.rs
View File

@@ -2,6 +2,7 @@
// Licensed under the MIT License.
use crate::ast::*;
use crate::compiler::hoist::HoistedLoopsLookup;
use crate::engine::Engine;
use crate::scheduler::*;
use crate::utils::*;
@@ -190,7 +191,7 @@ pub(crate) struct TargetInfo {
#[derive(Debug, Clone, Default)]
pub(crate) struct CompiledPolicyData {
pub(crate) modules: Rc<Vec<Ref<Module>>>,
pub(crate) schedule: Option<Schedule>,
pub(crate) schedule: Option<Rc<Schedule>>,
pub(crate) rules: Map<String, Vec<Ref<Rule>>>,
pub(crate) default_rules: Map<String, Vec<DefaultRuleInfo>>,
pub(crate) imports: BTreeMap<String, Ref<Expr>>,
@@ -212,4 +213,7 @@ pub(crate) struct CompiledPolicyData {
// The semantics of extensions ought to be changes to be more Clone friendly.
pub(crate) extensions: Map<String, (u8, Rc<Box<dyn Extension>>)>,
// Pre-computed loop hoisting information
pub(crate) loop_hoisting_table: HoistedLoopsLookup,
}

10
src/compiler.rs Normal file
View File

@@ -0,0 +1,10 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
//! Compiler-related functionality for Regorus.
//!
//! This module contains utilities and data structures used during
//! the compilation phase to prepare policies for efficient execution.
pub mod context;
pub mod hoist;

153
src/compiler/context.rs Normal file
View File

@@ -0,0 +1,153 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
//! Compilation context types shared across compiler components.
//!
//! This module defines context structures used for tracking scope-level information
//! during compilation and analysis phases. These types are designed to be compatible
//! with both the interpreter's loop hoisting and the RVM compiler.
use crate::ast::ExprRef;
use alloc::collections::BTreeSet;
use alloc::string::{String, ToString};
/// Type of compilation context for tracking different scenarios
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ContextType {
/// Rule context (Complete, PartialSet, PartialObject, or Function)
Rule,
/// Comprehension context (Array, Set, or Object)
Comprehension,
/// Every quantifier context
Every,
/// Query/statement context (no output expressions)
Query,
}
/// Context for tracking variable bindings and output expressions within a scope.
/// Used during loop hoisting and compilation to determine what needs to be hoisted
/// and what's already bound.
///
/// This design is compatible with RVM's CompilationContext for potential future unification.
#[derive(Debug, Clone)]
pub struct ScopeContext {
/// Type of context (Rule, Comprehension, Every, Query)
pub context_type: ContextType,
/// Variables that are bound in the current scope
pub bound_vars: BTreeSet<String>,
/// Variables that are explicitly marked as unbound (from `some` declarations)
pub unbound_vars: BTreeSet<String>,
/// Key expression from rule head or object comprehension (for output expression hoisting)
pub key_expr: Option<ExprRef>,
/// Value expression from rule assignment or comprehension term (for output expression hoisting)
pub value_expr: Option<ExprRef>,
}
impl ScopeContext {
/// Create a new context with Query type (default, no output expressions)
pub fn new() -> Self {
Self {
context_type: ContextType::Query,
bound_vars: BTreeSet::new(),
unbound_vars: BTreeSet::new(),
key_expr: None,
value_expr: None,
}
}
/// Create a new context with a specific context type
#[allow(dead_code)]
pub fn with_context_type(context_type: ContextType) -> Self {
Self {
context_type,
bound_vars: BTreeSet::new(),
unbound_vars: BTreeSet::new(),
key_expr: None,
value_expr: None,
}
}
/// Create a new context with output expressions (for rules and comprehensions)
#[allow(dead_code)]
pub fn with_output_exprs(
context_type: ContextType,
key_expr: Option<ExprRef>,
value_expr: Option<ExprRef>,
) -> Self {
Self {
context_type,
bound_vars: BTreeSet::new(),
unbound_vars: BTreeSet::new(),
key_expr,
value_expr,
}
}
/// Create a child context that inherits bindings but overrides context type and output expressions
pub fn child_with_output_exprs(
&self,
context_type: ContextType,
key_expr: Option<ExprRef>,
value_expr: Option<ExprRef>,
) -> Self {
Self {
context_type,
bound_vars: self.bound_vars.clone(),
unbound_vars: self.unbound_vars.clone(),
key_expr,
value_expr,
}
}
/// Add a variable to the bound set
pub fn bind_variable(&mut self, var_name: &str) {
if var_name != "_" {
self.bound_vars.insert(var_name.to_string());
self.unbound_vars.remove(var_name);
}
}
/// Mark a variable as unbound
pub fn add_unbound_variable(&mut self, var_name: &str) {
if var_name != "_" && !self.bound_vars.contains(var_name) {
self.unbound_vars.insert(var_name.to_string());
}
}
/// Check if a variable is known to be unbound
pub fn is_unbound(&self, var_name: &str) -> bool {
self.unbound_vars.contains(var_name)
}
/// Check if we can determine that a variable should be treated as a loop iterator
/// (either it's unbound or explicitly marked as such)
pub fn should_hoist_as_loop(&self, var_name: &str) -> bool {
if var_name == "_" || self.is_unbound(var_name) {
true
} else {
// Treat variables that haven't been bound in this scope as potential loop iterators
!self.bound_vars.contains(var_name)
}
}
/// Create a child context inheriting parent bindings, output expressions, and context type
pub fn child(&self) -> Self {
Self {
context_type: self.context_type.clone(),
bound_vars: self.bound_vars.clone(),
unbound_vars: self.unbound_vars.clone(),
key_expr: self.key_expr.clone(),
value_expr: self.value_expr.clone(),
}
}
}
impl Default for ScopeContext {
fn default() -> Self {
Self::new()
}
}

939
src/compiler/hoist.rs Normal file
View File

@@ -0,0 +1,939 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
//! Loop hoisting functionality for compilation.
//!
//! This module contains code adapted from the RVM compiler to support
//! pre-computing loop hoisting information that can be stored in the
//! compiled policy and reused by the interpreter.
use crate::ast::{Expr, ExprRef, Literal, LiteralStmt, Module, Query, Ref, Rule, RuleHead};
use crate::compiler::context::{ContextType, ScopeContext};
use crate::lookup::Lookup;
use crate::*;
use anyhow::Result;
use alloc::vec::Vec;
/// Type of loop that was hoisted
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LoopType {
/// `array[_]` or `object[idx]` patterns
IndexIteration,
Walk,
}
/// Information about a hoisted loop
#[derive(Debug, Clone)]
pub struct HoistedLoop {
/// The loop expression itself (e.g., `array[_]`)
pub loop_expr: Option<ExprRef>,
/// Key/index variable (e.g., `_` or `idx`)
pub key: Option<ExprRef>,
/// Value expression (the result of indexing)
pub value: ExprRef,
/// Collection being iterated
pub collection: ExprRef,
/// Type of loop
#[allow(dead_code)]
pub loop_type: LoopType,
}
/// Lookup table mapping statements/expressions to their hoisted loops
#[derive(Debug, Clone, Default)]
pub struct HoistedLoopsLookup {
/// Maps (module_index, statement_index) -> Vec<HoistedLoop>
/// Stores pre-computed loops for each statement in rules/queries
statement_loops: Lookup<Vec<HoistedLoop>>,
/// Maps (module_index, expr_index) -> Vec<HoistedLoop>
/// For output expressions in comprehensions and rule values
expr_loops: Lookup<Vec<HoistedLoop>>,
/// Maps (module_index, query_index) -> ScopeContext
/// Stores compilation contexts for queries (rules, comprehensions, every)
query_contexts: Lookup<ScopeContext>,
}
impl HoistedLoopsLookup {
/// Create a new empty lookup table
pub fn new() -> Self {
Self::default()
}
/// Ensure capacity for a given module and statement
pub fn ensure_statement_capacity(&mut self, module_idx: u32, stmt_idx: u32) {
self.statement_loops.ensure_capacity(module_idx, stmt_idx);
self.expr_loops.ensure_capacity(module_idx, 0);
self.query_contexts.ensure_capacity(module_idx, 0);
}
/// Ensure capacity for a given module and expression
pub fn ensure_expr_capacity(&mut self, module_idx: u32, expr_idx: u32) {
self.expr_loops.ensure_capacity(module_idx, expr_idx);
self.statement_loops.ensure_capacity(module_idx, 0);
self.query_contexts.ensure_capacity(module_idx, 0);
}
/// Ensure capacity for a given module and query
pub fn ensure_query_capacity(&mut self, module_idx: u32, query_idx: u32) {
self.query_contexts.ensure_capacity(module_idx, query_idx);
self.statement_loops.ensure_capacity(module_idx, 0);
self.expr_loops.ensure_capacity(module_idx, 0);
}
/// Store hoisted loops for a statement
pub fn set_statement_loops(&mut self, module_idx: u32, stmt_idx: u32, loops: Vec<HoistedLoop>) {
self.statement_loops.set(module_idx, stmt_idx, loops);
}
/// Get hoisted loops for a statement
pub fn get_statement_loops(&self, module_idx: u32, stmt_idx: u32) -> Option<&Vec<HoistedLoop>> {
self.statement_loops.get_checked(module_idx, stmt_idx)
}
/// Store hoisted loops for an expression (output expressions)
pub fn set_expr_loops(&mut self, module_idx: u32, expr_idx: u32, loops: Vec<HoistedLoop>) {
self.expr_loops.set(module_idx, expr_idx, loops);
}
/// Get hoisted loops for an expression
pub fn get_expr_loops(&self, module_idx: u32, expr_idx: u32) -> Option<&Vec<HoistedLoop>> {
self.expr_loops.get_checked(module_idx, expr_idx)
}
/// Store the compilation context for a query
pub fn set_query_context(&mut self, module_idx: u32, query_idx: u32, context: ScopeContext) {
self.query_contexts.set(module_idx, query_idx, context);
}
/// Get the compilation context for a query
pub fn get_query_context(&self, module_idx: u32, query_idx: u32) -> Option<&ScopeContext> {
self.query_contexts.get_checked(module_idx, query_idx)
}
/// Merge another loop hoisting table into this one
/// This is used to add query module loops to the existing table
pub fn merge_query_loops(&mut self, mut other: HoistedLoopsLookup, module_idx: usize) {
while self.statement_loops.module_len() < module_idx {
self.statement_loops.push_module(Vec::new());
}
while self.expr_loops.module_len() < module_idx {
self.expr_loops.push_module(Vec::new());
}
while self.query_contexts.module_len() < module_idx {
self.query_contexts.push_module(Vec::new());
}
let query_module_idx = other.statement_loops.module_len().saturating_sub(1);
if let Some(module) = other.statement_loops.remove_module(query_module_idx) {
self.statement_loops.push_module(module);
}
if let Some(module) = other.expr_loops.remove_module(query_module_idx) {
self.expr_loops.push_module(module);
}
if let Some(module) = other.query_contexts.remove_module(query_module_idx) {
self.query_contexts.push_module(module);
}
}
pub fn truncate_modules(&mut self, module_count: usize) {
self.statement_loops.truncate_modules(module_count);
self.expr_loops.truncate_modules(module_count);
self.query_contexts.truncate_modules(module_count);
}
pub fn module_len(&self) -> usize {
self.statement_loops.module_len()
}
}
// Note: ScopeContext is now defined in src/compiler/context.rs
/// Loop hoister that populates the HoistedLoopsLookup table
pub struct LoopHoister {
lookup: HoistedLoopsLookup,
schedule: Option<crate::Rc<crate::scheduler::Schedule>>,
}
impl LoopHoister {
/// Create a new loop hoister
pub fn new() -> Self {
Self {
lookup: HoistedLoopsLookup::new(),
schedule: None,
}
}
/// Create a new loop hoister with a schedule
pub fn new_with_schedule(schedule: crate::Rc<crate::scheduler::Schedule>) -> Self {
Self {
lookup: HoistedLoopsLookup::new(),
schedule: Some(schedule),
}
}
/// Populate loop hoisting information for all modules
/// Returns the populated lookup table
pub fn populate(mut self, modules: &[Ref<Module>]) -> Result<HoistedLoopsLookup> {
for (module_idx, module) in modules.iter().enumerate() {
self.populate_module(module_idx as u32, module)?;
}
Ok(self.lookup)
}
/// Populate loop hoisting information for all modules, with extra capacity
/// for additional modules that will be added later (e.g., query modules)
///
/// # Arguments
/// * `modules` - The modules to populate
/// * `extra_capacity` - Number of additional module slots to reserve
pub fn populate_with_extra_capacity(
mut self,
modules: &[Ref<Module>],
extra_capacity: u32,
) -> Result<HoistedLoopsLookup> {
for (module_idx, module) in modules.iter().enumerate() {
self.populate_module(module_idx as u32, module)?;
}
// Ensure capacity for extra modules by ensuring capacity for a dummy statement
// in each extra module (this will resize the module vector)
let last_module_idx = modules.len() as u32;
for i in 0..extra_capacity {
self.lookup
.ensure_statement_capacity(last_module_idx + i, 0);
self.lookup.ensure_expr_capacity(last_module_idx + i, 0);
}
Ok(self.lookup)
}
/// Populate loop information for a single module
pub fn populate_module(&mut self, module_idx: u32, module: &Module) -> Result<()> {
// Process all rules in the module
for rule in &module.policy {
self.populate_rule(module_idx, rule)?;
}
Ok(())
}
/// Finalize and return the populated lookup table
pub fn finalize(self) -> HoistedLoopsLookup {
self.lookup
}
/// Populate loop hoisting information for a query snippet
/// Query snippets are treated like they're in a module appended at the end
/// This matches how the analyzer handles query snippets
///
/// # Arguments
/// * `module_idx` - The module index to use (typically modules.len())
/// * `query` - The query to populate
/// * `num_statements` - Total number of statements in the query module
/// * `num_expressions` - Total number of expressions in the query module
pub fn populate_query_snippet(
&mut self,
module_idx: u32,
query: &Query,
num_statements: u32,
num_expressions: u32,
) -> Result<()> {
// Ensure capacity for all possible statement and expression indices
// Indices are 0-based, so max index is count - 1
if num_statements > 0 {
self.lookup
.ensure_statement_capacity(module_idx, num_statements - 1);
}
if num_expressions > 0 {
self.lookup
.ensure_expr_capacity(module_idx, num_expressions - 1);
}
// Populate the query with default context
let context = ScopeContext::new();
self.lookup.ensure_query_capacity(module_idx, query.qidx);
self.populate_query(module_idx, query, &context).map(|_| ())
}
/// Populate loop information for a single rule
fn populate_rule(&mut self, module_idx: u32, rule: &Rule) -> Result<()> {
match rule {
Rule::Spec { head, bodies, .. } => {
// Create a context for this rule
let mut context = ScopeContext::new();
// Bind function parameters if this is a function rule
if let RuleHead::Func { args, .. } = head {
for param in args {
// Extract variable name from parameter expression
if let Expr::Var { span, .. } = param.as_ref() {
context.bind_variable(span.text());
}
}
}
// Extract key and value expressions from the rule head (matching RVM compiler pattern)
let (key_expr, value_expr) = match head {
RuleHead::Compr { refr, assign, .. } => {
let output_expr = assign.as_ref().map(|a| a.value.clone());
let key_expr = match refr.as_ref() {
Expr::RefBrack { index, .. } => {
// For RefBrack (e.g., p[key]), the index is the key expression
Some(index.clone())
}
_ => {
// For non-RefBrack (e.g., p), no key expression
None
}
};
(key_expr, output_expr)
}
RuleHead::Set { key, .. } => {
// For set rules, no separate key_expr, output_expr is the key
(None, key.clone())
}
RuleHead::Func { assign, .. } => {
// Function rules return the assignment value (if any)
(None, assign.as_ref().map(|a| a.value.clone()))
}
};
// Process each rule body (definitions)
for body in bodies {
// Create a context with the output expressions (using Rule context type)
let body_context = context.child_with_output_exprs(
ContextType::Rule,
key_expr.clone(),
value_expr.clone(),
);
// Store the context for this query
let populated_body_context =
self.populate_query(module_idx, &body.query, &body_context)?;
self.lookup
.ensure_query_capacity(module_idx, body.query.qidx);
self.lookup.set_query_context(
module_idx,
body.query.qidx,
populated_body_context.clone(),
);
// Process the key expression if present
if let Some(ref key) = key_expr {
self.populate_output_expr(module_idx, key, &populated_body_context)?;
}
// Process the head value expression (if present)
if let Some(ref head_value) = value_expr {
self.populate_output_expr(module_idx, head_value, &populated_body_context)?;
}
// Process the value expression if present in the assign
if let Some(ref assign) = body.assign {
self.populate_output_expr(
module_idx,
&assign.value,
&populated_body_context,
)?;
}
}
// Handle rules with head assignments but no bodies (e.g., `y := "string"`)
if bodies.is_empty() {
let body_context = context.child_with_output_exprs(
ContextType::Rule,
key_expr.clone(),
value_expr.clone(),
);
if let Some(ref key) = key_expr {
self.populate_output_expr(module_idx, key, &body_context)?;
}
if let Some(ref value) = value_expr {
self.populate_output_expr(module_idx, value, &body_context)?;
}
}
}
Rule::Default { value, .. } => {
// For default rules, just process the value expression
let context = ScopeContext::new();
self.populate_output_expr(module_idx, value, &context)?;
}
}
Ok(())
}
/// Populate loop information for a query (sequence of statements)
fn populate_query(
&mut self,
module_idx: u32,
query: &Query,
parent_context: &ScopeContext,
) -> Result<ScopeContext> {
let mut context = parent_context.child();
// Get the scheduled order if available
let stmt_order: Vec<usize> = if let Some(ref schedule) = self.schedule {
if let Some(query_schedule) = schedule.queries.get(module_idx, query.qidx) {
query_schedule
.order
.iter()
.map(|&idx| idx as usize)
.collect()
} else {
// No schedule for this query, use source order
(0..query.stmts.len()).collect()
}
} else {
// No schedule available, use source order
(0..query.stmts.len()).collect()
};
// Process statements in scheduled order
for &stmt_idx in &stmt_order {
let stmt = &query.stmts[stmt_idx];
self.populate_statement(module_idx, stmt.sidx, stmt, &mut context)?;
}
Ok(context)
}
/// Populate loop information for a single statement
fn populate_statement(
&mut self,
module_idx: u32,
stmt_idx: u32,
stmt: &LiteralStmt,
context: &mut ScopeContext,
) -> Result<()> {
// Handle SomeVars to mark variables as unbound
if let Literal::SomeVars { vars, .. } = &stmt.literal {
for var in vars {
context.add_unbound_variable(var.text());
}
}
// Traverse literal expressions to populate nested contexts (comprehensions, every, etc.)
self.process_literal_for_contexts(module_idx, &stmt.literal, context)?;
for with_mod in &stmt.with_mods {
self.process_expr_for_contexts(module_idx, &with_mod.refr, context)?;
self.process_expr_for_contexts(module_idx, &with_mod.r#as, context)?;
}
// Hoist loops from this statement using populated contexts
let loops =
self.hoist_loops_from_literal_with_context(module_idx, &stmt.literal, context)?;
// Always store in lookup table, even if no loops (store empty vec)
// This ensures the interpreter can always find an entry
self.lookup.ensure_statement_capacity(module_idx, stmt_idx);
self.lookup.set_statement_loops(module_idx, stmt_idx, loops);
// Update context based on variable bindings in this statement
self.update_context_from_literal(&stmt.literal, context);
Ok(())
}
/// Hoist loops from a literal with variable binding context
fn hoist_loops_from_literal_with_context(
&self,
module_idx: u32,
literal: &Literal,
context: &ScopeContext,
) -> Result<Vec<HoistedLoop>> {
let mut loops = Vec::new();
use Literal::*;
match literal {
SomeIn {
key,
value,
collection,
..
} => {
// Recursively hoist from sub-expressions first
if let Some(key) = key {
self.hoist_loops_from_expr_with_context(module_idx, key, &mut loops, context)?;
}
self.hoist_loops_from_expr_with_context(module_idx, value, &mut loops, context)?;
self.hoist_loops_from_expr_with_context(
module_idx, collection, &mut loops, context,
)?;
}
Expr { expr, .. } => {
// Hoist loops from expressions (like array[_] patterns)
self.hoist_loops_from_expr_with_context(module_idx, expr, &mut loops, context)?;
}
Every { domain, query, .. } => {
// Hoist from domain expression
self.hoist_loops_from_expr_with_context(module_idx, domain, &mut loops, context)?;
// Process the Every query in a child context
let child_context = self
.lookup
.get_query_context(module_idx, query.qidx)
.cloned()
.unwrap_or_else(|| context.child());
for stmt in &query.stmts {
self.hoist_loops_from_literal_with_context(
module_idx,
&stmt.literal,
&child_context,
)?;
}
}
NotExpr { expr, .. } => {
self.hoist_loops_from_expr_with_context(module_idx, expr, &mut loops, context)?;
}
_ => {
// Other literal types don't have loops to hoist
}
}
Ok(loops)
}
/// Traverse literals to populate nested contexts (comprehensions, every, etc.)
fn process_literal_for_contexts(
&mut self,
module_idx: u32,
literal: &Literal,
context: &ScopeContext,
) -> Result<()> {
use Literal::*;
match literal {
SomeIn {
key,
value,
collection,
..
} => {
if let Some(key_expr) = key {
self.process_expr_for_contexts(module_idx, key_expr, context)?;
}
self.process_expr_for_contexts(module_idx, value, context)?;
self.process_expr_for_contexts(module_idx, collection, context)?;
}
Expr { expr, .. } | NotExpr { expr, .. } => {
self.process_expr_for_contexts(module_idx, expr, context)?;
}
Every { domain, query, .. } => {
// Process the domain expression for nested contexts
self.process_expr_for_contexts(module_idx, domain, context)?;
// Create a child context for the Every quantifier
let every_context = context.child_with_output_exprs(ContextType::Every, None, None);
let populated_every_context =
self.populate_query(module_idx, query.as_ref(), &every_context)?;
self.lookup.ensure_query_capacity(module_idx, query.qidx);
self.lookup.set_query_context(
module_idx,
query.qidx,
populated_every_context.clone(),
);
// Nested query already processed for hoisting via populated context
}
_ => {}
}
Ok(())
}
/// Traverse expressions to populate nested contexts (comprehensions, function params, etc.)
fn process_expr_for_contexts(
&mut self,
module_idx: u32,
expr: &ExprRef,
context: &ScopeContext,
) -> Result<()> {
use crate::ast::Expr as E;
match expr.as_ref() {
E::Array { items, .. } | E::Set { items, .. } => {
for item in items {
self.process_expr_for_contexts(module_idx, item, context)?;
}
}
E::Object { fields, .. } => {
for (_, key_expr, value_expr) in fields {
self.process_expr_for_contexts(module_idx, key_expr, context)?;
self.process_expr_for_contexts(module_idx, value_expr, context)?;
}
}
E::ArrayCompr { term, query, .. } | E::SetCompr { term, query, .. } => {
let compr_context = context.child_with_output_exprs(
ContextType::Comprehension,
None,
Some(term.clone()),
);
let populated_compr_context =
self.populate_query(module_idx, query.as_ref(), &compr_context)?;
self.lookup.ensure_query_capacity(module_idx, query.qidx);
self.lookup.set_query_context(
module_idx,
query.qidx,
populated_compr_context.clone(),
);
self.populate_output_expr(module_idx, term, &populated_compr_context)?;
}
E::ObjectCompr {
key, value, query, ..
} => {
let compr_context = context.child_with_output_exprs(
ContextType::Comprehension,
Some(key.clone()),
Some(value.clone()),
);
let populated_compr_context =
self.populate_query(module_idx, query.as_ref(), &compr_context)?;
self.lookup.ensure_query_capacity(module_idx, query.qidx);
self.lookup.set_query_context(
module_idx,
query.qidx,
populated_compr_context.clone(),
);
self.populate_output_expr(module_idx, key, &populated_compr_context)?;
self.populate_output_expr(module_idx, value, &populated_compr_context)?;
}
E::Call { fcn, params, .. } => {
self.process_expr_for_contexts(module_idx, fcn, context)?;
for param in params {
self.process_expr_for_contexts(module_idx, param, context)?;
}
}
E::UnaryExpr { expr, .. } => {
self.process_expr_for_contexts(module_idx, expr, context)?;
}
E::RefDot { refr, .. } => {
self.process_expr_for_contexts(module_idx, refr, context)?;
}
E::RefBrack { refr, index, .. } => {
self.process_expr_for_contexts(module_idx, refr, context)?;
self.process_expr_for_contexts(module_idx, index, context)?;
}
E::BinExpr { lhs, rhs, .. }
| E::BoolExpr { lhs, rhs, .. }
| E::ArithExpr { lhs, rhs, .. } => {
self.process_expr_for_contexts(module_idx, lhs, context)?;
self.process_expr_for_contexts(module_idx, rhs, context)?;
}
E::AssignExpr { lhs, rhs, .. } => {
self.process_expr_for_contexts(module_idx, lhs, context)?;
self.process_expr_for_contexts(module_idx, rhs, context)?;
}
E::Membership {
key,
value,
collection,
..
} => {
if let Some(key_expr) = key {
self.process_expr_for_contexts(module_idx, key_expr, context)?;
}
self.process_expr_for_contexts(module_idx, value, context)?;
self.process_expr_for_contexts(module_idx, collection, context)?;
}
#[cfg(feature = "rego-extensions")]
E::OrExpr { lhs, rhs, .. } => {
self.process_expr_for_contexts(module_idx, lhs, context)?;
self.process_expr_for_contexts(module_idx, rhs, context)?;
}
_ => {}
}
Ok(())
}
/// Hoist loops from expressions with variable binding context
fn hoist_loops_from_expr_with_context(
&self,
module_idx: u32,
expr: &ExprRef,
loops: &mut Vec<HoistedLoop>,
context: &ScopeContext,
) -> Result<()> {
use Expr::*;
match expr.as_ref() {
// Primitive types - no loops to hoist
String { .. }
| RawString { .. }
| Number { .. }
| Bool { .. }
| Null { .. }
| Var { .. } => {
// No sub-expressions to process
}
// Collection types - hoist from items
Array { items, .. } => {
for item in items {
self.hoist_loops_from_expr_with_context(module_idx, item, loops, context)?;
}
}
Set { items, .. } => {
for item in items {
self.hoist_loops_from_expr_with_context(module_idx, item, loops, context)?;
}
}
Object { fields, .. } => {
for (_, key_expr, value_expr) in fields {
self.hoist_loops_from_expr_with_context(module_idx, key_expr, loops, context)?;
self.hoist_loops_from_expr_with_context(
module_idx, value_expr, loops, context,
)?;
}
}
// Comprehensions - process their queries
// Note: Comprehension contexts and output expressions will be handled
// by populate_comprehension called from the parent expression processing
ArrayCompr { term, query, .. } | SetCompr { term, query, .. } => {
let child_context = self
.lookup
.get_query_context(module_idx, query.qidx)
.cloned()
.unwrap_or_else(|| context.child());
for stmt in &query.stmts {
self.hoist_loops_from_literal_with_context(
module_idx,
&stmt.literal,
&child_context,
)?;
}
self.hoist_loops_from_expr_with_context(module_idx, term, loops, &child_context)?;
}
ObjectCompr {
key, value, query, ..
} => {
let child_context = self
.lookup
.get_query_context(module_idx, query.qidx)
.cloned()
.unwrap_or_else(|| context.child());
for stmt in &query.stmts {
self.hoist_loops_from_literal_with_context(
module_idx,
&stmt.literal,
&child_context,
)?;
}
self.hoist_loops_from_expr_with_context(module_idx, key, loops, &child_context)?;
self.hoist_loops_from_expr_with_context(module_idx, value, loops, &child_context)?;
}
// Function calls - check for walk() builtin which generates loops
Call { fcn, params, .. } => {
// First hoist loops in parameters.
for param in params {
self.hoist_loops_from_expr_with_context(module_idx, param, loops, context)?;
}
// Check if this is a walk() call
let is_walk = if let Var {
value: Value::String(name),
..
} = fcn.as_ref()
{
name.as_ref() == "walk"
} else {
false
};
if is_walk {
loops.push(HoistedLoop {
loop_expr: Some(expr.clone()),
key: None, // walk doesn't have an index
value: expr.clone(),
collection: expr.clone(), // The walk call itself
loop_type: LoopType::Walk,
});
return Ok(());
}
// For other function calls, hoist loops in parameters
}
// Unary expressions - hoist from operand
UnaryExpr { expr, .. } => {
self.hoist_loops_from_expr_with_context(module_idx, expr, loops, context)?;
}
// Reference expressions - check for array[_] patterns
RefDot { refr, .. } => {
self.hoist_loops_from_expr_with_context(module_idx, refr, loops, context)?;
}
RefBrack { refr, index, .. } => {
// Recursively hoist from sub-expressions
self.hoist_loops_from_expr_with_context(module_idx, refr, loops, context)?;
self.hoist_loops_from_expr_with_context(module_idx, index, loops, context)?;
// Check if the index expression contains unbound variables
// This handles both simple cases like array[x] and complex cases like array[[x, y]]
if Self::expr_contains_unbound_vars(index, context) {
// This index contains unbound variables - create a loop to iterate
loops.push(HoistedLoop {
loop_expr: Some(expr.clone()),
key: Some(index.clone()),
value: expr.clone(),
collection: refr.clone(),
loop_type: LoopType::IndexIteration,
});
return Ok(());
}
}
// Binary expressions - hoist from both operands
BinExpr { lhs, rhs, .. } => {
self.hoist_loops_from_expr_with_context(module_idx, lhs, loops, context)?;
self.hoist_loops_from_expr_with_context(module_idx, rhs, loops, context)?;
}
BoolExpr { lhs, rhs, .. } => {
self.hoist_loops_from_expr_with_context(module_idx, lhs, loops, context)?;
self.hoist_loops_from_expr_with_context(module_idx, rhs, loops, context)?;
}
ArithExpr { lhs, rhs, .. } => {
self.hoist_loops_from_expr_with_context(module_idx, lhs, loops, context)?;
self.hoist_loops_from_expr_with_context(module_idx, rhs, loops, context)?;
}
AssignExpr { lhs, rhs, .. } => {
self.hoist_loops_from_expr_with_context(module_idx, lhs, loops, context)?;
self.hoist_loops_from_expr_with_context(module_idx, rhs, loops, context)?;
}
// Membership expressions - hoist from key, value, and collection
Membership {
key,
value,
collection,
..
} => {
if let Some(key_expr) = key {
self.hoist_loops_from_expr_with_context(module_idx, key_expr, loops, context)?;
}
self.hoist_loops_from_expr_with_context(module_idx, value, loops, context)?;
self.hoist_loops_from_expr_with_context(module_idx, collection, loops, context)?;
}
// Handle conditionally compiled expression types
#[cfg(feature = "rego-extensions")]
OrExpr { lhs, rhs, .. } => {
self.hoist_loops_from_expr_with_context(module_idx, lhs, loops, context)?;
self.hoist_loops_from_expr_with_context(module_idx, rhs, loops, context)?;
}
}
Ok(())
}
/// Update context based on variable bindings in a literal
fn update_context_from_literal(&self, literal: &Literal, context: &mut ScopeContext) {
use crate::ast::Expr as E;
use Literal::*;
match literal {
SomeIn { key, value, .. } => {
// Bind the loop variables
if let Some(key_expr) = key {
if let E::Var { span, .. } = key_expr.as_ref() {
context.bind_variable(span.text());
}
}
if let E::Var { span, .. } = value.as_ref() {
context.bind_variable(span.text());
}
}
Expr { expr, .. } => {
// Look for assignment expressions that bind variables
if let E::AssignExpr { lhs, .. } = expr.as_ref() {
Self::bind_variables_from_expr(lhs, context);
}
}
_ => {}
}
}
/// Recursively bind variables from an expression (for assignments)
fn bind_variables_from_expr(expr: &ExprRef, context: &mut ScopeContext) {
use crate::ast::Expr as E;
match expr.as_ref() {
E::Var { span, .. } => {
context.bind_variable(span.text());
}
E::Array { items, .. } => {
for item in items {
Self::bind_variables_from_expr(item, context);
}
}
E::Object { fields, .. } => {
for (_, key_expr, value_expr) in fields {
Self::bind_variables_from_expr(key_expr, context);
Self::bind_variables_from_expr(value_expr, context);
}
}
_ => {}
}
}
/// Check if an expression contains any unbound variables that should trigger loop hoisting
fn expr_contains_unbound_vars(expr: &ExprRef, context: &ScopeContext) -> bool {
use crate::ast::Expr as E;
match expr.as_ref() {
E::Var {
value: Value::String(var_name),
..
} => context.should_hoist_as_loop(var_name.as_ref()),
E::Array { items, .. } | E::Set { items, .. } => items
.iter()
.any(|item| Self::expr_contains_unbound_vars(item, context)),
E::Object { fields, .. } => fields.iter().any(|(_, _, value_expr)| {
// For objects check only the value expression can be bound
Self::expr_contains_unbound_vars(value_expr, context)
}),
// Other expressions don't contribute unbound vars from parent scope
_ => false,
}
}
/// Populate loop information for output expressions (rule values, comprehension terms)
fn populate_output_expr(
&mut self,
module_idx: u32,
expr: &ExprRef,
context: &ScopeContext,
) -> Result<()> {
let mut loops = Vec::new();
self.hoist_loops_from_expr_with_context(module_idx, expr, &mut loops, context)?;
// Always store expression loops, even if empty
// This ensures the interpreter can always find an entry
let expr_idx = expr.as_ref().eidx();
self.lookup.ensure_expr_capacity(module_idx, expr_idx);
self.lookup.set_expr_loops(module_idx, expr_idx, loops);
// Traverse child expressions to populate any nested contexts (e.g., comprehensions)
self.process_expr_for_contexts(module_idx, expr, context)?;
Ok(())
}
}
impl Default for LoopHoister {
fn default() -> Self {
Self::new()
}
}

74
src/engine.rs
View File

@@ -846,6 +846,67 @@ impl Engine {
query_module.num_statements = parser.num_statements();
let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;
// Populate loop hoisting for the query snippet
// Query snippets are treated as if they're in a module appended at the end (same as analyzer)
// The loop hoisting table already has capacity for this (ensured in prepare_for_eval)
let module_idx = self.modules.len() as u32;
use crate::compiler::hoist::LoopHoister;
let query_schedule_rc = Rc::new(query_schedule.clone());
let mut hoister = LoopHoister::new_with_schedule(query_schedule_rc);
hoister.populate_query_snippet(
module_idx,
&query_node,
query_module.num_statements,
query_module.num_expressions,
)?;
let query_loops = hoister.finalize();
#[cfg(debug_assertions)]
{
debug_assert_eq!(
query_loops.module_len(),
module_idx as usize + 1,
"query hoisting table missing expected module slot {}",
module_idx
);
for stmt in &query_node.stmts {
debug_assert!(
query_loops
.get_statement_loops(module_idx, stmt.sidx)
.is_some(),
"missing hoisted loop entry for query statement index {}",
stmt.sidx
);
}
}
// Get the existing table, merge in the query loops, and set it back
let mut existing_table = self.interpreter.take_loop_hoisting_table();
existing_table.truncate_modules(self.modules.len());
#[cfg(debug_assertions)]
{
debug_assert!(
existing_table.module_len() <= self.modules.len(),
"loop hoisting table should not retain extra modules before merge"
);
}
existing_table.merge_query_loops(query_loops, self.modules.len());
#[cfg(debug_assertions)]
{
for stmt in &query_node.stmts {
debug_assert!(
existing_table
.get_statement_loops(module_idx, stmt.sidx)
.is_some(),
"missing hoisted loop entry after merge for module {} stmt {}",
module_idx,
stmt.sidx
);
}
}
self.interpreter.set_loop_hoisting_table(existing_table);
Ok((Ref::new(query_module), query_node, query_schedule))
}
@@ -873,9 +934,8 @@ impl Engine {
if !self.prepared {
// Analyze the modules and determine how statements must be scheduled.
let analyzer = Analyzer::new();
let schedule = analyzer.analyze(&self.modules)?;
let schedule = Rc::new(analyzer.analyze(&self.modules)?);
self.interpreter.set_schedule(Some(schedule));
self.interpreter.set_modules(self.modules.clone());
self.interpreter.clear_builtins_cache();
@@ -889,6 +949,16 @@ impl Engine {
self.interpreter.gather_rules()?;
self.interpreter.process_imports()?;
// Populate loop hoisting table for efficient evaluation
// Reserve capacity for 1 extra module (for query modules)
use crate::compiler::hoist::LoopHoister;
let hoister = LoopHoister::new_with_schedule(schedule.clone());
let loop_lookup = hoister.populate_with_extra_capacity(&self.modules, 0)?;
self.interpreter.set_loop_hoisting_table(loop_lookup);
// Set schedule after hoisting completes
self.interpreter.set_schedule(Some(schedule));
#[cfg(feature = "azure_policy")]
if for_target {
// Resolve and validate target specifications across all modules

195
src/interpreter.rs
View File

@@ -6,6 +6,7 @@ use crate::builtins::{self, BuiltinFcn};
use crate::compiled_policy::CompiledPolicyData;
#[cfg(feature = "azure_policy")]
use crate::compiled_policy::TargetInfo;
use crate::compiler::hoist::HoistedLoop;
use crate::lexer::*;
use crate::lookup::Lookup;
use crate::parser::Parser;
@@ -23,10 +24,6 @@ use core::ops::Bound::*;
type Scope = BTreeMap<SourceStr, Value>;
type ExprLookup = Lookup<Value>;
mod loops;
use loops::LoopExpr;
#[cfg(feature = "azure_policy")]
pub mod error;
#[cfg(feature = "azure_policy")]
@@ -270,7 +267,7 @@ impl Interpreter {
Rc::make_mut(&mut self.compiled_policy)
}
pub fn set_schedule(&mut self, schedule: Option<Schedule>) {
pub fn set_schedule(&mut self, schedule: Option<Rc<Schedule>>) {
self.compiled_policy_mut().schedule = schedule;
}
@@ -282,6 +279,17 @@ impl Interpreter {
self.compiled_policy_mut().modules = modules;
}
pub fn set_loop_hoisting_table(&mut self, table: crate::compiler::hoist::HoistedLoopsLookup) {
self.compiled_policy_mut().loop_hoisting_table = table;
}
pub fn take_loop_hoisting_table(&mut self) -> crate::compiler::hoist::HoistedLoopsLookup {
core::mem::replace(
&mut self.compiled_policy_mut().loop_hoisting_table,
crate::compiler::hoist::HoistedLoopsLookup::new(),
)
}
pub fn get_data_mut(&mut self) -> &mut Value {
&mut self.data
}
@@ -368,6 +376,26 @@ impl Interpreter {
self.get_loop_var_value(expr).is_some()
}
#[inline]
fn loop_assignment_expr(loop_info: &HoistedLoop) -> &ExprRef {
loop_info.loop_expr.as_ref().unwrap_or(&loop_info.value)
}
#[inline]
fn loop_index_expr(loop_info: &HoistedLoop) -> Option<&ExprRef> {
loop_info.key.as_ref().or(None)
}
#[inline]
fn loop_collection_expr(loop_info: &HoistedLoop) -> &ExprRef {
&loop_info.collection
}
#[inline]
fn loop_span(loop_info: &HoistedLoop) -> Span {
loop_info.value.span().clone()
}
fn ensure_loop_var_values_capacity(&mut self) {
for (module_idx, module) in self.compiled_policy.modules.iter().enumerate() {
self.loop_var_values
@@ -1132,6 +1160,7 @@ impl Interpreter {
.push(Self::make_expression_result(span, &Value::Bool(true)))
}
}
// https://github.com/open-policy-agent/opa/issues/1622#issuecomment-520547385
matches!(value, Value::Bool(false) | Value::Undefined)
}
@@ -1338,7 +1367,11 @@ impl Interpreter {
}
}
fn eval_stmts_in_loop(&mut self, stmts: &[&LiteralStmt], loops: &[LoopExpr]) -> Result<bool> {
fn eval_stmts_in_loop(
&mut self,
stmts: &[&LiteralStmt],
loops: &[HoistedLoop],
) -> Result<bool> {
if loops.is_empty() {
if !stmts.is_empty() {
// Evaluate the current statement whose loop expressions have been hoisted.
@@ -1355,20 +1388,20 @@ impl Interpreter {
self.eval_stmts(stmts)
}
} else {
let loop_expr = &loops[0];
let loop_info = &loops[0];
let mut result = false;
// Apply with modifiers before evaluating the loop expression.
let (saved_state, _) = self.apply_with_modifiers(stmts[0])?;
let loop_expr_value = loop_expr.value();
let loop_expr_value = if let Expr::Call {
let collection_expr = Self::loop_collection_expr(loop_info).clone();
let loop_value = if let Expr::Call {
span, fcn, params, ..
} = loop_expr_value.as_ref()
} = collection_expr.as_ref()
{
// Handle walk(obj, output_param)
let extra_arg = get_extra_arg(
&loop_expr_value,
&collection_expr,
Some(self.current_module_path.as_str()),
&self.compiled_policy.functions,
);
@@ -1378,9 +1411,9 @@ impl Interpreter {
} else {
&params[..]
};
self.eval_call_impl(span, &loop_expr_value, fcn, params)?
self.eval_call_impl(span, &collection_expr, fcn, params)?
} else {
self.eval_expr(&loop_expr_value)?
self.eval_expr(&collection_expr)?
};
// Restore with modifiers.
@@ -1390,13 +1423,13 @@ impl Interpreter {
// If the loop's index variable h<as already been assigned a value
// (this can happen if the same index is used for two different collections),
// then evaluate statements only if the index applies to this collection.
let loop_expr_index = loop_expr.index();
let index_expr = Self::loop_index_expr(loop_info);
if let Some(Expr::Var {
span: index_var, ..
}) = loop_expr_index.as_ref().map(|r| r.as_ref())
}) = index_expr.map(|r| r.as_ref())
{
if let Some(idx) = self.lookup_local_var(&index_var.source_str()) {
if loop_expr_value[&idx] != Value::Undefined {
if loop_value[&idx] != Value::Undefined {
result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
return Ok(result);
} else if idx != Value::Undefined {
@@ -1410,19 +1443,20 @@ impl Interpreter {
self.scopes.push(Scope::default());
let query_result = self.get_current_context()?.result.clone();
match loop_expr_value {
let loop_target_expr = Self::loop_assignment_expr(loop_info);
match loop_value {
Value::Array(items) => {
for (idx, v) in items.iter().enumerate() {
self.set_loop_var_value(&loop_expr.expr(), v.clone());
self.set_loop_var_value(loop_target_expr, v.clone());
let exec = if let Some(index) = loop_expr.index() {
let exec = if let Some(index) = index_expr {
let mut type_match = BTreeSet::new();
let mut cache = BTreeMap::new();
self.make_bindings(
false,
&mut type_match,
&mut cache,
&index,
index,
&Value::from(idx),
true,
)?
@@ -1442,17 +1476,17 @@ impl Interpreter {
}
}
self.remove_loop_var_value(&loop_expr.expr());
self.remove_loop_var_value(loop_target_expr);
}
Value::Set(items) => {
for v in items.iter() {
self.set_loop_var_value(&loop_expr.expr(), v.clone());
self.set_loop_var_value(loop_target_expr, v.clone());
// For sets, index is also the value.
let exec = if let Some(index) = loop_expr.index() {
let exec = if let Some(index) = index_expr {
let mut type_match = BTreeSet::new();
let mut cache = BTreeMap::new();
self.make_bindings(false, &mut type_match, &mut cache, &index, v, true)?
self.make_bindings(false, &mut type_match, &mut cache, index, v, true)?
} else {
true
};
@@ -1468,16 +1502,16 @@ impl Interpreter {
}
}
}
self.remove_loop_var_value(&loop_expr.expr());
self.remove_loop_var_value(loop_target_expr);
}
Value::Object(obj) => {
for (k, v) in obj.iter() {
self.set_loop_var_value(&loop_expr.expr(), v.clone());
self.set_loop_var_value(loop_target_expr, v.clone());
// For objects, index is key.
let exec = if let Some(index) = loop_expr.index() {
let exec = if let Some(index) = index_expr {
let mut type_match = BTreeSet::new();
let mut cache = BTreeMap::new();
self.make_bindings(false, &mut type_match, &mut cache, &index, k, true)?
self.make_bindings(false, &mut type_match, &mut cache, index, k, true)?
} else {
true
};
@@ -1493,7 +1527,7 @@ impl Interpreter {
}
}
}
self.remove_loop_var_value(&loop_expr.expr());
self.remove_loop_var_value(loop_target_expr);
}
Value::Undefined => {
result = false;
@@ -1646,7 +1680,7 @@ impl Interpreter {
Ok(is_const && self.is_simple_literal(output_expr)?)
}
fn eval_output_expr_in_loop(&mut self, loops: &[LoopExpr]) -> Result<bool> {
fn eval_output_expr_in_loop(&mut self, loops: &[HoistedLoop]) -> Result<bool> {
if loops.is_empty() {
let (key_expr, output_expr) = self.get_exprs_from_context()?;
@@ -1818,36 +1852,36 @@ impl Interpreter {
}
// Try out values in current loop expr.
let loop_expr = &loops[0];
let loop_info = &loops[0];
let mut result = false;
match self.eval_expr(&loop_expr.value())? {
let loop_target_expr = Self::loop_assignment_expr(loop_info);
match self.eval_expr(Self::loop_collection_expr(loop_info))? {
Value::Array(items) => {
for v in items.iter() {
self.set_loop_var_value(&loop_expr.expr(), v.clone());
self.set_loop_var_value(loop_target_expr, v.clone());
result = self.eval_output_expr_in_loop(&loops[1..])? || result;
}
}
Value::Set(items) => {
for v in items.iter() {
self.set_loop_var_value(&loop_expr.expr(), v.clone());
self.set_loop_var_value(loop_target_expr, v.clone());
result = self.eval_output_expr_in_loop(&loops[1..])? || result;
}
}
Value::Object(obj) => {
for (_, v) in obj.iter() {
self.set_loop_var_value(&loop_expr.expr(), v.clone());
self.set_loop_var_value(loop_target_expr, v.clone());
result = self.eval_output_expr_in_loop(&loops[1..])? || result;
}
}
_ => {
return Err(loop_expr.span().source.error(
loop_expr.span().line,
loop_expr.span().col,
"item cannot be indexed",
));
let span = Self::loop_span(loop_info);
return Err(span
.source
.error(span.line, span.col, "item cannot be indexed"));
}
}
self.remove_loop_var_value(&loop_expr.expr());
self.remove_loop_var_value(loop_target_expr);
Ok(result)
}
@@ -1867,13 +1901,38 @@ impl Interpreter {
// Evaluate output expression after all the statements have been executed.
let (key_expr, output_expr) = self.get_exprs_from_context()?;
let mut loops = vec![];
let mut loops: Vec<HoistedLoop> = Vec::new();
// Get pre-computed loops for key expression
if let Some(ke) = &key_expr {
self.hoist_loops_impl(ke, &mut loops);
match self
.compiled_policy
.loop_hoisting_table
.get_expr_loops(self.current_module_index, ke.as_ref().eidx())
{
Some(hoisted_loops) => {
loops.extend(hoisted_loops.iter().cloned());
}
None => {
bail!(ke.span().error("Loop hoisting information not found for key expression. This is likely a bug in the compilation phase."));
}
}
}
// Get pre-computed loops for output expression
if let Some(oe) = &output_expr {
self.hoist_loops_impl(oe, &mut loops);
match self
.compiled_policy
.loop_hoisting_table
.get_expr_loops(self.current_module_index, oe.as_ref().eidx())
{
Some(hoisted_loops) => {
loops.extend(hoisted_loops.iter().cloned());
}
None => {
bail!(oe.span().error("Loop hoisting information not found for output expression. This is likely a bug in the compilation phase."));
}
}
}
let r = self.eval_output_expr_in_loop(&loops[..])?;
@@ -1895,7 +1954,23 @@ impl Interpreter {
break;
}
let loop_exprs = self.hoist_loops(&stmt.literal);
// Get pre-computed hoisted loops from compilation phase
let loop_exprs = match self
.compiled_policy
.loop_hoisting_table
.get_statement_loops(self.current_module_index, stmt.sidx)
{
Some(hoisted_loops) => {
// Use pre-computed loops from compilation phase
hoisted_loops.clone()
}
None => {
// Loop hoisting should have been done during compilation
// If we reach here, it means the hoisting pass didn't process this statement
bail!(stmt.span.error("Loop hoisting information not found for statement. This is likely a bug in the compilation phase."));
}
};
if !loop_exprs.is_empty() {
// If there are hoisted loop expressions, execute subsequent statements
// within loops.
@@ -1903,6 +1978,7 @@ impl Interpreter {
}
result = self.eval_stmt(stmt, &stmts[idx + 1..])?;
if matches!(&stmt.literal, Literal::SomeIn { .. }) {
return Ok(result);
}
@@ -2774,7 +2850,12 @@ impl Interpreter {
&& !self.compiled_policy.default_rules.contains_key(&rule_path)
&& !self.compiled_policy.imports.contains_key(&rule_path)
{
bail!(span.error("var is unsafe"));
bail!(span.error(&format!(
"var {} is unsafe (path {:?}, scopes {:?})",
name.text(),
path,
self.scopes
)));
}
// Find the rule to which the var being looked up corresponds to. This is the prefix for
@@ -3119,12 +3200,22 @@ impl Interpreter {
}
fn find_module_index(&self, module: &Ref<Module>) -> u32 {
self.compiled_policy
if let Some(idx) = self
.compiled_policy
.modules
.iter()
.position(|m| core::ptr::eq(m.as_ref(), module.as_ref()))
.map(|i| i as u32)
.unwrap_or(0)
{
idx as u32
} else if let Some(query_module) = &self.query_module {
if core::ptr::eq(query_module.as_ref(), module.as_ref()) {
self.compiled_policy.modules.len() as u32
} else {
0
}
} else {
0
}
}
fn get_rule_refr(rule: &Rule) -> &ExprRef {
@@ -3951,6 +4042,12 @@ impl Interpreter {
compiled_policy.rule_to_evaluate = "".into();
}
// Populate loop hoisting lookup table
use crate::compiler::hoist::LoopHoister;
let hoister = LoopHoister::new();
let loop_lookup = hoister.populate(compiled_policy.modules.as_ref())?;
compiled_policy.loop_hoisting_table = loop_lookup;
Ok(self.compiled_policy.clone())
}
}

204
src/interpreter/loops.rs
View File

@@ -1,204 +0,0 @@
use crate::ast::*;
use crate::interpreter::Interpreter;
use crate::lexer::*;
use crate::scheduler::traverse;
use crate::utils::get_path_string;
use crate::value::Value;
use crate::*;
#[derive(Debug)]
pub enum LoopExpr {
Loop {
span: Span,
expr: Ref<Expr>,
value: Ref<Expr>,
index: Ref<Expr>,
},
Walk {
span: Span,
expr: Ref<Expr>,
},
}
impl LoopExpr {
pub fn span(&self) -> Span {
match self {
Self::Loop { span, .. } => span.clone(),
Self::Walk { span, .. } => span.clone(),
}
}
pub fn value(&self) -> Ref<Expr> {
match self {
Self::Loop { value, .. } => value.clone(),
Self::Walk { expr, .. } => expr.clone(),
}
}
pub fn expr(&self) -> Ref<Expr> {
match self {
Self::Loop { expr, .. } => expr.clone(),
Self::Walk { expr, .. } => expr.clone(),
}
}
pub fn index(&self) -> Option<Ref<Expr>> {
match self {
Self::Loop { index, .. } => Some(index.clone()),
Self::Walk { .. } => None,
}
}
}
impl Interpreter {
pub(super) fn hoist_loops_impl(&self, expr: &ExprRef, loops: &mut Vec<LoopExpr>) {
use Expr::*;
match expr.as_ref() {
RefBrack {
refr, index, span, ..
} => {
// First hoist any loops in refr
self.hoist_loops_impl(refr, loops);
// hoist any loops in index expression.
self.hoist_loops_impl(index, loops);
// Then hoist the current bracket operation.
let mut indices = Vec::with_capacity(1);
let _ = traverse(index, &mut |e| match e.as_ref() {
Var { span: ident, .. } if self.is_loop_index_var(&ident.source_str()) => {
indices.push(ident.source_str());
Ok(false)
}
Array { .. } | Object { .. } => Ok(true),
_ => Ok(false),
});
if !indices.is_empty() {
loops.push(LoopExpr::Loop {
span: span.clone(),
expr: expr.clone(),
value: refr.clone(),
index: index.clone(),
})
}
}
// Primitives
String { .. }
| RawString { .. }
| Number { .. }
| Bool { .. }
| Null { .. }
| Var { .. } => (),
// Recurse into expressions in other variants.
Array { items, .. } | Set { items, .. } | Call { params: items, .. } => {
for item in items {
self.hoist_loops_impl(item, loops);
}
// Handle walk builtin which acts as a generator.
// TODO: Handle with modifier on the walk builtin.
if let Expr::Call { fcn, .. } = expr.as_ref() {
if let Ok(fcn_path) = get_path_string(fcn, None) {
if fcn_path == "walk" {
// TODO: Use an enum for LoopExpr to handle walk
loops.push(LoopExpr::Walk {
span: expr.span().clone(),
expr: expr.clone(),
})
}
}
}
}
Object { fields, .. } => {
for (_, key, value) in fields {
self.hoist_loops_impl(key, loops);
self.hoist_loops_impl(value, loops);
}
}
RefDot { refr: expr, .. } | UnaryExpr { expr, .. } => {
self.hoist_loops_impl(expr, loops)
}
BinExpr { lhs, rhs, .. }
| BoolExpr { lhs, rhs, .. }
| ArithExpr { lhs, rhs, .. }
| AssignExpr { lhs, rhs, .. } => {
self.hoist_loops_impl(lhs, loops);
self.hoist_loops_impl(rhs, loops);
}
#[cfg(feature = "rego-extensions")]
OrExpr { lhs, rhs, .. } => {
self.hoist_loops_impl(lhs, loops);
self.hoist_loops_impl(rhs, loops);
}
Membership {
key,
value,
collection,
..
} => {
if let Some(key) = key.as_ref() {
self.hoist_loops_impl(key, loops);
}
self.hoist_loops_impl(value, loops);
self.hoist_loops_impl(collection, loops);
}
// The output expressions of comprehensions must be subject to hoisting
// only after evaluating the body of the comprehensions since the output
// expressions may depend on variables defined within the body.
ArrayCompr { .. } | SetCompr { .. } | ObjectCompr { .. } => (),
}
}
pub(super) fn hoist_loops(&self, literal: &Literal) -> Vec<LoopExpr> {
let mut loops = vec![];
use Literal::*;
match literal {
SomeVars { .. } => (),
SomeIn {
key,
value,
collection,
..
} => {
if let Some(key) = key {
self.hoist_loops_impl(key, &mut loops);
}
self.hoist_loops_impl(value, &mut loops);
self.hoist_loops_impl(collection, &mut loops);
}
Every {
domain: collection, ..
} => self.hoist_loops_impl(collection, &mut loops),
Expr { expr, .. } | NotExpr { expr, .. } => self.hoist_loops_impl(expr, &mut loops),
}
loops
}
pub(super) fn is_loop_index_var(&self, ident: &SourceStr) -> bool {
// TODO: check for vars that are declared using some-vars
match ident.text() {
"_" => true,
_ => match self.lookup_local_var(ident) {
// Vars declared using `some v` can be loop vars.
// They are initialized to undefined.
Some(Value::Undefined) => true,
// If ident is a local var (in current or parent scopes),
// then it is not a loop var.
Some(_) => false,
None => {
// Check if ident is a rule.
let path = self.current_module_path.clone() + "." + ident.text();
!self.compiled_policy.rules.contains_key(&path)
}
},
}
}
}

1
src/lib.rs
View File

@@ -26,6 +26,7 @@ mod ast;
mod builtins;
mod compile;
mod compiled_policy;
mod compiler;
mod engine;
mod indexchecker;
mod interpreter;

29
src/lookup.rs
View File

@@ -53,6 +53,15 @@ impl<T: Clone> Lookup<T> {
self.slots[module_idx as usize][node_idx as usize].as_ref()
}
/// Get data using direct indices with bounds checking.
/// Returns None if the module or node index is out of range or unset.
pub fn get_checked(&self, module_idx: u32, node_idx: u32) -> Option<&T> {
self.slots
.get(module_idx as usize)
.and_then(|module| module.get(node_idx as usize))
.and_then(|slot| slot.as_ref())
}
/// Get mutable reference to data using direct indices (bounds assumed to be ensured).
pub fn get_mut(&mut self, module_idx: u32, node_idx: u32) -> &mut Option<T> {
&mut self.slots[module_idx as usize][node_idx as usize]
@@ -66,4 +75,24 @@ impl<T: Clone> Lookup<T> {
self.slots[module_idx as usize][node_idx as usize] = None;
}
}
pub fn truncate_modules(&mut self, module_count: usize) {
self.slots.truncate(module_count);
}
pub fn module_len(&self) -> usize {
self.slots.len()
}
pub fn push_module(&mut self, module: Vec<Option<T>>) {
self.slots.push(module);
}
pub fn remove_module(&mut self, module_idx: usize) -> Option<Vec<Option<T>>> {
if module_idx < self.slots.len() {
Some(self.slots.remove(module_idx))
} else {
None
}
}
}

1
src/scheduler.rs
View File

@@ -417,7 +417,6 @@ pub struct Analyzer {
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct Schedule {
pub queries: Lookup<QuerySchedule>,
}

2
src/schema/tests/validate/resource.rs
View File

@@ -857,7 +857,7 @@ fn test_complex_nested_azure_template_validation() {
let value = Value::from(valid_template);
let result = SchemaValidator::validate(&value, &schema);
std::dbg!(&result);
assert!(
result.is_ok(),
"Complex valid template should pass validation"

3
src/target/tests/deserialize.rs
View File

@@ -268,8 +268,7 @@ fn test_target_deserialization_with_registry_schemas() {
// We expect this to fail with a "not found in registry" error since we haven't
// populated the registries with test data
if result.is_err() {
let error = result.unwrap_err();
if let Err(error) = result {
assert!(matches!(
error,
TargetError::JsonParseError(_) | TargetError::DeserializationError(_)

13
src/tests/interpreter/mod.rs
View File

@@ -82,7 +82,7 @@ mod load_target_definitions {
use crate::registry::targets;
// Load target definitions
let _ = load()?;
load()?;
// Check that the sample targets were loaded
assert!(
@@ -184,12 +184,13 @@ pub fn process_value(v: &Value) -> Result<Value> {
fn match_values(computed: &Value, expected: &Value) -> Result<()> {
if computed != expected {
let expected_yaml = serde_yaml::to_string(&expected)?;
let computed_yaml = serde_yaml::to_string(&computed)?;
panic!(
"{}",
prettydiff::diff_chars(
&serde_yaml::to_string(&expected)?,
&serde_yaml::to_string(&computed)?
)
"expected:\n{}computed:\n{}diff:\n{}",
expected_yaml,
computed_yaml,
prettydiff::diff_chars(&expected_yaml, &computed_yaml)
);
}
Ok(())