Accept --piece-length arguments with SI units

Valid units include MiB, KiB, and GiB.

type: changed

src/bytes.rs

@@ -0,0 +1,162 @@
+use crate::common::*;
+
+const KI: u128 = 1 << 10;
+const MI: u128 = KI << 10;
+const GI: u128 = MI << 10;
+const TI: u128 = GI << 10;
+const PI: u128 = TI << 10;
+const EI: u128 = PI << 10;
+const ZI: u128 = EI << 10;
+const YI: u128 = ZI << 10;
+
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub(crate) struct Bytes(pub(crate) u128);
+
+fn float_to_int(x: f64) -> u128 {
+  #![allow(
+    clippy::as_conversions,
+    clippy::cast_sign_loss,
+    clippy::cast_possible_truncation
+  )]
+  x as u128
+}
+
+fn int_to_float(x: u128) -> f64 {
+  #![allow(clippy::as_conversions, clippy::cast_precision_loss)]
+  x as f64
+}
+
+impl FromStr for Bytes {
+  type Err = Error;
+
+  fn from_str(text: &str) -> Result<Self, Self::Err> {
+    #[allow(clippy::trivially_copy_pass_by_ref)]
+    fn is_digit(c: &char) -> bool {
+      match c {
+        '0'..='9' | '.' => true,
+        _ => false,
+      }
+    }
+
+    let digits = text.chars().take_while(is_digit).collect::<String>();
+
+    let suffix = text.chars().skip_while(is_digit).collect::<String>();
+
+    let value = digits.parse::<f64>().map_err(|source| Error::ByteParse {
+      text: text.to_owned(),
+      source,
+    })?;
+
+    let multiple = match suffix.to_lowercase().as_str() {
+      "" | "b" | "byte" | "bytes" => 1,
+      "kib" => KI,
+      "mib" => MI,
+      "gib" => GI,
+      "tib" => TI,
+      "pib" => PI,
+      "eib" => EI,
+      "zib" => ZI,
+      "yib" => YI,
+      _ => {
+        return Err(Error::ByteSuffix {
+          text: text.to_owned(),
+          suffix: suffix.to_owned(),
+        })
+      }
+    };
+
+    Ok(Bytes(float_to_int(value * int_to_float(multiple))))
+  }
+}
+
+impl Display for Bytes {
+  fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+    const DISPLAY_SUFFIXES: &[&str] = &["KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB"];
+
+    let mut value = int_to_float(self.0);
+
+    let mut i = 0;
+
+    while value >= 1024.0 {
+      value /= 1024.0;
+      i += 1;
+    }
+
+    let suffix = if i == 0 {
+      if value == 1.0 {
+        "byte"
+      } else {
+        "bytes"
+      }
+    } else {
+      DISPLAY_SUFFIXES[i - 1]
+    };
+
+    let formatted = format!("{:.2}", value);
+    let trimmed = formatted.trim_end_matches('0').trim_end_matches('.');
+    write!(f, "{} {}", trimmed, suffix)
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn ok() {
+    const CASES: &[(&str, u128)] = &[
+      ("0", 0),
+      ("0kib", 0),
+      ("1", 1),
+      ("1b", 1),
+      ("1byte", 1),
+      ("1bytes", 1),
+      ("1kib", KI),
+      ("1KiB", KI),
+      ("12kib", 12 * KI),
+      ("1.5mib", 1 * MI + 512 * KI),
+      ("1yib", 1 * YI),
+    ];
+
+    for (text, value) in CASES {
+      assert_eq!(
+        text.parse::<Bytes>().unwrap(),
+        Bytes(*value),
+        "text: {}",
+        text
+      );
+    }
+  }
+
+  #[test]
+  fn err() {
+    assert_matches!(
+      "100foo".parse::<Bytes>().unwrap_err(),
+      Error::ByteSuffix { text, suffix }
+      if text == "100foo" && suffix == "foo"
+    );
+
+    assert_matches!(
+      "1.0.0foo".parse::<Bytes>().unwrap_err(),
+      Error::ByteParse { .. }
+    );
+  }
+
+  #[test]
+  fn display() {
+    assert_eq!(Bytes(0).to_string(), "0 bytes");
+    assert_eq!(Bytes(1).to_string(), "1 byte");
+    assert_eq!(Bytes(2).to_string(), "2 bytes");
+    assert_eq!(Bytes(KI).to_string(), "1 KiB");
+    assert_eq!(Bytes(512 * KI).to_string(), "512 KiB");
+    assert_eq!(Bytes(MI).to_string(), "1 MiB");
+    assert_eq!(Bytes(MI + 512 * KI).to_string(), "1.5 MiB");
+    assert_eq!(Bytes(1024 * MI + 512 * MI).to_string(), "1.5 GiB");
+    assert_eq!(Bytes(GI).to_string(), "1 GiB");
+    assert_eq!(Bytes(TI).to_string(), "1 TiB");
+    assert_eq!(Bytes(PI).to_string(), "1 PiB");
+    assert_eq!(Bytes(EI).to_string(), "1 EiB");
+    assert_eq!(Bytes(ZI).to_string(), "1 ZiB");
+    assert_eq!(Bytes(YI).to_string(), "1 YiB");
+  }
+}

src/common.rs

