From 22bc5ad3dbc1e3a19450fdc6a6e409607b92a5d0 Mon Sep 17 00:00:00 2001 From: Maciej Hirsz <1096222+maciejhirsz@users.noreply.github.com> Date: Wed, 17 Oct 2018 16:52:54 +0200 Subject: [PATCH] Perf and Rust idiomatic stuff (#1) --- Cargo.toml | 1 - benches/validate.rs | 24 +++++++ src/crypto.rs | 19 +++-- src/language.rs | 92 +++++++++++------------- src/lib.rs | 1 - src/mnemonic.rs | 165 +++++++++++++++++++++++-------------------- src/mnemonic_type.rs | 58 +++++---------- src/seed.rs | 12 ++-- src/util.rs | 62 ++++++++++++++-- 9 files changed, 241 insertions(+), 193 deletions(-) create mode 100644 benches/validate.rs diff --git a/Cargo.toml b/Cargo.toml index 7363c4c..08ebe72 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,6 @@ path = "src/lib.rs" [dependencies] error-chain ="^0.11.0" bitreader = "^0.3.0" -bit-vec = "^0.4.3" ring = "^0.12" rand = "^0.3.15" data-encoding = "^2.0" diff --git a/benches/validate.rs b/benches/validate.rs new file mode 100644 index 0000000..79a21cc --- /dev/null +++ b/benches/validate.rs @@ -0,0 +1,24 @@ +#![feature(test)] + +extern crate test; +extern crate bip39; + +use test::Bencher; + +use bip39::{Mnemonic, MnemonicType, Language}; + +#[bench] +fn validate(b: &mut Bencher) { + let phrase = "silly laptop awake length nature thunder category claim reveal supply attitude drip"; + + b.iter(|| { + let _ = Mnemonic::validate(phrase, Language::English); + }); +} + +#[bench] +fn new_mnemonic(b: &mut Bencher) { + b.iter(|| { + let _ = Mnemonic::new(MnemonicType::Type12Words, Language::English, ""); + }) +} diff --git a/src/crypto.rs b/src/crypto.rs index 06c4414..fcd84ea 100644 --- a/src/crypto.rs +++ b/src/crypto.rs @@ -6,13 +6,13 @@ //! -use ring::digest::{self, digest}; +use ring::digest::{self, digest, Digest}; use ring::pbkdf2; extern crate rand; use self::rand::{OsRng, Rng}; -use ::error::Error; +use ::error::Result; static PBKDF2_ROUNDS: u32 = 2048; static PBKDF2_BYTES: usize = 64; @@ -20,21 +20,19 @@ static PBKDF2_BYTES: usize = 64; /// SHA256 helper function, internal to the crate /// -pub(crate) fn sha256(input: &[u8]) -> Vec { +pub(crate) fn sha256(input: &[u8]) -> Digest { static DIGEST_ALG: &'static digest::Algorithm = &digest::SHA256; - let hash = digest(DIGEST_ALG, input); - - hash.as_ref().to_vec() + digest(DIGEST_ALG, input) } /// Random byte generator, used to create new mnemonics /// -pub(crate) fn gen_random_bytes(byte_length: usize) -> Result, Error> { +pub(crate) fn gen_random_bytes(byte_length: usize) -> Result> { let mut rng = OsRng::new()?; - let entropy = rng.gen_iter::().take(byte_length).collect::>(); + let entropy = rng.gen_iter().take(byte_length).collect(); Ok(entropy) } @@ -43,9 +41,8 @@ pub(crate) fn gen_random_bytes(byte_length: usize) -> Result, Error> { /// /// [Mnemonic]: ../mnemonic/struct.Mnemonic.html /// [Seed]: ../seed/struct.Seed.html -/// -pub(crate) fn pbkdf2(input: &[u8], - salt: String) -> Vec { +/// +pub(crate) fn pbkdf2(input: &[u8], salt: &str) -> Vec { let mut seed = vec![0u8; PBKDF2_BYTES]; diff --git a/src/language.rs b/src/language.rs index 6df6184..69bb56e 100644 --- a/src/language.rs +++ b/src/language.rs @@ -1,37 +1,30 @@ -use ::error::{Error, ErrorKind}; +use ::error::{ErrorKind, Result}; use std::collections::HashMap; mod lazy { - use std::collections::HashMap; - - /// lazy generation of the word list - fn gen_wordlist(lang_words: &str) -> Vec { + use super::HashMap; - lang_words.split_whitespace() - .map(|s| s.into()) - .collect() - } - - /// lazy generation of the word map - fn gen_wordmap(word_list: &Vec) -> HashMap { - - let mut word_map: HashMap = HashMap::new(); - for (i, item) in word_list.into_iter().enumerate() { - word_map.insert(item.to_owned(), i as u16); - } + /// lazy generation of the word list + fn gen_wordlist(lang_words: &str) -> Vec<&str> { + lang_words.split_whitespace().collect() + } - word_map - } + /// lazy generation of the word map + fn gen_wordmap(wordlist: &[&'static str]) -> HashMap<&'static str, u16> { + wordlist + .iter() + .enumerate() + .map(|(i, item)| (*item, i as u16)) + .collect() + } - static BIP39_WORDLIST_ENGLISH: &'static str = include_str!("bip39_english.txt"); + static BIP39_WORDLIST_ENGLISH: &str = include_str!("bip39_english.txt"); - lazy_static! { - pub static ref VEC_BIP39_WORDLIST_ENGLISH: Vec = { gen_wordlist(BIP39_WORDLIST_ENGLISH) }; - } + lazy_static! { + pub static ref VEC_BIP39_WORDLIST_ENGLISH: Vec<&'static str> = gen_wordlist(BIP39_WORDLIST_ENGLISH); - lazy_static! { - pub static ref HASHMAP_BIP39_WORDMAP_ENGLISH: HashMap = { gen_wordmap(&VEC_BIP39_WORDLIST_ENGLISH) }; - } + pub static ref HASHMAP_BIP39_WORDMAP_ENGLISH: HashMap<&'static str, u16> = gen_wordmap(&VEC_BIP39_WORDLIST_ENGLISH); + } } /// The language determines which words will be used in a mnemonic phrase, but also indirectly @@ -60,45 +53,42 @@ impl Language { /// /// ``` /// [Language]: ../language/struct.Language.html - pub fn for_locale(locale: S) -> Result where S: Into { - - let l = locale.into(); - - let lang = match &*l { + pub fn for_locale(locale: &str) -> Result { + let lang = match locale { "en_US.UTF-8" => Language::English, "en_GB.UTF-8" => Language::English, - _ => { return Err(ErrorKind::LanguageUnavailable.into()) } + _ => bail!(ErrorKind::LanguageUnavailable) }; Ok(lang) } - /// Get the word list for this language - pub fn get_wordlist(&self) -> &'static Vec { + /// Get the word list for this language + pub fn get_wordlist(&self) -> &'static [&'static str] { - match *self { - Language::English => &lazy::VEC_BIP39_WORDLIST_ENGLISH + match *self { + Language::English => &*lazy::VEC_BIP39_WORDLIST_ENGLISH } - } + } - /// Get a [`HashMap`][HashMap] that allows word -> index lookups in the word list - /// - /// The index of an individual word in the word list is used as the binary value of that word - /// when the phrase is turned into a [`Seed`][Seed]. - /// - /// [HashMap]: https://doc.rust-lang.org/std/collections/struct.HashMap.html - /// [Seed]: ../seed/struct.Seed.html - pub fn get_wordmap(&self) -> &'static HashMap { + /// Get a [`HashMap`][HashMap] that allows word -> index lookups in the word list + /// + /// The index of an individual word in the word list is used as the binary value of that word + /// when the phrase is turned into a [`Seed`][Seed]. + /// + /// [HashMap]: https://doc.rust-lang.org/std/collections/struct.HashMap.html + /// [Seed]: ../seed/struct.Seed.html + pub fn get_wordmap(&self) -> &'static HashMap<&'static str, u16> { - match *self { - Language::English => &lazy::HASHMAP_BIP39_WORDMAP_ENGLISH + match *self { + Language::English => &*lazy::HASHMAP_BIP39_WORDMAP_ENGLISH } - } + } } impl Default for Language { - fn default() -> Language { - Language::English - } + fn default() -> Language { + Language::English + } } diff --git a/src/lib.rs b/src/lib.rs index c680c45..2b83538 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -41,7 +41,6 @@ #[macro_use] extern crate lazy_static; extern crate data_encoding; extern crate bitreader; -extern crate bit_vec; extern crate ring; mod mnemonic; diff --git a/src/mnemonic.rs b/src/mnemonic.rs index 2990196..7d44b19 100644 --- a/src/mnemonic.rs +++ b/src/mnemonic.rs @@ -1,13 +1,12 @@ use bitreader::BitReader; -use bit_vec::BitVec; use data_encoding::HEXUPPER; +use util::{truncate, checksum, BitWriter}; use ::crypto::{gen_random_bytes, sha256}; -use ::error::{Error, ErrorKind}; +use ::error::{ErrorKind, Result}; use ::mnemonic_type::MnemonicType; use ::language::Language; -use ::util::bit_from_u16_as_u11; use ::seed::Seed; /// The primary type in this crate, most tasks require creating or using one. @@ -37,7 +36,7 @@ use ::seed::Seed; /// #[derive(Debug, Clone)] pub struct Mnemonic { - string: String, + phrase: String, seed: Seed, lang: Language, entropy: Vec, @@ -75,9 +74,9 @@ impl Mnemonic { /// [Mnemonic::get_seed()]: ./mnemonic/struct.Mnemonic.html#method.get_seed /// [Mnemonic::as_entropy()]: ./mnemonic/struct.Mnemonic.html#method.as_entropy /// [Mnemonic::get_entropy()]: ./mnemonic/struct.Mnemonic.html#method.get_entropy - pub fn new(mnemonic_type: MnemonicType, - lang: Language, - password: S) -> Result where S: Into { + pub fn new(mnemonic_type: MnemonicType, + lang: Language, + password: &str) -> Result { let entropy_bits = mnemonic_type.entropy_bits(); @@ -100,18 +99,17 @@ impl Mnemonic { /// ``` /// /// [Mnemonic]: ../mnemonic/struct.Mnemonic.html - pub fn from_entropy(entropy: &[u8], - mnemonic_type: MnemonicType, - lang: Language, - password: S) -> Result where S: Into { + pub fn from_entropy(entropy: &[u8], + mnemonic_type: MnemonicType, + lang: Language, + password: &str) -> Result { let entropy_length_bits = entropy.len() * 8; if entropy_length_bits != mnemonic_type.entropy_bits() { - return Err(ErrorKind::InvalidEntropyLength(entropy_length_bits, mnemonic_type).into()) + bail!(ErrorKind::InvalidEntropyLength(entropy_length_bits, mnemonic_type)); } - let num_words = mnemonic_type.word_count(); - + let word_count = mnemonic_type.word_count(); let word_list = lang.get_wordlist(); let entropy_hash = sha256(entropy); @@ -125,20 +123,38 @@ impl Mnemonic { // // ... and so on. It grabs the entropy and then the right number of hash bits and no more. - let mut combined = Vec::from(entropy); - combined.extend(&entropy_hash); + let mut combined = Vec::with_capacity(entropy.len() + 1); - let mut reader = BitReader::new(&combined); + combined.extend_from_slice(entropy); + combined.push(entropy_hash.as_ref()[0]); - let mut words: Vec<&str> = Vec::new(); - for _ in 0..num_words { - let n = reader.read_u16(11); - words.push(word_list[n.unwrap() as usize].as_ref()); - } + let phrase = { + let mut reader = BitReader::new(&combined); + let mut phrase = String::with_capacity(128); - let string = words.join(" "); + let n = reader.read_u16(11).expect("We are guaranteed to have enough bits to read; qed"); + phrase.push_str(word_list[n as usize]); - Mnemonic::from_string(string, lang, password.into()) + for _ in 1..word_count { + let n = reader.read_u16(11).expect("We are guaranteed to have enough bits to read; qed"); + phrase.push(' '); + phrase.push_str(word_list[n as usize]); + } + + phrase + }; + + let entropy = truncate(combined, entropy.len()); + let seed = Seed::generate(phrase.as_bytes(), password); + + let mnemonic = Mnemonic { + phrase, + seed, + lang, + entropy + }; + + Ok(mnemonic) } /// Create a [`Mnemonic`][Mnemonic] from generated entropy hexadecimal representation @@ -155,10 +171,10 @@ impl Mnemonic { /// ``` /// /// [Mnemonic]: ../mnemonic/struct.Mnemonic.html - pub fn from_entropy_hex(entropy: &str, + pub fn from_entropy_hex(entropy: &str, mnemonic_type: MnemonicType, lang: Language, - password: S) -> Result where S: Into { + password: &str) -> Result { Mnemonic::from_entropy(&HEXUPPER.decode(entropy.as_ref())?, mnemonic_type, lang, password) } @@ -179,25 +195,25 @@ impl Mnemonic { /// ``` /// /// [Mnemonic]: ../mnemonic/struct.Mnemonic.html - pub fn from_string(string: S, + pub fn from_string(phrase: S, lang: Language, - password: S) -> Result where S: Into { + password: S) -> Result where S: Into { - let m = string.into(); - let p = password.into(); + let phrase = phrase.into(); + let password = password.into(); // this also validates the checksum and phrase length before returning the entropy so we // can store it. We don't use the validate function here to avoid having a public API that // takes a phrase string and returns the entropy directly. See the Mnemonic::entropy() // docs for the reason. - let entropy = Mnemonic::entropy(&*m, lang)?; - let seed = Seed::generate(&m.as_bytes(), &p); + let entropy = Mnemonic::entropy(&phrase, lang)?; + let seed = Seed::generate(phrase.as_bytes(), &password); let mnemonic = Mnemonic { - string: (&m).clone(), - seed: seed, - lang: lang, - entropy: entropy + phrase, + seed, + lang, + entropy, }; Ok(mnemonic) @@ -226,9 +242,8 @@ impl Mnemonic { /// ``` /// /// [Mnemonic::from_string()]: ../mnemonic/struct.Mnemonic.html#method.from_string - pub fn validate(string: S, - lang: Language) -> Result<(), Error> where S: Into { - Mnemonic::entropy(string, lang).and(Ok(())) + pub fn validate(phrase: &str, lang: Language) -> Result<()> { + Mnemonic::entropy(phrase, lang).map(|_| ()) } /// Calculate the checksum, verify it and return the entropy @@ -236,49 +251,38 @@ impl Mnemonic { /// Only intended for internal use, as returning a `Vec` that looks a bit like it could be /// used as the seed is likely to cause problems for someone eventually. All the other functions /// that return something like that are explicit about what it is and what to use it for. - fn entropy(string: S, - lang: Language) -> Result, Error> where S: Into { - let m = string.into(); + fn entropy(phrase: &str, lang: Language) -> Result> { - let mnemonic_type = MnemonicType::for_phrase(&*m)?; + let mnemonic_type = MnemonicType::for_phrase(phrase)?; let entropy_bits = mnemonic_type.entropy_bits(); let checksum_bits = mnemonic_type.checksum_bits(); + let total_bits = mnemonic_type.total_bits(); - let word_map = lang.get_wordmap(); + let wordmap = lang.get_wordmap(); - let mut to_validate: BitVec = BitVec::new(); + let mut to_validate = BitWriter::with_capacity(total_bits); - for word in m.split(" ").into_iter() { - let n = match word_map.get(word) { - Some(n) => n, - None => return Err(ErrorKind::InvalidWord.into()) + for word in phrase.split(" ") { + let mut n = match wordmap.get(&word) { + Some(n) => *n, + None => bail!(ErrorKind::InvalidWord) }; - for i in 0..11 { - let bit = bit_from_u16_as_u11(*n, i); - to_validate.push(bit); - } - } - - let mut checksum_to_validate = BitVec::new(); - &checksum_to_validate.extend((&to_validate).into_iter().skip(entropy_bits).take(checksum_bits)); - assert!(checksum_to_validate.len() == checksum_bits, "invalid checksum size"); - let mut entropy_to_validate = BitVec::new(); - &entropy_to_validate.extend((&to_validate).into_iter().take(entropy_bits)); - assert!(entropy_to_validate.len() == entropy_bits, "invalid entropy size"); - - let entropy = entropy_to_validate.to_bytes(); + to_validate.push(n, 11); + } - let hash = sha256(entropy.as_ref()); + assert!(to_validate.len() == total_bits, "Insufficient amount of bits to validate"); - let entropy_hash_to_validate_bits = BitVec::from_bytes(hash.as_ref()); + let to_validate = to_validate.into_bytes(); + let entropy_bytes = entropy_bits / 8; + let checksum_to_validate = checksum(to_validate[entropy_bytes], checksum_bits); + let entropy = truncate(to_validate, entropy_bytes); + let hash = sha256(&entropy); + let new_checksum = checksum(hash.as_ref()[0], checksum_bits); - let mut new_checksum = BitVec::new(); - &new_checksum.extend(entropy_hash_to_validate_bits.into_iter().take(checksum_bits)); - assert!(new_checksum.len() == checksum_bits, "invalid new checksum size"); - if !(new_checksum == checksum_to_validate) { - return Err(ErrorKind::InvalidChecksum.into()) + if new_checksum != checksum_to_validate { + bail!(ErrorKind::InvalidChecksum) } Ok(entropy) @@ -305,14 +309,14 @@ impl Mnemonic { /// Get the mnemonic phrase as a string reference pub fn as_str(&self) -> &str { - self.string.as_ref() + &self.phrase } /// Get the mnemonic phrase as an owned string /// /// Note: this clones the internal Mnemonic String instance pub fn get_string(&self) -> String { - self.string.clone() + self.phrase.clone() } /// Get the [`Language`][Language] @@ -360,10 +364,8 @@ impl Mnemonic { /// /// let entropy: &[u8] = mnemonic.as_entropy(); /// ``` - /// - /// Note: this function clones the internal entropy bytes pub fn as_entropy(&self) -> &[u8] { - self.entropy.as_ref() + &self.entropy } } @@ -372,3 +374,16 @@ impl AsRef for Mnemonic { self.as_str() } } + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn new_same_as_from_phrase() { + let m1 = Mnemonic::new(MnemonicType::Type12Words, Language::English, "").unwrap(); + let m2 = Mnemonic::from_string(m1.as_str(), Language::English, "").unwrap(); + + assert_eq!(m1.entropy, m2.entropy, "Entropy must be the same"); + } +} diff --git a/src/mnemonic_type.rs b/src/mnemonic_type.rs index b0d5cc3..ff003ac 100644 --- a/src/mnemonic_type.rs +++ b/src/mnemonic_type.rs @@ -1,4 +1,4 @@ -use ::error::{Error, ErrorKind}; +use ::error::{ErrorKind, Result}; use std::fmt; /// Determines the number of words that will be present in a [`Mnemonic`][Mnemonic] phrase @@ -48,15 +48,14 @@ impl MnemonicType { /// /// let mnemonic_type = MnemonicType::for_word_count(12).unwrap(); /// ``` - pub fn for_word_count(size: usize) -> Result { - + pub fn for_word_count(size: usize) -> Result { let mnemonic_type = match size { 12 => MnemonicType::Type12Words, 15 => MnemonicType::Type15Words, 18 => MnemonicType::Type18Words, 21 => MnemonicType::Type21Words, 24 => MnemonicType::Type24Words, - _ => { return Err(ErrorKind::InvalidWordLength.into()) } + _ => bail!(ErrorKind::InvalidWordLength) }; Ok(mnemonic_type) @@ -73,15 +72,14 @@ impl MnemonicType { /// /// let mnemonic_type = MnemonicType::for_key_size(128).unwrap(); /// ``` - pub fn for_key_size(size: usize) -> Result { - + pub fn for_key_size(size: usize) -> Result { let mnemonic_type = match size { 128 => MnemonicType::Type12Words, 160 => MnemonicType::Type15Words, 192 => MnemonicType::Type18Words, 224 => MnemonicType::Type21Words, 256 => MnemonicType::Type24Words, - _ => { return Err(ErrorKind::InvalidKeysize.into()) } + _ => bail!(ErrorKind::InvalidKeysize) }; Ok(mnemonic_type) @@ -108,22 +106,10 @@ impl MnemonicType { /// ``` /// /// [MnemonicType::entropy_bits()]: ../mnemonic_type/struct.MnemonicType.html#method.entropy_bits - pub fn for_phrase(phrase: S) -> Result where S: Into { - - let m = phrase.into(); - - let v: Vec<&str> = m.split(" ").into_iter().collect(); - - let mnemonic_type = match v.len() { - 12 => MnemonicType::Type12Words, - 15 => MnemonicType::Type15Words, - 18 => MnemonicType::Type18Words, - 21 => MnemonicType::Type21Words, - 24 => MnemonicType::Type24Words, - _ => { return Err(ErrorKind::InvalidWordLength.into()) } - }; + pub fn for_phrase(phrase: &str) -> Result { + let word_count = phrase.split(" ").count(); - Ok(mnemonic_type) + Self::for_word_count(word_count) } /// Return the number of entropy+checksum bits @@ -140,16 +126,13 @@ impl MnemonicType { /// let total_bits = mnemonic_type.total_bits(); /// ``` pub fn total_bits(&self) -> usize { - - let total_bits: usize = match *self { + match *self { MnemonicType::Type12Words => 132, MnemonicType::Type15Words => 165, MnemonicType::Type18Words => 198, MnemonicType::Type21Words => 231, MnemonicType::Type24Words => 264 - }; - - total_bits + } } /// Return the number of entropy bits @@ -166,16 +149,13 @@ impl MnemonicType { /// let entropy_bits = mnemonic_type.entropy_bits(); /// ``` pub fn entropy_bits(&self) -> usize { - - let entropy_bits: usize = match *self { + match *self { MnemonicType::Type12Words => 128, MnemonicType::Type15Words => 160, MnemonicType::Type18Words => 192, MnemonicType::Type21Words => 224, MnemonicType::Type24Words => 256 - }; - - entropy_bits + } } /// Return the number of checksum bits @@ -192,16 +172,13 @@ impl MnemonicType { /// let checksum_bits = mnemonic_type.checksum_bits(); /// ``` pub fn checksum_bits(&self) -> usize { - - let checksum_bits: usize = match *self { + match *self { MnemonicType::Type12Words => 4, MnemonicType::Type15Words => 5, MnemonicType::Type18Words => 6, MnemonicType::Type21Words => 7, MnemonicType::Type24Words => 8 - }; - - checksum_bits + } } /// Return the number of words @@ -216,16 +193,13 @@ impl MnemonicType { /// let word_count = mnemonic_type.word_count(); /// ``` pub fn word_count(&self) -> usize { - - let word_count: usize = match *self { + match *self { MnemonicType::Type12Words => 12, MnemonicType::Type15Words => 15, MnemonicType::Type18Words => 18, MnemonicType::Type21Words => 21, MnemonicType::Type24Words => 24 - }; - - word_count + } } } diff --git a/src/seed.rs b/src/seed.rs index 29dc6ee..1812537 100644 --- a/src/seed.rs +++ b/src/seed.rs @@ -1,4 +1,4 @@ -use ::crypto::{pbkdf2}; +use ::crypto::pbkdf2; use data_encoding::HEXUPPER; @@ -26,7 +26,6 @@ use data_encoding::HEXUPPER; pub struct Seed { bytes: Vec, hex: String, - } impl Seed { @@ -35,12 +34,11 @@ impl Seed { /// /// Cannot be used outside the crate, in order to guarantee correctness /// [Mnemonic]: ../mnemonic/struct.Mnemonic.html - pub(crate) fn generate(entropy: &[u8], - password: &str) -> Seed { + pub(crate) fn generate(entropy: &[u8], password: &str) -> Seed { let salt = format!("mnemonic{}", password); - let seed_value = pbkdf2(entropy, salt); - let hex = HEXUPPER.encode(seed_value.as_ref()); + let seed_value = pbkdf2(entropy, &salt); + let hex = HEXUPPER.encode(&seed_value); Seed { bytes: seed_value, @@ -83,7 +81,7 @@ impl AsRef<[u8]> for Seed { impl AsRef for Seed { fn as_ref(&self) -> &str { - + self.as_hex() } } diff --git a/src/util.rs b/src/util.rs index 1080af0..9a79ea4 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,8 +1,60 @@ -pub(crate) fn bit_from_u16_as_u11(input: u16, position: u16) -> bool { - if position < 11 { - input & (1 << (10 - position)) != 0 - } else { - false +pub(crate) struct BitWriter { + offset: usize, + remainder: u32, + inner: Vec, +} + +impl BitWriter { + pub fn with_capacity(capacity: usize) -> Self { + let mut bytes = capacity / 8; + + if capacity % 8 != 0 { + bytes += 1; + } + + Self { + offset: 0, + remainder: 0, + inner: Vec::with_capacity(bytes) + } + } + + pub fn push(&mut self, source: u16, bits: usize) { + debug_assert!(bits > 0 && bits <= 16, "bits out of range"); + + let shift = 32 - bits; + + self.remainder |= ((source as u32) << shift) >> self.offset; + self.offset += bits; + + while self.offset >= 8 { + self.inner.push((self.remainder >> 24) as u8); + self.remainder <<= 8; + self.offset -= 8; + } } + + pub fn len(&self) -> usize { + self.inner.len() * 8 + self.offset + } + + pub fn into_bytes(mut self) -> Vec { + if self.offset != 0 { + self.inner.push((self.remainder >> 24) as u8); + } + + self.inner + } +} + +pub(crate) fn truncate(mut source: Vec, size: usize) -> Vec { + source.truncate(size); + source +} + +pub(crate) fn checksum(source: u8, bits: usize) -> u8 { + debug_assert!(bits <= 8, "Can operate on 8-bit integers only"); + + source >> (8 - bits) }