<commit_before>extern crate crypto;
extern crate rustc_serialize;
extern crate time;

use crypto::sha1::Sha1;
use crypto::hmac::Hmac;
use crypto::mac::Mac;
use crypto::digest::Digest;
use rustc_serialize::hex::FromHex;

pub fn from_hex(data: &str) -> Result<Vec<u8>, &str> {
    match data.from_hex() {
        Ok(d) => Ok(d),
        Err(_) => Err("Unable to decode hex")
    }
}

fn u64_from_be_bytes_4(bytes: &[u8], start: usize) -> u64 {
    let mut val = 0u64;

    val += (bytes[start]   as u64) << ((3 * 8) as u64);
    val += (bytes[start+1] as u64) << ((2 * 8) as u64);
    val += (bytes[start+2] as u64) << ((1 * 8) as u64);
    val += (bytes[start+3] as u64) << ((0 * 8) as u64);

    val
}

fn u64_to_be_bytes_8(data: u64) -> Vec<u8> {
    let mut v = Vec::new();
    v.push((data >> 56) as u8);
    v.push((data >> 48) as u8);
    v.push((data >> 40) as u8);
    v.push((data >> 32) as u8);
    v.push((data >> 24) as u8);
    v.push((data >> 16) as u8);
    v.push((data >> 8) as u8);
    v.push(data as u8);

    v
}

fn dynamic_truncation(hs: &[u8]) -> u64 {
    let offset_bits = (hs[19] & 0xf) as usize;
    let p = u64_from_be_bytes_4(hs, offset_bits);

    p & 0x7fffffff
}

pub fn hotp_custom<D: Digest>(key: &[u8], counter: u64, digits: u32,
                              hash: D) -> u64 {
    let mut hmac = Hmac::new(hash, key);
    let bytes = u64_to_be_bytes_8(counter);
    hmac.input(bytes.as_ref());
    let result = hmac.result();
    let hs = result.code();

    dynamic_truncation(hs) % 10_u64.pow(digits)
}

pub fn hotp_raw(key: &[u8], counter: u64, digits: u32) -> u64 {
    let hash = Sha1::new();
    hotp_custom(key, counter, digits, hash)
}

pub fn hotp(key: &str, counter: u64, digits: u32) -> Result<u64, &str> {
    match key.from_hex() {
        Ok(bytes) => Ok(hotp_raw(bytes.as_ref(), counter, digits)),
        Err(_) => Err("Unable to parse hex.")
    }
}

pub fn totp_custom<D: Digest>(key: &[u8], digits: u32, epoch: u64,
                              time_step: u64, current_time: u64,
                              hash: D) -> u64 {
    let counter = (current_time - epoch) / time_step;
    hotp_custom(key, counter, digits, hash)
}

pub fn totp_raw(key: &[u8], digits: u32, epoch: u64, time_step: u64) -> u64 {
    let hash = Sha1::new();
    let current_time = time::get_time();
    totp_custom(key, digits, epoch, time_step, current_time.sec as u64, hash)
}

pub fn totp(key: &str, digits: u32, epoch: u64,
            time_step: u64) -> Result<u64, &str> {
    match key.from_hex() {
        Ok(bytes) => Ok(totp_raw(bytes.as_ref(), digits, epoch, time_step)),
        Err(_) => Err("Unable to parse hex.")
    }
}

#[test]
fn test_hotp() {
    //use crypto::sha2::Sha256;

    assert_eq!(hotp_raw(b"\xff", 23, 6), 330795);
    assert_eq!(hotp("ff", 23, 6).unwrap(), 330795);
    assert_eq!(hotp_custom(b"\xff", 23, 6, Sha1::new()), 330795);
    assert_eq!(hotp_custom(from_hex("ff").unwrap().as_ref(), 23, 6, Sha1::new()), 330795);
    // test values from RFC 4226
    assert_eq!(hotp_raw(b"12345678901234567890", 0, 6), 755224);
    assert_eq!(hotp_raw(b"12345678901234567890", 1, 6), 287082);
    assert_eq!(hotp_raw(b"12345678901234567890", 2, 6), 359152);
    assert_eq!(hotp_raw(b"12345678901234567890", 3, 6), 969429);
    assert_eq!(hotp_raw(b"12345678901234567890", 4, 6), 338314);
    assert_eq!(hotp_raw(b"12345678901234567890", 5, 6), 254676);
    assert_eq!(hotp_raw(b"12345678901234567890", 6, 6), 287922);
    assert_eq!(hotp_raw(b"12345678901234567890", 7, 6), 162583);
    assert_eq!(hotp_raw(b"12345678901234567890", 8, 6), 399871);
    assert_eq!(hotp_raw(b"12345678901234567890", 9, 6), 520489);
    //assert_eq!(hotp_custom(from_hex("ff").unwrap().as_ref(), 23, 6, Sha256::new()), 225210);
    assert_eq!(hotp_custom(from_hex("3f906a54263361fccf").unwrap().as_ref(), 10, 7, Sha1::new()), 7615146);
    //assert_eq!(hotp_custom(from_hex("3f906a54263361fccf").unwrap().as_ref(), 10, 7, Sha256::new()), 6447746);
}

#[test]
fn test_totp() {
    //use crypto::sha2::Sha256;
    //use crypto::sha2::Sha512;
    assert_eq!(totp_custom(b"\xff", 6, 0, 1, 23, Sha1::new()), 330795);

    // test values from RFC 6238
    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 59, Sha1::new()), 94287082);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 59, Sha256::new()), 46119246);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 59, Sha512::new()), 90693936);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111109, Sha1::new()), 07081804);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111109, Sha256::new()), 68084774);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111109, Sha512::new()), 25091201);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111111, Sha1::new()), 14050471);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111111, Sha256::new()), 67062674);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111111, Sha512::new()), 99943326);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1234567890, Sha1::new()), 89005924);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1234567890, Sha256::new()), 91819424);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1234567890, Sha512::new()), 93441116);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 2000000000, Sha1::new()), 69279037);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 2000000000, Sha256::new()), 90698825);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 2000000000, Sha512::new()), 38618901);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 20000000000, Sha1::new()), 65353130);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 20000000000, Sha256::new()), 77737706);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 20000000000, Sha512::new()), 47863826);
}

#[cfg(test)]
mod ocra_tests {
    use crypto::sha1::Sha1;
    use crypto::digest::Digest;

    static STANDARD_KEY_20: &'static str = "12345678901234567890";
    static STANDARD_KEY_32: &'static str = "12345678901234567890123456789012";
    static STANDARD_KEY_64: &'static str = "1234567890123456789012345678901234567890123456789012345678901234";
    static PIN_1234_SHA1: [u8; 20] = [0x71, 0x10, 0xed, 0xa4, 0xd0, 0x9e, 0x06, 0x2a, 0xa5, 0xe4,
                                      0xa3, 0x90, 0xb0, 0xa5, 0x72, 0xac, 0x0d, 0x2c, 0x02, 0x20];
    #[test]
    fn sha1_pin_correct() {
        let mut sha: Sha1 = Sha1::new();
        sha.input_str("1234");

        let mut output = [0u8; 20];
        sha.result(&mut output);

        assert!(&PIN_1234_SHA1[..] == output);
    }
}<commit_msg>Remove useless function<commit_after>extern crate crypto;
extern crate rustc_serialize;
extern crate time;

use crypto::sha1::Sha1;
use crypto::hmac::Hmac;
use crypto::mac::Mac;
use crypto::digest::Digest;
use rustc_serialize::hex::FromHex;

pub fn from_hex(data: &str) -> Result<Vec<u8>, &str> {
    match data.from_hex() {
        Ok(d) => Ok(d),
        Err(_) => Err("Unable to decode hex")
    }
}

fn u64_from_be_bytes_4(bytes: &[u8], start: usize) -> u64 {
    let mut val = 0u64;

    val += (bytes[start]   as u64) << ((3 * 8) as u64);
    val += (bytes[start+1] as u64) << ((2 * 8) as u64);
    val += (bytes[start+2] as u64) << ((1 * 8) as u64);
    val += (bytes[start+3] as u64) << ((0 * 8) as u64);

    val
}

fn dynamic_truncation(hs: &[u8]) -> u64 {
    let offset_bits = (hs[19] & 0xf) as usize;
    let p = u64_from_be_bytes_4(hs, offset_bits);

    p & 0x7fffffff
}

pub fn hotp_custom<D: Digest>(key: &[u8], counter: u64, digits: u32,
                              hash: D) -> u64 {
    let mut hmac = Hmac::new(hash, key);
    let message = counter.to_be();
    let msg_ptr: &[u8] = unsafe { ::std::slice::from_raw_parts(&message as *const u64 as *const u8, 8) };
    hmac.input(msg_ptr);
    let result = hmac.result();
    let hs = result.code();

    dynamic_truncation(hs) % 10_u64.pow(digits)
}

pub fn hotp_raw(key: &[u8], counter: u64, digits: u32) -> u64 {
    let hash = Sha1::new();
    hotp_custom(key, counter, digits, hash)
}

pub fn hotp(key: &str, counter: u64, digits: u32) -> Result<u64, &str> {
    match key.from_hex() {
        Ok(bytes) => Ok(hotp_raw(bytes.as_ref(), counter, digits)),
        Err(_) => Err("Unable to parse hex.")
    }
}

pub fn totp_custom<D: Digest>(key: &[u8], digits: u32, epoch: u64,
                              time_step: u64, current_time: u64,
                              hash: D) -> u64 {
    let counter = (current_time - epoch) / time_step;
    hotp_custom(key, counter, digits, hash)
}

pub fn totp_raw(key: &[u8], digits: u32, epoch: u64, time_step: u64) -> u64 {
    let hash = Sha1::new();
    let current_time = time::get_time();
    totp_custom(key, digits, epoch, time_step, current_time.sec as u64, hash)
}

pub fn totp(key: &str, digits: u32, epoch: u64,
            time_step: u64) -> Result<u64, &str> {
    match key.from_hex() {
        Ok(bytes) => Ok(totp_raw(bytes.as_ref(), digits, epoch, time_step)),
        Err(_) => Err("Unable to parse hex.")
    }
}

#[test]
fn test_hotp() {
    //use crypto::sha2::Sha256;

    assert_eq!(hotp_raw(b"\xff", 23, 6), 330795);
    assert_eq!(hotp("ff", 23, 6).unwrap(), 330795);
    assert_eq!(hotp_custom(b"\xff", 23, 6, Sha1::new()), 330795);
    assert_eq!(hotp_custom(from_hex("ff").unwrap().as_ref(), 23, 6, Sha1::new()), 330795);
    // test values from RFC 4226
    assert_eq!(hotp_raw(b"12345678901234567890", 0, 6), 755224);
    assert_eq!(hotp_raw(b"12345678901234567890", 1, 6), 287082);
    assert_eq!(hotp_raw(b"12345678901234567890", 2, 6), 359152);
    assert_eq!(hotp_raw(b"12345678901234567890", 3, 6), 969429);
    assert_eq!(hotp_raw(b"12345678901234567890", 4, 6), 338314);
    assert_eq!(hotp_raw(b"12345678901234567890", 5, 6), 254676);
    assert_eq!(hotp_raw(b"12345678901234567890", 6, 6), 287922);
    assert_eq!(hotp_raw(b"12345678901234567890", 7, 6), 162583);
    assert_eq!(hotp_raw(b"12345678901234567890", 8, 6), 399871);
    assert_eq!(hotp_raw(b"12345678901234567890", 9, 6), 520489);
    //assert_eq!(hotp_custom(from_hex("ff").unwrap().as_ref(), 23, 6, Sha256::new()), 225210);
    assert_eq!(hotp_custom(from_hex("3f906a54263361fccf").unwrap().as_ref(), 10, 7, Sha1::new()), 7615146);
    //assert_eq!(hotp_custom(from_hex("3f906a54263361fccf").unwrap().as_ref(), 10, 7, Sha256::new()), 6447746);
}

#[test]
fn test_totp() {
    //use crypto::sha2::Sha256;
    //use crypto::sha2::Sha512;
    assert_eq!(totp_custom(b"\xff", 6, 0, 1, 23, Sha1::new()), 330795);

    // test values from RFC 6238
    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 59, Sha1::new()), 94287082);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 59, Sha256::new()), 46119246);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 59, Sha512::new()), 90693936);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111109, Sha1::new()), 07081804);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111109, Sha256::new()), 68084774);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111109, Sha512::new()), 25091201);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111111, Sha1::new()), 14050471);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111111, Sha256::new()), 67062674);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1111111111, Sha512::new()), 99943326);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1234567890, Sha1::new()), 89005924);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1234567890, Sha256::new()), 91819424);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 1234567890, Sha512::new()), 93441116);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 2000000000, Sha1::new()), 69279037);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 2000000000, Sha256::new()), 90698825);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 2000000000, Sha512::new()), 38618901);

    assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 20000000000, Sha1::new()), 65353130);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 20000000000, Sha256::new()), 77737706);
    //assert_eq!(totp_custom(b"12345678901234567890", 8, 0, 30, 20000000000, Sha512::new()), 47863826);
}

