Repository: microsoft/regorus (language: python)

Issue:
[Lib] graph builtins
- [x] [graph.reachable](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-graph-graphreachable)
- [x] [graph.reachable_paths](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-graph-graphreachable_paths)
- [x] [walk](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-graph-walk)

Note: graph.reachable_paths has bugs in upstream OPA:
- https://github.com/open-policy-agent/opa/issues/5871
- https://github.com/open-policy-agent/opa/issues/6128

Get `cargo test -r --test opa -- reachable walkbuiltin` to pass.


Discussion / hints:

Fix patch:
diff --git a/src/interpreter.rs b/src/interpreter.rs
index bf55d483..2829a0fd 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<Expr>,
+enum LoopExpr {
+    Loop {
+        span: Span,
+        expr: Ref<Expr>,
+        value: Ref<Expr>,
+        index: Ref<Expr>,
+    },
+    Walk {
+        span: Span,
+        expr: Ref<Expr>,
+    },
+}
+
+impl LoopExpr {
+    fn span(&self) -> Span {
+        match self {
+            Self::Loop { span, .. } => span.clone(),
+            Self::Walk { span, .. } => span.clone(),
+        }
+    }
+
+    fn value(&self) -> Ref<Expr> {
+        match self {
+            Self::Loop { value, .. } => value.clone(),
+            Self::Walk { expr, .. } => expr.clone(),
+        }
+    }
+
+    fn expr(&self) -> Ref<Expr> {
+        match self {
+            Self::Loop { expr, .. } => expr.clone(),
+            Self::Walk { expr, .. } => expr.clone(),
+        }
+    }
+
+    fn index(&self) -> Option<Ref<Expr>> {
+        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() {
+                    &params[..params.len() - 1]
+                } else {
+                    &params[..]
+                };
+                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());
-
-                        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,
-                        )? {
+                        self.loop_var_values.insert(loop_expr.expr(), v.clone());
+
+                        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<Value> {
+    fn eval_call_impl(
+        &mut self,
+        span: &Span,
+        expr: &ExprRef,
+        fcn: &ExprRef,
+        params: &[ExprRef],
+    ) -> Result<Value> {
+        // 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<ExprRef>,
@@ -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, &params[..params.len() - 1])?;
+                    let value =
+                        self.eval_call_impl(span, expr, fcn, &params[..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, &params[..params.len() - 1])?;
+                    let ret_value =
+                        self.eval_call_impl(span, expr, fcn, &params[..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(&params[params.len() - 1])?;
-                    let ret_value = self.eval_call_impl(span, fcn, &params[..params.len() - 1])?;
+                    let ret_value =
+                        self.eval_call_impl(span, expr, fcn, &params[..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)
+            }
         }
     }
 


Test patch:
diff --git a/tests/opa.passing b/tests/opa.passing
index 7b11c7ef..97ed1b72 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 d68a1355..05d3599a 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) =>
                 {