Minimize PR 22 (microsoft#26)

* specific functions added to eval different rego components Signed-off-by: eric-therond <[email protected]> * allow multiple inputs and results Signed-off-by: eric-therond <[email protected]> * prepare_for_eval is necessary to be called Signed-off-by: eric-therond <[email protected]> * test with the suggested code examples and clean scopes Signed-off-by: eric-therond <[email protected]> * try to refactor first steps of evaluations Signed-off-by: eric-therond <[email protected]> * improve coverage and fix clean state internal evaluation Signed-off-by: eric-therond <[email protected]> * add getters and setters and fix clean function Signed-off-by: eric-therond <[email protected]> * Tests are single input by default. Multi input specified via "many!" marker. Signed-off-by: Anand Krishnamoorthi <[email protected]> --------- Signed-off-by: eric-therond <[email protected]> Signed-off-by: Anand Krishnamoorthi <[email protected]> Co-authored-by: eric-therond <[email protected]>
anakrish · Oct 9, 2023 · 2436467 · 2436467
1 parent 2ba718b
commit 2436467
Show file tree

Hide file tree

Showing 11 changed files with 430 additions and 66 deletions.
diff --git a/src/interpreter.rs b/src/interpreter.rs
@@ -20,8 +20,10 @@ pub struct Interpreter<'source> {
     module: Option<&'source Module<'source>>,
     schedule: Option<&'source Schedule<'source>>,
     current_module_path: String,
+    prepared: bool,
     input: Value,
     data: Value,
+    init_data: Value,
     scopes: Vec<Scope>,
     // TODO: handle recursive calls where same expr could have different values.
     loop_var_values: BTreeMap<&'source Expr<'source>, Value>,
@@ -58,8 +60,10 @@ impl<'source> Interpreter<'source> {
             module: None,
             schedule: None,
             current_module_path: String::default(),
+            prepared: false,
             input: Value::new_object(),
             data: Value::new_object(),
+            init_data: Value::new_object(),
             scopes: vec![Scope::new()],
             contexts: vec![],
             loop_var_values: BTreeMap::new(),
@@ -74,6 +78,44 @@ impl<'source> Interpreter<'source> {
         })
     }
 
