add option to limit rng output (#224)

camshaft · Jun 8, 2024 · e45edd3 · e45edd3
1 parent 8b7539d
commit e45edd3
Show file tree

Hide file tree

Showing 4 changed files with 185 additions and 128 deletions.
diff --git a/lib/bolero-generator/src/driver.rs b/lib/bolero-generator/src/driver.rs
@@ -178,10 +178,12 @@ pub struct Options {
     driver_mode: Option<DriverMode>,
     shrink_time: Option<core::time::Duration>,
     max_depth: Option<usize>,
+    max_len: Option<usize>,
 }
 
 impl Options {
-    pub const DEFAULT_MAX_DEPTH: usize = 32;
+    pub const DEFAULT_MAX_DEPTH: usize = 5;
+    pub const DEFAULT_MAX_LEN: usize = 4096;
     pub const DEFAULT_SHRINK_TIME: core::time::Duration = core::time::Duration::from_secs(1);
 
     pub fn with_driver_mode(mut self, driver_mode: DriverMode) -> Self {
@@ -199,6 +201,11 @@ impl Options {
         self
     }
 
+    pub fn with_max_len(mut self, max_len: usize) -> Self {
+        self.max_len = Some(max_len);
+        self
+    }
+
     pub fn set_driver_mode(&mut self, driver_mode: DriverMode) -> &mut Self {
         self.driver_mode = Some(driver_mode);
         self
@@ -214,11 +221,21 @@ impl Options {
         self
     }
 
+    pub fn set_max_len(&mut self, max_len: usize) -> &mut Self {
+        self.max_len = Some(max_len);
+        self
+    }
+
     #[inline]
     pub fn max_depth(&self) -> Option<usize> {
         self.max_depth
     }
 
+    #[inline]
+    pub fn max_len(&self) -> Option<usize> {
+        self.max_len
+    }
+
     #[inline]
     pub fn shrink_time(&self) -> Option<core::time::Duration> {
         self.shrink_time
@@ -234,6 +251,11 @@ impl Options {
         self.max_depth.unwrap_or(Self::DEFAULT_MAX_DEPTH)
     }
 
+    #[inline]
+    pub fn max_len_or_default(&self) -> usize {
+        self.max_len.unwrap_or(Self::DEFAULT_MAX_LEN)
+    }
+
     #[inline]
     pub fn shrink_time_or_default(&self) -> core::time::Duration {
         self.shrink_time.unwrap_or(Self::DEFAULT_SHRINK_TIME)
@@ -251,6 +273,7 @@ impl Options {
 
         merge!(driver_mode);
         merge!(max_depth);
+        merge!(max_len);
         merge!(shrink_time);
     }
 }
diff --git a/lib/bolero-generator/src/driver/bytes.rs b/lib/bolero-generator/src/driver/bytes.rs
@@ -12,6 +12,8 @@ impl<'a> ByteSliceDriver<'a> {
     pub fn new(input: &'a [u8], options: &Options) -> Self {
         let mode = options.driver_mode.unwrap_or(DriverMode::Direct);
         let max_depth = options.max_depth_or_default();
+        let len = options.max_len_or_default().min(input.len());
+        let input = &input[..len];
 
         Self {
             input,