regorus/src/languages/rego/compiler/mod.rs

#![allow(
    missing_debug_implementations,
    clippy::missing_const_for_fn,
    clippy::option_if_let_else,
    clippy::if_then_some_else_none,
    clippy::unused_self
)] // compiler internals do not require Debug

mod comprehensions;
mod core;
mod destructuring;
mod error;
mod expressions;
mod function_calls;
mod loops;
mod program;
mod queries;
mod references;
mod rules;

pub use error::{CompilerError, Result, SpannedCompilerError};

use crate::ast::ExprRef;
use crate::lexer::Span;
use crate::rvm::program::{Program, RuleType, SpanInfo};
use crate::CompiledPolicy;
use crate::Value;
use alloc::collections::{BTreeMap, BTreeSet};
use alloc::format;
use alloc::string::String;
use alloc::string::ToString as _;
use alloc::vec;
use alloc::vec::Vec;
use indexmap::IndexMap;

pub type Register = u8;

#[derive(Debug, Clone, Default)]
struct Scope {
    bound_vars: BTreeMap<String, Register>,
    unbound_vars: BTreeSet<String>,
}

#[derive(Debug, Clone)]
pub enum ComprehensionType {
    Array,
    Object,
    Set,
}

#[derive(Debug, Clone)]
pub enum ContextType {
    Comprehension(ComprehensionType),
    Rule(RuleType),
    Every,
}

/// Compilation context for handling different types of rule bodies and comprehensions
#[derive(Debug, Clone)]
pub struct CompilationContext {
    pub(super) context_type: ContextType,
    pub(super) dest_register: Register,
    pub(super) key_expr: Option<ExprRef>,
    pub(super) value_expr: Option<ExprRef>,
    pub(super) span: Span,
    pub(super) key_value_loops_hoisted: bool,
}

/// Entry in the rule compilation worklist that tracks both rule path and full call stack for recursion detection
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct WorklistEntry {
    /// Rule path to be compiled (e.g., "data.package.rule")
    pub rule_path: String,
    /// Call stack of rule indices leading to this rule (empty for entry point)
    pub call_stack: Vec<u16>,
}

impl WorklistEntry {
    pub fn new(rule_path: String, call_stack: Vec<u16>) -> Self {
        Self {
            rule_path,
            call_stack,
        }
    }

    pub fn entry_point(rule_path: String) -> Self {
        Self {
            rule_path,
            call_stack: Vec::new(),
        }
    }

    /// Create a new entry by extending the call stack with the caller's rule index
    pub fn with_caller(
        rule_path: String,
        current_call_stack: &[u16],
        caller_rule_index: u16,
    ) -> Self {
        let mut new_call_stack = current_call_stack.to_vec();
        new_call_stack.push(caller_rule_index);
        Self {
            rule_path,
            call_stack: new_call_stack,
        }
    }

    /// Check if this entry would create a recursive call
    pub fn would_create_recursion(&self, target_rule_index: u16) -> bool {
        self.call_stack.contains(&target_rule_index)
    }
}

pub struct Compiler<'a> {
    program: Program,
    spans: Vec<SpanInfo>,
    register_counter: Register,
    scopes: Vec<Scope>,
    policy: &'a CompiledPolicy,
    current_package: String,
    current_module_index: u32,
    rule_index_map: BTreeMap<String, u16>,
    rule_worklist: Vec<WorklistEntry>,
    rule_definitions: Vec<Vec<Vec<u32>>>,
    rule_definition_function_params: Vec<Vec<Option<Vec<String>>>>,
    rule_definition_destructuring_patterns: Vec<Vec<Option<u32>>>,
    /// Per-rule, per-definition: the static value produced by this definition,
    /// or `None` if the value is dynamic or differs across else-branches.
    /// Used to compute `RuleInfo::early_exit_on_first_success`.
    rule_definition_static_values: Vec<Vec<Option<Value>>>,
    rule_types: Vec<RuleType>,
    rule_function_param_count: Vec<Option<usize>>,
    rule_result_registers: Vec<u8>,
    rule_num_registers: Vec<u8>,
    context_stack: Vec<CompilationContext>,
    loop_expr_register_map: BTreeMap<ExprRef, Register>,
    source_to_index: BTreeMap<String, usize>,
    builtin_index_map: BTreeMap<String, u16>,
    current_input_register: Option<Register>,
    current_data_register: Option<Register>,
    current_rule_path: String,
    current_call_stack: Vec<u16>,
    entry_points: IndexMap<String, usize>,
    soft_assert_mode: bool,
    /// Registered host-awaitable builtins: name → expected arg count.
    /// When the compiler encounters a call to one of these names, it emits a
    /// `HostAwait` instruction instead of a regular function or builtin call.
    host_await_builtins: BTreeMap<String, usize>,
}

