chore: release (#226)

* chore: release

.github/workflows/pr.yml

@@ -30,6 +30,8 @@ jobs:
       run: cargo test -r --verbose
     - name: Run tests (ACI)
       run: cargo test -r --test aci
+    - name: Run tests (KATA)
+      run: cargo test -r --test kata
     - name: Run tests (OPA Conformance)
       run: >-
         cargo test -r --test opa --features opa-testutil,serde_json/arbitrary_precision  -- $(tr '\n' ' ' < tests/opa.passing)
@@ -39,6 +41,8 @@ jobs:
       run: cargo test -r --verbose --target x86_64-unknown-linux-musl
     - name: Run tests (MUSL ACI)
       run: cargo test -r --test aci --target x86_64-unknown-linux-musl
+    - name: Run tests (KATA ACI)
+      run: cargo test -r --test kata --target x86_64-unknown-linux-musl
     - name: Run tests (MUSL OPA Conformance)
       run: >-
         cargo test -r --test opa --features opa-testutil,serde_json/arbitrary_precision --target x86_64-unknown-linux-musl  -- $(tr '\n' ' ' < tests/opa.passing)

CHANGELOG.md

@@ -6,6 +6,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.1.5](https://github.com/microsoft/regorus/compare/regorus-v0.1.4...regorus-v0.1.5) - 2024-05-07
+
+### Added
+- Support policy files greater than 64KB in size ([#217](https://github.com/microsoft/regorus/pull/217))
+- Add tests for kata containers policies ([#221](https://github.com/microsoft/regorus/pull/221))
+- Support for OPA v0.64.0 ([#219](https://github.com/microsoft/regorus/pull/219))
+  - New builtin `json.marshal_with_options`
+### Changed
+- Improve example in readme ([#224](https://github.com/microsoft/regorus/pull/224))
+### Fixed
+- OPA Conformance: Do not interpret # within regular string ([#216](https://github.com/microsoft/regorus/pull/216))
+
+## [0.1.4](https://github.com/microsoft/regorus/compare/regorus-v0.1.3...regorus-v0.1.4) - 2024-04-22
+
+### Other
+- early return ([#189](https://github.com/microsoft/regorus/pull/189))
+- Fix anyhow dependency issues ([#208](https://github.com/microsoft/regorus/pull/208))
+- remove unused compact-rc dependency ([#207](https://github.com/microsoft/regorus/pull/207))
+
 ## [0.1.3](https://github.com/microsoft/regorus/compare/regorus-v0.1.2...regorus-v0.1.3) - 2024-04-11
 
 ### Other

Cargo.toml

@@ -11,7 +11,7 @@ members = [
 [package]
 name = "regorus"
 description = "A fast, lightweight Rego (OPA policy language) interpreter"
-version = "0.1.3"
+version = "0.1.5"
 edition = "2021"
 license-file = "LICENSE"
 repository = "https://github.com/microsoft/regorus"
@@ -69,7 +69,7 @@ full-opa = [
 opa-testutil = []
 
 [dependencies]
-anyhow = {version = "1.0.66", features = ["backtrace"] }
+anyhow = { version = "1.0.45", default-features=false }
 serde = {version = "1.0.150", features = ["derive", "rc"] }
 serde_json = "1.0.89"
 serde_yaml = {version = "0.9.16", optional = true }
@@ -96,20 +96,19 @@ uuid = { version = "1.6.1", features = ["v4", "fast-rng"], optional = true }
 jsonschema = { version = "0.17.1", default-features = false, optional = true }
 chrono = { version = "0.4.31", optional = true }
 chrono-tz = { version = "0.8.5", optional = true }
-compact-rc = "0.5.2"
 jsonwebtoken = { version = "9.2.0", optional = true }
 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"
 
 [build-dependencies]
-anyhow = "1.0.66"
+anyhow = "1.0"
 
 [profile.release]
 debug = true
@@ -125,6 +124,11 @@ name="aci"
 harness=false
 test=false
 
+[[test]]
+name="kata"
+harness=false
+test=false
+
 [package.metadata.docs.rs]
 # To build locally:
 # RUSTDOCFLAGS="--cfg docsrs" cargo +nightly doc --all-features --no-deps

README.md

@@ -9,7 +9,7 @@
 Regorus is also
   - *cross-platform* - Written in platform-agnostic Rust.
   - *current* - We strive to keep Regorus up to date with latest OPA release. Regorus supports `import rego.v1`.
-  - *compliant* - Regorus is mostly compliant with the latest [OPA release v0.63.0](https://github.com/open-policy-agent/opa/releases/tag/v0.63.0). See [OPA Conformance](#opa-conformance) for details. Note that while we behaviorally produce the same results, we don't yet support all the builtins.
+  - *compliant* - Regorus is mostly compliant with the latest [OPA release v0.64.0](https://github.com/open-policy-agent/opa/releases/tag/v0.64.0). See [OPA Conformance](#opa-conformance) for details. Note that while we behaviorally produce the same results, we don't yet support all the builtins.
   - *extensible* - Extend the Rego language by implementing custom stateful builtins in Rust.
     See [add_extension](https://github.com/microsoft/regorus/blob/fc68bf9c8bea36427dae9401a7d1f6ada771f7ab/src/engine.rs#L352).
     Support for extensibility using other languages coming soon.
@@ -24,31 +24,60 @@ Regorus is available as a library that can be easily integrated into your Rust p
 Here is an example of evaluating a simple Rego policy:
 
 ```rust
-use anyhow::Result;
-use regorus::*;
-use serde_json;
+fn main() -> anyhow::Result<()> {
+    // Create an engine for evaluating Rego policies.
+    let mut engine = regorus::Engine::new();
 
-fn main() -> Result<()> {
-  // Create an engine for evaluating Rego policies.
-  let mut engine = Engine::new();
+    let policy = String::from(
+        r#"
+       package example
+       import rego.v1
 
-  // Add policy to the engine.
-  engine.add_policy(
-    // Filename to be associated with the policy.
-    "hello.rego".to_string(),
+       allow if {
+          ## All actions are allowed for admins.
+          input.principal == "admin"
+       } else if {
+          ## Check if action is allowed for given user.
+          input.action in data.allowed_actions[input.principal]
+       }
+	"#,
+    );
 
-    // Rego policy that just sets a message.
-    r#"
-       package test
-       message = "Hello, World!"
-    "#.to_string()
-  )?;
+    // Add policy to the engine.
+    engine.add_policy(String::from("policy.rego"), policy)?;
 
-  // Evaluate the policy, fetch the message and print it.
-  let results = engine.eval_query("data.test.message".to_string(), false)?;
-  println!("{}", serde_json::to_string_pretty(&results)?);
+    // Add data to engine.
+    engine.add_data(regorus::Value::from_json_str(
+        r#"{
+     "allowed_actions": {
+        "user1" : ["read", "write"],
+        "user2" : ["read"]
+     }}"#,
+    )?)?;
 
-  Ok(())
+    // Set input and evaluate whether user1 can write.
+    engine.set_input(regorus::Value::from_json_str(
+        r#"{
+      "principal": "user1",
+      "action": "write"
+    }"#,
+    )?);
+
+    let r = engine.eval_rule(String::from("data.example.allow"))?;
+    assert_eq!(r, regorus::Value::from(true));
+
+    // Set input and evaluate whether user2 can write.
+    engine.set_input(regorus::Value::from_json_str(
+        r#"{
+      "principal": "user2",
+      "action": "write"
+    }"#,
+    )?);
+
+    let r = engine.eval_rule(String::from("data.example.allow"))?;
+    assert_eq!(r, regorus::Value::Undefined);
+
+    Ok(())
 }
 ```
 
@@ -69,7 +98,7 @@ $ cargo build -r --example regorus --features "yaml" --no-default-features; stri
 -rwxr-xr-x  1 anand  staff   2.9M Jan 19 11:26 target/release/examples/regorus*
 ```
 
-Regorus passes the [OPA v0.63.0 test-suite](https://www.openpolicyagent.org/docs/latest/ir/#test-suite) barring a few
+Regorus passes the [OPA v0.64.0 test-suite](https://www.openpolicyagent.org/docs/latest/ir/#test-suite) barring a few
 builtins. See [OPA Conformance](#opa-conformance) below.
 
 ## Bindings
@@ -245,7 +274,7 @@ Benchmark 1: opa eval -b tests/aci -d tests/aci/data.json -i tests/aci/input.jso
 ```
 ## OPA Conformance
 
-Regorus has been verified to be compliant with [OPA v0.63.0](https://github.com/open-policy-agent/opa/releases/tag/v0.63.0)
+Regorus has been verified to be compliant with [OPA v0.64.0](https://github.com/open-policy-agent/opa/releases/tag/v0.64.0)
 using a [test driver](https://github.com/microsoft/regorus/blob/main/tests/opa.rs) that loads and runs the OPA testsuite using Regorus, and verifies that expected outputs are produced.
 
 The test driver can be invoked by running:

bindings/ffi/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2021"
 crate-type = ["cdylib"]
 
 [dependencies]
-anyhow = "1.0.79"
+anyhow = "1.0"
 regorus = { path  = "../.." }
 serde_json = "1.0.113"
 

bindings/java/Cargo.toml

@@ -11,7 +11,7 @@ keywords = ["interpreter", "opa", "policy-as-code", "rego"]
 crate-type = ["cdylib"]
 
 [dependencies]
-anyhow = "1.0.79"
+anyhow = "1.0"
 serde_json = "1.0.112"
 jni = "0.21.1"
 regorus = { path = "../.." }

bindings/python/Cargo.toml

@@ -12,7 +12,7 @@ keywords = ["interpreter", "opa", "policy-as-code", "rego"]
 crate-type = ["cdylib"]
 
 [dependencies]
-anyhow = "1.0.79"
+anyhow = "1.0"
 ordered-float = "4.2.0"
 pyo3 = {version = "0.21.0", features = ["anyhow", "extension-module"] }
 regorus = { path  = "../.." }

docs/builtins.md

@@ -159,24 +159,25 @@ In future, each builtin will be associated with a feature (many builtins could b
   | [type_name](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-types-type_name)   | _       |
   
 - [Encoding](https://www.openpolicyagent.org/docs/latest/policy-reference/#encoding)
-  | Builtin                                                                                                                          | Feature     |
-  |----------------------------------------------------------------------------------------------------------------------------------|-------------|
-  | [base64.is_valid](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-base64is_valid)                 | `base64`    |
-  | [base64url.decode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-base64urldecode)               | `base64`    |
-  | [base64url.encode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-base64urlencode)               | `base64url` |
-  | [base64url.encode_no_pad](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-base64urlencode_no_pad) | `base64url` |
-  | [hex.decode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-hexdecode)                           | `hex`       |
-  | [hex.encode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-hexencode)                           | `hex`       |
-  | [json.is_valid](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-jsonis_valid)                     | _           |
-  | [json.marshal](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-jsonmarshal)                       | _           |
-  | [json.unmarshal](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-jsonunmarshal)                   | _           |
-  | [urlquery.decode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-urlquerydecode)                 | `urlquery`  |
-  | [urlquery.decode_object](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-urlquerydecode_object)   | `urlquery`  |
-  | [urlquery.encode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-urlqueryencode)                 | `urlquery`  |
-  | [urlquery.encode_object](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-urlqueryencode_object)   | `urlquery`  |
-  | [yaml.is_valid](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-yamlis_valid)                     | `yaml`      |
-  | [yaml.marshal](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-yamlmarshal)                       | `yaml`      |
-  | [yaml.unmarshal](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-yamlunmarshal)                   | `yaml`            |
+  | Builtin                                                                                                                              | Feature     |
+  |--------------------------------------------------------------------------------------------------------------------------------------|-------------|
+  | [base64.is_valid](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-base64is_valid)                     | `base64`    |
+  | [base64url.decode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-base64urldecode)                   | `base64`    |
+  | [base64url.encode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-base64urlencode)                   | `base64url` |
+  | [base64url.encode_no_pad](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-base64urlencode_no_pad)     | `base64url` |
+  | [hex.decode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-hexdecode)                               | `hex`       |
+  | [hex.encode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-hexencode)                               | `hex`       |
+  | [json.is_valid](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-jsonis_valid)                         | _           |
+  | [json.marshal](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-jsonmarshal)                           | _           |
+  | [json.marshal_with_options](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-jsonmarshal_with_options) | _           |
+  | [json.unmarshal](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-jsonunmarshal)                       | _           |
+  | [urlquery.decode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-urlquerydecode)                     | `urlquery`  |
+  | [urlquery.decode_object](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-urlquerydecode_object)       | `urlquery`  |
+  | [urlquery.encode](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-urlqueryencode)                     | `urlquery`  |
+  | [urlquery.encode_object](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-urlqueryencode_object)       | `urlquery`  |
+  | [yaml.is_valid](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-yamlis_valid)                         | `yaml`      |
+  | [yaml.marshal](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-yamlmarshal)                           | `yaml`      |
+  | [yaml.unmarshal](https://www.openpolicyagent.org/docs/latest/policy-reference/#builtin-encoding-yamlunmarshal)                       | `yaml`      |
 
 - [Time](https://www.openpolicyagent.org/docs/latest/policy-reference/#time)
    | Builtin                                                                                                                    | Feature |

scripts/pre-push

@@ -15,6 +15,7 @@ if [ -f Cargo.toml ]; then
    # Ensure that all tests pass
    cargo test -r
    cargo test -r --test aci
+   cargo test -r --test kata
 
    # Ensure that OPA conformance tests don't regress.
    cargo test -r --features opa-testutil,serde_json/arbitrary_precision --test opa -- $(tr '\n' ' ' < tests/opa.passing)

src/builtins/encoding.rs

@@ -42,6 +42,7 @@ pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {
     }
     m.insert("json.is_valid", (json_is_valid, 1));
     m.insert("json.marshal", (json_marshal, 1));
+    m.insert("json.marshal_with_options", (json_marshal_with_options, 2));
     m.insert("json.unmarshal", (json_unmarshal, 1));
 
     #[cfg(feature = "yaml")]
@@ -368,6 +369,74 @@ fn json_marshal(span: &Span, params: &[Ref<Expr>], args: &[Value], _strict: bool
     ))
 }
 
+fn json_marshal_with_options(
+    span: &Span,
+    params: &[Ref<Expr>],
+    args: &[Value],
+    _strict: bool,
+) -> Result<Value> {
+    let name = "json.marshal_with_options";
+    ensure_args_count(span, name, params, args, 2)?;
+
+    let options = ensure_object(name, &params[1], args[1].clone())?;
+    let (mut pretty, mut indent, mut prefix) = (true, Some("\t".to_owned()), None);
+    for (option, option_value) in options.iter() {
+        match option {
+            Value::String(s) if s.as_ref() == "pretty" && option_value.as_bool().is_ok() => {
+                pretty = option_value == &Value::Bool(true);
+            }
+            Value::String(s) if s.as_ref() == "pretty" => bail!(params[1]
+                .span()
+                .error("marshaling option `pretty` must be true or false")),
+            Value::String(s) if s.as_ref() == "prefix" && option_value.as_string().is_ok() => {
+                prefix = Some(option_value.as_string()?.as_ref().to_string());
+            }
+            Value::String(s) if s.as_ref() == "prefix" => bail!(params[1]
+                .span()
+                .error("marshaling option `pretty` must be string")),
+            Value::String(s) if s.as_ref() == "indent" && option_value.as_string().is_ok() => {
+                indent = Some(option_value.as_string()?.as_ref().to_string());
+            }
+            Value::String(s) if s.as_ref() == "indent" => bail!(params[1]
+                .span()
+                .error("marshaling option `pretty` must be string")),
+            _ => bail!(params[1]
+                .span()
+                .error("marshaling option must be one of `indent`, `prefix` or `pretty`")),
+        }
+    }
+
+    if !pretty || options.is_empty() {
+        return Ok(Value::String(
+            serde_json::to_string(&args[0])
+                .with_context(|| span.error("could not serialize to json"))?
+                .into(),
+        ));
+    }
+
+    let lines: Vec<String> = serde_json::to_string_pretty(&args[0])
+        .with_context(|| span.error("could not serialize to json"))?
+        .split('\n')
+        .map(|line| {
+            let mut line = line.to_string();
+
+            if let Some(indent) = &indent {
+                let start_trimmed = line.trim_start();
+                let leading_spaces = line.len() - start_trimmed.len();
+                let indentation_level = leading_spaces / 2;
+                line = indent.repeat(indentation_level) + start_trimmed;
+            }
+
+            if let Some(prefix) = &prefix {
+                line = prefix.to_owned() + &line;
+            }
+            line
+        })
+        .collect();
+
+    Ok(Value::from(lines.join("\n")))
+}
+
 fn json_unmarshal(
     span: &Span,
     params: &[Ref<Expr>],

src/builtins/semver.rs

@@ -12,7 +12,7 @@ use semver::Version;
 use std::cmp::Ordering;
 use std::collections::HashMap;
 
-use anyhow::{Ok, Result};
+use anyhow::Result;
 
 pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {
     m.insert("semver.compare", (compare, 2));

src/builtins/uuid.rs

@@ -9,7 +9,7 @@ use crate::value::Value;
 
 use std::collections::{BTreeMap, HashMap};
 
-use anyhow::{Ok, Result};
+use anyhow::Result;
 use uuid::{Timestamp, Uuid};
 
 pub fn register(m: &mut HashMap<&'static str, builtins::BuiltinFcn>) {

src/engine.rs

@@ -65,7 +65,7 @@ impl Engine {
     /// ```
     ///
     pub fn add_policy(&mut self, path: String, rego: String) -> Result<()> {
-        let source = Source::new(path, rego);
+        let source = Source::from_contents(path, rego)?;
         let mut parser = Parser::new(&source)?;
         self.modules.push(Ref::new(parser.parse()?));
         // if policies change, interpreter needs to be prepared again
@@ -294,15 +294,15 @@ impl Engine {
 
         self.interpreter.create_rule_prefixes()?;
         let query_module = {
-            let source = Source::new(
+            let source = Source::from_contents(
                 "<query_module.rego>".to_owned(),
                 "package __internal_query_module".to_owned(),
-            );
+            )?;
             Ref::new(Parser::new(&source)?.parse()?)
         };
 
         // Parse the query.
-        let query_source = Source::new("<query.rego>".to_string(), query);
+        let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
         let mut parser = Parser::new(&query_source)?;
         let query_node = parser.parse_user_query()?;
         if query_node.span.text() == "data" {
@@ -411,15 +411,15 @@ impl Engine {
         self.eval_modules(enable_tracing)?;
 
         let query_module = {
-            let source = Source::new(
+            let source = Source::from_contents(
                 "<query_module.rego>".to_owned(),
                 "package __internal_query_module".to_owned(),
-            );
+            )?;
             Ref::new(Parser::new(&source)?.parse()?)
         };
 
         // Parse the query.
-        let query_source = Source::new("<query.rego>".to_string(), query);
+        let query_source = Source::from_contents("<query.rego>".to_string(), query)?;
         let mut parser = Parser::new(&query_source)?;
         let query_node = parser.parse_user_query()?;
         let query_schedule = Analyzer::new().analyze_query_snippet(&self.modules, &query_node)?;

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/lexer.rs

@@ -18,7 +18,7 @@ use anyhow::{anyhow, bail, Result};
 struct SourceInternal {
     pub file: String,
     pub contents: String,
-    pub lines: Vec<(u16, u16)>,
+    pub lines: Vec<(u32, u32)>,
 }
 
 #[derive(Clone)]
@@ -61,8 +61,8 @@ impl Debug for Source {
 #[derive(Clone)]
 pub struct SourceStr {
     source: Source,
-    start: u16,
-    end: u16,
+    start: u32,
+    end: u32,
 }
 
 impl Debug for SourceStr {
@@ -78,7 +78,7 @@ impl std::fmt::Display for SourceStr {
 }
 
 impl SourceStr {
-    pub fn new(source: Source, start: u16, end: u16) -> Self {
+    pub fn new(source: Source, start: u32, end: u32) -> Self {
         Self { source, start, end }
     }
 
@@ -116,39 +116,43 @@ impl std::cmp::Ord for SourceStr {
 }
 
 impl Source {
-    pub fn new(file: String, contents: String) -> Source {
+    pub fn from_contents(file: String, contents: String) -> Result<Source> {
+        let max_size = u32::MAX as usize - 2; // Account for rows, cols possibly starting at 1, EOF etc.
+        if contents.len() > max_size {
+            bail!("{file} exceeds maximum allowed policy file size {max_size}");
+        }
         let mut lines = vec![];
         let mut prev_ch = ' ';
-        let mut prev_pos = 0u16;
-        let mut start = 0u16;
+        let mut prev_pos = 0u32;
+        let mut start = 0u32;
         for (i, ch) in contents.char_indices() {
             if ch == '\n' {
                 let end = match prev_ch {
                     '\r' => prev_pos,
-                    _ => i as u16,
+                    _ => i as u32,
                 };
                 lines.push((start, end));
-                start = i as u16 + 1;
+                start = i as u32 + 1;
             }
             prev_ch = ch;
-            prev_pos = i as u16;
+            prev_pos = i as u32;
         }
 
         if (start as usize) < contents.len() {
-            lines.push((start, contents.len() as u16));
+            lines.push((start, contents.len() as u32));
         } else if contents.is_empty() {
             lines.push((0, 0));
         } else {
-            let s = (contents.len() - 1) as u16;
+            let s = (contents.len() - 1) as u32;
             lines.push((s, s));
         }
-        Self {
+        Ok(Self {
             src: Rc::new(SourceInternal {
                 file,
                 contents,
                 lines,
             }),
-        }
+        })
     }
 
     pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Source> {
@@ -157,10 +161,7 @@ impl Source {
             Err(e) => bail!("Failed to read {}. {e}", path.as_ref().display()),
         };
         // TODO: retain path instead of converting to string
-        Ok(Self::new(
-            path.as_ref().to_string_lossy().to_string(),
-            contents,
-        ))
+        Self::from_contents(path.as_ref().to_string_lossy().to_string(), contents)
     }
 
     pub fn file(&self) -> &String {
@@ -169,7 +170,7 @@ impl Source {
     pub fn contents(&self) -> &String {
         &self.src.contents
     }
-    pub fn line(&self, idx: u16) -> &str {
+    pub fn line(&self, idx: u32) -> &str {
         let idx = idx as usize;
         if idx < self.src.lines.len() {
             let (start, end) = self.src.lines[idx];
@@ -179,7 +180,7 @@ impl Source {
         }
     }
 
-    pub fn message(&self, line: u16, col: u16, kind: &str, msg: &str) -> String {
+    pub fn message(&self, line: u32, col: u32, kind: &str, msg: &str) -> String {
         if line as usize > self.src.lines.len() {
             return format!("{}: invalid line {} specified", self.src.file, line);
         }
@@ -206,7 +207,7 @@ impl Source {
         )
     }
 
-    pub fn error(&self, line: u16, col: u16, msg: &str) -> anyhow::Error {
+    pub fn error(&self, line: u32, col: u32, msg: &str) -> anyhow::Error {
         anyhow!(self.message(line, col, "error", msg))
     }
 }
@@ -214,10 +215,10 @@ impl Source {
 #[derive(Clone)]
 pub struct Span {
     pub source: Source,
-    pub line: u16,
-    pub col: u16,
-    pub start: u16,
-    pub end: u16,
+    pub line: u32,
+    pub col: u32,
+    pub start: u32,
+    pub end: u32,
 }
 
 impl Span {
@@ -272,8 +273,8 @@ pub struct Token(pub TokenKind, pub Span);
 pub struct Lexer<'source> {
     source: Source,
     iter: Peekable<CharIndices<'source>>,
-    line: u16,
-    col: u16,
+    line: u32,
+    col: u32,
 }
 
 impl<'source> Lexer<'source> {
@@ -312,15 +313,15 @@ impl<'source> Lexer<'source> {
             }
         }
         let end = self.peek().0;
-        self.col += (end - start) as u16;
+        self.col += (end - start) as u32;
         Ok(Token(
             TokenKind::Ident,
             Span {
                 source: self.source.clone(),
                 line: self.line,
                 col,
-                start: start as u16,
-                end: end as u16,
+                start: start as u32,
+                end: end as u32,
             },
         ))
     }
@@ -363,7 +364,7 @@ impl<'source> Lexer<'source> {
         }
 
         let end = self.peek().0;
-        self.col += (end - start) as u16;
+        self.col += (end - start) as u32;
 
         // Check for invalid number.Valid number cannot be followed by
         // these characters:
@@ -403,8 +404,8 @@ impl<'source> Lexer<'source> {
                 source: self.source.clone(),
                 line: self.line,
                 col,
-                start: start as u16,
-                end: end as u16,
+                start: start as u32,
+                end: end as u32,
             },
         ))
     }
@@ -440,8 +441,8 @@ impl<'source> Lexer<'source> {
                 source: self.source.clone(),
                 line,
                 col,
-                start: start as u16,
-                end: end as u16 - 1,
+                start: start as u32,
+                end: end as u32 - 1,
             },
         ))
     }
@@ -453,9 +454,9 @@ impl<'source> Lexer<'source> {
         let (start, _) = self.peek();
         loop {
             let (offset, ch) = self.peek();
-            let col = self.col + (offset - start) as u16;
+            let col = self.col + (offset - start) as u32;
             match ch {
-                '"' | '#' | '\x00' => {
+                '"' | '\x00' => {
                     break;
                 }
                 '\\' => {
@@ -468,7 +469,7 @@ impl<'source> Lexer<'source> {
                         'u' => {
                             for _i in 0..4 {
                                 let (offset, ch) = self.peek();
-                                let col = self.col + (offset - start) as u16;
+                                let col = self.col + (offset - start) as u32;
                                 if !ch.is_ascii_hexdigit() {
                                     return Err(self.source.error(
                                         line,
@@ -484,7 +485,7 @@ impl<'source> Lexer<'source> {
                 }
                 _ => {
                     // check for valid json chars
-                    let col = self.col + (offset - start) as u16;
+                    let col = self.col + (offset - start) as u32;
                     if !('\u{0020}'..='\u{10FFFF}').contains(&ch) {
                         return Err(self.source.error(line, col, "invalid character in string"));
                     }
@@ -499,7 +500,7 @@ impl<'source> Lexer<'source> {
 
         self.iter.next();
         let end = self.peek().0;
-        self.col += (end - start) as u16;
+        self.col += (end - start) as u32;
 
         // Ensure that the string is parsable in Rust.
         match serde_json::from_str::<String>(&self.source.contents()[start - 1..end]) {
@@ -522,8 +523,8 @@ impl<'source> Lexer<'source> {
                 source: self.source.clone(),
                 line,
                 col: col + 1,
-                start: start as u16,
-                end: end as u16 - 1,
+                start: start as u32,
+                end: end as u32 - 1,
             },
         ))
     }
@@ -593,14 +594,14 @@ impl<'source> Lexer<'source> {
 		    source: self.source.clone(),
 		    line: self.line,
 		    col,
-		    start: start as u16,
-		    end: start as u16 + 1,
+		    start: start as u32,
+		    end: start as u32 + 1,
 		}))
 	    }
 	    ':' => {
 		self.col += 1;
 		self.iter.next();
-		let mut end = start as u16 + 1;
+		let mut end = start as u32 + 1;
 		if self.peek().1 == '=' {
 		    self.col += 1;
 		    self.iter.next();
@@ -610,7 +611,7 @@ impl<'source> Lexer<'source> {
 		    source: self.source.clone(),
 		    line: self.line,
 		    col,
-		    start: start as u16,
+		    start: start as u32,
 		    end
 		}))
 	    }
@@ -626,8 +627,8 @@ impl<'source> Lexer<'source> {
 		    source: self.source.clone(),
 		    line: self.line,
 		    col,
-		    start: start as u16,
-		    end: self.peek().0 as u16,
+		    start: start as u32,
+		    end: self.peek().0 as u32,
 		}))
 	    }
 	    '!' if self.peekahead(1).1 == '=' => {
@@ -638,8 +639,8 @@ impl<'source> Lexer<'source> {
 		    source: self.source.clone(),
 		    line: self.line,
 		    col,
-		    start: start as u16,
-		    end: self.peek().0 as u16,
+		    start: start as u32,
+		    end: self.peek().0 as u32,
 		}))
 	    }
 	    '"' => self.read_string(),
@@ -648,8 +649,8 @@ impl<'source> Lexer<'source> {
 		source: self.source.clone(),
 		line:self.line,
 		col,
-		start: start as u16,
-		end: start as u16
+		start: start as u32,
+		end: start as u32
 	    })),
 	    _ if chr.is_ascii_digit() => self.read_number(),
 	    _ if chr.is_ascii_alphabetic() || chr == '_' => {

src/lib.rs

@@ -47,9 +47,9 @@ use std::rc::Rc;
 #[derive(Debug, Clone, Serialize, Eq, PartialEq)]
 pub struct Location {
     /// Line number. Starts at 1.
-    pub row: u16,
+    pub row: u32,
     /// Column number. Starts at 1.
-    pub col: u16,
+    pub col: u32,
 }
 
 /// An expression in a Rego query.

src/parser.rs

@@ -16,8 +16,8 @@ pub struct Parser<'source> {
     source: Source,
     lexer: Lexer<'source>,
     tok: Token,
-    line: u16,
-    end: u16,
+    line: u32,
+    end: u32,
     future_keywords: BTreeMap<String, Span>,
     rego_v1: bool,
 }
@@ -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(())
     }
@@ -743,7 +753,7 @@ impl<'source> Parser<'source> {
 
     fn parse_membership_tail(
         &mut self,
-        start: u16,
+        start: u32,
         mut expr1: Expr,
         mut expr2: Option<Expr>,
     ) -> Result<Expr> {

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}");
         }
     }
 }

src/tests/scheduler/analyzer/mod.rs

@@ -38,7 +38,10 @@ fn analyze_file(regos: &[String], expected_scopes: &[Scope]) -> Result<()> {
     let mut sources = vec![];
     let mut modules = vec![];
     for (idx, _) in regos.iter().enumerate() {
-        sources.push(Source::new(format!("rego_{idx}"), regos[idx].clone()));
+        sources.push(Source::from_contents(
+            format!("rego_{idx}"),
+            regos[idx].clone(),
+        )?);
     }
 
     for source in &sources {

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"]
+        

tests/kata/data/k8s-policy-job/output.json

@@ -0,0 +1,19 @@
+[
+  false,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true
+]

tests/kata/data/k8s-policy-pod/output.json

@@ -0,0 +1,19 @@
+[
+  false,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true
+]

tests/kata/data/k8s-policy-rc/output.json

@@ -0,0 +1,19 @@
+[
+  false,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true,
+  true
+]

tests/kata/main.rs

@@ -0,0 +1,115 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+use regorus::*;
+
+use std::path::Path;
+
+use anyhow::Result;
+use clap::Parser;
+use walkdir::WalkDir;
+
+fn run_kata_tests(tests_dir: &Path, generate: bool) -> Result<()> {
+    for entry in WalkDir::new(tests_dir)
+        .max_depth(1) // Do not recurse
+        .sort_by_file_name()
+        .into_iter()
+        .filter_map(|e| e.ok())
+    {
+        let path = entry.path();
+        if path == tests_dir || !path.is_dir() {
+            continue;
+        }
+
+        let policy_file = path.join("policy.rego");
+        let inputs_file = path.join("inputs.txt");
+        let outputs_file = path.join("output.json");
+
+        let mut engine = Engine::new();
+        engine.add_policy_from_file(&policy_file)?;
+        engine.set_gather_prints(true);
+        engine.set_strict_builtin_errors(false);
+
+        #[cfg(feature = "coverage")]
+        engine.set_enable_coverage(true);
+
+        // Keep a copy of the engine.
+        let engine_base = engine.clone();
+        let mut results = if generate {
+            vec![]
+        } else {
+            Value::from_json_str(&std::fs::read_to_string(&outputs_file)?)?
+                .as_array()?
+                .iter()
+                .cloned()
+                .rev()
+                .collect()
+        };
+
+        let inputs = std::fs::read_to_string(&inputs_file)?;
+        for (lineno, line) in inputs.split('\n').enumerate() {
+            let line = line.trim();
+            if line.is_empty() {
+                continue;
+            }
+            // Remove "ep":
+            let line = line.replace("\"ep\":", "");
+            // Remove trailing ,
+            let line = &line[0..line.len() - 1];
+
+            let request = Value::from_json_str(line)?;
+
+            let rule = format!("data.agent_policy.{}", request[0].as_string()?.as_ref());
+            let input = request[1].clone();
+
+            // Evaluate using engine.
+            engine.set_input(input.clone());
+            let r = engine.eval_rule(rule.clone())?;
+
+            // Evaluate using fresh engine.
+            let mut new_engine = engine_base.clone();
+            new_engine.set_input(input);
+            let r_new = new_engine.eval_rule(rule)?;
+
+            // Ensure that both evaluations produced the same result.
+            assert_eq!(r, r_new);
+
+            if generate {
+                results.push(r);
+            } else {
+                let expected = results.pop().unwrap();
+                assert_eq!(r, expected, "{lineno} failed in {}", inputs_file.display());
+            }
+        }
+
+        if generate {
+            std::fs::write(outputs_file, Value::from(results).to_json_str()?)?;
+        }
+        #[cfg(feature = "coverage")]
+        {
+            let report = engine.get_coverage_report()?;
+            println!("{}", report.to_colored_string()?);
+        }
+    }
+
+    println!("kata tests passed");
+    Ok(())
+}
+
+#[derive(clap::Parser)]
+#[command(author, version, about, long_about = None)]
+struct Cli {
+    /// Path to Kata test suite.
+    #[arg(long, short)]
+    #[clap(default_value = "tests/kata/data")]
+    test_dir: String,
+
+    /// Generate outputs instead of testing.
+    #[arg(long, short)]
+    #[clap(default_value = "false")]
+    generate: bool,
+}
+
+fn main() -> Result<()> {
+    let cli = Cli::parse();
+    run_kata_tests(&Path::new(&cli.test_dir), cli.generate)
+}

tests/lexer/cases/comment.yaml

@@ -56,11 +56,9 @@ cases:
 
   - note: comments-within-rstring
     rego: |
-      #Comments aren't allowed within strings,
-      "#This is not a comment
-      "#c
-    tokens:
-    error: unmatched "
+      #Comments aren't lexed within raw strings.
+      "#This is not a comment"#c
+    tokens: ["#This is not a comment", ""]
 
   - note: comment-integer-break
     rego: |

tests/lexer/mod.rs

@@ -77,7 +77,7 @@ fn yaml_test_impl(file: &str) -> Result<()> {
     let test: Test = serde_yaml::from_str(&yaml)?;
 
     for case in &test.cases {
-        let source = Source::new("case.rego".to_string(), case.rego.clone());
+        let source = Source::from_contents("case.rego".to_string(), case.rego.clone())?;
         print!("case {} ", &case.note);
 
         match get_tokens(&source) {
@@ -162,7 +162,7 @@ fn run(path: &str) {
 #[test]
 fn debug() -> Result<()> {
     let rego = "\"This string is 35 characters long.\"\"short string\"";
-    let source = Source::new("case.rego".to_string(), rego.to_string());
+    let source = Source::from_contents("case.rego".to_string(), rego.to_string())?;
 
     let mut lexer = Lexer::new(&source);
     let tok = lexer.next_token()?;
@@ -184,7 +184,7 @@ fn debug() -> Result<()> {
 #[test]
 fn tab() -> Result<()> {
     let rego = r#"	"This string is 35 characters long."`raw	string`p"#;
-    let source = Source::new("case.rego".to_string(), rego.to_string());
+    let source = Source::from_contents("case.rego".to_string(), rego.to_string())?;
 
     let mut lexer = Lexer::new(&source);
 
@@ -214,7 +214,7 @@ fn tab() -> Result<()> {
 #[test]
 fn invalid_line() -> Result<()> {
     let rego = "";
-    let source = Source::new("case.rego".to_string(), rego.to_string());
+    let source = Source::from_contents("case.rego".to_string(), rego.to_string())?;
 
     assert_eq!(
         source.message(2, 0, "", ""),
@@ -223,3 +223,21 @@ fn invalid_line() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn file_more_than_64_kb_size() -> Result<()> {
+    let source = Source::from_file("tests/kata/data/large.rego")?;
+    let mut lexer = Lexer::new(&source);
+
+    let mut count = 0;
+    // Read tokens until EOF.
+    loop {
+        let token = lexer.next_token()?;
+        count += 1;
+        if token.0 == TokenKind::Eof {
+            break;
+        }
+    }
+    assert_eq!(count, 8789);
+    Ok(())
+}

tests/opa.rs

@@ -13,7 +13,7 @@ use serde::{Deserialize, Serialize};
 use walkdir::WalkDir;
 
 const OPA_REPO: &str = "https://github.com/open-policy-agent/opa";
-const OPA_BRANCH: &str = "v0.63.0";
+const OPA_BRANCH: &str = "v0.64.0";
 
 #[derive(Serialize, Deserialize, PartialEq, Debug)]
 #[serde(deny_unknown_fields)]

tests/parser/mod.rs

@@ -630,7 +630,7 @@ fn yaml_test_impl(file: &str) -> Result<()> {
 
     for case in &test.cases {
         print!("\ncase {} ", case.note);
-        let source = Source::new("case.rego".to_string(), case.rego.clone());
+        let source = Source::from_contents("case.rego".to_string(), case.rego.clone())?;
         let mut parser = Parser::new(&source)?;
         match parser.parse() {
             Ok(module) => {