#[cfg(test)]
mod ocra_tests {
    use crypto::sha1::Sha1;
    use crypto::digest::Digest;

    static STANDARD_KEY_20: &'static str = "12345678901234567890";
    static STANDARD_KEY_32: &'static str = "12345678901234567890123456789012";
    static STANDARD_KEY_64: &'static str = "1234567890123456789012345678901234567890123456789012345678901234";
    static PIN_1234_SHA1: [u8; 20] = [0x71, 0x10, 0xed, 0xa4, 0xd0, 0x9e, 0x06, 0x2a, 0xa5, 0xe4,
                                      0xa3, 0x90, 0xb0, 0xa5, 0x72, 0xac, 0x0d, 0x2c, 0x02, 0x20];
    #[test]
    fn sha1_pin_correct() {
        let mut sha: Sha1 = Sha1::new();
        sha.input_str("1234");

        let mut output = [0u8; 20];
        sha.result(&mut output);

        assert!(&PIN_1234_SHA1[..] == output);
    }
}<|endoftext|>


<commit_before>#![feature(plugin)]
#![plugin(regex_macros)]
extern crate regex;

use std::ops::Index;
use regex::Regex;

#[stable(since="0.1.0")]
pub struct ReIdx(Regex);

#[unstable]
impl ReIdx {
    #[unstable]
    pub fn new(re: Regex) -> ReIdx {
        ReIdx(re)
    }
}

#[stable(since="0.1.0")]
impl Index<ReIdx> for str {
    type Output = str;
    fn index<'a>(&'a self, index: &ReIdx) -> &'a str {
        match index.0.find(self) {
            Some((start, end)) => &self[start..end],
            None => &self[..0]
        }
    }
}

#[test]
fn test_index_by_regex() {
    let re = ReIdx::new(regex!("^ab+"));
    assert_eq!("abb", &"abbcccdddd"[re]);
    assert_eq!("", &"acccdddd"[re]);
}
<commit_msg>accept into cow types<commit_after>#![feature(plugin)]
#![plugin(regex_macros)]
extern crate regex;

use std::ops::Index;
use regex::Regex;
use std::borrow::IntoCow;

#[stable(since="0.1.0")]
pub struct ReIdx(Regex);

impl ReIdx {
    #[unstable]
    pub fn from_regex(re: Regex) -> ReIdx {
        ReIdx(re)
    }

    #[unstable]
    pub fn new<'a, S: IntoCow<'a, str>>(s: S) -> Result<ReIdx, regex::Error> {
        Regex::new(&s.into_cow()).map(|re| ReIdx(re))
    }
}

#[stable(since="0.1.0")]
impl Index<ReIdx> for str {
    type Output = str;
    fn index<'a>(&'a self, index: &ReIdx) -> &'a str {
        match index.0.find(self) {
            Some((start, end)) => &self[start..end],
            None => &self[..0]
        }
    }
}

#[cfg(test)]
mod tests {
    use super::ReIdx;

    #[test]
    fn test_index_by_regex() {
        let re = ReIdx::new("^ab+").unwrap();
        assert_eq!("abb", &"abbcccdddd"[re]);
        assert_eq!("", &"acccdddd"[re]);
    }
}
<|endoftext|>


<commit_before>//! # quicksort
//!
//! Quicksort is an efficient in-line sorting algorithm created by Tony Hoare.

/// In-line sorting of a slice of T.
///
/// ```rust
/// # use ironsort::quicksort;
/// let presorted: Vec<u64> = vec![1, 1, 2, 3, 3, 4, 5, 5, 6, 9];
/// let mut vector: Vec<u64> = vec![3, 1, 4, 1, 5, 9, 2, 6, 5, 3];
/// quicksort(&mut vector);
///
/// assert_eq!(vector, presorted.as_slice());
/// ```

pub fn quicksort<T: PartialOrd>(vec: &mut [T]) {
    if vec.len() <= 1 {
        return;
    }
    let pivot: usize = 0;
    let mut first_opened = Vec::new();
    let mut last_closed: usize = 0;

    for i in 1..vec.len() {
        first_opened.push(i);
        if vec[i] < vec[pivot] {
            let to = first_opened.remove(0);
            vec.swap(i, to);
            last_closed = to;
        }
    }

    vec.swap(pivot, last_closed);
    quicksort(&mut vec[0..last_closed]);
    quicksort(&mut vec[last_closed+1..]);
}

#[cfg(test)]
extern crate rand;

#[cfg(test)]
mod tests {
    use rand::{self, Rng};
    use super::*;

    #[test]
    fn test_vec_of_u64() {
        let presorted: Vec<u64> = vec![1, 1, 2, 3, 3, 4, 5, 5, 6, 9];
        let mut vector: Vec<u64> = vec![3, 1, 4, 1, 5, 9, 2, 6, 5, 3];
        quicksort(&mut vector);

        assert_eq!(vector, presorted.as_slice());
    }

    #[test]
    fn test_vec_of_u8() {
        let presorted = "        Tabcdeeefghhijklmnoooopqrrstuuvwxyz".to_string().into_bytes();
        let mut vector = "The quick brown fox jumps over the lazy dog".to_string().into_bytes();
        quicksort(&mut vector);

        assert_eq!(presorted, vector);
    }

    #[test]
    fn test_array_of_u8() {
        let presorted: [u8; 10] = [1, 2, 3, 4, 5, 5, 5, 7, 9, 10];
        let mut vector: [u8; 10] = [7, 5, 3, 2, 5, 1, 4, 5, 9, 10];
        quicksort(&mut vector);

        assert_eq!(presorted, vector);
    }

    #[test]
    fn test_random() {
        let mut rng = rand::thread_rng();

        for _ in 0u64 .. 10_000u64 {
            let len: usize = rng.gen();
            let mut vector = rng.gen_iter()
                                .take((len % 64) + 1)
                                .collect::<Vec<usize>>();
            quicksort(&mut vector);

            for i in 0 .. vector.len() - 1 {
                assert!(vector[i] <= vector[i + 1])
            }
        }
    }
}
<commit_msg>Split quicksort and quicksort_by<commit_after>//! # quicksort
//!
//! Quicksort is an efficient in-line sorting algorithm created by Tony Hoare.
//!

use std::cmp::Ordering;

/// In-line sorting of a slice of T.
///
/// ```rust
/// # use ironsort::quicksort;
/// let presorted: Vec<u64> = vec![1, 1, 2, 3, 3, 4, 5, 5, 6, 9];
/// let mut vector: Vec<u64> = vec![3, 1, 4, 1, 5, 9, 2, 6, 5, 3];
/// quicksort(&mut vector);
///
/// assert_eq!(vector, presorted.as_slice());
/// ```
pub fn quicksort<T: Ord>(vec: &mut [T]) {
    quicksort_by(vec, &Ord::cmp)
}

pub fn quicksort_by<T: PartialOrd, F>(vec: &mut [T], cmp: &F)
    where F: Fn(&T, &T) -> Ordering {
    if vec.len() <= 1 {
        return;
    }
    let pivot: usize = 0;
    let mut first_opened = Vec::new();
    let mut last_closed: usize = 0;

    for i in 1..vec.len() {
        first_opened.push(i);
        match cmp(&vec[i], &vec[pivot]) {
            Ordering::Less => {
                let to = first_opened.remove(0);
                vec.swap(i, to);
                last_closed = to;
            }
            _ => { }
        }
    }

    vec.swap(pivot, last_closed);
    quicksort_by(&mut vec[0..last_closed], cmp);
    quicksort_by(&mut vec[last_closed+1..], cmp);
}

#[cfg(test)]
extern crate rand;

#[cfg(test)]
mod tests {
    use rand::{self, Rng};
    use super::*;

    #[test]
    fn test_vec_of_u64() {
        let presorted: Vec<u64> = vec![1, 1, 2, 3, 3, 4, 5, 5, 6, 9];
        let mut vector: Vec<u64> = vec![3, 1, 4, 1, 5, 9, 2, 6, 5, 3];
        quicksort(&mut vector);

        assert_eq!(vector, presorted.as_slice());
    }

    #[test]
    fn test_vec_of_u8() {
        let presorted = "        Tabcdeeefghhijklmnoooopqrrstuuvwxyz".to_string().into_bytes();
        let mut vector = "The quick brown fox jumps over the lazy dog".to_string().into_bytes();
        quicksort(&mut vector);

        assert_eq!(presorted, vector);
    }

    #[test]
    fn test_array_of_u8() {
        let presorted: [u8; 10] = [1, 2, 3, 4, 5, 5, 5, 7, 9, 10];
        let mut vector: [u8; 10] = [7, 5, 3, 2, 5, 1, 4, 5, 9, 10];
        quicksort(&mut vector);

        assert_eq!(presorted, vector);
    }

    #[test]
    fn test_random() {
        let mut rng = rand::thread_rng();

        for _ in 0u64 .. 10_000u64 {
            let len: usize = rng.gen();
            let mut vector = rng.gen_iter()
                                .take((len % 64) + 1)
                                .collect::<Vec<usize>>();
            quicksort(&mut vector);

            for i in 0 .. vector.len() - 1 {
                assert!(vector[i] <= vector[i + 1])
            }
        }
    }
}
<|endoftext|>


<commit_before>#![feature(macro_rules)]

use std::num::Int;

pub static BLOCK_BYTES  : uint  = 128;
pub static KEY_BYTES    : uint  = 64;
pub static OUT_BYTES    : uint  = 64;

static IV : [u64, ..8] = [
    0x6a09e667f3bcc908, 0xbb67ae8584caa73b, 0x3c6ef372fe94f82b, 0xa54ff53a5f1d36f1,
    0x510e527fade682d1, 0x9b05688c2b3e6c1f, 0x1f83d9abfb41bd6b, 0x5be0cd19137e2179,
];

static SIGMA : [[u8, ..16], ..12] = [
    [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15 ],
    [ 14, 10,  4,  8,  9, 15, 13,  6,  1, 12,  0,  2, 11,  7,  5,  3 ],
    [ 11,  8, 12,  0,  5,  2, 15, 13, 10, 14,  3,  6,  7,  1,  9,  4 ],
    [  7,  9,  3,  1, 13, 12, 11, 14,  2,  6,  5, 10,  4,  0, 15,  8 ],
    [  9,  0,  5,  7,  2,  4, 10, 15, 14,  1, 11, 12,  6,  8,  3, 13 ],
    [  2, 12,  6, 10,  0, 11,  8,  3,  4, 13,  7,  5, 15, 14,  1,  9 ],
    [ 12,  5,  1, 15, 14, 13,  4, 10,  0,  7,  6,  3,  9,  2,  8, 11 ],
    [ 13, 11,  7, 14, 12,  1,  3,  9,  5,  0, 15,  4,  8,  6,  2, 10 ],
    [  6, 15, 14,  9, 11,  3,  0,  8, 12,  2, 13,  7,  1,  4, 10,  5 ],
    [ 10,  2,  8,  4,  7,  6,  1,  5, 15, 11,  9, 14,  3, 12, 13 , 0 ],
    [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15 ],
    [ 14, 10,  4,  8,  9, 15, 13,  6,  1, 12,  0,  2, 11,  7,  5,  3 ],
];

pub struct Blake2b {
    h: [u64, ..8],
    t: [u64, ..2],
    f: [u64, ..2],
    buf: [u8, ..2*BLOCK_BYTES],
    buf_len: uint,
}

impl Blake2b {
    pub fn new(size: uint) -> Blake2b {
        assert!(size > 0 && size <= OUT_BYTES);

        let param = encode_params(size as u8, 0);
        let mut state = IV;

        for i in range(0, state.len()) {
            state[i] ^= load64(param.slice_from(i*8));
        }

        Blake2b{
            h: state,
            t: [0, 0],
            f: [0, 0],
            buf: [0u8, ..2*BLOCK_BYTES],
            buf_len: 0,
        }
    }

