early return

Cargo.toml

@@ -102,7 +102,7 @@ itertools = "0.12.1"
 [dev-dependencies]
 cfg-if = "1.0.0"
 clap = { version = "4.4.7", features = ["derive"] }
-colored-diff = "0.2.3"
+prettydiff = { version = "0.6.4", default-features = false }
 serde_yaml = "0.9.16"
 test-generator = "0.3.1"
 walkdir = "2.3.2"

src/interpreter.rs

@@ -94,6 +94,8 @@ struct Context {
     rule_value: Value,
     is_set: bool,
     is_old_style_set: bool,
+    output_constness_determined: bool,
+    early_return: bool,
 }
 
 impl Default for Context {
@@ -109,6 +111,8 @@ impl Default for Context {
             rule_value: Value::new_object(),
             is_set: false,
             is_old_style_set: false,
+            output_constness_determined: false,
+            early_return: false,
         }
     }
 }
@@ -1428,6 +1432,9 @@ impl Interpreter {
                         Self::clear_scope(self.current_scope_mut()?);
                         if let Some(ctx) = self.contexts.last_mut() {
                             ctx.result = query_result.clone();
+                            if ctx.early_return {
+                                break;
+                            }
                         }
                     }
 
@@ -1452,6 +1459,9 @@ impl Interpreter {
                         Self::clear_scope(self.current_scope_mut()?);
                         if let Some(ctx) = self.contexts.last_mut() {
                             ctx.result = query_result.clone();
+                            if ctx.early_return {
+                                break;
+                            }
                         }
                     }
                     self.loop_var_values.remove(&loop_expr.expr());
@@ -1474,6 +1484,9 @@ impl Interpreter {
                         Self::clear_scope(self.current_scope_mut()?);
                         if let Some(ctx) = self.contexts.last_mut() {
                             ctx.result = query_result.clone();
+                            if ctx.early_return {
+                                break;
+                            }
                         }
                     }
                     self.loop_var_values.remove(&loop_expr.expr());
@@ -1586,30 +1599,95 @@ impl Interpreter {
         Ok(())
     }
 
