Implement import keyword (#101)

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

src/engine.rs

@@ -100,6 +100,7 @@ impl Engine {
             self.interpreter
                 .set_functions(gather_functions(&self.modules)?);
             self.interpreter.gather_rules()?;
+            self.interpreter.process_imports()?;
             self.prepared = true;
         }
 

src/interpreter.rs

@@ -64,6 +64,7 @@ pub struct Interpreter {
     traces: Option<Vec<std::rc::Rc<str>>>,
     allow_deprecated: bool,
     strict_builtin_errors: bool,
+    imports: BTreeMap<String, Ref<Expr>>,
 }
 
 impl Default for Interpreter {
@@ -209,6 +210,7 @@ impl Interpreter {
             traces: None,
             allow_deprecated: true,
             strict_builtin_errors: true,
+            imports: BTreeMap::default(),
         }
     }
 
@@ -2437,13 +2439,14 @@ impl Interpreter {
                 }
             }
         }
+
         Ok(())
     }
 
     fn lookup_var(&mut self, span: &Span, fields: &[&str], no_error: bool) -> Result<Value> {
         let name = span.source_str();
-
         debug_new_group!("lookup_var: name={name}, fields={fields:?}, no_error={no_error}");
+
         // Return local variable/argument.
         if let Some(v) = self.lookup_local_var(&name) {
             return Ok(Self::get_value_chained(v, fields));
@@ -2516,6 +2519,7 @@ impl Interpreter {
             if !no_error
                 && self.rules.get(&rule_path).is_none()
                 && self.default_rules.get(&rule_path).is_none()
+                && self.imports.get(&rule_path).is_none()
             {
                 bail!(span.error("var is unsafe"));
             }
@@ -2538,8 +2542,14 @@ impl Interpreter {
                 }
             }
 
-            // TODO: Is found needed?
-            let _ = found;
+            if !found {
+                if let Some(imported_var) = self.imports.get(&rule_path).cloned() {
+                    return Ok(Self::get_value_chained(
+                        self.eval_expr(&imported_var)?,
+                        fields,
+                    ));
+                }
+            }
 
             let value = Self::get_value_chained(self.data.clone(), &path[..]);
             Ok(Self::get_value_chained(value, fields))
@@ -3323,6 +3333,31 @@ impl Interpreter {
         Ok(())
     }
 
+    pub fn process_imports(&mut self) -> Result<()> {
+        for module in &self.modules {
+            let module_path = get_path_string(&module.package.refr, Some("data"))?;
+            for import in &module.imports {
+                let target = match &import.r#as {
+                    Some(s) => s.text(),
+                    _ => match import.refr.as_ref() {
+                        Expr::RefDot { field, .. } => field.text(),
+                        Expr::RefBrack { index, .. } => match index.as_ref() {
+                            Expr::String(s) => s.text(),
+                            _ => "",
+                        },
+                        _ => "",
+                    },
+                };
+                if target.is_empty() {
+                    bail!(import.refr.span().error("invalid ref in import"));
+                }
+                self.imports
+                    .insert(module_path.clone() + "." + target, import.refr.clone());
+            }
+        }
+        Ok(())
+    }
+
     pub fn gather_rules(&mut self) -> Result<()> {
         for module in self.modules.clone() {
             let prev_module = self.set_current_module(Some(module.clone()))?;

tests/interpreter/cases/import/tests.yaml

@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: import input
+    input:
+      a: 10
+    modules:
+      - |      
+        package a
+        import input as foo
+        b = foo
+      - |
+        package b
+        import input.a
+        import input.a as A
+
+        c = a
+        d = A
+    query: data
+    want_result:
+      a:
+        b:
+          a: 10
+      b:
+        c: 10
+        d: 10
+  - note: import data, cross ref
+    modules:
+      - |
+        package a
+        import data.b.a
+        b = a + 1
+        
+      - |
+        package b
+        # Both the following imports are overridden by rules
+        #import data.a.b as a
+        #import data.a.b
+
+        import data.a.b as C
+        import data.a.b
+
+        a = 10
+        c = C + b
+    query: data
+    want_result:
+      a:
+        b: 11
+      b:
+        a: 10
+        c: 22
+        
+  - note: import overridden by rule
+    modules:
+      - |
+        package a
+        a = 10
+      - |
+        package b
+        import data.a.a
+
+        a = 20
+    query: data.b.a
+    want_result: 20
+        
+  - note: invalid import ref
+    modules:
+      - |
+        package a
+        import foo
+    query: data
+    error: "import path must begin with one of"