diff --git a/examples/regorus.rs b/examples/regorus.rs index c335480..2f5ec5a 100644 --- a/examples/regorus.rs +++ b/examples/regorus.rs @@ -69,16 +69,14 @@ fn rego_eval( None }; + let modules_ref: Vec<®orus::Module> = modules.iter().collect(); + // Analyze the modules and determine how statements must be schedules. let analyzer = regorus::Analyzer::new(); - let schedule = analyzer.analyze(&modules)?; + let schedule = analyzer.analyze(&modules_ref)?; // Create interpreter object. - let modules_ref: Vec<®orus::Module> = modules.iter().collect(); - let mut interpreter = regorus::Interpreter::new(modules_ref)?; - - // Prepare for evalution. - interpreter.prepare_for_eval(Some(schedule.clone()), &Some(data.clone()))?; + let mut interpreter = regorus::Interpreter::new(&modules_ref)?; // Evaluate all the modules. interpreter.eval(&Some(data), &input, false, Some(schedule))?; @@ -98,7 +96,8 @@ fn rego_eval( }; let mut parser = regorus::Parser::new(&query_source)?; let query_node = parser.parse_query(query_span, "")?; - let query_schedule = regorus::Analyzer::new().analyze_query_snippet(&modules, &query_node)?; + let query_schedule = + regorus::Analyzer::new().analyze_query_snippet(&modules_ref, &query_node)?; let results = interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?; println!("{}", serde_json::to_string_pretty(&results)?); diff --git a/src/interpreter.rs b/src/interpreter.rs index a497e90..e04d8d4 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -31,7 +31,7 @@ pub struct Interpreter<'source> { // TODO: handle recursive calls where same expr could have different values. loop_var_values: BTreeMap<&'source Expr<'source>, Value>, contexts: Vec>, - functions: HashMap>>, + functions: FunctionTable<'source>, rules: HashMap>>, default_rules: HashMap, Option)>>, processed: BTreeSet<&'source Rule<'source>>, @@ -82,12 +82,12 @@ struct LoopExpr<'source> { } impl<'source> Interpreter<'source> { - pub fn new(modules: Vec<&'source Module<'source>>) -> Result> { + pub fn new(modules: &[&'source Module<'source>]) -> Result> { let mut with_document = Value::new_object(); *Self::make_or_get_value_mut(&mut with_document, &["data"])? = Value::new_object(); *Self::make_or_get_value_mut(&mut with_document, &["input"])? = Value::new_object(); Ok(Interpreter { - modules, + modules: modules.to_vec(), module: None, schedule: None, current_module_path: String::default(), @@ -99,7 +99,7 @@ impl<'source> Interpreter<'source> { scopes: vec![Scope::new()], contexts: vec![], loop_var_values: BTreeMap::new(), - functions: HashMap::new(), + functions: FunctionTable::new(), rules: HashMap::new(), default_rules: HashMap::new(), processed: BTreeSet::new(), @@ -830,7 +830,7 @@ impl<'source> Interpreter<'source> { span, fcn, params, - get_extra_arg(expr, &HashMap::new()), + get_extra_arg(expr, &self.functions), true, )?, _ => self.eval_expr(expr)?, @@ -860,7 +860,7 @@ impl<'source> Interpreter<'source> { span, fcn, params, - get_extra_arg(expr, &HashMap::new()), + get_extra_arg(expr, &self.functions), false, )?, _ => self.eval_expr(expr)?, @@ -1443,7 +1443,7 @@ impl<'source> Interpreter<'source> { } match self.functions.get(&path) { - Some(r) => Ok(r), + Some((r, _)) => Ok(r), _ => { bail!(fcn.span().error("function not found")) } @@ -1871,18 +1871,25 @@ impl<'source> Interpreter<'source> { self.eval_output_expr() } else { let mut result = Ok(true); - for body in bodies { - self.contexts.push(ctx.clone()); + for (idx, body) in bodies.iter().enumerate() { + if idx == 0 { + self.contexts.push(ctx.clone()); + } else { + self.contexts.pop(); + let output_expr = body.assign.as_ref().map(|e| &e.value); + self.contexts.push(Context { + key_expr: None, + output_expr, + value: Value::new_array(), + result: None, + results: QueryResults::default(), + }); + } result = self.eval_query(&body.query); if matches!(&result, Ok(true) | Err(_)) { break; } - - // TODO: Manage other scoped data. - if bodies.len() > 1 { - unimplemented!("else bodies"); - } } result }; @@ -2006,7 +2013,7 @@ impl<'source> Interpreter<'source> { Ok(m) } - pub fn update_function_table(&mut self) -> Result<()> { + /* pub fn update_function_table(&mut self) -> Result<()> { for module in self.modules.clone() { let prev_module = self.set_current_module(Some(module))?; let module_path = @@ -2034,6 +2041,7 @@ impl<'source> Interpreter<'source> { let full_path = Self::get_path_string(refr, Some(module_path.as_str()))?; if let Some(functions) = self.functions.get_mut(&full_path) { + // TODO: check function arity. functions.push(rule); } else { self.functions.insert(full_path, vec![rule]); @@ -2043,7 +2051,7 @@ impl<'source> Interpreter<'source> { self.set_current_module(prev_module)?; } Ok(()) - } + }*/ fn get_rule_refr(rule: &'source Rule<'source>) -> &'source Expr<'source> { match rule { @@ -2234,6 +2242,22 @@ impl<'source> Interpreter<'source> { } self.processed.insert(rule); + } else if let RuleHead::Func { refr, .. } = rule_head { + let mut path = + Parser::get_path_ref_components(&self.current_module()?.package.refr)?; + + Parser::get_path_ref_components_into(refr, &mut path)?; + let path: Vec<&str> = path.iter().map(|s| s.text()).collect(); + + // Ensure that for functions with a nesting level (e.g: a.foo), + // `a` is created as an empty object. + if path.len() > 1 { + let value = + Self::make_or_get_value_mut(&mut self.data, &path[0..path.len() - 1])?; + if value == &Value::Undefined { + *value = Value::new_object(); + } + } } } _ => bail!("internal error: unexpected"), @@ -2283,7 +2307,8 @@ impl<'source> Interpreter<'source> { } self.check_default_rules()?; - self.update_function_table()?; + self.functions = gather_functions(&self.modules)?; + self.gather_rules()?; self.init_data = self.data.clone(); @@ -2469,7 +2494,7 @@ impl<'source> Interpreter<'source> { if old == new { bail!(refr.span().error("multiple default rules for the variable with the same index")); } - } else { + } else if index.is_some() || i.is_some() { bail!(refr .span() .error("conflict type with the default rules")); diff --git a/src/scheduler.rs b/src/scheduler.rs index 826958e..1f16a52 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -4,9 +4,9 @@ use crate::ast::Expr::*; use crate::ast::*; use crate::lexer::Span; -use crate::utils; +use crate::utils::*; -use std::collections::{BTreeMap, BTreeSet, HashMap, VecDeque}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::string::String; use anyhow::{bail, Result}; @@ -380,6 +380,7 @@ pub struct Analyzer<'a> { locals: BTreeMap<&'a Query<'a>, Scope<'a>>, scopes: Vec>, order: BTreeMap<&'a Query<'a>, Vec>, + functions: FunctionTable<'a>, } #[derive(Clone)] @@ -401,11 +402,13 @@ impl<'a> Analyzer<'a> { locals: BTreeMap::new(), scopes: vec![], order: BTreeMap::new(), + functions: FunctionTable::new(), } } - pub fn analyze(mut self, modules: &'a [Module<'a>]) -> Result { + pub fn analyze(mut self, modules: &'a [&'a Module<'a>]) -> Result { self.add_rules(modules)?; + self.functions = gather_functions(modules)?; for m in modules { self.analyze_module(m)?; @@ -419,7 +422,7 @@ impl<'a> Analyzer<'a> { pub fn analyze_query_snippet( mut self, - modules: &'a [Module<'a>], + modules: &'a [&'a Module<'a>], query: &'a Query<'a>, ) -> Result> { self.add_rules(modules)?; @@ -431,9 +434,9 @@ impl<'a> Analyzer<'a> { }) } - fn add_rules(&mut self, modules: &'a [Module<'a>]) -> Result<()> { + fn add_rules(&mut self, modules: &'a [&'a Module<'a>]) -> Result<()> { for m in modules { - let path = utils::get_path_string(&m.package.refr, Some("data"))?; + let path = get_path_string(&m.package.refr, Some("data"))?; let scope: &mut Scope = self.packages.entry(path).or_default(); for r in &m.policy { let var = match r { @@ -454,7 +457,7 @@ impl<'a> Analyzer<'a> { } fn analyze_module(&mut self, m: &'a Module<'a>) -> Result<()> { - let path = utils::get_path_string(&m.package.refr, Some("data"))?; + let path = get_path_string(&m.package.refr, Some("data"))?; let scope = match self.packages.get(&path) { Some(s) => s, _ => bail!("internal error: package scope missing"), @@ -956,8 +959,7 @@ impl<'a> Analyzer<'a> { // TODO: vars in compr } Literal::Expr { expr, .. } => { - if let Some(Expr::Var(return_arg)) = utils::get_extra_arg(expr, &HashMap::new()) - { + if let Some(Expr::Var(return_arg)) = get_extra_arg(expr, &self.functions) { let (mut used_vars, comprs) = Self::gather_used_vars_comprs_index_vars( expr, &mut scope, diff --git a/src/utils.rs b/src/utils.rs index 072c483..5f824da 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -4,7 +4,7 @@ use crate::ast::*; use crate::builtins::*; -use std::collections::HashMap; +use std::collections::BTreeMap; use anyhow::{bail, Result}; @@ -39,10 +39,12 @@ pub fn get_path_string(refr: &Expr, document: Option<&str>) -> Result { Ok(comps.join(".")) } -pub fn get_extra_arg<'a>(expr: &'a Expr, arities: &HashMap) -> Option<&'a Expr<'a>> { +pub type FunctionTable<'a> = BTreeMap>, u8)>; + +pub fn get_extra_arg<'a>(expr: &'a Expr, functions: &FunctionTable) -> Option<&'a Expr<'a>> { if let Expr::Call { fcn, params, .. } = expr { if let Ok(path) = get_path_string(fcn, None) { - let n_args = if let Some(n_args) = arities.get(&path) { + let n_args = if let Some((_, n_args)) = functions.get(&path) { *n_args } else if let Some((_, n_args)) = BUILTINS.get(path.as_str()) { *n_args @@ -59,3 +61,34 @@ pub fn get_extra_arg<'a>(expr: &'a Expr, arities: &HashMap) -> Optio None } + +pub fn gather_functions<'a>(modules: &[&'a Module<'a>]) -> Result> { + let mut table = FunctionTable::new(); + + for module in modules { + let module_path = get_path_string(&module.package.refr, Some("data"))?; + for rule in &module.policy { + if let Rule::Spec { + span, + head: RuleHead::Func { refr, args, .. }, + .. + } = rule + { + let full_path = get_path_string(refr, Some(module_path.as_str()))?; + + if let Some((functions, arity)) = table.get_mut(&full_path) { + if args.len() as u8 != *arity { + bail!(span.error( + format!("{full_path} was previously defined with {arity} arguments.") + .as_str() + )); + } + functions.push(rule); + } else { + table.insert(full_path, (vec![rule], args.len() as u8)); + } + } + } + } + Ok(table) +} diff --git a/tests/interpreter/mod.rs b/tests/interpreter/mod.rs index 5d765be..7e98ed6 100644 --- a/tests/interpreter/mod.rs +++ b/tests/interpreter/mod.rs @@ -240,11 +240,12 @@ pub fn eval_file_first_rule( }; let mut parser = regorus::Parser::new(&query_source)?; let query_node = parser.parse_query(query_span, "")?; - let query_schedule = regorus::Analyzer::new().analyze_query_snippet(&modules, &query_node)?; + let query_schedule = + regorus::Analyzer::new().analyze_query_snippet(&modules_ref, &query_node)?; let analyzer = Analyzer::new(); - let schedule = analyzer.analyze(&modules)?; + let schedule = analyzer.analyze(&modules_ref)?; - let mut interpreter = interpreter::Interpreter::new(modules_ref)?; + let mut interpreter = interpreter::Interpreter::new(&modules_ref)?; if let Some(input) = input_opt { // if inputs are defined then first the evaluation if prepared interpreter.prepare_for_eval(Some(schedule), &data_opt)?; @@ -332,12 +333,13 @@ pub fn eval_file( }; let mut parser = regorus::Parser::new(&query_source)?; let query_node = parser.parse_query(query_span, "")?; - let query_schedule = regorus::Analyzer::new().analyze_query_snippet(&modules, &query_node)?; + let query_schedule = + regorus::Analyzer::new().analyze_query_snippet(&modules_ref, &query_node)?; let analyzer = Analyzer::new(); - let schedule = analyzer.analyze(&modules)?; + let schedule = analyzer.analyze(&modules_ref)?; - let mut interpreter = interpreter::Interpreter::new(modules_ref)?; + let mut interpreter = interpreter::Interpreter::new(&modules_ref)?; if let Some(input) = input_opt { // if inputs are defined then first the evaluation if prepared interpreter.prepare_for_eval(Some(schedule), &data_opt)?; diff --git a/tests/scheduler/analyzer/mod.rs b/tests/scheduler/analyzer/mod.rs index aa552f7..5340774 100644 --- a/tests/scheduler/analyzer/mod.rs +++ b/tests/scheduler/analyzer/mod.rs @@ -52,9 +52,10 @@ fn analyze_file(regos: &[String], expected_scopes: &[Scope]) -> Result<()> { let mut parser = Parser::new(source)?; modules.push(parser.parse()?); } + let modules_ref: Vec<&Module> = modules.iter().collect(); let analyzer = Analyzer::new(); - let schedule = analyzer.analyze(&modules)?; + let schedule = analyzer.analyze(&modules_ref)?; for (idx, (_, scope)) in schedule.scopes.iter().enumerate() { if idx > expected_scopes.len() { bail!("extra scope generated.")