    pub fn new_with_key(size: uint, key: &[u8]) -> Blake2b {
        assert!(size > 0 && size <= OUT_BYTES);
        assert!(key.len() > 0 && key.len() <= KEY_BYTES);

        let param = encode_params(size as u8, key.len() as u8);
        let mut state = IV;

        for i in range(0, state.len()) {
            state[i] ^= load64(param.slice_from(i*8));
        }

        let mut b = Blake2b{
            h: state,
            t: [0, 0],
            f: [0, 0],
            buf: [0u8, ..2*BLOCK_BYTES],
            buf_len: 0,
        };

        let mut block = [0u8, ..BLOCK_BYTES];
        for i in range(0, key.len()) {
            block[i] = key[i];
        }
        b.update(block.as_slice());
        b
    }

    pub fn update(&mut self, m: &[u8]) {
        let mut m = m;

        while m.len() > 0 {
            let left = self.buf_len;
            let fill = 2 * BLOCK_BYTES - left;

            if m.len() > fill {
                for i in range(0, fill) {
                    self.buf[left+i] = m[i];
                }
                self.buf_len += fill;
                m = m.slice_from(fill);
                self.increment_counter(BLOCK_BYTES as u64);
                self.compress();
                for i in range(0, BLOCK_BYTES) {
                    self.buf[i] = self.buf[i+BLOCK_BYTES];
                }
                self.buf_len -= BLOCK_BYTES;
            } else {
                for i in range(0, m.len()) {
                    self.buf[left+i] = m[i];
                }
                self.buf_len += m.len();
                m = m.slice_from(m.len());
            }
        }
    }

    pub fn final(&mut self, out: &mut [u8]) {
        let mut buf = [0u8, ..OUT_BYTES];
        if self.buf_len > BLOCK_BYTES {
            self.increment_counter(BLOCK_BYTES as u64);
            self.compress();
            for i in range(0, BLOCK_BYTES) {
                self.buf[i] = self.buf[i+BLOCK_BYTES];
            }
            self.buf_len -= BLOCK_BYTES;
        }
        let n = self.buf_len as u64;
        self.increment_counter(n);
        self.f[0] = !0;
        for i in range(self.buf_len, self.buf.len()) {
            self.buf[i] = 0;
        }
        self.compress();
        for i in range(0, self.h.len()) {
            store64(buf.mut_slice_from(i*8), self.h[i]);
        }

        for i in range(0, out.len()) {
            out[i] = buf[i];
        }
    }

    fn increment_counter(&mut self, inc: u64) {
        self.t[0] += inc;
        self.t[1] += if self.t[0] < inc {1} else {0};
    }

    fn compress(&mut self) {
        let mut m = [0u64, ..16];
        let mut v = [0u64, ..16];
        let block = self.buf.as_slice();

        assert!(block.len() >= BLOCK_BYTES);

        for i in range(0, m.len()) {
            m[i] = load64(block.slice_from(i*8));
        }

        for i in range(0, 8) {
            v[i] = self.h[i];
        }

        v[ 8] = IV[0];
        v[ 9] = IV[1];
        v[10] = IV[2];
        v[11] = IV[3];
        v[12] = self.t[0] ^ IV[4];
        v[13] = self.t[1] ^ IV[5];
        v[14] = self.f[0] ^ IV[6];
        v[15] = self.f[1] ^ IV[7];

        macro_rules! g(
            ($r: expr, $i: expr, $a: expr, $b: expr, $c: expr, $d: expr) => ({
                $a = $a + $b + m[SIGMA[$r][2*$i+0] as uint];
                $d = ($d ^ $a).rotate_right(32);
                $c = $c + $d;
                $b = ($b ^ $c).rotate_right(24);
                $a = $a + $b + m[SIGMA[$r][2*$i+1] as uint];
                $d = ($d ^ $a).rotate_right(16);
                $c = $c + $d;
                $b = ($b ^ $c).rotate_right(63);
            });
        )

        macro_rules! round(
            ($r: expr) => ({
                g!($r, 0, v[ 0], v[ 4], v[ 8], v[12]);
                g!($r, 1, v[ 1], v[ 5], v[ 9], v[13]);
                g!($r, 2, v[ 2], v[ 6], v[10], v[14]);
                g!($r, 3, v[ 3], v[ 7], v[11], v[15]);
                g!($r, 4, v[ 0], v[ 5], v[10], v[15]);
                g!($r, 5, v[ 1], v[ 6], v[11], v[12]);
                g!($r, 6, v[ 2], v[ 7], v[ 8], v[13]);
                g!($r, 7, v[ 3], v[ 4], v[ 9], v[14]);
            });
        )

        for i in range(0, 12) {
            round!(i);
        }

        for i in range(0, 8) {
            self.h[i] = self.h[i] ^ v[i] ^ v[i+8];
        }
    }
}

fn encode_params(size: u8, keylen: u8) -> [u8, ..64] {
    let mut param = [0u8, ..64];
    param[0] = size as u8;
    param[1] = keylen as u8;
    param[2] = 1; // fanout
    param[3] = 1; // depth
    param
}

fn load64(b: &[u8]) -> u64 {
    let mut v = 0u64;
    for i in range(0, 8) {
        v |= b[i] as u64 << (8*i);
    }
    v
}

fn store64(b: &mut [u8], v: u64) {
    let mut w = v;
    for i in range(0, 8) {
        b[i] = w as u8;
        w >>= 8;
    }
}

#[cfg(test)]
mod test {
    use super::{Blake2b, KEY_BYTES, OUT_BYTES};

    mod kat;

    #[test]
    fn test_blake2b_kat() {
        let mut input = [0u8, ..256];
        for i in range(0, input.len()) {
            input[i] = i as u8;
        }

        for i in range(0, kat::blake2b.len()) {
            let mut h = Blake2b::new(OUT_BYTES);
            let mut out = [0u8, ..OUT_BYTES];
            h.update(input.slice(0, i));
            h.final(out);
            assert_eq!(out.as_slice(), kat::blake2b[i].as_slice());
        }
    }

    #[test]
    fn test_blake2b_keyed_kat() {
        let mut input = [0u8, ..256];
        let mut key = [0u8, ..KEY_BYTES];

        for i in range(0, input.len()) {
            input[i] = i as u8;
        }

        for i in range(0, key.len()) {
            key[i] = i as u8;
        }

        for i in range(0, kat::blake2b_keyed.len()) {
            let mut h = Blake2b::new_with_key(OUT_BYTES, key.as_slice());
            let mut out = [0u8, ..OUT_BYTES];
            h.update(input.slice(0, i));
            h.final(out);
            assert_eq!(out.as_slice(), kat::blake2b_keyed[i].as_slice());
        }
    }
}
<commit_msg>remove unused import<commit_after>#![feature(macro_rules)]

pub static BLOCK_BYTES  : uint  = 128;
pub static KEY_BYTES    : uint  = 64;
pub static OUT_BYTES    : uint  = 64;

static IV : [u64, ..8] = [
    0x6a09e667f3bcc908, 0xbb67ae8584caa73b, 0x3c6ef372fe94f82b, 0xa54ff53a5f1d36f1,
    0x510e527fade682d1, 0x9b05688c2b3e6c1f, 0x1f83d9abfb41bd6b, 0x5be0cd19137e2179,
];

static SIGMA : [[u8, ..16], ..12] = [
    [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15 ],
    [ 14, 10,  4,  8,  9, 15, 13,  6,  1, 12,  0,  2, 11,  7,  5,  3 ],
    [ 11,  8, 12,  0,  5,  2, 15, 13, 10, 14,  3,  6,  7,  1,  9,  4 ],
    [  7,  9,  3,  1, 13, 12, 11, 14,  2,  6,  5, 10,  4,  0, 15,  8 ],
    [  9,  0,  5,  7,  2,  4, 10, 15, 14,  1, 11, 12,  6,  8,  3, 13 ],
    [  2, 12,  6, 10,  0, 11,  8,  3,  4, 13,  7,  5, 15, 14,  1,  9 ],
    [ 12,  5,  1, 15, 14, 13,  4, 10,  0,  7,  6,  3,  9,  2,  8, 11 ],
    [ 13, 11,  7, 14, 12,  1,  3,  9,  5,  0, 15,  4,  8,  6,  2, 10 ],
    [  6, 15, 14,  9, 11,  3,  0,  8, 12,  2, 13,  7,  1,  4, 10,  5 ],
    [ 10,  2,  8,  4,  7,  6,  1,  5, 15, 11,  9, 14,  3, 12, 13 , 0 ],
    [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15 ],
    [ 14, 10,  4,  8,  9, 15, 13,  6,  1, 12,  0,  2, 11,  7,  5,  3 ],
];

pub struct Blake2b {
    h: [u64, ..8],
    t: [u64, ..2],
    f: [u64, ..2],
    buf: [u8, ..2*BLOCK_BYTES],
    buf_len: uint,
}

impl Blake2b {
    pub fn new(size: uint) -> Blake2b {
        assert!(size > 0 && size <= OUT_BYTES);

        let param = encode_params(size as u8, 0);
        let mut state = IV;

        for i in range(0, state.len()) {
            state[i] ^= load64(param.slice_from(i*8));
        }

        Blake2b{
            h: state,
            t: [0, 0],
            f: [0, 0],
            buf: [0u8, ..2*BLOCK_BYTES],
            buf_len: 0,
        }
    }

    pub fn new_with_key(size: uint, key: &[u8]) -> Blake2b {
        assert!(size > 0 && size <= OUT_BYTES);
        assert!(key.len() > 0 && key.len() <= KEY_BYTES);

        let param = encode_params(size as u8, key.len() as u8);
        let mut state = IV;

        for i in range(0, state.len()) {
            state[i] ^= load64(param.slice_from(i*8));
        }

        let mut b = Blake2b{
            h: state,
            t: [0, 0],
            f: [0, 0],
            buf: [0u8, ..2*BLOCK_BYTES],
            buf_len: 0,
        };

        let mut block = [0u8, ..BLOCK_BYTES];
        for i in range(0, key.len()) {
            block[i] = key[i];
        }
        b.update(block.as_slice());
        b
    }

    pub fn update(&mut self, m: &[u8]) {
        let mut m = m;

        while m.len() > 0 {
            let left = self.buf_len;
            let fill = 2 * BLOCK_BYTES - left;

            if m.len() > fill {
                for i in range(0, fill) {
                    self.buf[left+i] = m[i];
                }
                self.buf_len += fill;
                m = m.slice_from(fill);
                self.increment_counter(BLOCK_BYTES as u64);
                self.compress();
                for i in range(0, BLOCK_BYTES) {
                    self.buf[i] = self.buf[i+BLOCK_BYTES];
                }
                self.buf_len -= BLOCK_BYTES;
            } else {
                for i in range(0, m.len()) {
                    self.buf[left+i] = m[i];
                }
                self.buf_len += m.len();
                m = m.slice_from(m.len());
            }
        }
    }

    pub fn final(&mut self, out: &mut [u8]) {
        let mut buf = [0u8, ..OUT_BYTES];
        if self.buf_len > BLOCK_BYTES {
            self.increment_counter(BLOCK_BYTES as u64);
            self.compress();
            for i in range(0, BLOCK_BYTES) {
                self.buf[i] = self.buf[i+BLOCK_BYTES];
            }
            self.buf_len -= BLOCK_BYTES;
        }
        let n = self.buf_len as u64;
        self.increment_counter(n);
        self.f[0] = !0;
        for i in range(self.buf_len, self.buf.len()) {
            self.buf[i] = 0;
        }
        self.compress();
        for i in range(0, self.h.len()) {
            store64(buf.mut_slice_from(i*8), self.h[i]);
        }

        for i in range(0, out.len()) {
            out[i] = buf[i];
        }
    }

    fn increment_counter(&mut self, inc: u64) {
        self.t[0] += inc;
        self.t[1] += if self.t[0] < inc {1} else {0};
    }

