Repository: microsoft/regorus (language: python)

Issue:
Ability to add custom builtin functions
Provide a way for users to implement and use custom builtin function.

Similar to OPA's [Custom Built-in Functions in Go](https://www.openpolicyagent.org/docs/latest/extensions/#custom-built-in-functions-in-go)


Discussion / hints:

Fix patch:
diff --git a/src/engine.rs b/src/engine.rs
index 8877d635..f1d64c6e 100644
--- a/src/engine.rs
+++ b/src/engine.rs
@@ -8,7 +8,7 @@ use crate::parser::*;
 use crate::scheduler::*;
 use crate::utils::gather_functions;
 use crate::value::*;
-use crate::QueryResults;
+use crate::{Extension, QueryResults};
 
 use std::convert::AsRef;
 use std::path::Path;
@@ -348,4 +348,92 @@ impl Engine {
         self.interpreter.create_rule_prefixes()?;
         Ok(self.interpreter.get_data_mut().clone())
     }
+
+    /// Add a custom builtin (extension).
+    ///
+    /// * `path`: The fully qualified path of the builtin.
+    /// * `nargs`: The number of arguments the builtin takes.
+    /// * `extension`: The [`Extension`] instance.
+    ///
+    /// ```rust
+    /// # use regorus::*;
+    /// # use anyhow::{bail, Result};
+    /// # fn main() -> Result<()> {
+    /// let mut engine = Engine::new();
+    ///
+    /// // Policy uses `do_magic` custom builtin.
+    /// engine.add_policy(
+    ///    "test.rego".to_string(),
+    ///    r#"package test
+    ///       x = do_magic(1)
+    ///    "#.to_string(),
+    /// )?;
+    ///
+    /// // Evaluating fails since `do_magic` is not defined.
+    /// assert!(engine.eval_query("data.test.x".to_string(), false).is_err());
+    ///
+    /// // Add extension to implement `do_magic`. The extension can be stateful.
+    /// let mut magic = 8;
+    /// engine.add_extension("do_magic".to_string(), 1 , Box::new(move | mut params: Vec<Value> | {
+    ///   // params is mut and therefore individual values can be removed from it and modified.
+    ///   // The number of parameters (1) has already been validated.
+    ///
+    ///   match &params[0].as_i64() {
+    ///      Ok(i) => {
+    ///         // Compute value
+    ///         let v = *i + magic;
+    ///         // Update extension state.
+    ///         magic += 1;
+    ///         Ok(Value::from(v))
+    ///      }
+    ///      // Extensions can raise errors. Regorus will add location information to
+    ///      // the error.
+    ///      _ => bail!("do_magic expects i64 value")
+    ///   }
+    /// }))?;
+    ///
+    /// // Evaluation will now succeed.
+    /// let r = engine.eval_query("data.test.x".to_string(), false)?;
+    /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 9);
+    ///
+    /// // Cloning the engine will also clone the extension.
+    /// let mut engine1 = engine.clone();
+    ///
+    /// // Evaluating again will return a different value since the extension is stateful.
+    /// let r = engine.eval_query("data.test.x".to_string(), false)?;
+    /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 10);
+    ///
+    /// // The second engine has a clone of the extension.
+    /// let r = engine1.eval_query("data.test.x".to_string(), false)?;
+    /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 10);
+    ///
+    /// // Once added, the extension cannot be replaced or removed.
+    /// assert!(engine.add_extension("do_magic".to_string(), 1, Box::new(|_:Vec<Value>| {
+    ///   Ok(Value::Undefined)
+    /// })).is_err());
+    ///
+    /// // Extensions don't support out-parameter syntax.
+    /// engine.add_policy(
+    ///   "policy.rego".to_string(),
+    ///   r#"package invalid
+    ///      x = y {
+    ///       # y = do_magic(2)
+    ///       do_magic(2, y)  # y is supplied as an out parameter.
+    ///     }
+    ///    "#.to_string()
+    /// )?;
+    ///
+    /// // Evaluation fails since y is not defined.
+    /// assert!(engine.eval_query("data.invalid.y".to_string(), false).is_err());
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn add_extension(
+        &mut self,
+        path: String,
+        nargs: u8,
+        extension: Box<dyn Extension>,
+    ) -> Result<()> {
+        self.interpreter.add_extension(path, nargs, extension)
+    }
 }
