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: Implement efficient node lookup table using node indices (#463)

Browse Source 
Major Changes:
- Add generic Lookup<T> structure for efficient O(1) module-level data access
- Combine separate scope and order lookups into unified QuerySchedule structure
- Add query_schedule field to Interpreter for dedicated user query scheduling
- Refactor loop hoising to separate module
- Use efficient lookup for loop vars
- Also added more tests for loops

Key Concept:
- Ensure module context and indexing stay synchronized during function calls

Testing:
- All scheduler and interpreter tests passing

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2025-08-26 15:01:45 -05:00
committed by GitHub
parent c43c94559a
commit 85753aaf37
8 changed files with 599 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
src/engine.rs
View File

@@ -737,28 +737,13 @@ impl Engine {
self.interpreter.clean_internal_evaluation_state();
self.interpreter.create_rule_prefixes()?;
let query_module = {
let source = Source::from_contents(
"<query_module.rego>".to_owned(),
"package __internal_query_module".to_owned(),
)?;
Ref::new(Parser::new(&source)?.parse()?)
};
// Parse the query.
let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
let mut parser = self.make_parser(&query_source)?;
let query_node = parser.parse_user_query()?;
let (query_module, query_node, query_schedule) = self.make_query(query)?;
if query_node.span.text() == "data" {
self.eval_modules(enable_tracing)?;
}
let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;
self.interpreter.eval_user_query(
&query_module,
&query_node,
&query_schedule,
enable_tracing,
)
self.interpreter
.eval_user_query(&query_module, &query_node, query_schedule, enable_tracing)
}
/// Evaluate a Rego query that produces a boolean value.
@@ -843,6 +828,27 @@ impl Engine {
!matches!(self.eval_bool_query(query, enable_tracing), Ok(false))
}
fn make_query(&mut self, query: String) -> Result<(NodeRef<Module>, NodeRef<Query>, Schedule)> {
let mut query_module = {
let source = Source::from_contents(
"<query_module.rego>".to_owned(),
"package __internal_query_module".to_owned(),
)?;
Parser::new(&source)?.parse()?
};
// Parse the query.
let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
let mut parser = self.make_parser(&query_source)?;
let query_node = parser.parse_user_query()?;
query_module.num_expressions = parser.num_expressions();
query_module.num_queries = parser.num_queries();
query_module.num_statements = parser.num_statements();
let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;
Ok((Ref::new(query_module), query_node, query_schedule))
}
#[doc(hidden)]
/// Evaluate the given query and all the rules in the supplied policies.
///
@@ -854,25 +860,9 @@ impl Engine {
) -> Result<QueryResults> {
self.eval_modules(enable_tracing)?;
let query_module = {
let source = Source::from_contents(
"<query_module.rego>".to_owned(),
"package __internal_query_module".to_owned(),
)?;
Ref::new(Parser::new(&source)?.parse()?)
};
// Parse the query.
let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
let mut parser = self.make_parser(&query_source)?;
let query_node = parser.parse_user_query()?;
let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;
self.interpreter.eval_user_query(
&query_module,
&query_node,
&query_schedule,
enable_tracing,
)
let (query_module, query_node, query_schedule) = self.make_query(query)?;
self.interpreter
.eval_user_query(&query_module, &query_node, query_schedule, enable_tracing)
}
#[doc(hidden)]

421
src/interpreter.rs
View File

@@ -7,6 +7,7 @@ use crate::compiled_policy::CompiledPolicyData;
#[cfg(feature = "azure_policy")]
use crate::compiled_policy::TargetInfo;
use crate::lexer::*;
use crate::lookup::Lookup;
use crate::parser::Parser;
use crate::scheduler::*;
use crate::utils::*;
@@ -20,6 +21,11 @@ use anyhow::{anyhow, bail, Result};
use core::ops::Bound::*;
type Scope = BTreeMap<SourceStr, Value>;
type ExprLookup = Lookup<Value>;
mod loops;
use loops::LoopExpr;
#[cfg(feature = "azure_policy")]
pub mod error;
@@ -68,6 +74,9 @@ pub struct Interpreter {
module: Option<Ref<Module>>,
current_module_path: String,
current_module_index: u32,
query_schedule: Option<Schedule>,
query_module: Option<NodeRef<Module>>,
input: Value,
init_data: Value,
@@ -75,7 +84,7 @@ pub struct Interpreter {
with_functions: BTreeMap<String, FunctionModifier>,
scopes: Vec<Scope>,
// TODO: handle recursive calls where same expr could have different values.
loop_var_values: BTreeMap<ExprRef, Value>,
loop_var_values: ExprLookup,
contexts: Vec<Context>,
processed: BTreeSet<Ref<Rule>>,
@@ -118,7 +127,7 @@ impl Clone for Interpreter {
// Hence, they need not be copied.
processed: BTreeSet::default(),
processed_paths: Value::new_object(),
loop_var_values: BTreeMap::default(),
loop_var_values: ExprLookup::new(),
scopes: Vec::default(),
rule_values: BTreeMap::default(),
@@ -126,6 +135,9 @@ impl Clone for Interpreter {
active_rules: Vec::default(),
contexts: Vec::default(),
current_module_path: String::default(),
current_module_index: 0,
query_schedule: None,
query_module: None,
module: None,
no_rules_lookup: false,
}
@@ -167,50 +179,6 @@ impl Default for Context {
}
}
#[derive(Debug)]
enum LoopExpr {
Loop {
span: Span,
expr: Ref<Expr>,
value: Ref<Expr>,
index: Ref<Expr>,
},
Walk {
span: Span,
expr: Ref<Expr>,
},
}
impl LoopExpr {
fn span(&self) -> Span {
match self {
Self::Loop { span, .. } => span.clone(),
Self::Walk { span, .. } => span.clone(),
}
}
fn value(&self) -> Ref<Expr> {
match self {
Self::Loop { value, .. } => value.clone(),
Self::Walk { expr, .. } => expr.clone(),
}
}
fn expr(&self) -> Ref<Expr> {
match self {
Self::Loop { expr, .. } => expr.clone(),
Self::Walk { expr, .. } => expr.clone(),
}
}
fn index(&self) -> Option<Ref<Expr>> {
match self {
Self::Loop { index, .. } => Some(index.clone()),
Self::Walk { .. } => None,
}
}
}
impl Interpreter {
pub fn new() -> Interpreter {
let compiled_policy = compiled_policy::CompiledPolicyData {
@@ -224,6 +192,9 @@ impl Interpreter {
module: None,
current_module_path: String::default(),
current_module_index: 0,
query_schedule: None,
query_module: None,
input: Value::Undefined,
init_data: Value::new_object(),
@@ -231,7 +202,7 @@ impl Interpreter {
with_functions: BTreeMap::default(),
scopes: Vec::default(),
contexts: Vec::default(),
loop_var_values: BTreeMap::default(),
loop_var_values: Lookup::default(),
processed: BTreeSet::default(),
processed_paths: Value::new_object(),
@@ -259,13 +230,16 @@ impl Interpreter {
module: None,
current_module_path: String::default(),
current_module_index: 0,
query_module: None,
query_schedule: None,
input: Value::Undefined,
with_document: Value::new_object(),
with_functions: BTreeMap::default(),
scopes: Vec::default(),
contexts: Vec::default(),
loop_var_values: BTreeMap::default(),
loop_var_values: Lookup::default(),
processed: BTreeSet::default(),
processed_paths: Value::new_object(),
@@ -360,13 +334,49 @@ impl Interpreter {
self.data = self.init_data.clone();
self.processed.clear();
self.processed_paths = Value::new_object();
self.loop_var_values.clear();
self.ensure_loop_var_values_capacity();
self.scopes = vec![Scope::new()];
self.contexts = vec![];
self.rule_values.clear();
self.builtins_cache.clear();
}
// Helper methods for working with ExprLookup
fn set_loop_var_value(&mut self, expr: &ExprRef, value: Value) {
let module_idx = self.current_module_index;
let expr_idx = expr.eidx();
self.loop_var_values.set(module_idx, expr_idx, value);
}
fn get_loop_var_value(&self, expr: &ExprRef) -> Option<&Value> {
let module_idx = self.current_module_index;
let expr_idx = expr.eidx();
self.loop_var_values.get(module_idx, expr_idx)
}
fn remove_loop_var_value(&mut self, expr: &ExprRef) {
let module_idx = self.current_module_index;
let expr_idx = expr.eidx();
self.loop_var_values.clear(module_idx, expr_idx);
}
fn has_loop_var_value(&self, expr: &ExprRef) -> bool {
self.get_loop_var_value(expr).is_some()
}
fn ensure_loop_var_values_capacity(&mut self) {
for (module_idx, module) in self.compiled_policy.modules.iter().enumerate() {
self.loop_var_values
.ensure_capacity(module_idx as u32, module.num_expressions);
}
if let Some(query_module) = &self.query_module {
let query_module_idx = self.compiled_policy.modules.len();
let eidx = query_module.num_expressions;
self.loop_var_values
.ensure_capacity(query_module_idx as u32, eidx);
}
}
fn current_module(&self) -> Result<Ref<Module>> {
self.module
.clone()
@@ -422,7 +432,7 @@ impl Interpreter {
// Collect a chaing of '.field' or '["field"]'
let mut path = vec![];
loop {
if let Some(v) = self.loop_var_values.get(expr) {
if let Some(v) = self.get_loop_var_value(expr) {
path.reverse();
return Ok(Self::get_value_chained(v.clone(), &path[..]));
}
@@ -490,157 +500,6 @@ impl Interpreter {
}
}
fn is_loop_index_var(&self, ident: &SourceStr) -> bool {
// TODO: check for vars that are declared using some-vars
match ident.text() {
"_" => true,
_ => match self.lookup_local_var(ident) {
// Vars declared using `some v` can be loop vars.
// They are initialized to undefined.
Some(Value::Undefined) => true,
// If ident is a local var (in current or parent scopes),
// then it is not a loop var.
Some(_) => false,
None => {
// Check if ident is a rule.
let path = self.current_module_path.clone() + "." + ident.text();
!self.compiled_policy.rules.contains_key(&path)
}
},
}
}
fn hoist_loops_impl(&self, expr: &ExprRef, loops: &mut Vec<LoopExpr>) {
use Expr::*;
match expr.as_ref() {
RefBrack {
refr, index, span, ..
} => {
// First hoist any loops in refr
self.hoist_loops_impl(refr, loops);
// hoist any loops in index expression.
self.hoist_loops_impl(index, loops);
// Then hoist the current bracket operation.
let mut indices = Vec::with_capacity(1);
let _ = traverse(index, &mut |e| match e.as_ref() {
Var { span: ident, .. } if self.is_loop_index_var(&ident.source_str()) => {
indices.push(ident.source_str());
Ok(false)
}
Array { .. } | Object { .. } => Ok(true),
_ => Ok(false),
});
if !indices.is_empty() {
loops.push(LoopExpr::Loop {
span: span.clone(),
expr: expr.clone(),
value: refr.clone(),
index: index.clone(),
})
}
}
// Primitives
String { .. }
| RawString { .. }
| Number { .. }
| Bool { .. }
| Null { .. }
| Var { .. } => (),
// Recurse into expressions in other variants.
Array { items, .. } | Set { items, .. } | Call { params: items, .. } => {
for item in items {
self.hoist_loops_impl(item, loops);
}
// Handle walk builtin which acts as a generator.
// TODO: Handle with modifier on the walk builtin.
if let Expr::Call { fcn, .. } = expr.as_ref() {
if let Ok(fcn_path) = get_path_string(fcn, None) {
if fcn_path == "walk" {
// TODO: Use an enum for LoopExpr to handle walk
loops.push(LoopExpr::Walk {
span: expr.span().clone(),
expr: expr.clone(),
})
}
}
}
}
Object { fields, .. } => {
for (_, key, value) in fields {
self.hoist_loops_impl(key, loops);
self.hoist_loops_impl(value, loops);
}
}
RefDot { refr: expr, .. } | UnaryExpr { expr, .. } => {
self.hoist_loops_impl(expr, loops)
}
BinExpr { lhs, rhs, .. }
| BoolExpr { lhs, rhs, .. }
| ArithExpr { lhs, rhs, .. }
| AssignExpr { lhs, rhs, .. } => {
self.hoist_loops_impl(lhs, loops);
self.hoist_loops_impl(rhs, loops);
}
#[cfg(feature = "rego-extensions")]
OrExpr { lhs, rhs, .. } => {
self.hoist_loops_impl(lhs, loops);
self.hoist_loops_impl(rhs, loops);
}
Membership {
key,
value,
collection,
..
} => {
if let Some(key) = key.as_ref() {
self.hoist_loops_impl(key, loops);
}
self.hoist_loops_impl(value, loops);
self.hoist_loops_impl(collection, loops);
}
// The output expressions of comprehensions must be subject to hoisting
// only after evaluating the body of the comprehensions since the output
// expressions may depend on variables defined within the body.
ArrayCompr { .. } | SetCompr { .. } | ObjectCompr { .. } => (),
}
}
fn hoist_loops(&self, literal: &Literal) -> Vec<LoopExpr> {
let mut loops = vec![];
use Literal::*;
match literal {
SomeVars { .. } => (),
SomeIn {
key,
value,
collection,
..
} => {
if let Some(key) = key {
self.hoist_loops_impl(key, &mut loops);
}
self.hoist_loops_impl(value, &mut loops);
self.hoist_loops_impl(collection, &mut loops);
}
Every {
domain: collection, ..
} => self.hoist_loops_impl(collection, &mut loops),
Expr { expr, .. } | NotExpr { expr, .. } => self.hoist_loops_impl(expr, &mut loops),
}
loops
}
fn eval_bool_expr(
&mut self,
op: &BoolOp,
@@ -832,7 +691,7 @@ impl Interpreter {
// Allow variable overwritten inside a loop
let lhs_val = self.lookup_local_var(&name);
if !matches!(lhs_val, None | Some(Value::Undefined))
&& !self.loop_var_values.contains_key(rhs)
&& !self.has_loop_var_value(rhs)
{
bail!(rhs
.span()
@@ -1550,7 +1409,7 @@ impl Interpreter {
match loop_expr_value {
Value::Array(items) => {
for (idx, v) in items.iter().enumerate() {
self.loop_var_values.insert(loop_expr.expr(), v.clone());
self.set_loop_var_value(&loop_expr.expr(), v.clone());
let exec = if let Some(index) = loop_expr.index() {
let mut type_match = BTreeSet::new();
@@ -1579,11 +1438,11 @@ impl Interpreter {
}
}
self.loop_var_values.remove(&loop_expr.expr());
self.remove_loop_var_value(&loop_expr.expr());
}
Value::Set(items) => {
for v in items.iter() {
self.loop_var_values.insert(loop_expr.expr(), v.clone());
self.set_loop_var_value(&loop_expr.expr(), v.clone());
// For sets, index is also the value.
let exec = if let Some(index) = loop_expr.index() {
@@ -1605,11 +1464,11 @@ impl Interpreter {
}
}
}
self.loop_var_values.remove(&loop_expr.expr());
self.remove_loop_var_value(&loop_expr.expr());
}
Value::Object(obj) => {
for (k, v) in obj.iter() {
self.loop_var_values.insert(loop_expr.expr(), v.clone());
self.set_loop_var_value(&loop_expr.expr(), v.clone());
// For objects, index is key.
let exec = if let Some(index) = loop_expr.index() {
let mut type_match = BTreeSet::new();
@@ -1630,7 +1489,7 @@ impl Interpreter {
}
}
}
self.loop_var_values.remove(&loop_expr.expr());
self.remove_loop_var_value(&loop_expr.expr());
}
Value::Undefined => {
result = false;
@@ -1960,19 +1819,19 @@ impl Interpreter {
match self.eval_expr(&loop_expr.value())? {
Value::Array(items) => {
for v in items.iter() {
self.loop_var_values.insert(loop_expr.expr(), v.clone());
self.set_loop_var_value(&loop_expr.expr(), v.clone());
result = self.eval_output_expr_in_loop(&loops[1..])? || result;
}
}
Value::Set(items) => {
for v in items.iter() {
self.loop_var_values.insert(loop_expr.expr(), v.clone());
self.set_loop_var_value(&loop_expr.expr(), v.clone());
result = self.eval_output_expr_in_loop(&loops[1..])? || result;
}
}
Value::Object(obj) => {
for (_, v) in obj.iter() {
self.loop_var_values.insert(loop_expr.expr(), v.clone());
self.set_loop_var_value(&loop_expr.expr(), v.clone());
result = self.eval_output_expr_in_loop(&loops[1..])? || result;
}
}
@@ -1984,7 +1843,7 @@ impl Interpreter {
));
}
}
self.loop_var_values.remove(&loop_expr.expr());
self.remove_loop_var_value(&loop_expr.expr());
Ok(result)
}
@@ -2079,18 +1938,50 @@ impl Interpreter {
fn eval_query(&mut self, query: &Ref<Query>) -> Result<bool> {
// Execute the query in a new scope
self.scopes.push(Scope::new());
let ordered_stmts: Vec<&LiteralStmt> =
if let Some(schedule) = &self.compiled_policy.schedule {
match schedule.order.get(query) {
Some(ord) => ord.iter().map(|i| &query.stmts[*i as usize]).collect(),
// TODO
_ => bail!(query
.span
.error("statements not scheduled in query {query:?}")),
let order_indices = {
let query_module_index = self.compiled_policy.modules.len() as u32;
if self.current_module_index == query_module_index {
// Use query schedule for the current module
match self
.query_schedule
.as_ref()
.and_then(|s| s.queries.get(query_module_index, query.qidx))
{
Some(schedule) => Some(&schedule.order),
None => {
if self.query_schedule.is_some() {
bail!(query
.span
.error("statements not scheduled in query {query:?}"));
}
None
}
}
} else {
query.stmts.iter().collect()
};
// Use compiled policy schedule for other modules
match self
.compiled_policy
.schedule
.as_ref()
.and_then(|s| s.queries.get(self.current_module_index, query.qidx))
{
Some(schedule) => Some(&schedule.order),
None => {
if self.compiled_policy.schedule.is_some() {
bail!(query
.span
.error("statements not scheduled in query {query:?}"));
}
None
}
}
}
};
let ordered_stmts: Vec<&LiteralStmt> = match order_indices {
Some(order) => order.iter().map(|i| &query.stmts[*i as usize]).collect(),
None => query.stmts.iter().collect(),
};
let r = self.eval_stmts(&ordered_stmts);
self.scopes.pop();
@@ -2406,7 +2297,7 @@ impl Interpreter {
params: &[ExprRef],
) -> Result<Value> {
// Return generated values of walk builtin.
if let Some(v) = self.loop_var_values.get(expr) {
if let Some(v) = self.get_loop_var_value(expr) {
return Ok(v.clone());
}
@@ -3217,11 +3108,21 @@ impl Interpreter {
let m = self.module.clone();
if let Some(m) = &module {
self.current_module_path = Self::get_path_string(&m.package.refr, Some("data"))?;
self.current_module_index = self.find_module_index(m);
}
self.module = module;
Ok(m.clone())
}
fn find_module_index(&self, module: &Ref<Module>) -> u32 {
self.compiled_policy
.modules
.iter()
.position(|m| core::ptr::eq(m.as_ref(), module.as_ref()))
.map(|i| i as u32)
.unwrap_or(0)
}
fn get_rule_refr(rule: &Rule) -> &ExprRef {
match rule {
Rule::Spec { head, .. } => match &head {
@@ -3505,6 +3406,9 @@ impl Interpreter {
}
pub fn eval_rule(&mut self, module: &Ref<Module>, rule: &Ref<Rule>) -> Result<()> {
// Set current module index
self.current_module_index = self.find_module_index(module);
// Skip reprocessing rule
if self.processed.contains(rule) {
return Ok(());
@@ -3556,7 +3460,7 @@ impl Interpreter {
&mut self,
module: &Ref<Module>,
query: &Ref<Query>,
schedule: &Schedule,
query_schedule: Schedule,
enable_tracing: bool,
) -> Result<QueryResults> {
self.traces = match enable_tracing {
@@ -3564,12 +3468,8 @@ impl Interpreter {
false => None,
};
// Add schedules for queries.
if let Some(ref mut self_schedule) = &mut self.compiled_policy_mut().schedule {
for (k, v) in schedule.order.iter() {
self_schedule.order.insert(k.clone(), v.clone());
}
}
// Store the query schedule for lookup during evaluation
self.query_schedule = Some(query_schedule);
// Push new context.
self.contexts.push(Context {
@@ -3579,46 +3479,59 @@ impl Interpreter {
..Context::default()
});
self.query_module = Some(module.clone());
let prev_module = self.set_current_module(Some(module.clone()))?;
// For user queries, set the module index to match the schedule
// Query snippets are scheduled as if they're in a module at the end
let prev_module_index = self.current_module_index;
let compiled_modules_len = self.compiled_policy.modules.len() as u32;
self.current_module_index = compiled_modules_len;
self.ensure_loop_var_values_capacity();
// Eval the query.
let query_r = self.eval_query(query);
self.query_module = None;
let mut results = match self.contexts.pop() {
Some(ctx) => ctx.results,
_ => bail!("internal error: no context"),
};
// Restore schedules.
if let Some(ref mut self_schedule) = &mut self.compiled_policy_mut().schedule {
for (k, ord) in schedule.order.iter() {
if k == query {
for idx in 0..results.result.len() {
let e = Expression {
value: Value::Undefined,
text: "".into(),
location: Location { row: 0, col: 0 },
};
let mut ordered_expressions =
vec![e; results.result[idx].expressions.len()];
for (expr_idx, value) in results.result[idx].expressions.iter().enumerate()
{
let orig_idx = ord[expr_idx] as usize;
ordered_expressions[orig_idx] = value.clone();
}
if !ordered_expressions
.iter()
.any(|v| v.value == Value::Undefined)
{
results.result[idx].expressions = ordered_expressions;
}
// Apply expression ordering from the schedule if needed
let current_module_idx = compiled_modules_len;
let current_query_idx = query.qidx;
if let Some(ref self_schedule) = &self.query_schedule {
if let Some(query_schedule) = self_schedule
.queries
.get(current_module_idx, current_query_idx)
{
for idx in 0..results.result.len() {
let e = Expression {
value: Value::Undefined,
text: "".into(),
location: Location { row: 0, col: 0 },
};
let mut ordered_expressions = vec![e; results.result[idx].expressions.len()];
for (expr_idx, value) in results.result[idx].expressions.iter().enumerate() {
let orig_idx = query_schedule.order[expr_idx] as usize;
ordered_expressions[orig_idx] = value.clone();
}
if !ordered_expressions
.iter()
.any(|v| v.value == Value::Undefined)
{
results.result[idx].expressions = ordered_expressions;
}
}
self_schedule.order.remove(k);
}
}
// Clear the query schedule
self.query_schedule = None;
self.set_current_module(prev_module)?;
self.current_module_index = prev_module_index;
if let Some(r) = results.result.last() {
if matches!(&r.bindings, Value::Object(obj) if obj.is_empty())

204
src/interpreter/loops.rs Normal file
View File

@@ -0,0 +1,204 @@
use crate::ast::*;
use crate::interpreter::Interpreter;
use crate::lexer::*;
use crate::scheduler::traverse;
use crate::utils::get_path_string;
use crate::value::Value;
use crate::*;
#[derive(Debug)]
pub enum LoopExpr {
Loop {
span: Span,
expr: Ref<Expr>,
value: Ref<Expr>,
index: Ref<Expr>,
},
Walk {
span: Span,
expr: Ref<Expr>,
},
}
impl LoopExpr {
pub fn span(&self) -> Span {
match self {
Self::Loop { span, .. } => span.clone(),
Self::Walk { span, .. } => span.clone(),
}
}
pub fn value(&self) -> Ref<Expr> {
match self {
Self::Loop { value, .. } => value.clone(),
Self::Walk { expr, .. } => expr.clone(),
}
}
pub fn expr(&self) -> Ref<Expr> {
match self {
Self::Loop { expr, .. } => expr.clone(),
Self::Walk { expr, .. } => expr.clone(),
}
}
pub fn index(&self) -> Option<Ref<Expr>> {
match self {
Self::Loop { index, .. } => Some(index.clone()),
Self::Walk { .. } => None,
}
}
}
impl Interpreter {
pub(super) fn hoist_loops_impl(&self, expr: &ExprRef, loops: &mut Vec<LoopExpr>) {
use Expr::*;
match expr.as_ref() {
RefBrack {
refr, index, span, ..
} => {
// First hoist any loops in refr
self.hoist_loops_impl(refr, loops);
// hoist any loops in index expression.
self.hoist_loops_impl(index, loops);
// Then hoist the current bracket operation.
let mut indices = Vec::with_capacity(1);
let _ = traverse(index, &mut |e| match e.as_ref() {
Var { span: ident, .. } if self.is_loop_index_var(&ident.source_str()) => {
indices.push(ident.source_str());
Ok(false)
}
Array { .. } | Object { .. } => Ok(true),
_ => Ok(false),
});
if !indices.is_empty() {
loops.push(LoopExpr::Loop {
span: span.clone(),
expr: expr.clone(),
value: refr.clone(),
index: index.clone(),
})
}
}
// Primitives
String { .. }
| RawString { .. }
| Number { .. }
| Bool { .. }
| Null { .. }
| Var { .. } => (),
// Recurse into expressions in other variants.
Array { items, .. } | Set { items, .. } | Call { params: items, .. } => {
for item in items {
self.hoist_loops_impl(item, loops);
}
// Handle walk builtin which acts as a generator.
// TODO: Handle with modifier on the walk builtin.
if let Expr::Call { fcn, .. } = expr.as_ref() {
if let Ok(fcn_path) = get_path_string(fcn, None) {
if fcn_path == "walk" {
// TODO: Use an enum for LoopExpr to handle walk
loops.push(LoopExpr::Walk {
span: expr.span().clone(),
expr: expr.clone(),
})
}
}
}
}
Object { fields, .. } => {
for (_, key, value) in fields {
self.hoist_loops_impl(key, loops);
self.hoist_loops_impl(value, loops);
}
}
RefDot { refr: expr, .. } | UnaryExpr { expr, .. } => {
self.hoist_loops_impl(expr, loops)
}
BinExpr { lhs, rhs, .. }
| BoolExpr { lhs, rhs, .. }
| ArithExpr { lhs, rhs, .. }
| AssignExpr { lhs, rhs, .. } => {
self.hoist_loops_impl(lhs, loops);
self.hoist_loops_impl(rhs, loops);
}
#[cfg(feature = "rego-extensions")]
OrExpr { lhs, rhs, .. } => {
self.hoist_loops_impl(lhs, loops);
self.hoist_loops_impl(rhs, loops);
}
Membership {
key,
value,
collection,
..
} => {
if let Some(key) = key.as_ref() {
self.hoist_loops_impl(key, loops);
}
self.hoist_loops_impl(value, loops);
self.hoist_loops_impl(collection, loops);
}
// The output expressions of comprehensions must be subject to hoisting
// only after evaluating the body of the comprehensions since the output
// expressions may depend on variables defined within the body.
ArrayCompr { .. } | SetCompr { .. } | ObjectCompr { .. } => (),
}
}
pub(super) fn hoist_loops(&self, literal: &Literal) -> Vec<LoopExpr> {
let mut loops = vec![];
use Literal::*;
match literal {
SomeVars { .. } => (),
SomeIn {
key,
value,
collection,
..
} => {
if let Some(key) = key {
self.hoist_loops_impl(key, &mut loops);
}
self.hoist_loops_impl(value, &mut loops);
self.hoist_loops_impl(collection, &mut loops);
}
Every {
domain: collection, ..
} => self.hoist_loops_impl(collection, &mut loops),
Expr { expr, .. } | NotExpr { expr, .. } => self.hoist_loops_impl(expr, &mut loops),
}
loops
}
pub(super) fn is_loop_index_var(&self, ident: &SourceStr) -> bool {
// TODO: check for vars that are declared using some-vars
match ident.text() {
"_" => true,
_ => match self.lookup_local_var(ident) {
// Vars declared using `some v` can be loop vars.
// They are initialized to undefined.
Some(Value::Undefined) => true,
// If ident is a local var (in current or parent scopes),
// then it is not a loop var.
Some(_) => false,
None => {
// Check if ident is a rule.
let path = self.current_module_path.clone() + "." + ident.text();
!self.compiled_policy.rules.contains_key(&path)
}
},
}
}
}

1
src/lib.rs
View File

@@ -30,6 +30,7 @@ mod engine;
mod indexchecker;
mod interpreter;
mod lexer;
mod lookup;
mod number;
mod parser;
mod policy_info;

69
src/lookup.rs Normal file
View File

@@ -0,0 +1,69 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
//! Lookup table for associating data structures with AST nodes.
//!
//! This module provides efficient lookup tables for associating custom data
//! with expressions, queries, and statements using their respective indices.
use crate::*;
/// Generic lookup table that stores data indexed by module and node indices.
#[derive(Debug, Clone)]
pub struct Lookup<T: Clone> {
/// Array of slots indexed by node index within each module
slots: Vec<Vec<Option<T>>>,
}
impl<T: Clone> Default for Lookup<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Clone> Lookup<T> {
/// Create a new lookup table.
pub fn new() -> Self {
Self { slots: Vec::new() }
}
/// Ensure the lookup table has capacity for the given module and node index.
pub fn ensure_capacity(&mut self, module_idx: u32, node_idx: u32) {
let module_idx = module_idx as usize;
let node_idx = node_idx as usize;
// Ensure we have enough modules
if self.slots.len() <= module_idx {
self.slots.resize(module_idx + 1, Vec::new());
}
// Ensure the module has enough capacity for this node index
if self.slots[module_idx].len() <= node_idx {
self.slots[module_idx].resize(node_idx + 1, None);
}
}
/// Set data using direct indices (bounds assumed to be ensured).
pub fn set(&mut self, module_idx: u32, node_idx: u32, value: T) {
*self.get_mut(module_idx, node_idx) = Some(value);
}
/// Get data using direct indices (bounds assumed to be ensured).
pub fn get(&self, module_idx: u32, node_idx: u32) -> Option<&T> {
self.slots[module_idx as usize][node_idx as usize].as_ref()
}
/// Get mutable reference to data using direct indices (bounds assumed to be ensured).
pub fn get_mut(&mut self, module_idx: u32, node_idx: u32) -> &mut Option<T> {
&mut self.slots[module_idx as usize][node_idx as usize]
}
/// Clear data at the given indices by setting it to None.
pub fn clear(&mut self, module_idx: u32, node_idx: u32) {
if (module_idx as usize) < self.slots.len()
&& (node_idx as usize) < self.slots[module_idx as usize].len()
{
self.slots[module_idx as usize][node_idx as usize] = None;
}
}
}

10
src/parser.rs
View File

@@ -1972,4 +1972,14 @@ impl<'source> Parser<'source> {
}
Ok(query)
}
pub fn num_expressions(&self) -> u32 {
self.eidx
}
pub fn num_statements(&self) -> u32 {
self.sidx
}
pub fn num_queries(&self) -> u32 {
self.qidx
}
}

124
src/scheduler.rs
View File

@@ -4,6 +4,7 @@
use crate::ast::Expr::{Set, *};
use crate::ast::*;
use crate::lexer::*;
use crate::lookup::*;
use crate::utils::*;
use crate::*;
@@ -212,6 +213,13 @@ pub struct Scope {
pub locals: BTreeMap<SourceStr, Span>,
pub unscoped: BTreeSet<SourceStr>,
pub inputs: BTreeSet<SourceStr>,
pub uses_input: bool,
}
#[derive(Clone, Default, Debug)]
pub struct QuerySchedule {
pub scope: Scope,
pub order: Vec<u16>,
}
pub fn traverse(expr: &Ref<Expr>, f: &mut dyn FnMut(&Ref<Expr>) -> Result<bool>) -> Result<()> {
@@ -313,7 +321,12 @@ fn gather_assigned_vars(
) -> Result<()> {
traverse(expr, &mut |e| match e.as_ref() {
// Ignore _, input, data.
Var { span: v, .. } if matches!(v.text(), "_" | "input" | "data") => Ok(false),
Var { span: v, .. } if matches!(v.text(), "_" | "input" | "data") => {
if v.text() == "input" {
scope.uses_input = true;
}
Ok(false)
}
// Record local var that can shadow input var.
Var { span: v, .. } if can_shadow => {
@@ -341,10 +354,13 @@ fn gather_assigned_vars(
fn gather_input_vars(expr: &Ref<Expr>, parent_scopes: &[Scope], scope: &mut Scope) -> Result<()> {
traverse(expr, &mut |e| match e.as_ref() {
Var { span: v, .. }
if !scope.unscoped.contains(&v.source_str()) && var_exists(v, parent_scopes) =>
{
scope.inputs.insert(v.source_str());
Var { span: v, .. } => {
let name = v.source_str();
if name.text() == "input" {
scope.uses_input = true;
} else if !scope.unscoped.contains(&name) && var_exists(v, parent_scopes) {
scope.inputs.insert(name);
}
Ok(false)
}
_ => Ok(true),
@@ -353,6 +369,10 @@ fn gather_input_vars(expr: &Ref<Expr>, parent_scopes: &[Scope], scope: &mut Scop
fn gather_loop_vars(expr: &Ref<Expr>, parent_scopes: &[Scope], scope: &mut Scope) -> Result<()> {
traverse(expr, &mut |e| match e.as_ref() {
Var { span: v, .. } if v.text() == "input" => {
scope.uses_input = true;
Ok(false)
}
RefBrack { index, .. } => {
gather_assigned_vars(index, false, parent_scopes, scope)?;
Ok(true)
@@ -389,19 +409,17 @@ fn gather_vars(
pub struct Analyzer {
packages: BTreeMap<String, Scope>,
scope_table: BTreeMap<Ref<Query>, Scope>,
scopes: Vec<Scope>,
order: BTreeMap<Ref<Query>, Vec<u16>>,
schedule_table: Lookup<QuerySchedule>,
functions: FunctionTable,
current_module_path: String,
current_module_index: u32,
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct Schedule {
#[allow(unused)]
pub scopes: BTreeMap<Ref<Query>, Scope>,
pub order: BTreeMap<Ref<Query>, Vec<u16>>,
pub queries: Lookup<QuerySchedule>,
}
impl Default for Analyzer {
@@ -414,11 +432,11 @@ impl Analyzer {
pub fn new() -> Analyzer {
Analyzer {
packages: BTreeMap::new(),
scope_table: BTreeMap::new(),
schedule_table: Lookup::new(),
scopes: vec![],
order: BTreeMap::new(),
functions: FunctionTable::new(),
current_module_path: String::default(),
current_module_index: 0,
}
}
@@ -426,13 +444,23 @@ impl Analyzer {
self.add_rules_and_aliases(modules)?;
self.functions = gather_functions(modules)?;
for m in modules {
// Pre-allocate capacity for all modules based on their num_queries
for (module_index, m) in modules.iter().enumerate() {
let module_idx = module_index as u32;
if m.num_queries > 0 {
// Reserve capacity for all queries in this module (0 to num_queries-1)
self.schedule_table
.ensure_capacity(module_idx, m.num_queries - 1);
}
}
for (module_index, m) in modules.iter().enumerate() {
self.current_module_index = module_index as u32;
self.analyze_module(m)?;
}
Ok(Schedule {
scopes: self.scope_table,
order: self.order,
queries: self.schedule_table,
})
}
@@ -442,11 +470,27 @@ impl Analyzer {
query: &Ref<Query>,
) -> Result<Schedule> {
self.add_rules_and_aliases(modules)?;
// Pre-allocate capacity for all modules based on their num_queries
for (module_index, m) in modules.iter().enumerate() {
let module_idx = module_index as u32;
if m.num_queries > 0 {
// Reserve capacity for all queries in this module (0 to num_queries-1)
self.schedule_table
.ensure_capacity(module_idx, m.num_queries - 1);
}
}
// Query snippets are treated as if they're part of a module appended at the end
let snippet_module_index = modules.len() as u32;
self.schedule_table
.ensure_capacity(snippet_module_index, query.qidx);
self.current_module_index = snippet_module_index;
self.analyze_query(None, None, query, Scope::default())?;
Ok(Schedule {
scopes: self.scope_table,
order: self.order,
queries: self.schedule_table,
})
}
@@ -672,6 +716,11 @@ impl Analyzer {
Ok(false)
}
Var { span: v, .. } if v.text() == "input" => {
scope.uses_input = true;
Ok(false)
}
RefBrack { refr, index, .. } => {
traverse(index, &mut |e| match e.as_ref() {
Var { span: v, .. } => {
@@ -724,7 +773,9 @@ impl Analyzer {
let compr_scope = match compr.as_ref() {
Expr::ArrayCompr { query, term, .. } | Expr::SetCompr { query, term, .. } => {
self.analyze_query(None, Some(term.clone()), query, Scope::default())?;
self.scope_table.get(query)
self.schedule_table
.get(self.current_module_index, query.qidx)
.map(|qs| &qs.scope)
}
Expr::ObjectCompr {
query, key, value, ..
@@ -735,13 +786,20 @@ impl Analyzer {
query,
Scope::default(),
)?;
self.scope_table.get(query)
self.schedule_table
.get(self.current_module_index, query.qidx)
.map(|qs| &qs.scope)
}
_ => break,
};
// Record vars used by the comprehension scope.
if let Some(compr_scope) = compr_scope {
// Propagate input usage from nested scope to parent
if compr_scope.uses_input {
scope.uses_input = true;
}
for iv in &compr_scope.inputs {
if scope.locals.contains_key(iv) || scope.unscoped.contains(iv) {
// Record possible first use of current scope's local var.
@@ -958,6 +1016,11 @@ impl Analyzer {
non_vars: &mut Vec<Ref<Expr>>,
) -> Result<()> {
traverse(expr, &mut |e| match e.as_ref() {
Var { span: v, .. } if v.text() == "input" => {
// Note: We can't modify scope here since it's not mutable,
// but input usage will be tracked elsewhere
Ok(false)
}
Var { span: v, .. } if scope.locals.contains_key(&v.source_str()) => {
vars.push(v.source_str());
Ok(false)
@@ -1209,16 +1272,27 @@ impl Analyzer {
}
let res = schedule(&mut infos[..], &query.span.source_str().clone_empty());
match res {
Ok(SortResult::Order(ord)) => {
self.order.insert(query.clone(), ord);
}
let order = match res {
Ok(SortResult::Order(ord)) => ord,
Err(err) => {
bail!(query.span.error(&err.to_string()))
}
_ => (),
_ => Vec::new(),
};
let query_schedule = QuerySchedule {
scope: scope.clone(),
order,
};
self.schedule_table
.set(self.current_module_index, query.qidx, query_schedule);
// Propagate input usage to parent scopes
if scope.uses_input && !self.scopes.is_empty() {
if let Some(parent_scope) = self.scopes.last_mut() {
parent_scope.uses_input = true;
}
}
self.scope_table.insert(query.clone(), scope);
Ok(())
}

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

@@ -52,11 +52,27 @@ fn analyze_file(regos: &[String], expected_scopes: &[Scope]) -> Result<()> {
let analyzer = Analyzer::new();
let schedule = analyzer.analyze(&modules)?;
let mut scopes: Vec<(Ref<Query>, &crate::scheduler::Scope)> = schedule
.scopes
.iter()
.map(|(r, s)| (r.clone(), s))
.collect();
// Collect all queries from modules to create the mapping
let mut all_queries = Vec::new();
for (module_idx, module) in modules.iter().enumerate() {
for rule in &module.policy {
if let Rule::Spec { bodies, .. } = rule.as_ref() {
for body in bodies {
all_queries.push((module_idx as u32, body.query.qidx, body.query.clone()));
}
}
}
}
let mut scopes = Vec::new();
for (module_idx, qidx, query) in all_queries.iter() {
// Find the corresponding query schedule
if let Some(query_schedule) = schedule.queries.get(*module_idx, *qidx) {
scopes.push((query.clone(), &query_schedule.scope));
}
}
scopes.sort_by(|a, b| a.0.span.line.cmp(&b.0.span.line));
for (idx, (_, scope)) in scopes.iter().enumerate() {
if idx > expected_scopes.len() {