diff --git a/src/engine.rs b/src/engine.rs index 19c14e4..21d1578 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -737,28 +737,13 @@ impl Engine { self.interpreter.clean_internal_evaluation_state(); self.interpreter.create_rule_prefixes()?; - let query_module = { - let source = Source::from_contents( - "".to_owned(), - "package __internal_query_module".to_owned(), - )?; - Ref::new(Parser::new(&source)?.parse()?) - }; - - // Parse the query. - let query_source = Source::from_contents("".to_string(), query)?; - let mut parser = self.make_parser(&query_source)?; - let query_node = parser.parse_user_query()?; + let (query_module, query_node, query_schedule) = self.make_query(query)?; if query_node.span.text() == "data" { self.eval_modules(enable_tracing)?; } - let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?; - self.interpreter.eval_user_query( - &query_module, - &query_node, - &query_schedule, - enable_tracing, - ) + + self.interpreter + .eval_user_query(&query_module, &query_node, query_schedule, enable_tracing) } /// Evaluate a Rego query that produces a boolean value. @@ -843,6 +828,27 @@ impl Engine { !matches!(self.eval_bool_query(query, enable_tracing), Ok(false)) } + fn make_query(&mut self, query: String) -> Result<(NodeRef, NodeRef, Schedule)> { + let mut query_module = { + let source = Source::from_contents( + "".to_owned(), + "package __internal_query_module".to_owned(), + )?; + Parser::new(&source)?.parse()? + }; + + // Parse the query. + let query_source = Source::from_contents("".to_string(), query)?; + let mut parser = self.make_parser(&query_source)?; + let query_node = parser.parse_user_query()?; + query_module.num_expressions = parser.num_expressions(); + query_module.num_queries = parser.num_queries(); + query_module.num_statements = parser.num_statements(); + let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?; + + Ok((Ref::new(query_module), query_node, query_schedule)) + } + #[doc(hidden)] /// Evaluate the given query and all the rules in the supplied policies. /// @@ -854,25 +860,9 @@ impl Engine { ) -> Result { self.eval_modules(enable_tracing)?; - let query_module = { - let source = Source::from_contents( - "".to_owned(), - "package __internal_query_module".to_owned(), - )?; - Ref::new(Parser::new(&source)?.parse()?) - }; - - // Parse the query. - let query_source = Source::from_contents("".to_string(), query)?; - let mut parser = self.make_parser(&query_source)?; - let query_node = parser.parse_user_query()?; - let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?; - self.interpreter.eval_user_query( - &query_module, - &query_node, - &query_schedule, - enable_tracing, - ) + let (query_module, query_node, query_schedule) = self.make_query(query)?; + self.interpreter + .eval_user_query(&query_module, &query_node, query_schedule, enable_tracing) } #[doc(hidden)] diff --git a/src/interpreter.rs b/src/interpreter.rs index 1c2052e..80f8b04 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -7,6 +7,7 @@ use crate::compiled_policy::CompiledPolicyData; #[cfg(feature = "azure_policy")] use crate::compiled_policy::TargetInfo; use crate::lexer::*; +use crate::lookup::Lookup; use crate::parser::Parser; use crate::scheduler::*; use crate::utils::*; @@ -20,6 +21,11 @@ use anyhow::{anyhow, bail, Result}; use core::ops::Bound::*; type Scope = BTreeMap; +type ExprLookup = Lookup; + +mod loops; + +use loops::LoopExpr; #[cfg(feature = "azure_policy")] pub mod error; @@ -68,6 +74,9 @@ pub struct Interpreter { module: Option>, current_module_path: String, + current_module_index: u32, + query_schedule: Option, + query_module: Option>, input: Value, init_data: Value, @@ -75,7 +84,7 @@ pub struct Interpreter { with_functions: BTreeMap, scopes: Vec, // TODO: handle recursive calls where same expr could have different values. - loop_var_values: BTreeMap, + loop_var_values: ExprLookup, contexts: Vec, processed: BTreeSet>, @@ -118,7 +127,7 @@ impl Clone for Interpreter { // Hence, they need not be copied. processed: BTreeSet::default(), processed_paths: Value::new_object(), - loop_var_values: BTreeMap::default(), + loop_var_values: ExprLookup::new(), scopes: Vec::default(), rule_values: BTreeMap::default(), @@ -126,6 +135,9 @@ impl Clone for Interpreter { active_rules: Vec::default(), contexts: Vec::default(), current_module_path: String::default(), + current_module_index: 0, + query_schedule: None, + query_module: None, module: None, no_rules_lookup: false, } @@ -167,50 +179,6 @@ impl Default for Context { } } -#[derive(Debug)] -enum LoopExpr { - Loop { - span: Span, - expr: Ref, - value: Ref, - index: Ref, - }, - Walk { - span: Span, - expr: Ref, - }, -} - -impl LoopExpr { - fn span(&self) -> Span { - match self { - Self::Loop { span, .. } => span.clone(), - Self::Walk { span, .. } => span.clone(), - } - } - - fn value(&self) -> Ref { - match self { - Self::Loop { value, .. } => value.clone(), - Self::Walk { expr, .. } => expr.clone(), - } - } - - fn expr(&self) -> Ref { - match self { - Self::Loop { expr, .. } => expr.clone(), - Self::Walk { expr, .. } => expr.clone(), - } - } - - fn index(&self) -> Option> { - match self { - Self::Loop { index, .. } => Some(index.clone()), - Self::Walk { .. } => None, - } - } -} - impl Interpreter { pub fn new() -> Interpreter { let compiled_policy = compiled_policy::CompiledPolicyData { @@ -224,6 +192,9 @@ impl Interpreter { module: None, current_module_path: String::default(), + current_module_index: 0, + query_schedule: None, + query_module: None, input: Value::Undefined, init_data: Value::new_object(), @@ -231,7 +202,7 @@ impl Interpreter { with_functions: BTreeMap::default(), scopes: Vec::default(), contexts: Vec::default(), - loop_var_values: BTreeMap::default(), + loop_var_values: Lookup::default(), processed: BTreeSet::default(), processed_paths: Value::new_object(), @@ -259,13 +230,16 @@ impl Interpreter { module: None, current_module_path: String::default(), + current_module_index: 0, + query_module: None, + query_schedule: None, input: Value::Undefined, with_document: Value::new_object(), with_functions: BTreeMap::default(), scopes: Vec::default(), contexts: Vec::default(), - loop_var_values: BTreeMap::default(), + loop_var_values: Lookup::default(), processed: BTreeSet::default(), processed_paths: Value::new_object(), @@ -360,13 +334,49 @@ impl Interpreter { self.data = self.init_data.clone(); self.processed.clear(); self.processed_paths = Value::new_object(); - self.loop_var_values.clear(); + self.ensure_loop_var_values_capacity(); self.scopes = vec![Scope::new()]; self.contexts = vec![]; self.rule_values.clear(); self.builtins_cache.clear(); } + // Helper methods for working with ExprLookup + fn set_loop_var_value(&mut self, expr: &ExprRef, value: Value) { + let module_idx = self.current_module_index; + let expr_idx = expr.eidx(); + self.loop_var_values.set(module_idx, expr_idx, value); + } + + fn get_loop_var_value(&self, expr: &ExprRef) -> Option<&Value> { + let module_idx = self.current_module_index; + let expr_idx = expr.eidx(); + self.loop_var_values.get(module_idx, expr_idx) + } + + fn remove_loop_var_value(&mut self, expr: &ExprRef) { + let module_idx = self.current_module_index; + let expr_idx = expr.eidx(); + self.loop_var_values.clear(module_idx, expr_idx); + } + + fn has_loop_var_value(&self, expr: &ExprRef) -> bool { + self.get_loop_var_value(expr).is_some() + } + + fn ensure_loop_var_values_capacity(&mut self) { + for (module_idx, module) in self.compiled_policy.modules.iter().enumerate() { + self.loop_var_values + .ensure_capacity(module_idx as u32, module.num_expressions); + } + if let Some(query_module) = &self.query_module { + let query_module_idx = self.compiled_policy.modules.len(); + let eidx = query_module.num_expressions; + self.loop_var_values + .ensure_capacity(query_module_idx as u32, eidx); + } + } + fn current_module(&self) -> Result> { self.module .clone() @@ -422,7 +432,7 @@ impl Interpreter { // Collect a chaing of '.field' or '["field"]' let mut path = vec![]; loop { - if let Some(v) = self.loop_var_values.get(expr) { + if let Some(v) = self.get_loop_var_value(expr) { path.reverse(); return Ok(Self::get_value_chained(v.clone(), &path[..])); } @@ -490,157 +500,6 @@ impl Interpreter { } } - 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) - } - }, - } - } - - 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 { .. } => (), - } - } - - 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 - } - fn eval_bool_expr( &mut self, op: &BoolOp, @@ -832,7 +691,7 @@ impl Interpreter { // Allow variable overwritten inside a loop let lhs_val = self.lookup_local_var(&name); if !matches!(lhs_val, None | Some(Value::Undefined)) - && !self.loop_var_values.contains_key(rhs) + && !self.has_loop_var_value(rhs) { bail!(rhs .span() @@ -1550,7 +1409,7 @@ impl Interpreter { match loop_expr_value { Value::Array(items) => { for (idx, v) in items.iter().enumerate() { - self.loop_var_values.insert(loop_expr.expr(), v.clone()); + self.set_loop_var_value(&loop_expr.expr(), v.clone()); let exec = if let Some(index) = loop_expr.index() { let mut type_match = BTreeSet::new(); @@ -1579,11 +1438,11 @@ impl Interpreter { } } - self.loop_var_values.remove(&loop_expr.expr()); + self.remove_loop_var_value(&loop_expr.expr()); } Value::Set(items) => { for v in items.iter() { - self.loop_var_values.insert(loop_expr.expr(), v.clone()); + self.set_loop_var_value(&loop_expr.expr(), v.clone()); // For sets, index is also the value. let exec = if let Some(index) = loop_expr.index() { @@ -1605,11 +1464,11 @@ impl Interpreter { } } } - self.loop_var_values.remove(&loop_expr.expr()); + self.remove_loop_var_value(&loop_expr.expr()); } Value::Object(obj) => { for (k, v) in obj.iter() { - self.loop_var_values.insert(loop_expr.expr(), v.clone()); + self.set_loop_var_value(&loop_expr.expr(), v.clone()); // For objects, index is key. let exec = if let Some(index) = loop_expr.index() { let mut type_match = BTreeSet::new(); @@ -1630,7 +1489,7 @@ impl Interpreter { } } } - self.loop_var_values.remove(&loop_expr.expr()); + self.remove_loop_var_value(&loop_expr.expr()); } Value::Undefined => { result = false; @@ -1960,19 +1819,19 @@ impl Interpreter { match self.eval_expr(&loop_expr.value())? { Value::Array(items) => { for v in items.iter() { - self.loop_var_values.insert(loop_expr.expr(), v.clone()); + self.set_loop_var_value(&loop_expr.expr(), v.clone()); result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } Value::Set(items) => { for v in items.iter() { - self.loop_var_values.insert(loop_expr.expr(), v.clone()); + self.set_loop_var_value(&loop_expr.expr(), v.clone()); result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } Value::Object(obj) => { for (_, v) in obj.iter() { - self.loop_var_values.insert(loop_expr.expr(), v.clone()); + self.set_loop_var_value(&loop_expr.expr(), v.clone()); result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } @@ -1984,7 +1843,7 @@ impl Interpreter { )); } } - self.loop_var_values.remove(&loop_expr.expr()); + self.remove_loop_var_value(&loop_expr.expr()); Ok(result) } @@ -2079,18 +1938,50 @@ impl Interpreter { fn eval_query(&mut self, query: &Ref) -> Result { // Execute the query in a new scope self.scopes.push(Scope::new()); - let ordered_stmts: Vec<&LiteralStmt> = - if let Some(schedule) = &self.compiled_policy.schedule { - match schedule.order.get(query) { - Some(ord) => ord.iter().map(|i| &query.stmts[*i as usize]).collect(), - // TODO - _ => bail!(query - .span - .error("statements not scheduled in query {query:?}")), + let order_indices = { + let query_module_index = self.compiled_policy.modules.len() as u32; + if self.current_module_index == query_module_index { + // Use query schedule for the current module + match self + .query_schedule + .as_ref() + .and_then(|s| s.queries.get(query_module_index, query.qidx)) + { + Some(schedule) => Some(&schedule.order), + None => { + if self.query_schedule.is_some() { + bail!(query + .span + .error("statements not scheduled in query {query:?}")); + } + None + } } } else { - query.stmts.iter().collect() - }; + // Use compiled policy schedule for other modules + match self + .compiled_policy + .schedule + .as_ref() + .and_then(|s| s.queries.get(self.current_module_index, query.qidx)) + { + Some(schedule) => Some(&schedule.order), + None => { + if self.compiled_policy.schedule.is_some() { + bail!(query + .span + .error("statements not scheduled in query {query:?}")); + } + None + } + } + } + }; + + let ordered_stmts: Vec<&LiteralStmt> = match order_indices { + Some(order) => order.iter().map(|i| &query.stmts[*i as usize]).collect(), + None => query.stmts.iter().collect(), + }; let r = self.eval_stmts(&ordered_stmts); self.scopes.pop(); @@ -2406,7 +2297,7 @@ impl Interpreter { params: &[ExprRef], ) -> Result { // Return generated values of walk builtin. - if let Some(v) = self.loop_var_values.get(expr) { + if let Some(v) = self.get_loop_var_value(expr) { return Ok(v.clone()); } @@ -3217,11 +3108,21 @@ impl Interpreter { let m = self.module.clone(); if let Some(m) = &module { self.current_module_path = Self::get_path_string(&m.package.refr, Some("data"))?; + self.current_module_index = self.find_module_index(m); } self.module = module; Ok(m.clone()) } + fn find_module_index(&self, module: &Ref) -> u32 { + self.compiled_policy + .modules + .iter() + .position(|m| core::ptr::eq(m.as_ref(), module.as_ref())) + .map(|i| i as u32) + .unwrap_or(0) + } + fn get_rule_refr(rule: &Rule) -> &ExprRef { match rule { Rule::Spec { head, .. } => match &head { @@ -3505,6 +3406,9 @@ impl Interpreter { } pub fn eval_rule(&mut self, module: &Ref, rule: &Ref) -> Result<()> { + // Set current module index + self.current_module_index = self.find_module_index(module); + // Skip reprocessing rule if self.processed.contains(rule) { return Ok(()); @@ -3556,7 +3460,7 @@ impl Interpreter { &mut self, module: &Ref, query: &Ref, - schedule: &Schedule, + query_schedule: Schedule, enable_tracing: bool, ) -> Result { self.traces = match enable_tracing { @@ -3564,12 +3468,8 @@ impl Interpreter { false => None, }; - // Add schedules for queries. - if let Some(ref mut self_schedule) = &mut self.compiled_policy_mut().schedule { - for (k, v) in schedule.order.iter() { - self_schedule.order.insert(k.clone(), v.clone()); - } - } + // Store the query schedule for lookup during evaluation + self.query_schedule = Some(query_schedule); // Push new context. self.contexts.push(Context { @@ -3579,46 +3479,59 @@ impl Interpreter { ..Context::default() }); + self.query_module = Some(module.clone()); let prev_module = self.set_current_module(Some(module.clone()))?; + // For user queries, set the module index to match the schedule + // Query snippets are scheduled as if they're in a module at the end + let prev_module_index = self.current_module_index; + let compiled_modules_len = self.compiled_policy.modules.len() as u32; + self.current_module_index = compiled_modules_len; + + self.ensure_loop_var_values_capacity(); // Eval the query. let query_r = self.eval_query(query); + self.query_module = None; let mut results = match self.contexts.pop() { Some(ctx) => ctx.results, _ => bail!("internal error: no context"), }; - // Restore schedules. - if let Some(ref mut self_schedule) = &mut self.compiled_policy_mut().schedule { - for (k, ord) in schedule.order.iter() { - if k == query { - for idx in 0..results.result.len() { - let e = Expression { - value: Value::Undefined, - text: "".into(), - location: Location { row: 0, col: 0 }, - }; - let mut ordered_expressions = - vec![e; results.result[idx].expressions.len()]; - for (expr_idx, value) in results.result[idx].expressions.iter().enumerate() - { - let orig_idx = ord[expr_idx] as usize; - ordered_expressions[orig_idx] = value.clone(); - } - if !ordered_expressions - .iter() - .any(|v| v.value == Value::Undefined) - { - results.result[idx].expressions = ordered_expressions; - } + // Apply expression ordering from the schedule if needed + let current_module_idx = compiled_modules_len; + let current_query_idx = query.qidx; + if let Some(ref self_schedule) = &self.query_schedule { + if let Some(query_schedule) = self_schedule + .queries + .get(current_module_idx, current_query_idx) + { + for idx in 0..results.result.len() { + let e = Expression { + value: Value::Undefined, + text: "".into(), + location: Location { row: 0, col: 0 }, + }; + let mut ordered_expressions = vec![e; results.result[idx].expressions.len()]; + for (expr_idx, value) in results.result[idx].expressions.iter().enumerate() { + let orig_idx = query_schedule.order[expr_idx] as usize; + ordered_expressions[orig_idx] = value.clone(); + } + if !ordered_expressions + .iter() + .any(|v| v.value == Value::Undefined) + { + results.result[idx].expressions = ordered_expressions; } } - self_schedule.order.remove(k); } } + // Clear the query schedule + self.query_schedule = None; + self.set_current_module(prev_module)?; + self.current_module_index = prev_module_index; if let Some(r) = results.result.last() { if matches!(&r.bindings, Value::Object(obj) if obj.is_empty()) diff --git a/src/interpreter/loops.rs b/src/interpreter/loops.rs new file mode 100644 index 0000000..238831f --- /dev/null +++ b/src/interpreter/loops.rs @@ -0,0 +1,204 @@ +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 abeff61..1429cd1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -30,6 +30,7 @@ mod engine; mod indexchecker; mod interpreter; mod lexer; +mod lookup; mod number; mod parser; mod policy_info; diff --git a/src/lookup.rs b/src/lookup.rs new file mode 100644 index 0000000..caa31e8 --- /dev/null +++ b/src/lookup.rs @@ -0,0 +1,69 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Lookup table for associating data structures with AST nodes. +//! +//! This module provides efficient lookup tables for associating custom data +//! with expressions, queries, and statements using their respective indices. + +use crate::*; + +/// Generic lookup table that stores data indexed by module and node indices. +#[derive(Debug, Clone)] +pub struct Lookup { + /// Array of slots indexed by node index within each module + slots: Vec>>, +} + +impl Default for Lookup { + fn default() -> Self { + Self::new() + } +} + +impl Lookup { + /// Create a new lookup table. + pub fn new() -> Self { + Self { slots: Vec::new() } + } + + /// Ensure the lookup table has capacity for the given module and node index. + pub fn ensure_capacity(&mut self, module_idx: u32, node_idx: u32) { + let module_idx = module_idx as usize; + let node_idx = node_idx as usize; + + // Ensure we have enough modules + if self.slots.len() <= module_idx { + self.slots.resize(module_idx + 1, Vec::new()); + } + + // Ensure the module has enough capacity for this node index + if self.slots[module_idx].len() <= node_idx { + self.slots[module_idx].resize(node_idx + 1, None); + } + } + + /// Set data using direct indices (bounds assumed to be ensured). + pub fn set(&mut self, module_idx: u32, node_idx: u32, value: T) { + *self.get_mut(module_idx, node_idx) = Some(value); + } + + /// Get data using direct indices (bounds assumed to be ensured). + pub fn get(&self, module_idx: u32, node_idx: u32) -> Option<&T> { + self.slots[module_idx as usize][node_idx as usize].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] + } + + /// Clear data at the given indices by setting it to None. + pub fn clear(&mut self, module_idx: u32, node_idx: u32) { + if (module_idx as usize) < self.slots.len() + && (node_idx as usize) < self.slots[module_idx as usize].len() + { + self.slots[module_idx as usize][node_idx as usize] = None; + } + } +} diff --git a/src/parser.rs b/src/parser.rs index c46e3ac..db1ecd7 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -1972,4 +1972,14 @@ impl<'source> Parser<'source> { } Ok(query) } + + pub fn num_expressions(&self) -> u32 { + self.eidx + } + pub fn num_statements(&self) -> u32 { + self.sidx + } + pub fn num_queries(&self) -> u32 { + self.qidx + } } diff --git a/src/scheduler.rs b/src/scheduler.rs index 16c1d1a..3cc2e3c 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -4,6 +4,7 @@ use crate::ast::Expr::{Set, *}; use crate::ast::*; use crate::lexer::*; +use crate::lookup::*; use crate::utils::*; use crate::*; @@ -212,6 +213,13 @@ pub struct Scope { pub locals: BTreeMap, pub unscoped: BTreeSet, pub inputs: BTreeSet, + pub uses_input: bool, +} + +#[derive(Clone, Default, Debug)] +pub struct QuerySchedule { + pub scope: Scope, + pub order: Vec, } pub fn traverse(expr: &Ref, f: &mut dyn FnMut(&Ref) -> Result) -> Result<()> { @@ -313,7 +321,12 @@ fn gather_assigned_vars( ) -> Result<()> { traverse(expr, &mut |e| match e.as_ref() { // Ignore _, input, data. - Var { span: v, .. } if matches!(v.text(), "_" | "input" | "data") => Ok(false), + Var { span: v, .. } if matches!(v.text(), "_" | "input" | "data") => { + if v.text() == "input" { + scope.uses_input = true; + } + Ok(false) + } // Record local var that can shadow input var. Var { span: v, .. } if can_shadow => { @@ -341,10 +354,13 @@ fn gather_assigned_vars( fn gather_input_vars(expr: &Ref, parent_scopes: &[Scope], scope: &mut Scope) -> Result<()> { traverse(expr, &mut |e| match e.as_ref() { - Var { span: v, .. } - if !scope.unscoped.contains(&v.source_str()) && var_exists(v, parent_scopes) => - { - scope.inputs.insert(v.source_str()); + Var { span: v, .. } => { + let name = v.source_str(); + if name.text() == "input" { + scope.uses_input = true; + } else if !scope.unscoped.contains(&name) && var_exists(v, parent_scopes) { + scope.inputs.insert(name); + } Ok(false) } _ => Ok(true), @@ -353,6 +369,10 @@ fn gather_input_vars(expr: &Ref, parent_scopes: &[Scope], scope: &mut Scop fn gather_loop_vars(expr: &Ref, parent_scopes: &[Scope], scope: &mut Scope) -> Result<()> { traverse(expr, &mut |e| match e.as_ref() { + Var { span: v, .. } if v.text() == "input" => { + scope.uses_input = true; + Ok(false) + } RefBrack { index, .. } => { gather_assigned_vars(index, false, parent_scopes, scope)?; Ok(true) @@ -389,19 +409,17 @@ fn gather_vars( pub struct Analyzer { packages: BTreeMap, - scope_table: BTreeMap, Scope>, scopes: Vec, - order: BTreeMap, Vec>, + schedule_table: Lookup, functions: FunctionTable, current_module_path: String, + current_module_index: u32, } #[derive(Debug, Clone)] #[allow(dead_code)] pub struct Schedule { - #[allow(unused)] - pub scopes: BTreeMap, Scope>, - pub order: BTreeMap, Vec>, + pub queries: Lookup, } impl Default for Analyzer { @@ -414,11 +432,11 @@ impl Analyzer { pub fn new() -> Analyzer { Analyzer { packages: BTreeMap::new(), - scope_table: BTreeMap::new(), + schedule_table: Lookup::new(), scopes: vec![], - order: BTreeMap::new(), functions: FunctionTable::new(), current_module_path: String::default(), + current_module_index: 0, } } @@ -426,13 +444,23 @@ impl Analyzer { self.add_rules_and_aliases(modules)?; self.functions = gather_functions(modules)?; - for m in modules { + // Pre-allocate capacity for all modules based on their num_queries + for (module_index, m) in modules.iter().enumerate() { + let module_idx = module_index as u32; + if m.num_queries > 0 { + // Reserve capacity for all queries in this module (0 to num_queries-1) + self.schedule_table + .ensure_capacity(module_idx, m.num_queries - 1); + } + } + + for (module_index, m) in modules.iter().enumerate() { + self.current_module_index = module_index as u32; self.analyze_module(m)?; } Ok(Schedule { - scopes: self.scope_table, - order: self.order, + queries: self.schedule_table, }) } @@ -442,11 +470,27 @@ impl Analyzer { query: &Ref, ) -> Result { self.add_rules_and_aliases(modules)?; + + // Pre-allocate capacity for all modules based on their num_queries + for (module_index, m) in modules.iter().enumerate() { + let module_idx = module_index as u32; + if m.num_queries > 0 { + // Reserve capacity for all queries in this module (0 to num_queries-1) + self.schedule_table + .ensure_capacity(module_idx, m.num_queries - 1); + } + } + + // Query snippets are treated as if they're part of a module appended at the end + let snippet_module_index = modules.len() as u32; + self.schedule_table + .ensure_capacity(snippet_module_index, query.qidx); + + self.current_module_index = snippet_module_index; self.analyze_query(None, None, query, Scope::default())?; Ok(Schedule { - scopes: self.scope_table, - order: self.order, + queries: self.schedule_table, }) } @@ -672,6 +716,11 @@ impl Analyzer { Ok(false) } + Var { span: v, .. } if v.text() == "input" => { + scope.uses_input = true; + Ok(false) + } + RefBrack { refr, index, .. } => { traverse(index, &mut |e| match e.as_ref() { Var { span: v, .. } => { @@ -724,7 +773,9 @@ impl Analyzer { let compr_scope = match compr.as_ref() { Expr::ArrayCompr { query, term, .. } | Expr::SetCompr { query, term, .. } => { self.analyze_query(None, Some(term.clone()), query, Scope::default())?; - self.scope_table.get(query) + self.schedule_table + .get(self.current_module_index, query.qidx) + .map(|qs| &qs.scope) } Expr::ObjectCompr { query, key, value, .. @@ -735,13 +786,20 @@ impl Analyzer { query, Scope::default(), )?; - self.scope_table.get(query) + self.schedule_table + .get(self.current_module_index, query.qidx) + .map(|qs| &qs.scope) } _ => break, }; // Record vars used by the comprehension scope. if let Some(compr_scope) = compr_scope { + // Propagate input usage from nested scope to parent + if compr_scope.uses_input { + scope.uses_input = true; + } + for iv in &compr_scope.inputs { if scope.locals.contains_key(iv) || scope.unscoped.contains(iv) { // Record possible first use of current scope's local var. @@ -958,6 +1016,11 @@ impl Analyzer { non_vars: &mut Vec>, ) -> Result<()> { traverse(expr, &mut |e| match e.as_ref() { + Var { span: v, .. } if v.text() == "input" => { + // Note: We can't modify scope here since it's not mutable, + // but input usage will be tracked elsewhere + Ok(false) + } Var { span: v, .. } if scope.locals.contains_key(&v.source_str()) => { vars.push(v.source_str()); Ok(false) @@ -1209,16 +1272,27 @@ impl Analyzer { } let res = schedule(&mut infos[..], &query.span.source_str().clone_empty()); - match res { - Ok(SortResult::Order(ord)) => { - self.order.insert(query.clone(), ord); - } + let order = match res { + Ok(SortResult::Order(ord)) => ord, Err(err) => { bail!(query.span.error(&err.to_string())) } - _ => (), + _ => Vec::new(), + }; + + let query_schedule = QuerySchedule { + scope: scope.clone(), + order, + }; + self.schedule_table + .set(self.current_module_index, query.qidx, query_schedule); + + // Propagate input usage to parent scopes + if scope.uses_input && !self.scopes.is_empty() { + if let Some(parent_scope) = self.scopes.last_mut() { + parent_scope.uses_input = true; + } } - self.scope_table.insert(query.clone(), scope); Ok(()) } diff --git a/src/tests/scheduler/analyzer/mod.rs b/src/tests/scheduler/analyzer/mod.rs index 7b2360f..1e24e00 100644 --- a/src/tests/scheduler/analyzer/mod.rs +++ b/src/tests/scheduler/analyzer/mod.rs @@ -52,11 +52,27 @@ fn analyze_file(regos: &[String], expected_scopes: &[Scope]) -> Result<()> { let analyzer = Analyzer::new(); let schedule = analyzer.analyze(&modules)?; - let mut scopes: Vec<(Ref, &crate::scheduler::Scope)> = schedule - .scopes - .iter() - .map(|(r, s)| (r.clone(), s)) - .collect(); + + // Collect all queries from modules to create the mapping + let mut all_queries = Vec::new(); + for (module_idx, module) in modules.iter().enumerate() { + for rule in &module.policy { + if let Rule::Spec { bodies, .. } = rule.as_ref() { + for body in bodies { + all_queries.push((module_idx as u32, body.query.qidx, body.query.clone())); + } + } + } + } + + let mut scopes = Vec::new(); + for (module_idx, qidx, query) in all_queries.iter() { + // Find the corresponding query schedule + if let Some(query_schedule) = schedule.queries.get(*module_idx, *qidx) { + scopes.push((query.clone(), &query_schedule.scope)); + } + } + scopes.sort_by(|a, b| a.0.span.line.cmp(&b.0.span.line)); for (idx, (_, scope)) in scopes.iter().enumerate() { if idx > expected_scopes.len() {