Minimize PR 22 (#26)

* specific functions added to eval different rego components

Signed-off-by: eric-therond <eric.therond.fr@gmail.com>

* allow multiple inputs and results

Signed-off-by: eric-therond <eric.therond.fr@gmail.com>

* prepare_for_eval is necessary to be called

Signed-off-by: eric-therond <eric.therond.fr@gmail.com>

* test with the suggested code examples and clean scopes

Signed-off-by: eric-therond <eric.therond.fr@gmail.com>

* try to refactor first steps of evaluations

Signed-off-by: eric-therond <eric.therond.fr@gmail.com>

* improve coverage and fix clean state internal evaluation

Signed-off-by: eric-therond <eric.therond.fr@gmail.com>

* add getters and setters and fix clean function

Signed-off-by: eric-therond <eric.therond.fr@gmail.com>

* Tests are single input by default. Multi input specified via "many!" marker.

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

---------

Signed-off-by: eric-therond <eric.therond.fr@gmail.com>
Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
Co-authored-by: eric-therond <eric.therond.fr@gmail.com>
This commit is contained in:
Anand Krishnamoorthi
2023-10-09 14:27:31 -07:00
committed by GitHub
parent 2ba718ba72
commit 2436467fbd
11 changed files with 430 additions and 66 deletions
+123 -33
View File
@@ -20,8 +20,10 @@ pub struct Interpreter<'source> {
module: Option<&'source Module<'source>>,
schedule: Option<&'source Schedule<'source>>,
current_module_path: String,
prepared: bool,
input: Value,
data: Value,
init_data: Value,
scopes: Vec<Scope>,
// TODO: handle recursive calls where same expr could have different values.
loop_var_values: BTreeMap<&'source Expr<'source>, Value>,
@@ -58,8 +60,10 @@ impl<'source> Interpreter<'source> {
module: None,
schedule: None,
current_module_path: String::default(),
prepared: false,
input: Value::new_object(),
data: Value::new_object(),
init_data: Value::new_object(),
scopes: vec![Scope::new()],
contexts: vec![],
loop_var_values: BTreeMap::new(),
@@ -74,6 +78,44 @@ impl<'source> Interpreter<'source> {
})
}
pub fn get_modules(&mut self) -> &mut Vec<&'source Module<'source>> {
&mut self.modules
}
pub fn set_data(&mut self, data: Value) {
self.data = data;
}
pub fn get_data(&mut self) -> &mut Value {
&mut self.data
}
fn clean_internal_evaluation_state(&mut self) {
self.data = self.init_data.clone();
self.processed.clear();
self.loop_var_values.clear();
self.scopes = vec![Scope::new()];
self.contexts = vec![];
}
fn checks_for_eval(&mut self, input: &Option<Value>, enable_tracing: bool) -> Result<()> {
if !self.prepared {
bail!("prepare_for_eval should be called before eval_modules");
}
self.traces = match enable_tracing {
true => Some(vec![]),
false => None,
};
if let Some(input) = input {
self.input = input.clone();
info!("input: {:#?}", self.input);
}
Ok(())
}
fn current_module(&self) -> Result<&'source Module<'source>> {
self.module
.ok_or_else(|| anyhow!("internal error: current module not set"))
@@ -1435,6 +1477,7 @@ impl<'source> Interpreter<'source> {
let module_path =
Self::get_path_string(&self.current_module()?.package.refr, Some("data"))?;
let path = module_path + "." + name;
self.ensure_rule_evaluated(path)?;
let mut path: Vec<&str> =
@@ -1933,62 +1976,65 @@ impl<'source> Interpreter<'source> {
head: rule_head,
bodies: rule_body,
} => {
if matches!(rule_head, RuleHead::Func { .. }) {
return Ok(());
}
let (ctx, mut path) = self.make_rule_context(rule_head)?;
let special_set = matches!((ctx.output_expr, &ctx.value), (None, Value::Set(_)));
let value = match self.eval_rule_bodies(ctx, span, rule_body)? {
Value::Set(_) if special_set => {
let entry = path[path.len() - 1].text();
let mut s = BTreeSet::new();
s.insert(Value::String(entry.to_owned()));
path = path[0..path.len() - 1].to_vec();
Value::from_set(s)
if !matches!(rule_head, RuleHead::Func { .. }) {
let (ctx, mut path) = self.make_rule_context(rule_head)?;
let special_set =
matches!((ctx.output_expr, &ctx.value), (None, Value::Set(_)));
let value = match self.eval_rule_bodies(ctx, span, rule_body)? {
Value::Set(_) if special_set => {
let entry = path[path.len() - 1].text();
let mut s = BTreeSet::new();
s.insert(Value::String(entry.to_owned()));
path = path[0..path.len() - 1].to_vec();
Value::from_set(s)
}
v => v,
};
if value != Value::Undefined {
let paths: Vec<&str> = path.iter().map(|s| s.text()).collect();
let vref = Self::make_or_get_value_mut(&mut self.data, &paths[..])?;
Self::merge_value(span, vref, value)?;
}
v => v,
};
if value != Value::Undefined {
let paths: Vec<&str> = path.iter().map(|s| s.text()).collect();
let vref = Self::make_or_get_value_mut(&mut self.data, &paths[..])?;
Self::merge_value(span, vref, value)?;
self.processed.insert(rule);
}
}
_ => bail!("internal error: unexpected"),
}
self.set_current_module(prev_module)?;
self.processed.insert(rule);
match self.active_rules.pop() {
Some(r) if r == rule => Ok(()),
_ => bail!("internal error: current rule not active"),
}
}
pub fn eval(
pub fn eval_rule_with_input(
&mut self,
data: &Option<Value>,
module: &'source Module<'source>,
rule: &'source Rule<'source>,
input: &Option<Value>,
enable_tracing: bool,
schedule: Option<&'source Schedule<'source>>,
) -> Result<Value> {
self.schedule = schedule;
self.traces = match enable_tracing {
true => Some(vec![]),
false => None,
};
self.checks_for_eval(input, enable_tracing)?;
self.clean_internal_evaluation_state();
self.eval_rule(module, rule)?;
Ok(self.data.clone())
}
pub fn prepare_for_eval(
&mut self,
schedule: Option<&'source Schedule<'source>>,
data: &Option<Value>,
) -> Result<()> {
self.schedule = schedule;
self.builtins_cache.clear();
if let Some(input) = input {
self.input = input.clone();
info!("input: {:#?}", self.input);
}
if let Some(data) = data {
self.data = data.clone();
}
// Ensure that each module has an empty object
for m in &self.modules {
let path = Parser::get_path_ref_components(&m.package.refr)?;
@@ -2003,6 +2049,39 @@ impl<'source> Interpreter<'source> {
self.update_function_table()?;
self.gather_rules()?;
self.init_data = self.data.clone();
self.prepared = true;
Ok(())
}
pub fn eval_module(
&mut self,
module: &'source Module<'source>,
input: &Option<Value>,
enable_tracing: bool,
) -> Result<Value> {
self.checks_for_eval(input, enable_tracing)?;
self.clean_internal_evaluation_state();
for rule in &module.policy {
self.eval_rule(module, rule)?;
}
// Defer the evaluation of the default rules to here
let prev_module = self.set_current_module(Some(module))?;
for rule in &module.policy {
self.eval_default_rule(rule)?;
}
self.set_current_module(prev_module)?;
Ok(self.data.clone())
}
pub fn eval_modules(&mut self, input: &Option<Value>, enable_tracing: bool) -> Result<Value> {
self.checks_for_eval(input, enable_tracing)?;
self.clean_internal_evaluation_state();
for module in self.modules.clone() {
for rule in &module.policy {
self.eval_rule(module, rule)?;
@@ -2021,6 +2100,17 @@ impl<'source> Interpreter<'source> {
Ok(self.data.clone())
}
pub fn eval(
&mut self,
data: &Option<Value>,
input: &Option<Value>,
enable_tracing: bool,
schedule: Option<&'source Schedule<'source>>,
) -> Result<Value> {
self.prepare_for_eval(schedule, data)?;
self.eval_modules(input, enable_tracing)
}
pub fn eval_query_snippet(
&mut self,
snippet: &'source Expr<'source>,