add b3sum --seek

2 files changed, 58 insertions, 21 deletions

diff --git a/b3sum/src/main.rs b/b3sum/src/main.rs
index 2b61a67..fd35f68 100644
--- a/b3sum/src/main.rs
+++ b/b3sum/src/main.rs
@@ -26,6 +26,16 @@ struct Inner {
     /// When no file is given, or when - is given, read standard input.
     file: Vec<PathBuf>,
 
+    /// Use the keyed mode, reading the 32-byte key from stdin
+    #[arg(long, requires("file"))]
+    keyed: bool,
+
+    /// Use the key derivation mode, with the given context string
+    ///
+    /// Cannot be used with --keyed.
+    #[arg(long, value_name("CONTEXT"), conflicts_with(KEYED_ARG))]
+    derive_key: Option<String>,
+
     /// The number of output bytes, before hex encoding
     #[arg(
         short,
@@ -35,6 +45,10 @@ struct Inner {
     )]
     length: u64,
 
+    /// The starting output byte offset, before hex encoding
+    #[arg(long, default_value_t = 0, value_name("SEEK"))]
+    seek: u64,
+
     /// The maximum number of threads to use
     ///
     /// By default, this is the number of logical cores. If this flag is
@@ -42,16 +56,6 @@ struct Inner {
     #[arg(long, value_name("NUM"))]
     num_threads: Option<usize>,
 
-    /// Use the keyed mode, reading the 32-byte key from stdin
-    #[arg(long, requires("file"))]
-    keyed: bool,
-
-    /// Use the key derivation mode, with the given context string
-    ///
-    /// Cannot be used with --keyed.
-    #[arg(long, value_name("CONTEXT"), conflicts_with(KEYED_ARG))]
-    derive_key: Option<String>,
-
     /// Disable memory mapping
     ///
     /// Currently this also disables multithreading.
@@ -80,7 +84,7 @@ struct Inner {
     )]
     check: bool,
 
-    /// Skip printing OK for each successfully verified file
+    /// Skip printing OK for each checked file
     ///
     /// Must be used with --check.
     #[arg(long, requires(CHECK_ARG))]
@@ -146,6 +150,10 @@ impl Args {
         self.inner.length
     }
 
+    fn seek(&self) -> u64 {
+        self.inner.seek
+    }
+
     fn keyed(&self) -> bool {
         self.inner.keyed
     }
@@ -208,7 +216,9 @@ impl Input {
                 copy_wide(lock, &mut hasher)?;
             }
         }
-        Ok(hasher.finalize_xof())
+        let mut output_reader = hasher.finalize_xof();
+        output_reader.set_position(args.seek());
+        Ok(output_reader)
     }
 }
 
@@ -275,7 +285,9 @@ fn maybe_memmap_file(file: &File) -> Result<Option<memmap2::Mmap>> {
 }
 
 fn write_hex_output(mut output: blake3::OutputReader, args: &Args) -> Result<()> {
-    // Encoding multiples of the block size is most efficient.
+    // Encoding multiples of the 64 bytes is most efficient.
+    // TODO: This computes each output block twice when the --seek argument isn't a multiple of 64.
+    // We'll refactor all of this soon anyway, once SIMD optimizations are available for the XOF.
     let mut len = args.len();
     let mut block = [0; blake3::guts::BLOCK_LEN];
     while len > 0 {
diff --git a/b3sum/tests/cli_tests.rs b/b3sum/tests/cli_tests.rs
index 62602ec..d5d4efa 100644
--- a/b3sum/tests/cli_tests.rs
+++ b/b3sum/tests/cli_tests.rs
@@ -87,18 +87,43 @@ fn test_missing_files() {
 }
 
 #[test]
-fn test_hash_length() {
-    let mut buf = [0; 100];
+fn test_hash_length_and_seek() {
+    let mut expected = [0; 100];
     blake3::Hasher::new()
         .update(b"foo")
         .finalize_xof()
-        .fill(&mut buf);
-    let expected = format!("{}  -", hex::encode(&buf[..]));
-    let output = cmd!(b3sum_exe(), "--length=100")
+        .fill(&mut expected);
+    let output = cmd!(b3sum_exe(), "--raw", "--length=100")
         .stdin_bytes("foo")
-        .read()
-        .unwrap();
-    assert_eq!(&*expected, &*output);
+        .stdout_capture()
+        .run()
+        .unwrap()
+        .stdout;
+    assert_eq!(expected[..], output);
+
+    let short_output = cmd!(b3sum_exe(), "--raw", "--length=99")
+        .stdin_bytes("foo")
+        .stdout_capture()
+        .run()
+        .unwrap()
+        .stdout;
+    assert_eq!(expected[..99], short_output);
+
+    let seek1_output = cmd!(b3sum_exe(), "--raw", "--length=99", "--seek=1")
+        .stdin_bytes("foo")
+        .stdout_capture()
+        .run()
+        .unwrap()
+        .stdout;
+    assert_eq!(expected[1..], seek1_output);
+
+    let seek99_output = cmd!(b3sum_exe(), "--raw", "--length=1", "--seek=99")
+        .stdin_bytes("foo")
+        .stdout_capture()
+        .run()
+        .unwrap()
+        .stdout;
+    assert_eq!(expected[99..], seek99_output);
 }
 
 #[test]