From 6bc1249dc807e08bc1fe171c89f0a8048de2a032 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Fri, 19 Dec 2025 07:53:02 -0600 Subject: [PATCH] feat: Safeguard lookup use Detect invalid indexes and raise internal errors. Signed-off-by: Anand Krishnamoorthi --- src/compiler/hoist.rs | 130 ++++++++++++++++++------- src/engine.rs | 4 + src/interpreter.rs | 99 +++++++++++-------- src/languages/rego/compiler/core.rs | 12 ++- src/languages/rego/compiler/loops.rs | 51 +++++++--- src/languages/rego/compiler/queries.rs | 15 ++- src/lookup.rs | 104 ++++++++++++++++---- src/scheduler.rs | 15 ++- src/tests/scheduler/analyzer/mod.rs | 8 +- 9 files changed, 319 insertions(+), 119 deletions(-) diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index bba31e0..8f6e392 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -13,6 +13,7 @@ use super::destructuring_planner::{ use crate::ast::{Expr, ExprRef, Literal, LiteralStmt, Module, Query, Ref, Rule, RuleHead}; use crate::compiler::context::{ContextType, ScopeContext}; use crate::lookup::Lookup; +use crate::lookup::LookupResult; use crate::scheduler::compute_module_globals; use crate::*; use anyhow::{anyhow, Result}; @@ -147,43 +148,87 @@ impl HoistedLoopsLookup { } /// 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); + pub fn set_statement_loops( + &mut self, + module_idx: u32, + stmt_idx: u32, + loops: Vec, + ) -> Result<()> { + self.statement_loops + .set_checked(module_idx, stmt_idx, loops) + .map_err(|err| anyhow!("statement_loops out of bounds: {err}")) } /// Get hoisted loops for a statement - pub fn get_statement_loops(&self, module_idx: u32, stmt_idx: u32) -> Option<&Vec> { + pub fn get_statement_loops( + &self, + module_idx: u32, + stmt_idx: u32, + ) -> LookupResult>> { 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); + pub fn set_expr_loops( + &mut self, + module_idx: u32, + expr_idx: u32, + loops: Vec, + ) -> Result<()> { + self.expr_loops + .set_checked(module_idx, expr_idx, loops) + .map_err(|err| anyhow!("expr_loops out of bounds: {err}")) } /// Get hoisted loops for an expression - pub fn get_expr_loops(&self, module_idx: u32, expr_idx: u32) -> Option<&Vec> { + pub fn get_expr_loops( + &self, + module_idx: u32, + expr_idx: u32, + ) -> LookupResult>> { 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); + pub fn set_query_context( + &mut self, + module_idx: u32, + query_idx: u32, + context: ScopeContext, + ) -> Result<()> { + self.query_contexts + .set_checked(module_idx, query_idx, context) + .map_err(|err| anyhow!("query_contexts out of bounds: {err}")) } /// Store a binding plan for an expression - pub fn set_expr_binding_plan(&mut self, module_idx: u32, expr_idx: u32, plan: BindingPlan) { - self.expr_binding_plans.set(module_idx, expr_idx, plan); + pub fn set_expr_binding_plan( + &mut self, + module_idx: u32, + expr_idx: u32, + plan: BindingPlan, + ) -> Result<()> { + self.expr_binding_plans + .set_checked(module_idx, expr_idx, plan) + .map_err(|err| anyhow!("expr_binding_plans out of bounds: {err}")) } /// Get the compilation context for a query #[allow(dead_code)] - pub fn get_query_context(&self, module_idx: u32, query_idx: u32) -> Option<&ScopeContext> { + pub fn get_query_context( + &self, + module_idx: u32, + query_idx: u32, + ) -> LookupResult> { self.query_contexts.get_checked(module_idx, query_idx) } /// Get the binding plan for an expression - pub fn get_expr_binding_plan(&self, module_idx: u32, expr_idx: u32) -> Option<&BindingPlan> { + pub fn get_expr_binding_plan( + &self, + module_idx: u32, + expr_idx: u32, + ) -> LookupResult> { self.expr_binding_plans.get_checked(module_idx, expr_idx) } @@ -275,14 +320,18 @@ impl LoopHoister { Ok(self.lookup) } - fn create_scope_context(&self, module_idx: u32) -> ScopeContext { + fn create_scope_context(&self, module_idx: u32) -> Result { let mut context = ScopeContext::new(); - if let Some(globals) = self.module_globals.get_checked(module_idx, 0) { + if let Some(globals) = self + .module_globals + .get_checked(module_idx, 0) + .map_err(|err| anyhow!("module_globals out of bounds: {err}"))? + { context.module_globals = Some(globals.clone()); } - context + Ok(context) } /// Populate loop hoisting information for all modules, with extra capacity @@ -309,7 +358,8 @@ impl LoopHoister { self.lookup.ensure_expr_capacity(last_module_idx + i, 0); self.module_globals.ensure_capacity(last_module_idx + i, 0); self.module_globals - .set(last_module_idx + i, 0, crate::Rc::new(BTreeSet::new())); + .set_checked(last_module_idx + i, 0, crate::Rc::new(BTreeSet::new())) + .map_err(|err| anyhow!("module_globals out of bounds: {err}"))?; } Ok(self.lookup) } @@ -360,10 +410,11 @@ impl LoopHoister { reserved_globals.insert("data".to_string()); reserved_globals.insert("input".to_string()); self.module_globals - .set(module_idx, 0, crate::Rc::new(reserved_globals)); + .set_checked(module_idx, 0, crate::Rc::new(reserved_globals)) + .map_err(|err| anyhow!("module_globals out of bounds: {err}"))?; // Populate the query with default context - let context = self.create_scope_context(module_idx); + let context = self.create_scope_context(module_idx)?; self.lookup.ensure_query_capacity(module_idx, query.qidx); self.populate_query(module_idx, query, &context)?; Ok(()) @@ -374,7 +425,7 @@ impl LoopHoister { match rule { Rule::Spec { head, bodies, .. } => { // Create a context for this rule - let mut context = self.create_scope_context(module_idx); + let mut context = self.create_scope_context(module_idx)?; // Bind function parameters if this is a function rule if let RuleHead::Func { args, .. } = head { @@ -396,7 +447,7 @@ impl LoopHoister { module_idx, expr_idx, binding_plan, - ); + )?; } Err(err) => return Err(map_binding_error(err)), } @@ -454,7 +505,7 @@ impl LoopHoister { module_idx, body.query.qidx, populated_body_context.clone(), - ); + )?; // Process the key expression if present if let Some(ref key) = key_expr { @@ -497,7 +548,7 @@ impl LoopHoister { } Rule::Default { value, .. } => { // For default rules, just process the value expression - let context = self.create_scope_context(module_idx); + let context = self.create_scope_context(module_idx)?; self.populate_output_expr(module_idx, value, &context)?; } } @@ -518,7 +569,11 @@ impl LoopHoister { // 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) { + if let Some(query_schedule) = schedule + .queries + .get_checked(module_idx, query.qidx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))? + { query_schedule .order .iter() @@ -566,7 +621,8 @@ impl LoopHoister { } self.lookup.ensure_statement_capacity(module_idx, stmt_idx); - self.lookup.set_statement_loops(module_idx, stmt_idx, loops); + self.lookup + .set_statement_loops(module_idx, stmt_idx, loops)?; Ok(()) } @@ -596,7 +652,7 @@ impl LoopHoister { self.lookup.ensure_expr_capacity(module_idx, expr_idx); Self::bind_vars_from_plan_to_context(&binding_plan, context); self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan); + .set_expr_binding_plan(module_idx, expr_idx, binding_plan)?; if let Some(key_expr) = key { self.analyze_expr(module_idx, key_expr, context, loops)?; @@ -615,7 +671,7 @@ impl LoopHoister { 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_context); + .set_query_context(module_idx, query.qidx, populated_context)?; } NotExpr { expr, .. } => { self.analyze_expr(module_idx, expr, context, loops)?; @@ -663,7 +719,7 @@ impl LoopHoister { 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_context.clone()); + .set_query_context(module_idx, query.qidx, populated_context.clone())?; self.populate_output_expr_with_context(module_idx, term, &populated_context)?; } E::ObjectCompr { @@ -678,7 +734,7 @@ impl LoopHoister { 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_context.clone()); + .set_query_context(module_idx, query.qidx, populated_context.clone())?; self.populate_output_expr_with_context(module_idx, key, &populated_context)?; self.populate_output_expr_with_context(module_idx, value, &populated_context)?; } @@ -721,8 +777,11 @@ impl LoopHoister { // Immediately bind variables from the plan to context Self::bind_vars_from_plan_to_context(&binding_plan, context); - self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan); + self.lookup.set_expr_binding_plan( + module_idx, + expr_idx, + binding_plan, + )?; } Err(err) => return Err(map_binding_error(err)), } @@ -746,8 +805,11 @@ impl LoopHoister { let expr_idx = index.as_ref().eidx(); self.lookup.ensure_expr_capacity(module_idx, expr_idx); Self::bind_vars_from_plan_to_context(&binding_plan, context); - self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan); + self.lookup.set_expr_binding_plan( + module_idx, + expr_idx, + binding_plan, + )?; } Err(err) => return Err(map_binding_error(err)), } @@ -780,7 +842,7 @@ impl LoopHoister { self.lookup.ensure_expr_capacity(module_idx, expr_idx); Self::bind_vars_from_plan_to_context(&binding_plan, context); self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan); + .set_expr_binding_plan(module_idx, expr_idx, binding_plan)?; self.analyze_expr(module_idx, lhs, context, loops)?; self.analyze_expr(module_idx, rhs, context, loops)?; @@ -857,7 +919,7 @@ impl LoopHoister { 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); + self.lookup.set_expr_loops(module_idx, expr_idx, loops)?; Ok(()) } diff --git a/src/engine.rs b/src/engine.rs index 3f90f04..c11a212 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -871,6 +871,8 @@ impl Engine { debug_assert!( query_lookup .get_statement_loops(module_idx, stmt.sidx) + .ok() + .and_then(|entry| entry) .is_some(), "missing hoisted loop entry for query statement index {}", stmt.sidx @@ -895,6 +897,8 @@ impl Engine { debug_assert!( existing_table .get_statement_loops(module_idx, stmt.sidx) + .ok() + .and_then(|entry| entry) .is_some(), "missing hoisted loop entry after merge for module {} stmt {}", module_idx, diff --git a/src/interpreter.rs b/src/interpreter.rs index d0c3bf3..860489b 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -359,16 +359,20 @@ impl Interpreter { } // Helper methods for working with ExprLookup - fn set_loop_var_value(&mut self, expr: &ExprRef, value: Value) { + fn set_loop_var_value(&mut self, expr: &ExprRef, value: Value) -> Result<()> { let module_idx = self.current_module_index; let expr_idx = expr.eidx(); - self.loop_var_values.set(module_idx, expr_idx, value); + self.loop_var_values + .set_checked(module_idx, expr_idx, value) + .map_err(|err| anyhow!("internal error: loop var indices out of bounds: {err}")) } - fn get_loop_var_value(&self, expr: &ExprRef) -> Option<&Value> { + fn get_loop_var_value(&self, expr: &ExprRef) -> Result> { let module_idx = self.current_module_index; let expr_idx = expr.eidx(); - self.loop_var_values.get(module_idx, expr_idx) + self.loop_var_values + .get_checked(module_idx, expr_idx) + .map_err(|err| anyhow!("internal error: loop var indices out of bounds: {err}")) } fn remove_loop_var_value(&mut self, expr: &ExprRef) { @@ -406,12 +410,13 @@ impl Interpreter { if let Some(last_param) = params.last() { let module_idx = self.current_module_index; let expr_idx = last_param.as_ref().eidx(); - return match self + let binding_plan = self .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) - .cloned() - { + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; + + return match binding_plan.cloned() { Some(BindingPlan::Parameter { destructuring_plan, .. }) => Ok(Some(destructuring_plan)), @@ -483,7 +488,7 @@ impl Interpreter { // Collect a chaing of '.field' or '["field"]' let mut path = vec![]; loop { - if let Some(v) = self.get_loop_var_value(expr) { + if let Some(v) = self.get_loop_var_value(expr)? { path.reverse(); return Ok(Self::get_value_chained(v.clone(), &path[..])); } @@ -876,15 +881,17 @@ impl Interpreter { let module_idx = self.current_module_index; let expr_idx = collection.as_ref().eidx(); + let binding_plan = self + .compiled_policy + .loop_hoisting_table + .get_expr_binding_plan(module_idx, expr_idx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; + let Some(BindingPlan::SomeIn { key_plan, value_plan, .. - }) = self - .compiled_policy - .loop_hoisting_table - .get_expr_binding_plan(module_idx, expr_idx) - .cloned() + }) = binding_plan.cloned() else { bail!("internal error: missing binding plan for some..in expression"); }; @@ -1367,7 +1374,7 @@ impl Interpreter { match loop_value { Value::Array(items) => { for item in items.iter() { - self.set_loop_var_value(loop_target_expr, item.clone()); + self.set_loop_var_value(loop_target_expr, item.clone())?; if self.execute_destructuring_plan(&walk_plan, item)? == Value::from(true) @@ -1417,12 +1424,13 @@ impl Interpreter { let index_plan = if let Some(index) = index_expr { let module_idx = self.current_module_index; let expr_idx = index.as_ref().eidx(); - match self + let plan = self .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) - .cloned() - { + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; + + match plan.cloned() { Some(BindingPlan::LoopIndex { destructuring_plan, .. }) => destructuring_plan, @@ -1445,7 +1453,7 @@ impl Interpreter { match loop_value { Value::Array(items) => { for (idx, v) in items.iter().enumerate() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; if self.execute_destructuring_plan(&index_plan, &Value::from(idx))? == Value::from(true) @@ -1466,7 +1474,7 @@ impl Interpreter { } Value::Set(items) => { for v in items.iter() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; // For sets, index is also the value. if self.execute_destructuring_plan(&index_plan, v)? == Value::from(true) { @@ -1485,7 +1493,7 @@ impl Interpreter { } Value::Object(obj) => { for (k, v) in obj.iter() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; // For objects, index is key. if self.execute_destructuring_plan(&index_plan, k)? == Value::from(true) { result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result; @@ -1830,19 +1838,19 @@ impl Interpreter { match self.eval_expr(Self::loop_collection_expr(loop_info))? { Value::Array(items) => { for v in items.iter() { - self.set_loop_var_value(loop_target_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_target_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_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } @@ -1881,9 +1889,10 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_loops(self.current_module_index, ke.as_ref().eidx()) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? { Some(hoisted_loops) => { - loops.extend(hoisted_loops.iter().cloned()); + loops.extend(hoisted_loops.clone()); } None => { bail!(ke.span().error("Loop hoisting information not found for key expression. This is likely a bug in the compilation phase.")); @@ -1897,9 +1906,10 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_loops(self.current_module_index, oe.as_ref().eidx()) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? { Some(hoisted_loops) => { - loops.extend(hoisted_loops.iter().cloned()); + loops.extend(hoisted_loops.clone()); } None => { bail!(oe.span().error("Loop hoisting information not found for output expression. This is likely a bug in the compilation phase.")); @@ -1927,10 +1937,11 @@ impl Interpreter { } // Get pre-computed hoisted loops from compilation phase - let loop_exprs = match self + let loop_exprs: Vec = match self .compiled_policy .loop_hoisting_table .get_statement_loops(self.current_module_index, stmt.sidx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? { Some(hoisted_loops) => { // Use pre-computed loops from compilation phase @@ -1994,11 +2005,14 @@ impl Interpreter { 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)) - { + let schedule = match self.query_schedule.as_ref() { + Some(s) => s + .queries + .get_checked(query_module_index, query.qidx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))?, + None => None, + }; + match schedule { Some(schedule) => Some(&schedule.order), None => { if self.query_schedule.is_some() { @@ -2011,12 +2025,15 @@ impl Interpreter { } } else { // 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)) - { + let schedule = match self.compiled_policy.schedule.as_ref() { + Some(s) => s + .queries + .get_checked(self.current_module_index, query.qidx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))?, + None => None, + }; + + match schedule { Some(schedule) => Some(&schedule.order), None => { if self.compiled_policy.schedule.is_some() { @@ -2349,7 +2366,7 @@ impl Interpreter { params: &[ExprRef], ) -> Result { // Return generated values of walk builtin. - if let Some(v) = self.get_loop_var_value(expr) { + if let Some(v) = self.get_loop_var_value(expr)? { return Ok(v.clone()); } @@ -2497,6 +2514,7 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? .cloned() { // Execute the destructuring plan with the parameter value @@ -2642,6 +2660,7 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? .cloned() { // Execute the destructuring plan with the return value @@ -2933,6 +2952,7 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? .cloned() .ok_or_else(|| { expr.span().error( @@ -3644,7 +3664,8 @@ impl Interpreter { if let Some(ref self_schedule) = &self.query_schedule { if let Some(query_schedule) = self_schedule .queries - .get(current_module_idx, current_query_idx) + .get_checked(current_module_idx, current_query_idx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))? { for idx in 0..results.result.len() { let e = Expression { diff --git a/src/languages/rego/compiler/core.rs b/src/languages/rego/compiler/core.rs index 5d0d72a..4718d4f 100644 --- a/src/languages/rego/compiler/core.rs +++ b/src/languages/rego/compiler/core.rs @@ -186,14 +186,20 @@ impl<'a> Compiler<'a> { self.lookup_local_var(var_name) } - pub(super) fn get_binding_plan_for_expr(&self, expr: &ExprRef) -> Option { + pub(super) fn get_binding_plan_for_expr(&self, expr: &ExprRef) -> Result> { let module_idx = self.current_module_index; let expr_idx = expr.as_ref().eidx(); self.policy .inner .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) - .cloned() + .map_err(|err| { + CompilerError::General { + message: format!("loop hoisting table out of bounds: {err}"), + } + .at(expr.span()) + }) + .map(|plan: Option<&BindingPlan>| plan.cloned()) } pub(super) fn expect_binding_plan_for_expr( @@ -201,7 +207,7 @@ impl<'a> Compiler<'a> { expr: &ExprRef, context: &str, ) -> Result { - self.get_binding_plan_for_expr(expr).ok_or_else(|| { + self.get_binding_plan_for_expr(expr)?.ok_or_else(|| { CompilerError::MissingBindingPlan { context: context.to_string(), } diff --git a/src/languages/rego/compiler/loops.rs b/src/languages/rego/compiler/loops.rs index f02300c..4809627 100644 --- a/src/languages/rego/compiler/loops.rs +++ b/src/languages/rego/compiler/loops.rs @@ -15,31 +15,54 @@ use alloc::vec::Vec; impl<'a> Compiler<'a> { pub(super) fn get_statement_loops(&self, stmt: &LiteralStmt) -> Result> { - self.policy + let loops = self + .policy .inner .loop_hoisting_table .get_statement_loops(self.current_module_index, stmt.sidx) - .cloned() - .ok_or_else(|| { + .map_err(|err| { CompilerError::General { - message: format!( - "missing loop hoisting data for statement at {}:{}", - stmt.span.line, stmt.span.col - ), + message: format!("loop hoisting table out of bounds: {err}"), } .at(&stmt.span) - }) + })?; + + loops.cloned().ok_or_else(|| { + CompilerError::General { + message: format!( + "missing loop hoisting data for statement at {}:{}", + stmt.span.line, stmt.span.col + ), + } + .at(&stmt.span) + }) } - pub(super) fn get_expr_loops(&self, expr: &ExprRef) -> Vec { + pub(super) fn get_expr_loops(&self, expr: &ExprRef) -> Result> { let module_idx = self.current_module_index; let expr_idx = expr.as_ref().eidx(); - self.policy + let loops = self + .policy .inner .loop_hoisting_table .get_expr_loops(module_idx, expr_idx) - .cloned() - .unwrap_or_default() + .map_err(|err| { + CompilerError::General { + message: format!("loop hoisting table out of bounds: {err}"), + } + .at(expr.span()) + })?; + + loops.cloned().ok_or_else(|| { + CompilerError::General { + message: format!( + "missing loop hoisting data for expression at {}:{}", + expr.span().line, + expr.span().col + ), + } + .at(expr.span()) + }) } pub(super) fn compile_hoisted_loops( @@ -192,7 +215,7 @@ impl<'a> Compiler<'a> { let mut key_binding_plan: Option<(BindingPlan, Span)> = None; if let Some(key_var) = key_var { - if let Some(binding_plan) = self.get_binding_plan_for_expr(key_var) { + if let Some(binding_plan) = self.get_binding_plan_for_expr(key_var)? { if let BindingPlan::LoopIndex { .. } = &binding_plan { key_binding_plan = Some((binding_plan, key_var.span().clone())); } else { @@ -327,7 +350,7 @@ impl<'a> Compiler<'a> { let body_start = self.program.instructions.len() as u16; - if let Some(binding_plan) = self.get_binding_plan_for_expr(collection) { + if let Some(binding_plan) = self.get_binding_plan_for_expr(collection)? { if let BindingPlan::SomeIn { key_plan, value_plan, diff --git a/src/languages/rego/compiler/queries.rs b/src/languages/rego/compiler/queries.rs index a6990a5..c8c3cac 100644 --- a/src/languages/rego/compiler/queries.rs +++ b/src/languages/rego/compiler/queries.rs @@ -5,6 +5,7 @@ use super::{Compiler, CompilerError, ComprehensionType, ContextType, Result}; use crate::ast::{self, LiteralStmt, Query}; use crate::rvm::program::RuleType; use crate::rvm::Instruction; +use alloc::format; use alloc::vec::Vec; impl<'a> Compiler<'a> { @@ -13,7 +14,15 @@ impl<'a> Compiler<'a> { let result = { let schedule = match &self.policy.inner.schedule { - Some(s) => s.queries.get(self.current_module_index, query.qidx), + Some(s) => s + .queries + .get_checked(self.current_module_index, query.qidx) + .map_err(|err| { + CompilerError::General { + message: format!("schedule out of bounds: {err}"), + } + .at(&query.span) + })?, None => None, }; @@ -94,11 +103,11 @@ impl<'a> Compiler<'a> { let mut key_value_loops = Vec::new(); if let Some(expr) = key_expr.as_ref() { - key_value_loops.extend(self.get_expr_loops(expr)); + key_value_loops.extend(self.get_expr_loops(expr)?); } if let Some(expr) = value_expr.as_ref() { - key_value_loops.extend(self.get_expr_loops(expr)); + key_value_loops.extend(self.get_expr_loops(expr)?); } if !key_value_loops.is_empty() { diff --git a/src/lookup.rs b/src/lookup.rs index 7a8a234..0cdafef 100644 --- a/src/lookup.rs +++ b/src/lookup.rs @@ -7,6 +7,49 @@ //! with expressions, queries, and statements using their respective indices. use crate::*; +use core::fmt; + +/// Error indicating that lookup indices are out of bounds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LookupIndexError { + /// The requested module index exceeds the available modules. + ModuleOutOfBounds { module_idx: u32, modules: usize }, + /// The requested node index exceeds the available nodes for the module. + NodeOutOfBounds { + module_idx: u32, + node_idx: u32, + nodes: usize, + }, +} + +impl fmt::Display for LookupIndexError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + LookupIndexError::ModuleOutOfBounds { + module_idx, + modules, + } => { + write!( + f, + "module_idx {module_idx} out of bounds (modules={modules})" + ) + } + LookupIndexError::NodeOutOfBounds { + module_idx, + node_idx, + nodes, + } => write!( + f, + "node_idx {node_idx} out of bounds for module {module_idx} (nodes={nodes})" + ), + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for LookupIndexError {} + +pub type LookupResult = core::result::Result; /// Generic lookup table that stores data indexed by module and node indices. #[derive(Debug, Clone)] @@ -43,28 +86,19 @@ impl Lookup { } } - /// 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() + /// Set data using direct indices with bounds checking. + /// Returns Ok(()) if written, Err if either index is out of bounds. + pub fn set_checked(&mut self, module_idx: u32, node_idx: u32, value: T) -> LookupResult<()> { + let (m, n) = self.validate_indices(module_idx, node_idx)?; + self.slots[m][n] = Some(value); + Ok(()) } /// 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] + /// Returns Ok(None) if the entry is unset, Err if indices are out of range. + pub fn get_checked(&self, module_idx: u32, node_idx: u32) -> LookupResult> { + let (m, n) = self.validate_indices(module_idx, node_idx)?; + Ok(self.slots[m][n].as_ref()) } /// Clear data at the given indices by setting it to None. @@ -95,4 +129,36 @@ impl Lookup { None } } + + /// Validate indices and return them as usize on success. + fn validate_indices(&self, module_idx: u32, node_idx: u32) -> LookupResult<(usize, usize)> { + let m = module_idx as usize; + if m >= self.slots.len() { + debug_assert!( + m < self.slots.len(), + "module_idx {m} out of bounds (modules={})", + self.slots.len() + ); + return Err(LookupIndexError::ModuleOutOfBounds { + module_idx, + modules: self.slots.len(), + }); + } + + let n = node_idx as usize; + if n >= self.slots[m].len() { + debug_assert!( + n < self.slots[m].len(), + "node_idx {n} out of bounds for module {m} (nodes={})", + self.slots[m].len() + ); + return Err(LookupIndexError::NodeOutOfBounds { + module_idx, + node_idx, + nodes: self.slots[m].len(), + }); + } + + Ok((m, n)) + } } diff --git a/src/scheduler.rs b/src/scheduler.rs index 888228c..c4722df 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -17,7 +17,7 @@ use alloc::string::String; use core::cmp; use core::fmt; -use anyhow::{bail, Result}; +use anyhow::{anyhow, bail, Result}; #[derive(Debug)] pub struct Definition { @@ -584,7 +584,8 @@ impl Analyzer { Expr::ArrayCompr { query, term, .. } | Expr::SetCompr { query, term, .. } => { self.analyze_query(None, Some(term.clone()), query, Scope::default())?; self.schedule_table - .get(self.current_module_index, query.qidx) + .get_checked(self.current_module_index, query.qidx) + .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))? .map(|qs| &qs.scope) } Expr::ObjectCompr { @@ -597,7 +598,8 @@ impl Analyzer { Scope::default(), )?; self.schedule_table - .get(self.current_module_index, query.qidx) + .get_checked(self.current_module_index, query.qidx) + .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))? .map(|qs| &qs.scope) } _ => break, @@ -1099,7 +1101,8 @@ impl Analyzer { order, }; self.schedule_table - .set(self.current_module_index, query.qidx, query_schedule); + .set_checked(self.current_module_index, query.qidx, query_schedule) + .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))?; // Propagate input usage to parent scopes if scope.uses_input && !self.scopes.is_empty() { @@ -1174,7 +1177,9 @@ pub fn compute_module_globals( } result.ensure_capacity(module_idx as u32, 0); - result.set(module_idx as u32, 0, module_globals); + result + .set_checked(module_idx as u32, 0, module_globals) + .map_err(|err| anyhow!("module globals out of bounds: {err}"))?; } Ok(result) diff --git a/src/tests/scheduler/analyzer/mod.rs b/src/tests/scheduler/analyzer/mod.rs index 1e24e00..90aa009 100644 --- a/src/tests/scheduler/analyzer/mod.rs +++ b/src/tests/scheduler/analyzer/mod.rs @@ -3,7 +3,7 @@ use crate::*; use crate::{ast::*, lexer::*, parser::*, scheduler::*}; -use anyhow::{bail, Result}; +use anyhow::{anyhow, bail, Result}; use serde::{Deserialize, Serialize}; use test_generator::test_resources; @@ -68,7 +68,11 @@ fn analyze_file(regos: &[String], expected_scopes: &[Scope]) -> Result<()> { 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) { + if let Some(query_schedule) = schedule + .queries + .get_checked(*module_idx, *qidx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))? + { scopes.push((query.clone(), &query_schedule.scope)); } }