Ability to add custom builtin functions (#132)

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

tests/engine/mod.rs

@@ -0,0 +1,103 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use anyhow::{bail, Result};
+use regorus::*;
+
+#[test]
+fn extension() -> Result<()> {
+    fn repeat(mut params: Vec<Value>) -> Result<Value> {
+        match params.remove(0) {
+            Value::String(s) => {
+                let s = s.as_ref().to_owned();
+                Ok(Value::from(s.clone() + &s))
+            }
+            _ => bail!("param must be string"),
+        }
+    }
+    let mut engine = Engine::new();
+    engine.add_policy(
+        "test.rego".to_string(),
+        r#"package test
+               x = repeat("hello")
+             "#
+        .to_string(),
+    )?;
+
+    // Raises error since repeat is not defined.
+    assert!(engine.eval_query("data.test.x".to_string(), false).is_err());
+
+    // Register extension.
+    engine.add_extension("repeat".to_string(), 1, Box::new(repeat))?;
+
+    // Adding extension twice is error.
+    assert!(engine
+        .add_extension(
+            "repeat".to_string(),
+            1,
+            Box::new(|_| { Ok(Value::Undefined) })
+        )
+        .is_err());
+
+    let r = engine.eval_query("data.test.x".to_string(), false)?;
+    assert_eq!(
+        r.result[0].expressions[0].value.as_string()?.as_ref(),
+        "hellohello"
+    );
+
+    Ok(())
+}
+
+#[test]
+fn extension_with_state() -> Result<()> {
+    #[derive(Clone)]
+    struct Gen {
+        n: i64,
+    }
+
+    let mut engine = Engine::new();
+    engine.add_policy(
+        "test.rego".to_string(),
+        r#"package test
+               x = gen()
+        "#
+        .to_string(),
+    )?;
+
+    let mut g = Box::new(Gen { n: 5 });
+    engine.add_extension(
+        "gen".to_string(),
+        0,
+        Box::new(move |_: Vec<Value>| {
+            let v = Value::from(g.n);
+            g.n += 1;
+            Ok(v)
+        }),
+    )?;
+
+    // First eval.
+    let r = engine.eval_query("data.test.x".to_string(), false)?;
+    assert_eq!(r.result[0].expressions[0].value.as_i64()?, 5);
+
+    // Second eval will produce a new value since for each query, the
+    // internal evaluation state of the interpreter is cleared.
+    // This might change in the future.
+    let r = engine.eval_query("data.test.x".to_string(), false)?;
+    assert_eq!(r.result[0].expressions[0].value.as_i64()?, 6);
+
+    // Clone the engine.
+    // This should also clone the stateful extension.
+    let mut engine1 = engine.clone();
+
+    // Both the engines should produce the same value.
+    let r = engine.eval_query("data.test.x".to_string(), false)?;
+    let r1 = engine1.eval_query("data.test.x".to_string(), false)?;
+    assert_eq!(
+        r.result[0].expressions[0].value,
+        r1.result[0].expressions[0].value
+    );
+
+    assert_eq!(r.result[0].expressions[0].value.as_i64()?, 7);
+
+    Ok(())
+}