    fn compress(&mut self) {
        let mut m = [0u64, ..16];
        let mut v = [0u64, ..16];
        let block = self.buf.as_slice();

        assert!(block.len() >= BLOCK_BYTES);

        for i in range(0, m.len()) {
            m[i] = load64(block.slice_from(i*8));
        }

        for i in range(0, 8) {
            v[i] = self.h[i];
        }

        v[ 8] = IV[0];
        v[ 9] = IV[1];
        v[10] = IV[2];
        v[11] = IV[3];
        v[12] = self.t[0] ^ IV[4];
        v[13] = self.t[1] ^ IV[5];
        v[14] = self.f[0] ^ IV[6];
        v[15] = self.f[1] ^ IV[7];

        macro_rules! g(
            ($r: expr, $i: expr, $a: expr, $b: expr, $c: expr, $d: expr) => ({
                $a = $a + $b + m[SIGMA[$r][2*$i+0] as uint];
                $d = ($d ^ $a).rotate_right(32);
                $c = $c + $d;
                $b = ($b ^ $c).rotate_right(24);
                $a = $a + $b + m[SIGMA[$r][2*$i+1] as uint];
                $d = ($d ^ $a).rotate_right(16);
                $c = $c + $d;
                $b = ($b ^ $c).rotate_right(63);
            });
        )

        macro_rules! round(
            ($r: expr) => ({
                g!($r, 0, v[ 0], v[ 4], v[ 8], v[12]);
                g!($r, 1, v[ 1], v[ 5], v[ 9], v[13]);
                g!($r, 2, v[ 2], v[ 6], v[10], v[14]);
                g!($r, 3, v[ 3], v[ 7], v[11], v[15]);
                g!($r, 4, v[ 0], v[ 5], v[10], v[15]);
                g!($r, 5, v[ 1], v[ 6], v[11], v[12]);
                g!($r, 6, v[ 2], v[ 7], v[ 8], v[13]);
                g!($r, 7, v[ 3], v[ 4], v[ 9], v[14]);
            });
        )

        for i in range(0, 12) {
            round!(i);
        }

        for i in range(0, 8) {
            self.h[i] = self.h[i] ^ v[i] ^ v[i+8];
        }
    }
}

fn encode_params(size: u8, keylen: u8) -> [u8, ..64] {
    let mut param = [0u8, ..64];
    param[0] = size as u8;
    param[1] = keylen as u8;
    param[2] = 1; // fanout
    param[3] = 1; // depth
    param
}

fn load64(b: &[u8]) -> u64 {
    let mut v = 0u64;
    for i in range(0, 8) {
        v |= b[i] as u64 << (8*i);
    }
    v
}

fn store64(b: &mut [u8], v: u64) {
    let mut w = v;
    for i in range(0, 8) {
        b[i] = w as u8;
        w >>= 8;
    }
}

#[cfg(test)]
mod test {
    use super::{Blake2b, KEY_BYTES, OUT_BYTES};

    mod kat;

    #[test]
    fn test_blake2b_kat() {
        let mut input = [0u8, ..256];
        for i in range(0, input.len()) {
            input[i] = i as u8;
        }

        for i in range(0, kat::blake2b.len()) {
            let mut h = Blake2b::new(OUT_BYTES);
            let mut out = [0u8, ..OUT_BYTES];
            h.update(input.slice(0, i));
            h.final(out);
            assert_eq!(out.as_slice(), kat::blake2b[i].as_slice());
        }
    }

    #[test]
    fn test_blake2b_keyed_kat() {
        let mut input = [0u8, ..256];
        let mut key = [0u8, ..KEY_BYTES];

        for i in range(0, input.len()) {
            input[i] = i as u8;
        }

        for i in range(0, key.len()) {
            key[i] = i as u8;
        }

        for i in range(0, kat::blake2b_keyed.len()) {
            let mut h = Blake2b::new_with_key(OUT_BYTES, key.as_slice());
            let mut out = [0u8, ..OUT_BYTES];
            h.update(input.slice(0, i));
            h.final(out);
            assert_eq!(out.as_slice(), kat::blake2b_keyed[i].as_slice());
        }
    }
}
<|endoftext|>


<commit_before>#[cfg(test)]
mod test;

use smallvec::{Array, SmallVec};

use std::convert::{From, Into};
use std::fmt::{self, Debug, Formatter};
use std::iter::FromIterator;

/// A `NibbleVec` backed by a `SmallVec` with 64 inline element slots.
/// This will not allocate until more than 64 elements are added.
pub type Nibblet = NibbleVec<[u8; 64]>;

/// A data-structure for storing a sequence of 4-bit values.
///
/// Values are stored in a `Vec<u8>`, with two values per byte.
///
/// Values at even indices are stored in the most-significant half of their byte,
/// while values at odd indices are stored in the least-significant half.
///
/// Imagine a vector of [MSB][msb-wiki] first bytes, and you'll be right.
///
/// n = [_ _ | _ _ | _ _]
///
/// [msb-wiki]: http://en.wikipedia.org/wiki/Most_significant_bit
#[derive(Clone, Default)]
pub struct NibbleVec<A: Array<Item = u8>> {
    length: usize,
    data: SmallVec<A>,
}

impl<A: Array<Item = u8>> NibbleVec<A> {
    /// Create an empty nibble vector.
    pub fn new() -> NibbleVec<A> {
        NibbleVec {
            length: 0,
            data: SmallVec::new(),
        }
    }

    /// Create a nibble vector from a vector of bytes.
    ///
    /// Each byte is split into two 4-bit entries (MSB, LSB).
    #[inline]
    pub fn from_byte_vec(vec: Vec<u8>) -> NibbleVec<A> {
        let length = 2 * vec.len();
        NibbleVec {
            length,
            data: SmallVec::from_iter(vec),
        }
    }

    /// Create a nibble vector from a `SmallVec` of Size.
    ///
    /// Each byte is split into two 4-bit entries (MSB, LSB).
    #[inline]
    pub fn from_small_vec(length: usize) -> NibbleVec<A> {
        NibbleVec {
            length,
            data: SmallVec::new(),
        }
    }

    /// Returns a byte slice of the nibble vector's contents.
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        &self.data[..]
    }

    /// Converts a nibble vector into a byte vector.
    ///
    /// This consumes the nibble vector, so we do not need to copy its contents.
    #[inline]
    pub fn into_bytes(self) -> Vec<u8> {
        self.data.to_vec()
    }

    /// Get the number of elements stored in the vector.
    #[inline]
    pub fn len(&self) -> usize {
        self.length
    }

    /// Returns `true` if the nibble vector has a length of 0.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Fetch a single entry from the vector.
    ///
    /// Guaranteed to be a value in the interval [0, 15].
    ///
    /// **Panics** if `idx >= self.len()`.
    #[inline]
    pub fn get(&self, idx: usize) -> u8 {
        if idx >= self.length {
            panic!(
                "NibbleVec index out of bounds: len is {}, index is {}",
                self.length, idx
            );
        }
        let vec_idx = idx / 2;
        match idx % 2 {
            // If the index is even, take the first (most significant) half of the stored byte.
            0 => self.data[vec_idx] >> 4,
            // If the index is odd, take the second (least significant) half.
            _ => self.data[vec_idx] & 0x0F,
        }
    }

    /// Add a single nibble to the vector.
    ///
    /// Only the 4 least-significant bits of the value are used.
    #[inline]
    pub fn push(&mut self, val: u8) {
        if self.length % 2 == 0 {
            self.data.push(val << 4);
        } else {
            let vec_len = self.data.len();

            // Zero the second half of the last byte just to be safe.
            self.data[vec_len - 1] &= 0xF0;

            // Write the new value.
            self.data[vec_len - 1] |= val & 0x0F;
        }
        self.length += 1;
    }

    /// Split the vector into two parts.
    ///
    /// All elements at or following the given index are returned in a new `NibbleVec`,
    /// with exactly `idx` elements remaining in this vector.
    ///
    /// **Panics** if `idx > self.len()`.
    pub fn split(&mut self, idx: usize) -> NibbleVec<A> {
        if idx > self.length {
            panic!(
                "attempted to split past vector end. len is {}, index is {}",
                self.length, idx
            );
        } else if idx == self.length {
            NibbleVec::new()
        } else if idx % 2 == 0 {
            self.split_even(idx)
        } else {
            self.split_odd(idx)
        }
    }

    /// Split function for odd *indices*.
    #[inline]
    fn split_odd(&mut self, idx: usize) -> NibbleVec<A> {
        // let tail_vec_size = (self.length - idx) / 2;
        let mut tail = NibbleVec::from_small_vec(0);

        // Perform an overlap copy, copying the last nibble of the original vector only if
        // the length of the new tail is *odd*.
        let tail_length = self.length - idx;
        let take_last = tail_length % 2 == 1;
        self.overlap_copy(
            idx / 2,
            self.data.len(),
            &mut tail.data,
            &mut tail.length,
            take_last,
        );

        // Remove the copied bytes, being careful to skip the idx byte.
        for _ in (idx / 2 + 1)..self.data.len() {
            self.data.pop();
        }

        // Zero the second half of the index byte so as to maintain the last-nibble invariant.
        self.data[idx / 2] &= 0xF0;

        // Update the length of the first NibbleVec.
        self.length = idx;

        tail
    }

    /// Split function for even *indices*.
    #[inline]
    fn split_even(&mut self, idx: usize) -> NibbleVec<A> {
        // Avoid allocating a temporary vector by copying all the bytes in order, then popping them.

        // Possible to prove: l_d - ⌊i / 2⌋ = ⌊(l_v - i + 1) / 2⌋
        //  where l_d = self.data.len()
        //        l_v = self.length

        // let tail_vec_size = (self.length - idx + 1) / 2;
        let half_idx = idx / 2;
        let mut tail = NibbleVec::from_small_vec(0);

        // Copy the bytes.
        for i in half_idx..self.data.len() {
            tail.data.push(self.data[i]);
        }

        // Pop the same bytes.
        for _ in half_idx..self.data.len() {
            self.data.pop();
        }

        // Update lengths.
        tail.length = self.length - idx;
        self.length = idx;

        tail
    }

    /// Copy data between the second half of self.data[start] and
    /// self.data[end - 1]. The second half of the last entry is included
    /// if include_last is true.
    #[inline]
    fn overlap_copy(
        &self,
        start: usize,
        end: usize,
        vec: &mut SmallVec<A>,
        length: &mut usize,
        include_last: bool,
    ) {
        // Copy up to the first half of the last byte.
        for i in start..(end - 1) {
            // The first half is the second half of the old entry.
            let first_half = self.data[i] & 0x0f;

            // The second half is the first half of the next entry.
            let second_half = self.data[i + 1] >> 4;

            vec.push((first_half << 4) | second_half);
            *length += 2;
        }

        if include_last {
            let last = self.data[end - 1] & 0x0f;
            vec.push(last << 4);
            *length += 1;
        }
    }

    /// Append another nibble vector whilst consuming this vector.
    #[inline]
    pub fn join(mut self, other: &NibbleVec<A>) -> NibbleVec<A> {
        // If the length is even, we can append directly.
        if self.length % 2 == 0 {
            self.length += other.length;
            self.data.extend_from_slice(&other.data);
            return self;
        }

        // If the other vector is empty, bail out.
        if other.is_empty() {
            return self;
        }

        // If the length is odd, we have to perform an overlap copy.
        // Copy the first half of the first element, to make the vector an even length.
        self.push(other.get(0));

        // Copy the rest of the vector using an overlap copy.
        let take_last = other.len() % 2 == 0;
        other.overlap_copy(
            0,
            other.data.len(),
            &mut self.data,
            &mut self.length,
            take_last,
        );

        self
    }
}

impl<A: Array<Item = u8>> PartialEq<NibbleVec<A>> for NibbleVec<A> {
    #[inline]
    fn eq(&self, other: &NibbleVec<A>) -> bool {
        self.length == other.length && self.data == other.data
    }
}

impl<A: Array<Item = u8>> Eq for NibbleVec<A> {}

/// Compare a `NibbleVec` and a slice of bytes *element-by-element*.
/// Bytes are **not** interpreted as two `NibbleVec` entries.
impl<A: Array<Item = u8>> PartialEq<[u8]> for NibbleVec<A> {
    #[inline]
    fn eq(&self, other: &[u8]) -> bool {
        if other.len() != self.len() {
            return false;
        }

        for (i, x) in other.iter().enumerate() {
            if self.get(i) != *x {
                return false;
            }
        }
        true
    }
}

impl<A: Array<Item = u8>> Debug for NibbleVec<A> {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "NibbleVec [")?;

        if !self.is_empty() {
            write!(fmt, "{}", self.get(0))?;
        }

        for i in 1..self.len() {
            write!(fmt, ", {}", self.get(i))?;
        }
        write!(fmt, "]")
    }
}

impl<A: Array<Item = u8>> From<Vec<u8>> for NibbleVec<A> {
    #[inline]
    fn from(v: Vec<u8>) -> NibbleVec<A> {
        NibbleVec::from_byte_vec(v)
    }
}

impl<'a, A: Array<Item = u8>> From<&'a [u8]> for NibbleVec<A> {
    #[inline]
    fn from(v: &[u8]) -> NibbleVec<A> {
        NibbleVec::from_byte_vec(v.into())
    }
}

impl<A: Array<Item = u8>> Into<Vec<u8>> for NibbleVec<A> {
    #[inline]
    fn into(self) -> Vec<u8> {
        self.data.to_vec()
    }
}

