#![allow(
clippy::indexing_slicing,
clippy::arithmetic_side_effects,
clippy::as_conversions,
clippy::unused_trait_names,
clippy::pattern_type_mismatch
)]
use super::{CompilationContext, Compiler, CompilerError, ContextType, Register, Result};
use crate::ast::{self, ExprRef, LiteralStmt, Query};
use crate::compiler::destructuring_planner::plans::BindingPlan;
use crate::compiler::hoist::{HoistedLoop, LoopType};
use crate::lexer::Span;
use crate::rvm::instructions::{LoopMode, LoopStartParams};
use crate::rvm::Instruction;
use crate::Value;
use alloc::format;
use alloc::string::ToString;
use alloc::vec::Vec;
impl<'a> Compiler<'a> {
pub(super) fn get_statement_loops(&self, stmt: &LiteralStmt) -> Result<Vec<HoistedLoop>> {
let loops = self
.policy
.inner
.loop_hoisting_table
.get_statement_loops(self.current_module_index, stmt.sidx)
.map_err(|err| {
CompilerError::General {
message: format!("loop hoisting table out of bounds: {err}"),
}
.at(&stmt.span)
})?;
loops.cloned().ok_or_else(|| {
CompilerError::General {
message: format!(
"missing loop hoisting data for statement at {}:{}",
stmt.span.line, stmt.span.col
),
}
.at(&stmt.span)
})
}
pub(super) fn get_expr_loops(&self, expr: &ExprRef) -> Result<Vec<HoistedLoop>> {
let module_idx = self.current_module_index;
let expr_idx = expr.as_ref().eidx();
let loops = self
.policy
.inner
.loop_hoisting_table
.get_expr_loops(module_idx, expr_idx)
.map_err(|err| {
CompilerError::General {
message: format!("loop hoisting table out of bounds: {err}"),
}
.at(expr.span())
})?;
loops.cloned().ok_or_else(|| {
CompilerError::General {
message: format!(
"missing loop hoisting data for expression at {}:{}",
expr.span().line,
expr.span().col
),
}
.at(expr.span())
})
}
pub(super) fn compile_hoisted_loops(
&mut self,
stmts: &[&LiteralStmt],
loops: &[HoistedLoop],
) -> Result<()> {
if loops.is_empty() {
if !stmts.is_empty() {
self.compile_single_statement(stmts[0])?;
return self.hoist_loops_and_compile_statements(&stmts[1..]);
} else {
self.hoist_loops_and_emit_context_yield()?;
}
}
let current_loop = &loops[0];
let remaining_loops = &loops[1..];
match current_loop.loop_type {
LoopType::IndexIteration => {
self.compile_index_iteration_loop(
¤t_loop.loop_expr,
¤t_loop.key,
¤t_loop.value,
¤t_loop.collection,
stmts,
remaining_loops,
)?;
Ok(())
}
LoopType::Walk => Err(CompilerError::General {
message: "walk loops are not yet supported in the RVM compiler".to_string(),
}
.into()),
}
}
pub(super) fn compile_every_quantifier(
&mut self,
key: &Option<Span>,
value: &Span,
domain: &ExprRef,
query: &Query,
span: &Span,
) -> Result<()> {
let collection_reg = self.compile_rego_expr(domain)?;
let key_reg = self.alloc_register();
let value_reg = self.alloc_register();
let result_reg = self.alloc_register();
let value_var_name = value.text().to_string();
let key_var_name = key.as_ref().map(|k| k.text().to_string());
let actual_key_reg = if key_var_name.is_none() || key_var_name.as_deref() == Some("_") {
value_reg
} else {
key_reg
};
let loop_params_index = self.program.add_loop_params(LoopStartParams {
mode: LoopMode::Every,
collection: collection_reg,
key_reg: actual_key_reg,
value_reg,
result_reg,
body_start: 0,
loop_end: 0,
});
self.emit_instruction(
Instruction::LoopStart {
params_index: loop_params_index,
},
span,
);
let body_start = self.program.instructions.len() as u16;
self.push_scope();
let every_context = CompilationContext {
context_type: ContextType::Every,
dest_register: result_reg,
key_expr: None,
value_expr: None,
span: span.clone(),
key_value_loops_hoisted: false,
};
self.push_context(every_context);
self.add_variable(&value_var_name, value_reg);
if let Some(ref key_name) = key_var_name {
self.add_variable(key_name, key_reg);
}
self.compile_query(query)?;
self.pop_context();
self.pop_scope();
self.emit_instruction(
Instruction::LoopNext {
body_start,
loop_end: 0,
},
span,
);
let loop_end = self.program.instructions.len() as u16;
self.program
.update_loop_params(loop_params_index, |params| {
params.body_start = body_start;
params.loop_end = loop_end;
});
let loop_next_idx = self.program.instructions.len() - 1;
if let Instruction::LoopNext {
loop_end: ref mut end,
..
} = &mut self.program.instructions[loop_next_idx]
{
*end = loop_end;
}
Ok(())
}
fn compile_index_iteration_loop(
&mut self,
loop_expr: &Option<ExprRef>,
key_var: &Option<ExprRef>,
_value_var: &ExprRef,
collection: &ExprRef,
remaining_stmts: &[&LiteralStmt],
remaining_loops: &[HoistedLoop],
) -> Result<()> {
let collection_reg = self.compile_rego_expr(collection)?;
let key_reg = self.alloc_register();
let value_reg = self.alloc_register();
let result_reg = self.alloc_register();
if let Some(loop_expr) = loop_expr {
self.loop_expr_register_map
.insert(loop_expr.clone(), value_reg);
}
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 BindingPlan::LoopIndex { .. } = &binding_plan {
key_binding_plan = Some((binding_plan, key_var.span().clone()));
} else {
return Err(CompilerError::UnexpectedBindingPlan {
context: format!("loop index pattern {}", key_var.span().text()),
found: format!("{binding_plan:?}"),
}
.at(key_var.span()));
}
} else {
match key_var.as_ref() {
ast::Expr::Var { value, .. } => {
let var_name = match value {
Value::String(s) => {
if s.as_ref() == "_" {
"".to_string()
} else {
s.to_string()
}
}
_ => value.to_string(),
};
if !var_name.is_empty() && var_name != "_" {
self.store_variable(var_name, key_reg);
}
}
_ => {
return Err(CompilerError::MissingBindingPlan {
context: format!("loop index pattern {}", key_var.span().text()),
}
.at(key_var.span()));
}
}
}
self.loop_expr_register_map.insert(key_var.clone(), key_reg);
}
let loop_params_index = self.program.add_loop_params(LoopStartParams {
mode: LoopMode::ForEach,
collection: collection_reg,
key_reg,
value_reg,
result_reg,
body_start: 0,
loop_end: 0,
});
self.emit_instruction(
Instruction::LoopStart {
params_index: loop_params_index,
},
collection.span(),
);
let body_start = self.program.instructions.len() as u16;
if let Some((binding_plan, plan_span)) = key_binding_plan.as_ref() {
let _ = self
.apply_binding_plan(binding_plan, key_reg, plan_span)
.map_err(|e| CompilerError::from(e).at(plan_span))?;
}
let body_stmts = &remaining_stmts[0..];
self.compile_hoisted_loops(body_stmts, remaining_loops)?;
self.emit_instruction(
Instruction::LoopNext {
body_start,
loop_end: 0,
},
collection.span(),
);
let loop_end = self.program.instructions.len() as u16;
self.program
.update_loop_params(loop_params_index, |params| {
params.body_start = body_start;
params.loop_end = loop_end;
});
let loop_next_idx = self.program.instructions.len() - 1;
if let Instruction::LoopNext {
loop_end: ref mut end,
..
} = &mut self.program.instructions[loop_next_idx]
{
*end = loop_end;
}
Ok(())
}
pub(super) fn compile_some_in_loop_with_remaining_statements(
&mut self,
key: &Option<ExprRef>,
value: &ExprRef,
collection: &ExprRef,
remaining_stmts: &[&LiteralStmt],
) -> Result<Register> {
let loop_body_stmts = &remaining_stmts[1..];
self.compile_some_in_loop_with_body(key, value, collection, loop_body_stmts)
}
fn compile_some_in_loop_with_body(
&mut self,
key: &Option<ExprRef>,
value: &ExprRef,
collection: &ExprRef,
loop_body_stmts: &[&LiteralStmt],
) -> Result<Register> {
let collection_reg = self.compile_rego_expr(collection)?;
let key_reg = self.alloc_register();
let value_reg = self.alloc_register();
let result_reg = self.alloc_register();
let loop_params_index = self.program.add_loop_params(LoopStartParams {
mode: LoopMode::ForEach,
collection: collection_reg,
key_reg,
value_reg,
result_reg,
body_start: 0,
loop_end: 0,
});
self.emit_instruction(
Instruction::LoopStart {
params_index: loop_params_index,
},
collection.span(),
);
let body_start = self.program.instructions.len() as u16;
if let Some(binding_plan) = self.get_binding_plan_for_expr(collection)? {
if let BindingPlan::SomeIn {
key_plan,
value_plan,
..
} = &binding_plan
{
let key_register = key_plan.as_ref().map(|_| key_reg);
self.apply_some_in_binding_plan(
key_plan.as_ref(),
key_register,
value_plan,
value_reg,
collection.span(),
)
.map_err(|e| CompilerError::from(e).at(collection.span()))?;
} else {
return Err(CompilerError::UnexpectedBindingPlan {
context: format!("some-in binding {}", collection.span().text()),
found: format!("{binding_plan:?}"),
}
.at(collection.span()));
}
} else {
if let Some(key_expr) = key {
match key_expr.as_ref() {
ast::Expr::Var {
value: var_name, ..
} => {
let var_name = var_name
.as_string()
.map_err(|err| {
CompilerError::General {
message: format!(
"Failed to read some-in key variable name: {err}"
),
}
.at(key_expr.span())
})?
.to_string();
self.store_variable(var_name, key_reg);
}
_ => {
return Err(CompilerError::MissingBindingPlan {
context: format!("some-in key pattern {}", key_expr.span().text()),
}
.at(key_expr.span()));
}
}
}
match value.as_ref() {
ast::Expr::Var {
value: var_name, ..
} => {
let var_name = var_name
.as_string()
.map_err(|err| {
CompilerError::General {
message: format!(
"Failed to read some-in value variable name: {err}"
),
}
.at(value.span())
})?
.to_string();
self.store_variable(var_name, value_reg);
}
_ => {
return Err(CompilerError::MissingBindingPlan {
context: format!("some-in value pattern {}", value.span().text()),
}
.at(value.span()));
}
}
}
self.hoist_loops_and_compile_statements(loop_body_stmts)?;
self.emit_instruction(
Instruction::LoopNext {
body_start,
loop_end: 0,
},
collection.span(),
);
let loop_end = self.program.instructions.len() as u16;
self.program
.update_loop_params(loop_params_index, |params| {
params.body_start = body_start;
params.loop_end = loop_end;
});
let loop_next_idx = self.program.instructions.len() - 1;
if let Instruction::LoopNext {
loop_end: ref mut end,
..
} = &mut self.program.instructions[loop_next_idx]
{
*end = loop_end;
}
Ok(result_reg)
}
}