Improvements (microsoft#33)

1. Skip recording undefined variables 2. Parse `in` correctly if it is not imported. 3. base64.decode 4. Handle `with` modifier for qualified data and input. Signed-off-by: Anand Krishnamoorthi <[email protected]>
anakrish · Nov 2, 2023 · c53d002 · c53d002
1 parent 6228eaa
commit c53d002
Show file tree

Hide file tree

Showing 8 changed files with 150 additions and 84 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -14,6 +14,7 @@ log = "0.4.17"
 env_logger="0.10.0"
 lazy_static = "1.4.0"
 rand = "0.8.5"
+data-encoding = "2.4.0"
 
 [dev-dependencies]
 clap = { version = "4.4.7", features = ["derive"] }

diff --git a/src/builtins/encoding.rs b/src/builtins/encoding.rs
@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use crate::ast::Expr;
+use crate::builtins;
+use crate::builtins::utils::{ensure_args_count, ensure_string};
+use crate::lexer::Span;
+use crate::value::Value;
+
+use std::collections::HashMap;
+
+use anyhow::Result;
+use data_encoding::BASE64;
+
+pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {
+    m.insert("base64.decode", (base64_decode, 1));
+}
+
+fn base64_decode(span: &Span, params: &[Expr], args: &[Value]) -> Result<Value> {
+    let name = "base64.decode";
+    ensure_args_count(span, name, params, args, 1)?;
+
+    let encoded_str = ensure_string(name, &params[0], &args[0])?;
+    let decoded_bytes = BASE64.decode(encoded_str.as_bytes())?;
+    Ok(Value::String(
+        String::from_utf8_lossy(&decoded_bytes).to_string(),
+    ))
+}
diff --git a/src/builtins/mod.rs b/src/builtins/mod.rs
@@ -7,6 +7,7 @@ mod bitwise;
 pub mod comparison;
 mod conversions;
 mod debugging;
+mod encoding;
 pub mod numbers;
 mod objects;
 pub mod sets;
@@ -44,7 +45,7 @@ lazy_static! {
 	conversions::register(&mut m);
 	//units::register(&mut m);
 	types::register(&mut m);
-	//encoding::register(&mut m);
+	encoding::register(&mut m);
 	//token_signing::register(&mut m);
 	//token_verification::register(&mut m);
 	//time::register(&mut m);

diff --git a/src/interpreter.rs b/src/interpreter.rs
@@ -26,6 +26,7 @@ pub struct Interpreter<'source> {
     input: Value,
     data: Value,
     init_data: Value,
+    with_document: Value,
     scopes: Vec<Scope>,
     // TODO: handle recursive calls where same expr could have different values.
     loop_var_values: BTreeMap<&'source Expr<'source>, Value>,
@@ -81,6 +82,9 @@ struct LoopExpr<'source> {
 
 impl<'source> Interpreter<'source> {
     pub fn new(modules: Vec<&'source Module<'source>>) -> Result<Interpreter<'source>> {
+        let mut with_document = Value::new_object();
+        *Self::make_or_get_value_mut(&mut with_document, &["data"])? = Value::new_object();
+        *Self::make_or_get_value_mut(&mut with_document, &["input"])? = Value::new_object();
         Ok(Interpreter {
             modules,
             module: None,
@@ -90,6 +94,7 @@ impl<'source> Interpreter<'source> {
             input: Value::new_object(),
             data: Value::new_object(),
             init_data: Value::new_object(),
+            with_document,
             scopes: vec![Scope::new()],
             contexts: vec![],
             loop_var_values: BTreeMap::new(),
@@ -241,6 +246,9 @@ impl<'source> Interpreter<'source> {
         match ident {
             "_" => true,
             _ => match self.lookup_local_var(ident) {
+                // Vars declared using `some v` can be loop vars.
+                // They are initialized to undefined.
+                Some(Value::Undefined) => true,
                 // If ident is a local var (in current or parent scopes),
                 // then it is not a loop var.
                 Some(_) => false,
@@ -478,6 +486,11 @@ impl<'source> Interpreter<'source> {
             }
         };
 
+        // Omit recording undefined values.
+        if value == Value::Undefined {
+            return Ok(Value::Bool(false));
+        }
+
         self.add_variable_or(name)?;
 
         // TODO: optimize this
@@ -803,44 +816,12 @@ impl<'source> Interpreter<'source> {
         Value::from_map(expr)
     }
 
-    fn eval_stmt(
+    fn eval_stmt_impl(
         &mut self,
         stmt: &'source LiteralStmt<'source>,
         stmts: &[&'source LiteralStmt<'source>],
     ) -> Result<bool> {
-        let mut to_restore = vec![];
-        for wm in &stmt.with_mods {
-            // Evaluate value and ref
-            let value = self.eval_expr(&wm.r#as)?;
-            let path = Parser::get_path_ref_components(&wm.refr)?;
-            let mut path: Vec<&str> = path.iter().map(|s| s.text()).collect();
-
-            // TODO: multiple modules and qualified path
-            if path.len() > 2 && format!("{}.{}", path[0], path[1]) == self.current_module_path {
-                path = path[1..].to_vec();
-            }
-
-            // Set new values in modifications table
-            let mut saved = false;
-            for (i, _) in path.iter().enumerate() {
-                let vref = Self::make_or_get_value_mut(&mut self.data, &path[0..i])?;
-                if vref == &Value::Undefined {
-                    to_restore.push((path[0..i].to_vec(), vref.clone()));
-                    saved = false;
-                    break;
-                }
-            }
-
-            // TODO: input
-            let vref = Self::make_or_get_value_mut(&mut self.data, &path[..])?;
-            if !saved {
-                to_restore.push((path, vref.clone()));
-            }
-
-            *vref = value;
-        }
-
-        let r = Ok(match &stmt.literal {
+        Ok(match &stmt.literal {
             Literal::Expr { span, expr, .. } => {
                 let value = match expr {
                     Expr::Call { span, fcn, params } => self.eval_call(
@@ -944,16 +925,59 @@ impl<'source> Interpreter<'source> {
                 }
                 self.eval_every(span, key, value, domain, query)?
             }
-        });
+        })
+    }
 
-        for (path, value) in to_restore.into_iter().rev() {
-            if value == Value::Undefined {
-                unimplemented!("handle undefined restore");
-            } else {
-                let vref = Self::make_or_get_value_mut(&mut self.data, &path[..])?;
-                *vref = value;
-            }
+    fn eval_stmt(
+        &mut self,
+        stmt: &'source LiteralStmt<'source>,
+        stmts: &[&'source LiteralStmt<'source>],
+    ) -> Result<bool> {
+        let saved_state = if !stmt.with_mods.is_empty() {
+            // Save state;
+            let with_document = self.with_document.clone();
+            let input = self.input.clone();
+            let data = self.data.clone();
+            let processed = self.processed.clone();
+
+            // Apply with modifiers.
+            for wm in &stmt.with_mods {
+                // Evaluate value and ref
+                let value = self.eval_expr(&wm.r#as)?;
+                let path = Parser::get_path_ref_components(&wm.refr)?;
+                let path: Vec<&str> = path.iter().map(|s| s.text()).collect();
+
+                if path[0] == "input" || path[0] == "data" {
+                    *Self::make_or_get_value_mut(&mut self.with_document, &path[..])? = value;
+                } /* else if path.len() == 1 {
+                      // TODO: handle var in current module.
+                  } else {
+                      // TODO: error about input, data
+                  } */
+                // TODO: functions
+            }
+
+            self.data = self.with_document["data"].clone();
+            self.input = self.with_document["input"].clone();
+            self.processed.clear();
+
+            (with_document, input, data, processed)
+        } else {
+            (
+                Value::Undefined,
+                Value::Undefined,
+                Value::Undefined,
+                BTreeSet::new(),
+            )
+        };
+
+        let r = self.eval_stmt_impl(stmt, stmts);
+
+        // Restore state.
+        if saved_state.0 != Value::Undefined {
+            (self.with_document, self.input, self.data, self.processed) = saved_state;
         }
+
         r
     }
 
@@ -989,8 +1013,8 @@ impl<'source> Interpreter<'source> {
             if let Some(idx) = self.lookup_local_var(loop_expr.index) {
                 if loop_expr_value[&idx] != Value::Undefined {
                     result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result;
+                    return Ok(result);
                 }
-                return Ok(result);
             }
 
             // Save the current scope and restore it after evaluating the statements so
@@ -1672,6 +1696,14 @@ impl<'source> Interpreter<'source> {
 
         // Ensure that rules are evaluated
         if name == "data" {
+            let v = Self::get_value_chained(self.data.clone(), fields);
+
+            // If the rule has already been evaluated or specified via a with modifier,
+            // use that value.
+            if v != Value::Undefined {
+                return Ok(v);
+            }
+
             // Evaluate rule corresponding to longest matching path.
             for i in (1..fields.len() + 1).rev() {
                 let path = "data.".to_owned() + &fields[0..i].join(".");
@@ -1680,22 +1712,31 @@ impl<'source> Interpreter<'source> {
                     break;
                 }
             }
+
             Ok(Self::get_value_chained(self.data.clone(), fields))
         } else if !self.modules.is_empty() {
-            // Add module prefix and ensure that any matching rule is evaluated.
-            let module_path =
-                Self::get_path_string(&self.current_module()?.package.refr, Some("data"))?;
-            let path = module_path + "." + name;
-
-            self.ensure_rule_evaluated(path)?;
-
             let mut path: Vec<&str> =
                 Parser::get_path_ref_components(&self.module.unwrap().package.refr)?
                     .iter()
                     .map(|s| s.text())
                     .collect();
             path.push(name);
 
+            let v = Self::get_value_chained(self.data.clone(), &path);
+
+            // If the rule has already been evaluated or specified via a with modifier,
+            // use that value.
+            if v != Value::Undefined {
+                return Ok(v);
+            }
+
+            // Add module prefix and ensure that any matching rule is evaluated.
+            let module_path =
+                Self::get_path_string(&self.current_module()?.package.refr, Some("data"))?;
+            let rule_path = module_path + "." + name;
+
+            self.ensure_rule_evaluated(rule_path)?;
+
             let value = Self::get_value_chained(self.data.clone(), &path[..]);
             Ok(Self::get_value_chained(value, fields))
         } else {

diff --git a/src/parser.rs b/src/parser.rs
@@ -714,7 +714,7 @@ impl<'source> Parser<'source> {
         let start = self.tok.1.start;
         let mut expr = self.parse_bool_expr()?;
 
-        while self.tok.1.text() == "in" {
+        while self.tok.1.text() == "in" && self.future_keywords.get("in").is_some() {
             expr = self.parse_membership_tail(start, expr, None)?;
         }
 
@@ -731,7 +731,7 @@ impl<'source> Parser<'source> {
             expr = self.parse_membership_tail(start, expr, Some(value))?;
         }
 
-        while self.tok.1.text() == "in" {
+        while self.tok.1.text() == "in" && self.is_imported_future_keyword("in") {
             expr = self.parse_membership_tail(start, expr, None)?;
         }
 
@@ -834,7 +834,7 @@ impl<'source> Parser<'source> {
             vars.push(span);
         }
 
-        if self.tok.1.text() != "in" || self.future_keywords.get("in").is_none() {
+        if self.tok.1.text() != "in" || !self.is_imported_future_keyword("in") {
             if self.tok.1.text() == "in" {
                 self.warn_future_keyword();
             }

diff --git a/src/scheduler.rs b/src/scheduler.rs
@@ -913,8 +913,10 @@ impl<'a> Analyzer<'a> {
                         collection,
                         &mut scope,
                         &mut first_use,
-                        &mut col_definitions, // TODO: handle these definitions
+                        &mut col_definitions,
                     )?;
+                    definitions.append(&mut col_definitions);
+
                     self.process_comprs(
                         &col_comprs[..],
                         &mut scope,