diff --git a/crates/nu-command/src/random/binary.rs b/crates/nu-command/src/random/binary.rs index 47e6429eba..663affbc12 100644 --- a/crates/nu-command/src/random/binary.rs +++ b/crates/nu-command/src/random/binary.rs @@ -1,5 +1,5 @@ +use super::byte_stream::{random_byte_stream, RandomDistribution}; use nu_engine::command_prelude::*; -use rand::{thread_rng, RngCore}; #[derive(Clone)] pub struct RandomBinary; @@ -57,12 +57,12 @@ impl Command for RandomBinary { }), }?; - let mut rng = thread_rng(); - - let mut out = vec![0u8; length]; - rng.fill_bytes(&mut out); - - Ok(Value::binary(out, call.head).into_pipeline_data()) + Ok(random_byte_stream( + RandomDistribution::Binary, + length, + call.head, + engine_state.signals().clone(), + )) } fn examples(&self) -> Vec { diff --git a/crates/nu-command/src/random/byte_stream.rs b/crates/nu-command/src/random/byte_stream.rs new file mode 100644 index 0000000000..baff8fbc6c --- /dev/null +++ b/crates/nu-command/src/random/byte_stream.rs @@ -0,0 +1,49 @@ +use nu_engine::command_prelude::*; +use nu_protocol::Signals; +use rand::{ + distributions::{Alphanumeric, Standard}, + thread_rng, Rng, +}; + +pub(super) enum RandomDistribution { + Binary, + Alphanumeric, +} + +pub(super) fn random_byte_stream( + distribution: RandomDistribution, + length: usize, + span: Span, + signals: Signals, +) -> PipelineData { + let stream_type = match distribution { + RandomDistribution::Binary => ByteStreamType::Binary, + RandomDistribution::Alphanumeric => ByteStreamType::String, + }; + + const OUTPUT_CHUNK_SIZE: usize = 8192; + let mut remaining_bytes = length; + PipelineData::ByteStream( + ByteStream::from_fn(span, signals.clone(), stream_type, move |out| { + if remaining_bytes == 0 || signals.interrupted() { + return Ok(false); + } + + let bytes_to_write = std::cmp::min(remaining_bytes, OUTPUT_CHUNK_SIZE); + + let rng = thread_rng(); + let byte_iter: Box> = match distribution { + RandomDistribution::Binary => Box::new(rng.sample_iter(Standard)), + RandomDistribution::Alphanumeric => Box::new(rng.sample_iter(Alphanumeric)), + }; + out.extend(byte_iter.take(bytes_to_write)); + + remaining_bytes -= bytes_to_write; + + Ok(true) + }) + .with_known_size(Some(length as u64)), + None, + ) + .with_span(span) +} diff --git a/crates/nu-command/src/random/chars.rs b/crates/nu-command/src/random/chars.rs index 5ef76603b7..ad9d8ac2ac 100644 --- a/crates/nu-command/src/random/chars.rs +++ b/crates/nu-command/src/random/chars.rs @@ -1,8 +1,5 @@ +use super::byte_stream::{random_byte_stream, RandomDistribution}; use nu_engine::command_prelude::*; -use rand::{ - distributions::{Alphanumeric, Distribution}, - thread_rng, -}; const DEFAULT_CHARS_LENGTH: usize = 25; @@ -71,7 +68,6 @@ fn chars( stack: &mut Stack, call: &Call, ) -> Result { - let span = call.head; let length: Option = call.get_flag(engine_state, stack, "length")?; let length = if let Some(length_val) = length { match length_val { @@ -97,17 +93,11 @@ fn chars( DEFAULT_CHARS_LENGTH }; - let mut rng = thread_rng(); - - let random_string = Alphanumeric - .sample_iter(&mut rng) - .take(length) - .map(char::from) - .collect::(); - - Ok(PipelineData::Value( - Value::string(random_string, span), - None, + Ok(random_byte_stream( + RandomDistribution::Alphanumeric, + length, + call.head, + engine_state.signals().clone(), )) } diff --git a/crates/nu-command/src/random/mod.rs b/crates/nu-command/src/random/mod.rs index 43c8af862f..c4b635449b 100644 --- a/crates/nu-command/src/random/mod.rs +++ b/crates/nu-command/src/random/mod.rs @@ -1,5 +1,6 @@ mod binary; mod bool; +mod byte_stream; mod chars; mod dice; mod float;