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

Author: michaelstaib

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)