impl<'a, A: Array<Item = u8>> Into<Vec<u8>> for &'a NibbleVec<A> {
    #[inline]
    fn into(self) -> Vec<u8> {
        self.data.to_vec()
    }
}
<commit_msg>make note of removal of assert! for perf<commit_after>#[cfg(test)]
mod test;

use smallvec::{Array, SmallVec};

use std::convert::{From, Into};
use std::fmt::{self, Debug, Formatter};
use std::iter::FromIterator;

/// A `NibbleVec` backed by a `SmallVec` with 64 inline element slots.
/// This will not allocate until more than 64 elements are added.
pub type Nibblet = NibbleVec<[u8; 64]>;

/// A data-structure for storing a sequence of 4-bit values.
///
/// Values are stored in a `Vec<u8>`, with two values per byte.
///
/// Values at even indices are stored in the most-significant half of their byte,
/// while values at odd indices are stored in the least-significant half.
///
/// Imagine a vector of [MSB][msb-wiki] first bytes, and you'll be right.
///
/// n = [_ _ | _ _ | _ _]
///
/// [msb-wiki]: http://en.wikipedia.org/wiki/Most_significant_bit
#[derive(Clone, Default)]
pub struct NibbleVec<A: Array<Item = u8>> {
    length: usize,
    data: SmallVec<A>,
}

impl<A: Array<Item = u8>> NibbleVec<A> {
    /// Create an empty nibble vector.
    pub fn new() -> NibbleVec<A> {
        NibbleVec {
            length: 0,
            data: SmallVec::new(),
        }
    }

    /// Create a nibble vector from a vector of bytes.
    ///
    /// Each byte is split into two 4-bit entries (MSB, LSB).
    #[inline]
    pub fn from_byte_vec(vec: Vec<u8>) -> NibbleVec<A> {
        let length = 2 * vec.len();
        NibbleVec {
            length,
            data: SmallVec::from_iter(vec),
        }
    }

    /// Create a nibble vector from a `SmallVec` of Size.
    ///
    /// Each byte is split into two 4-bit entries (MSB, LSB).
    #[inline]
    pub fn from_small_vec(length: usize) -> NibbleVec<A> {
        NibbleVec {
            length,
            data: SmallVec::new(),
        }
    }

    /// Returns a byte slice of the nibble vector's contents.
    #[inline]
    pub fn as_bytes(&self) -> &[u8] {
        &self.data[..]
    }

    /// Converts a nibble vector into a byte vector.
    ///
    /// This consumes the nibble vector, so we do not need to copy its contents.
    #[inline]
    pub fn into_bytes(self) -> Vec<u8> {
        self.data.to_vec()
    }

    /// Get the number of elements stored in the vector.
    #[inline]
    pub fn len(&self) -> usize {
        self.length
    }

    /// Returns `true` if the nibble vector has a length of 0.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    /// Fetch a single entry from the vector.
    ///
    /// Guaranteed to be a value in the interval [0, 15].
    ///
    /// **Panics** if `idx >= self.len()`.
    #[inline]
    pub fn get(&self, idx: usize) -> u8 {
        if idx >= self.length {
            panic!(
                "NibbleVec index out of bounds: len is {}, index is {}",
                self.length, idx
            );
        }
        let vec_idx = idx / 2;
        match idx % 2 {
            // If the index is even, take the first (most significant) half of the stored byte.
            0 => self.data[vec_idx] >> 4,
            // If the index is odd, take the second (least significant) half.
            _ => self.data[vec_idx] & 0x0F,
        }
    }

    /// Add a single nibble to the vector.
    ///
    /// Only the 4 least-significant bits of the value are used.
    #[inline]
    pub fn push(&mut self, val: u8) {
        if self.length % 2 == 0 {
            self.data.push(val << 4);
        } else {
            let vec_len = self.data.len();

            // Zero the second half of the last byte just to be safe.
            self.data[vec_len - 1] &= 0xF0;

            // Write the new value.
            self.data[vec_len - 1] |= val & 0x0F;
        }
        self.length += 1;
    }

    /// Split the vector into two parts.
    ///
    /// All elements at or following the given index are returned in a new `NibbleVec`,
    /// with exactly `idx` elements remaining in this vector.
    ///
    /// **Panics** if `idx > self.len()`.
    pub fn split(&mut self, idx: usize) -> NibbleVec<A> {
        // assert! is a few percent slower surprisingly
        if idx > self.length {
            panic!(
                "attempted to split past vector end. len is {}, index is {}",
                self.length, idx
            );
        } else if idx == self.length {
            NibbleVec::new()
        } else if idx % 2 == 0 {
            self.split_even(idx)
        } else {
            self.split_odd(idx)
        }
    }

    /// Split function for odd *indices*.
    #[inline]
    fn split_odd(&mut self, idx: usize) -> NibbleVec<A> {
        // let tail_vec_size = (self.length - idx) / 2;
        let mut tail = NibbleVec::from_small_vec(0);

        // Perform an overlap copy, copying the last nibble of the original vector only if
        // the length of the new tail is *odd*.
        let tail_length = self.length - idx;
        let take_last = tail_length % 2 == 1;
        self.overlap_copy(
            idx / 2,
            self.data.len(),
            &mut tail.data,
            &mut tail.length,
            take_last,
        );

        // Remove the copied bytes, being careful to skip the idx byte.
        for _ in (idx / 2 + 1)..self.data.len() {
            self.data.pop();
        }

        // Zero the second half of the index byte so as to maintain the last-nibble invariant.
        self.data[idx / 2] &= 0xF0;

        // Update the length of the first NibbleVec.
        self.length = idx;

        tail
    }

    /// Split function for even *indices*.
    #[inline]
    fn split_even(&mut self, idx: usize) -> NibbleVec<A> {
        // Avoid allocating a temporary vector by copying all the bytes in order, then popping them.

        // Possible to prove: l_d - ⌊i / 2⌋ = ⌊(l_v - i + 1) / 2⌋
        //  where l_d = self.data.len()
        //        l_v = self.length

        // let tail_vec_size = (self.length - idx + 1) / 2;
        let half_idx = idx / 2;
        let mut tail = NibbleVec::from_small_vec(0);

        // Copy the bytes.
        for i in half_idx..self.data.len() {
            tail.data.push(self.data[i]);
        }

        // Pop the same bytes.
        for _ in half_idx..self.data.len() {
            self.data.pop();
        }

        // Update lengths.
        tail.length = self.length - idx;
        self.length = idx;

        tail
    }

    /// Copy data between the second half of self.data[start] and
    /// self.data[end - 1]. The second half of the last entry is included
    /// if include_last is true.
    #[inline]
    fn overlap_copy(
        &self,
        start: usize,
        end: usize,
        vec: &mut SmallVec<A>,
        length: &mut usize,
        include_last: bool,
    ) {
        // Copy up to the first half of the last byte.
        for i in start..(end - 1) {
            // The first half is the second half of the old entry.
            let first_half = self.data[i] & 0x0f;

            // The second half is the first half of the next entry.
            let second_half = self.data[i + 1] >> 4;

            vec.push((first_half << 4) | second_half);
            *length += 2;
        }

        if include_last {
            let last = self.data[end - 1] & 0x0f;
            vec.push(last << 4);
            *length += 1;
        }
    }

    /// Append another nibble vector whilst consuming this vector.
    #[inline]
    pub fn join(mut self, other: &NibbleVec<A>) -> NibbleVec<A> {
        // If the length is even, we can append directly.
        if self.length % 2 == 0 {
            self.length += other.length;
            self.data.extend_from_slice(&other.data);
            return self;
        }

        // If the other vector is empty, bail out.
        if other.is_empty() {
            return self;
        }

        // If the length is odd, we have to perform an overlap copy.
        // Copy the first half of the first element, to make the vector an even length.
        self.push(other.get(0));

        // Copy the rest of the vector using an overlap copy.
        let take_last = other.len() % 2 == 0;
        other.overlap_copy(
            0,
            other.data.len(),
            &mut self.data,
            &mut self.length,
            take_last,
        );

        self
    }
}

impl<A: Array<Item = u8>> PartialEq<NibbleVec<A>> for NibbleVec<A> {
    #[inline]
    fn eq(&self, other: &NibbleVec<A>) -> bool {
        self.length == other.length && self.data == other.data
    }
}

impl<A: Array<Item = u8>> Eq for NibbleVec<A> {}

/// Compare a `NibbleVec` and a slice of bytes *element-by-element*.
/// Bytes are **not** interpreted as two `NibbleVec` entries.
impl<A: Array<Item = u8>> PartialEq<[u8]> for NibbleVec<A> {
    #[inline]
    fn eq(&self, other: &[u8]) -> bool {
        if other.len() != self.len() {
            return false;
        }

        for (i, x) in other.iter().enumerate() {
            if self.get(i) != *x {
                return false;
            }
        }
        true
    }
}

impl<A: Array<Item = u8>> Debug for NibbleVec<A> {
    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "NibbleVec [")?;

        if !self.is_empty() {
            write!(fmt, "{}", self.get(0))?;
        }

        for i in 1..self.len() {
            write!(fmt, ", {}", self.get(i))?;
        }
        write!(fmt, "]")
    }
}

impl<A: Array<Item = u8>> From<Vec<u8>> for NibbleVec<A> {
    #[inline]
    fn from(v: Vec<u8>) -> NibbleVec<A> {
        NibbleVec::from_byte_vec(v)
    }
}

impl<'a, A: Array<Item = u8>> From<&'a [u8]> for NibbleVec<A> {
    #[inline]
    fn from(v: &[u8]) -> NibbleVec<A> {
        NibbleVec::from_byte_vec(v.into())
    }
}

impl<A: Array<Item = u8>> Into<Vec<u8>> for NibbleVec<A> {
    #[inline]
    fn into(self) -> Vec<u8> {
        self.data.to_vec()
    }
}

impl<'a, A: Array<Item = u8>> Into<Vec<u8>> for &'a NibbleVec<A> {
    #[inline]
    fn into(self) -> Vec<u8> {
        self.data.to_vec()
    }
}
<|endoftext|>


<commit_before>//! proxer api library
#[macro_use]
extern crate serde_derive;
extern crate chrono;
extern crate reqwest;
extern crate serde_json;
extern crate serde;
extern crate serde_urlencoded;
#[macro_use]
extern crate log;
#[macro_use]
extern crate hyper;

pub mod error;
pub mod response;
pub mod api;
pub mod prelude;
pub mod client;
pub mod tests;

pub use client::Client;
pub use prelude::*;
use serde::Serialize;
use serde::de::DeserializeOwned;

/// Every struct that is an endpoint, implements this trait.
pub trait Endpoint: Serialize {
	type ResponseType: std::fmt::Debug + Clone + DeserializeOwned;
	const URL: &'static str;
}

pub trait PageableEndpoint
where
	Self: Endpoint + std::marker::Sized,
	<Self as Endpoint>::ResponseType: IntoIterator + Clone,
{
	fn pager(self, client: Client) -> Pager<Self>;

	fn page_mut(&mut self) -> &mut Option<usize>;
	fn limit_mut(&mut self) -> &mut Option<usize>;
}

//#[derive(Debug, Clone)]
pub struct Pager<T>
where
	T: Endpoint,
	<T as Endpoint>::ResponseType: IntoIterator + Clone + std::fmt::Debug,
{
	client: Client,
	shifted: usize,
	endpoint: T,
	data: Vec<<<T as Endpoint>::ResponseType as IntoIterator>::Item>,
}

impl<T> Pager<T>
where
	T: Endpoint + PageableEndpoint,
	T: Clone + std::fmt::Debug,
	<T as Endpoint>::ResponseType: IntoIterator + Clone + std::fmt::Debug,
{
	pub fn new(client: Client, mut endpoint: T, start: Option<usize>, limit: Option<usize>)
		-> Self
	{

		match (endpoint.page_mut(), start)
		{
			(&mut None, None) => *endpoint.page_mut() = Some(0),
			(&mut None, Some(_)) => *endpoint.page_mut() = start,
			_ => {}
		}

		match (endpoint.limit_mut(), limit)
		{
			(&mut None, None) => *endpoint.limit_mut() = Some(750),
			(&mut None, Some(_)) => *endpoint.limit_mut() = limit,
			_ => {}
		}

		debug!(
			"initialize new pager: page: {}, limit {}",
			endpoint.page_mut().unwrap(),
			endpoint.limit_mut().unwrap(),
		);

		Self {
			client: client,
			data: Vec::new(),
			shifted: 0,
			endpoint: endpoint,
		}
	}

	fn page_mut(&mut self) -> usize
	{
		self.endpoint.page_mut().unwrap()
	}

	fn limit_mut(&mut self) -> usize
	{
		self.endpoint.limit_mut().unwrap()
	}
}

