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

feat: Safeguard lookup use

Browse Source 
Detect invalid indexes and raise internal errors.

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2025-12-19 07:53:02 -06:00
parent 5d0cf95332
commit 6bc1249dc8
9 changed files with 319 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
src/compiler/hoist.rs
View File

@@ -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<HoistedLoop>) {
self.statement_loops.set(module_idx, stmt_idx, loops);
pub fn set_statement_loops(
&mut self,
module_idx: u32,
stmt_idx: u32,
loops: Vec<HoistedLoop>,
) -> 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<HoistedLoop>> {
pub fn get_statement_loops(
&self,
module_idx: u32,
stmt_idx: u32,
) -> LookupResult<Option<&Vec<HoistedLoop>>> {
self.statement_loops.get_checked(module_idx, stmt_idx)
}
/// Store hoisted loops for an expression (output expressions)
pub fn set_expr_loops(&mut self, module_idx: u32, expr_idx: u32, loops: Vec<HoistedLoop>) {
self.expr_loops.set(module_idx, expr_idx, loops);
pub fn set_expr_loops(
&mut self,
module_idx: u32,
expr_idx: u32,
loops: Vec<HoistedLoop>,
) -> 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<HoistedLoop>> {
pub fn get_expr_loops(
&self,
module_idx: u32,
expr_idx: u32,
) -> LookupResult<Option<&Vec<HoistedLoop>>> {
self.expr_loops.get_checked(module_idx, expr_idx)
}
/// Store the compilation context for a query
pub fn set_query_context(&mut self, module_idx: u32, query_idx: u32, context: ScopeContext) {
self.query_contexts.set(module_idx, query_idx, context);
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<Option<&ScopeContext>> {
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<Option<&BindingPlan>> {
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<ScopeContext> {
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<usize> = 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(())
}

4
src/engine.rs
View File

@@ -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,

99
src/interpreter.rs
View File

@@ -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<Option<&Value>> {
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<HoistedLoop> = 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<Value> {
// 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 {

12
src/languages/rego/compiler/core.rs
View File

@@ -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<BindingPlan> {
pub(super) fn get_binding_plan_for_expr(&self, expr: &ExprRef) -> Result<Option<BindingPlan>> {
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<BindingPlan> {
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(),
}

51
src/languages/rego/compiler/loops.rs
View File

@@ -15,31 +15,54 @@ use alloc::vec::Vec;
impl<'a> Compiler<'a> {
pub(super) fn get_statement_loops(&self, stmt: &LiteralStmt) -> Result<Vec<HoistedLoop>> {
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<HoistedLoop> {
pub(super) fn get_expr_loops(&self, expr: &ExprRef) -> Result<Vec<HoistedLoop>> {
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,

15
src/languages/rego/compiler/queries.rs
View File

@@ -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() {

104
src/lookup.rs
View File

@@ -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<T> = core::result::Result<T, LookupIndexError>;
/// Generic lookup table that stores data indexed by module and node indices.
#[derive(Debug, Clone)]
@@ -43,28 +86,19 @@ impl<T: Clone> Lookup<T> {
}
}
/// 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<T> {
&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<Option<&T>> {
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<T: Clone> Lookup<T> {
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))
}
}

15
src/scheduler.rs
View File

@@ -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<Str: Clone + cmp::Ord> {
@@ -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)

8
src/tests/scheduler/analyzer/mod.rs
View File

@@ -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));
}
}