+    pub fn get_modules(&mut self) -> &mut Vec<&'source Module<'source>> {
+        &mut self.modules
+    }
+
+    pub fn set_data(&mut self, data: Value) {
+        self.data = data;
+    }
+
+    pub fn get_data(&mut self) -> &mut Value {
+        &mut self.data
+    }
+
+    fn clean_internal_evaluation_state(&mut self) {
+        self.data = self.init_data.clone();
+        self.processed.clear();
+        self.loop_var_values.clear();
+        self.scopes = vec![Scope::new()];
+        self.contexts = vec![];
+    }
+
+    fn checks_for_eval(&mut self, input: &Option<Value>, enable_tracing: bool) -> Result<()> {
+        if !self.prepared {
+            bail!("prepare_for_eval should be called before eval_modules");
+        }
+
+        self.traces = match enable_tracing {
+            true => Some(vec![]),
+            false => None,
+        };
+
+        if let Some(input) = input {
+            self.input = input.clone();
+            info!("input: {:#?}", self.input);
+        }
+
+        Ok(())
+    }
+
     fn current_module(&self) -> Result<&'source Module<'source>> {
         self.module
             .ok_or_else(|| anyhow!("internal error: current module not set"))
@@ -1435,6 +1477,7 @@ impl<'source> Interpreter<'source> {
             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> =
@@ -1933,62 +1976,65 @@ impl<'source> Interpreter<'source> {
                 head: rule_head,
                 bodies: rule_body,
             } => {
-                if matches!(rule_head, RuleHead::Func { .. }) {
-                    return Ok(());
-                }
-
-                let (ctx, mut path) = self.make_rule_context(rule_head)?;
-                let special_set = matches!((ctx.output_expr, &ctx.value), (None, Value::Set(_)));
-                let value = match self.eval_rule_bodies(ctx, span, rule_body)? {
-                    Value::Set(_) if special_set => {
-                        let entry = path[path.len() - 1].text();
-                        let mut s = BTreeSet::new();
-                        s.insert(Value::String(entry.to_owned()));
-                        path = path[0..path.len() - 1].to_vec();
-                        Value::from_set(s)
+                if !matches!(rule_head, RuleHead::Func { .. }) {
+                    let (ctx, mut path) = self.make_rule_context(rule_head)?;
+                    let special_set =
+                        matches!((ctx.output_expr, &ctx.value), (None, Value::Set(_)));
+                    let value = match self.eval_rule_bodies(ctx, span, rule_body)? {
+                        Value::Set(_) if special_set => {
+                            let entry = path[path.len() - 1].text();
+                            let mut s = BTreeSet::new();
+                            s.insert(Value::String(entry.to_owned()));
+                            path = path[0..path.len() - 1].to_vec();
+                            Value::from_set(s)
+                        }
+                        v => v,
+                    };
+                    if value != Value::Undefined {
+                        let paths: Vec<&str> = path.iter().map(|s| s.text()).collect();
+                        let vref = Self::make_or_get_value_mut(&mut self.data, &paths[..])?;
+                        Self::merge_value(span, vref, value)?;
                     }
-                    v => v,
-                };
 
-                if value != Value::Undefined {
-                    let paths: Vec<&str> = path.iter().map(|s| s.text()).collect();
-                    let vref = Self::make_or_get_value_mut(&mut self.data, &paths[..])?;
-                    Self::merge_value(span, vref, value)?;
+                    self.processed.insert(rule);
                 }
             }
             _ => bail!("internal error: unexpected"),
         }
         self.set_current_module(prev_module)?;
-        self.processed.insert(rule);
         match self.active_rules.pop() {
             Some(r) if r == rule => Ok(()),
             _ => bail!("internal error: current rule not active"),
         }
     }
 
-    pub fn eval(
+    pub fn eval_rule_with_input(
         &mut self,
-        data: &Option<Value>,
+        module: &'source Module<'source>,
+        rule: &'source Rule<'source>,
         input: &Option<Value>,
         enable_tracing: bool,
-        schedule: Option<&'source Schedule<'source>>,
     ) -> Result<Value> {
-        self.schedule = schedule;
-        self.traces = match enable_tracing {
-            true => Some(vec![]),
-            false => None,
-        };
+        self.checks_for_eval(input, enable_tracing)?;
+        self.clean_internal_evaluation_state();
 
-        self.builtins_cache.clear();
+        self.eval_rule(module, rule)?;
 
-        if let Some(input) = input {
-            self.input = input.clone();
+        Ok(self.data.clone())
+    }
+
+    pub fn prepare_for_eval(
+        &mut self,
+        schedule: Option<&'source Schedule<'source>>,
+        data: &Option<Value>,
+    ) -> Result<()> {
+        self.schedule = schedule;
+        self.builtins_cache.clear();
 
-            info!("input: {:#?}", self.input);
-        }
         if let Some(data) = data {
             self.data = data.clone();
         }
+
         // Ensure that each module has an empty object
         for m in &self.modules {
             let path = Parser::get_path_ref_components(&m.package.refr)?;
@@ -2003,6 +2049,39 @@ impl<'source> Interpreter<'source> {
         self.update_function_table()?;
         self.gather_rules()?;
 
+        self.init_data = self.data.clone();
+        self.prepared = true;
+
+        Ok(())
+    }
+
+    pub fn eval_module(
+        &mut self,
+        module: &'source Module<'source>,
+        input: &Option<Value>,
+        enable_tracing: bool,
+    ) -> Result<Value> {
+        self.checks_for_eval(input, enable_tracing)?;
+        self.clean_internal_evaluation_state();
+
+        for rule in &module.policy {
+            self.eval_rule(module, rule)?;
+        }
+
+        // Defer the evaluation of the default rules to here
+        let prev_module = self.set_current_module(Some(module))?;
+        for rule in &module.policy {
+            self.eval_default_rule(rule)?;
+        }
+        self.set_current_module(prev_module)?;
+
+        Ok(self.data.clone())
+    }
+
+    pub fn eval_modules(&mut self, input: &Option<Value>, enable_tracing: bool) -> Result<Value> {
+        self.checks_for_eval(input, enable_tracing)?;
+        self.clean_internal_evaluation_state();
+
         for module in self.modules.clone() {
             for rule in &module.policy {
                 self.eval_rule(module, rule)?;
@@ -2021,6 +2100,17 @@ impl<'source> Interpreter<'source> {
         Ok(self.data.clone())
     }
 
+    pub fn eval(
+        &mut self,
+        data: &Option<Value>,
+        input: &Option<Value>,
+        enable_tracing: bool,
+        schedule: Option<&'source Schedule<'source>>,
+    ) -> Result<Value> {
+        self.prepare_for_eval(schedule, data)?;
+        self.eval_modules(input, enable_tracing)
+    }
+
     pub fn eval_query_snippet(
         &mut self,
         snippet: &'source Expr<'source>,

diff --git a/tests/interpreter/cases/arithmetic/mod.rs b/tests/interpreter/cases/arithmetic/mod.rs
@@ -29,14 +29,14 @@ fn basic() -> Result<()> {
     }
 "#;
 
-    let expected = Value::from_json_str(
+    let expected = vec![Value::from_json_str(
         r#" {
     "add" : true,
     "sub" : true,
     "mul" : true,
     "div" : true
 }"#,
-    )?;
+    )?];
 
     assert_eq!(
         eval_file(&[rego.to_owned()], None, None, "data.test", false)?,

diff --git a/tests/interpreter/cases/compr/mod.rs b/tests/interpreter/cases/compr/mod.rs
@@ -32,7 +32,7 @@ fn basic_array() -> Result<()> {
     array_compr_7 = [ 1 | [1, 2, 3][_]; [1, 2][_] >= 2 ]
 "#;
 
-    let expected = Value::from_json_str(
+    let expected = vec![Value::from_json_str(
         r#" {
             "array": [1, 2, 3],
             "array_compr_0": [1],
@@ -44,7 +44,7 @@ fn basic_array() -> Result<()> {
             "array_compr_6": [1, 1, 1, 1, 1, 1],
             "array_compr_7": [1, 1, 1]
 }"#,
-    )?;
+    )?];
 
     assert_match(
         eval_file(&[rego.to_owned()], None, None, "data.test", false)?,
@@ -81,7 +81,7 @@ fn basic_set() -> Result<()> {
     set_compr_7 = { a | a = [1, 2, 3][_]; [1, 2][_] >= 2 }
 "#;
 
-    let expected = Value::from_json_str(
+    let expected = vec![Value::from_json_str(
         r#" {
             "set": {
                 "set!": [1, "string", [2, 3, 4], 567, false]
@@ -113,7 +113,7 @@ fn basic_set() -> Result<()> {
                 "set!": [1, 2, 3]
             }
  }"#,
-    )?;
+    )?];
 
     assert_match(
         eval_file(&[rego.to_owned()], None, None, "data.test", false)?,

diff --git a/tests/interpreter/cases/compr/object.yaml b/tests/interpreter/cases/compr/object.yaml
@@ -64,4 +64,3 @@ cases:
         x = { k:v | k = ["Hello", "world", 1][_]; v = [1, 2][_] }
     query: data.test
     error: "value for key `\"Hello\"` generated multiple times: `1` and `2`"
-    want_result:
diff --git a/tests/interpreter/cases/in/mod.rs b/tests/interpreter/cases/in/mod.rs
@@ -99,7 +99,7 @@ fn basic() -> Result<()> {
     }
 "#;
 
-    let expected = Value::from_json_str(
+    let expected = vec![Value::from_json_str(
         r#" {
             "array": [1, 2, 3],
             "in_array_key_value": true,
@@ -117,7 +117,7 @@ fn basic() -> Result<()> {
             "in_set_value": true,
             "some_decl_set_value": true
 }"#,
-    )?;
+    )?];
 
     assert_match(
         eval_file(&[rego.to_owned()], None, None, "data.test", false)?,

diff --git a/tests/interpreter/cases/input/mod.rs b/tests/interpreter/cases/input/mod.rs
@@ -0,0 +1,48 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+#![cfg(test)]
+
+use crate::interpreter::*;
+use anyhow::Result;
+
+#[test]
+fn basic() -> Result<()> {
+    let rego = r#"
+    package test
+        
+    x[a] {
+      a = y
+    }
+
+    y[a] {
+      a = input.x + 5
+    }
+"#;
+
+    let input = ValueOrVec::Many(vec![
+        Value::from_json_str(r#"{"x": 1}"#)?,
+        Value::from_json_str(r#"{"x": 6}"#)?,
+    ]);
+
+    let expected = vec![
+        Value::from_json_str(
+            r#" {
+        "y": {"set!": [6]},
+        "x": {"set!": [{"set!":[6]}]}
+}"#,
+        )?,
+        Value::from_json_str(
+            r#" {
+        "y": {"set!": [11]},
+        "x": {"set!": [{"set!":[11]}]}
+}"#,
+        )?,
+    ];
+
+    assert_match(
+        eval_file_first_rule(&[rego.to_owned()], None, Some(input), "data.test", false)?,
+        expected,
+    );
+    Ok(())
+}
diff --git a/tests/interpreter/cases/input/multiple.yaml b/tests/interpreter/cases/input/multiple.yaml
@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: input-multiple-1
+    data: {}
+    input:
+      many!: 
+        - { x: 1 }
+        - { x: 5 }
+    modules:
+      - |
+        package test
+        
+        x[a] {
+          a = y
+        }
+
+        y[a] {
+          a = input.x + 5
+        }
+
+    query: data.test
+    want_result:
+      many!:
+        - y: 
+            set!: [6]
+          x: 
+            set!: [ set!: [6] ]
+        - y: 
+            set!: [10]
+          x: 
+            set!: [ set!: [10] ]