impl<T> Iterator for Pager<T>
where
	T: Endpoint,
	T: PageableEndpoint + Clone + std::fmt::Debug,
	<T as Endpoint>::ResponseType: IntoIterator,
{
	type Item = Result<<<T as Endpoint>::ResponseType as IntoIterator>::Item, error::Error>;

	fn next(&mut self) -> Option<Self::Item>
	{
		match self.data.pop()
		{
			Some(i) => {
				debug!("returning from buffer");
				self.shifted += 1;

				debug!("buffer: remaining/shifted: {}/{}",
					self.data.len(),
					self.shifted,
				);
				Some(Ok(i))
			}
			None => {
				debug!("expected shifted: {}", self.page_mut() * self.limit_mut());
				debug!("actually shifted: {}", self.shifted);

				if self.page_mut() * self.limit_mut() > self.shifted {
					debug!("reached the last page");
					return None;
				}

				debug!("fetching new data");
				let res = self.client.execute(self.endpoint.clone()).unwrap();

				for var in res.into_iter() {
					self.data.push(var);
				}

				debug!("filled buffer with {} entries", self.data.len());

				*self.endpoint.page_mut() = Some(self.page_mut() + 1);

				self.next()

			}
		}
	}
}
<commit_msg>improve error handling for pageable endpoints<commit_after>//! proxer api library
#[macro_use]
extern crate serde_derive;
extern crate chrono;
extern crate reqwest;
extern crate serde_json;
extern crate serde;
extern crate serde_urlencoded;
#[macro_use]
extern crate log;
#[macro_use]
extern crate hyper;

pub mod error;
pub mod response;
pub mod api;
pub mod prelude;
pub mod client;
pub mod tests;

pub use client::Client;
pub use prelude::*;
use serde::Serialize;
use serde::de::DeserializeOwned;

/// Every struct that is an endpoint, implements this trait.
pub trait Endpoint: Serialize {
	type ResponseType: std::fmt::Debug + Clone + DeserializeOwned;
	const URL: &'static str;
}

pub trait PageableEndpoint
where
	Self: Endpoint + std::marker::Sized,
	<Self as Endpoint>::ResponseType: IntoIterator + Clone,
{
	fn pager(self, client: Client) -> Pager<Self>;

	fn page_mut(&mut self) -> &mut Option<usize>;
	fn limit_mut(&mut self) -> &mut Option<usize>;
}

//#[derive(Debug, Clone)]
pub struct Pager<T>
where
	T: Endpoint,
	<T as Endpoint>::ResponseType: IntoIterator + Clone + std::fmt::Debug,
{
	client: Client,
	shifted: usize,
	endpoint: T,
	data: Vec<<<T as Endpoint>::ResponseType as IntoIterator>::Item>,
}

impl<T> Pager<T>
where
	T: Endpoint + PageableEndpoint,
	T: Clone + std::fmt::Debug,
	<T as Endpoint>::ResponseType: IntoIterator + Clone + std::fmt::Debug,
{
	pub fn new(client: Client, mut endpoint: T, start: Option<usize>, limit: Option<usize>)
		-> Self
	{

		match (endpoint.page_mut(), start)
		{
			(&mut None, None) => *endpoint.page_mut() = Some(0),
			(&mut None, Some(_)) => *endpoint.page_mut() = start,
			_ => {}
		}

		match (endpoint.limit_mut(), limit)
		{
			(&mut None, None) => *endpoint.limit_mut() = Some(750),
			(&mut None, Some(_)) => *endpoint.limit_mut() = limit,
			_ => {}
		}

		debug!(
			"initialize new pager: page: {}, limit {}",
			endpoint.page_mut().unwrap(),
			endpoint.limit_mut().unwrap(),
		);

		Self {
			client: client,
			data: Vec::new(),
			shifted: 0,
			endpoint: endpoint,
		}
	}

	fn page_mut(&mut self) -> usize
	{
		self.endpoint.page_mut().unwrap()
	}

	fn limit_mut(&mut self) -> usize
	{
		self.endpoint.limit_mut().unwrap()
	}

	fn next_page(&mut self) {
		*self.endpoint.page_mut() = Some(self.page_mut() + 1);
	}
}

impl<T> Iterator for Pager<T>
where
	T: Endpoint,
	T: PageableEndpoint + Clone + std::fmt::Debug,
	<T as Endpoint>::ResponseType: IntoIterator,
{
	type Item = Result<<<T as Endpoint>::ResponseType as IntoIterator>::Item, error::Error>;

	fn next(&mut self) -> Option<Self::Item>
	{
		match self.data.pop()
		{
			Some(i) => {
				debug!("returning from buffer");
				self.shifted += 1;

				debug!("buffer: remaining/shifted: {}/{}",
					self.data.len(),
					self.shifted,
				);
				Some(Ok(i))
			}
			None => {
				debug!("expected shifted: {}", self.page_mut() * self.limit_mut());
				debug!("actually shifted: {}", self.shifted);

				if self.page_mut() * self.limit_mut() > self.shifted {
					debug!("reached the last page");
					return None;
				}

				debug!("fetching new data");
				let res = self.client
					.execute(self.endpoint.clone());

				self.next_page();

				match res {
					Ok(res) => {
						for var in res.into_iter() {
							self.data.push(var);
						}
						debug!("filled buffer with {} entries", self.data.len());

						self.next()
					},
					Err(e) => {
						Some(Err(e))
					}
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>#![no_std]

const PRIME32_1: u32 = 2_654_435_761;
const PRIME32_2: u32 = 2_246_822_519;
const PRIME32_3: u32 = 3_266_489_917;
const PRIME32_4: u32 = 668_265_263;
const PRIME32_5: u32 = 374_761_393;

pub struct XXH32 {
    seed: u32,
    total_len: usize,
    v1: u32,
    v2: u32,
    v3: u32,
    v4: u32,
    memory: [u8; 16],
    memsize: usize,
}

fn calc_next_chunk(val: u32, bytes: &[u8]) -> u32 {
    u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
        .wrapping_mul(PRIME32_2)
        .wrapping_add(val)
        .rotate_left(13)
        .wrapping_mul(PRIME32_1)
}

impl Default for XXH32 {
    fn default() -> XXH32 {
        XXH32::new_with_seed(0)
    }
}

impl XXH32 {
    pub fn new_with_seed(seed: u32) -> XXH32 {
        XXH32 {
            seed,
            total_len: 0,
            v1: seed.wrapping_add(PRIME32_1).wrapping_add(PRIME32_2),
            v2: seed.wrapping_add(PRIME32_2),
            v3: seed,
            v4: seed.wrapping_sub(PRIME32_1),
            memory: [0; 16],
            memsize: 0,
        }
    }
    pub fn write(&mut self, bytes: &[u8]) {
        self.total_len += bytes.len();

        if self.memsize + bytes.len() < 16 {
            self.memory[self.memsize..self.memsize + bytes.len()].copy_from_slice(bytes);
            self.memsize += bytes.len();
            return;
        }

        let bytesview = &bytes[(16 - self.memsize) % 16..];

        if self.memsize > 0 {
            self.memory[self.memsize..].copy_from_slice(&bytes[..16 - self.memsize]);

            self.v1 = calc_next_chunk(self.v1, &self.memory[0..4]);
            self.v2 = calc_next_chunk(self.v2, &self.memory[4..8]);
            self.v3 = calc_next_chunk(self.v3, &self.memory[8..12]);
            self.v4 = calc_next_chunk(self.v4, &self.memory[12..16]);

            self.memsize = 0;
        }

        let mut iter = bytesview.chunks_exact(16);
        for chunk in iter.by_ref() {
            self.v1 = calc_next_chunk(self.v1, &chunk[0..4]);
            self.v2 = calc_next_chunk(self.v2, &chunk[4..8]);
            self.v3 = calc_next_chunk(self.v3, &chunk[8..12]);
            self.v4 = calc_next_chunk(self.v4, &chunk[12..16]);
        }

        let bytesview = iter.remainder();

        self.memory[..bytesview.len()].copy_from_slice(bytesview);
        self.memsize += bytesview.len();
    }

    pub fn finish(&self) -> u32 {
        let mut h32 = if self.total_len >= 16 {
            self.v1
                .rotate_left(1)
                .wrapping_add(self.v2.rotate_left(7))
                .wrapping_add(self.v3.rotate_left(12))
                .wrapping_add(self.v4.rotate_left(18))
        } else {
            self.seed.wrapping_add(PRIME32_5)
        };

        h32 = h32.wrapping_add(self.total_len as u32);

        let mut iter = self.memory[..self.memsize].chunks_exact(4);
        for chunk in iter.by_ref() {
            h32 = h32.wrapping_add(
                u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])
                    .wrapping_mul(PRIME32_3),
            );
            h32 = h32.rotate_left(17).wrapping_mul(PRIME32_4);
        }

        for byte in iter.remainder() {
            h32 = h32.wrapping_add(u32::from(*byte).wrapping_mul(PRIME32_5));
            h32 = h32.rotate_left(11).wrapping_mul(PRIME32_1);
        }

        h32 = (h32 ^ (h32 >> 15)).wrapping_mul(PRIME32_2);
        h32 = (h32 ^ (h32 >> 13)).wrapping_mul(PRIME32_3);
        h32 ^ (h32 >> 16)
    }
}
<commit_msg>Move operations on chunks into methods.<commit_after>#![no_std]

const PRIME32_1: u32 = 2_654_435_761;
const PRIME32_2: u32 = 2_246_822_519;
const PRIME32_3: u32 = 3_266_489_917;
const PRIME32_4: u32 = 668_265_263;
const PRIME32_5: u32 = 374_761_393;

pub struct XXH32 {
    seed: u32,
    total_len: usize,
    v1: u32,
    v2: u32,
    v3: u32,
    v4: u32,
    memory: [u8; 16],
    memsize: usize,
}

fn calc_next_chunk(val: u32, bytes: &[u8]) -> u32 {
    u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
        .wrapping_mul(PRIME32_2)
        .wrapping_add(val)
        .rotate_left(13)
        .wrapping_mul(PRIME32_1)
}

impl Default for XXH32 {
    fn default() -> XXH32 {
        XXH32::new_with_seed(0)
    }
}

impl XXH32 {
    pub fn new_with_seed(seed: u32) -> XXH32 {
        XXH32 {
            seed,
            total_len: 0,
            v1: seed.wrapping_add(PRIME32_1).wrapping_add(PRIME32_2),
            v2: seed.wrapping_add(PRIME32_2),
            v3: seed,
            v4: seed.wrapping_sub(PRIME32_1),
            memory: [0; 16],
            memsize: 0,
        }
    }
    fn calc_chunk(&mut self, bytes: &[u8]) {
        self.v1 = calc_next_chunk(self.v1, &bytes[0..4]);
        self.v2 = calc_next_chunk(self.v2, &bytes[4..8]);
        self.v3 = calc_next_chunk(self.v3, &bytes[8..12]);
        self.v4 = calc_next_chunk(self.v4, &bytes[12..16]);
    }
    fn calc_chunk_memory(&mut self) {
        self.v1 = calc_next_chunk(self.v1, &self.memory[0..4]);
        self.v2 = calc_next_chunk(self.v2, &self.memory[4..8]);
        self.v3 = calc_next_chunk(self.v3, &self.memory[8..12]);
        self.v4 = calc_next_chunk(self.v4, &self.memory[12..16]);
    }
    pub fn write(&mut self, bytes: &[u8]) {
        self.total_len += bytes.len();

        if self.memsize + bytes.len() < 16 {
            self.memory[self.memsize..self.memsize + bytes.len()].copy_from_slice(bytes);
            self.memsize += bytes.len();
            return;
        }

        let bytesview = &bytes[(16 - self.memsize) % 16..];

        if self.memsize > 0 {
            self.memory[self.memsize..].copy_from_slice(&bytes[..16 - self.memsize]);
            self.calc_chunk_memory();
            self.memsize = 0;
        }

        let mut iter = bytesview.chunks_exact(16);
        for chunk in iter.by_ref() {
            self.calc_chunk(&chunk);
        }

        let bytesview = iter.remainder();

        self.memory[..bytesview.len()].copy_from_slice(bytesview);
        self.memsize += bytesview.len();
    }

    pub fn finish(&self) -> u32 {
        let mut h32 = if self.total_len >= 16 {
            self.v1
                .rotate_left(1)
                .wrapping_add(self.v2.rotate_left(7))
                .wrapping_add(self.v3.rotate_left(12))
                .wrapping_add(self.v4.rotate_left(18))
        } else {
            self.seed.wrapping_add(PRIME32_5)
        };

        h32 = h32.wrapping_add(self.total_len as u32);

        let mut iter = self.memory[..self.memsize].chunks_exact(4);
        for chunk in iter.by_ref() {
            h32 = h32.wrapping_add(
                u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])
                    .wrapping_mul(PRIME32_3),
            );
            h32 = h32.rotate_left(17).wrapping_mul(PRIME32_4);
        }

        for byte in iter.remainder() {
            h32 = h32.wrapping_add(u32::from(*byte).wrapping_mul(PRIME32_5));
            h32 = h32.rotate_left(11).wrapping_mul(PRIME32_1);
        }

        h32 = (h32 ^ (h32 >> 15)).wrapping_mul(PRIME32_2);
        h32 = (h32 ^ (h32 >> 13)).wrapping_mul(PRIME32_3);
        h32 ^ (h32 >> 16)
    }
}
<|endoftext|>


<commit_before>extern crate prettytable;
use prettytable::Table;
use prettytable::row::Row;
use prettytable::cell::Cell;

use std::ops::{Index, IndexMut};
use self::Direction::{Down, Left, Right, Up};
use self::Error::*;

#[cfg(test)]
mod test;
// Possible directions for cycle
#[derive(Clone,Copy,PartialEq,Debug)]
enum Direction {
    Up,
    Down,
    Left,
    Right,
    None,
}
// The structure for emulation a two-dimensional array
struct Matrix<T> {
    cols: usize,
    data: Vec<T>,
}

impl<T> Index<usize> for Matrix<T> {
    type Output = [T];

    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        let i = index * self.cols;
        &self.data[i..i + self.cols]
    }
}