+    // A ref is a constant ref, if it does not contain any local variables.
+    // For now, we restrict constant refs to those that contain only simple literals.
+    fn is_constant_ref(&self, mut expr: &Ref<Expr>) -> Result<bool> {
+        loop {
+            match expr.as_ref() {
+                Expr::Var(_) => break,
+                Expr::RefDot { refr, .. } => expr = refr,
+                Expr::RefBrack { refr, index, .. } if self.is_simple_literal(index)? => expr = refr,
+                _ => return Ok(false),
+            }
+        }
+        Ok(true)
+    }
+
+    fn is_simple_literal(&self, expr: &Ref<Expr>) -> Result<bool> {
+        Ok(matches!(
+            expr.as_ref(),
+            Expr::String(_)
+                | Expr::RawString(_)
+                | Expr::True(_)
+                | Expr::False(_)
+                | Expr::Null(_)
+                | Expr::Number(_)
+        ))
+    }
+
+    // A rule's output expression is constant if it does not contain local variables.
+    // For now, we restrict output expressions to those that contain only simple literals.
+    fn is_constant_output(
+        &self,
+        key_expr: &Option<Ref<Expr>>,
+        output_expr: &Ref<Expr>,
+    ) -> Result<bool> {
+        let mut is_const = true;
+        if let Some(key_expr) = key_expr {
+            is_const = self.is_simple_literal(key_expr)?;
+        }
+        Ok(is_const && self.is_simple_literal(output_expr)?)
+    }
+
     fn eval_output_expr_in_loop(&mut self, loops: &[LoopExpr]) -> Result<bool> {
         if loops.is_empty() {
             let (key_expr, output_expr) = self.get_exprs_from_context()?;
 
             let ctx = self.get_current_context()?;
-            let (is_set, is_old_style_set) = (ctx.is_set, ctx.is_old_style_set);
+            let (is_set, is_old_style_set, is_rule, constness_determined) = (
+                ctx.is_set,
+                ctx.is_old_style_set,
+                !ctx.is_compr,
+                ctx.output_constness_determined,
+            );
+
             if let Some(rule_ref) = ctx.rule_ref.clone() {
+                let mut is_const_rule = if is_rule && !constness_determined {
+                    self.is_constant_ref(&rule_ref)?
+                } else {
+                    // Constness has already been determined or is not a rule.
+                    // Treat the expression as not constant.
+                    false
+                };
+
                 let mut comps = self.eval_rule_ref(&rule_ref)?;
                 if let Some(ke) = &key_expr {
                     comps.push(self.eval_expr(ke)?);
                 }
                 let output = if let Some(oe) = &output_expr {
+                    // Rule is constant only if its ref, key and output are constant.
+                    is_const_rule = is_const_rule && self.is_constant_output(&key_expr, oe)?;
                     self.eval_expr(oe)?
                 } else if is_old_style_set && !comps.is_empty() {
+                    // Rule's constness is determined only by its ref.
                     let output = comps[comps.len() - 1].clone();
                     comps.pop();
                     output
                 } else {
+                    // Rule's constness is determined only by its ref.
                     Value::Bool(true)
                 };
 
                 let comps_defined = comps.iter().all(|v| v != &Value::Undefined);
                 let ctx = self.contexts.last_mut().expect("no current context");
 
+                if is_const_rule {
+                    ctx.early_return = true;
+                }
+                if is_rule {
+                    ctx.output_constness_determined = true;
+                }
+
                 if output == Value::Undefined || !comps_defined {
                     return Ok(false);
                 }

src/parser.rs

@@ -114,7 +114,17 @@ impl<'source> Parser<'source> {
             }
             Expr::Var(v) => comps.push(v.0.clone()),
             Expr::String(s) => comps.push(s.0.clone()),
-            _ => bail!("internal error: not a simple ref"),
+            Expr::True(s) | Expr::False(s) | Expr::Null(s) => comps.push(s.clone()),
+            Expr::Number(s) => {
+                // Ensure that the span will be the serialized representation.
+                if *s.0.text() == s.1.to_json_str()? {
+                    comps.push(s.0.clone());
+                } else {
+                    bail!(refr.span().error("not a valid ref"));
+                }
+            }
+
+            _ => bail!(refr.span().error("not a valid ref")),
         }
         Ok(())
     }

src/tests/interpreter/mod.rs

@@ -73,10 +73,10 @@ fn match_values(computed: &Value, expected: &Value) -> Result<()> {
     if computed != expected {
         panic!(
             "{}",
-            colored_diff::PrettyDifference {
-                expected: &serde_yaml::to_string(&expected)?,
-                actual: &serde_yaml::to_string(&computed)?
-            }
+            prettydiff::diff_chars(
+                &serde_yaml::to_string(&expected)?,
+                &serde_yaml::to_string(&computed)?
+            )
         );
     }
     Ok(())
@@ -347,7 +347,7 @@ fn yaml_test(file: &str) -> Result<()> {
         Err(e) => {
             // If Err is returned, it doesn't always get printed by cargo test.
             // Therefore, panic with the error.
-            panic!("{}", e);
+            panic!("{e}");
         }
     }
 }

tests/aci/main.rs

@@ -92,10 +92,10 @@ fn run_aci_tests(dir: &Path) -> Result<()> {
                 Ok(actual) => {
                     println!(
                         "DIFF {}",
-                        colored_diff::PrettyDifference {
-                            expected: &serde_yaml::to_string(&case.want_result)?,
-                            actual: &serde_yaml::to_string(&actual)?
-                        }
+                        prettydiff::diff_chars(
+                            &serde_yaml::to_string(&case.want_result)?,
+                            &serde_yaml::to_string(&actual)?
+                        )
                     );
 
                     nfailures += 1;

tests/interpreter/cases/loop/basic.yaml

@@ -21,3 +21,87 @@ cases:
       x1: [[1, 0], [2, 1], [3, 2], [4, 3]]
       x2: [[1, 1], [2, 2], [3, 3], [4, 4]]
       x3: [["q", "p"], ["s", "r"]]
+
+  - note: early return
+    data: {}
+    modules:
+      - |
+        package test
+        import future.keywords
+
+        a = [1, "hello"]
+        # Implicit value
+        b1 {
+            a[_] + 1
+        }
+
+        # Literals
+        b2 := true { a[_] + 1 }
+        b3 := false { a[_] + 1 }
+        b4 := 1 { a[_] + 1 }
+        b5 := null { a[_] + 1 }
+        b6 := "hello" { a[_] + 1 }
+        b7 := `world` { a[_] + 1 }
+
+        # constant refs
+        c[null] := true { a[_] + 1 }
+        c["hello"] := false { a[_] + 1 }
+        c[`world`] := 1 { a[_] + 1 }
+        c[true] := null { a[_] + 1 }
+        c[false] := "hello" { a[_] + 1 }
+        c[7] := `world` { a[_] + 1 }
+
+        # Old style set must should also be considered for early return.
+        old.style { a[_] + 1 }
+
+        # Multi part constant refactor
+        multi[1]["hello"] := 5 { a[_] +1 }
+
+        # Two elements must be produced
+        d = [1 | [1,2][_] ]
+
+        # Non simple ref must not result in early return.
+        f[p] = 5 {
+          p := a[_]
+        }
+
+        f1[p] = 5 {
+          a[p]
+        }
+
+        # Contains syntax
+        g contains p if {
+          p := a[_]
+        }
+    query: data.test
+    want_result:
+      a: [1, "hello"]
+      b1: true
+      b2: true
+      b3: false
+      b4: 1
+      b5: null
+      b6: "hello"
+      b7: "world"
+      c:
+        null: true
+        "hello": false
+        "world": 1
+        true: null
+        false: "hello"
+        7: "world"
+      d: [1, 1]
+      f:
+        1: 5
+        "hello": 5
+      f1:
+        0: 5
+        1: 5
+      g:
+        set!: [1, "hello"]
+      multi:
+        1:
+          "hello": 5
+      old:
+        set!: ["style"]
+