Allow with modifier for builtin and user functions (#42)

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
This commit is contained in:
Anand Krishnamoorthi
2023-11-12 21:39:59 -08:00
committed by GitHub
parent 1c492144b3
commit a1d0f8576a
9 changed files with 216 additions and 78 deletions

View File

@@ -13,6 +13,7 @@ pub mod numbers;
mod objects;
pub mod sets;
mod strings;
mod time;
mod tracing;
pub mod types;
mod utils;
@@ -51,7 +52,7 @@ lazy_static! {
encoding::register(&mut m);
//token_signing::register(&mut m);
//token_verification::register(&mut m);
//time::register(&mut m);
time::register(&mut m);
//cryptography::register(&mut m);
//graphs::register(&mut m);
//graphql::register(&mut m);

30
src/builtins/time.rs Normal file
View File

@@ -0,0 +1,30 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
use crate::ast::Expr;
use crate::builtins;
use crate::builtins::utils::ensure_args_count;
use crate::lexer::Span;
use crate::value::Value;
use std::collections::HashMap;
use std::time::SystemTime;
use anyhow::{bail, Result};
pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {
m.insert("time.now_ns", (now_ns, 0));
}
fn now_ns(span: &Span, params: &[Expr], args: &[Value]) -> Result<Value> {
let name = "time.now_ns";
ensure_args_count(span, name, params, args, 0)?;
let now = SystemTime::now();
let elapsed = match now.duration_since(SystemTime::UNIX_EPOCH) {
Ok(e) => e,
Err(e) => bail!(span.error(format!("could not fetch elapsed time. {e}").as_str())),
};
let nanos = elapsed.as_nanos();
Ok(Value::from_u128(nanos))
}

View File

@@ -2,7 +2,7 @@
// Licensed under the MIT License.
use crate::ast::*;
use crate::builtins;
use crate::builtins::{self, BuiltinFcn};
use crate::lexer::Span;
use crate::parser::Parser;
use crate::scheduler::*;
@@ -27,6 +27,7 @@ pub struct Interpreter<'source> {
data: Value,
init_data: Value,
with_document: Value,
with_functions: BTreeMap<String, String>,
scopes: Vec<Scope>,
// TODO: handle recursive calls where same expr could have different values.
loop_var_values: BTreeMap<Ref<'source, Expr>, Value>,
@@ -96,6 +97,7 @@ impl<'source> Interpreter<'source> {
data: Value::new_object(),
init_data: Value::new_object(),
with_document,
with_functions: BTreeMap::new(),
scopes: vec![Scope::new()],
contexts: vec![],
loop_var_values: BTreeMap::new(),
@@ -842,7 +844,11 @@ impl<'source> Interpreter<'source> {
span,
fcn,
params,
get_extra_arg(expr, &self.functions),
get_extra_arg(
expr,
Some(self.current_module_path.as_str()),
&self.functions,
),
true,
)?,
_ => self.eval_expr(expr)?,
@@ -872,7 +878,11 @@ impl<'source> Interpreter<'source> {
span,
fcn,
params,
get_extra_arg(expr, &self.functions),
get_extra_arg(
expr,
Some(self.current_module_path.as_str()),
&self.functions,
),
false,
)?,
_ => self.eval_expr(expr)?,
@@ -947,19 +957,43 @@ impl<'source> Interpreter<'source> {
stmt: &'source LiteralStmt,
stmts: &[&'source LiteralStmt],
) -> Result<bool> {
let mut skip_exec = false;
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();
let with_functions = self.with_functions.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();
let target = path.join(".");
let value = match self.eval_expr(&wm.r#as) {
Ok(Value::Undefined) => {
if let Ok(fcn_path) = get_path_string(&wm.r#as, None) {
let span = wm.r#as.span();
if self.lookup_function_by_name(&fcn_path).is_some() {
let mut fcn_path = get_path_string(&wm.r#as, None)?;
if !fcn_path.starts_with("data.") {
fcn_path = self.current_module_path.clone() + "." + &fcn_path;
}
self.with_functions.insert(target, fcn_path);
continue;
} else if self.lookup_builtin(span, &fcn_path).is_ok() {
self.with_functions.insert(target, fcn_path);
continue;
}
}
skip_exec = true;
continue;
}
Ok(v) => v,
Err(e) => return Err(e),
};
if path[0] == "input" || path[0] == "data" {
*Self::make_or_get_value_mut(&mut self.with_document, &path[..])? = value;
@@ -968,28 +1002,38 @@ impl<'source> Interpreter<'source> {
} 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)
(with_document, input, data, processed, with_functions)
} else {
(
Value::Undefined,
Value::Undefined,
Value::Undefined,
BTreeSet::new(),
BTreeMap::new(),
)
};
let r = self.eval_stmt_impl(stmt, stmts);
let r = if !skip_exec {
self.eval_stmt_impl(stmt, stmts)
} else {
Ok(false)
};
// Restore state.
if saved_state.0 != Value::Undefined {
(self.with_document, self.input, self.data, self.processed) = saved_state;
(
self.with_document,
self.input,
self.data,
self.processed,
self.with_functions,
) = saved_state;
}
r
@@ -1443,24 +1487,22 @@ impl<'source> Interpreter<'source> {
}
}
fn lookup_function(&self, fcn: &'source Expr) -> Result<&Vec<&'source Rule>> {
let mut path = Self::get_path_string(fcn, None)?;
fn lookup_function_by_name(&self, path: &str) -> Option<&Vec<&'source Rule>> {
let mut path = path.to_owned();
if !path.starts_with("data.") {
path = self.current_module_path.clone() + "." + &path;
}
match self.functions.get(&path) {
Some((r, _)) => Ok(r),
_ => {
bail!(fcn.span().error("function not found"))
}
Some((f, _)) => Some(f),
_ => None,
}
}
fn eval_builtin_call(
&mut self,
span: &'source Span,
name: String,
name: &str,
builtin: builtins::BuiltinFcn,
params: &'source [Expr],
) -> Result<Value> {
@@ -1474,7 +1516,7 @@ impl<'source> Interpreter<'source> {
}
}
let cache = builtins::must_cache(name.as_str());
let cache = builtins::must_cache(name);
if let Some(name) = &cache {
if let Some(v) = self.builtins_cache.get(&(name, args.clone())) {
return Ok(v.clone());
@@ -1496,30 +1538,47 @@ impl<'source> Interpreter<'source> {
Ok(v)
}
fn lookup_builtin(&self, span: &'source Span, path: &str) -> Result<Option<&BuiltinFcn>> {
Ok(if let Some(builtin) = builtins::BUILTINS.get(path) {
Some(builtin)
} else if let Some(builtin) = builtins::DEPRECATED.get(path) {
if !self.allow_deprecated {
bail!(span.error(format!("{path} is deprecated").as_str()))
}
Some(builtin)
} else {
None
})
}
fn eval_call_impl(
&mut self,
span: &'source Span,
fcn: &'source Expr,
params: &'source [Expr],
) -> Result<Value> {
let fcns_rules = match self.lookup_function(fcn) {
Ok(r) => r,
let fcn_path = match get_path_string(fcn, None) {
Ok(p) => p,
_ => bail!(span.error("invalid function expression")),
};
let orig_fcn_path = fcn_path;
let fcn_path = match self.with_functions.remove(&orig_fcn_path) {
Some(p) => p,
_ => orig_fcn_path.clone(),
};
let fcns_rules = match self.lookup_function_by_name(&fcn_path) {
Some(r) => r,
_ => {
// Look up builtin function.
// TODO: handle with modifier
if let Ok(path) = Self::get_path_string(fcn, None) {
if let Some(builtin) = builtins::BUILTINS.get(path.as_str()) {
return self.eval_builtin_call(span, path, *builtin, params);
}
if let Some(builtin) = builtins::DEPRECATED.get(path.as_str()) {
if self.allow_deprecated {
return self.eval_builtin_call(span, path, *builtin, params);
} else {
bail!(span.error(format!("{path} is deprecated").as_str()))
}
if let Ok(Some(builtin)) = self.lookup_builtin(span, &fcn_path) {
let r = self.eval_builtin_call(span, &fcn_path.clone(), *builtin, params);
if orig_fcn_path != fcn_path {
self.with_functions.insert(orig_fcn_path, fcn_path);
}
return r;
}
return Err(span
.source
.error(span.line, span.col, "could not find function"));
@@ -1629,6 +1688,11 @@ impl<'source> Interpreter<'source> {
"functions must not produce multiple outputs for same inputs",
));
}
if orig_fcn_path != fcn_path {
self.with_functions.insert(orig_fcn_path, fcn_path);
}
Ok(results[0].clone())
}

View File

@@ -522,19 +522,16 @@ impl<'source> Parser<'source> {
}
"(" if possible_fcn => {
self.next_token()?;
if *self.tok.1.text() == ")" {
return Err(self
.tok
.1
.error("at least one argument required for function calls"));
}
let mut args = vec![self.parse_in_expr()?];
while *self.tok.1.text() == "," {
self.next_token()?;
match *self.tok.1.text() {
")" => break,
"" if self.tok.0 == TokenKind::Eof => break,
_ => args.push(self.parse_in_expr()?),
let mut args = vec![];
if *self.tok.1.text() != ")" {
args.push(self.parse_in_expr()?);
while *self.tok.1.text() == "," {
self.next_token()?;
match *self.tok.1.text() {
")" => break,
"" if self.tok.0 == TokenKind::Eof => break,
_ => args.push(self.parse_in_expr()?),
}
}
}
self.expect(")", "while parsing call expr")?;
@@ -925,19 +922,19 @@ impl<'source> Parser<'source> {
pub fn parse_query(&mut self, mut span: Span, end_delim: &str) -> Result<Query> {
let state = self.clone();
let _is_definite_query = matches!(*self.tok.1.text(), "some" | "every");
let is_definite_query = matches!(*self.tok.1.text(), "some" | "every");
// TODO: empty query?
let mut literals = vec![];
let stmt = match self.parse_literal_stmt() {
Ok(stmt) => stmt,
Err(e) if _is_definite_query => return Err(e),
_ => {
Err(e) if is_definite_query => return Err(e),
Err(_) => {
// There was error parsing the first literal
// Restore the state and return.
*self = state;
return Err(anyhow!("encountered , when expecting {}", end_delim));
bail!(span.error(format!("expecting {end_delim}").as_str()));
}
};
@@ -1170,13 +1167,16 @@ impl<'source> Parser<'source> {
"(" => {
self.check_rule_ref(&rule_ref)?;
self.next_token()?;
let mut args = vec![self.parse_term()?];
while *self.tok.1.text() == "," {
self.next_token()?;
match *self.tok.1.text() {
")" => break,
"" if self.tok.0 == TokenKind::Eof => break,
_ => args.push(self.parse_term()?),
let mut args = vec![];
if *self.tok.1.text() != ")" {
args.push(self.parse_term()?);
while *self.tok.1.text() == "," {
self.next_token()?;
match *self.tok.1.text() {
")" => break,
"" if self.tok.0 == TokenKind::Eof => break,
_ => args.push(self.parse_term()?),
}
}
}
self.expect(")", "while parsing function rule args")?;

View File

@@ -222,8 +222,6 @@ fn traverse<'a>(expr: &'a Expr, f: &mut dyn FnMut(&'a Expr) -> Result<bool>) ->
ArrayCompr { .. } | SetCompr { .. } | ObjectCompr { .. } => (),
Call { params, .. } => {
// TODO: is traversing function needed?
// traverse(fcn, f)?;
for p in params {
traverse(p, f)?;
}
@@ -372,6 +370,7 @@ pub struct Analyzer<'a> {
scopes: Vec<Scope<'a>>,
order: BTreeMap<Ref<'a, Query>, Vec<u16>>,
functions: FunctionTable<'a>,
current_module_path: String,
}
#[derive(Clone)]
@@ -394,6 +393,7 @@ impl<'a> Analyzer<'a> {
scopes: vec![],
order: BTreeMap::new(),
functions: FunctionTable::new(),
current_module_path: String::default(),
}
}
@@ -453,7 +453,7 @@ impl<'a> Analyzer<'a> {
Some(s) => s,
_ => bail!("internal error: package scope missing"),
};
self.current_module_path = path;
self.scopes.push(scope.clone());
for r in &m.policy {
self.analyze_rule(r)?;
@@ -597,6 +597,7 @@ impl<'a> Analyzer<'a> {
scope: &mut Scope<'a>,
first_use: &mut BTreeMap<&'a str, Span>,
definitions: &mut Vec<Definition<'a>>,
return_arg: &Option<Ref<Expr>>,
) -> Result<(Vec<&'a str>, Vec<&'a Expr>)> {
let mut used_vars = vec![];
let mut comprs = vec![];
@@ -608,8 +609,8 @@ impl<'a> Analyzer<'a> {
{
used_vars.push(name);
first_use.entry(name).or_insert(v.clone());
} else if !scope.inputs.contains(name) {
bail!(v.error(format!("Use of undefined variable `{name}` is unsafe").as_str()));
} else if !scope.inputs.contains(name) && Some(Ref::make(e)) != *return_arg {
bail!(v.error(format!("use of undefined variable `{name}` is unsafe").as_str()));
}
Ok(false)
}
@@ -623,6 +624,7 @@ impl<'a> Analyzer<'a> {
scope,
first_use,
definitions,
return_arg,
)?;
definitions.push(Definition {
var,
@@ -764,6 +766,7 @@ impl<'a> Analyzer<'a> {
scope,
first_use,
definitions,
&None,
)?;
self.process_comprs(&comprs[..], scope, first_use, &mut used_vars)?;
let check_first_use = *op == AssignOp::ColEq;
@@ -789,6 +792,7 @@ impl<'a> Analyzer<'a> {
scope,
first_use,
definitions,
&None,
)?;
let check_first_use = false;
self.process_comprs(&comprs[..], scope, first_use, &mut used_vars)?;
@@ -825,8 +829,13 @@ impl<'a> Analyzer<'a> {
self.process_assign_expr(op, lhs, rhs, scope, first_use, definitions)
}
_ => {
let (mut used_vars, comprs) =
Self::gather_used_vars_comprs_index_vars(expr, scope, first_use, definitions)?;
let (mut used_vars, comprs) = Self::gather_used_vars_comprs_index_vars(
expr,
scope,
first_use,
definitions,
&None,
)?;
self.process_comprs(&comprs[..], scope, first_use, &mut used_vars)?;
definitions.push(Definition { var: "", used_vars });
Ok(())
@@ -922,6 +931,7 @@ impl<'a> Analyzer<'a> {
&mut scope,
&mut first_use,
&mut col_definitions,
&None,
)?;
definitions.append(&mut col_definitions);
@@ -948,6 +958,7 @@ impl<'a> Analyzer<'a> {
&mut scope,
&mut first_use,
&mut definitions,
&None,
)?;
if !definitions.is_empty() {
bail!("internal error: non empty definitions");
@@ -964,12 +975,18 @@ impl<'a> Analyzer<'a> {
// TODO: vars in compr
}
Literal::Expr { expr, .. } => {
if let Some(Expr::Var(return_arg)) = get_extra_arg(expr, &self.functions) {
let extra_arg = get_extra_arg(
expr,
Some(self.current_module_path.as_str()),
&self.functions,
);
if let Some(ra @ Expr::Var(return_arg)) = &extra_arg {
let (mut used_vars, comprs) = Self::gather_used_vars_comprs_index_vars(
expr,
&mut scope,
&mut first_use,
&mut definitions,
&Some(Ref::make(ra)),
)?;
let var = if *return_arg.text() != "_" {
// The var in the return argument slot would have been processed as
@@ -1010,6 +1027,7 @@ impl<'a> Analyzer<'a> {
&mut scope,
&mut first_use,
&mut definitions,
&None,
)?;
self.process_comprs(&comprs[..], &mut scope, &mut first_use, &mut uv)?;
definitions.push(Definition {

View File

@@ -39,25 +39,41 @@ pub fn get_path_string(refr: &Expr, document: Option<&str>) -> Result<String> {
pub type FunctionTable<'a> = BTreeMap<String, (Vec<&'a Rule>, u8)>;
pub fn get_extra_arg<'a>(expr: &'a Expr, functions: &FunctionTable) -> Option<&'a Expr> {
fn get_extra_arg_impl<'a>(
expr: &'a Expr,
module: Option<&str>,
functions: &FunctionTable,
) -> Result<Option<&'a Expr>> {
if let Expr::Call { fcn, params, .. } = expr {
if let Ok(path) = get_path_string(fcn, None) {
let n_args = if let Some((_, n_args)) = functions.get(&path) {
*n_args
} else if let Some((_, n_args)) = BUILTINS.get(path.as_str()) {
let full_path = get_path_string(fcn, module)?;
let n_args = if let Some((_, n_args)) = functions.get(&full_path) {
*n_args
} else {
let path = get_path_string(fcn, None)?;
if let Some((_, n_args)) = BUILTINS.get(path.as_str()) {
*n_args
} else if let Some((_, n_args)) = DEPRECATED.get(path.as_str()) {
*n_args
} else {
return None;
};
if n_args as usize == params.len() - 1 {
return params.last();
return Ok(None);
}
};
if (n_args as usize) + 1 == params.len() {
return Ok(params.last());
}
}
Ok(None)
}
None
pub fn get_extra_arg<'a>(
expr: &'a Expr,
module: Option<&str>,
functions: &FunctionTable,
) -> Option<&'a Expr> {
match get_extra_arg_impl(expr, module, functions) {
Ok(a) => a,
_ => None,
}
}
pub fn gather_functions<'a>(modules: &[&'a Module]) -> Result<FunctionTable<'a>> {

View File

@@ -190,6 +190,11 @@ impl Value {
Value::Number(Number(OrderedFloat(v)))
}
pub fn from_u128(v: u128) -> Value {
// TODO: fix precision loss
Value::Number(Number(OrderedFloat(v as f64)))
}
pub fn from_array(a: Vec<Value>) -> Value {
Value::Array(Rc::new(a))
}

View File

@@ -39,7 +39,7 @@ cases:
package test
x = array.reverse()
query: data.test
error: "at least one argument required for function calls"
error: "expects 1 argument"
- note: reverse-more-args
data: {}

View File

@@ -72,7 +72,7 @@ fn run_opa_tests() -> Result<()> {
return Ok(());
}
};
dbg!(&opa_tests_dir);
let tests_path = Path::new(&opa_tests_dir);
let mut status = BTreeMap::<String, (u32, u32)>::new();
let mut n = 0;
@@ -152,14 +152,18 @@ fn run_opa_tests() -> Result<()> {
}
println!("TESTSUITE STATUS");
println!(" {:30} {:4} {:4}", "FOLDER", "PASS", "FAIL");
println!(" {:40} {:4} {:4}", "FOLDER", "PASS", "FAIL");
let (mut npass, mut nfail) = (0, 0);
for (dir, (pass, fail)) in status {
if fail == 0 {
println!("\x1b[32m {dir:40}: {pass:4} {fail:4}\x1b[0m");
} else {
println!("\x1b[31m {dir:40}: {pass:4} {fail:4}\x1b[0m");
}
npass += pass;
nfail += fail;
}
println!();
println!("\x1b[31m {:40}: {npass:4} {nfail:4}\x1b[0m", "TOTAL");
Ok(())
}