fix: Merge data to init document (#293)

Init document is the aggregated data documen that the user has
specified using multiple `add_data` calls. Each query evaluation
starts of by initializing the current data to the init document.

Previously `add_data` was incorrectly added to the current document,
causing the added data to be lost if the addition happened after query
evaluation.

With this fix, scenarios where data addition may be interspersed with
query evaluation calls are supported.

Also provide a get_data method to obtain the (init) data document.

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

tests/kata/main.rs

@@ -172,6 +172,64 @@ struct Cli {
     generate: bool,
 }
 
+fn stateful_policy_test() -> Result<()> {
+    // Create an engine for evaluating Rego policies.
+    let mut engine = regorus::Engine::new();
+
+    let policy = String::from(
+        r#"
+       package example
+       import rego.v1
+
+       default allow := false
+
+       allow if {
+         print("data.allowed_actions = ", data.allowed_actions)
+         input.action in data.allowed_actions["user1"]
+         print("This rule should be allowed")
+       }
+	"#,
+    );
+
+    // Add policy to the engine.
+    engine.add_policy(String::from("policy.rego"), policy)?;
+
+    // Evaluate first input. Expect to evaluate to false, since state is not set
+    engine.set_input(regorus::Value::from_json_str(
+        r#"{
+      "action": "write"
+    }"#,
+    )?);
+
+    let r = engine.eval_bool_query(String::from("data.example.allow"), false)?;
+    println!("Received result: {:?}", r);
+    assert_eq!(r, false);
+
+    // Add data to engine. Set state
+    engine.add_data(regorus::Value::from_json_str(
+        r#"{
+     "allowed_actions": {
+        "user1" : ["read", "write"]
+     }}"#,
+    )?)?;
+
+    // Evaluate second input. Expect to evaluate to true, since state has been set now
+    engine.set_input(regorus::Value::from_json_str(
+        r#"{
+      "action": "write"
+    }"#,
+    )?);
+
+    let r = engine.eval_bool_query(String::from("data.example.allow"), false)?;
+    println!("Received result: {:?}", r);
+    assert_eq!(
+        r, true,
+        "expect result to be true since rule evaluates to true after state has been updated, per rego logs"
+    );
+
+    Ok(())
+}
+
 fn main() -> Result<()> {
     let cli = Cli::parse();
     run_kata_tests(
@@ -179,5 +237,6 @@ fn main() -> Result<()> {
         &cli.name,
         cli.coverage,
         cli.generate,
-    )
+    )?;
+    stateful_policy_test()
 }