From 7095e269b72149726190b681e2bcf73e11ec27c5 Mon Sep 17 00:00:00 2001
From: Anand Krishnamoorthi <35780660+anakrish@users.noreply.github.com>
Date: Sat, 27 Jul 2024 23:26:53 +0530
Subject: [PATCH] fix: Handle aliases in scheduler (#285)

Earlier scheduler only recognized rules and would raise an
`unsafe var` error on alias.

Register alias var names to fix this.

fixes #284

Also fix clippy warning treated as error

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>
---
 src/lib.rs                                |  2 ++
 src/scheduler.rs                          | 12 +++++++++---
 tests/interpreter/cases/import/tests.yaml |  6 ++++++
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 39f84b2..cf0b7f9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+#![allow(unknown_lints)]
+#![allow(clippy::doc_lazy_continuation)]
 // Use README.md as crate documentation.
 #![doc = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/README.md"))]
 // We'll default to building for no_std - use core, alloc instead of std.
diff --git a/src/scheduler.rs b/src/scheduler.rs
index b5469d9..1781542 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -412,7 +412,7 @@ impl Analyzer {
     }
 
     pub fn analyze(mut self, modules: &[Ref<Module>]) -> Result<Schedule> {
-        self.add_rules(modules)?;
+        self.add_rules_and_aliases(modules)?;
         self.functions = gather_functions(modules)?;
 
         for m in modules {
@@ -430,7 +430,7 @@ impl Analyzer {
         modules: &[Ref<Module>],
         query: &Ref<Query>,
     ) -> Result<Schedule> {
-        self.add_rules(modules)?;
+        self.add_rules_and_aliases(modules)?;
         self.analyze_query(None, None, query, Scope::default())?;
 
         Ok(Schedule {
@@ -439,7 +439,7 @@ impl Analyzer {
         })
     }
 
-    fn add_rules(&mut self, modules: &[Ref<Module>]) -> Result<()> {
+    fn add_rules_and_aliases(&mut self, modules: &[Ref<Module>]) -> Result<()> {
         for m in modules {
             let path = get_path_string(&m.package.refr, Some("data"))?;
             let scope: &mut Scope = self.packages.entry(path).or_default();
@@ -456,6 +456,12 @@ impl Analyzer {
                 };
                 scope.unscoped.insert(var);
             }
+
+            for import in &m.imports {
+                if let Some(var) = &import.r#as {
+                    scope.unscoped.insert(var.source_str());
+                }
+            }
         }
 
         Ok(())
diff --git a/tests/interpreter/cases/import/tests.yaml b/tests/interpreter/cases/import/tests.yaml
index b67db95..721a4ad 100644
--- a/tests/interpreter/cases/import/tests.yaml
+++ b/tests/interpreter/cases/import/tests.yaml
@@ -34,6 +34,7 @@ cases:
         
       - |
         package b
+        import rego.v1
         # Both the following imports are overridden by rules
         #import data.a.b as a
         #import data.a.b
@@ -43,6 +44,10 @@ cases:
 
         a = 10
         c = C + b
+
+        r if {
+          some v in [C]
+        }
     query: data
     want_result:
       a:
@@ -50,6 +55,7 @@ cases:
       b:
         a: 10
         c: 22
+        r: true
         
   - note: import overridden by rule
     modules: