fix(interpreter,rvm): resolve function calls through import aliases (#769)

* fix(interpreter,rvm): resolve function calls through import aliases

* fix(rvm): assert every-quantifier results so failing cases don't pass

src/interpreter.rs

@@ -1258,7 +1258,34 @@ impl Interpreter {
             // Apply with modifiers.
             for wm in &stmt.with_mods {
                 let path = Parser::get_path_ref_components(&wm.refr)?;
-                let path: Vec<&str> = path.iter().map(|s| s.text()).collect();
+                let mut path: Vec<String> = path.iter().map(|s| s.text().to_string()).collect();
+
+                // Matching OPA, a leading import alias is rewritten before
+                // any lookups: functions register as overrides below,
+                // anything else becomes a data override. Only the alias
+                // component is replaced so bracketed keys containing dots
+                // survive the rewrite.
+                let rewritten: Option<Vec<String>> = match path.split_first() {
+                    Some((head, rest)) if head.as_str() != "data" => {
+                        self.lookup_import(head).and_then(|import_expr| {
+                            // Use the import target's parsed components, not
+                            // its dot-joined string, so bracketed keys
+                            // containing dots survive in the import path too.
+                            let comps = Parser::get_path_ref_components(import_expr).ok()?;
+                            Some(
+                                comps
+                                    .iter()
+                                    .map(|s| s.text().to_string())
+                                    .chain(rest.iter().cloned())
+                                    .collect(),
+                            )
+                        })
+                    }
+                    _ => None,
+                };
+                if let Some(new_path) = rewritten {
+                    path = new_path;
+                }
                 let mut target = path.join(".");
 
                 let mut target_is_function = self.lookup_function_by_name(&target).is_some()
@@ -1297,11 +1324,17 @@ impl Interpreter {
                             if self.lookup_function_by_name(&function_path).is_none() {
                                 // Lookup without current module path prefixed.
                                 function_path = get_path_string(&wm.r#as, None)?;
-                                if self.lookup_function_by_name(&function_path).is_none()
-                                    && !Self::is_builtin(wm.r#as.span(), &function_path)
-                                {
-                                    // bail!(wm.r#as.span().error("could not evaluate expression"));
-                                    skip_exec = true;
+                                if self.lookup_function_by_name(&function_path).is_none() {
+                                    // Resolve an aliased replacement before builtins.
+                                    let resolved = self
+                                        .resolve_fcn_path_through_imports(&function_path)
+                                        .filter(|r| self.compiled_policy.functions.contains_key(r));
+                                    if let Some(resolved) = resolved {
+                                        function_path = resolved;
+                                    } else if !Self::is_builtin(wm.r#as.span(), &function_path) {
+                                        // bail!(wm.r#as.span().error("could not evaluate expression"));
+                                        skip_exec = true;
+                                    }
                                 }
                             }
                             self.with_functions
@@ -2371,6 +2404,72 @@ impl Interpreter {
         }
     }
 
+    /// Look up the import of the current module with the given alias, e.g.
+    /// the `data.a.b` import expression for `b` after `import data.a.b`.
+    fn lookup_import(&self, alias: &str) -> Option<&Ref<Expr>> {
+        if self.compiled_policy.imports.is_empty() {
+            return None;
+        }
+        let import_key = format!("{}.{}", self.current_module_path, alias);
+        self.compiled_policy.imports.get(&import_key)
+    }
+
+    /// Look up the dot-joined target path of an import of the current module
+    /// with the given alias, e.g. `data.a.b` for `b` after `import data.a.b`.
+    fn lookup_import_alias(&self, alias: &str) -> Option<String> {
+        get_path_string(self.lookup_import(alias)?, None).ok()
+    }
+
+    /// Rewrite a path whose leading component is an import alias of the
+    /// current module to the import's target, e.g. `b.f` to `data.a.b.f`
+    /// after `import data.a.b`.
+    fn rewrite_path_through_imports(&self, path: &str) -> Option<String> {
+        if path.starts_with("data.") {
+            return None;
+        }
+
+        let (alias, rest) = match path.split_once('.') {
+            Some((alias, rest)) => (alias, Some(rest)),
+            None => (path, None),
+        };
+        let target = self.lookup_import_alias(alias)?;
+        Some(match rest {
+            Some(rest) => format!("{target}.{rest}"),
+            None => target,
+        })
+    }
+
+    /// Rewrite an import-aliased call path to its target, e.g. `b.f(1)` to
+    /// `data.a.b.f` after `import data.a.b`. Resolves only when the target is
+    /// a known function or default function, so an alias whose target defines
+    /// the called function shadows a like-named builtin namespace, while other
+    /// spellings keep their prior meaning (e.g. a builtin call). OPA instead
+    /// rewrites aliases unconditionally and rejects calls to a missing target
+    /// at compile time.
+    fn resolve_fcn_path_through_imports(&self, path: &str) -> Option<String> {
+        let candidate = self.rewrite_path_through_imports(path)?;
+        (self.compiled_policy.functions.contains_key(&candidate)
+            || self.is_default_function(&candidate))
+        .then_some(candidate)
+    }
+
+    /// True if `path` is the exact path of a `default` function rule.
+    /// `default_rules` also indexes every prefix of a rule path, so it cannot
+    /// be consulted alone: `rule_paths` holds only exact rule paths, and the
+    /// non-empty argument list distinguishes functions from value rules.
+    fn is_default_function(&self, path: &str) -> bool {
+        self.compiled_policy.rule_paths.contains(path)
+            && self
+                .compiled_policy
+                .default_rules
+                .get(path)
+                .is_some_and(|rules| {
+                    rules.iter().any(|(rule, _)| {
+                        matches!(rule.as_ref(), Rule::Default { args, .. } if !args.is_empty())
+                    })
+                })
+    }
+
     fn eval_builtin_call(
         &mut self,
         span: &Span,
@@ -2542,6 +2641,13 @@ impl Interpreter {
             param_values.push(self.eval_expr(p)?);
         }
 
+        // Resolve a leading import alias before the `with` override and builtin
+        // lookups, so an override keyed by the full path reaches aliased calls
+        // and the alias shadows a like-named builtin namespace (matching OPA).
+        let fcn_path = self
+            .resolve_fcn_path_through_imports(&fcn_path)
+            .unwrap_or(fcn_path);
+
         let orig_fcn_path = fcn_path.clone();
 
         let mut with_functions_saved = None;
@@ -2718,7 +2824,12 @@ impl Interpreter {
             let value = match self.eval_rule_bodies(ctx, span, bodies) {
                 Ok(v) => v,
                 Err(e) => {
-                    // If the rule produces an error, save the error.
+                    // If the rule produces an error, save the error. Restore
+                    // the caller's module even so: leaving the callee's module
+                    // in place would make the rest of the caller's body
+                    // resolve paths through the wrong module's imports when
+                    // the error is swallowed below in non-strict mode.
+                    self.set_current_module(prev_module)?;
                     errors.push(e);
                     self.scopes = scopes;
                     continue;

src/languages/rego/compiler/function_calls.rs

@@ -63,6 +63,9 @@ impl<'a> Compiler<'a> {
         let original_fcn_path = fcn_path.clone();
         let full_fcn_path = if self.policy.inner.rules.contains_key(&fcn_path) {
             fcn_path
+        } else if let Some(resolved) = self.resolve_fcn_path_through_imports(&original_fcn_path) {
+            // Resolve a leading import alias before module-prefixing and builtins.
+            resolved
         } else {
             get_path_string(fcn, Some(&self.current_package))
                 .map_err(|_| CompilerError::InvalidFunctionExpressionWithPackage.at(&span))?
@@ -220,6 +223,37 @@ impl<'a> Compiler<'a> {
         Ok(dest)
     }
 
+    /// Rewrite an import-aliased call path to its target, e.g. `b.f(1)` to
+    /// `data.a.b.f` after `import data.a.b`. Resolves only when the target is
+    /// a known function (the `rules` map cannot be used here: it also indexes
+    /// value rules and every rule-path prefix, which must not become callable
+    /// through an alias), so an alias whose target defines the called
+    /// function shadows a like-named builtin namespace, while other spellings
+    /// keep their prior meaning (e.g. a builtin call). OPA instead rewrites
+    /// aliases unconditionally and rejects calls to a missing target at
+    /// compile time.
+    fn resolve_fcn_path_through_imports(&self, path: &str) -> Option<String> {
+        if self.policy.inner.imports.is_empty() || path.starts_with("data.") {
+            return None;
+        }
+        let (alias, rest) = match path.split_once('.') {
+            Some((alias, rest)) => (alias, Some(rest)),
+            None => (path, None),
+        };
+        let import_key = format!("{}.{}", self.current_package, alias);
+        let import_expr = self.policy.inner.imports.get(&import_key)?;
+        let target = get_path_string(import_expr, None).ok()?;
+        let candidate = match rest {
+            Some(rest) => format!("{target}.{rest}"),
+            None => target,
+        };
+        self.policy
+            .inner
+            .functions
+            .contains_key(&candidate)
+            .then_some(candidate)
+    }
+
     fn lookup_builtin_arity(&self, name: &str) -> Option<usize> {
         if name == "print" {
             Some(2)

tests/interpreter/cases/import/tests.yaml

@@ -188,3 +188,435 @@ cases:
         import input
     query: data.test
     want_result: {}
+  - note: function call through import
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: function call through import alias
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib as l
+
+        result := l.double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: function imported directly
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib.double
+
+        result := double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: function call through import, nested package
+    modules:
+      - |
+        package a.b.c
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.a
+
+        result := a.b.c.double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: default function value via import
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        default pick(_) := "fallback"
+
+        pick(x) := x if x == "a"
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        a := lib.pick("a")
+        b := lib.pick("z")
+    query: data.rules
+    want_result:
+      a: a
+      b: fallback
+  - note: with override applies to function call through import
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package lib2
+        import rego.v1
+
+        fake(x) := 1000 + x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.double(5)
+
+        main := x if {
+          x := result with data.lib.double as data.lib2.fake
+        }
+    query: data.rules.main
+    want_result: 1005
+  - note: with override target written through import alias
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package lib2
+        import rego.v1
+
+        fake(x) := 1000 + x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := data.lib.double(5)
+
+        main := x if {
+          x := result with lib.double as data.lib2.fake
+        }
+    query: data.rules.main
+    want_result: 1005
+  - note: with override replacement written through import alias
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package lib2
+        import rego.v1
+
+        fake(x) := 1000 + x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+        import data.lib2
+
+        result := data.lib.double(5)
+
+        main := x if {
+          x := result with data.lib.double as lib2.fake
+        }
+    query: data.rules.main
+    want_result: 1005
+  - note: import alias shadows builtin namespace
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        unmarshal(_) := {"from": "lib"}
+      - |
+        package rules
+        import rego.v1
+        import data.lib as json
+
+        result := json.unmarshal(`[1,2]`)
+    query: data.rules.result
+    want_result:
+      from: lib
+  - note: with override target written through builtin-shadowing alias
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        unmarshal(_) := {"from": "lib"}
+      - |
+        package lib2
+        import rego.v1
+
+        unmarshal(_) := {"from": "lib2"}
+      - |
+        package rules
+        import rego.v1
+        import data.lib as json
+
+        result := json.unmarshal(`[1]`)
+
+        main := x if {
+          x := result with json.unmarshal as data.lib2.unmarshal
+        }
+    query: data.rules.main
+    want_result:
+      from: lib2
+  - note: with override replacement written through builtin-shadowing alias
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package lib2
+        import rego.v1
+
+        unmarshal(_) := {"from": "lib2"}
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+        import data.lib2 as json
+
+        result := lib.double(5)
+
+        main := x if {
+          x := result with data.lib.double as json.unmarshal
+        }
+    query: data.rules.main
+    want_result:
+      from: lib2
+  - note: with override target aliasing a non-function rule overrides the value
+    modules:
+      - |
+        package flib
+        import rego.v1
+
+        default flag := false
+
+        flag := true if input.x == 1
+      - |
+        package rules
+        import rego.v1
+        import data.flib as fl
+
+        result := data.flib.flag
+
+        main := x if {
+          x := result with fl.flag as 5
+        }
+    query: data.rules.main
+    want_result: 5
+  - note: import shadows module-local rule for function calls
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        lib.double(x) := 100 * x
+
+        main := lib.double(3)
+    query: data.rules.main
+    want_result: 6
+  - note: import shadows module-local value rule for function calls
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        lib.double := 7
+
+        main := lib.double(3)
+    query: data.rules.main
+    want_result: 6
+  - note: fully qualified function call with import present
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := data.lib.double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: unknown function through import is an error
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.missing(1)
+    query: data.rules.result
+    error: could not find function
+  - note: function call in module without imports
+    modules:
+      - |
+        package rules
+
+        double(x) := 2 * x
+
+        result := double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: call to non-function value rule through import is an error
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        x := 5
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.x(1)
+    query: data.rules.result
+    error: could not find function
+  - note: call to rule path prefix through import is an error
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        a.c := 9
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.a(7)
+    query: data.rules.result
+    error: could not find function
+  - note: call to prefix of default function through import is an error
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        default a.b(_) := 1
+
+        a.b(x) := x * 2 if x > 10
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.a(7)
+    query: data.rules.result
+    error: could not find function
+  - note: default-only function call through import
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        default dfl(_) := 42
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.dfl(5)
+    query: data.rules.result
+    want_result: 42
+  - note: overload error does not leak callee module imports
+    strict: false
+    modules:
+      - |
+        package libx
+        import rego.v1
+
+        h(_) := "WRONG"
+      - |
+        package liby
+        import rego.v1
+
+        h(_) := "RIGHT"
+      - |
+        package lib
+        import rego.v1
+        import data.libx as ali
+
+        m(_) := 1
+
+        m(_) := 2
+
+        f(x) := 100 if x == 1
+
+        f(x) := y if {
+          x == 1
+          y := m(1)
+        }
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+        import data.liby as ali
+
+        main := [v, w] if {
+          v := lib.f(1)
+          w := ali.h(1)
+        }
+    query: data.rules.main
+    want_result: [100, "RIGHT"]
+  - note: with target through dotted bracket key import
+    data:
+      "a.b":
+        x: 1
+    modules:
+      - |
+        package rules
+        import rego.v1
+        import data["a.b"] as ab
+
+        main := x if {
+          x := data["a.b"].x with ab.x as 2
+        }
+    query: data.rules.main
+    want_result: 2

tests/rvm/rego/cases/imports.yaml

@@ -95,3 +95,204 @@ cases:
         }
     query: data.rules.present
     want_result: true
+  - note: import_function_call
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import data.lib
+        import rego.v1
+
+        result := lib.double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: import_function_call_alias
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import data.lib as mylib
+        import rego.v1
+
+        result := mylib.double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: import_function_directly
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import data.lib.double
+        import rego.v1
+
+        result := double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: import_function_call_nested_package
+    modules:
+      - |
+        package a.b.c
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import data.a
+        import rego.v1
+
+        result := a.b.c.double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: import_alias_shadows_builtin_namespace
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        unmarshal(_) := {"from": "lib"}
+      - |
+        package rules
+        import rego.v1
+        import data.lib as json
+
+        result := json.unmarshal(`[1,2]`)
+    query: data.rules.result
+    want_result:
+      from: lib
+  - note: default_function_value_via_import
+    skip: true  # TODO: cross-package calls to functions with default values fail to compile ("not a valid rule path"), even with fully-qualified paths
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        default pick(_) := "fallback"
+
+        pick(x) := x if x == "a"
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        a := lib.pick("a")
+        b := lib.pick("z")
+    query: data.rules
+    want_result:
+      a: a
+      b: fallback
+  - note: import_shadows_module_local_rule_for_function_calls
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        lib.double(x) := 100 * x
+
+        main := lib.double(3)
+    query: data.rules.main
+    want_result: 6
+  - note: import_shadows_module_local_value_rule_for_function_calls
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        lib.double := 7
+
+        main := lib.double(3)
+    query: data.rules.main
+    want_result: 6
+  - note: import_fully_qualified_function_call
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := data.lib.double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: unknown_function_through_import
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        double(x) := 2 * x
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.missing(1)
+    query: data.rules.result
+    want_error: "Unknown function"
+  - note: function_call_without_imports
+    modules:
+      - |
+        package rules
+
+        double(x) := 2 * x
+
+        result := double(21)
+    query: data.rules.result
+    want_result: 42
+  - note: non_function_value_rule_through_import
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        x := 5
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.x(1)
+    query: data.rules.result
+    want_error: "Unknown function"
+  - note: rule_path_prefix_through_import
+    modules:
+      - |
+        package lib
+        import rego.v1
+
+        a.c := 9
+      - |
+        package rules
+        import rego.v1
+        import data.lib
+
+        result := lib.a(7)
+    query: data.rules.result
+    want_error: "Unknown function"