Improvements (#33)

1. Skip recording undefined variables
2. Parse `in` correctly if it is not imported.
3. base64.decode
4. Handle `with` modifier for qualified data and input.

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2023-11-01 21:57:06 -07:00
committed by GitHub
parent 6228eaab4a
commit c53d002347
8 changed files with 149 additions and 83 deletions

View File

@@ -14,6 +14,7 @@ log = "0.4.17"
env_logger="0.10.0"
lazy_static = "1.4.0"
rand = "0.8.5"
data-encoding = "2.4.0"
[dev-dependencies]
clap = { version = "4.4.7", features = ["derive"] }

28
src/builtins/encoding.rs Normal file
View File

@@ -0,0 +1,28 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
use crate::ast::Expr;
use crate::builtins;
use crate::builtins::utils::{ensure_args_count, ensure_string};
use crate::lexer::Span;
use crate::value::Value;
use std::collections::HashMap;
use anyhow::Result;
use data_encoding::BASE64;
pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {
m.insert("base64.decode", (base64_decode, 1));
}
fn base64_decode(span: &Span, params: &[Expr], args: &[Value]) -> Result<Value> {
let name = "base64.decode";
ensure_args_count(span, name, params, args, 1)?;
let encoded_str = ensure_string(name, &params[0], &args[0])?;
let decoded_bytes = BASE64.decode(encoded_str.as_bytes())?;
Ok(Value::String(
String::from_utf8_lossy(&decoded_bytes).to_string(),
))
}

View File

@@ -7,6 +7,7 @@ mod bitwise;
pub mod comparison;
mod conversions;
mod debugging;
mod encoding;
pub mod numbers;
mod objects;
pub mod sets;
@@ -44,7 +45,7 @@ lazy_static! {
conversions::register(&mut m);
//units::register(&mut m);
types::register(&mut m);
//encoding::register(&mut m);
encoding::register(&mut m);
//token_signing::register(&mut m);
//token_verification::register(&mut m);
//time::register(&mut m);

View File

@@ -26,6 +26,7 @@ pub struct Interpreter<'source> {
input: Value,
data: Value,
init_data: Value,
with_document: Value,
scopes: Vec<Scope>,
// TODO: handle recursive calls where same expr could have different values.
loop_var_values: BTreeMap<&'source Expr<'source>, Value>,
@@ -81,6 +82,9 @@ struct LoopExpr<'source> {
impl<'source> Interpreter<'source> {
pub fn new(modules: Vec<&'source Module<'source>>) -> Result<Interpreter<'source>> {
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,
module: None,
@@ -90,6 +94,7 @@ impl<'source> Interpreter<'source> {
input: Value::new_object(),
data: Value::new_object(),
init_data: Value::new_object(),
with_document,
scopes: vec![Scope::new()],
contexts: vec![],
loop_var_values: BTreeMap::new(),
@@ -241,6 +246,9 @@ impl<'source> Interpreter<'source> {
match ident {
"_" => 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,
@@ -478,6 +486,11 @@ impl<'source> Interpreter<'source> {
}
};
// Omit recording undefined values.
if value == Value::Undefined {
return Ok(Value::Bool(false));
}
self.add_variable_or(name)?;
// TODO: optimize this
@@ -803,44 +816,12 @@ impl<'source> Interpreter<'source> {
Value::from_map(expr)
}
fn eval_stmt(
fn eval_stmt_impl(
&mut self,
stmt: &'source LiteralStmt<'source>,
stmts: &[&'source LiteralStmt<'source>],
) -> Result<bool> {
let mut to_restore = vec![];
for wm in &stmt.with_mods {
// Evaluate value and ref
let value = self.eval_expr(&wm.r#as)?;
let path = Parser::get_path_ref_components(&wm.refr)?;
let mut path: Vec<&str> = path.iter().map(|s| s.text()).collect();
// TODO: multiple modules and qualified path
if path.len() > 2 && format!("{}.{}", path[0], path[1]) == self.current_module_path {
path = path[1..].to_vec();
}
// Set new values in modifications table
let mut saved = false;
for (i, _) in path.iter().enumerate() {
let vref = Self::make_or_get_value_mut(&mut self.data, &path[0..i])?;
if vref == &Value::Undefined {
to_restore.push((path[0..i].to_vec(), vref.clone()));
saved = false;
break;
}
}
// TODO: input
let vref = Self::make_or_get_value_mut(&mut self.data, &path[..])?;
if !saved {
to_restore.push((path, vref.clone()));
}
*vref = value;
}
let r = Ok(match &stmt.literal {
Ok(match &stmt.literal {
Literal::Expr { span, expr, .. } => {
let value = match expr {
Expr::Call { span, fcn, params } => self.eval_call(
@@ -944,16 +925,59 @@ impl<'source> Interpreter<'source> {
}
self.eval_every(span, key, value, domain, query)?
}
});
})
}
for (path, value) in to_restore.into_iter().rev() {
if value == Value::Undefined {
unimplemented!("handle undefined restore");
} else {
let vref = Self::make_or_get_value_mut(&mut self.data, &path[..])?;
*vref = value;
fn eval_stmt(
&mut self,
stmt: &'source LiteralStmt<'source>,
stmts: &[&'source LiteralStmt<'source>],
) -> Result<bool> {
let saved_state = if !stmt.with_mods.is_empty() {
// Save state;
let with_document = self.with_document.clone();
let input = self.input.clone();
let data = self.data.clone();
let processed = self.processed.clone();
// Apply with modifiers.
for wm in &stmt.with_mods {
// Evaluate value and ref
let value = self.eval_expr(&wm.r#as)?;
let path = Parser::get_path_ref_components(&wm.refr)?;
let path: Vec<&str> = path.iter().map(|s| s.text()).collect();
if path[0] == "input" || path[0] == "data" {
*Self::make_or_get_value_mut(&mut self.with_document, &path[..])? = value;
} /* else if path.len() == 1 {
// TODO: handle var in current module.
} else {
// TODO: error about input, data
} */
// TODO: functions
}
self.data = self.with_document["data"].clone();
self.input = self.with_document["input"].clone();
self.processed.clear();
(with_document, input, data, processed)
} else {
(
Value::Undefined,
Value::Undefined,
Value::Undefined,
BTreeSet::new(),
)
};
let r = self.eval_stmt_impl(stmt, stmts);
// Restore state.
if saved_state.0 != Value::Undefined {
(self.with_document, self.input, self.data, self.processed) = saved_state;
}
r
}
@@ -989,8 +1013,8 @@ impl<'source> Interpreter<'source> {
if let Some(idx) = self.lookup_local_var(loop_expr.index) {
if loop_expr_value[&idx] != Value::Undefined {
result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
return Ok(result);
}
return Ok(result);
}
// Save the current scope and restore it after evaluating the statements so
@@ -1672,6 +1696,14 @@ impl<'source> Interpreter<'source> {
// Ensure that rules are evaluated
if name == "data" {
let v = Self::get_value_chained(self.data.clone(), fields);
// If the rule has already been evaluated or specified via a with modifier,
// use that value.
if v != Value::Undefined {
return Ok(v);
}
// Evaluate rule corresponding to longest matching path.
for i in (1..fields.len() + 1).rev() {
let path = "data.".to_owned() + &fields[0..i].join(".");
@@ -1680,15 +1712,9 @@ impl<'source> Interpreter<'source> {
break;
}
}
Ok(Self::get_value_chained(self.data.clone(), fields))
} else if !self.modules.is_empty() {
// Add module prefix and ensure that any matching rule is evaluated.
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> =
Parser::get_path_ref_components(&self.module.unwrap().package.refr)?
.iter()
@@ -1696,6 +1722,21 @@ impl<'source> Interpreter<'source> {
.collect();
path.push(name);
let v = Self::get_value_chained(self.data.clone(), &path);
// If the rule has already been evaluated or specified via a with modifier,
// use that value.
if v != Value::Undefined {
return Ok(v);
}
// Add module prefix and ensure that any matching rule is evaluated.
let module_path =
Self::get_path_string(&self.current_module()?.package.refr, Some("data"))?;
let rule_path = module_path + "." + name;
self.ensure_rule_evaluated(rule_path)?;
let value = Self::get_value_chained(self.data.clone(), &path[..]);
Ok(Self::get_value_chained(value, fields))
} else {

View File

@@ -714,7 +714,7 @@ impl<'source> Parser<'source> {
let start = self.tok.1.start;
let mut expr = self.parse_bool_expr()?;
while self.tok.1.text() == "in" {
while self.tok.1.text() == "in" && self.future_keywords.get("in").is_some() {
expr = self.parse_membership_tail(start, expr, None)?;
}
@@ -731,7 +731,7 @@ impl<'source> Parser<'source> {
expr = self.parse_membership_tail(start, expr, Some(value))?;
}
while self.tok.1.text() == "in" {
while self.tok.1.text() == "in" && self.is_imported_future_keyword("in") {
expr = self.parse_membership_tail(start, expr, None)?;
}
@@ -834,7 +834,7 @@ impl<'source> Parser<'source> {
vars.push(span);
}
if self.tok.1.text() != "in" || self.future_keywords.get("in").is_none() {
if self.tok.1.text() != "in" || !self.is_imported_future_keyword("in") {
if self.tok.1.text() == "in" {
self.warn_future_keyword();
}

View File

@@ -913,8 +913,10 @@ impl<'a> Analyzer<'a> {
collection,
&mut scope,
&mut first_use,
&mut col_definitions, // TODO: handle these definitions
&mut col_definitions,
)?;
definitions.append(&mut col_definitions);
self.process_comprs(
&col_comprs[..],
&mut scope,

View File

@@ -181,18 +181,14 @@ pub fn check_output(computed_results: &[Value], expected_results: &[Value]) -> R
Ok(())
}
fn query_results_to_value(query_results: QueryResults) -> Result<Value> {
fn push_query_results(query_results: QueryResults, results: &mut Vec<Value>) {
if let Some(query_result) = query_results.result.last() {
if !query_result.bindings.is_empty_object() {
return Ok(query_result.bindings.clone());
} else {
return match query_result.expressions.last() {
Some(v) => Ok(v["value"].clone()),
_ => bail!("no expressions in query results"),
};
results.push(query_result.bindings.clone());
} else if let Some(v) = query_result.expressions.last() {
results.push(v["value"].clone());
}
}
bail!("query result incomplete")
}
pub fn eval_file_first_rule(
@@ -268,22 +264,20 @@ pub fn eval_file_first_rule(
}
// Now eval the query.
results.push(query_results_to_value(interpreter.eval_user_query(
&query_node,
&query_schedule,
enable_tracing,
)?)?);
push_query_results(
interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?,
&mut results,
);
}
} else {
// it no input is defined then one evaluation of all modules is performed
interpreter.eval(&data_opt, &None, enable_tracing, Some(schedule))?;
// Now eval the query.
results.push(query_results_to_value(interpreter.eval_user_query(
&query_node,
&query_schedule,
enable_tracing,
)?)?);
// Now eval the query
push_query_results(
interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?,
&mut results,
);
}
Ok(results)
@@ -359,22 +353,20 @@ pub fn eval_file(
interpreter.eval_modules(&Some(input), enable_tracing)?;
// Now eval the query.
results.push(query_results_to_value(interpreter.eval_user_query(
&query_node,
&query_schedule,
enable_tracing,
)?)?);
push_query_results(
interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?,
&mut results,
);
}
} else {
// it no input is defined then one evaluation of all modules is performed
interpreter.eval(&data_opt, &None, enable_tracing, Some(schedule))?;
// Now eval the query.
results.push(query_results_to_value(interpreter.eval_user_query(
&query_node,
&query_schedule,
enable_tracing,
)?)?);
push_query_results(
interpreter.eval_user_query(&query_node, &query_schedule, enable_tracing)?,
&mut results,
);
}
Ok(results)

View File

@@ -134,6 +134,7 @@ cases:
- note: same-line
rego: |
package test
import future.keywords.in
x = y{
some a b in {4, 5}
[1, 2, 3][a] == 3