diff --git a/Cargo.toml b/Cargo.toml index d1143ee..fb45997 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -189,6 +189,11 @@ harness = false name = "aci_benchmark" harness = false +[[bench]] +name = "rvm_benchmark" +harness = false +required-features = ["rvm"] + [[example]] name="regorus" harness=false diff --git a/benches/rvm_benchmark.rs b/benches/rvm_benchmark.rs new file mode 100644 index 0000000..2604be0 --- /dev/null +++ b/benches/rvm_benchmark.rs @@ -0,0 +1,677 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Comprehensive RVM benchmarks covering all aspects of the Rego Virtual Machine. +//! +//! # Policy families +//! +//! | Family | Source | Policies | Inputs/policy | +//! |------------|-------------------------------|----------|---------------| +//! | Synthetic | `benches/evaluation/test_data`| 9 | 3 each | +//! | ACI | `tests/aci` | 9 | 1 each | +//! +//! # Benchmark groups +//! +//! | Group | What it measures | +//! |--------------------------|-------------------------------------------------------| +//! | `cold/{case}/{config}` | Cold: new VM + load + data + input + execute | +//! | `hot/{case}/{config}` | Hot: set_input + execute (VM reused across iters) | +//! | `compilation` | Rego CompiledPolicy → RVM Program | +//! | `serialization` | Program binary serialize / deserialize roundtrip | +//! | `startup` | Isolated VM creation & setup overhead | +//! | `stats` | Instruction/literal counts (reported as throughput) | +//! | `end_to_end` | Full roundtrip: compile → serialize → deserialize → eval | +//! +//! # Running subsets +//! +//! ```sh +//! cargo bench --bench rvm_benchmark # everything +//! cargo bench --bench rvm_benchmark -- cold # all cold eval +//! cargo bench --bench rvm_benchmark -- hot # all hot eval +//! cargo bench --bench rvm_benchmark -- regular_with_limits # one config across cases +//! cargo bench --bench rvm_benchmark -- cold/aci/ # all ACI cold benchmarks +//! cargo bench --bench rvm_benchmark -- rbac # one policy family +//! cargo bench --bench rvm_benchmark -- compilation # compilation only +//! cargo bench --bench rvm_benchmark -- serialization # serialization only +//! cargo bench --bench rvm_benchmark -- startup # startup overhead +//! ``` + +use std::hint::black_box; +use std::num::NonZeroU32; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use serde::{Deserialize, Serialize}; +use walkdir::WalkDir; + +use regorus::languages::rego::compiler::Compiler; +use regorus::rvm::program::Program; +use regorus::rvm::vm::{ExecutionMode, RegoVM}; +use regorus::utils::limits::ExecutionTimerConfig; +use regorus::{Engine, Rc, Value}; + +// --------------------------------------------------------------------------- +// Limit constants – generous ceilings that still exercise the limit-checking +// hot path (memory_check, execution_timer_tick, instruction-limit compare). +// --------------------------------------------------------------------------- + +const MEMORY_LIMIT_BYTES: u64 = 256 * 1024 * 1024; +const TIME_LIMIT: Duration = Duration::from_secs(30); +const TIMER_CHECK_INTERVAL: NonZeroU32 = NonZeroU32::new(16).unwrap(); +const INSTRUCTION_LIMIT: usize = 10_000_000; + +#[derive(Clone, Copy)] +struct EvalConfig { + name: &'static str, + mode: ExecutionMode, + limits: bool, +} + +const EVAL_CONFIGS: [EvalConfig; 4] = [ + EvalConfig { + name: "regular_no_limits", + mode: ExecutionMode::RunToCompletion, + limits: false, + }, + EvalConfig { + name: "regular_with_limits", + mode: ExecutionMode::RunToCompletion, + limits: true, + }, + EvalConfig { + name: "suspendable_no_limits", + mode: ExecutionMode::Suspendable, + limits: false, + }, + EvalConfig { + name: "suspendable_with_limits", + mode: ExecutionMode::Suspendable, + limits: true, + }, +]; + +// --------------------------------------------------------------------------- +// Data types +// --------------------------------------------------------------------------- + +/// A compiled benchmark program ready for RVM execution. +struct BenchmarkProgram { + /// Human-readable name (e.g. "rbac_policy" or "aci/create_container"). + name: String, + /// Pre-compiled RVM program. + program: Arc, + /// Compiled policy (kept for compilation benchmarks). + compiled_policy: regorus::CompiledPolicy, + /// Entry-point rule path. + entry_point: String, + /// Data object (Some for policies that require external data like ACI). + data: Option, + /// Named inputs for this policy. + inputs: Vec<(String, Value)>, +} + +// --------------------------------------------------------------------------- +// ACI YAML types +// --------------------------------------------------------------------------- + +#[derive(Serialize, Deserialize, Debug)] +struct AciTestCase { + note: String, + data: Value, + input: Value, + modules: Vec, + query: String, + want_result: Value, +} + +#[derive(Serialize, Deserialize, Debug)] +struct AciYamlTest { + cases: Vec, +} + +// --------------------------------------------------------------------------- +// Synthetic policy loading +// --------------------------------------------------------------------------- + +/// Policy ↔ input file mapping for synthetic policies. +const SYNTHETIC_POLICIES: &[(&str, &str, &[&str])] = &[ + ( + "rbac_policy", + "rbac_policy.rego", + &["rbac_input.json", "rbac_input2.json", "rbac_input3.json"], + ), + ( + "api_access", + "api_access_policy.rego", + &[ + "api_access_input.json", + "api_access_input2.json", + "api_access_input3.json", + ], + ), + ( + "data_sensitivity", + "data_sensitivity_policy.rego", + &[ + "data_sensitivity_input.json", + "data_sensitivity_input2.json", + "data_sensitivity_input3.json", + ], + ), + ( + "time_based", + "time_based_policy.rego", + &[ + "time_based_input.json", + "time_based_input2.json", + "time_based_input3.json", + ], + ), + ( + "data_processing", + "data_processing_policy.rego", + &[ + "data_processing_input.json", + "data_processing_input2.json", + "data_processing_input3.json", + ], + ), + ( + "azure_vm", + "azure_vm_policy.rego", + &[ + "azure_vm_input.json", + "azure_vm_input2.json", + "azure_vm_input3.json", + ], + ), + ( + "azure_storage", + "azure_storage_policy.rego", + &[ + "azure_storage_input.json", + "azure_storage_input2.json", + "azure_storage_input3.json", + ], + ), + ( + "azure_keyvault", + "azure_keyvault_policy.rego", + &[ + "azure_keyvault_input.json", + "azure_keyvault_input2.json", + "azure_keyvault_input3.json", + ], + ), + ( + "azure_nsg", + "azure_nsg_policy.rego", + &[ + "azure_nsg_input.json", + "azure_nsg_input2.json", + "azure_nsg_input3.json", + ], + ), +]; + +/// Compile synthetic Rego policies into RVM programs. +fn compile_synthetic_programs() -> Vec { + let base_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("benches") + .join("evaluation") + .join("test_data"); + + let entry_point = "data.bench.allow"; + let entry_point_rc: Rc = entry_point.into(); + + SYNTHETIC_POLICIES + .iter() + .map(|(name, policy_file, input_files)| { + let policy_path = base_dir.join("policies").join(policy_file); + let policy_content = std::fs::read_to_string(&policy_path) + .unwrap_or_else(|e| panic!("Failed to read {policy_path:?}: {e}")); + + let mut engine = Engine::new(); + engine + .add_policy("policy.rego".to_string(), policy_content) + .expect("failed to add policy"); + + let compiled_policy = engine + .compile_with_entrypoint(&entry_point_rc) + .expect("failed to compile policy"); + + let program = Compiler::compile_from_policy(&compiled_policy, &[entry_point]) + .expect("failed to compile to RVM program"); + + let inputs: Vec<(String, Value)> = input_files + .iter() + .map(|input_file| { + let input_path = base_dir.join("inputs").join(input_file); + let json = std::fs::read_to_string(&input_path) + .unwrap_or_else(|e| panic!("Failed to read {input_path:?}: {e}")); + let value = Value::from_json_str(&json).expect("failed to parse input JSON"); + let display = input_file.trim_end_matches(".json").to_string(); + (display, value) + }) + .collect(); + + BenchmarkProgram { + name: name.to_string(), + program, + compiled_policy, + entry_point: entry_point.to_string(), + data: None, + inputs, + } + }) + .collect() +} + +// --------------------------------------------------------------------------- +// ACI policy loading +// --------------------------------------------------------------------------- + +/// Load all ACI test cases from YAML files. +fn load_aci_cases(dir: &Path) -> Vec { + let mut cases = Vec::new(); + for entry in WalkDir::new(dir) + .sort_by_file_name() + .into_iter() + .filter_map(|e| e.ok()) + { + let path = entry.path(); + if !path.to_string_lossy().ends_with(".yaml") { + continue; + } + let yaml = std::fs::read(path).expect("failed to read yaml"); + let yaml = String::from_utf8_lossy(&yaml); + let test: AciYamlTest = serde_yaml::from_str(&yaml).expect("failed to deserialize yaml"); + cases.extend(test.cases); + } + cases +} + +/// Build an Engine with policies loaded for a given ACI test case. +fn build_aci_engine(dir: &Path, case: &AciTestCase) -> Engine { + let mut engine = Engine::new(); + engine.set_rego_v0(true); + engine + .add_data(case.data.clone()) + .expect("failed to add data"); + engine.set_input(case.input.clone()); + for (idx, rego) in case.modules.iter().enumerate() { + if rego.ends_with(".rego") { + engine + .add_policy_from_file(dir.join(rego).to_str().expect("invalid path")) + .expect("failed to add policy"); + } else { + engine + .add_policy(format!("rego{idx}.rego"), rego.clone()) + .expect("failed to add policy"); + } + } + engine +} + +/// Compile ACI test cases into RVM programs. +fn compile_aci_programs() -> Vec { + let dir = Path::new("tests/aci"); + load_aci_cases(dir) + .into_iter() + .map(|case| { + let mut engine = build_aci_engine(dir, &case); + let rule = case.query.replace("=x", ""); + let rule_rc: Rc = rule.clone().into(); + let compiled_policy = engine + .compile_with_entrypoint(&rule_rc) + .expect("failed to compile"); + let program = Compiler::compile_from_policy(&compiled_policy, &[rule.as_str()]) + .expect("failed to compile to RVM"); + + BenchmarkProgram { + name: format!("aci/{}", case.note), + program, + compiled_policy, + entry_point: rule, + data: Some(case.data), + inputs: vec![("input".to_string(), case.input)], + } + }) + .collect() +} + +// --------------------------------------------------------------------------- +// Compile all policies +// --------------------------------------------------------------------------- + +/// Compile all policies (synthetic + ACI) into RVM programs. +fn compile_all_programs() -> Vec { + let mut programs = compile_synthetic_programs(); + programs.extend(compile_aci_programs()); + programs +} + +// --------------------------------------------------------------------------- +// Limit helpers +// --------------------------------------------------------------------------- + +/// Apply or remove production-style limits based on a boolean flag. +fn configure_limits(vm: &mut RegoVM, limits: bool) { + if limits { + regorus::set_global_memory_limit(Some(MEMORY_LIMIT_BYTES)); + vm.set_execution_timer_config(Some(ExecutionTimerConfig { + limit: TIME_LIMIT, + check_interval: TIMER_CHECK_INTERVAL, + })); + vm.set_max_instructions(INSTRUCTION_LIMIT); + } else { + regorus::set_global_memory_limit(None); + vm.set_execution_timer_config(None); + vm.set_max_instructions(usize::MAX); + } +} + +// --------------------------------------------------------------------------- +// Cold evaluation — new VM per iteration (full setup + execute) +// +// Benchmarks are registered case-first so each workload is shown with all +// config variants adjacent to one another, making per-case comparisons easier. +// --------------------------------------------------------------------------- + +fn bench_cold(c: &mut Criterion) { + let programs = compile_all_programs(); + let mut group = c.benchmark_group("cold"); + + for bp in &programs { + for (input_name, input_value) in &bp.inputs { + let case_id = if bp.inputs.len() == 1 { + bp.name.clone() + } else { + format!("{}/{}", bp.name, input_name) + }; + let program = bp.program.clone(); + let data = bp.data.clone(); + let input = input_value.clone(); + + for config in EVAL_CONFIGS { + group.bench_function(BenchmarkId::new(&case_id, config.name), |b| { + b.iter(|| { + let mut vm = RegoVM::new(); + vm.set_execution_mode(config.mode); + vm.load_program(black_box(program.clone())); + if let Some(ref d) = data { + vm.set_data(black_box(d.clone())).unwrap(); + } + vm.set_input(black_box(input.clone())); + configure_limits(&mut vm, config.limits); + black_box(vm.execute().unwrap()) + }) + }); + } + } + } + group.finish(); +} + +// --------------------------------------------------------------------------- +// Hot evaluation — VM reused across iterations +// +// The VM is created once with program, data, mode, and limits. Each +// iteration only calls set_input + execute, measuring pure execution +// overhead with minimal setup. A warm-up execution fills the register +// window pool so all iterations benefit from pooled allocations. +// --------------------------------------------------------------------------- + +fn bench_hot(c: &mut Criterion) { + let programs = compile_all_programs(); + let mut group = c.benchmark_group("hot"); + + for bp in &programs { + let program = bp.program.clone(); + let data = bp.data.clone(); + let inputs: Vec = bp.inputs.iter().map(|(_, v)| v.clone()).collect(); + let num_inputs = inputs.len(); + + for config in EVAL_CONFIGS { + group.bench_function(BenchmarkId::new(&bp.name, config.name), |b| { + let mut vm = RegoVM::new(); + vm.set_execution_mode(config.mode); + vm.load_program(program.clone()); + if let Some(ref d) = data { + vm.set_data(d.clone()).unwrap(); + } + configure_limits(&mut vm, config.limits); + + // Warm up: fill register window pools, caches, etc. + vm.set_input(inputs[0].clone()); + vm.execute().expect("warm-up failed"); + + let mut i = 0usize; + b.iter(|| { + let input = &inputs[i % num_inputs]; + vm.set_input(black_box(input.clone())); + black_box(vm.execute().unwrap()); + i += 1; + }) + }); + } + } + group.finish(); +} + +// --------------------------------------------------------------------------- +// Compilation — Rego CompiledPolicy → RVM Program +// --------------------------------------------------------------------------- + +fn bench_compilation(c: &mut Criterion) { + let programs = compile_all_programs(); + let mut group = c.benchmark_group("compilation"); + + for bp in &programs { + let entry_point: &str = &bp.entry_point; + group.bench_with_input( + BenchmarkId::new("rego_to_rvm", &bp.name), + &bp.compiled_policy, + |b, compiled_policy| { + b.iter(|| { + Compiler::compile_from_policy( + black_box(compiled_policy), + black_box(&[entry_point]), + ) + .unwrap(); + }) + }, + ); + } + group.finish(); +} + +// --------------------------------------------------------------------------- +// Serialization — binary serialize / deserialize roundtrip +// --------------------------------------------------------------------------- + +fn bench_serialization(c: &mut Criterion) { + let programs = compile_all_programs(); + let mut group = c.benchmark_group("serialization"); + + for bp in &programs { + let program = &bp.program; + let serialized = program + .serialize_binary() + .expect("failed to serialize program"); + let byte_len = serialized.len() as u64; + + group.throughput(Throughput::Bytes(byte_len)); + group.bench_function(BenchmarkId::new("serialize", &bp.name), |b| { + b.iter(|| black_box(program.serialize_binary().unwrap())) + }); + + group.throughput(Throughput::Bytes(byte_len)); + group.bench_function(BenchmarkId::new("deserialize", &bp.name), |b| { + b.iter(|| black_box(Program::deserialize_binary(black_box(&serialized)).unwrap())) + }); + } + group.finish(); +} + +// --------------------------------------------------------------------------- +// Startup — isolated VM creation & setup overhead +// --------------------------------------------------------------------------- + +fn bench_startup(c: &mut Criterion) { + let programs = compile_all_programs(); + let mut group = c.benchmark_group("startup"); + + // Use the first program as representative for startup overhead. + let bp = &programs[0]; + let program = bp.program.clone(); + let input = bp.inputs[0].1.clone(); + + // Bare VM creation + group.bench_function("new", |b| b.iter(|| black_box(RegoVM::new()))); + + // load_program (Arc clone + internal setup) + group.bench_function("load_program", |b| { + b.iter(|| { + let mut vm = RegoVM::new(); + vm.load_program(black_box(program.clone())); + black_box(&vm); + }) + }); + + // set_input + group.bench_function("set_input", |b| { + let mut vm = RegoVM::new(); + vm.load_program(program.clone()); + b.iter(|| { + vm.set_input(black_box(input.clone())); + }) + }); + + group.finish(); +} + +// --------------------------------------------------------------------------- +// Stats — instruction / literal counts (reported as throughput) +// --------------------------------------------------------------------------- + +fn bench_stats(c: &mut Criterion) { + let programs = compile_all_programs(); + + eprintln!(); + eprintln!( + "{:<30} {:>8} {:>8} {:>8} {:>10}", + "program", "instrs", "lits", "entries", "bytes" + ); + eprintln!("{}", "-".repeat(70)); + + let mut group = c.benchmark_group("stats"); + for bp in &programs { + let serialized = bp.program.serialize_binary().expect("serialize failed"); + let byte_len = serialized.len(); + let instr_count = bp.program.instructions.len(); + let lit_count = bp.program.literals.len(); + let entry_count = bp.program.entry_points.len(); + + eprintln!( + "{:<30} {:>8} {:>8} {:>8} {:>10}", + bp.name, instr_count, lit_count, entry_count, byte_len, + ); + + group.throughput(Throughput::Elements(instr_count as u64)); + group.bench_function(BenchmarkId::new("serialize", &bp.name), |b| { + b.iter(|| black_box(bp.program.serialize_binary().unwrap())) + }); + } + group.finish(); +} + +// --------------------------------------------------------------------------- +// End-to-end roundtrip (compile + serialize + deserialize + eval) +// +// Only runs for synthetic policies where we have direct access to rego +// source files. ACI policies are loaded from YAML with module references +// which makes the setup pipeline different. +// --------------------------------------------------------------------------- + +fn bench_end_to_end(c: &mut Criterion) { + let base_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("benches") + .join("evaluation") + .join("test_data"); + + let entry_point = "data.bench.allow"; + let entry_point_rc: Rc = entry_point.into(); + + let mut group = c.benchmark_group("end_to_end"); + + for &(name, policy_file, input_files) in SYNTHETIC_POLICIES { + let policy_path = base_dir.join("policies").join(policy_file); + let policy_content = std::fs::read_to_string(&policy_path) + .unwrap_or_else(|e| panic!("Failed to read {policy_path:?}: {e}")); + + // Use just the first input for end-to-end + let input_path = base_dir.join("inputs").join(input_files[0]); + let input_json = std::fs::read_to_string(&input_path) + .unwrap_or_else(|e| panic!("Failed to read {input_path:?}: {e}")); + + group.bench_function(BenchmarkId::new("roundtrip", name), |b| { + b.iter(|| { + // 1. Engine + parse + let mut engine = Engine::new(); + engine + .add_policy("policy.rego".to_string(), policy_content.clone()) + .unwrap(); + + // 2. Compile to CompiledPolicy + let compiled_policy = engine.compile_with_entrypoint(&entry_point_rc).unwrap(); + + // 3. Compile to RVM Program + let program = + Compiler::compile_from_policy(&compiled_policy, &[entry_point]).unwrap(); + + // 4. Serialize + let bytes = program.serialize_binary().unwrap(); + + // 5. Deserialize + let deserialized = Program::deserialize_binary(&bytes).unwrap(); + let program = match deserialized { + regorus::rvm::program::DeserializationResult::Complete(p) => Arc::new(p), + regorus::rvm::program::DeserializationResult::Partial(p) => { + Arc::new(Program::compile_from_partial(p).unwrap()) + } + }; + + // 6. Execute + let mut vm = RegoVM::new(); + vm.load_program(program); + let input = Value::from_json_str(&input_json).unwrap(); + vm.set_input(input); + black_box(vm.execute().unwrap()); + }) + }); + } + + group.finish(); +} + +// --------------------------------------------------------------------------- +// Criterion groups — organised for selective runs +// --------------------------------------------------------------------------- + +criterion_group!(cold_benches, bench_cold); + +criterion_group!(hot_benches, bench_hot); + +criterion_group!( + misc_benches, + bench_compilation, + bench_serialization, + bench_startup, + bench_stats, + bench_end_to_end, +); + +criterion_main!(cold_benches, hot_benches, misc_benches); diff --git a/src/rvm/vm/comprehension.rs b/src/rvm/vm/comprehension.rs index 298b52b..9df9074 100644 --- a/src/rvm/vm/comprehension.rs +++ b/src/rvm/vm/comprehension.rs @@ -255,26 +255,22 @@ impl RegoVM { }; let result_reg = comprehension_context.result_reg; - let current_result = self.get_register(result_reg)?.clone(); - let mode = comprehension_context.mode.clone(); + // Take ownership of the result register so Rc refcount stays at 1, + // allowing Rc::make_mut to mutate in-place instead of deep-cloning. + let mut current_result = self.take_register(result_reg)?; - let updated_result = match (mode, current_result) { - (ComprehensionMode::Set, Value::Set(set)) => { - let mut new_set = set.as_ref().clone(); - new_set.insert(value_to_add); - Value::Set(crate::Rc::new(new_set)) + match (&comprehension_context.mode, &mut current_result) { + (&ComprehensionMode::Set, &mut Value::Set(ref mut set)) => { + crate::Rc::make_mut(set).insert(value_to_add); } - (ComprehensionMode::Array, Value::Array(arr)) => { - let mut new_arr = arr.as_ref().to_vec(); - new_arr.push(value_to_add); - Value::Array(crate::Rc::new(new_arr)) + (&ComprehensionMode::Array, &mut Value::Array(ref mut arr)) => { + crate::Rc::make_mut(arr).push(value_to_add); } - (ComprehensionMode::Object, Value::Object(obj)) => { + (&ComprehensionMode::Object, &mut Value::Object(ref mut obj)) => { if let Some(key) = key_value { - let mut new_obj = obj.as_ref().clone(); - new_obj.insert(key, value_to_add); - Value::Object(crate::Rc::new(new_obj)) + crate::Rc::make_mut(obj).insert(key, value_to_add); } else { + self.set_register(result_reg, current_result)?; self.comprehension_stack.push(comprehension_context); return Err(VmError::InvalidIteration { value: Value::String(Arc::from("Object comprehension requires key")), @@ -283,15 +279,17 @@ impl RegoVM { } } (_mode, other) => { + let offending = core::mem::replace(other, Value::Undefined); + self.set_register(result_reg, current_result)?; self.comprehension_stack.push(comprehension_context); return Err(VmError::InvalidIteration { - value: other, + value: offending, pc: self.pc, }); } - }; + } - self.set_register(result_reg, updated_result)?; + self.set_register(result_reg, current_result)?; if let Some(iter_state) = comprehension_context.iteration_state.as_mut() { match *iter_state { @@ -357,7 +355,6 @@ impl RegoVM { let ( value_to_add, key_value, - current_result, mode, result_reg_idx, key_reg_idx, @@ -384,7 +381,6 @@ impl RegoVM { }; let result_reg_idx = context.result_reg; - let current_result = self.get_register(result_reg_idx)?.clone(); let mode = context.mode.clone(); let iteration_key = self.get_register(context.key_reg)?.clone(); let iteration_value = self.get_register(context.value_reg)?.clone(); @@ -392,7 +388,6 @@ impl RegoVM { ( value_to_add, key_value, - current_result, mode, result_reg_idx, context.key_reg, @@ -408,23 +403,22 @@ impl RegoVM { } }; - let updated_result = match (mode, current_result) { - (ComprehensionMode::Set, Value::Set(set)) => { - let mut new_set = set.as_ref().clone(); - new_set.insert(value_to_add); - Value::Set(crate::Rc::new(new_set)) + // Take ownership of the result register so Rc refcount stays at 1, + // allowing Rc::make_mut to mutate in-place instead of deep-cloning. + let mut current_result = self.take_register(result_reg_idx)?; + + match (&mode, &mut current_result) { + (&ComprehensionMode::Set, &mut Value::Set(ref mut set)) => { + crate::Rc::make_mut(set).insert(value_to_add); } - (ComprehensionMode::Array, Value::Array(arr)) => { - let mut new_arr = arr.as_ref().to_vec(); - new_arr.push(value_to_add); - Value::Array(crate::Rc::new(new_arr)) + (&ComprehensionMode::Array, &mut Value::Array(ref mut arr)) => { + crate::Rc::make_mut(arr).push(value_to_add); } - (ComprehensionMode::Object, Value::Object(obj)) => { + (&ComprehensionMode::Object, &mut Value::Object(ref mut obj)) => { if let Some(key) = key_value { - let mut new_obj = obj.as_ref().clone(); - new_obj.insert(key, value_to_add); - Value::Object(crate::Rc::new(new_obj)) + crate::Rc::make_mut(obj).insert(key, value_to_add); } else { + self.set_register(result_reg_idx, current_result)?; return Err(VmError::InvalidIteration { value: Value::String(Arc::from("Object comprehension requires key")), pc: self.pc, @@ -432,12 +426,14 @@ impl RegoVM { } } (_mode, other) => { + let offending = core::mem::replace(other, Value::Undefined); + self.set_register(result_reg_idx, current_result)?; return Err(VmError::InvalidIteration { - value: other, + value: offending, pc: self.pc, }); } - }; + } let (iteration_state_snapshot, body_start, comprehension_end) = { let frame = self.execution_stack.get_mut(comprehension_index).ok_or( @@ -489,7 +485,7 @@ impl RegoVM { } }; - self.set_register(result_reg_idx, updated_result)?; + self.set_register(result_reg_idx, current_result)?; if let Some(state) = iteration_state_snapshot.as_ref() { let has_next = self.setup_next_iteration(state, key_reg_idx, value_reg_idx)?; diff --git a/src/rvm/vm/dispatch.rs b/src/rvm/vm/dispatch.rs index f86e7d1..02b9313 100644 --- a/src/rvm/vm/dispatch.rs +++ b/src/rvm/vm/dispatch.rs @@ -434,7 +434,8 @@ impl RegoVM { let key_value = self.get_register(key)?.clone(); let value_value = self.get_register(value)?.clone(); - let mut obj_value = self.get_register(obj)?.clone(); + // Take ownership so Rc refcount stays at 1 and make_mut is a no-op. + let mut obj_value = self.take_register(obj)?; if let Ok(obj_mut) = obj_value.as_object_mut() { obj_mut.insert(key_value, value_value); @@ -574,7 +575,8 @@ impl RegoVM { ArrayPush { arr, value } => { let value_to_push = self.get_register(value)?.clone(); - let mut arr_value = self.get_register(arr)?.clone(); + // Take ownership so Rc refcount stays at 1 and make_mut is a no-op. + let mut arr_value = self.take_register(arr)?; if let Ok(arr_mut) = arr_value.as_array_mut() { arr_mut.push(value_to_push); @@ -632,7 +634,8 @@ impl RegoVM { SetAdd { set, value } => { let value_to_add = self.get_register(value)?.clone(); - let mut set_value = self.get_register(set)?.clone(); + // Take ownership so Rc refcount stays at 1 and make_mut is a no-op. + let mut set_value = self.take_register(set)?; if let Ok(set_mut) = set_value.as_set_mut() { set_mut.insert(value_to_add); diff --git a/src/rvm/vm/execution.rs b/src/rvm/vm/execution.rs index 9fd20d7..8400714 100644 --- a/src/rvm/vm/execution.rs +++ b/src/rvm/vm/execution.rs @@ -4,7 +4,6 @@ use crate::rvm::instructions::Instruction; use crate::rvm::program::Program; use crate::value::Value; use alloc::string::String; -use alloc::vec::Vec; use core::convert::TryFrom as _; use super::dispatch::InstructionOutcome; @@ -24,35 +23,22 @@ impl RegoVM { } pub fn execute_entry_point_by_index(&mut self, index: usize) -> Result { - let entry_points: Vec<(String, usize)> = self - .program - .entry_points - .iter() - .map(|(name, pc)| (name.clone(), *pc)) - .collect(); - - if index >= entry_points.len() { - return Err(VmError::InvalidEntryPointIndex { - index, - max_index: entry_points.len().saturating_sub(1), - pc: self.pc, - }); - } - - let &(ref entry_point_name, entry_point_pc) = - entry_points - .get(index) - .ok_or(VmError::InvalidEntryPointIndex { + let (entry_point_name, entry_point_pc) = { + let (name, &pc) = self.program.entry_points.get_index(index).ok_or( + VmError::InvalidEntryPointIndex { index, - max_index: entry_points.len().saturating_sub(1), + max_index: self.program.entry_points.len().saturating_sub(1), pc: self.pc, - })?; + }, + )?; + (name.clone(), pc) + }; if entry_point_pc >= self.program.instructions.len() { return Err(VmError::EntryPointPcOutOfBounds { pc: entry_point_pc, instruction_count: self.program.instructions.len(), - entry_point: entry_point_name.clone(), + entry_point: entry_point_name, }); } @@ -215,15 +201,18 @@ impl RegoVM { } pub fn resume(&mut self, resume_value: Option) -> Result { - let (reason, mut last_result) = match self.execution_state.clone() { + let old_state = core::mem::replace(&mut self.execution_state, ExecutionState::Running); + let (reason, mut last_result) = match old_state { ExecutionState::Suspended { reason, last_result, .. } => (reason, last_result), current_state => { + let desc = alloc::format!("{:?}", current_state); + self.execution_state = current_state; return Err(VmError::InvalidResumeState { - state: alloc::format!("{:?}", current_state), + state: desc, pc: self.pc, }); } diff --git a/src/rvm/vm/functions.rs b/src/rvm/vm/functions.rs index e892430..9a4f9aa 100644 --- a/src/rvm/vm/functions.rs +++ b/src/rvm/vm/functions.rs @@ -2,8 +2,6 @@ // Licensed under the MIT License. use crate::builtins; use crate::value::Value; -use alloc::string::String; -use alloc::vec::Vec; use super::errors::{Result, VmError}; use super::execution_model::ExecutionMode; @@ -39,24 +37,26 @@ impl RegoVM { } pub(super) fn execute_builtin_call(&mut self, params_index: u16) -> Result<()> { - let params = self - .program + let program = self.program.clone(); + let params = program .instruction_data .get_builtin_call_params(params_index) .ok_or(VmError::InvalidBuiltinCallParamsIndex { index: params_index, pc: self.pc, - available: self.program.instruction_data.builtin_call_params.len(), - })?; - let builtin_info = self.program.get_builtin_info(params.builtin_index).ok_or( + available: program.instruction_data.builtin_call_params.len(), + })? + .clone(); + let builtin_info = program.get_builtin_info(params.builtin_index).ok_or( VmError::InvalidBuiltinInfoIndex { index: params.builtin_index, pc: self.pc, - available: self.program.builtin_info_table.len(), + available: program.builtin_info_table.len(), }, )?; - let mut args = Vec::new(); + let mut args = core::mem::take(&mut self.cached_builtin_args); + args.clear(); for &arg_reg in params.arg_registers().iter() { let arg_value = self.get_register(arg_reg)?.clone(); args.push(arg_value); @@ -65,6 +65,7 @@ impl RegoVM { let expected_args = builtin_info.num_args; let actual_args = args.len(); if u16::try_from(actual_args).unwrap_or(u16::MAX) != expected_args { + self.cached_builtin_args = args; return Err(VmError::BuiltinArgumentMismatch { expected: expected_args, actual: actual_args, @@ -73,65 +74,84 @@ impl RegoVM { } if args.iter().any(|a| a == &Value::Undefined) { + self.cached_builtin_args = args; self.set_register(params.dest, Value::Undefined)?; self.memory_check()?; return Ok(()); } - if let Some(builtin_fcn) = self.program.get_resolved_builtin(params.builtin_index) { - let dummy_source = crate::lexer::Source::from_contents("arg".into(), String::new())?; - let dummy_span = crate::lexer::Span { - source: dummy_source, - line: 1, - col: 1, - start: 0, - end: 3, - }; - - let mut dummy_exprs: Vec> = Vec::new(); - for _ in 0..args.len() { - let dummy_expr = crate::ast::Expr::Null { - span: dummy_span.clone(), - value: Value::Null, - eidx: 0, - }; - dummy_exprs.push(crate::ast::Ref::new(dummy_expr)); + // Extract everything we need from program before releasing the borrow. + let builtin_fn = match program.get_resolved_builtin(params.builtin_index) { + Some(fcn) => fcn.0, + None => { + self.cached_builtin_args = args; + return Err(VmError::BuiltinNotResolved { + name: builtin_info.name.clone(), + pc: self.pc, + }); } + }; + let cache_name = builtins::must_cache(builtin_info.name.as_str()); + drop(program); - let cache_name = builtins::must_cache(builtin_info.name.as_str()); - if let Some(name) = cache_name { - if let Some(value) = self.builtins_cache.get(&(name, args.clone())) { - self.set_register(params.dest, value.clone())?; - return Ok(()); + self.ensure_dummy_exprs(args.len())?; + let dummy_span = self.get_dummy_span()?.clone(); + // Take the dummy_exprs vec out of self so we can pass it to the builtin + // while still calling &mut self methods afterwards. + let dummy_exprs = core::mem::take(&mut self.dummy_exprs); + + if let Some(name) = cache_name { + if let Some(entries) = self.builtins_cache.get(name) { + for entry in entries { + if entry.0.as_slice() == args.as_slice() { + let cached = entry.1.clone(); + self.dummy_exprs = dummy_exprs; + self.cached_builtin_args = args; + self.set_register(params.dest, cached)?; + self.memory_check()?; + return Ok(()); + } } } - - let result = - match (builtin_fcn.0)(&dummy_span, &dummy_exprs, &args, self.strict_builtin_errors) - { - Ok(value) => value, - Err(_) if !self.strict_builtin_errors => Value::Undefined, - Err(err) => return Err(err.into()), - }; - - if result == Value::Undefined { - self.set_register(params.dest, Value::Undefined)?; - } else { - self.set_register(params.dest, result.clone())?; - } - - if let Some(name) = cache_name { - self.builtins_cache.insert((name, args), result); - } - - self.memory_check()?; - } else { - return Err(VmError::BuiltinNotResolved { - name: builtin_info.name.clone(), - pc: self.pc, - }); } + let result = match builtin_fn( + &dummy_span, + dummy_exprs.get(..args.len()).unwrap_or(&[]), + &args, + self.strict_builtin_errors, + ) { + Ok(value) => value, + Err(_) if !self.strict_builtin_errors => Value::Undefined, + Err(err) => { + self.dummy_exprs = dummy_exprs; + self.cached_builtin_args = args; + return Err(err.into()); + } + }; + + // Put dummy_exprs back for reuse. + self.dummy_exprs = dummy_exprs; + + if let Some(name) = cache_name { + // Move args into the cache. The now-empty (zero-capacity) Vec is + // stored back in cached_builtin_args; the next call will re-allocate. + // This is acceptable because cache inserts are rare (once per unique + // argument set) while the hot path (cache hit above) reuses the Vec. + let cache_args = core::mem::take(&mut args); + self.cached_builtin_args = args; + self.builtins_cache + .entry(name) + .or_default() + .push((cache_args, result.clone())); + self.set_register(params.dest, result)?; + } else { + self.cached_builtin_args = args; + self.set_register(params.dest, result)?; + } + + self.memory_check()?; + Ok(()) } } diff --git a/src/rvm/vm/machine.rs b/src/rvm/vm/machine.rs index 670bea3..9d931be 100644 --- a/src/rvm/vm/machine.rs +++ b/src/rvm/vm/machine.rs @@ -109,8 +109,15 @@ pub struct RegoVM { /// Whether builtins should raise errors strictly or return undefined on failure pub(super) strict_builtin_errors: bool, - /// Cache for builtin calls that must stay deterministic across a single evaluation - pub(super) builtins_cache: BTreeMap<(&'static str, Vec), Value>, + /// Cache for builtin calls that must stay deterministic across a single evaluation. + /// + /// Two-level structure: outer BTreeMap keyed by builtin name, inner Vec of + /// (args, result) pairs scanned linearly. This avoids allocating a composite + /// key on every lookup (which a single-level BTreeMap<(name, Vec), Value> + /// would require). Linear scan is fast for the small number of entries per + /// builtin (typically <10). Can be revisited with a HashMap if `Value` gains + /// a `Hash` implementation. + pub(super) builtins_cache: BTreeMap<&'static str, Vec<(Vec, Value)>>, /// Optional override for the execution timer configuration pub(super) execution_timer_config: Option, @@ -120,6 +127,15 @@ pub struct RegoVM { /// Elapsed wall-clock time recorded when the VM entered a suspended state pub(super) execution_timer_elapsed_at_suspend: Option, + + /// Cached dummy span for builtin calls (avoids Source::from_contents per call) + pub(super) dummy_span: Option, + + /// Cached dummy expressions for builtin calls (avoids Rc allocs per call) + pub(super) dummy_exprs: Vec>, + + /// Cached args Vec for builtin calls (avoids Vec allocation per call) + pub(super) cached_builtin_args: Vec, } impl Default for RegoVM { @@ -163,6 +179,9 @@ impl RegoVM { execution_timer_config: None, execution_timer: ExecutionTimer::new(fallback_timer), execution_timer_elapsed_at_suspend: None, + dummy_span: None, + dummy_exprs: Vec::new(), + cached_builtin_args: Vec::new(), } } @@ -451,6 +470,24 @@ impl RegoVM { }) } + /// Take ownership of a register value, replacing it with `Value::Undefined`. + /// This avoids bumping the Rc refcount that a clone would cause, keeping the + /// refcount at 1 so that subsequent `Rc::make_mut` calls can mutate in place. + #[inline] + #[allow(dead_code)] + pub(super) fn take_register(&mut self, index: u8) -> Result { + let register_count = self.registers.len(); + + let slot = self.registers.get_mut(usize::from(index)).ok_or( + VmError::RegisterIndexOutOfBounds { + index, + pc: self.pc, + register_count, + }, + )?; + Ok(core::mem::replace(slot, Value::Undefined)) + } + #[inline] #[allow(dead_code)] pub(super) fn set_register(&mut self, index: u8, value: Value) -> Result<()> { @@ -486,4 +523,44 @@ impl RegoVM { pub(super) fn memory_check(&mut self) -> Result<()> { Ok(()) } + + /// Get or create the cached dummy span for builtin calls. + pub(super) fn get_dummy_span(&mut self) -> Result<&crate::lexer::Span> { + if self.dummy_span.is_none() { + let source = crate::lexer::Source::from_contents("".into(), String::new()) + .map_err(|e| VmError::Internal { + message: alloc::format!("failed to create dummy source: {e}"), + pc: self.pc, + })?; + self.dummy_span = Some(crate::lexer::Span { + source, + line: 1, + col: 1, + start: 0, + end: 0, + }); + } + // SAFETY: we just ensured it's Some above + self.dummy_span.as_ref().ok_or(VmError::Internal { + message: String::from("dummy span not initialized"), + pc: self.pc, + }) + } + + /// Ensure the cached dummy_exprs vec has at least `count` elements. + pub(super) fn ensure_dummy_exprs(&mut self, count: usize) -> Result<()> { + if self.dummy_exprs.len() >= count { + return Ok(()); + } + let span = self.get_dummy_span()?.clone(); + while self.dummy_exprs.len() < count { + self.dummy_exprs + .push(crate::ast::Ref::new(crate::ast::Expr::Null { + span: span.clone(), + value: Value::Null, + eidx: 0, + })); + } + Ok(()) + } } diff --git a/src/rvm/vm/rules.rs b/src/rvm/vm/rules.rs index 5173a23..d9388bc 100644 --- a/src/rvm/vm/rules.rs +++ b/src/rvm/vm/rules.rs @@ -153,16 +153,17 @@ impl RegoVM { }); } - let rule_info = self - .program + // Clone the Arc (cheap atomic increment) so we can borrow &RuleInfo + // without holding an immutable borrow on self. + let program = self.program.clone(); + let rule_info = program .rule_infos .get(rule_idx) .ok_or(VmError::RuleInfoMissing { index: rule_index, pc: self.pc, - available: self.program.rule_infos.len(), - })? - .clone(); + available: program.rule_infos.len(), + })?; let is_function_rule = rule_info.function_info.is_some(); @@ -214,7 +215,7 @@ impl RegoVM { }); let (final_result, rule_failed_due_to_inconsistency) = self - .execute_rule_definitions_common(&rule_definitions, &rule_info, function_call_params)?; + .execute_rule_definitions_common(&rule_definitions, rule_info, function_call_params)?; self.set_register(dest, Value::Undefined)?; @@ -230,7 +231,7 @@ impl RegoVM { Value::Undefined }; - self.set_register(dest, result_from_rule.clone())?; + self.set_register(dest, result_from_rule)?; if self.get_register(dest)? == &Value::Undefined && !rule_failed_due_to_inconsistency { match call_context.rule_type { @@ -272,7 +273,7 @@ impl RegoVM { pc: self.pc, available, })?; - *entry = (true, final_value.clone()); + *entry = (true, final_value); } Ok(()) } @@ -301,16 +302,15 @@ impl RegoVM { }); } - let rule_info = self - .program + let program = self.program.clone(); + let rule_info = program .rule_infos .get(rule_idx) .ok_or(VmError::RuleInfoMissing { index: rule_index, pc: self.pc, - available: self.program.rule_infos.len(), - })? - .clone(); + available: program.rule_infos.len(), + })?; let is_function_rule = rule_info.function_info.is_some(); @@ -425,7 +425,7 @@ impl RegoVM { }; let initial_pc = self - .prepare_rule_frame_initial_pc(&mut frame_data, &rule_info)? + .prepare_rule_frame_initial_pc(&mut frame_data, rule_info)? .ok_or(VmError::RuleFrameMissingInitialPc { pc: self.pc })?; let frame = ExecutionFrame::new(initial_pc, FrameKind::Rule(frame_data)); @@ -639,16 +639,15 @@ impl RegoVM { } = frame_data; let rule_idx = usize::from(rule_index); - let rule_info = self - .program + let program = self.program.clone(); + let rule_info = program .rule_infos .get(rule_idx) .ok_or(VmError::RuleInfoMissing { index: rule_index, pc: self.pc, - available: self.program.rule_infos.len(), - })? - .clone(); + available: program.rule_infos.len(), + })?; let result_from_rule = if rule_failed_due_to_inconsistency { Value::Undefined @@ -760,6 +759,7 @@ impl RegoVM { pc: self.pc, available, })?; + // Clone into cache; return the original below. *entry = (true, final_value.clone()); } @@ -778,12 +778,20 @@ impl RegoVM { &mut self, frame_data: &mut RuleFrameData, ) -> Result> { - let rule_info = self.get_rule_info(frame_data.rule_index)?; + let program = self.program.clone(); + let rule_info = program + .rule_infos + .get(usize::from(frame_data.rule_index)) + .ok_or(VmError::RuleInfoMissing { + index: frame_data.rule_index, + pc: self.pc, + available: program.rule_infos.len(), + })?; match frame_data.phase { RuleFramePhase::ExecutingDestructuring => { - self.rule_frame_after_destructuring_success(frame_data, &rule_info) + self.rule_frame_after_destructuring_success(frame_data, rule_info) } - RuleFramePhase::ExecutingBody => self.rule_frame_after_success(frame_data, &rule_info), + RuleFramePhase::ExecutingBody => self.rule_frame_after_success(frame_data, rule_info), RuleFramePhase::Initializing | RuleFramePhase::Finalizing => Ok(None), } } @@ -792,21 +800,16 @@ impl RegoVM { &mut self, frame_data: &mut RuleFrameData, ) -> Result> { - let rule_info = self.get_rule_info(frame_data.rule_index)?; - self.rule_frame_after_failure(frame_data, &rule_info) - } - - fn get_rule_info(&self, rule_index: u16) -> Result { - let idx = usize::from(rule_index); - self.program + let program = self.program.clone(); + let rule_info = program .rule_infos - .get(idx) - .cloned() + .get(usize::from(frame_data.rule_index)) .ok_or(VmError::RuleInfoMissing { - index: rule_index, + index: frame_data.rule_index, pc: self.pc, - available: self.program.rule_infos.len(), - }) + available: program.rule_infos.len(), + })?; + self.rule_frame_after_failure(frame_data, rule_info) } pub(super) fn checked_add_one(&self, value: usize, context: &'static str) -> Result { diff --git a/src/rvm/vm/virtual_data.rs b/src/rvm/vm/virtual_data.rs index bb66ef9..1212509 100644 --- a/src/rvm/vm/virtual_data.rs +++ b/src/rvm/vm/virtual_data.rs @@ -21,14 +21,15 @@ impl RegoVM { root_path.push(key_value); } - let mut data_subobject = self.data.clone(); + // Walk data tree by reference, only clone the leaf. + let mut data_ref = &self.data; for path_component in &root_path { - data_subobject = data_subobject[path_component].clone(); + data_ref = &data_ref[path_component]; } - let mut result_subobject = match data_subobject { + let mut result_subobject = match *data_ref { Value::Undefined => Value::new_object(), - _ => data_subobject, + _ => data_ref.clone(), }; self.traverse_rule_tree_subobject(rule_tree_subobject, &mut result_subobject, &root_path)?; @@ -74,11 +75,12 @@ impl RegoVM { result_subobject: &mut Value, root_path: &[Value], ) -> Result<()> { + let mut relative_path = Vec::new(); self.traverse_rule_tree_subobject_with_path( rule_tree_node, result_subobject, root_path, - &[], + &mut relative_path, ) } @@ -87,19 +89,18 @@ impl RegoVM { rule_tree_node: &Value, result_subobject: &mut Value, root_path: &[Value], - relative_path: &[Value], + relative_path: &mut Vec, ) -> Result<()> { match *rule_tree_node { Value::Number(ref rule_idx) => { if let Some(rule_index) = rule_idx.as_u64() { - let mut full_cache_path = root_path.to_vec(); - full_cache_path.extend_from_slice(relative_path); - + // Walk the evaluated cache using root_path then relative_path, + // without allocating a combined full_cache_path Vec. let cached_result = { let mut cache_lookup = &self.evaluated; let mut path_exists = true; - for path_component in &full_cache_path { + for path_component in root_path.iter().chain(relative_path.iter()) { if let Value::Object(ref map) = *cache_lookup { if let Some(next_value) = map.get(path_component) { cache_lookup = next_value; @@ -153,7 +154,15 @@ impl RegoVM { register_count, })?; - let mut cache_path = full_cache_path.clone(); + // Build cache_path from root_path + relative_path + Undefined sentinel. + let mut cache_path: Vec = Vec::with_capacity( + root_path + .len() + .saturating_add(relative_path.len()) + .saturating_add(1), + ); + cache_path.extend_from_slice(root_path); + cache_path.extend_from_slice(relative_path); cache_path.push(Value::Undefined); Self::set_nested_value_static( &mut self.evaluated, @@ -174,14 +183,14 @@ impl RegoVM { } Value::Object(ref obj) => { for (key, value) in obj.iter() { - let mut new_relative_path = relative_path.to_vec(); - new_relative_path.push(key.clone()); + relative_path.push(key.clone()); self.traverse_rule_tree_subobject_with_path( value, result_subobject, root_path, - &new_relative_path, + relative_path, )?; + relative_path.pop(); } } _ => {} @@ -232,15 +241,16 @@ impl RegoVM { self.execute_call_rule_common(params.dest, rule_index, None)?; if components_consumed < params.path_components.len() { - let mut rule_result = self.get_register(params.dest)?.clone(); + // Walk remaining path by reference, clone only the leaf. + let mut ref_val = self.get_register(params.dest)?; for component in params.path_components.iter().skip(components_consumed) { let key_value = self.literal_or_register_value(component)?; - - rule_result = rule_result[&key_value].clone(); + ref_val = &ref_val[&key_value]; } - self.set_register(params.dest, rule_result)?; + let leaf = ref_val.clone(); + self.set_register(params.dest, leaf)?; } } else { return Err(VmError::InvalidRuleIndex { @@ -252,14 +262,15 @@ impl RegoVM { Value::Undefined | Value::Object(_) if components_consumed != params.path_components.len() => { - let mut result = self.data.clone(); + // Walk data tree by reference, clone only the leaf. + let mut data_ref = &self.data; for component in ¶ms.path_components { let key_value = self.literal_or_register_value(component)?; - - result = result[&key_value].clone(); + data_ref = &data_ref[&key_value]; } + let result = data_ref.clone(); self.set_register(params.dest, result)?; } Value::Object(_) => {