From b470c3fb7fe1ae01aa664b80ab4a72b38dbb9dc3 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi <35780660+anakrish@users.noreply.github.com> Date: Mon, 1 Jan 2024 15:16:57 -0800 Subject: [PATCH] Handle walk builtin as a loop expression (#86) The walk builtin generates values and implicitly creates a loop over the values. Hoist walk calls as loops and handle them. Also handle cases where return value is bound to an extra parameter. Closes #83 --- src/interpreter.rs | 211 +++++++++++++++++++++++++++++++-------------- tests/opa.passing | 4 +- tests/opa.rs | 20 ++++- 3 files changed, 168 insertions(+), 67 deletions(-) diff --git a/src/interpreter.rs b/src/interpreter.rs index bf55d48..2829a0f 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -102,11 +102,47 @@ struct Context { } #[derive(Debug)] -struct LoopExpr { - span: Span, - expr: ExprRef, - value: ExprRef, - index: Ref, +enum LoopExpr { + Loop { + span: Span, + expr: Ref, + value: Ref, + index: Ref, + }, + Walk { + span: Span, + expr: Ref, + }, +} + +impl LoopExpr { + fn span(&self) -> Span { + match self { + Self::Loop { span, .. } => span.clone(), + Self::Walk { span, .. } => span.clone(), + } + } + + fn value(&self) -> Ref { + match self { + Self::Loop { value, .. } => value.clone(), + Self::Walk { expr, .. } => expr.clone(), + } + } + + fn expr(&self) -> Ref { + match self { + Self::Loop { expr, .. } => expr.clone(), + Self::Walk { expr, .. } => expr.clone(), + } + } + + fn index(&self) -> Option> { + match self { + Self::Loop { index, .. } => Some(index.clone()), + Self::Walk { .. } => None, + } + } } impl Interpreter { @@ -364,7 +400,7 @@ impl Interpreter { _ => Ok(false), }); if !indices.is_empty() { - loops.push(LoopExpr { + loops.push(LoopExpr::Loop { span: span.clone(), expr: expr.clone(), value: refr.clone(), @@ -381,6 +417,20 @@ impl Interpreter { for item in items { self.hoist_loops_impl(item, loops); } + + // Handle walk builtin which acts as a generator. + // TODO: Handle with modifier on the walk builtin. + if let Expr::Call { fcn, .. } = expr.as_ref() { + if let Ok(fcn_path) = get_path_string(fcn, None) { + if fcn_path == "walk" { + // TODO: Use an enum for LoopExpr to handle walk + loops.push(LoopExpr::Walk { + span: expr.span().clone(), + expr: expr.clone(), + }) + } + } + } } Object { fields, .. } => { @@ -998,6 +1048,7 @@ impl Interpreter { let value = match expr.as_ref() { Expr::Call { span, fcn, params } => self.eval_call( span, + expr, fcn, params, get_extra_arg( @@ -1037,6 +1088,7 @@ impl Interpreter { // Extra parameter is allowed; but a return argument is not allowed. Expr::Call { span, fcn, params } => self.eval_call( span, + expr, fcn, params, get_extra_arg( @@ -1230,12 +1282,31 @@ impl Interpreter { let loop_expr = &loops[0]; let mut result = false; - let loop_expr_value = self.eval_expr(&loop_expr.value)?; + let loop_expr_value = loop_expr.value(); + let loop_expr_value = if let Expr::Call { span, fcn, params } = loop_expr_value.as_ref() + { + // Handle walk(obj, output_param) + let extra_arg = get_extra_arg( + &loop_expr_value, + Some(self.current_module_path.as_str()), + &self.functions, + ); + // If there is an extra arg, ignore it while computing the loop value. + let params = if extra_arg.is_some() { + ¶ms[..params.len() - 1] + } else { + ¶ms[..] + }; + self.eval_call_impl(span, &loop_expr_value, fcn, params)? + } else { + self.eval_expr(&loop_expr_value)? + }; // If the loop's index variable has already been assigned a value // (this can happen if the same index is used for two different collections), // then evaluate statements only if the index applies to this collection. - if let Expr::Var(index_var) = loop_expr.index.as_ref() { + let loop_expr_index = loop_expr.index(); + if let Some(Expr::Var(index_var)) = loop_expr_index.as_ref().map(|r| r.as_ref()) { if let Some(idx) = self.lookup_local_var(&index_var.source_str()) { if loop_expr_value[&idx] != Value::Undefined { result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result; @@ -1255,23 +1326,27 @@ impl Interpreter { match loop_expr_value { Value::Array(items) => { for (idx, v) in items.iter().enumerate() { - self.loop_var_values - .insert(loop_expr.expr.clone(), v.clone()); + self.loop_var_values.insert(loop_expr.expr(), v.clone()); - let mut type_match = BTreeSet::new(); - let mut cache = BTreeMap::new(); - if self.make_bindings( - false, - &mut type_match, - &mut cache, - &loop_expr.index, - &Value::from(idx), - true, - )? { + let exec = if let Some(index) = loop_expr.index() { + let mut type_match = BTreeSet::new(); + let mut cache = BTreeMap::new(); + self.make_bindings( + false, + &mut type_match, + &mut cache, + &index, + &Value::from(idx), + true, + )? + } else { + true + }; + if exec { result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result; } - self.loop_var_values.remove(&loop_expr.expr); + self.loop_var_values.remove(&loop_expr.expr()); *self.current_scope_mut()? = scope_saved.clone(); if let Some(ctx) = self.contexts.last_mut() { ctx.result = query_result.clone(); @@ -1280,24 +1355,21 @@ impl Interpreter { } Value::Set(items) => { for v in items.iter() { - self.loop_var_values - .insert(loop_expr.expr.clone(), v.clone()); + self.loop_var_values.insert(loop_expr.expr(), v.clone()); // For sets, index is also the value. - let mut type_match = BTreeSet::new(); - let mut cache = BTreeMap::new(); - if self.make_bindings( - false, - &mut type_match, - &mut cache, - &loop_expr.index, - v, - true, - )? { + let exec = if let Some(index) = loop_expr.index() { + let mut type_match = BTreeSet::new(); + let mut cache = BTreeMap::new(); + self.make_bindings(false, &mut type_match, &mut cache, &index, v, true)? + } else { + true + }; + if exec { result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result; } - self.loop_var_values.remove(&loop_expr.expr); + self.loop_var_values.remove(&loop_expr.expr()); *self.current_scope_mut()? = scope_saved.clone(); if let Some(ctx) = self.contexts.last_mut() { ctx.result = query_result.clone(); @@ -1306,22 +1378,19 @@ impl Interpreter { } Value::Object(obj) => { for (k, v) in obj.iter() { - self.loop_var_values - .insert(loop_expr.expr.clone(), v.clone()); + self.loop_var_values.insert(loop_expr.expr(), v.clone()); // For objects, index is key. - let mut type_match = BTreeSet::new(); - let mut cache = BTreeMap::new(); - if self.make_bindings( - false, - &mut type_match, - &mut cache, - &loop_expr.index, - k, - true, - )? { + let exec = if let Some(index) = loop_expr.index() { + let mut type_match = BTreeSet::new(); + let mut cache = BTreeMap::new(); + self.make_bindings(false, &mut type_match, &mut cache, &index, k, true)? + } else { + true + }; + if exec { result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result; } - self.loop_var_values.remove(&loop_expr.expr); + self.loop_var_values.remove(&loop_expr.expr()); *self.current_scope_mut()? = scope_saved.clone(); if let Some(ctx) = self.contexts.last_mut() { ctx.result = query_result.clone(); @@ -1431,37 +1500,34 @@ impl Interpreter { // Try out values in current loop expr. let loop_expr = &loops[0]; let mut result = false; - match self.eval_expr(&loop_expr.value)? { + match self.eval_expr(&loop_expr.value())? { Value::Array(items) => { for v in items.iter() { - self.loop_var_values - .insert(loop_expr.expr.clone(), v.clone()); + self.loop_var_values.insert(loop_expr.expr(), v.clone()); result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } Value::Set(items) => { for v in items.iter() { - self.loop_var_values - .insert(loop_expr.expr.clone(), v.clone()); + self.loop_var_values.insert(loop_expr.expr(), v.clone()); result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } Value::Object(obj) => { for (_, v) in obj.iter() { - self.loop_var_values - .insert(loop_expr.expr.clone(), v.clone()); + self.loop_var_values.insert(loop_expr.expr(), v.clone()); result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } _ => { - return Err(loop_expr.span.source.error( - loop_expr.span.line, - loop_expr.span.col, + return Err(loop_expr.span().source.error( + loop_expr.span().line, + loop_expr.span().col, "item cannot be indexed", )); } } - self.loop_var_values.remove(&loop_expr.expr); + self.loop_var_values.remove(&loop_expr.expr()); Ok(result) } @@ -1791,7 +1857,18 @@ impl Interpreter { Ok(None) } - fn eval_call_impl(&mut self, span: &Span, fcn: &ExprRef, params: &[ExprRef]) -> Result { + fn eval_call_impl( + &mut self, + span: &Span, + expr: &ExprRef, + fcn: &ExprRef, + params: &[ExprRef], + ) -> Result { + // Return generated values of walk builtin. + if let Some(v) = self.loop_var_values.get(expr) { + return Ok(v.clone()); + } + let fcn_path = match get_path_string(fcn, None) { Ok(p) => p, _ => bail!(span.error("invalid function expression")), @@ -1994,6 +2071,7 @@ impl Interpreter { fn eval_call( &mut self, span: &Span, + expr: &ExprRef, fcn: &ExprRef, params: &[ExprRef], extra_arg: Option, @@ -2005,14 +2083,16 @@ impl Interpreter { Expr::Var(var) if allow_return_arg && self.lookup_local_var(&var.source_str()).is_none() => { - let value = self.eval_call_impl(span, fcn, ¶ms[..params.len() - 1])?; + let value = + self.eval_call_impl(span, expr, fcn, ¶ms[..params.len() - 1])?; if var.text() != "_" { self.add_variable(&var.source_str(), value)?; } Ok(Value::Bool(true)) } _ if allow_return_arg => { - let ret_value = self.eval_call_impl(span, fcn, ¶ms[..params.len() - 1])?; + let ret_value = + self.eval_call_impl(span, expr, fcn, ¶ms[..params.len() - 1])?; let mut cache = BTreeMap::new(); let mut type_match = BTreeSet::new(); self.make_bindings(false, &mut type_match, &mut cache, &ea, &ret_value, false) @@ -2020,12 +2100,13 @@ impl Interpreter { } _ => { let expected = self.eval_expr(¶ms[params.len() - 1])?; - let ret_value = self.eval_call_impl(span, fcn, ¶ms[..params.len() - 1])?; + let ret_value = + self.eval_call_impl(span, expr, fcn, ¶ms[..params.len() - 1])?; Ok(Value::Bool(ret_value == expected)) } } } else { - self.eval_call_impl(span, fcn, params) + self.eval_call_impl(span, expr, fcn, params) } } @@ -2254,7 +2335,9 @@ impl Interpreter { } => self.eval_object_compr(key, value, query), Expr::SetCompr { term, query, .. } => self.eval_set_compr(term, query), Expr::UnaryExpr { .. } => unimplemented!("unar expr is umplemented"), - Expr::Call { span, fcn, params } => self.eval_call(span, fcn, params, None, false), + Expr::Call { span, fcn, params } => { + self.eval_call(span, expr, fcn, params, None, false) + } } } diff --git a/tests/opa.passing b/tests/opa.passing index 7b11c7e..97ed1b7 100644 --- a/tests/opa.passing +++ b/tests/opa.passing @@ -76,6 +76,7 @@ partialobjectdoc partialsetdoc planner-ir rand +reachable regexfind regexfindallstringsubmatch regexisvalid @@ -108,4 +109,5 @@ units urlbuiltins uuid varreferences -virtualdocs \ No newline at end of file +virtualdocs +walkbuiltin \ No newline at end of file diff --git a/tests/opa.rs b/tests/opa.rs index d68a135..05d3599 100644 --- a/tests/opa.rs +++ b/tests/opa.rs @@ -149,11 +149,27 @@ fn run_opa_tests(opa_tests_dir: String, folders: &[String]) -> Result<()> { let yaml_str = std::fs::read_to_string(&path_str)?; let test: YamlTest = serde_yaml::from_str(&yaml_str)?; - for case in &test.cases { + for mut case in test.cases { let is_json_schema_test = case.note.starts_with("json_verify_schema") || case.note.starts_with("json_match_schema"); - match (eval_test_case(case), &case.want_result) { + if case.note == "reachable_paths/cycle_1022_3" { + // The OPA behavior is not well-defined. + // See: https://github.com/open-policy-agent/opa/issues/5871 + // https://github.com/open-policy-agent/opa/issues/6128 + // We lock down all the paths leading to leaf nodes instead. + case.want_result = serde_json::from_str( + r#" [{ + "x" : [ + ["one", "five", "seven", "eight", "three"], + ["one", "five", "six", "nine"], + ["one", "five", "six", "seven", "eight", "three"], + ["one", "two", "four", "three"] + ] + }]"#, + )?; + }; + match (eval_test_case(&case), &case.want_result) { (Ok(actual), Some(expected)) if is_json_schema_test && json_schema_tests_check(&actual, &expected) => {