impl<T> IndexMut<usize> for Matrix<T> {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        let i = index * self.cols;
        &mut self.data[i..i + self.cols]
    }
}

impl<T> Matrix<T> {
    fn new(columns: usize) -> Matrix<T> {
        Matrix {
            cols: columns,
            data: Vec::new(),
        }
    }
    // Add a line to the matrix
    fn push(&mut self, mut vector: Vec<T>) {
        assert_eq!(self.cols, vector.len());
        self.data.append(&mut vector);
    }
    // The number of rows in the matrix
    fn rows(&self) -> usize {
        self.data.len() / self.cols
    }
    // The number of columns in the matrix
    fn cols(&self) -> usize {
        self.cols
    }
}

impl Matrix<Direction> {
    fn iter<'a>(&'a self, i: usize, j: usize) -> MatrixIter<'a> {
        MatrixIter {
            state: self,
            start_i: i,
            start_j: j,
            i: i,
            j: j,
            f: false,
        }
    }
}

struct MatrixIter<'a> {
    state: &'a Matrix<Direction>,
    start_i: usize,
    start_j: usize,
    i: usize,
    j: usize,
    f: bool,
}

impl<'a> MatrixIter<'a> {
    fn move_by(&mut self, d: Direction) {
        match d {
            Up => self.i += 1,
            Down => self.i -= 1,
            Left => self.j += 1,
            Right => self.j -= 1,
            Direction::None => unreachable!(),
        }
    }
}

impl<'a> Iterator for MatrixIter<'a> {
    type Item = (usize,usize);
    fn next(&mut self) -> Option<Self::Item> {
        if self.i == self.start_i && self.j == self.start_j {
            if self.f {
                return None;
            }
            self.f = true;
        }

        let buf = (self.i, self.j);
        let d = self.state[self.i][self.j];
        while self.state[self.i][self.j] == d || self.state[self.i][self.j] == Direction::None {
            self.move_by(d);
        }
        Some(buf)
    }
}

// Initialization errors
#[derive(Debug)]
pub enum Error {
    NumOfSupOrCust,
    NumOfRows,
    NumOfCols,
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use std::error::Error;
        f.write_str(self.description())
    }
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            NumOfSupOrCust => "invalid number of suppliers or customers",
            NumOfRows => "invalid number of rows in the matrix of costs",
            NumOfCols => "invalid number of columns in the matrix of costs",
        }
    }
}

// The main structure
pub struct Transportation {
    supply: Vec<u64>,
    demand: Vec<u64>,
    cost: Matrix<u64>,
    trans: Matrix<Option<u64>>,
}

