Merge pull request #527 from anakrish/checked-indexing

Checked indexing

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<HoistedLoop>) {
+    pub fn set_statement_loops(
-        self.statement_loops.set(module_idx, stmt_idx, 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>) {
+    pub fn set_expr_loops(
-        self.expr_loops.set(module_idx, expr_idx, 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) {
+    pub fn set_query_context(
-        self.query_contexts.set(module_idx, query_idx, 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) {
+    pub fn set_expr_binding_plan(
-        self.expr_binding_plans.set(module_idx, expr_idx, 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
+                            self.lookup.set_expr_binding_plan(
-                                .set_expr_binding_plan(module_idx, expr_idx, 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
+                            self.lookup.set_expr_binding_plan(
-                                .set_expr_binding_plan(module_idx, expr_idx, 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(())
     }

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,

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<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
+        }) = binding_plan.cloned()
-            .compiled_policy
-            .loop_hoisting_table
-            .get_expr_binding_plan(module_idx, expr_idx)
-            .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;
@@ -1692,7 +1700,7 @@ impl Interpreter {
                 };
 
                 let comps_defined = comps.iter().all(|v| v != &Value::Undefined);
-                let ctx = self.contexts.last_mut().expect("no current context");
+                let ctx = self.get_current_context_mut()?;
 
                 if is_const_rule {
                     ctx.early_return = true;
@@ -1745,7 +1753,7 @@ impl Interpreter {
                     let key = self.eval_expr(&ke)?;
                     let value = self.eval_expr(&oe)?;
 
-                    let ctx = self.contexts.last_mut().unwrap();
+                    let ctx = self.get_current_context_mut()?;
                     if key != Value::Undefined && value != Value::Undefined {
                         let map = ctx.value.as_object_mut()?;
                         match map.get(&key) {
@@ -1774,7 +1782,7 @@ impl Interpreter {
                 }
                 (None, Some(oe)) => {
                     let output = self.eval_expr(&oe)?;
-                    let ctx = self.contexts.last_mut().unwrap();
+                    let ctx = self.get_current_context_mut()?;
                     if output != Value::Undefined {
                         match &mut ctx.value {
                             Value::Array(a) => {
@@ -1798,9 +1806,12 @@ impl Interpreter {
             }
 
             // If a query snippet is being run, gather results.
-            let ctx = self.contexts.last_mut().expect("no current context");
+            let result_opt = {
-            if let Some(result) = &ctx.result {
+                let ctx = self.get_current_context_mut()?;
-                let mut result = result.clone();
+                ctx.result.clone()
+            };
+
+            if let Some(mut result) = result_opt {
                 if let Some(scope) = self.scopes.last() {
                     for (name, value) in scope.iter() {
                         result
@@ -1816,6 +1827,7 @@ impl Interpreter {
                        .all(|v| v.value != Value::Undefined && v.value != Value::Bool(false))
                        && !result.expressions.is_empty()
                 {
+                    let ctx = self.get_current_context_mut()?;
                     ctx.results.result.push(result);
                 }
             }
@@ -1830,19 +1842,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;
                 }
             }
@@ -1864,6 +1876,13 @@ impl Interpreter {
         }
     }
 
+    fn get_current_context_mut(&mut self) -> Result<&mut Context> {
+        match self.contexts.last_mut() {
+            Some(ctx) => Ok(ctx),
+            _ => bail!("internal error: no active context found"),
+        }
+    }
+
     fn get_exprs_from_context(&self) -> Result<ContextExprs> {
         let ctx = self.get_current_context()?;
         Ok((ctx.key_expr.clone(), ctx.output_expr.clone()))
@@ -1881,9 +1900,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 +1917,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 +1948,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
@@ -1960,9 +1982,12 @@ impl Interpreter {
             result = self.eval_output_expr()?;
         } else {
             // If a query snippet is being run, gather results.
-            let ctx = self.contexts.last_mut().expect("no current context");
+            let result_opt = {
-            if let Some(result) = &ctx.result {
+                let ctx = self.get_current_context_mut()?;
-                let mut result = result.clone();
+                ctx.result.clone()
+            };
+
+            if let Some(mut result) = result_opt {
                 if let Some(scope) = self.scopes.last() {
                     for (name, value) in scope.iter() {
                         result
@@ -1979,6 +2004,7 @@ impl Interpreter {
                     .all(|v| v.value != Value::Undefined && v.value != Value::Bool(false))
                     && !result.expressions.is_empty()
                 {
+                    let ctx = self.get_current_context_mut()?;
                     ctx.results.result.push(result);
                 }
             }
@@ -1994,11 +2020,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
+                let schedule = match self.query_schedule.as_ref() {
-                    .query_schedule
+                    Some(s) => s
-                    .as_ref()
+                        .queries
-                    .and_then(|s| s.queries.get(query_module_index, query.qidx))
+                        .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 +2040,15 @@ impl Interpreter {
                 }
             } else {
                 // Use compiled policy schedule for other modules
-                match self
+                let schedule = match self.compiled_policy.schedule.as_ref() {
-                    .compiled_policy
+                    Some(s) => s
-                    .schedule
+                        .queries
-                    .as_ref()
+                        .get_checked(self.current_module_index, query.qidx)
-                    .and_then(|s| s.queries.get(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() {
@@ -2031,7 +2063,30 @@ impl Interpreter {
         };
 
         let ordered_stmts: Vec<&LiteralStmt> = match order_indices {
-            Some(order) => order.iter().map(|i| &query.stmts[*i as usize]).collect(),
+            Some(order) => {
+                let stmts_len = query.stmts.len();
+                if order.len() != stmts_len {
+                    let msg = format!(
+                        "invalid schedule: expected {stmts_len} statement indices, found {}",
+                        order.len()
+                    );
+                    bail!(query.span.error(msg.as_str()));
+                }
+
+                let mut ordered = Vec::with_capacity(stmts_len);
+                for idx in order {
+                    let stmt_idx = *idx as usize;
+                    if stmt_idx >= stmts_len {
+                        let msg = format!(
+                            "invalid schedule index {stmt_idx} for {} statements",
+                            stmts_len
+                        );
+                        bail!(query.span.error(msg.as_str()));
+                    }
+                    ordered.push(&query.stmts[stmt_idx]);
+                }
+                ordered
+            }
             None => query.stmts.iter().collect(),
         };
 
@@ -2349,7 +2404,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 +2552,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 +2698,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
@@ -2824,7 +2881,8 @@ impl Interpreter {
 
             Ok(Self::get_value_chained(self.data.clone(), fields))
         } else if !self.compiled_policy.modules.is_empty() {
-            let path = Parser::get_path_ref_components(&self.module.clone().unwrap().package.refr)?;
+            let module = self.current_module()?;
+            let path = Parser::get_path_ref_components(&module.package.refr)?;
             let mut path: Vec<&str> = path.iter().map(|s| s.text()).collect();
             path.push(name.text());
 
@@ -2933,6 +2991,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(
@@ -3001,7 +3060,8 @@ impl Interpreter {
     }
 
     fn make_rule_context(&self, head: &RuleHead) -> Result<(Context, Vec<Span>)> {
-        let mut path = Parser::get_path_ref_components(&self.module.clone().unwrap().package.refr)?;
+        let module = self.current_module()?;
+        let mut path = Parser::get_path_ref_components(&module.package.refr)?;
         match head {
             RuleHead::Compr { refr, assign, .. } => {
                 let output_expr = assign.as_ref().map(|assign| assign.value.clone());
@@ -3321,8 +3381,8 @@ impl Interpreter {
 
             let scopes = core::mem::take(&mut self.scopes);
 
-            let mut path =
+            let module = self.current_module()?;
-                Parser::get_path_ref_components(&self.module.clone().unwrap().package.refr)?;
+            let mut path = Parser::get_path_ref_components(&module.package.refr)?;
 
             let (refr, index) = match refr.as_ref() {
                 Expr::RefBrack { refr, index, .. } => (refr, Some(index.clone())),
@@ -3638,28 +3698,49 @@ impl Interpreter {
             _ => bail!("internal error: no context"),
         };
 
-        // Apply expression ordering from the schedule if needed
+        // Apply expression ordering from the schedule when it is safe to do so.
+        // If the schedule references statements that did not produce expressions, the lengths
+        // will differ; in that case we keep the collected order to avoid spurious errors.
+        // Example: a schedule for `1 == 2; 1 == 1` may list both statements. The first produces
+        // an expression (false), but evaluation stops and the second never runs, so only one
+        // expression is collected. In that case `order.len()` can be 2 while one expression is
+        // available. In these cases, we avoid reordering - doing so requires maintaining additional
+        // data during evaluation and is not worth the complexity for now.
         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)
+                .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 {
+                    let exprs_len = results.result[idx].expressions.len();
+                    // Skip reordering when the schedule length does not match produced expressions.
+                    if query_schedule.order.len() != exprs_len {
+                        continue;
+                    }
+
+                    let placeholder = Expression {
                         value: Value::Undefined,
                         text: "".into(),
                         location: Location { row: 0, col: 0 },
                     };
-                    let mut ordered_expressions = vec![e; results.result[idx].expressions.len()];
+                    let mut ordered_expressions = vec![placeholder; exprs_len];
+                    let mut invalid = false;
                     for (expr_idx, value) in results.result[idx].expressions.iter().enumerate() {
                         let orig_idx = query_schedule.order[expr_idx] as usize;
+                        if orig_idx >= exprs_len {
+                            invalid = true;
+                            break;
+                        }
                         ordered_expressions[orig_idx] = value.clone();
                     }
-                    if !ordered_expressions
+
-                        .iter()
+                    if !invalid
-                        .any(|v| v.value == Value::Undefined)
+                        && !ordered_expressions
+                            .iter()
+                            .any(|v| v.value == Value::Undefined)
                     {
                         results.result[idx].expressions = ordered_expressions;
                     }

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<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(),
             }

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<Vec<HoistedLoop>> {
-        self.policy
+        let loops = self
+            .policy
             .inner
             .loop_hoisting_table
             .get_statement_loops(self.current_module_index, stmt.sidx)
-            .cloned()
+            .map_err(|err| {
-            .ok_or_else(|| {
                 CompilerError::General {
-                    message: format!(
+                    message: format!("loop hoisting table out of bounds: {err}"),
-                        "missing loop hoisting data for statement at {}:{}",
-                        stmt.span.line, stmt.span.col
-                    ),
                 }
                 .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()
+            .map_err(|err| {
-            .unwrap_or_default()
+                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,

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

src/lexer.rs

@@ -98,7 +98,13 @@ impl SourceStr {
     }
 
     pub fn text(&self) -> &str {
-        &self.source.contents()[self.start as usize..self.end as usize]
+        let start = self.start as usize;
+        let end = self.end as usize;
+        // Use safe slicing to avoid panics on malformed spans
+        self.source
+            .contents()
+            .get(start..end)
+            .unwrap_or("<invalid-span>")
     }
 
     pub fn clone_empty(&self) -> SourceStr {
@@ -241,7 +247,13 @@ pub struct Span {
 
 impl Span {
     pub fn text(&self) -> &str {
-        &self.source.contents()[self.start as usize..self.end as usize]
+        let start = self.start as usize;
+        let end = self.end as usize;
+        // Use safe slicing to avoid panics on malformed spans
+        self.source
+            .contents()
+            .get(start..end)
+            .unwrap_or("<invalid-span>")
     }
 
     pub fn source_str(&self) -> SourceStr {
@@ -445,7 +457,13 @@ impl<'source> Lexer<'source> {
         }
 
         // Ensure that the number is parsable in Rust.
-        match serde_json::from_str::<Value>(&self.source.contents()[start..end]) {
+        let num_slice = self
+            .source
+            .contents()
+            .get(start..end)
+            .ok_or_else(|| self.source.error(self.line, col, "invalid number span"))?;
+
+        match serde_json::from_str::<Value>(num_slice) {
             Ok(_) => (),
             Err(e) => {
                 let serde_msg = &e.to_string();
@@ -506,6 +524,10 @@ impl<'source> Lexer<'source> {
             }
         }
         let end = self.peek().0;
+        if end <= start {
+            // Guard against invalid span that would underflow end - 1
+            return Err(self.source.error(line, col, "invalid raw string span"));
+        }
         Ok(Token(
             TokenKind::RawString,
             Span {
@@ -574,8 +596,19 @@ impl<'source> Lexer<'source> {
         let end = self.peek().0;
         self.col += (end - start) as u32;
 
+        if start == 0 || end <= start {
+            // Reject invalid spans before slicing/serde to avoid panic
+            return Err(self.source.error(line, col, "invalid string span"));
+        }
+
+        let str_slice = self
+            .source
+            .contents()
+            .get(start - 1..end)
+            .ok_or_else(|| self.source.error(line, col, "invalid string span"))?;
+
         // Ensure that the string is parsable in Rust.
-        match serde_json::from_str::<String>(&self.source.contents()[start - 1..end]) {
+        match serde_json::from_str::<String>(str_slice) {
             Ok(_) => (),
             Err(e) => {
                 let serde_msg = &e.to_string();

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<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).
+    /// Set data using direct indices with bounds checking.
-    pub fn set(&mut self, module_idx: u32, node_idx: u32, value: T) {
+    /// Returns Ok(()) if written, Err if either index is out of bounds.
-        *self.get_mut(module_idx, node_idx) = Some(value);
+    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);
-    /// Get data using direct indices (bounds assumed to be ensured).
+        Ok(())
-    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 data using direct indices with bounds checking.
-    /// Returns None if the module or node index is out of range or unset.
+    /// 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) -> Option<&T> {
+    pub fn get_checked(&self, module_idx: u32, node_idx: u32) -> LookupResult<Option<&T>> {
-        self.slots
+        let (m, n) = self.validate_indices(module_idx, node_idx)?;
-            .get(module_idx as usize)
+        Ok(self.slots[m][n].as_ref())
-            .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]
     }
 
     /// 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))
+    }
 }

src/parser.rs

@@ -22,6 +22,11 @@ pub struct Parser<'source> {
     future_keywords: BTreeMap<String, Option<Span>>,
     rego_v1: bool,
 
+    // Tracks current expression/comprehension/query nesting to enforce a recursion limit.
+    expr_depth: usize,
+    max_expr_depth: usize,
+    expr_depth_overflow: bool,
+
     // The index of the last expression that was parsed.
     eidx: u32,
     // The index of the last statement that was parsed.
@@ -31,6 +36,7 @@ pub struct Parser<'source> {
 }
 
 const FUTURE_KEYWORDS: [&str; 4] = ["contains", "every", "if", "in"];
+const DEFAULT_MAX_EXPR_DEPTH: usize = 32;
 
 impl<'source> Parser<'source> {
     pub fn new(source: &'source Source) -> Result<Self> {
@@ -44,6 +50,9 @@ impl<'source> Parser<'source> {
             end: 0,
             future_keywords: BTreeMap::new(),
             rego_v1: false,
+            expr_depth: 0,
+            max_expr_depth: DEFAULT_MAX_EXPR_DEPTH,
+            expr_depth_overflow: false,
             eidx: 0,
             sidx: 0,
             qidx: 0,
@@ -94,6 +103,32 @@ impl<'source> Parser<'source> {
         Ok(())
     }
 
+    fn with_expr_depth<T>(&mut self, f: impl FnOnce(&mut Self) -> Result<T>) -> Result<T> {
+        // Increment expression depth.
+        self.expr_depth = self.expr_depth.saturating_add(1);
+        let current_depth = self.expr_depth;
+
+        // Enforce recursion limit.
+        if self.expr_depth > self.max_expr_depth {
+            self.expr_depth = current_depth.saturating_sub(1);
+            self.expr_depth_overflow = true;
+            bail!(self.tok.1.error(&format!(
+                "expression nesting too deep (>{})",
+                self.max_expr_depth
+            )));
+        }
+
+        let res = f(self);
+
+        // Upon return, ensure that expression depth is still current_depth.
+        if self.expr_depth != current_depth {
+            bail!("internal error: expression depth imbalance");
+        }
+
+        self.expr_depth = current_depth.saturating_sub(1);
+        res
+    }
+
     fn expect(&mut self, text: &str, context: &str) -> Result<()> {
         if self.token_text() == text {
             self.next_token()
@@ -180,19 +215,27 @@ impl<'source> Parser<'source> {
 
     fn handle_import_future_keywords(&mut self, comps: &[Span]) -> Result<bool> {
         if comps.len() >= 2 && comps[0].text() == "future" && comps[1].text() == "keywords" {
-            match comps.len() - 2 {
+            match comps.len().saturating_sub(2) {
-                1 => self.set_future_keyword(comps[2].text(), &Some(comps[2].clone()))?,
+                1 if comps.len() >= 3 => {
+                    self.set_future_keyword(comps[2].text(), &Some(comps[2].clone()))?
+                }
                 0 => {
                     let span = &comps[1];
                     for kw in FUTURE_KEYWORDS.iter() {
                         self.set_future_keyword(kw, &Some(span.clone()))?;
                     }
                 }
-                _ => {
+                _ if comps.len() >= 4 => {
                     let s = &comps[3];
-                    return Err(self
+                    return Err(self.source.error(
-                        .source
+                        s.line,
-                        .error(s.line, s.col - 1, "invalid future keyword"));
+                        s.col.saturating_sub(1),
+                        "invalid future keyword",
+                    ));
+                }
+                _ => {
+                    let s = &comps[1];
+                    return Err(self.source.error(s.line, s.col, "invalid future keyword"));
                 }
             }
             Ok(true)
@@ -384,7 +427,12 @@ impl<'source> Parser<'source> {
                 span.end = self.end;
                 Ok((term, query))
             }
-            Err(_) if self.end == pos => {
+            Err(err) if self.end == pos => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
+
                 // No progress was made in parsing the query.
                 // Restore state and try parsing as set, array or object.
                 *self = state;
@@ -410,7 +458,11 @@ impl<'source> Parser<'source> {
                     eidx: self.next_eidx(),
                 })
             }
-            Err(_) if self.end == pos => {
+            Err(err) if self.end == pos => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
                 // No progress was made in parsing comprehension.
                 // Parse as array.
                 let mut items = vec![];
@@ -453,11 +505,20 @@ impl<'source> Parser<'source> {
                 });
             }
             Err(err) if self.end != pos => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
                 // Some progress was made parsing the set comprehension.
                 // Report errors.
                 return Err(err);
             }
-            _ => (),
+            Err(err) => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
+            }
         }
 
         // It could be a set, object or object comprehension.
@@ -511,11 +572,20 @@ impl<'source> Parser<'source> {
                 });
             }
             Err(err) if self.end != pos => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
                 // Some progress was made parsing the object comprehension.
                 // Report errors.
                 return Err(err);
             }
-            _ => (),
+            Err(err) => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
+            }
         }
 
         // Parse object
@@ -890,11 +960,17 @@ impl<'source> Parser<'source> {
     }
 
     pub fn parse_expr(&mut self) -> Result<Expr> {
-        #[cfg(feature = "rego-extensions")]
+        self.with_expr_depth(|this| {
-        return self.parse_or_expr();
+            #[cfg(feature = "rego-extensions")]
+            {
+                this.parse_or_expr()
+            }
 
-        #[cfg(not(feature = "rego-extensions"))]
+            #[cfg(not(feature = "rego-extensions"))]
-        return self.parse_membership_expr();
+            {
+                this.parse_membership_expr()
+            }
+        })
     }
 
     #[cfg(feature = "rego-extensions")]
@@ -1067,16 +1143,14 @@ impl<'source> Parser<'source> {
             }
 
             span.end = self.end;
-            // Since exprs are discarded, adjust the expression index counter.
+            // Since exprs are discarded, adjust the expression index counter (saturating to avoid underflow).
-            self.eidx -= vars.len() as u32;
+            self.eidx = self.eidx.saturating_sub(vars.len() as u32);
             return Ok(Literal::SomeVars { span, vars });
         }
 
-        let (key, value) = match refs.len() {
+        if refs.len() >= 3 {
-            2 => (Some(refs[0].clone()), refs[1].clone()),
+            // Too many identifiers before `in`.
-            1 => (None, refs[0].clone()),
+            if let Some(span) = vars.get(2).or_else(|| vars.last()) {
-            _ => {
-                let span = &vars[2];
                 return Err(anyhow!(
                     "{}:{}:{} error: encountered `{}` while expecting `in`",
                     span.source.file(),
@@ -1085,6 +1159,21 @@ impl<'source> Parser<'source> {
                     span.text()
                 ));
             }
+            return Err(anyhow!(
+                "invalid some-decl: expected `in` after variable names"
+            ));
+        }
+
+        let (key, value) = match refs.len() {
+            2 => (Some(refs[0].clone()), refs[1].clone()),
+            1 => (None, refs[0].clone()),
+            _ => {
+                // We always parse at least one identifier before `in`; guard defensively.
+                // parse_ident rejects `in` when no vars are present, so this is effectively unreachable.
+                return Err(anyhow!(
+                    "invalid some-decl: expected variable names before `in`"
+                ));
+            }
         };
 
         self.parse_future_keyword("in", false, "while parsing some-decl")?;
@@ -1904,7 +1993,6 @@ impl<'source> Parser<'source> {
     }
 
     fn parse_target_rule(&mut self) -> Result<Option<String>> {
-        // Check if the current token starts a target rule: __target__
         if self.tok.0 == TokenKind::Ident && self.token_text() == "__target__" {
             // Parse __target__
             self.next_token()?;

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<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)

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

tests/lexer/mod.rs

@@ -224,6 +224,34 @@ fn invalid_line() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn invalid_span_text_fallbacks() -> Result<()> {
+    let rego = "abc";
+    let source = Source::from_contents("case.rego".to_string(), rego.to_string())?;
+
+    let ss = SourceStr::new(source.clone(), 100, 200);
+    assert_eq!(
+        ss.text(),
+        "<invalid-span>",
+        "SourceStr should return fallback for out-of-bounds span"
+    );
+
+    let span = Span {
+        source: source.clone(),
+        line: 1,
+        col: 1,
+        start: 5,
+        end: 2,
+    };
+    assert_eq!(
+        span.text(),
+        "<invalid-span>",
+        "Span should return fallback for malformed span"
+    );
+
+    Ok(())
+}
+
 #[test]
 #[cfg(feature = "std")]
 fn file_more_than_64_kb_size() -> Result<()> {

tests/parser/cases/expressions/recursion.yaml

@@ -0,0 +1,17 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: deep-parens-over-limit
+    rego: |
+      package test
+      x = ((((((((((((((((((((((((((((((((1))))))))))))))))))))))))))))))))
+    error: "expression nesting too deep"
+    skip: true
+
+  - note: compr-mixed-deep
+    rego: |
+      package test
+      x = (1 + ( 2 + ( 3 + ( 4 + ((((((( [ [ 1 + (2 + ((((((((((((((((((((1))))))))))))))))))))) | true ] | true ] )))))))))))))
+    error: "expression nesting too deep"
+

tests/parser/cases/import/future.yaml

@@ -163,6 +163,12 @@ cases:
     num_statements: 0
     want_result: {}
 
+  - note: invalid-future-component
+    rego: |
+      package test
+      import future.keywords.foo.bar
+    error: "invalid future keyword"
+
   - note: shadow/1
     rego: |
       package test

tests/parser/cases/some/some.vars.yaml

@@ -78,3 +78,21 @@ cases:
            some a, 5
        }
     error: encountered `5` while expecting identifier
+
+  - note: no-vars-before-in
+    rego: |
+      package test
+      import future.keywords.in
+      x = y {
+          some in xs
+      }
+    error: "unexpected keyword `in`"
+
+  - note: too-many-before-in
+    rego: |
+      package test
+      import future.keywords.in
+      x = y {
+          some a, b, c in xs
+      }
+    error: "encountered `c` while expecting `in`"