From 5d8387f4d96b69e8ab4e4037fd95608bb2263177 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi <35780660+anakrish@users.noreply.github.com> Date: Wed, 8 Oct 2025 11:11:30 -0500 Subject: [PATCH] feat(hoist): pre-compute loop hoisting metadata at compilation time (#483) 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 --- benches/schema_validation_benchmark.rs | 2 +- src/compiled_policy.rs | 6 +- src/compiler.rs | 10 + src/compiler/context.rs | 153 ++++ src/compiler/hoist.rs | 939 +++++++++++++++++++++++++ src/engine.rs | 74 +- src/interpreter.rs | 195 +++-- src/interpreter/loops.rs | 204 ------ src/lib.rs | 1 + src/lookup.rs | 29 + src/scheduler.rs | 1 - src/schema/tests/validate/resource.rs | 2 +- src/target/tests/deserialize.rs | 3 +- src/tests/interpreter/mod.rs | 13 +- 14 files changed, 1365 insertions(+), 267 deletions(-) create mode 100644 src/compiler.rs create mode 100644 src/compiler/context.rs create mode 100644 src/compiler/hoist.rs delete mode 100644 src/interpreter/loops.rs diff --git a/benches/schema_validation_benchmark.rs b/benches/schema_validation_benchmark.rs index f13e199..9d6d277 100644 --- a/benches/schema_validation_benchmark.rs +++ b/benches/schema_validation_benchmark.rs @@ -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(|| { diff --git a/src/compiled_policy.rs b/src/compiled_policy.rs index 213e448..6a9db74 100644 --- a/src/compiled_policy.rs +++ b/src/compiled_policy.rs @@ -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>>, - pub(crate) schedule: Option, + pub(crate) schedule: Option>, pub(crate) rules: Map>>, pub(crate) default_rules: Map>, pub(crate) imports: BTreeMap>, @@ -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>)>, + + // Pre-computed loop hoisting information + pub(crate) loop_hoisting_table: HoistedLoopsLookup, } diff --git a/src/compiler.rs b/src/compiler.rs new file mode 100644 index 0000000..2999ca9 --- /dev/null +++ b/src/compiler.rs @@ -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; diff --git a/src/compiler/context.rs b/src/compiler/context.rs new file mode 100644 index 0000000..6bb1c45 --- /dev/null +++ b/src/compiler/context.rs @@ -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, + + /// Variables that are explicitly marked as unbound (from `some` declarations) + pub unbound_vars: BTreeSet, + + /// Key expression from rule head or object comprehension (for output expression hoisting) + pub key_expr: Option, + + /// Value expression from rule assignment or comprehension term (for output expression hoisting) + pub value_expr: Option, +} + +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, + value_expr: Option, + ) -> 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, + value_expr: Option, + ) -> 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() + } +} diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs new file mode 100644 index 0000000..0b15fa5 --- /dev/null +++ b/src/compiler/hoist.rs @@ -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, + + /// Key/index variable (e.g., `_` or `idx`) + pub key: Option, + + /// 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 + /// Stores pre-computed loops for each statement in rules/queries + statement_loops: Lookup>, + + /// Maps (module_index, expr_index) -> Vec + /// For output expressions in comprehensions and rule values + expr_loops: Lookup>, + + /// Maps (module_index, query_index) -> ScopeContext + /// Stores compilation contexts for queries (rules, comprehensions, every) + query_contexts: Lookup, +} + +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) { + 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> { + 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) { + 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> { + 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>, +} + +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) -> 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]) -> Result { + 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], + extra_capacity: u32, + ) -> Result { + 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 { + let mut context = parent_context.child(); + + // Get the scheduled order if available + let stmt_order: Vec = 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> { + 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, + 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() + } +} diff --git a/src/engine.rs b/src/engine.rs index 21d1578..81df4c3 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -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 diff --git a/src/interpreter.rs b/src/interpreter.rs index 830424c..2c70d02 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -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; type ExprLookup = Lookup; -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) { + pub fn set_schedule(&mut self, schedule: Option>) { 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 { + fn eval_stmts_in_loop( + &mut self, + stmts: &[&LiteralStmt], + loops: &[HoistedLoop], + ) -> Result { 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 { ¶ms[..] }; - 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 { 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 { + fn eval_output_expr_in_loop(&mut self, loops: &[HoistedLoop]) -> Result { 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 = 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) -> 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()) } } diff --git a/src/interpreter/loops.rs b/src/interpreter/loops.rs deleted file mode 100644 index 238831f..0000000 --- a/src/interpreter/loops.rs +++ /dev/null @@ -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, - value: Ref, - index: Ref, - }, - Walk { - span: Span, - expr: Ref, - }, -} - -impl LoopExpr { - pub fn span(&self) -> Span { - match self { - Self::Loop { span, .. } => span.clone(), - Self::Walk { span, .. } => span.clone(), - } - } - - pub fn value(&self) -> Ref { - match self { - Self::Loop { value, .. } => value.clone(), - Self::Walk { expr, .. } => expr.clone(), - } - } - - pub fn expr(&self) -> Ref { - match self { - Self::Loop { expr, .. } => expr.clone(), - Self::Walk { expr, .. } => expr.clone(), - } - } - - pub fn index(&self) -> Option> { - 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) { - 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 { - 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) - } - }, - } - } -} diff --git a/src/lib.rs b/src/lib.rs index 1429cd1..6e8bfcd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,6 +26,7 @@ mod ast; mod builtins; mod compile; mod compiled_policy; +mod compiler; mod engine; mod indexchecker; mod interpreter; diff --git a/src/lookup.rs b/src/lookup.rs index caa31e8..7a8a234 100644 --- a/src/lookup.rs +++ b/src/lookup.rs @@ -53,6 +53,15 @@ impl Lookup { 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 { &mut self.slots[module_idx as usize][node_idx as usize] @@ -66,4 +75,24 @@ impl Lookup { 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>) { + self.slots.push(module); + } + + pub fn remove_module(&mut self, module_idx: usize) -> Option>> { + if module_idx < self.slots.len() { + Some(self.slots.remove(module_idx)) + } else { + None + } + } } diff --git a/src/scheduler.rs b/src/scheduler.rs index 3cc2e3c..17cab72 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -417,7 +417,6 @@ pub struct Analyzer { } #[derive(Debug, Clone)] -#[allow(dead_code)] pub struct Schedule { pub queries: Lookup, } diff --git a/src/schema/tests/validate/resource.rs b/src/schema/tests/validate/resource.rs index 07bc8a6..3f54e5d 100644 --- a/src/schema/tests/validate/resource.rs +++ b/src/schema/tests/validate/resource.rs @@ -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" diff --git a/src/target/tests/deserialize.rs b/src/target/tests/deserialize.rs index 1385023..bc31bd8 100644 --- a/src/target/tests/deserialize.rs +++ b/src/target/tests/deserialize.rs @@ -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(_) diff --git a/src/tests/interpreter/mod.rs b/src/tests/interpreter/mod.rs index c0a2548..984d184 100644 --- a/src/tests/interpreter/mod.rs +++ b/src/tests/interpreter/mod.rs @@ -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 { 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(())