Skip to content

Commit 604591a

Browse files
anakrishanakrish
authored
Merge pull request microsoft#527 from anakrish/checked-indexing
Checked indexing
2 parents 52b56f4 + dbfb8e3 commit 604591a

15 files changed

Lines changed: 615 additions & 165 deletions

File tree

‎src/compiler/hoist.rs‎

Lines changed: 96 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ use super::destructuring_planner::{
1313
use crate::ast::{Expr, ExprRef, Literal, LiteralStmt, Module, Query, Ref, Rule, RuleHead};
1414
use crate::compiler::context::{ContextType, ScopeContext};
1515
use crate::lookup::Lookup;
16+
use crate::lookup::LookupResult;
1617
use crate::scheduler::compute_module_globals;
1718
use crate::*;
1819
use anyhow::{anyhow, Result};
@@ -147,43 +148,87 @@ impl HoistedLoopsLookup {
147148
}
148149

149150
/// Store hoisted loops for a statement
150-
pub fn set_statement_loops(&mut self, module_idx: u32, stmt_idx: u32, loops: Vec<HoistedLoop>) {
151-
self.statement_loops.set(module_idx, stmt_idx, loops);
151+
pub fn set_statement_loops(
152+
&mut self,
153+
module_idx: u32,
154+
stmt_idx: u32,
155+
loops: Vec<HoistedLoop>,
156+
) -> Result<()> {
157+
self.statement_loops
158+
.set_checked(module_idx, stmt_idx, loops)
159+
.map_err(|err| anyhow!("statement_loops out of bounds: {err}"))
152160
}
153161

154162
/// Get hoisted loops for a statement
155-
pub fn get_statement_loops(&self, module_idx: u32, stmt_idx: u32) -> Option<&Vec<HoistedLoop>> {
163+
pub fn get_statement_loops(
164+
&self,
165+
module_idx: u32,
166+
stmt_idx: u32,
167+
) -> LookupResult<Option<&Vec<HoistedLoop>>> {
156168
self.statement_loops.get_checked(module_idx, stmt_idx)
157169
}
158170

159171
/// Store hoisted loops for an expression (output expressions)
160-
pub fn set_expr_loops(&mut self, module_idx: u32, expr_idx: u32, loops: Vec<HoistedLoop>) {
161-
self.expr_loops.set(module_idx, expr_idx, loops);
172+
pub fn set_expr_loops(
173+
&mut self,
174+
module_idx: u32,
175+
expr_idx: u32,
176+
loops: Vec<HoistedLoop>,
177+
) -> Result<()> {
178+
self.expr_loops
179+
.set_checked(module_idx, expr_idx, loops)
180+
.map_err(|err| anyhow!("expr_loops out of bounds: {err}"))
162181
}
163182

164183
/// Get hoisted loops for an expression
165-
pub fn get_expr_loops(&self, module_idx: u32, expr_idx: u32) -> Option<&Vec<HoistedLoop>> {
184+
pub fn get_expr_loops(
185+
&self,
186+
module_idx: u32,
187+
expr_idx: u32,
188+
) -> LookupResult<Option<&Vec<HoistedLoop>>> {
166189
self.expr_loops.get_checked(module_idx, expr_idx)
167190
}
168191

169192
/// Store the compilation context for a query
170-
pub fn set_query_context(&mut self, module_idx: u32, query_idx: u32, context: ScopeContext) {
171-
self.query_contexts.set(module_idx, query_idx, context);
193+
pub fn set_query_context(
194+
&mut self,
195+
module_idx: u32,
196+
query_idx: u32,
197+
context: ScopeContext,
198+
) -> Result<()> {
199+
self.query_contexts
200+
.set_checked(module_idx, query_idx, context)
201+
.map_err(|err| anyhow!("query_contexts out of bounds: {err}"))
172202
}
173203

174204
/// Store a binding plan for an expression
175-
pub fn set_expr_binding_plan(&mut self, module_idx: u32, expr_idx: u32, plan: BindingPlan) {
176-
self.expr_binding_plans.set(module_idx, expr_idx, plan);
205+
pub fn set_expr_binding_plan(
206+
&mut self,
207+
module_idx: u32,
208+
expr_idx: u32,
209+
plan: BindingPlan,
210+
) -> Result<()> {
211+
self.expr_binding_plans
212+
.set_checked(module_idx, expr_idx, plan)
213+
.map_err(|err| anyhow!("expr_binding_plans out of bounds: {err}"))
177214
}
178215

179216
/// Get the compilation context for a query
180217
#[allow(dead_code)]
181-
pub fn get_query_context(&self, module_idx: u32, query_idx: u32) -> Option<&ScopeContext> {
218+
pub fn get_query_context(
219+
&self,
220+
module_idx: u32,
221+
query_idx: u32,
222+
) -> LookupResult<Option<&ScopeContext>> {
182223
self.query_contexts.get_checked(module_idx, query_idx)
183224
}
184225

185226
/// Get the binding plan for an expression
186-
pub fn get_expr_binding_plan(&self, module_idx: u32, expr_idx: u32) -> Option<&BindingPlan> {
227+
pub fn get_expr_binding_plan(
228+
&self,
229+
module_idx: u32,
230+
expr_idx: u32,
231+
) -> LookupResult<Option<&BindingPlan>> {
187232
self.expr_binding_plans.get_checked(module_idx, expr_idx)
188233
}
189234

@@ -275,14 +320,18 @@ impl LoopHoister {
275320
Ok(self.lookup)
276321
}
277322

278-
fn create_scope_context(&self, module_idx: u32) -> ScopeContext {
323+
fn create_scope_context(&self, module_idx: u32) -> Result<ScopeContext> {
279324
let mut context = ScopeContext::new();
280325

281-
if let Some(globals) = self.module_globals.get_checked(module_idx, 0) {
326+
if let Some(globals) = self
327+
.module_globals
328+
.get_checked(module_idx, 0)
329+
.map_err(|err| anyhow!("module_globals out of bounds: {err}"))?
330+
{
282331
context.module_globals = Some(globals.clone());
283332
}
284333

285-
context
334+
Ok(context)
286335
}
287336

288337
/// Populate loop hoisting information for all modules, with extra capacity
@@ -309,7 +358,8 @@ impl LoopHoister {
309358
self.lookup.ensure_expr_capacity(last_module_idx + i, 0);
310359
self.module_globals.ensure_capacity(last_module_idx + i, 0);
311360
self.module_globals
312-
.set(last_module_idx + i, 0, crate::Rc::new(BTreeSet::new()));
361+
.set_checked(last_module_idx + i, 0, crate::Rc::new(BTreeSet::new()))
362+
.map_err(|err| anyhow!("module_globals out of bounds: {err}"))?;
313363
}
314364
Ok(self.lookup)
315365
}
@@ -360,10 +410,11 @@ impl LoopHoister {
360410
reserved_globals.insert("data".to_string());
361411
reserved_globals.insert("input".to_string());
362412
self.module_globals
363-
.set(module_idx, 0, crate::Rc::new(reserved_globals));
413+
.set_checked(module_idx, 0, crate::Rc::new(reserved_globals))
414+
.map_err(|err| anyhow!("module_globals out of bounds: {err}"))?;
364415

365416
// Populate the query with default context
366-
let context = self.create_scope_context(module_idx);
417+
let context = self.create_scope_context(module_idx)?;
367418
self.lookup.ensure_query_capacity(module_idx, query.qidx);
368419
self.populate_query(module_idx, query, &context)?;
369420
Ok(())
@@ -374,7 +425,7 @@ impl LoopHoister {
374425
match rule {
375426
Rule::Spec { head, bodies, .. } => {
376427
// Create a context for this rule
377-
let mut context = self.create_scope_context(module_idx);
428+
let mut context = self.create_scope_context(module_idx)?;
378429

379430
// Bind function parameters if this is a function rule
380431
if let RuleHead::Func { args, .. } = head {
@@ -396,7 +447,7 @@ impl LoopHoister {
396447
module_idx,
397448
expr_idx,
398449
binding_plan,
399-
);
450+
)?;
400451
}
401452
Err(err) => return Err(map_binding_error(err)),
402453
}
@@ -454,7 +505,7 @@ impl LoopHoister {
454505
module_idx,
455506
body.query.qidx,
456507
populated_body_context.clone(),
457-
);
508+
)?;
458509

459510
// Process the key expression if present
460511
if let Some(ref key) = key_expr {
@@ -497,7 +548,7 @@ impl LoopHoister {
497548
}
498549
Rule::Default { value, .. } => {
499550
// For default rules, just process the value expression
500-
let context = self.create_scope_context(module_idx);
551+
let context = self.create_scope_context(module_idx)?;
501552
self.populate_output_expr(module_idx, value, &context)?;
502553
}
503554
}
@@ -518,7 +569,11 @@ impl LoopHoister {
518569

519570
// Get the scheduled order if available
520571
let stmt_order: Vec<usize> = if let Some(ref schedule) = self.schedule {
521-
if let Some(query_schedule) = schedule.queries.get(module_idx, query.qidx) {
572+
if let Some(query_schedule) = schedule
573+
.queries
574+
.get_checked(module_idx, query.qidx)
575+
.map_err(|err| anyhow!("schedule out of bounds: {err}"))?
576+
{
522577
query_schedule
523578
.order
524579
.iter()
@@ -566,7 +621,8 @@ impl LoopHoister {
566621
}
567622

568623
self.lookup.ensure_statement_capacity(module_idx, stmt_idx);
569-
self.lookup.set_statement_loops(module_idx, stmt_idx, loops);
624+
self.lookup
625+
.set_statement_loops(module_idx, stmt_idx, loops)?;
570626

571627
Ok(())
572628
}
@@ -596,7 +652,7 @@ impl LoopHoister {
596652
self.lookup.ensure_expr_capacity(module_idx, expr_idx);
597653
Self::bind_vars_from_plan_to_context(&binding_plan, context);
598654
self.lookup
599-
.set_expr_binding_plan(module_idx, expr_idx, binding_plan);
655+
.set_expr_binding_plan(module_idx, expr_idx, binding_plan)?;
600656

601657
if let Some(key_expr) = key {
602658
self.analyze_expr(module_idx, key_expr, context, loops)?;
@@ -615,7 +671,7 @@ impl LoopHoister {
615671
self.populate_query(module_idx, query.as_ref(), &every_context)?;
616672
self.lookup.ensure_query_capacity(module_idx, query.qidx);
617673
self.lookup
618-
.set_query_context(module_idx, query.qidx, populated_context);
674+
.set_query_context(module_idx, query.qidx, populated_context)?;
619675
}
620676
NotExpr { expr, .. } => {
621677
self.analyze_expr(module_idx, expr, context, loops)?;
@@ -663,7 +719,7 @@ impl LoopHoister {
663719
self.populate_query(module_idx, query.as_ref(), &compr_context)?;
664720
self.lookup.ensure_query_capacity(module_idx, query.qidx);
665721
self.lookup
666-
.set_query_context(module_idx, query.qidx, populated_context.clone());
722+
.set_query_context(module_idx, query.qidx, populated_context.clone())?;
667723
self.populate_output_expr_with_context(module_idx, term, &populated_context)?;
668724
}
669725
E::ObjectCompr {
@@ -678,7 +734,7 @@ impl LoopHoister {
678734
self.populate_query(module_idx, query.as_ref(), &compr_context)?;
679735
self.lookup.ensure_query_capacity(module_idx, query.qidx);
680736
self.lookup
681-
.set_query_context(module_idx, query.qidx, populated_context.clone());
737+
.set_query_context(module_idx, query.qidx, populated_context.clone())?;
682738
self.populate_output_expr_with_context(module_idx, key, &populated_context)?;
683739
self.populate_output_expr_with_context(module_idx, value, &populated_context)?;
684740
}
@@ -721,8 +777,11 @@ impl LoopHoister {
721777
// Immediately bind variables from the plan to context
722778
Self::bind_vars_from_plan_to_context(&binding_plan, context);
723779

724-
self.lookup
725-
.set_expr_binding_plan(module_idx, expr_idx, binding_plan);
780+
self.lookup.set_expr_binding_plan(
781+
module_idx,
782+
expr_idx,
783+
binding_plan,
784+
)?;
726785
}
727786
Err(err) => return Err(map_binding_error(err)),
728787
}
@@ -746,8 +805,11 @@ impl LoopHoister {
746805
let expr_idx = index.as_ref().eidx();
747806
self.lookup.ensure_expr_capacity(module_idx, expr_idx);
748807
Self::bind_vars_from_plan_to_context(&binding_plan, context);
749-
self.lookup
750-
.set_expr_binding_plan(module_idx, expr_idx, binding_plan);
808+
self.lookup.set_expr_binding_plan(
809+
module_idx,
810+
expr_idx,
811+
binding_plan,
812+
)?;
751813
}
752814
Err(err) => return Err(map_binding_error(err)),
753815
}
@@ -780,7 +842,7 @@ impl LoopHoister {
780842
self.lookup.ensure_expr_capacity(module_idx, expr_idx);
781843
Self::bind_vars_from_plan_to_context(&binding_plan, context);
782844
self.lookup
783-
.set_expr_binding_plan(module_idx, expr_idx, binding_plan);
845+
.set_expr_binding_plan(module_idx, expr_idx, binding_plan)?;
784846

785847
self.analyze_expr(module_idx, lhs, context, loops)?;
786848
self.analyze_expr(module_idx, rhs, context, loops)?;
@@ -857,7 +919,7 @@ impl LoopHoister {
857919

858920
let expr_idx = expr.as_ref().eidx();
859921
self.lookup.ensure_expr_capacity(module_idx, expr_idx);
860-
self.lookup.set_expr_loops(module_idx, expr_idx, loops);
922+
self.lookup.set_expr_loops(module_idx, expr_idx, loops)?;
861923

862924
Ok(())
863925
}

‎src/engine.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -871,6 +871,8 @@ impl Engine {
871871
debug_assert!(
872872
query_lookup
873873
.get_statement_loops(module_idx, stmt.sidx)
874+
.ok()
875+
.and_then(|entry| entry)
874876
.is_some(),
875877
"missing hoisted loop entry for query statement index {}",
876878
stmt.sidx
@@ -895,6 +897,8 @@ impl Engine {
895897
debug_assert!(
896898
existing_table
897899
.get_statement_loops(module_idx, stmt.sidx)
900+
.ok()
901+
.and_then(|entry| entry)
898902
.is_some(),
899903
"missing hoisted loop entry after merge for module {} stmt {}",
900904
module_idx,

0 commit comments

Comments
 (0)