feat(compiler): support registered host-await builtins for natural function call syntax (#667)

* feat(compiler): support registered host-await builtins

Allow hosts to register function names at compile time so that calls to
those names emit HostAwait instructions directly, enabling natural syntax
like fetch(x) instead of __builtin_host_await(x, "fetch").

- Add host_await_builtins map and register_host_await_builtin() to Compiler
- Validate arg_count == 1 and reject reserved __builtin_host_await name
- Extend determine_call_target() resolution: explicit > registered > user > builtin
- Both explicit and registered paths emit identical HostAwait bytecode
- Add compile_from_policy_with_host_await() entry point in rules.rs
- Extended test harness with HostAwaitBuiltinSpec and args assertion
- 9 YAML test cases: suspend/resume, run-to-completion, multiple names,
  queue, shadowing, object packing, arg_count rejection, reserved name
  rejection, standard builtin override
- Documentation: instruction-set.md, architecture.md

* Update src/languages/rego/compiler/function_calls.rs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Mark Birger <birgerm@yandex.ru>

* fix(compiler): address PR #667 review feedback on host-await registration

- Compiler::register_host_await_builtin now rejects duplicate, empty,
  and whitespace-only names. Previously a duplicate registration would
  silently overwrite the existing entry, which could mask the host's
  own registration mistakes.
- YAML test cases added: empty registration list as no-op, duplicate
  name rejection, empty/whitespace name rejection, out-param (a, out)
  calling syntax with a single-arg registered builtin, and mixed
  __builtin_host_await + registered builtins in the same policy
  consuming from their respective identifier queues.
- Test harness: replace assert_eq! on HostAwait argument mismatch with
  anyhow::Error so mismatches propagate through the case reporter
  instead of panicking and skipping the harness's normal error path.
- YAML comment fix: "Registration panics" -> "Registration fails with
  an error" (registration returns Err, never panics).

Addresses anakrish + Copilot inline review comments on PR #667.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* compiler: split CallTarget::HostAwait into explicit and registered variants

Addresses PR #667 review item #8: at the emit site in
`compile_function_call`, the discrimination between explicit
`__builtin_host_await(arg, id)` and a registered host-awaitable
builtin was being recovered by string-comparing `original_fcn_path`
against `"__builtin_host_await"`. The information was already known
in `determine_call_target` and was being thrown away.

Replace the single `CallTarget::HostAwait` variant with two:

* `ExplicitHostAwait` (unit) — the two-argument call form. The
  identifier register comes from the user's second argument.
* `RegisteredHostAwait { identifier: String }` — the one-argument
  call form for registered builtins. The identifier is the registered
  name and is captured in the variant at recognition time, so the
  emit site never re-derives it from the function path.

This removes the magic-string comparison at the emit site (the source
of truth is now `determine_call_target`) and makes both match sites
in `compile_function_call` exhaustive over the two forms — adding a
third host-await form in the future would force a compile error at
every match site instead of silently falling through.

Arities are now hardcoded in the `expected_args` extraction
(`Some(2)` for explicit, `Some(1)` for registered) rather than
carried in the variant; registered builtins are constrained to
`arg_count == 1` at registration time, so there is no per-call
variability to carry.

Bytecode output is unchanged; the full RVM test suite (97 cases) and
the registered_host_await suite (15 cases) pass without modification.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* docs(compiler): clarify registered host-await intercepts unqualified calls only

PR #667 review (Medium): the docs implied registered host-await names
shadow user functions and builtins unconditionally, but
determine_call_target matches only the bare original_fcn_path. A
package-qualified call such as data.demo.resolve(x) is therefore not
intercepted -- it resolves through the normal path like any other call.

Rather than expand registration to qualified paths (which would let a
registered name leak into every package exposing a same-named rule),
document the unqualified-only behavior and pin it with tests.

- register_host_await_builtin: doc now states only the unqualified call
  form is intercepted; qualified calls resolve normally.
- determine_call_target: inline comment explaining the deliberate
  original_fcn_path-only match.
- docs/rvm/instruction-set.md: describe qualified-call resolution,
  including that builtins have no qualified form.
- tests: cross-package and same-package qualified calls resolve to the
  rule; bare-name shadowing of a standard builtin; Unknown-function
  outcome when no rule exists at the qualified path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(tests): compare host-await argument without re-running process_value

PR #667 review (Low): the suspendable test harness compared the
host-await argument via process_value(argument), but argument is already
a runtime Value. process_value is a YAML-fixture decoder -- it rewrites
"#undefined" to Undefined, {set!: [...]} to a set, and errors on a
runtime Value::Set. Re-running it on the runtime argument could coerce a
legitimate payload into a fixture sentinel (passing for the wrong
reason) or error outright on sets.

Compare the runtime argument directly against the expected value, which
is already decoded once at YAML load time.

Add a regression case (registered_builtin_suspendable_set_argument) that
passes a set payload: it fails under the old double-processing
("unexpected set in value read from json/yaml") and passes with the fix.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(tests): reject `args:` payload expectations in run-to-completion mode

PR #667 review (Low): a run-to-completion host-await response could carry
an `args:` payload expectation, but RTC execution pre-loads responses and
never surfaces the call argument to the harness, so the expectation was
parsed and silently dropped. A case with `args: "WRONG"` passed as long as
the result matched -- asserting a payload that was never checked.

Reject `args:` for run-to-completion fixtures at load time, directing the
author to suspendable mode where arguments are validated. Also only build
the run-to-completion response vector when the case actually runs in RTC
mode, so a suspendable case using the shared host_await_responses field
with `args:` is not wrongly rejected.

Route the fixture-load error through the same want_error handling used for
compilation errors, and add registered_builtin_run_to_completion_rejects_args
which now fails loudly instead of passing silently.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(compiler): reject host-await builtin names with surrounding whitespace

PR #667 review (Low): register_host_await_builtin rejected all-whitespace
names via name.trim().is_empty(), but accepted padded names like " lookup"
or "lookup ". Those were inserted into host_await_builtins, but Rego
function-call paths produce the trimmed identifier, so a padded
registration could never match -- a silent dead registration.

Reject any name that is not already trimmed (name != name.trim()) in
addition to empty names, and update the error message accordingly. Add
test cases for leading and trailing whitespace.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Signed-off-by: Mark Birger <birgerm@yandex.ru>
Co-authored-by: Mark Birger <markbirger@microsoft.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: 4 people

docs/rvm/architecture.md

@@ -254,7 +254,13 @@ include formatted state snapshots where possible.
 7. **Host await**: In run-to-completion mode, `HostAwait` consumes a response
    from `host_await_responses`. Suspendable mode yields control with a
    `SuspendReason::HostAwait { dest, argument, identifier }` that the host must
-   service.
+   service. The compiler supports two ways to emit `HostAwait`:
+   - **Explicit**: `__builtin_host_await(payload, identifier)` — raw 2-argument
+     form.
+   - **Registered**: `compile_from_policy_with_host_await` accepts a list of
+     `(name, arg_count)` pairs. Calls to registered names are compiled as
+     `HostAwait` with the function name as the identifier literal. Registered
+     names take precedence over user-defined functions and standard builtins.
 8. **Completion**: `Return` wraps the selected register value into
    `InstructionOutcome::Return`, unwinding frames until the entry frame is
    cleared. `RuleReturn` is a specialised variant used by rule execution

docs/rvm/instruction-set.md

@@ -177,6 +177,75 @@ Parameter tables:
 - Suspendable: emits `InstructionOutcome::Suspend` with `SuspendReason::HostAwait`.
   The host must resume with a value that will be written into `dest`.
 
+### Registered host-await builtins
+
+The compiler can be configured with a list of function names that map directly
+to `HostAwait` instructions. This allows policy authors to write natural
+function calls (e.g. `lookup(input.account_id)`) instead of the raw
+`__builtin_host_await(payload, identifier)` builtin.
+
+Registration is done at compile time via `Compiler::compile_from_policy_with_host_await`:
+
+```rust
+let builtins = [("lookup", 1), ("persist", 1)];
+let program = Compiler::compile_from_policy_with_host_await(
+    &compiled_policy, &entry_points, &builtins,
+)?;
+```
+
+Each registered name is a `(name, arg_count)` pair. When the compiler
+encounters a call to a registered name, it emits a `HostAwait` instruction
+with:
+- `arg` = the first argument register
+- `id` = a register loaded with a string literal containing the function name
+
+Both the explicit `__builtin_host_await(arg, id)` call and a registered
+builtin call produce the **same `HostAwait` bytecode instruction**. The only
+difference is how the `id` register is populated: explicit calls take it from
+the second user-supplied argument, while registered calls auto-generate a
+`Load` instruction for the function name string. The VM cannot distinguish
+between the two at runtime.
+
+**Resolution order** in `determine_call_target()`:
+1. `__builtin_host_await` (magic 2-argument form)
+2. Registered host-await builtins (matched by **bare** function name only)
+3. User-defined functions (matched by package-qualified path)
+4. Standard builtins (matched by bare function name)
+
+Registered names shadow both user-defined functions and standard builtins.
+This means `time.parse_duration_ns` can be overridden to route through the
+host instead of the built-in Rust implementation.
+
+**Only unqualified calls are intercepted.** Registration matches a call by
+the name *as written in the policy*. A bare call — `lookup(x)` — is
+intercepted and compiled to a `HostAwait`. A package-qualified call —
+`data.pkg.lookup(x)` — is **not** intercepted; it is resolved normally, as
+if the name were never registered.
+
+```rego
+# "lookup" is registered as a host-await builtin.
+
+package other
+import rego.v1
+lookup(k) := k            # an ordinary rule that happens to share the name
+
+package demo
+import rego.v1
+a := lookup(input.k)             # intercepted     -> HostAwait
+b := data.other.lookup(input.k)  # NOT intercepted -> calls other.lookup
+```
+
+The qualified form is resolved exactly as it would be without registration:
+if a rule exists at that path it is called, otherwise compilation fails with
+`Unknown function`. (A standard builtin like `count` has no qualified form at
+all, so `data.pkg.count(x)` is always an `Unknown function` error, registered
+or not.)
+
+**Argument handling**: The `HostAwait` instruction carries a single `arg`
+register. Registered builtins must use `arg_count: 1`; the compiler rejects
+`arg_count > 1` at registration time. To pass multiple values, use object
+packing: `lookup({"user": x, "resource": y})`.
+
 ---
 
 ## Halt instruction

src/languages/rego/compiler/function_calls.rs

@@ -17,8 +17,20 @@ use crate::lexer::Span;
 use crate::rvm::instructions::{BuiltinCallParams, FunctionCallParams};
 use crate::rvm::Instruction;
 use crate::utils::get_path_string;
-use alloc::{format, string::ToString, vec::Vec};
+use crate::value::Value;
+use alloc::{
+    format,
+    string::{String, ToString},
+    vec::Vec,
+};
 
+/// Resolved destination of a Rego function-call expression. Produced by
+/// [`Compiler::determine_call_target`] and consumed by
+/// [`Compiler::compile_function_call`] to choose which instruction to emit.
+/// Carrying the discrimination in the type (rather than re-matching on a
+/// magic name at the emit site) keeps the host-await handling honest under
+/// future refactors — the compiler will refuse to build if a new variant is
+/// added without updating every match site.
 enum CallTarget {
     User {
         rule_index: u16,
@@ -28,9 +40,14 @@ enum CallTarget {
         builtin_index: u16,
         expected_args: Option<usize>,
     },
-    HostAwait {
-        expected_args: Option<usize>,
-    },
+    /// Explicit `__builtin_host_await(arg, id)` call form (2 user args).
+    /// The identifier is supplied by the policy author at runtime via the
+    /// second argument register.
+    ExplicitHostAwait,
+    /// A registered host-awaitable builtin invoked by its registered name
+    /// (1 user arg). The identifier is the registered name itself and is
+    /// baked into the bytecode as a string literal at compile time.
+    RegisteredHostAwait { identifier: String },
 }
 
 impl<'a> Compiler<'a> {
@@ -59,7 +76,11 @@ impl<'a> Compiler<'a> {
         let expected_args = match &call_target {
             CallTarget::User { expected_args, .. } => *expected_args,
             CallTarget::Builtin { expected_args, .. } => *expected_args,
-            CallTarget::HostAwait { expected_args } => *expected_args,
+            // Both host-await variants have a known fixed arity; carrying it
+            // in the variant lets the rest of the compiler depend on the type
+            // rather than re-matching on the magic name `__builtin_host_await`.
+            CallTarget::ExplicitHostAwait => Some(2),
+            CallTarget::RegisteredHostAwait { .. } => Some(1),
         };
 
         if let Some(expected) = expected_args {
@@ -126,7 +147,8 @@ impl<'a> Compiler<'a> {
                 });
                 self.emit_instruction(Instruction::BuiltinCall { params_index }, &span);
             }
-            CallTarget::HostAwait { .. } => {
+            CallTarget::ExplicitHostAwait => {
+                // Explicit __builtin_host_await(arg, id) — 2 arguments
                 if arg_regs.len() != 2 {
                     return Err(CompilerError::General {
                         message: format!(
@@ -136,7 +158,6 @@ impl<'a> Compiler<'a> {
                     }
                     .at(&span));
                 }
-
                 self.emit_instruction(
                     Instruction::HostAwait {
                         dest,
@@ -146,6 +167,37 @@ impl<'a> Compiler<'a> {
                     &span,
                 );
             }
+            CallTarget::RegisteredHostAwait { identifier } => {
+                // Registered host-awaitable builtin — the identifier is the
+                // registered name and is baked into the bytecode as a literal.
+                if arg_regs.len() != 1 {
+                    return Err(CompilerError::General {
+                        message: format!(
+                            "host-awaitable builtin '{}' expects exactly 1 argument, got {}",
+                            identifier,
+                            arg_regs.len()
+                        ),
+                    }
+                    .at(&span));
+                }
+                let id_reg = self.alloc_register();
+                let literal_idx = self.add_literal(Value::String(identifier.into()));
+                self.emit_instruction(
+                    Instruction::Load {
+                        dest: id_reg,
+                        literal_idx,
+                    },
+                    &span,
+                );
+                self.emit_instruction(
+                    Instruction::HostAwait {
+                        dest,
+                        arg: arg_regs[0],
+                        id: id_reg,
+                    },
+                    &span,
+                );
+            }
         }
 
         if let Some((plan, plan_span)) = &out_param_plan {
@@ -187,8 +239,26 @@ impl<'a> Compiler<'a> {
         span: &Span,
     ) -> Result<CallTarget> {
         if original_fcn_path == "__builtin_host_await" {
-            return Ok(CallTarget::HostAwait {
-                expected_args: Some(2),
+            return Ok(CallTarget::ExplicitHostAwait);
+        }
+
+        // Check registered host-awaitable builtins. Registered builtins are
+        // restricted to arg_count == 1 at registration time (see
+        // `Compiler::register_host_await_builtin`), so the variant doesn't
+        // need to carry an arity — it's fixed at 1.
+        //
+        // We deliberately match against `original_fcn_path` only, not
+        // `full_fcn_path`. Registration intercepts the *unqualified* call
+        // form (e.g. `lookup(x)` inside the policy's own package). A
+        // package-qualified call like `data.other.lookup(x)` is left to
+        // resolve through the normal user-defined / builtin path, so a
+        // registered name does not leak into unrelated packages that
+        // happen to expose a rule with the same identifier. This is
+        // documented on `register_host_await_builtin`; the
+        // `registered_host_await.yaml` suite pins the behavior.
+        if self.host_await_builtins.contains_key(original_fcn_path) {
+            return Ok(CallTarget::RegisteredHostAwait {
+                identifier: original_fcn_path.to_string(),
             });
         }
 

src/languages/rego/compiler/mod.rs

@@ -26,7 +26,9 @@ 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;
@@ -139,6 +141,10 @@ pub struct Compiler<'a> {
     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> {
@@ -173,9 +179,75 @@ impl<'a> Compiler<'a> {
             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,

src/languages/rego/compiler/rules.rs

@@ -234,8 +234,20 @@ impl<'a> Compiler<'a> {
     pub fn compile_from_policy(
         policy: &CompiledPolicy,
         entry_points: &[&str],
+    ) -> Result<Arc<Program>> {
+        Self::compile_from_policy_with_host_await(policy, entry_points, &[])
+    }
+
+    /// Compile from a CompiledPolicy to RVM Program with registered host-awaitable builtins.
+    pub fn compile_from_policy_with_host_await(
+        policy: &CompiledPolicy,
+        entry_points: &[&str],
+        host_await_builtins: &[(&str, usize)],
     ) -> Result<Arc<Program>> {
         let mut compiler = Compiler::with_policy(policy);
+        for &(name, arg_count) in host_await_builtins {
+            compiler.register_host_await_builtin(name, arg_count)?;
+        }
         compiler.current_rule_path = "".to_string();
         let rules = policy.get_rules();