@@ -10,6 +10,7 @@ pub(crate) use std::{
   fs::{self, File},
   hash::Hash,
   io::{self, Read, Write},
+  num::ParseFloatError,
   path::{Path, PathBuf},
   process::{self, Command, ExitStatus},
   str::{self, FromStr},
@@ -42,9 +43,9 @@ pub(crate) use crate::{
 
 // structs and enums
 pub(crate) use crate::{
-  env::Env, error::Error, file_info::FileInfo, hasher::Hasher, info::Info, metainfo::Metainfo,
+  bytes::Bytes, env::Env, error::Error, file_info::FileInfo, hasher::Hasher, info::Info,
-  mode::Mode, opt::Opt, platform::Platform, style::Style, subcommand::Subcommand, torrent::Torrent,
+  metainfo::Metainfo, mode::Mode, opt::Opt, platform::Platform, style::Style,
-  use_color::UseColor,
+  subcommand::Subcommand, torrent::Torrent, use_color::UseColor,
 };
 
 // test modules

src/error.rs

@@ -11,6 +11,13 @@ pub(crate) enum Error {
   AnnounceUrlParse { source: url::ParseError },
   #[snafu(display("Failed to decode bencode: {}", source))]
   BencodeDecode { source: serde_bencode::Error },
+  #[snafu(display("Failed to parse byte count `{}`: {}", text, source))]
+  ByteParse {
+    text: String,
+    source: ParseFloatError,
+  },
+  #[snafu(display("Failed to parse byte count `{}`, invalid suffix: `{}`", text, suffix))]
+  ByteSuffix { text: String, suffix: String },
   #[snafu(display("{}", source))]
   Clap { source: clap::Error },
   #[snafu(display("Failed to invoke command `{}`: {}", command, source,))]
@@ -25,6 +32,12 @@ pub(crate) enum Error {
   Filesystem { source: io::Error, path: PathBuf },
   #[snafu(display("Failed to find opener utility, please install one of {}", tried.join(",")))]
   OpenerMissing { tried: &'static [&'static str] },
+  #[snafu(display(
+    "Piece length `{}` too large. The maximum supported piece length is {}.",
+    bytes,
+    Bytes(u32::max_value().into())
+  ))]
+  PieceLength { bytes: Bytes },
   #[snafu(display("Serialization failed: {}", source))]
   Serialize { source: serde_bencode::Error },
   #[snafu(display("Failed to write to standard error: {}", source))]

src/main.rs

@@ -2,6 +2,8 @@
 #![allow(
   clippy::else_if_without_else,
   clippy::enum_glob_use,
+  clippy::float_arithmetic,
+  clippy::float_cmp,
   clippy::implicit_return,
   clippy::indexing_slicing,
   clippy::integer_arithmetic,
@@ -40,6 +42,7 @@ mod test_env;
 mod capture;
 
 mod bencode;
+mod bytes;
 mod common;
 mod consts;
 mod env;

src/torrent/create.rs

@@ -84,10 +84,11 @@ Note: Many BitTorrent clients do not implement the behavior described in BEP 12.
   #[structopt(
     name = "PIECE-LENGTH",
     long = "piece-length",
-    default_value = "524288",
+    default_value = "512KiB",
-    help = "Set piece length to `PIECE-LENGTH` bytes."
+    help = "Set piece length to `PIECE-LENGTH` bytes.",
+    long_help = "Set piece length to `PIECE-LENGTH` bytes. Accepts SI units, e.g. kib, mib, and gib."
   )]
-  piece_length: u32,
+  piece_length: Bytes,
   #[structopt(
     name = "PRIVATE",
     long = "private",
@@ -99,6 +100,14 @@ Note: Many BitTorrent clients do not implement the behavior described in BEP 12.
 
 impl Create {
   pub(crate) fn run(self, env: &Env) -> Result<(), Error> {
+    let piece_length: u32 = self
+      .piece_length
+      .0
+      .try_into()
+      .map_err(|_| Error::PieceLength {
+        bytes: self.piece_length,
+      })?;
+
     let input = env.resolve(&self.input);
 
     let mut announce_list = Vec::new();
@@ -157,10 +166,10 @@ impl Create {
       Some(String::from(consts::CREATED_BY_DEFAULT))
     };
 
-    let (mode, pieces) = Hasher::hash(&input, self.md5sum, self.piece_length)?;
+    let (mode, pieces) = Hasher::hash(&input, self.md5sum, piece_length)?;
 
     let info = Info {
-      piece_length: self.piece_length,
+      piece_length,
       mode,
       pieces,
       name,
@@ -406,6 +415,24 @@ mod tests {
     assert_eq!(metainfo.info.piece_length, 1);
   }
 
+  #[test]
+  fn si_piece_size() {
+    let mut env = environment(&[
+      "--input",
+      "foo",
+      "--announce",
+      "http://bar",
+      "--piece-length",
+      "0.5MiB",
+    ]);
+    fs::write(env.resolve("foo"), "").unwrap();
+    env.run().unwrap();
+    let torrent = env.resolve("foo.torrent");
+    let bytes = fs::read(torrent).unwrap();
+    let metainfo = serde_bencode::de::from_bytes::<Metainfo>(&bytes).unwrap();
+    assert_eq!(metainfo.info.piece_length, 512 * 1024);
+  }
+
   #[test]
   fn name() {
     let mut env = environment(&[