Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ clap = { version = ">= 4.5.0, < 5.0.0", default-features = false, features = ["d
serde = "1"
serde_json = { version = "1", features = ["preserve_order"] }
rayon = "1.12"
memchr = "2"
mimalloc = { version = ">=0.1.40, <0.2.0", default-features = false, optional = true }
json5 = "1.3"
toml = { version = "=1.1.6", default-features = false, features = ["parse", "serde", "preserve_order"], optional = true }
Expand Down
84 changes: 59 additions & 25 deletions src/input_handler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,10 @@ fn handle_buffer<R: Read>(
program: &Program,
arg_variables: &BTreeMap<String, CelValue>,
input_params: &InputParameters,
reader: BufReader<R>,
mut reader: BufReader<R>,
) -> Result<Vec<(String, bool)>> {
let root_context = root_context(arg_variables, input_params)?;

// Check if we're processing JSON or NDJSON that is not slurped
if input_params.input_format == InputFormat::Json {
// Determine thread pool size
Expand All @@ -78,24 +80,22 @@ fn handle_buffer<R: Read>(
input_params.parallelism as usize
};

// Collect all non-empty lines first
let lines: Vec<String> = reader
.lines()
.collect::<std::io::Result<Vec<_>>>()
.context("Failed to read lines from input")?
.into_iter()
.filter(|line| !line.trim().is_empty())
.collect();
// Read the whole input once, then split it into non-empty lines without copying
let mut buffer = String::new();
reader
.read_to_string(&mut buffer)
.context("Failed to read input")?;
let lines = non_empty_lines(&buffer);

// If no lines were processed, execute with no input
if lines.is_empty() {
let result = handle_json(program, arg_variables, input_params, None)?;
let result = handle_json(program, &root_context, input_params, None)?;
return Ok(vec![result]);
}

// Try to process the last line
let last_idx = lines.len() - 1;
let last_result = handle_json(program, arg_variables, input_params, Some(&lines[last_idx]));
let last_result = handle_json(program, &root_context, input_params, Some(lines[last_idx]));

match last_result {
Ok(last_output) => {
Expand All @@ -108,7 +108,7 @@ fn handle_buffer<R: Read>(
// Use regular iterator for single-threaded execution
lines[..last_idx]
.iter()
.map(|line| handle_json(program, arg_variables, input_params, Some(line)))
.map(|&line| handle_json(program, &root_context, input_params, Some(line)))
.collect()
} else {
// Use Rayon for parallel execution
Expand All @@ -119,8 +119,8 @@ fn handle_buffer<R: Read>(
.install(|| {
lines[..last_idx]
.par_iter()
.map(|line| {
handle_json(program, arg_variables, input_params, Some(line))
.map(|&line| {
handle_json(program, &root_context, input_params, Some(line))
})
.collect()
})
Expand All @@ -132,8 +132,7 @@ fn handle_buffer<R: Read>(
}
Err(_) => {
// Last line failed, try reading entire input as single JSON document
let full_buffer = lines.join("\n");
let result = handle_json(program, arg_variables, input_params, Some(&full_buffer))?;
let result = handle_json(program, &root_context, input_params, Some(&buffer))?;
Ok(vec![result])
}
}
Expand All @@ -152,28 +151,40 @@ fn handle_buffer<R: Read>(
}

// Process the entire buffer as one document
let result = handle_json(program, arg_variables, input_params, Some(&buffer))?;
let result = handle_json(program, &root_context, input_params, Some(&buffer))?;
Ok(vec![result])
}
}

/// Execute the CEL program with given JSON input and argument variables
/// Split input into its non-blank lines
fn non_empty_lines(input: &str) -> Vec<&str> {
let mut lines = Vec::new();
let mut start = 0;
// '\n' is ASCII, so every split point is a valid UTF-8 boundary
// we use memchr to hopefully get some SIMD boost.
for end in memchr::memchr_iter(b'\n', input.as_bytes()).chain(std::iter::once(input.len())) {
let line = &input[start..end];
if !line.trim().is_empty() {
lines.push(line);
}
start = end + 1;
}
lines
}

/// Build the context shared by every input: standard library, extensions and argument variables
///
/// # Arguments
/// * `program` - The compiled CEL program
/// * `arg_variables` - BTreeMap of variables from CLI arguments
/// * `input_params` - Input configuration parameters
/// * `json_str` - Optional JSON string to process
///
/// # Returns
/// * Ok((output_string, is_truthy)) - The output and whether it's truthy
/// * Ok(Context) - The root context, meant to be reused via `new_inner_scope`
/// * Err(anyhow::Error) - Any error that occurred
fn handle_json(
program: &Program,
fn root_context(
arg_variables: &BTreeMap<String, CelValue>,
input_params: &InputParameters,
json_str: Option<&str>,
) -> Result<(String, bool)> {
) -> Result<Context<'static>> {
// Create context with default values
let mut context = Context::default();

Expand All @@ -188,6 +199,29 @@ fn handle_json(
.with_context(|| format!("Failed to add variable '{}'", name))?;
}

Ok(context)
}

/// Execute the CEL program with given JSON input
///
/// # Arguments
/// * `program` - The compiled CEL program
/// * `root_context` - Shared context from `root_context`
/// * `input_params` - Input configuration parameters
/// * `json_str` - Optional JSON string to process
///
/// # Returns
/// * Ok((output_string, is_truthy)) - The output and whether it's truthy
/// * Err(anyhow::Error) - Any error that occurred
fn handle_json(
program: &Program,
root_context: &Context,
input_params: &InputParameters,
json_str: Option<&str>,
) -> Result<(String, bool)> {
// Input variables live in a child scope, so the root context is never copied
let mut context = root_context.new_inner_scope();

// If we have input, parse it and add to context
if let Some(json) = json_str {
let json_variables = json_to_cel_variables(
Expand Down
113 changes: 97 additions & 16 deletions src/input_handler_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,13 @@ fn test_handle_json_null_input() {
let args = BTreeMap::new();
let params = default_params();

let (output, is_truthy) = handle_json(&program, &args, &params, None).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
)
.unwrap();

assert!(output.contains("5"));
assert!(is_truthy);
Expand All @@ -37,7 +43,13 @@ fn test_handle_json_with_json() {
let json = r#"{"x": 10, "y": 20}"#;
let params = default_params();

let (output, is_truthy) = handle_json(&program, &args, &params, Some(json)).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
Some(json),
)
.unwrap();

assert!(output.contains("30"));
assert!(is_truthy);
Expand All @@ -51,7 +63,13 @@ fn test_handle_json_with_args() {
args.insert("y".to_string(), CelValue::Int(7));
let params = default_params();

let (output, is_truthy) = handle_json(&program, &args, &params, None).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
)
.unwrap();

assert!(output.contains("12"));
assert!(is_truthy);
Expand All @@ -65,7 +83,13 @@ fn test_handle_json_input_overrides_arg_with_same_name() {
let json = r#"{"value": 50}"#;
let params = default_params();

let (output, is_truthy) = handle_json(&program, &args, &params, Some(json)).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
Some(json),
)
.unwrap();

assert!(output.contains("50"));
assert!(is_truthy);
Expand All @@ -79,7 +103,13 @@ fn test_handle_json_args_and_json() {
let json = r#"{"value": 50}"#;
let params = default_params();

let (output, is_truthy) = handle_json(&program, &args, &params, Some(json)).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
Some(json),
)
.unwrap();

assert!(output.contains("150"));
assert!(is_truthy);
Expand All @@ -91,7 +121,13 @@ fn test_handle_json_boolean_false() {
let args = BTreeMap::new();
let params = default_params();

let (output, is_truthy) = handle_json(&program, &args, &params, None).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
)
.unwrap();

assert!(output.contains("false"));
assert!(!is_truthy);
Expand All @@ -103,7 +139,13 @@ fn test_handle_json_boolean_true() {
let args = BTreeMap::new();
let params = default_params();

let (output, is_truthy) = handle_json(&program, &args, &params, None).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
)
.unwrap();

assert!(output.contains("true"));
assert!(is_truthy);
Expand All @@ -115,7 +157,13 @@ fn test_handle_json_truthiness_zero() {
let args = BTreeMap::new();
let params = default_params();

let (_output, is_truthy) = handle_json(&program, &args, &params, None).unwrap();
let (_output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
)
.unwrap();

assert!(!is_truthy);
}
Expand All @@ -126,7 +174,13 @@ fn test_handle_json_truthiness_empty_string() {
let args = BTreeMap::new();
let params = default_params();

let (_output, is_truthy) = handle_json(&program, &args, &params, None).unwrap();
let (_output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
)
.unwrap();

assert!(!is_truthy);
}
Expand Down Expand Up @@ -177,7 +231,13 @@ fn test_handle_json_raw_output_for_string() {
let mut params = default_params();
params.raw_output = true;

let (output, is_truthy) = handle_json(&program, &args, &params, None).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
)
.unwrap();

assert_eq!(output, "hello");
assert!(is_truthy);
Expand All @@ -191,7 +251,13 @@ fn test_handle_json_sorted_pretty_output() {
params.sort_keys = true;
params.pretty_print = true;

let (output, is_truthy) = handle_json(&program, &args, &params, None).unwrap();
let (output, is_truthy) = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
)
.unwrap();

assert_eq!(
output,
Expand Down Expand Up @@ -224,7 +290,12 @@ fn test_handle_json_invalid_json() {
let json = r#"not valid json"#;
let params = default_params();

let result = handle_json(&program, &args, &params, Some(json));
let result = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
Some(json),
);

assert!(result.is_err());
}
Expand All @@ -235,7 +306,12 @@ fn test_handle_json_missing_variable() {
let args = BTreeMap::new();
let params = default_params();

let result = handle_json(&program, &args, &params, None);
let result = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
None,
);

assert!(result.is_err());
}
Expand Down Expand Up @@ -385,9 +461,14 @@ fn test_handle_json_greppable_disabled() {
let mut params = default_params();
params.greppable = true;

let err = handle_json(&program, &args, &params, Some(r#"{"x": 42}"#))
.unwrap_err()
.to_string();
let err = handle_json(
&program,
&root_context(&args, &params).unwrap(),
&params,
Some(r#"{"x": 42}"#),
)
.unwrap_err()
.to_string();

assert!(err.contains("Binary was compiled without greppable support"));
}
Loading