diff --git a/Cargo.lock b/Cargo.lock index 2905803..5ce4fce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -160,6 +160,7 @@ dependencies = [ "cel", "clap", "json5", + "memchr", "mimalloc", "rayon", "serde", diff --git a/Cargo.toml b/Cargo.toml index 03b3a3d..f669df3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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 } diff --git a/src/input_handler.rs b/src/input_handler.rs index a50a931..5111ed6 100644 --- a/src/input_handler.rs +++ b/src/input_handler.rs @@ -60,8 +60,10 @@ fn handle_buffer( program: &Program, arg_variables: &BTreeMap, input_params: &InputParameters, - reader: BufReader, + mut reader: BufReader, ) -> Result> { + 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 @@ -78,24 +80,22 @@ fn handle_buffer( input_params.parallelism as usize }; - // Collect all non-empty lines first - let lines: Vec = reader - .lines() - .collect::>>() - .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) => { @@ -108,7 +108,7 @@ fn handle_buffer( // 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 @@ -119,8 +119,8 @@ fn handle_buffer( .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() }) @@ -132,8 +132,7 @@ fn handle_buffer( } 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]) } } @@ -152,28 +151,40 @@ fn handle_buffer( } // 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, input_params: &InputParameters, - json_str: Option<&str>, -) -> Result<(String, bool)> { +) -> Result> { // Create context with default values let mut context = Context::default(); @@ -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( diff --git a/src/input_handler_test.rs b/src/input_handler_test.rs index baa81d6..1903ab9 100644 --- a/src/input_handler_test.rs +++ b/src/input_handler_test.rs @@ -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, ¶ms, None).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ) + .unwrap(); assert!(output.contains("5")); assert!(is_truthy); @@ -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, ¶ms, Some(json)).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + Some(json), + ) + .unwrap(); assert!(output.contains("30")); assert!(is_truthy); @@ -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, ¶ms, None).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ) + .unwrap(); assert!(output.contains("12")); assert!(is_truthy); @@ -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, ¶ms, Some(json)).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + Some(json), + ) + .unwrap(); assert!(output.contains("50")); assert!(is_truthy); @@ -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, ¶ms, Some(json)).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + Some(json), + ) + .unwrap(); assert!(output.contains("150")); assert!(is_truthy); @@ -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, ¶ms, None).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ) + .unwrap(); assert!(output.contains("false")); assert!(!is_truthy); @@ -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, ¶ms, None).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ) + .unwrap(); assert!(output.contains("true")); assert!(is_truthy); @@ -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, ¶ms, None).unwrap(); + let (_output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ) + .unwrap(); assert!(!is_truthy); } @@ -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, ¶ms, None).unwrap(); + let (_output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ) + .unwrap(); assert!(!is_truthy); } @@ -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, ¶ms, None).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ) + .unwrap(); assert_eq!(output, "hello"); assert!(is_truthy); @@ -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, ¶ms, None).unwrap(); + let (output, is_truthy) = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ) + .unwrap(); assert_eq!( output, @@ -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, ¶ms, Some(json)); + let result = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + Some(json), + ); assert!(result.is_err()); } @@ -235,7 +306,12 @@ fn test_handle_json_missing_variable() { let args = BTreeMap::new(); let params = default_params(); - let result = handle_json(&program, &args, ¶ms, None); + let result = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + None, + ); assert!(result.is_err()); } @@ -385,9 +461,14 @@ fn test_handle_json_greppable_disabled() { let mut params = default_params(); params.greppable = true; - let err = handle_json(&program, &args, ¶ms, Some(r#"{"x": 42}"#)) - .unwrap_err() - .to_string(); + let err = handle_json( + &program, + &root_context(&args, ¶ms).unwrap(), + ¶ms, + Some(r#"{"x": 42}"#), + ) + .unwrap_err() + .to_string(); assert!(err.contains("Binary was compiled without greppable support")); }