regorus/
engine.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3#![allow(clippy::print_stderr)]
4
5use crate::ast::*;
6use crate::compiled_policy::CompiledPolicy;
7use crate::interpreter::*;
8use crate::lexer::*;
9use crate::parser::*;
10use crate::scheduler::*;
11use crate::utils::gather_functions;
12use crate::utils::limits::{self, fallback_execution_timer_config, ExecutionTimerConfig};
13use crate::value::*;
14use crate::*;
15use crate::{Extension, QueryResults};
16
17use crate::Rc;
18use anyhow::{anyhow, bail, Result};
19
20/// The Rego evaluation engine.
21///
22#[derive(Debug, Clone)]
23pub struct Engine {
24    modules: Rc<Vec<Ref<Module>>>,
25    interpreter: Interpreter,
26    prepared: bool,
27    rego_v1: bool,
28    execution_timer_config: Option<ExecutionTimerConfig>,
29}
30
31#[cfg(feature = "azure_policy")]
32#[derive(Debug, Clone, Serialize)]
33pub struct PolicyPackageNameDefinition {
34    pub source_file: String,
35    pub package_name: String,
36}
37
38#[cfg(feature = "azure_policy")]
39#[derive(Debug, Clone, Serialize)]
40pub struct PolicyParameter {
41    pub name: String,
42    pub modifiable: bool,
43    pub required: bool,
44}
45
46#[cfg(feature = "azure_policy")]
47#[derive(Debug, Clone, Serialize)]
48pub struct PolicyModifier {
49    pub name: String,
50}
51
52#[cfg(feature = "azure_policy")]
53#[derive(Debug, Clone, Serialize)]
54pub struct PolicyParameters {
55    pub source_file: String,
56    pub parameters: Vec<PolicyParameter>,
57    pub modifiers: Vec<PolicyModifier>,
58}
59
60/// Create a default engine.
61impl Default for Engine {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl Engine {
68    fn effective_execution_timer_config(&self) -> Option<ExecutionTimerConfig> {
69        self.execution_timer_config
70            .or_else(fallback_execution_timer_config)
71    }
72
73    fn apply_effective_execution_timer_config(&mut self) {
74        let config = self.effective_execution_timer_config();
75        self.interpreter.set_execution_timer_config(config);
76    }
77
78    /// Create an instance of [Engine].
79    pub fn new() -> Self {
80        let mut engine = Self {
81            modules: Rc::new(vec![]),
82            interpreter: Interpreter::new(),
83            prepared: false,
84            rego_v1: true,
85            execution_timer_config: None,
86        };
87        engine.apply_effective_execution_timer_config();
88        engine
89    }
90
91    /// Enable rego v0.
92    ///
93    /// Note that regorus now defaults to v1.
94    /// ```
95    /// # use regorus::*;
96    /// # fn main() -> anyhow::Result<()> {
97    /// let mut engine = Engine::new();
98    ///
99    /// // Enable v0 for old style policies.
100    /// engine.set_rego_v0(true);
101    ///
102    /// engine.add_policy(
103    ///    "test.rego".to_string(),
104    ///    r#"
105    ///    package test
106    ///
107    ///    allow { # v0 syntax does not require if keyword
108    ///       1 < 2
109    ///    }
110    ///    "#.to_string())?;
111    ///
112    /// # Ok(())
113    /// # }
114    /// ```
115    ///
116    pub const fn set_rego_v0(&mut self, rego_v0: bool) {
117        self.rego_v1 = !rego_v0;
118    }
119
120    /// Configure the execution timer.
121    ///
122    /// Stores the supplied configuration and ensures the next evaluation is checked against those
123    /// limits. Engines start without a time limit and otherwise fall back to the global
124    /// configuration (if provided).
125    ///
126    /// # Examples
127    ///
128    /// ```
129    /// use std::num::NonZeroU32;
130    /// use std::time::Duration;
131    /// use regorus::utils::limits::ExecutionTimerConfig;
132    /// use regorus::Engine;
133    ///
134    /// let mut engine = Engine::new();
135    /// let config = ExecutionTimerConfig {
136    ///     limit: Duration::from_millis(10),
137    ///     check_interval: NonZeroU32::new(1).unwrap(),
138    /// };
139    ///
140    /// engine.set_execution_timer_config(config);
141    /// ```
142    pub fn set_execution_timer_config(&mut self, config: ExecutionTimerConfig) {
143        self.execution_timer_config = Some(config);
144        self.interpreter.set_execution_timer_config(Some(config));
145    }
146
147    /// Clear the engine-specific execution timer configuration, falling back to the global value.
148    ///
149    /// # Examples
150    ///
151    /// ```
152    /// use std::num::NonZeroU32;
153    /// use std::time::Duration;
154    /// use regorus::utils::limits::{
155    ///     set_fallback_execution_timer_config,
156    ///     ExecutionTimerConfig,
157    /// };
158    /// use regorus::Engine;
159    ///
160    /// let mut engine = Engine::new();
161    /// let global = ExecutionTimerConfig {
162    ///     limit: Duration::from_millis(5),
163    ///     check_interval: NonZeroU32::new(1).unwrap(),
164    /// };
165    /// set_fallback_execution_timer_config(Some(global));
166    ///
167    /// engine.clear_execution_timer_config();
168    /// ```
169    pub fn clear_execution_timer_config(&mut self) {
170        self.execution_timer_config = None;
171        self.apply_effective_execution_timer_config();
172    }
173
174    /// Add a policy.
175    ///
176    /// The policy file will be parsed and converted to AST representation.
177    /// Multiple policy files may be added to the engine.
178    /// Returns the Rego package name declared in the policy.
179    ///
180    /// * `path`: A filename to be associated with the policy.
181    /// * `rego`: The rego policy code.
182    ///
183    /// ```
184    /// # use regorus::*;
185    /// # fn main() -> anyhow::Result<()> {
186    /// let mut engine = Engine::new();
187    ///
188    /// let package = engine.add_policy(
189    ///    "test.rego".to_string(),
190    ///    r#"
191    ///    package test
192    ///    allow = input.user == "root"
193    ///    "#.to_string())?;
194    ///
195    /// assert_eq!(package, "data.test");
196    /// # Ok(())
197    /// # }
198    /// ```
199    ///
200    pub fn add_policy(&mut self, path: String, rego: String) -> Result<String> {
201        let source = Source::from_contents(path, rego)?;
202        let mut parser = self.make_parser(&source)?;
203        let module = Ref::new(parser.parse()?);
204        limits::enforce_memory_limit().map_err(|err| anyhow!(err))?;
205        Rc::make_mut(&mut self.modules).push(module.clone());
206        // if policies change, interpreter needs to be prepared again
207        self.prepared = false;
208        Interpreter::get_path_string(&module.package.refr, Some("data"))
209    }
210
211    /// Add a policy from a given file.
212    ///
213    /// The policy file will be parsed and converted to AST representation.
214    /// Multiple policy files may be added to the engine.
215    /// Returns the Rego package name declared in the policy.
216    ///
217    /// * `path`: Path to the policy file (.rego).
218    ///
219    /// ```
220    /// # use regorus::*;
221    /// # fn main() -> anyhow::Result<()> {
222    /// let mut engine = Engine::new();
223    /// // framework.rego does not conform to v1.
224    /// engine.set_rego_v0(true);
225    ///
226    /// let package = engine.add_policy_from_file("tests/aci/framework.rego")?;
227    ///
228    /// assert_eq!(package, "data.framework");
229    /// # Ok(())
230    /// # }
231    /// ```
232    #[cfg(feature = "std")]
233    #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
234    pub fn add_policy_from_file<P: AsRef<std::path::Path>>(&mut self, path: P) -> Result<String> {
235        let source = Source::from_file(path)?;
236        let mut parser = self.make_parser(&source)?;
237        let module = Ref::new(parser.parse()?);
238        limits::enforce_memory_limit().map_err(|err| anyhow!(err))?;
239        Rc::make_mut(&mut self.modules).push(module.clone());
240        // if policies change, interpreter needs to be prepared again
241        self.prepared = false;
242        Interpreter::get_path_string(&module.package.refr, Some("data"))
243    }
244
245    /// Get the list of packages defined by loaded policies.
246    ///
247    /// ```
248    /// # use regorus::*;
249    /// # fn main() -> anyhow::Result<()> {
250    /// let mut engine = Engine::new();
251    /// // framework.rego does not conform to v1.
252    /// engine.set_rego_v0(true);
253    ///
254    /// let _ = engine.add_policy_from_file("tests/aci/framework.rego")?;
255    ///
256    /// // Package names can be different from file names.
257    /// let _ = engine.add_policy("policy.rego".into(), "package hello.world".into())?;
258    ///
259    /// assert_eq!(engine.get_packages()?, vec!["data.framework", "data.hello.world"]);
260    /// # Ok(())
261    /// # }
262    /// ```
263    pub fn get_packages(&self) -> Result<Vec<String>> {
264        self.modules
265            .iter()
266            .map(|m| Interpreter::get_path_string(&m.package.refr, Some("data")))
267            .collect()
268    }
269
270    /// Get the list of policy files.
271    /// ```
272    /// # use regorus::*;
273    /// # fn main() -> anyhow::Result<()> {
274    /// # let mut engine = Engine::new();
275    ///
276    /// let pkg = engine.add_policy("hello.rego".to_string(), "package test".to_string())?;
277    /// assert_eq!(pkg, "data.test");
278    ///
279    /// let policies = engine.get_policies()?;
280    ///
281    /// assert_eq!(policies[0].get_path(), "hello.rego");
282    /// assert_eq!(policies[0].get_contents(), "package test");
283    /// # Ok(())
284    /// # }
285    /// ```
286    pub fn get_policies(&self) -> Result<Vec<Source>> {
287        Ok(self
288            .modules
289            .iter()
290            .map(|m| m.package.refr.span().source.clone())
291            .collect())
292    }
293
294    /// Get the list of policy files as a JSON object.
295    /// ```
296    /// # use regorus::*;
297    /// # fn main() -> anyhow::Result<()> {
298    /// # let mut engine = Engine::new();
299    ///
300    /// let pkg = engine.add_policy("hello.rego".to_string(), "package test".to_string())?;
301    /// assert_eq!(pkg, "data.test");
302    ///
303    /// let policies = engine.get_policies_as_json()?;
304    ///
305    /// let v = Value::from_json_str(&policies)?;
306    /// assert_eq!(v[0]["path"].as_string()?.as_ref(), "hello.rego");
307    /// assert_eq!(v[0]["contents"].as_string()?.as_ref(), "package test");
308    /// # Ok(())
309    /// # }
310    /// ```
311    pub fn get_policies_as_json(&self) -> Result<String> {
312        #[derive(Serialize)]
313        struct Source<'a> {
314            path: &'a String,
315            contents: &'a String,
316        }
317
318        let mut sources = vec![];
319        for m in self.modules.iter() {
320            let source = &m.package.refr.span().source;
321            sources.push(Source {
322                path: source.get_path(),
323                contents: source.get_contents(),
324            });
325        }
326
327        serde_json::to_string_pretty(&sources).map_err(anyhow::Error::msg)
328    }
329
330    /// Set the input document.
331    ///
332    /// * `input`: Input documented. Typically this [Value] is constructed from JSON or YAML.
333    ///
334    /// ```
335    /// # use regorus::*;
336    /// # fn main() -> anyhow::Result<()> {
337    /// let mut engine = Engine::new();
338    ///
339    /// let input = Value::from_json_str(r#"
340    /// {
341    ///   "role" : "admin",
342    ///   "action": "delete"
343    /// }"#)?;
344    ///
345    /// engine.set_input(input);
346    /// # Ok(())
347    /// # }
348    /// ```
349    pub fn set_input(&mut self, input: Value) {
350        self.interpreter.set_input(input);
351    }
352
353    pub fn set_input_json(&mut self, input_json: &str) -> Result<()> {
354        self.set_input(Value::from_json_str(input_json)?);
355        Ok(())
356    }
357
358    /// Clear the data document.
359    ///
360    /// The data document will be reset to an empty object.
361    ///
362    /// ```
363    /// # use regorus::*;
364    /// # fn main() -> anyhow::Result<()> {
365    /// let mut engine = Engine::new();
366    ///
367    /// engine.clear_data();
368    ///
369    /// // Evaluate data.
370    /// let results = engine.eval_query("data".to_string(), false)?;
371    ///
372    /// // Assert that it is empty object.
373    /// assert_eq!(results.result.len(), 1);
374    /// assert_eq!(results.result[0].expressions.len(), 1);
375    /// assert_eq!(results.result[0].expressions[0].value, Value::new_object());
376    /// # Ok(())
377    /// # }
378    /// ```
379    pub fn clear_data(&mut self) {
380        self.interpreter.set_init_data(Value::new_object());
381        self.prepared = false;
382    }
383
384    /// Add data document.
385    ///
386    /// The specified data document is merged into existing data document.
387    ///
388    /// ```
389    /// # use regorus::*;
390    /// # fn main() -> anyhow::Result<()> {
391    /// let mut engine = Engine::new();
392    ///
393    /// // Only objects can be added.
394    /// assert!(engine.add_data(Value::from_json_str("[]")?).is_err());
395    ///
396    /// // Merge { "x" : 1, "y" : {} }
397    /// assert!(engine.add_data(Value::from_json_str(r#"{ "x" : 1, "y" : {}}"#)?).is_ok());
398    ///
399    /// // Merge { "z" : 2 }
400    /// assert!(engine.add_data(Value::from_json_str(r#"{ "z" : 2 }"#)?).is_ok());
401    ///
402    /// // Merge { "z" : 3 }. Conflict error.
403    /// assert!(engine.add_data(Value::from_json_str(r#"{ "z" : 3 }"#)?).is_err());
404    ///
405    /// assert_eq!(
406    ///   engine.eval_query("data".to_string(), false)?.result[0].expressions[0].value,
407    ///   Value::from_json_str(r#"{ "x": 1, "y": {}, "z": 2}"#)?
408    /// );
409    /// # Ok(())
410    /// # }
411    /// ```
412    pub fn add_data(&mut self, data: Value) -> Result<()> {
413        if data.as_object().is_err() {
414            bail!("data must be object");
415        }
416        self.prepared = false;
417        self.interpreter.get_init_data_mut().merge(data)
418    }
419
420    /// Get the data document.
421    ///
422    /// The returned value is the data document that has been constructed using
423    /// one or more calls to [`Engine::pre`]. The values of policy rules are
424    /// not included in the returned document.
425    ///
426    ///
427    /// ```
428    /// # use regorus::*;
429    /// # fn main() -> anyhow::Result<()> {
430    /// let mut engine = Engine::new();
431    ///
432    /// // If not set, data document is empty.
433    /// assert_eq!(engine.get_data(), Value::new_object());
434    ///
435    /// // Merge { "x" : 1, "y" : {} }
436    /// assert!(engine.add_data(Value::from_json_str(r#"{ "x" : 1, "y" : {}}"#)?).is_ok());
437    ///
438    /// // Merge { "z" : 2 }
439    /// assert!(engine.add_data(Value::from_json_str(r#"{ "z" : 2 }"#)?).is_ok());
440    ///
441    /// let data = engine.get_data();
442    /// assert_eq!(data["x"], Value::from(1));
443    /// assert_eq!(data["y"], Value::new_object());
444    /// assert_eq!(data["z"], Value::from(2));
445    ///
446    /// # Ok(())
447    /// # }
448    /// ```
449    pub fn get_data(&self) -> Value {
450        self.interpreter.get_init_data().clone()
451    }
452
453    pub fn add_data_json(&mut self, data_json: &str) -> Result<()> {
454        self.add_data(Value::from_json_str(data_json)?)
455    }
456
457    /// Set whether builtins should raise errors strictly or not.
458    ///
459    /// Regorus differs from OPA in that by default builtins will
460    /// raise errors instead of returning Undefined.
461    ///
462    /// ----
463    /// **_NOTE:_** Currently not all builtins honor this flag and will always strictly raise errors.
464    /// ----
465    pub fn set_strict_builtin_errors(&mut self, b: bool) {
466        self.interpreter.set_strict_builtin_errors(b);
467    }
468
469    #[doc(hidden)]
470    pub fn get_modules(&mut self) -> &Vec<Ref<Module>> {
471        &self.modules
472    }
473
474    /// Compiles a target-aware policy from the current engine state.
475    ///
476    /// This method creates a compiled policy that can work with Azure Policy targets,
477    /// enabling resource type inference and target-specific evaluation. The compiled
478    /// policy will automatically detect and handle `__target__` declarations in the
479    /// loaded modules.
480    ///
481    /// The engine must have been prepared with:
482    /// - Policy modules added via [`Engine::add_policy`]
483    /// - Data added via [`Engine::add_data`] (optional)
484    ///
485    /// # Returns
486    ///
487    /// Returns a [`CompiledPolicy`] that can be used for efficient policy evaluation
488    /// with target support, including resource type inference capabilities.
489    ///
490    /// # Examples
491    ///
492    /// ## Basic Target-Aware Compilation
493    ///
494    /// ```no_run
495    /// use regorus::*;
496    ///
497    /// # fn main() -> anyhow::Result<()> {
498    /// let mut engine = Engine::new();
499    /// engine.add_data(Value::from_json_str(r#"{"allowed_sizes": ["small", "medium"]}"#)?)?;
500    /// engine.add_policy("policy.rego".to_string(), r#"
501    ///     package policy.test
502    ///     import rego.v1
503    ///     __target__ := "target.tests.sample_test_target"
504    ///     
505    ///     default allow := false
506    ///     allow if {
507    ///         input.type == "vm"
508    ///         input.size in data.allowed_sizes
509    ///     }
510    /// "#.to_string())?;
511    ///
512    /// let compiled = engine.compile_for_target()?;
513    /// let result = compiled.eval_with_input(Value::from_json_str(r#"{"type": "vm", "size": "small"}"#)?)?;
514    /// # Ok(())
515    /// # }
516    /// ```
517    ///
518    /// ## Target Registration and Usage
519    ///
520    /// ```no_run
521    /// use regorus::*;
522    /// use regorus::registry::targets;
523    /// use regorus::target::Target;
524    /// use std::sync::Arc;
525    ///
526    /// # fn main() -> anyhow::Result<()> {
527    /// // Register a target first
528    /// let target_json = r#"
529    /// {
530    ///   "name": "target.example.vm_policy",
531    ///   "description": "Simple VM validation target",
532    ///   "version": "1.0.0",
533    ///   "resource_schema_selector": "type",
534    ///   "resource_schemas": [
535    ///     {
536    ///       "type": "object",
537    ///       "properties": {
538    ///         "name": { "type": "string" },
539    ///         "type": { "const": "vm" },
540    ///         "size": { "enum": ["small", "medium", "large"] }
541    ///       },
542    ///       "required": ["name", "type", "size"]
543    ///     }
544    ///   ],
545    ///   "effects": {
546    ///     "allow": { "type": "boolean" },
547    ///     "deny": { "type": "boolean" }
548    ///   }
549    /// }
550    /// "#;
551    ///
552    /// let target = Target::from_json_str(target_json)?;
553    /// targets::register(Arc::new(target))?;
554    ///
555    /// // Use the target in a policy
556    /// let mut engine = Engine::new();
557    /// engine.add_data(Value::from_json_str(r#"{"allowed_locations": ["us-east"]}"#)?)?;
558    /// engine.add_policy("vm_policy.rego".to_string(), r#"
559    ///     package vm.validation
560    ///     import rego.v1
561    ///     __target__ := "target.example.vm_policy"
562    ///     
563    ///     default allow := false
564    ///     allow if {
565    ///         input.type == "vm"
566    ///         input.size in ["small", "medium"]
567    ///     }
568    /// "#.to_string())?;
569    ///
570    /// let compiled = engine.compile_for_target()?;
571    /// let result = compiled.eval_with_input(Value::from_json_str(r#"
572    /// {
573    ///   "name": "test-vm",
574    ///   "type": "vm",
575    ///   "size": "small"
576    /// }"#)?)?;
577    /// assert_eq!(result, Value::from(true));
578    /// # Ok(())
579    /// # }
580    /// ```
581    ///
582    /// # Notes
583    ///
584    /// - This method is only available when the `azure_policy` feature is enabled
585    /// - Automatically enables print gathering for debugging purposes
586    /// - Requires that at least one module contains a `__target__` declaration
587    /// - The target referenced must be registered in the target registry
588    ///
589    /// # See Also
590    ///
591    /// - [`Engine::compile_with_entrypoint`] for explicit rule-based compilation
592    /// - [`crate::compile_policy_for_target`] for a higher-level convenience function
593    #[cfg(feature = "azure_policy")]
594    #[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))]
595    pub fn compile_for_target(&mut self) -> Result<CompiledPolicy> {
596        self.prepare_for_eval(false, true)?;
597        self.apply_effective_execution_timer_config();
598        self.interpreter.clean_internal_evaluation_state();
599        self.interpreter.compile(None).map(CompiledPolicy::new)
600    }
601
602    /// Compiles a policy with a specific entry point rule.
603    ///
604    /// This method creates a compiled policy that evaluates a specific rule as the entry point.
605    /// Unlike [`Engine::compile_for_target`], this method requires you to explicitly specify which
606    /// rule should be evaluated and does not automatically handle target-specific features.
607    ///
608    /// The engine must have been prepared with:
609    /// - Policy modules added via [`Engine::add_policy`]
610    /// - Data added via [`Engine::add_data`] (optional)
611    ///
612    /// # Arguments
613    ///
614    /// * `rule` - The specific rule path to evaluate (e.g., "data.policy.allow")
615    ///
616    /// # Returns
617    ///
618    /// Returns a [`CompiledPolicy`] that can be used for efficient policy evaluation
619    /// focused on the specified entry point rule.
620    ///
621    /// # Examples
622    ///
623    /// ## Basic Usage
624    ///
625    /// ```no_run
626    /// use regorus::*;
627    /// use std::rc::Rc;
628    ///
629    /// # fn main() -> anyhow::Result<()> {
630    /// let mut engine = Engine::new();
631    /// engine.add_data(Value::from_json_str(r#"{"allowed_users": ["alice", "bob"]}"#)?)?;
632    /// engine.add_policy("authz.rego".to_string(), r#"
633    ///     package authz
634    ///     import rego.v1
635    ///     
636    ///     default allow := false
637    ///     allow if {
638    ///         input.user in data.allowed_users
639    ///         input.action == "read"
640    ///     }
641    ///     
642    ///     deny if {
643    ///         input.user == "guest"
644    ///     }
645    /// "#.to_string())?;
646    ///
647    /// let compiled = engine.compile_with_entrypoint(&"data.authz.allow".into())?;
648    /// let result = compiled.eval_with_input(Value::from_json_str(r#"{"user": "alice", "action": "read"}"#)?)?;
649    /// assert_eq!(result, Value::from(true));
650    /// # Ok(())
651    /// # }
652    /// ```
653    ///
654    /// ## Multi-Module Policy
655    ///
656    /// ```no_run
657    /// use regorus::*;
658    /// use std::rc::Rc;
659    ///
660    /// # fn main() -> anyhow::Result<()> {
661    /// let mut engine = Engine::new();
662    /// engine.add_data(Value::from_json_str(r#"{"departments": {"engineering": ["alice"], "hr": ["bob"]}}"#)?)?;
663    ///
664    /// engine.add_policy("users.rego".to_string(), r#"
665    ///     package users
666    ///     import rego.v1
667    ///     
668    ///     user_department(user) := dept if {
669    ///         dept := [d | data.departments[d][_] == user][0]
670    ///     }
671    /// "#.to_string())?;
672    ///
673    /// engine.add_policy("permissions.rego".to_string(), r#"
674    ///     package permissions
675    ///     import rego.v1
676    ///     import data.users
677    ///     
678    ///     default allow := false
679    ///     allow if {
680    ///         users.user_department(input.user) == "engineering"
681    ///         input.resource.type == "code"
682    ///     }
683    ///     
684    ///     allow if {
685    ///         users.user_department(input.user) == "hr"
686    ///         input.resource.type == "personnel_data"
687    ///     }
688    /// "#.to_string())?;
689    ///
690    /// let compiled = engine.compile_with_entrypoint(&"data.permissions.allow".into())?;
691    ///
692    /// // Test engineering access to code
693    /// let result = compiled.eval_with_input(Value::from_json_str(r#"
694    /// {
695    ///   "user": "alice",
696    ///   "resource": {"type": "code", "name": "main.rs"}
697    /// }"#)?)?;
698    /// assert_eq!(result, Value::from(true));
699    /// # Ok(())
700    /// # }
701    /// ```
702    ///
703    /// # Entry Point Rule Format
704    ///
705    /// The `rule` parameter should follow the Rego rule path format:
706    /// - `"data.package.rule"` - For rules in a specific package
707    /// - `"data.package.subpackage.rule"` - For nested packages
708    /// - `"allow"` - For rules in the default package (though this is not recommended)
709    ///
710    /// # Notes
711    ///
712    /// - Automatically enables print gathering for debugging purposes
713    /// - If you need target-aware compilation with automatic `__target__` handling,
714    ///   consider using [`Engine::compile_for_target`] instead (requires `azure_policy` feature)
715    ///
716    /// # See Also
717    ///
718    /// - [`Engine::compile_for_target`] for target-aware compilation
719    /// - [`crate::compile_policy_with_entrypoint`] for a higher-level convenience function
720    pub fn compile_with_entrypoint(&mut self, rule: &Rc<str>) -> Result<CompiledPolicy> {
721        self.prepare_for_eval(false, false)?;
722        self.apply_effective_execution_timer_config();
723        self.interpreter.clean_internal_evaluation_state();
724        self.interpreter
725            .compile(Some(rule.clone()))
726            .map(CompiledPolicy::new)
727    }
728
729    /// Evaluate specified rule(s).
730    ///
731    /// [`Engine::eval_rule`] is often faster than [`Engine::eval_query`] and should be preferred if
732    /// OPA style [`QueryResults`] are not needed.
733    ///
734    /// ```
735    /// # use regorus::*;
736    /// # fn main() -> anyhow::Result<()> {
737    /// let mut engine = Engine::new();
738    ///
739    /// // Add policy
740    /// engine.add_policy(
741    ///   "policy.rego".to_string(),
742    ///   r#"
743    ///   package example
744    ///   import rego.v1
745    ///
746    ///   x = [1, 2]
747    ///
748    ///   y := 5 if input.a > 2
749    ///   "#.to_string())?;
750    ///
751    /// // Evaluate rule.
752    /// let v = engine.eval_rule("data.example.x".to_string())?;
753    /// assert_eq!(v, Value::from(vec![Value::from(1), Value::from(2)]));
754    ///
755    /// // y evaluates to undefined.
756    /// let v = engine.eval_rule("data.example.y".to_string())?;
757    /// assert_eq!(v, Value::Undefined);
758    ///
759    /// // Evaluating a non-existent rule is an error.
760    /// let r = engine.eval_rule("data.exaample.x".to_string());
761    /// assert!(r.is_err());
762    ///
763    /// // Path must be valid rule paths.
764    /// assert!( engine.eval_rule("data".to_string()).is_err());
765    /// assert!( engine.eval_rule("data.example".to_string()).is_err());
766    /// # Ok(())
767    /// # }
768    /// ```
769    pub fn eval_rule(&mut self, rule: String) -> Result<Value> {
770        self.prepare_for_eval(false, false)?;
771        self.apply_effective_execution_timer_config();
772        self.interpreter.clean_internal_evaluation_state();
773        self.interpreter.eval_rule_in_path(rule)
774    }
775
776    /// Evaluate a Rego query.
777    ///
778    /// ```
779    /// # use regorus::*;
780    /// # fn main() -> anyhow::Result<()> {
781    /// let mut engine = Engine::new();
782    ///
783    /// // Add policies
784    /// engine.set_rego_v0(true);
785    /// engine.add_policy_from_file("tests/aci/framework.rego")?;
786    /// engine.add_policy_from_file("tests/aci/api.rego")?;
787    /// engine.add_policy_from_file("tests/aci/policy.rego")?;
788    ///
789    /// // Add data document (if any).
790    /// // If multiple data documents can be added, they will be merged together.
791    /// engine.add_data(Value::from_json_file("tests/aci/data.json")?)?;
792    ///
793    /// // At this point the policies and data have been loaded.
794    /// // Either the same engine can be used to make multiple queries or the engine
795    /// // can be cloned to avoid having the reload the policies and data.
796    /// let _clone = engine.clone();
797    ///
798    /// // Evaluate a query.
799    /// // Load input and make query.
800    /// engine.set_input(Value::new_object());
801    /// let results = engine.eval_query("data.framework.mount_overlay.allowed".to_string(), false)?;
802    /// assert_eq!(results.result[0].expressions[0].value, Value::from(false));
803    ///
804    /// // Evaluate query with different inputs.
805    /// engine.set_input(Value::from_json_file("tests/aci/input.json")?);
806    /// let results = engine.eval_query("data.framework.mount_overlay.allowed".to_string(), false)?;
807    /// assert_eq!(results.result[0].expressions[0].value, Value::from(true));
808    /// # Ok(())
809    /// # }
810    /// ```
811    pub fn eval_query(&mut self, query: String, enable_tracing: bool) -> Result<QueryResults> {
812        self.prepare_for_eval(enable_tracing, false)?;
813        self.apply_effective_execution_timer_config();
814        self.interpreter.clean_internal_evaluation_state();
815
816        self.interpreter.create_rule_prefixes()?;
817        let (query_module, query_node, query_schedule) = self.make_query(query)?;
818        if query_node.span.text() == "data" {
819            self.eval_modules(enable_tracing)?;
820        }
821
822        self.interpreter
823            .eval_user_query(&query_module, &query_node, query_schedule, enable_tracing)
824    }
825
826    /// Evaluate a Rego query that produces a boolean value.
827    ///
828    ///
829    /// This function should be preferred over [`Engine::eval_query`] if just a `true`/`false`
830    /// value is desired instead of [`QueryResults`].
831    ///
832    /// ```
833    /// # use regorus::*;
834    /// # fn main() -> anyhow::Result<()> {
835    /// # let mut engine = Engine::new();
836    ///
837    /// let enable_tracing = false;
838    /// assert_eq!(engine.eval_bool_query("1 > 2".to_string(), enable_tracing)?, false);
839    /// assert_eq!(engine.eval_bool_query("1 < 2".to_string(), enable_tracing)?, true);
840    ///
841    /// // Non boolean queries will raise an error.
842    /// assert!(engine.eval_bool_query("1+1".to_string(), enable_tracing).is_err());
843    ///
844    /// // Queries producing multiple values will raise an error.
845    /// assert!(engine.eval_bool_query("true; true".to_string(), enable_tracing).is_err());
846    ///
847    /// // Queries producing no values will raise an error.
848    /// assert!(engine.eval_bool_query("true; false; true".to_string(), enable_tracing).is_err());
849    /// # Ok(())
850    /// # }
851    /// ```
852    pub fn eval_bool_query(&mut self, query: String, enable_tracing: bool) -> Result<bool> {
853        let results = self.eval_query(query, enable_tracing)?;
854        let entries = results.result.as_slice();
855        let entry = entries
856            .first()
857            .ok_or_else(|| anyhow!("query did not produce any values"))?;
858        if entries.len() > 1 {
859            bail!("query produced more than one value");
860        }
861
862        let expressions = entry.expressions.as_slice();
863        let expr = expressions
864            .first()
865            .ok_or_else(|| anyhow!("query result missing expression"))?;
866        if expressions.len() > 1 {
867            bail!("query produced more than one value");
868        }
869
870        expr.value.as_bool().copied()
871    }
872
873    /// Evaluate an `allow` query.
874    ///
875    /// This is a wrapper over [`Engine::eval_bool_query`] that returns true only if the
876    /// boolean query succeed and produced a `true` value.
877    ///
878    /// ```
879    /// # use regorus::*;
880    /// # fn main() -> anyhow::Result<()> {
881    /// # let mut engine = Engine::new();
882    ///
883    /// let enable_tracing = false;
884    /// assert_eq!(engine.eval_allow_query("1 > 2".to_string(), enable_tracing), false);
885    /// assert_eq!(engine.eval_allow_query("1 < 2".to_string(), enable_tracing), true);
886    /// assert_eq!(engine.eval_allow_query("1+1".to_string(), enable_tracing), false);
887    /// assert_eq!(engine.eval_allow_query("true; true".to_string(), enable_tracing), false);
888    /// assert_eq!(engine.eval_allow_query("true; false; true".to_string(), enable_tracing), false);
889    /// # Ok(())
890    /// # }
891    /// ```
892    pub fn eval_allow_query(&mut self, query: String, enable_tracing: bool) -> bool {
893        matches!(self.eval_bool_query(query, enable_tracing), Ok(true))
894    }
895
896    /// Evaluate a `deny` query.
897    ///
898    /// This is a wrapper over [`Engine::eval_bool_query`] that returns false only if the
899    /// boolean query succeed and produced a `false` value.
900    /// ```
901    /// # use regorus::*;
902    /// # fn main() -> anyhow::Result<()> {
903    /// # let mut engine = Engine::new();
904    ///
905    /// let enable_tracing = false;
906    /// assert_eq!(engine.eval_deny_query("1 > 2".to_string(), enable_tracing), false);
907    /// assert_eq!(engine.eval_deny_query("1 < 2".to_string(), enable_tracing), true);
908    ///
909    /// assert_eq!(engine.eval_deny_query("1+1".to_string(), enable_tracing), true);
910    /// assert_eq!(engine.eval_deny_query("true; true".to_string(), enable_tracing), true);
911    /// assert_eq!(engine.eval_deny_query("true; false; true".to_string(), enable_tracing), true);
912    /// # Ok(())
913    /// # }
914    /// ```
915    pub fn eval_deny_query(&mut self, query: String, enable_tracing: bool) -> bool {
916        !matches!(self.eval_bool_query(query, enable_tracing), Ok(false))
917    }
918
919    fn make_query(&mut self, query: String) -> Result<(NodeRef<Module>, NodeRef<Query>, Schedule)> {
920        let mut query_module = {
921            let source = Source::from_contents(
922                "<query_module.rego>".to_owned(),
923                "package __internal_query_module".to_owned(),
924            )?;
925            Parser::new(&source)?.parse()?
926        };
927
928        // Parse the query.
929        let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
930        let mut parser = self.make_parser(&query_source)?;
931        let query_node = parser.parse_user_query()?;
932        query_module.num_expressions = parser.num_expressions();
933        query_module.num_queries = parser.num_queries();
934        query_module.num_statements = parser.num_statements();
935        let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;
936
937        // Populate loop hoisting for the query snippet
938        // Query snippets are treated as if they're in a module appended at the end (same as analyzer)
939        // The loop hoisting table already has capacity for this (ensured in prepare_for_eval)
940        let module_idx = u32::try_from(self.modules.len())
941            .map_err(|_| anyhow!("module count exceeds u32::MAX"))?;
942
943        use crate::compiler::hoist::LoopHoister;
944
945        let query_schedule_rc = Rc::new(query_schedule.clone());
946
947        // Run loop hoisting for query snippet
948        let mut hoister = LoopHoister::new_with_schedule(query_schedule_rc.clone());
949        hoister.populate_query_snippet(
950            module_idx,
951            &query_node,
952            query_module.num_statements,
953            query_module.num_expressions,
954        )?;
955        let query_lookup = hoister.finalize();
956
957        #[cfg(debug_assertions)]
958        {
959            for stmt in &query_node.stmts {
960                debug_assert!(
961                    query_lookup
962                        .get_statement_loops(module_idx, stmt.sidx)
963                        .ok()
964                        .and_then(|entry| entry)
965                        .is_some(),
966                    "missing hoisted loop entry for query statement index {}",
967                    stmt.sidx
968                );
969            }
970        }
971
972        // Get the existing table, merge in the query loops, and set it back
973        let mut existing_table = self.interpreter.take_loop_hoisting_table();
974        existing_table.truncate_modules(self.modules.len());
975        #[cfg(debug_assertions)]
976        {
977            debug_assert!(
978                existing_table.module_len() <= self.modules.len(),
979                "loop hoisting table should not retain extra modules before merge"
980            );
981        }
982        existing_table.merge_query_loops(query_lookup, self.modules.len());
983        #[cfg(debug_assertions)]
984        {
985            for stmt in &query_node.stmts {
986                debug_assert!(
987                    existing_table
988                        .get_statement_loops(module_idx, stmt.sidx)
989                        .ok()
990                        .and_then(|entry| entry)
991                        .is_some(),
992                    "missing hoisted loop entry after merge for module {} stmt {}",
993                    module_idx,
994                    stmt.sidx
995                );
996            }
997        }
998        self.interpreter.set_loop_hoisting_table(existing_table);
999
1000        Ok((Ref::new(query_module), query_node, query_schedule))
1001    }
1002
1003    #[doc(hidden)]
1004    /// Evaluate the given query and all the rules in the supplied policies.
1005    ///
1006    /// This is mainly used for testing Regorus itself.
1007    pub fn eval_query_and_all_rules(
1008        &mut self,
1009        query: String,
1010        enable_tracing: bool,
1011    ) -> Result<QueryResults> {
1012        self.eval_modules(enable_tracing)?;
1013        // Restart the timer window for the user query after module evaluation.
1014        self.apply_effective_execution_timer_config();
1015
1016        let (query_module, query_node, query_schedule) = self.make_query(query)?;
1017        self.interpreter
1018            .eval_user_query(&query_module, &query_node, query_schedule, enable_tracing)
1019    }
1020
1021    #[doc(hidden)]
1022    fn prepare_for_eval(&mut self, enable_tracing: bool, for_target: bool) -> Result<()> {
1023        // Fail fast if the engine already exceeds the global memory limit before evaluation work.
1024        limits::enforce_memory_limit().map_err(|err| anyhow!(err))?;
1025
1026        self.interpreter.set_traces(enable_tracing);
1027
1028        // if the data/policies have changed or the interpreter has never been prepared
1029        if !self.prepared {
1030            // Analyze the modules and determine how statements must be scheduled.
1031            let analyzer = Analyzer::new();
1032            let schedule = Rc::new(analyzer.analyze(&self.modules)?);
1033
1034            self.interpreter.set_modules(self.modules.clone());
1035
1036            self.interpreter.clear_builtins_cache();
1037            // clean_internal_evaluation_state will set data to an efficient clont of use supplied init_data
1038            // Initialize the with-document with initial data values.
1039            // with-modifiers will be applied to this document.
1040            self.interpreter.init_with_document()?;
1041
1042            self.interpreter
1043                .set_functions(gather_functions(&self.modules)?);
1044            self.interpreter.gather_rules()?;
1045            self.interpreter.process_imports()?;
1046
1047            // Populate loop hoisting table for efficient evaluation
1048            // Reserve capacity for 1 extra module (for query modules)
1049            use crate::compiler::hoist::LoopHoister;
1050
1051            // Run loop hoisting pass first
1052            let hoister = LoopHoister::new_with_schedule(schedule.clone());
1053            let loop_lookup = hoister.populate_with_extra_capacity(&self.modules, 0)?;
1054
1055            self.interpreter.set_loop_hoisting_table(loop_lookup);
1056
1057            // Set schedule after hoisting completes
1058            self.interpreter.set_schedule(Some(schedule));
1059
1060            #[cfg(feature = "azure_policy")]
1061            if for_target {
1062                // Resolve and validate target specifications across all modules
1063                crate::interpreter::target::resolve::resolve_and_apply_target(
1064                    &mut self.interpreter,
1065                )?;
1066                // Infer resource types
1067                crate::interpreter::target::infer::infer_resource_type(&mut self.interpreter)?;
1068            }
1069
1070            if !for_target {
1071                // Check if any module specifies a target and warn if so
1072                #[cfg(feature = "azure_policy")]
1073                self.warn_if_targets_present();
1074            }
1075
1076            self.prepared = true;
1077        }
1078
1079        Ok(())
1080    }
1081
1082    #[doc(hidden)]
1083    pub fn eval_rule_in_module(
1084        &mut self,
1085        module: &Ref<Module>,
1086        rule: &Ref<Rule>,
1087        enable_tracing: bool,
1088    ) -> Result<Value> {
1089        self.prepare_for_eval(enable_tracing, false)?;
1090        self.apply_effective_execution_timer_config();
1091        self.interpreter.clean_internal_evaluation_state();
1092
1093        self.interpreter.eval_rule(module, rule)?;
1094
1095        Ok(self.interpreter.get_data_mut().clone())
1096    }
1097
1098    #[doc(hidden)]
1099    pub fn eval_modules(&mut self, enable_tracing: bool) -> Result<Value> {
1100        self.prepare_for_eval(enable_tracing, false)?;
1101        self.apply_effective_execution_timer_config();
1102        self.interpreter.clean_internal_evaluation_state();
1103
1104        // Ensure that empty modules are created.
1105        for m in self.modules.iter().filter(|m| m.policy.is_empty()) {
1106            let path = Parser::get_path_ref_components(&m.package.refr)?;
1107            let path: Vec<&str> = path.iter().map(|s| s.text()).collect();
1108            let vref =
1109                Interpreter::make_or_get_value_mut(self.interpreter.get_data_mut(), &path[..])?;
1110            if *vref == Value::Undefined {
1111                *vref = Value::new_object();
1112            }
1113        }
1114
1115        self.interpreter.check_default_rules()?;
1116        for module in self.modules.clone().iter() {
1117            for rule in &module.policy {
1118                self.interpreter.eval_rule(module, rule)?;
1119            }
1120        }
1121        // Defer the evaluation of the default rules to here
1122        for module in self.modules.clone().iter() {
1123            let prev_module = self.interpreter.set_current_module(Some(module.clone()))?;
1124            for rule in &module.policy {
1125                self.interpreter.eval_default_rule(rule)?;
1126            }
1127            self.interpreter.set_current_module(prev_module)?;
1128        }
1129
1130        // Ensure that all modules are created.
1131        for m in self.modules.iter() {
1132            let path = Parser::get_path_ref_components(&m.package.refr)?;
1133            let path: Vec<&str> = path.iter().map(|s| s.text()).collect();
1134            let vref =
1135                Interpreter::make_or_get_value_mut(self.interpreter.get_data_mut(), &path[..])?;
1136            if *vref == Value::Undefined {
1137                *vref = Value::new_object();
1138            }
1139        }
1140        self.interpreter.create_rule_prefixes()?;
1141        Ok(self.interpreter.get_data_mut().clone())
1142    }
1143
1144    /// Add a custom builtin (extension).
1145    ///
1146    /// * `path`: The fully qualified path of the builtin.
1147    /// * `nargs`: The number of arguments the builtin takes.
1148    /// * `extension`: The [`Extension`] instance.
1149    ///
1150    /// ```rust
1151    /// # use regorus::*;
1152    /// # use anyhow::{bail, Result};
1153    /// # fn main() -> Result<()> {
1154    /// let mut engine = Engine::new();
1155    ///
1156    /// // Policy uses `do_magic` custom builtin.
1157    /// engine.add_policy(
1158    ///    "test.rego".to_string(),
1159    ///    r#"package test
1160    ///       x = do_magic(1)
1161    ///    "#.to_string(),
1162    /// )?;
1163    ///
1164    /// // Evaluating fails since `do_magic` is not defined.
1165    /// assert!(engine.eval_query("data.test.x".to_string(), false).is_err());
1166    ///
1167    /// // Add extension to implement `do_magic`. The extension can be stateful.
1168    /// let mut magic = 8;
1169    /// engine.add_extension("do_magic".to_string(), 1 , Box::new(move | mut params: Vec<Value> | {
1170    ///   // params is mut and therefore individual values can be removed from it and modified.
1171    ///   // The number of parameters (1) has already been validated.
1172    ///
1173    ///   match &params[0].as_i64() {
1174    ///      Ok(i) => {
1175    ///         // Compute value
1176    ///         let v = *i + magic;
1177    ///         // Update extension state.
1178    ///         magic += 1;
1179    ///         Ok(Value::from(v))
1180    ///      }
1181    ///      // Extensions can raise errors. Regorus will add location information to
1182    ///      // the error.
1183    ///      _ => bail!("do_magic expects i64 value")
1184    ///   }
1185    /// }))?;
1186    ///
1187    /// // Evaluation will now succeed.
1188    /// let r = engine.eval_query("data.test.x".to_string(), false)?;
1189    /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 9);
1190    ///
1191    /// // Cloning the engine will also clone the extension.
1192    /// let mut engine1 = engine.clone();
1193    ///
1194    /// // Evaluating again will return a different value since the extension is stateful.
1195    /// let r = engine.eval_query("data.test.x".to_string(), false)?;
1196    /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 10);
1197    ///
1198    /// // The second engine has a clone of the extension.
1199    /// let r = engine1.eval_query("data.test.x".to_string(), false)?;
1200    /// assert_eq!(r.result[0].expressions[0].value.as_i64()?, 10);
1201    ///
1202    /// // Once added, the extension cannot be replaced or removed.
1203    /// assert!(engine.add_extension("do_magic".to_string(), 1, Box::new(|_:Vec<Value>| {
1204    ///   Ok(Value::Undefined)
1205    /// })).is_err());
1206    ///
1207    /// // Extensions don't support out-parameter syntax.
1208    /// engine.add_policy(
1209    ///   "policy.rego".to_string(),
1210    ///   r#"package invalid
1211    ///      x = y if {
1212    ///       # y = do_magic(2)
1213    ///       do_magic(2, y)  # y is supplied as an out parameter.
1214    ///     }
1215    ///    "#.to_string()
1216    /// )?;
1217    ///
1218    /// // Evaluation fails since rule x calls an extension with out parameter.
1219    /// assert!(engine.eval_query("data.invalid.x".to_string(), false).is_err());
1220    /// # Ok(())
1221    /// # }
1222    /// ```
1223    pub fn add_extension(
1224        &mut self,
1225        path: String,
1226        nargs: u8,
1227        extension: Box<dyn Extension>,
1228    ) -> Result<()> {
1229        self.interpreter.add_extension(path, nargs, extension)
1230    }
1231
1232    #[cfg(feature = "coverage")]
1233    #[cfg_attr(docsrs, doc(cfg(feature = "coverage")))]
1234    /// Get the coverage report.
1235    ///
1236    /// ```rust
1237    /// # use regorus::*;
1238    /// # use anyhow::{bail, Result};
1239    /// # fn main() -> Result<()> {
1240    /// let mut engine = Engine::new();
1241    ///
1242    /// engine.add_policy(
1243    ///    "policy.rego".to_string(),
1244    ///    r#"
1245    /// package test    # Line 2
1246    ///
1247    /// x = y if {         # Line 4
1248    ///   input.a > 2   # Line 5
1249    ///   y = 5         # Line 6
1250    /// }
1251    ///    "#.to_string()
1252    /// )?;
1253    ///
1254    /// // Enable coverage.
1255    /// engine.set_enable_coverage(true);
1256    ///
1257    /// engine.eval_query("data".to_string(), false)?;
1258    ///
1259    /// let report = engine.get_coverage_report()?;
1260    /// assert_eq!(report.files[0].path, "policy.rego");
1261    ///
1262    /// // Only line 5 is evaluated.
1263    /// assert_eq!(report.files[0].covered.iter().cloned().collect::<Vec<u32>>(), vec![5]);
1264    ///
1265    /// // Line 4 and 6 are not evaluated.
1266    /// assert_eq!(report.files[0].not_covered.iter().cloned().collect::<Vec<u32>>(), vec![4, 6]);
1267    /// # Ok(())
1268    /// # }
1269    /// ```
1270    ///
1271    /// See also [`crate::coverage::Report::to_colored_string`].
1272    pub fn get_coverage_report(&self) -> Result<crate::coverage::Report> {
1273        self.interpreter.get_coverage_report()
1274    }
1275
1276    #[cfg(feature = "coverage")]
1277    #[cfg_attr(docsrs, doc(cfg(feature = "coverage")))]
1278    /// Enable/disable policy coverage.
1279    ///
1280    /// If `enable` is different from the current value, then any existing coverage
1281    /// information will be cleared.
1282    pub fn set_enable_coverage(&mut self, enable: bool) {
1283        self.interpreter.set_enable_coverage(enable);
1284    }
1285
1286    #[cfg(feature = "coverage")]
1287    #[cfg_attr(docsrs, doc(cfg(feature = "coverage")))]
1288    /// Clear the gathered policy coverage data.
1289    pub fn clear_coverage_data(&mut self) {
1290        self.interpreter.clear_coverage_data();
1291    }
1292
1293    /// Gather output from print statements instead of emiting to stderr.
1294    ///
1295    /// See [`Engine::take_prints`].
1296    pub fn set_gather_prints(&mut self, b: bool) {
1297        self.interpreter.set_gather_prints(b);
1298    }
1299
1300    /// Take the gathered output of print statements.
1301    ///
1302    /// ```rust
1303    /// # use regorus::*;
1304    /// # use anyhow::{bail, Result};
1305    /// # fn main() -> Result<()> {
1306    /// let mut engine = Engine::new();
1307    ///
1308    /// // Print to stderr.
1309    /// engine.eval_query("print(\"Hello\")".to_string(), false)?;
1310    ///
1311    /// // Configure gathering print statements.
1312    /// engine.set_gather_prints(true);
1313    ///
1314    /// // Execute query.
1315    /// engine.eval_query("print(\"Hello\")".to_string(), false)?;
1316    ///
1317    /// // Take and clear prints.
1318    /// let prints = engine.take_prints()?;
1319    /// assert_eq!(prints.len(), 1);
1320    /// assert!(prints[0].contains("Hello"));
1321    ///
1322    /// for p in prints {
1323    ///   println!("{p}");
1324    /// }
1325    /// # Ok(())
1326    /// # }
1327    /// ```
1328    pub fn take_prints(&mut self) -> Result<Vec<String>> {
1329        self.interpreter.take_prints()
1330    }
1331
1332    /// Get the policies and corresponding AST.
1333    ///
1334    ///
1335    /// ```rust
1336    /// # use regorus::*;
1337    /// # use anyhow::{bail, Result};
1338    /// # fn main() -> Result<()> {
1339    /// # let mut engine = Engine::new();
1340    /// engine.add_policy("test.rego".to_string(), "package test\n x := 1".to_string())?;
1341    ///
1342    /// let ast = engine.get_ast_as_json()?;
1343    /// let value = Value::from_json_str(&ast)?;
1344    ///
1345    /// assert_eq!(value[0]["ast"]["package"]["refr"]["Var"][1].as_string()?.as_ref(), "test");
1346    /// # Ok(())
1347    /// # }
1348    /// ```
1349    #[cfg(feature = "ast")]
1350    #[cfg_attr(docsrs, doc(cfg(feature = "ast")))]
1351    pub fn get_ast_as_json(&self) -> Result<String> {
1352        #[derive(Serialize)]
1353        struct Policy<'a> {
1354            source: &'a Source,
1355            version: u32,
1356            ast: &'a Module,
1357        }
1358        let mut ast = vec![];
1359        for m in self.modules.iter() {
1360            ast.push(Policy {
1361                source: &m.package.span.source,
1362                version: 1,
1363                ast: m,
1364            });
1365        }
1366
1367        serde_json::to_string_pretty(&ast).map_err(anyhow::Error::msg)
1368    }
1369
1370    /// Get the package names of each policy added to the engine.
1371    ///
1372    ///
1373    /// ```rust
1374    /// # use regorus::*;
1375    /// # use anyhow::{bail, Result};
1376    /// # fn main() -> Result<()> {
1377    /// # let mut engine = Engine::new();
1378    /// engine.add_policy("test.rego".to_string(), "package test\n x := 1".to_string())?;
1379    /// engine.add_policy("test2.rego".to_string(), "package test.multi.segment\n x := 1".to_string())?;
1380    ///
1381    /// let package_names = engine.get_policy_package_names()?;
1382    ///
1383    /// assert_eq!("test", package_names[0].package_name);
1384    /// assert_eq!("test.multi.segment", package_names[1].package_name);
1385    /// # Ok(())
1386    /// # }
1387    /// ```
1388    #[cfg(feature = "azure_policy")]
1389    #[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))]
1390    pub fn get_policy_package_names(&self) -> Result<Vec<PolicyPackageNameDefinition>> {
1391        let mut package_names = vec![];
1392        for m in self.modules.iter() {
1393            let package_name = Interpreter::get_path_string(&m.package.refr, None)?;
1394            package_names.push(PolicyPackageNameDefinition {
1395                source_file: m.package.span.source.file().to_string(),
1396                package_name,
1397            });
1398        }
1399
1400        Ok(package_names)
1401    }
1402
1403    /// Get the parameters defined in each policy.
1404    ///
1405    ///
1406    /// ```rust
1407    /// # use regorus::*;
1408    /// # use anyhow::{bail, Result};
1409    /// # fn main() -> Result<()> {
1410    /// # let mut engine = Engine::new();
1411    /// engine.add_policy("test.rego".to_string(), "package test default parameters.a = 5 parameters.b = 10\n x := 1".to_string())?;
1412    ///
1413    /// let parameters = engine.get_policy_parameters()?;
1414    ///
1415    /// assert_eq!("a", parameters[0].parameters[0].name);
1416    /// assert_eq!("b", parameters[0].modifiers[0].name);
1417    ///
1418    /// # Ok(())
1419    /// # }
1420    /// ```
1421    #[cfg(feature = "azure_policy")]
1422    #[cfg_attr(docsrs, doc(cfg(feature = "azure_policy")))]
1423    pub fn get_policy_parameters(&self) -> Result<Vec<PolicyParameters>> {
1424        let mut policy_parameter_definitions = vec![];
1425        for m in self.modules.iter() {
1426            let mut parameters = vec![];
1427            let mut modifiers = vec![];
1428
1429            for rule in &m.policy {
1430                match *rule.as_ref() {
1431                    // Extract parameter definitions from the policy rule
1432                    // e.g. default parameters.a = 5
1433                    Rule::Default { ref refr, .. } => {
1434                        let path = Parser::get_path_ref_components(refr)?;
1435                        let paths: Vec<&str> = path.iter().map(|s| s.text()).collect();
1436
1437                        if paths.len() == 2 && paths.first().is_some_and(|p| *p == "parameters") {
1438                            if let Some(name) = paths.get(1) {
1439                                // Todo: Fetch fields other than name from rego metadoc for the parameter
1440                                parameters.push(PolicyParameter {
1441                                    name: (*name).to_string(),
1442                                    modifiable: false,
1443                                    required: false,
1444                                });
1445                            }
1446                        }
1447                    }
1448                    // Extract modifiers to the parameters from the policy rule
1449                    // e.g. parameters.a = 5
1450                    Rule::Spec { ref head, .. } => {
1451                        match *head {
1452                            RuleHead::Compr { ref refr, .. } => {
1453                                let path = Parser::get_path_ref_components(refr)?;
1454                                let paths: Vec<&str> = path.iter().map(|s| s.text()).collect();
1455
1456                                if paths.len() == 2
1457                                    && paths.first().is_some_and(|p| *p == "parameters")
1458                                {
1459                                    if let Some(name) = paths.get(1) {
1460                                        // Todo: Fetch fields other than name from rego metadoc for the parameter
1461                                        modifiers.push(PolicyModifier {
1462                                            name: (*name).to_string(),
1463                                        });
1464                                    }
1465                                }
1466                            }
1467                            RuleHead::Func { .. } => {}
1468                            RuleHead::Set { .. } => {}
1469                        }
1470                    }
1471                }
1472            }
1473
1474            policy_parameter_definitions.push(PolicyParameters {
1475                source_file: m.package.span.source.file().to_string(),
1476                parameters,
1477                modifiers,
1478            });
1479        }
1480
1481        Ok(policy_parameter_definitions)
1482    }
1483
1484    /// Emit a warning if any modules contain target specifications but we're not using target-aware compilation.
1485    #[cfg(feature = "azure_policy")]
1486    fn warn_if_targets_present(&self) {
1487        let mut has_target = false;
1488        let mut target_files = Vec::new();
1489
1490        for module in self.modules.iter() {
1491            if module.target.is_some() {
1492                has_target = true;
1493                target_files.push(module.package.span.source.get_path());
1494            }
1495        }
1496
1497        if has_target {
1498            std::eprintln!("Warning: Target specifications found in policy modules but not using target-aware compilation.");
1499            std::eprintln!("         The following files contain __target__ declarations:");
1500            for file in target_files {
1501                std::eprintln!("         - {}", file);
1502            }
1503            std::eprintln!("         Consider using compile_for_target() instead of compile_with_entrypoint() for target-aware evaluation.");
1504        }
1505    }
1506
1507    fn make_parser<'a>(&self, source: &'a Source) -> Result<Parser<'a>> {
1508        let mut parser = Parser::new(source)?;
1509        if self.rego_v1 {
1510            parser.enable_rego_v1()?;
1511        }
1512        Ok(parser)
1513    }
1514
1515    /// Create a new Engine from a compiled policy.
1516    #[doc(hidden)]
1517    pub(crate) fn new_from_compiled_policy(
1518        compiled_policy: Rc<crate::compiled_policy::CompiledPolicyData>,
1519    ) -> Self {
1520        let modules = compiled_policy.modules.clone();
1521        let mut engine = Self {
1522            modules,
1523            interpreter: Interpreter::new_from_compiled_policy(compiled_policy),
1524            rego_v1: true, // Value doesn't matter since this is used only for policy parsing
1525            prepared: true,
1526            execution_timer_config: None,
1527        };
1528        engine.apply_effective_execution_timer_config();
1529        engine
1530    }
1531}