impl<'a> Compiler<'a> {
    pub fn with_policy(policy: &'a CompiledPolicy) -> Self {
        let mut program = Program::new();
        program.rego_v0 = policy.is_rego_v0();
        Self {
            program,
            spans: Vec::new(),
            register_counter: 1,
            scopes: vec![Scope::default()],
            policy,
            current_package: String::new(),
            current_module_index: 0,
            rule_index_map: BTreeMap::new(),
            rule_worklist: Vec::new(),
            rule_definitions: Vec::new(),
            rule_definition_function_params: Vec::new(),
            rule_definition_destructuring_patterns: Vec::new(),
            rule_definition_static_values: Vec::new(),
            rule_types: Vec::new(),
            rule_function_param_count: Vec::new(),
            rule_result_registers: Vec::new(),
            rule_num_registers: Vec::new(),
            context_stack: vec![],
            loop_expr_register_map: BTreeMap::new(),
            source_to_index: BTreeMap::new(),
            builtin_index_map: BTreeMap::new(),
            current_input_register: None,
            current_data_register: None,
            current_rule_path: String::new(),
            current_call_stack: Vec::new(),
            entry_points: IndexMap::new(),
            soft_assert_mode: false,
            host_await_builtins: BTreeMap::new(),
        }
    }

    /// Register a function name as a host-awaitable builtin.
    ///
    /// When the compiler encounters an **unqualified** call to `name(arg)`
    /// (i.e. `name(arg)` from inside the policy's own package, not
    /// `data.pkg.name(arg)` or any other package-qualified form), it will
    /// emit a `HostAwait` instruction with the argument and `name` as the
    /// identifier, instead of treating it as a user-defined or standard
    /// builtin function.
    ///
    /// Package-qualified calls (e.g. `data.other.name(arg)`) are **not**
    /// intercepted by registration. Those resolve through the normal
    /// user-defined / builtin lookup against their fully-qualified path
    /// (`data.other.name`).
    ///
    /// `arg_count` must be exactly 1. The `HostAwait` instruction carries a
    /// single argument register; use object packing to pass multiple values
    /// (e.g. `name({"key1": v1, "key2": v2})`).
    ///
    /// Returns `Err` when:
    /// - `name` is the reserved identifier `__builtin_host_await`,
    /// - `name` is empty, only whitespace, or has leading/trailing
    ///   whitespace (whitespace-padded names would never match the
    ///   trimmed identifier produced by the Rego parser, creating dead
    ///   registrations),
    /// - `name` is already registered (duplicate registration is rejected
    ///   rather than silently overwritten),
    /// - `arg_count` is not exactly 1.
    pub fn register_host_await_builtin(&mut self, name: &str, arg_count: usize) -> Result<()> {
        if name == "__builtin_host_await" {
            return Err(CompilerError::General {
                message: "__builtin_host_await is a reserved name and cannot be registered as a host-await builtin"
                    .to_string(),
            }
            .into());
        }
        if name.is_empty() || name != name.trim() {
            return Err(CompilerError::General {
                message: format!(
                    "host-await builtin name {name:?} must not be empty or contain leading/trailing whitespace"
                ),
            }
            .into());
        }
        if self.host_await_builtins.contains_key(name) {
            return Err(CompilerError::General {
                message: format!(
                    "host-await builtin '{name}' is already registered; \
                     duplicate registration is not allowed"
                ),
            }
            .into());
        }
        if arg_count != 1 {
            return Err(CompilerError::General {
                message: format!(
                    "registered host-await builtin '{name}' must have arg_count == 1, got {arg_count}. \
                     Use object packing to pass multiple values."
                ),
            }
            .into());
        }
        self.host_await_builtins.insert(name.to_string(), arg_count);
        Ok(())
    }

    pub(super) fn with_soft_assert_mode<F, R>(&mut self, enabled: bool, f: F) -> R
    where
        F: FnOnce(&mut Self) -> R,
    {
        let previous = self.soft_assert_mode;
        self.soft_assert_mode = enabled;
        let result = f(self);
        self.soft_assert_mode = previous;
        result
    }
}