diff --git a/src/interpreter.rs b/src/interpreter.rs
index 0f69e688..f23aa3dd 100644
--- a/src/interpreter.rs
+++ b/src/interpreter.rs
@@ -9,7 +9,7 @@ use crate::parser::Parser;
 use crate::scheduler::*;
 use crate::utils::*;
 use crate::value::*;
-use crate::{Expression, Location, QueryResult, QueryResults};
+use crate::{Expression, Extension, Location, QueryResult, QueryResults};
 
 use anyhow::{anyhow, bail, Result};
 use log::info;
@@ -65,6 +65,7 @@ pub struct Interpreter {
     allow_deprecated: bool,
     strict_builtin_errors: bool,
     imports: BTreeMap<String, Ref<Expr>>,
+    extensions: HashMap<String, (u8, Box<dyn Extension>)>,
 }
 
 impl Default for Interpreter {
@@ -175,6 +176,7 @@ impl Interpreter {
             allow_deprecated: true,
             strict_builtin_errors: true,
             imports: BTreeMap::default(),
+            extensions: HashMap::new(),
         }
     }
 
@@ -2101,6 +2103,7 @@ impl Interpreter {
         }
 
         let orig_fcn_path = fcn_path;
+
         let mut with_functions_saved = None;
         let fcn_path = match self.with_functions.get(&orig_fcn_path) {
             Some(FunctionModifier::Function(p)) => {
@@ -2121,6 +2124,7 @@ impl Interpreter {
             _ => orig_fcn_path.clone(),
         };
 
+        let mut extension = None;
         let empty: Vec<Ref<Rule>> = vec![];
         let (fcns_rules, fcn_module) = match self.lookup_function_by_name(&fcn_path) {
             Some((fcns, m)) => (fcns, Some(m.clone())),
@@ -2134,6 +2138,11 @@ impl Interpreter {
                     // process default functions later.
                     (&empty, self.module.clone())
                 }
+                // Look up extension.
+                else if let Some(ext) = self.extensions.get_mut(&fcn_path) {
+                    extension = Some(ext);
+                    (&empty, None)
+                }
                 // Look up builtin function.
                 else if let Some(builtin) = self.lookup_builtin(span, &fcn_path)? {
                     let r = self.eval_builtin_call(span, &fcn_path.clone(), *builtin, params);
@@ -2153,6 +2162,21 @@ impl Interpreter {
             return Ok(Value::Undefined);
         }
 
+        if let Some((nargs, ext)) = extension {
+            if param_values.len() != *nargs as usize {
+                bail!(span.error("incorrect number of parameters supplied to extension"));
+            }
+            let r = ext(param_values);
+            // Restore with_functions.
+            if let Some(with_functions) = with_functions_saved {
+                self.with_functions = with_functions;
+            }
+            match r {
+                Ok(v) => return Ok(v),
+                Err(e) => bail!(span.error(&format!("{e}"))),
+            }
+        }
+
         let fcns = fcns_rules.clone();
 
         let mut results: Vec<Value> = Vec::new();
@@ -3016,43 +3040,7 @@ impl Interpreter {
         Ok(())
     }
 
-    pub fn eval_rule(&mut self, module: &Ref<Module>, rule: &Ref<Rule>) -> Result<()> {
-        // Skip reprocessing rule
-        if self.processed.contains(rule) {
-            return Ok(());
-        }
-
-        // Skip default rules
-        if let Rule::Default { .. } = rule.as_ref() {
-            return Ok(());
-        }
-
-        self.active_rules.push(rule.clone());
-        if self.active_rules.iter().filter(|&r| r == rule).count() == 2 {
-            let mut msg = String::default();
-            for r in &self.active_rules {
-                let refr = Self::get_rule_refr(r);
-                let span = refr.span();
-                msg += span
-                    .source
-                    .message(span.line, span.col, "depends on", "")
-                    .as_str();
-            }
-            msg += "cyclic evaluation";
-            let refr = Self::get_rule_refr(rule);
-            let span = refr.span();
-            return Err(span.source.error(
-                span.line,
-                span.col,
-                format!("recursion detected when evaluating rule:{msg}").as_str(),
-            ));
-        }
-
-        // Back up local variables of current function and empty
-        // the local variables of callee function.
-        let scopes = std::mem::take(&mut self.scopes);
-        let prev_module = self.set_current_module(Some(module.clone()))?;
-
+    fn eval_rule_impl(&mut self, module: &Ref<Module>, rule: &Ref<Rule>) -> Result<()> {
         match rule.as_ref() {
             Rule::Spec {
                 span,
@@ -3133,10 +3121,53 @@ impl Interpreter {
             }
             _ => bail!("internal error: unexpected"),
         }
+        Ok(())
+    }
+
+    pub fn eval_rule(&mut self, module: &Ref<Module>, rule: &Ref<Rule>) -> Result<()> {
+        // Skip reprocessing rule
+        if self.processed.contains(rule) {
+            return Ok(());
+        }
+
+        // Skip default rules
+        if let Rule::Default { .. } = rule.as_ref() {
+            return Ok(());
+        }
+
+        self.active_rules.push(rule.clone());
+        if self.active_rules.iter().filter(|&r| r == rule).count() == 2 {
+            let mut msg = String::default();
+            for r in &self.active_rules {
+                let refr = Self::get_rule_refr(r);
+                let span = refr.span();
+                msg += span
+                    .source
+                    .message(span.line, span.col, "depends on", "")
+                    .as_str();
+            }
+            msg += "cyclic evaluation";
+            self.active_rules.pop();
+            let refr = Self::get_rule_refr(rule);
+            let span = refr.span();
+            return Err(span.source.error(
+                span.line,
+                span.col,
+                format!("recursion detected when evaluating rule:{msg}").as_str(),
+            ));
+        }
+
+        // Back up local variables of current function and empty
+        // the local variables of callee function.
+        let scopes = std::mem::take(&mut self.scopes);
+        let prev_module = self.set_current_module(Some(module.clone()))?;
+
+        let res = self.eval_rule_impl(module, rule);
+
         self.set_current_module(prev_module)?;
         self.scopes = scopes;
         match self.active_rules.pop() {
-            Some(ref r) if r == rule => Ok(()),
+            Some(ref r) if r == rule => res,
             _ => bail!("internal error: current rule not active"),
         }
     }
@@ -3409,4 +3440,18 @@ impl Interpreter {
         }
         Ok(())
     }
+
+    pub fn add_extension(
+        &mut self,
+        path: String,
+        nargs: u8,
+        extension: Box<dyn Extension>,
+    ) -> Result<()> {
+        if let std::collections::hash_map::Entry::Vacant(v) = self.extensions.entry(path) {
+            v.insert((nargs, extension));
+            Ok(())
+        } else {
+            bail!("extension already added");
+        }
+    }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 03824a76..28082199 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -260,6 +260,37 @@ pub struct QueryResults {
     pub result: Vec<QueryResult>,
 }
 
+/// A user defined builtin function implementation.
+///
+/// It is not necessary to implement this trait directly.
+pub trait Extension: FnMut(Vec<Value>) -> anyhow::Result<Value> {
+    /// Fn, FnMut etc are not sized and cannot be cloned in their boxed form.
+    /// clone_box exists to overcome that.
+    fn clone_box<'a>(&self) -> Box<dyn 'a + Extension>
+    where
+        Self: 'a;
+}
+
+/// Automatically make matching closures a valid [`Extension`].
+impl<F> Extension for F
+where
+    F: FnMut(Vec<Value>) -> anyhow::Result<Value> + Clone,
+{
+    fn clone_box<'a>(&self) -> Box<dyn 'a + Extension>
+    where
+        Self: 'a,
+    {
+        Box::new(self.clone())
+    }
+}
+
+/// Implement clone for a boxed extension using [`Extension::clone_box`].
+impl<'a> Clone for Box<dyn 'a + Extension> {
+    fn clone(&self) -> Self {
+        (**self).clone_box()
+    }
+}
+
 /// Items in `unstable` are likely to change.
 #[doc(hidden)]
 pub mod unstable {


Test patch:
diff --git a/tests/engine/mod.rs b/tests/engine/mod.rs
new file mode 100644
index 00000000..7b8fa2b8
--- /dev/null
+++ b/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(())
+}
diff --git a/tests/mod.rs b/tests/mod.rs
index aa08f980..4782498e 100644
--- a/tests/mod.rs
+++ b/tests/mod.rs
@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+mod engine;
 mod lexer;
 mod parser;
 mod value;