More OPA conformance (#77)

- object.union
- object.union_n
- treat negative integers as two separate tokens (Sub and Number)
  when seen in arithmetic expressions
- Ensure that fully query string is parsed
- Evaluate queries in a separate module instead of the last read module.
  This correctly handles queries of the form `x = data.test.y` where x is
  already a ref in `data.test`
- Handle queries producing multiple outputs in test infrastructure
- Add tests for engine
- Add tests locking down valid queries
- Update opa.passing

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

src/builtins/objects.rs

@@ -3,7 +3,7 @@
 
 use crate::ast::{Expr, Ref};
 use crate::builtins;
-use crate::builtins::utils::{ensure_args_count, ensure_object};
+use crate::builtins::utils::{ensure_args_count, ensure_array, ensure_object};
 use crate::lexer::Span;
 use crate::value::Value;
 
@@ -21,6 +21,8 @@ pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {
     m.insert("object.keys", (keys, 1));
     m.insert("object.remove", (remove, 2));
     m.insert("object.subset", (subset, 2));
+    m.insert("object.union", (object_union, 2));
+    m.insert("object.union_n", (object_union_n, 1));
 }
 
 fn json_filter_impl(v: &Value, filter: &Value) -> Value {
@@ -324,3 +326,59 @@ fn subset(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool) -> R
 
     Ok(Value::Bool(is_subset(&args[0], &args[1])))
 }
+
+fn union(obj1: &Value, obj2: &Value) -> Result<Value> {
+    match (obj1, obj2) {
+        (Value::Object(m1), Value::Object(m2)) => {
+            let mut u = obj1.clone();
+            let um = u.as_object_mut()?;
+
+            for (key2, value2) in m2.iter() {
+                let vm = match m1.get(key2) {
+                    Some(value1) => union(value1, value2)?,
+                    _ => value2.clone(),
+                };
+                um.insert(key2.clone(), vm);
+            }
+            Ok(u)
+        }
+        _ => Ok(obj2.clone()),
+    }
+}
+
+fn object_union(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool) -> Result<Value> {
+    let name = "object.union";
+    ensure_args_count(span, name, params, args, 2)?;
+
+    let _ = ensure_object(name, &params[0], args[0].clone())?;
+    let _ = ensure_object(name, &params[1], args[1].clone())?;
+
+    union(&args[0], &args[1])
+}
+
+fn object_union_n(
+    span: &Span,
+    params: &[Ref<Expr>],
+    args: &[Value],
+    strict: bool,
+) -> Result<Value> {
+    let name = "object.union_n";
+    ensure_args_count(span, name, params, args, 1)?;
+
+    let arr = ensure_array(name, &params[0], args[0].clone())?;
+
+    let mut u = Value::new_object();
+    for (idx, a) in arr.iter().enumerate() {
+        if a.as_object().is_err() {
+            if strict {
+                bail!(params[0]
+                    .span()
+                    .error(&format!("item at index {idx} is not an object")));
+            }
+            return Ok(Value::Undefined);
+        }
+        u = union(&u, a)?;
+    }
+
+    Ok(u)
+}

src/engine.rs

@@ -164,23 +164,26 @@ impl Engine {
     pub fn eval_query(&mut self, query: String, enable_tracing: bool) -> Result<QueryResults> {
         self.eval_modules(false)?;
 
-        // Parse the query.
-        let query_len = query.len();
-        let query_source = Source::new("<query.rego>".to_string(), query);
-        let query_span = Span {
-            source: query_source.clone(),
-            line: 1,
-            col: 1,
-            start: 0,
-            end: query_len as u16,
+        let query_module = {
+            let source = Source::new(
+                "<query_module.rego>".to_owned(),
+                "package __internal_query_module".to_owned(),
+            );
+            Ref::new(Parser::new(&source)?.parse()?)
         };
+
+        // Parse the query.
+        let query_source = Source::new("<query.rego>".to_string(), query);
         let mut parser = Parser::new(&query_source)?;
-        let query_node = Ref::new(parser.parse_query(query_span, "")?);
+        let query_node = parser.parse_user_query()?;
         let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;
 
-        let results =
-            self.interpreter
-                .eval_user_query(&query_node, &query_schedule, enable_tracing)?;
+        let results = self.interpreter.eval_user_query(
+            &query_module,
+            &query_node,
+            &query_schedule,
+            enable_tracing,
+        )?;
         Ok(results)
     }
 }

src/interpreter.rs

@@ -2727,6 +2727,7 @@ impl Interpreter {
 
     pub fn eval_user_query(
         &mut self,
+        module: &Ref<Module>,
         query: &Ref<Query>,
         schedule: &Schedule,
         enable_tracing: bool,
@@ -2754,7 +2755,7 @@ impl Interpreter {
             is_compr: false,
         });
 
-        let prev_module = self.set_current_module(self.modules.last().cloned())?;
+        let prev_module = self.set_current_module(Some(module.clone()))?;
 
         // Eval the query.
         let query_r = self.eval_query(query);
@@ -2795,6 +2796,14 @@ impl Interpreter {
 
         self.set_current_module(prev_module)?;
 
+        if let Some(r) = results.result.last() {
+            if r.bindings.is_empty_object()
+                && r.expressions.iter().any(|e| e.value == Value::Bool(false))
+            {
+                results = QueryResults::default();
+            }
+        }
+
         match query_r {
             Ok(_) => Ok(results),
             Err(e) => Err(e),

src/parser.rs

@@ -620,10 +620,21 @@ impl<'source> Parser<'source> {
             let op = match self.token_text() {
                 "+" => ArithOp::Add,
                 "-" => ArithOp::Sub,
+                n if n.starts_with('-') && self.tok.0 == TokenKind::Number => ArithOp::Sub,
                 _ => return Ok(expr),
             };
-            self.next_token()?;
-            let right = self.parse_mul_div_mod_expr()?;
+            let right = if self.token_text().len() > 1 {
+                // Treat the - as a separate token
+                let mut rhs_span = self.tok.1.clone();
+                rhs_span.start += 1;
+                rhs_span.col += 1;
+
+                self.next_token()?;
+                Expr::Number(rhs_span)
+            } else {
+                self.next_token()?;
+                self.parse_mul_div_mod_expr()?
+            };
             span.end = self.end;
             expr = Expr::ArithExpr {
                 span,
@@ -953,7 +964,7 @@ impl<'source> Parser<'source> {
         })
     }
 
-    pub fn parse_query(&mut self, mut span: Span, end_delim: &str) -> Result<Query> {
+    fn parse_query(&mut self, mut span: Span, end_delim: &str) -> Result<Query> {
         let state = self.clone();
         let is_definite_query = matches!(self.token_text(), "some" | "every");
 
@@ -1485,7 +1496,7 @@ impl<'source> Parser<'source> {
         Ok(Rule::Spec { span, head, bodies })
     }
 
-    fn parse_package(&mut self) -> Result<Package> {
+    pub fn parse_package(&mut self) -> Result<Package> {
         let mut span = self.tok.1.clone();
         self.expect("package", "Missing package declaration.")?;
         let name = self.parse_path_ref()?;
@@ -1609,4 +1620,13 @@ impl<'source> Parser<'source> {
             policy,
         })
     }
+
+    pub fn parse_user_query(&mut self) -> Result<Ref<Query>> {
+        let span = self.tok.1.clone();
+        let query = Ref::new(self.parse_query(span, "")?);
+        if self.tok.0 != TokenKind::Eof {
+            bail!(self.tok.1.error("expecting EOF"));
+        }
+        Ok(query)
+    }
 }