special-case ExternalStream in bytes starts-with (#8203)

# Description
`bytes starts-with` converts the input into a `Value` before running
.starts_with to find if the binary matches. This has two side effects:
it makes the code simpler, only dealing in whole values, and simplifying
a lot of input pipeline handling and value transforming it would
otherwise have to do. _Especially_ in the presence of a cell path to
drill into. It also makes buffers the entire input into memory, which
can take up a lot of memory when dealing with large files, especially if
you only want to check the first few bytes (like for a magic number).

This PR adds a special branch on PipelineData::ExternalStream with a
streaming version of starts_with.

# User-Facing Changes
Opening large files and running bytes starts-with on them will not take
a long time.

# Tests + Formatting

Don't forget to add tests that cover your changes.

Make sure you've run and fixed any issues with these commands:

- `cargo fmt --all -- --check` to check standard code formatting (`cargo
fmt --all` applies these changes)
- `cargo clippy --workspace -- -D warnings -D clippy::unwrap_used -A
clippy::needless_collect` to check that you're using the standard code
style
- `cargo test --workspace` to check that all tests pass

# Drawbacks
Streaming checking is more complicated, and there may be bugs. I tested
it with multiple chunks with string data and binary data and it seems to
work alright up to 8k and over bytes, though.

The existing `operate` method still exists because the way it handles
cell paths and values is complicated. This causes some "code
duplication", or at least some intent duplication, between the value
code and the streaming code. This might be worthwhile considering the
performance gains (approaching infinity on larger inputs).

Another thing to consider is that my ExternalStream branch considers
string data as valid input. The operate branch only parses Binary
values, so it would fail. `open` is kind of unpredictable on whether it
returns string data or binary data, even when passing `--raw`. I think
this can be a problem but not really one I'm trying to tackle in this
PR, so, it's worth considering.
This commit is contained in:
Doru 2023-02-26 11:17:44 -03:00 committed by GitHub
parent 9bbb9711e4
commit c602b5a1e8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 240 additions and 7 deletions

View File

@ -4,6 +4,7 @@ use nu_protocol::ast::Call;
use nu_protocol::ast::CellPath;
use nu_protocol::engine::{Command, EngineState, Stack};
use nu_protocol::Category;
use nu_protocol::IntoPipelineData;
use nu_protocol::{Example, PipelineData, ShellError, Signature, Span, SyntaxShape, Type, Value};
struct Arguments {
@ -60,13 +61,63 @@ impl Command for BytesStartsWith {
pattern,
cell_paths,
};
operate(
starts_with,
arg,
input,
call.head,
engine_state.ctrlc.clone(),
)
match input {
PipelineData::ExternalStream {
stdout: Some(stream),
span,
..
} => {
let mut i = 0;
for item in stream {
let byte_slice = match &item {
// String and binary data are valid byte patterns
Ok(Value::String { val, .. }) => val.as_bytes(),
Ok(Value::Binary { val, .. }) => val,
// If any Error value is output, echo it back
Ok(v @ Value::Error { .. }) => return Ok(v.clone().into_pipeline_data()),
// Unsupported data
Ok(other) => {
return Ok(Value::Error {
error: ShellError::OnlySupportsThisInputType(
"string and binary".into(),
other.get_type().to_string(),
span,
// This line requires the Value::Error match above.
other.expect_span(),
),
}
.into_pipeline_data());
}
Err(err) => return Err(err.to_owned()),
};
let max = byte_slice.len().min(arg.pattern.len() - i);
if byte_slice[..max] == arg.pattern[i..i + max] {
i += max;
if i >= arg.pattern.len() {
return Ok(Value::boolean(true, span).into_pipeline_data());
}
} else {
return Ok(Value::boolean(false, span).into_pipeline_data());
}
}
// We reached the end of the stream and never returned,
// the pattern wasn't exhausted so it probably doesn't match
Ok(Value::boolean(false, span).into_pipeline_data())
}
_ => operate(
starts_with,
arg,
input,
call.head,
engine_state.ctrlc.clone(),
),
}
}
fn examples(&self) -> Vec<Example> {

View File

@ -0,0 +1 @@
mod starts_with;

View File

@ -0,0 +1,153 @@
use nu_test_support::nu;
#[test]
fn basic_binary_starts_with() {
let actual = nu!(
cwd: ".",
r#"
"hello world" | into binary | bytes starts-with 0x[68 65 6c 6c 6f]
"#
);
assert_eq!(actual.out, "true");
}
#[test]
fn basic_string_fails() {
let actual = nu!(
cwd: ".",
r#"
"hello world" | bytes starts-with 0x[68 65 6c 6c 6f]
"#
);
assert!(actual.err.contains("Input type not supported"));
assert_eq!(actual.out, "");
}
#[test]
fn short_stream_binary() {
let actual = nu!(
cwd: ".",
r#"
nu --testbin repeater (0x[01]) 5 | bytes starts-with 0x[010101]
"#
);
assert_eq!(actual.out, "true");
}
#[test]
fn short_stream_mismatch() {
let actual = nu!(
cwd: ".",
r#"
nu --testbin repeater (0x[010203]) 5 | bytes starts-with 0x[010204]
"#
);
assert_eq!(actual.out, "false");
}
#[test]
fn short_stream_binary_overflow() {
let actual = nu!(
cwd: ".",
r#"
nu --testbin repeater (0x[01]) 5 | bytes starts-with 0x[010101010101]
"#
);
assert_eq!(actual.out, "false");
}
#[test]
fn long_stream_binary() {
let actual = nu!(
cwd: ".",
r#"
nu --testbin repeater (0x[01]) 32768 | bytes starts-with 0x[010101]
"#
);
assert_eq!(actual.out, "true");
}
#[test]
fn long_stream_binary_overflow() {
// .. ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow
let actual = nu!(
cwd: ".",
r#"
nu --testbin repeater (0x[01]) 32768 | bytes starts-with (0..32768 | each { 0x[01] } | bytes collect)
"#
);
assert_eq!(actual.out, "false");
}
#[test]
fn long_stream_binary_exact() {
// ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow
let actual = nu!(
cwd: ".",
r#"
nu --testbin repeater (0x[01020304]) 8192 | bytes starts-with (0..<8192 | each { 0x[01020304] } | bytes collect)
"#
);
assert_eq!(actual.out, "true");
}
#[test]
fn long_stream_string_exact() {
// ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow
let actual = nu!(
cwd: ".",
r#"
nu --testbin repeater hell 8192 | bytes starts-with (0..<8192 | each { "hell" | into binary } | bytes collect)
"#
);
assert_eq!(actual.out, "true");
}
#[test]
fn long_stream_mixed_exact() {
// ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow
let actual = nu!(
cwd: ".",
r#"
let binseg = (0..<2048 | each { 0x[003d9fbf] } | bytes collect)
let strseg = (0..<2048 | each { "hell" | into binary } | bytes collect)
nu --testbin repeat_bytes 003d9fbf 2048 68656c6c 2048 | bytes starts-with (bytes build $binseg $strseg)
"#
);
assert_eq!(
actual.err, "",
"invocation failed. command line limit likely reached"
);
assert_eq!(actual.out, "true");
}
#[test]
fn long_stream_mixed_overflow() {
// ranges are inclusive..inclusive, so we don't need to +1 to check for an overflow
let actual = nu!(
cwd: ".",
r#"
let binseg = (0..<2048 | each { 0x[003d9fbf] } | bytes collect)
let strseg = (0..<2048 | each { "hell" | into binary } | bytes collect)
nu --testbin repeat_bytes 003d9fbf 2048 68656c6c 2048 | bytes starts-with (bytes build $binseg $strseg 0x[01])
"#
);
assert_eq!(
actual.err, "",
"invocation failed. command line limit likely reached"
);
assert_eq!(actual.out, "false");
}

View File

@ -4,6 +4,7 @@ mod any;
mod append;
mod assignment;
mod break_;
mod bytes;
mod cal;
mod cd;
mod compact;

View File

@ -168,6 +168,7 @@ fn main() -> Result<()> {
"nonu" => test_bins::nonu(),
"chop" => test_bins::chop(),
"repeater" => test_bins::repeater(),
"repeat_bytes" => test_bins::repeat_bytes(),
"nu_repl" => test_bins::nu_repl(),
"input_bytes_length" => test_bins::input_bytes_length(),
_ => std::process::exit(1),

View File

@ -84,6 +84,32 @@ pub fn repeater() {
let _ = stdout.flush();
}
// A version of repeater that can output binary data, even null bytes
pub fn repeat_bytes() {
let mut stdout = io::stdout();
let args = args();
let mut args = args.iter().skip(1);
while let (Some(binary), Some(count)) = (args.next(), args.next()) {
let bytes: Vec<u8> = (0..binary.len())
.step_by(2)
.map(|i| {
u8::from_str_radix(&binary[i..i + 2], 16)
.expect("binary string is valid hexadecimal")
})
.collect();
let count: u64 = count.parse().expect("repeat count must be a number");
for _ in 0..count {
stdout
.write_all(&bytes)
.expect("writing to stdout must not fail");
}
}
let _ = stdout.flush();
}
pub fn iecho() {
// println! panics if stdout gets closed, whereas writeln gives us an error
let mut stdout = io::stdout();