perf: Optimize clone (microsoft#443)

Introduce the notion of CompiledPolicy to hold stuff that
remains immutable during evaluation - e.g. rules, function,
schedules etc

Cloning takes about 60 nano seconds for an engine loaded with
ACI policies. Earlier it used to take 40 microseconds.
Thus there is easily more than 100x speedup.

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

Author: anakrish

benches/regorus_benchmark.rs

@@ -1,6 +1,6 @@
 use std::hint::black_box;
 
-use regorus::Engine;
+use regorus::{Engine, Value};
 
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use serde_json::json;
@@ -93,9 +93,59 @@ fn allow_with_simple_membership(c: &mut Criterion) {
     group.finish();
 }
 
+fn clone(c: &mut Criterion) {
+    // Use Arc<BtreeMap> as a reference. Clone will only increment
+    // the reference count.
+    let mut m = std::collections::BTreeMap::default();
+    m.insert(1, 2);
+    let m = std::sync::Arc::new(m);
+
+    c.bench_function("clone: Arc<BTreeMap>", |b| {
+        b.iter(|| {
+            let _ = m.clone();
+        })
+    });
+
+    let mut engine = Engine::new();
+    engine.set_rego_v0(true);
+    engine
+        .add_policy_from_file("tests/aci/framework.rego")
+        .unwrap();
+    engine.add_policy_from_file("tests/aci/api.rego").unwrap();
+    engine
+        .add_policy_from_file("tests/aci/policy.rego")
+        .unwrap();
+    engine
+        .add_data(Value::from_json_file("tests/aci/data.json").expect("failed to load data.json"))
+        .expect("failed to add data");
+    engine.set_input(
+        Value::from_json_file("tests/aci/input.json").expect("failed to load input.json"),
+    );
+
+    // An engine without preparation will not have processed fields populated.
+    c.bench_function("clone: engine with aci policies", |b| {
+        b.iter(|| {
+            let _ = engine.clone();
+        })
+    });
+
+    // Trigger engine preparation.
+    let _ = engine.eval_query("data.framework.mount_overlay".to_string(), false);
+
+    // Prepared engine will have many more fields populated. But the fields are
+    // immutable after preparation and will be shared between clones.
+    c.bench_function("clone: prepared engine with aci policies", |b| {
+        b.iter(|| {
+            let _ = engine.clone();
+        })
+    });
+}
+
 criterion_group!(
     benches,
     allow_with_simple_equality,
-    allow_with_simple_membership
+    allow_with_simple_membership,
+    clone
 );
+
 criterion_main!(benches);

src/engine.rs

@@ -17,7 +17,7 @@ use anyhow::{bail, Result};
 ///
 #[derive(Debug, Clone)]
 pub struct Engine {
-    modules: Vec<Ref<Module>>,
+    modules: Rc<Vec<Ref<Module>>>,
     interpreter: Interpreter,
     prepared: bool,
     rego_v1: bool,
@@ -59,7 +59,7 @@ impl Engine {
     /// Create an instance of [Engine].
     pub fn new() -> Self {
         Self {
-            modules: vec![],
+            modules: Rc::new(vec![]),
             interpreter: Interpreter::new(),
             prepared: false,
             rego_v1: true,
@@ -125,7 +125,7 @@ impl Engine {
         let source = Source::from_contents(path, rego)?;
         let mut parser = self.make_parser(&source)?;
         let module = Ref::new(parser.parse()?);
-        self.modules.push(module.clone());
+        Rc::make_mut(&mut self.modules).push(module.clone());
         // if policies change, interpreter needs to be prepared again
         self.prepared = false;
         Interpreter::get_path_string(&module.package.refr, Some("data"))
@@ -158,7 +158,7 @@ impl Engine {
         let source = Source::from_file(path)?;
         let mut parser = self.make_parser(&source)?;
         let module = Ref::new(parser.parse()?);
-        self.modules.push(module.clone());
+        Rc::make_mut(&mut self.modules).push(module.clone());
         // if policies change, interpreter needs to be prepared again
         self.prepared = false;
         Interpreter::get_path_string(&module.package.refr, Some("data"))
@@ -628,7 +628,7 @@ impl Engine {
             let schedule = analyzer.analyze(&self.modules)?;
 
             self.interpreter.set_schedule(Some(schedule));
-            self.interpreter.set_modules(&self.modules);
+            self.interpreter.set_modules(self.modules.clone());
 
             self.interpreter.clear_builtins_cache();
             // clean_internal_evaluation_state will set data to an efficient clont of use supplied init_data
@@ -678,13 +678,13 @@ impl Engine {
         }
 
         self.interpreter.check_default_rules()?;
-        for module in self.modules.clone() {
+        for module in self.modules.clone().iter() {
             for rule in &module.policy {
-                self.interpreter.eval_rule(&module, rule)?;
+                self.interpreter.eval_rule(module, rule)?;
             }
         }
         // Defer the evaluation of the default rules to here
-        for module in self.modules.clone() {
+        for module in self.modules.clone().iter() {
             let prev_module = self.interpreter.set_current_module(Some(module.clone()))?;
             for rule in &module.policy {
                 self.interpreter.eval_default_rule(rule)?;
@@ -693,7 +693,7 @@ impl Engine {
         }
 
         // Ensure that all modules are created.
-        for m in &self.modules {
+        for m in self.modules.iter() {
             let path = Parser::get_path_ref_components(&m.package.refr)?;
             let path: Vec<&str> = path.iter().map(|s| s.text()).collect();
             let vref =
@@ -921,7 +921,7 @@ impl Engine {
             ast: &'a Module,
         }
         let mut ast = vec![];
-        for m in &self.modules {
+        for m in self.modules.iter() {
             ast.push(Policy {
                 source: &m.package.span.source,
                 version: 1,