impl Transportation {
    fn least_cost_method(&mut self) {
        let mut a = self.supply.to_vec();
        let mut b = self.demand.to_vec();
        let mut min;

        loop {
            min = None;

            for i in 0..self.cost.rows() {
                if a[i] == 0 {
                    continue;
                }
                for j in 0..self.cost.cols() {
                    if b[j] == 0 {
                        continue;
                    }
                    if let Some((_, _, v)) = min {
                        if v > self.cost[i][j] {
                            min = Some((i, j, self.cost[i][j]));
                        }
                    } else {
                        min = Some((i, j, self.cost[i][j]));
                    }
                }
            }

            if let Some((i, j, _)) = min {
                if a[i] >= b[j] {
                    self.trans[i][j] = Some(b[j]);
                    a[i] -= b[j];
                    b[j] = 0;
                } else {
                    self.trans[i][j] = Some(a[i]);
                    b[j] -= a[i];
                    a[i] = 0;
                }
            } else {
                return;
            }
        }
    }
    // Detect cycle via dfs and build a matrix of directions
    // If the direction of the trans_state[i][j] isn't None, the cycle is found
    fn cycle_detection(&self, i: usize, j: usize) -> Matrix<Direction> {
        let mut trans_state = Matrix::new(self.trans.cols());
        for _ in 0..self.trans.rows() {
            trans_state.push(vec![Direction::None;self.trans.cols()]);
        }

        fn _cycle_detection(x: &Matrix<Option<u64>>, x_state: &mut Matrix<Direction>, i: usize, j: usize) {
            if x_state[i][j] != Down {
                for k in (0..i).rev() {
                    if x[k][j] != None {
                        if x_state[k][j] == Direction::None {
                            x_state[k][j] = Up;
                            _cycle_detection(x, x_state, k, j);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Up {
                for k in i + 1..x.rows() {
                    if x[k][j] != None {
                        if x_state[k][j] == Direction::None {
                            x_state[k][j] = Down;
                            _cycle_detection(x, x_state, k, j);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Right {
                for k in (0..j).rev() {
                    if x[i][k] != None {
                        if x_state[i][k] == Direction::None {
                            x_state[i][k] = Left;
                            _cycle_detection(x, x_state, i, k);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Left {
                for k in j + 1..x.cols() {
                    if x[i][k] != None {
                        if x_state[i][k] == Direction::None {
                            x_state[i][k] = Right;
                            _cycle_detection(x, x_state, i, k);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
        }

        _cycle_detection(&self.trans, &mut trans_state, i, j);

        trans_state
    }
    // If the number of basic cells in the transportation plan is less then
    // m + n - 1, it must be replenished
    fn replenish(&mut self) {
        let number = self.trans.data.iter().filter(|&&a| a != None).count();

        if number < self.trans.rows() + self.trans.cols() - 1 {
            for i in 0..self.trans.rows() {
                for j in 0..self.trans.cols() {
                    if self.trans[i][j] == None {
                        self.trans[i][j] = Some(0);
                        if self.cycle_detection(i, j)[i][j] != Direction::None {
                            self.trans[i][j] = None;
                        } else {
                            return;
                        }
                    }
                }
            }
        }
    }
    // Recursive calculation of potentials
    fn calculation_of_potentials(&self) -> (Vec<i64>, Vec<i64>) {
        let mut u = vec![0;self.trans.rows()];
        let mut v = vec![0;self.trans.cols()];

        fn calculation_of_potentials_h(c: &Matrix<u64>, x: &Matrix<Option<u64>>, u: &mut [i64], v: &mut [i64], i: usize, j: usize) {
            for k in 0..x.cols() {
                if x[i][k] != None {
                    v[k] = c[i][k] as i64 - u[i];
                    if k != j {
                        calculation_of_potentials_v(c, x, u, v, i, k);
                    }
                }
            }
        }

        fn calculation_of_potentials_v(c: &Matrix<u64>, x: &Matrix<Option<u64>>, u: &mut [i64], v: &mut [i64], i: usize, j: usize) {
            for k in 0..x.rows() {
                if x[k][j] != None {
                    u[k] = c[k][j] as i64 - v[j];
                    if k != i {
                        calculation_of_potentials_h(c, x, u, v, k, j);
                    }
                }
            }
        }

        calculation_of_potentials_h(&self.cost, &self.trans, &mut u, &mut v, 0, 0);
        calculation_of_potentials_v(&self.cost, &self.trans, &mut u, &mut v, 0, 0);

        (u, v)
    }
    // Check for optimality
    // If the transportation plan is not optimal, then return the minimum of difference
    fn check(&self) -> Option<(usize, usize, i64)> {
        let mut min = None;
        let (u, v) = self.calculation_of_potentials();

        for i in 0..self.cost.rows() {
            for j in 0..self.cost.cols() {
                let d = self.cost[i][j] as i64 - u[i] - v[j];
                if d < 0 {
                    if let Some((_, _, m)) = min {
                        if m > d {
                            min = Some((i, j, d));
                        }
                    } else {
                        min = Some((i, j, d));
                    }
                }
            }
        }
        min
    }
    // Calculate the total cost
    fn total_cost(&self) -> u64 {
        let mut z = 0;
        for (t, c) in self.trans.data.iter().zip(self.cost.data.iter()) {
            if !t.is_none() {
                z += t.unwrap() * c;
            }
        }
        z
    }

    pub fn potential_method(&mut self) {
        self.least_cost_method();
        self.replenish();

        loop {
            if let Some((i, j, _)) = self.check() {
                self.trans[i][j] = Some(0);
                let trans_state = self.cycle_detection(i, j);

                let mut min = std::u64::MAX;
                let mut f = true;
                for x in trans_state.iter(i, j) {
                    if f {
                        f = false;
                    } else {
                        if min > self.trans[x.0][x.1].unwrap() {
                            min = self.trans[x.0][x.1].unwrap();
                        }
                        f = true;
                    }
                }

                f = true;
                for x in trans_state.iter(i, j) {
                    if f {
                        self.trans[x.0][x.1] = Some(self.trans[x.0][x.1].unwrap() + min);
                        f = false;
                    } else {
                        self.trans[x.0][x.1] = Some(self.trans[x.0][x.1].unwrap() - min);
                        f = true;
                    }
                }

                let mut max = (0, 0, 0);
                for x in trans_state.iter(i, j) {
                    if self.trans[x.0][x.1].unwrap() == 0 && max.2 <= self.cost[x.0][x.1] {
                        max = (x.0, x.1, self.cost[x.0][x.1]);
                    }
                }
                self.trans[max.0][max.1] = None;
            } else {
                return; // the plan is optimal
            }
        }
    }

    pub fn printstd(&self) {
        let mut table = Table::new();
        table.add_row(Row::new(vec![Cell::new("")]));
        for i in 0..self.demand.len() {
            table[0].add_cell(Cell::new(&format!("B{}", i + 1)));
        }
        table[0].add_cell(Cell::new("Запасы"));
        for (i, s) in self.supply.iter().enumerate() {
            table.add_row(Row::new(vec![Cell::new(&format!("A{}", i + 1))]));
            for m in self.trans[i].iter() {
                if m.is_none() {
                    table[i + 1].add_cell(Cell::new("0"));
                } else {
                    table[i + 1].add_cell(Cell::new(&format!("{}", m.unwrap())));
                }
            }
            table[i + 1].add_cell(Cell::new(&format!("{}", s)));
        }
        table.add_row(Row::new(vec![Cell::new("Потребности")]));
        for d in self.demand.iter() {
            table[self.supply.len() + 1].add_cell(Cell::new(&format!("{}", d)));
        }
        table.add_row(Row::new(vec![Cell::new("Общая стоимость"), Cell::new(&format!("{}", self.total_cost()))]));
        table.printstd();
    }

    pub fn new(mut a: Vec<u64>, mut b: Vec<u64>, mut c: Vec<Vec<u64>>) -> Result<Transportation, Error> {
        if a.is_empty() || b.is_empty() {
            return Err(NumOfSupOrCust);
        }
        if a.len() != c.len() {
            return Err(NumOfRows);
        }
        for i in &c {
            if b.len() != i.len() {
                return Err(NumOfCols);
            }
        }

        let sum_s = a.iter().fold(0, std::ops::Add::add); //a.iter().sum();
        let sum_d = b.iter().fold(0, std::ops::Add::add); //b.iter().sum();
        if sum_s > sum_d {
            b.push(sum_s - sum_d);
            for i in &mut c {
                i.push(0);
            }
        }
        if sum_s < sum_d {
            a.push(sum_d - sum_s);
            c.push(vec![0;b.len()]);
        }
        let mut cost = Matrix::new(c[0].len());
        for i in c {
            cost.push(i);
        }

        Ok(Transportation {
            trans: Matrix {
                cols: b.len(),
                data: vec![None;a.len()*b.len()],
            },
            supply: a,
            demand: b,
            cost: cost,
        })
    }
}
<commit_msg>use while instead loop<commit_after>extern crate prettytable;
use prettytable::Table;
use prettytable::row::Row;
use prettytable::cell::Cell;

use std::ops::{Index, IndexMut};
use self::Direction::{Down, Left, Right, Up};
use self::Error::*;

#[cfg(test)]
mod test;
// Possible directions for cycle
#[derive(Clone,Copy,PartialEq,Debug)]
enum Direction {
    Up,
    Down,
    Left,
    Right,
    None,
}
// The structure for emulation a two-dimensional array
struct Matrix<T> {
    cols: usize,
    data: Vec<T>,
}

impl<T> Index<usize> for Matrix<T> {
    type Output = [T];

    #[inline]
    fn index(&self, index: usize) -> &Self::Output {
        let i = index * self.cols;
        &self.data[i..i + self.cols]
    }
}

impl<T> IndexMut<usize> for Matrix<T> {
    #[inline]
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        let i = index * self.cols;
        &mut self.data[i..i + self.cols]
    }
}

impl<T> Matrix<T> {
    fn new(columns: usize) -> Matrix<T> {
        Matrix {
            cols: columns,
            data: Vec::new(),
        }
    }
    // Add a line to the matrix
    fn push(&mut self, mut vector: Vec<T>) {
        assert_eq!(self.cols, vector.len());
        self.data.append(&mut vector);
    }
    // The number of rows in the matrix
    fn rows(&self) -> usize {
        self.data.len() / self.cols
    }
    // The number of columns in the matrix
    fn cols(&self) -> usize {
        self.cols
    }
}

impl Matrix<Direction> {
    fn iter<'a>(&'a self, i: usize, j: usize) -> MatrixIter<'a> {
        MatrixIter {
            state: self,
            start_i: i,
            start_j: j,
            i: i,
            j: j,
            f: false,
        }
    }
}

struct MatrixIter<'a> {
    state: &'a Matrix<Direction>,
    start_i: usize,
    start_j: usize,
    i: usize,
    j: usize,
    f: bool,
}

impl<'a> MatrixIter<'a> {
    fn move_by(&mut self, d: Direction) {
        match d {
            Up => self.i += 1,
            Down => self.i -= 1,
            Left => self.j += 1,
            Right => self.j -= 1,
            Direction::None => unreachable!(),
        }
    }
}

impl<'a> Iterator for MatrixIter<'a> {
    type Item = (usize,usize);
    fn next(&mut self) -> Option<Self::Item> {
        if self.i == self.start_i && self.j == self.start_j {
            if self.f {
                return None;
            }
            self.f = true;
        }

        let buf = (self.i, self.j);
        let d = self.state[self.i][self.j];
        while self.state[self.i][self.j] == d || self.state[self.i][self.j] == Direction::None {
            self.move_by(d);
        }
        Some(buf)
    }
}

// Initialization errors
#[derive(Debug)]
pub enum Error {
    NumOfSupOrCust,
    NumOfRows,
    NumOfCols,
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        use std::error::Error;
        f.write_str(self.description())
    }
}

impl std::error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            NumOfSupOrCust => "invalid number of suppliers or customers",
            NumOfRows => "invalid number of rows in the matrix of costs",
            NumOfCols => "invalid number of columns in the matrix of costs",
        }
    }
}

// The main structure
pub struct Transportation {
    supply: Vec<u64>,
    demand: Vec<u64>,
    cost: Matrix<u64>,
    trans: Matrix<Option<u64>>,
}

impl Transportation {
    fn least_cost_method(&mut self) {
        let mut a = self.supply.to_vec();
        let mut b = self.demand.to_vec();
        let mut min;

        loop {
            min = None;

            for i in 0..self.cost.rows() {
                if a[i] == 0 {
                    continue;
                }
                for j in 0..self.cost.cols() {
                    if b[j] == 0 {
                        continue;
                    }
                    if let Some((_, _, v)) = min {
                        if v > self.cost[i][j] {
                            min = Some((i, j, self.cost[i][j]));
                        }
                    } else {
                        min = Some((i, j, self.cost[i][j]));
                    }
                }
            }

            if let Some((i, j, _)) = min {
                if a[i] >= b[j] {
                    self.trans[i][j] = Some(b[j]);
                    a[i] -= b[j];
                    b[j] = 0;
                } else {
                    self.trans[i][j] = Some(a[i]);
                    b[j] -= a[i];
                    a[i] = 0;
                }
            } else {
                return;
            }
        }
    }
    // Detect cycle via dfs and build a matrix of directions
    // If the direction of the trans_state[i][j] isn't None, the cycle is found
    fn cycle_detection(&self, i: usize, j: usize) -> Matrix<Direction> {
        let mut trans_state = Matrix::new(self.trans.cols());
        for _ in 0..self.trans.rows() {
            trans_state.push(vec![Direction::None;self.trans.cols()]);
        }

        fn _cycle_detection(x: &Matrix<Option<u64>>, x_state: &mut Matrix<Direction>, i: usize, j: usize) {
            if x_state[i][j] != Down {
                for k in (0..i).rev() {
                    if x[k][j] != None {
                        if x_state[k][j] == Direction::None {
                            x_state[k][j] = Up;
                            _cycle_detection(x, x_state, k, j);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Up {
                for k in i + 1..x.rows() {
                    if x[k][j] != None {
                        if x_state[k][j] == Direction::None {
                            x_state[k][j] = Down;
                            _cycle_detection(x, x_state, k, j);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Right {
                for k in (0..j).rev() {
                    if x[i][k] != None {
                        if x_state[i][k] == Direction::None {
                            x_state[i][k] = Left;
                            _cycle_detection(x, x_state, i, k);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
            if x_state[i][j] != Left {
                for k in j + 1..x.cols() {
                    if x[i][k] != None {
                        if x_state[i][k] == Direction::None {
                            x_state[i][k] = Right;
                            _cycle_detection(x, x_state, i, k);
                        } else {
                            return;
                        }
                        break;
                    }
                }
            }
        }

        _cycle_detection(&self.trans, &mut trans_state, i, j);

        trans_state
    }
    // If the number of basic cells in the transportation plan is less then
    // m + n - 1, it must be replenished
    fn replenish(&mut self) {
        let number = self.trans.data.iter().filter(|&&a| a != None).count();

        if number < self.trans.rows() + self.trans.cols() - 1 {
            for i in 0..self.trans.rows() {
                for j in 0..self.trans.cols() {
                    if self.trans[i][j] == None {
                        self.trans[i][j] = Some(0);
                        if self.cycle_detection(i, j)[i][j] != Direction::None {
                            self.trans[i][j] = None;
                        } else {
                            return;
                        }
                    }
                }
            }
        }
    }
    // Recursive calculation of potentials
    fn calculation_of_potentials(&self) -> (Vec<i64>, Vec<i64>) {
        let mut u = vec![0;self.trans.rows()];
        let mut v = vec![0;self.trans.cols()];

        fn calculation_of_potentials_h(c: &Matrix<u64>, x: &Matrix<Option<u64>>, u: &mut [i64], v: &mut [i64], i: usize, j: usize) {
            for k in 0..x.cols() {
                if x[i][k] != None {
                    v[k] = c[i][k] as i64 - u[i];
                    if k != j {
                        calculation_of_potentials_v(c, x, u, v, i, k);
                    }
                }
            }
        }

        fn calculation_of_potentials_v(c: &Matrix<u64>, x: &Matrix<Option<u64>>, u: &mut [i64], v: &mut [i64], i: usize, j: usize) {
            for k in 0..x.rows() {
                if x[k][j] != None {
                    u[k] = c[k][j] as i64 - v[j];
                    if k != i {
                        calculation_of_potentials_h(c, x, u, v, k, j);
                    }
                }
            }
        }

        calculation_of_potentials_h(&self.cost, &self.trans, &mut u, &mut v, 0, 0);
        calculation_of_potentials_v(&self.cost, &self.trans, &mut u, &mut v, 0, 0);

        (u, v)
    }
    // Check for optimality
    // If the transportation plan is not optimal, then return the minimum of difference
    fn check(&self) -> Option<(usize, usize, i64)> {
        let mut min = None;
        let (u, v) = self.calculation_of_potentials();

        for i in 0..self.cost.rows() {
            for j in 0..self.cost.cols() {
                let d = self.cost[i][j] as i64 - u[i] - v[j];
                if d < 0 {
                    if let Some((_, _, m)) = min {
                        if m > d {
                            min = Some((i, j, d));
                        }
                    } else {
                        min = Some((i, j, d));
                    }
                }
            }
        }
        min
    }
    // Calculate the total cost
    fn total_cost(&self) -> u64 {
        let mut z = 0;
        for (t, c) in self.trans.data.iter().zip(self.cost.data.iter()) {
            if !t.is_none() {
                z += t.unwrap() * c;
            }
        }
        z
    }

    pub fn potential_method(&mut self) {
        self.least_cost_method();
        self.replenish();

        while let Some((i, j, _)) = self.check() {
            self.trans[i][j] = Some(0);
            let trans_state = self.cycle_detection(i, j);

            let mut min = std::u64::MAX;
            let mut f = true;
            for x in trans_state.iter(i, j) {
                if f {
                    f = false;
                } else {
                    if min > self.trans[x.0][x.1].unwrap() {
                        min = self.trans[x.0][x.1].unwrap();
                    }
                    f = true;
                }
            }

            f = true;
            for x in trans_state.iter(i, j) {
                if f {
                    self.trans[x.0][x.1] = Some(self.trans[x.0][x.1].unwrap() + min);
                    f = false;
                } else {
                    self.trans[x.0][x.1] = Some(self.trans[x.0][x.1].unwrap() - min);
                    f = true;
                }
            }

            let mut max = (0, 0, 0);
            for x in trans_state.iter(i, j) {
                if self.trans[x.0][x.1].unwrap() == 0 && max.2 <= self.cost[x.0][x.1] {
                    max = (x.0, x.1, self.cost[x.0][x.1]);
                }
            }
            self.trans[max.0][max.1] = None;
        }
    }

    pub fn printstd(&self) {
        let mut table = Table::new();
        table.add_row(Row::new(vec![Cell::new("")]));
        for i in 0..self.demand.len() {
            table[0].add_cell(Cell::new(&format!("B{}", i + 1)));
        }
        table[0].add_cell(Cell::new("Запасы"));
        for (i, s) in self.supply.iter().enumerate() {
            table.add_row(Row::new(vec![Cell::new(&format!("A{}", i + 1))]));
            for m in self.trans[i].iter() {
                if m.is_none() {
                    table[i + 1].add_cell(Cell::new("0"));
                } else {
                    table[i + 1].add_cell(Cell::new(&format!("{}", m.unwrap())));
                }
            }
            table[i + 1].add_cell(Cell::new(&format!("{}", s)));
        }
        table.add_row(Row::new(vec![Cell::new("Потребности")]));
        for d in self.demand.iter() {
            table[self.supply.len() + 1].add_cell(Cell::new(&format!("{}", d)));
        }
        table.add_row(Row::new(vec![Cell::new("Общая стоимость"), Cell::new(&format!("{}", self.total_cost()))]));
        table.printstd();
    }

    pub fn new(mut a: Vec<u64>, mut b: Vec<u64>, mut c: Vec<Vec<u64>>) -> Result<Transportation, Error> {
        if a.is_empty() || b.is_empty() {
            return Err(NumOfSupOrCust);
        }
        if a.len() != c.len() {
            return Err(NumOfRows);
        }
        for i in &c {
            if b.len() != i.len() {
                return Err(NumOfCols);
            }
        }

        let sum_s = a.iter().fold(0, std::ops::Add::add); //a.iter().sum();
        let sum_d = b.iter().fold(0, std::ops::Add::add); //b.iter().sum();
        if sum_s > sum_d {
            b.push(sum_s - sum_d);
            for i in &mut c {
                i.push(0);
            }
        }
        if sum_s < sum_d {
            a.push(sum_d - sum_s);
            c.push(vec![0;b.len()]);
        }
        let mut cost = Matrix::new(c[0].len());
        for i in c {
            cost.push(i);
        }

        Ok(Transportation {
            trans: Matrix {
                cols: b.len(),
                data: vec![None;a.len()*b.len()],
            },
            supply: a,
            demand: b,
            cost: cost,
        })
    }
}
<|endoftext|>