<commit_before>extern crate clap;
extern crate regex;
extern crate itertools;
extern crate permutohedron;
use clap::{Arg, App};
use std::fs;
use std::io;

pub mod common;
pub mod day1;
pub mod day12;
pub mod day15;
pub mod day16;
pub mod day23;
pub mod day24;

fn get_impl(day: i32) -> Box<common::Solution> {
    match day {
        1 => { Box::new(day1::Day1::new()) }
        12 => { Box::new(day12::Day12::new()) }
        15 => { Box::new(day15::Day15::new()) }
        16 => { Box::new(day16::Day16::new()) }
        23 => { Box::new(day23::Day23::new()) }
        24 => { Box::new(day24::Day24::new()) }
        _ => {
            panic!("Unimplemented day {}", day)
        }
    }
}

fn main() {
    let matches = App::new("Advent of Code")
        .version("2016")
        .author("Bert Peters <bert@bertptrs.nl>")
        .arg(Arg::with_name("day")
             .value_name("DAY")
             .help("Number of the day to execute")
             .required(true)
             .takes_value(true))
        .arg(Arg::with_name("part2")
             .short("2")
             .help("Whether to run part 2")
             .long("part2"))
        .arg(Arg::with_name("input")
             .short("i")
             .long("input")
             .help("Optional input file, stdin otherwise")
             .takes_value(true))
        .get_matches();

    let day: i32 = (&matches.value_of("day").unwrap()).parse()
        .expect("Invalid int");
    let mut implementation = get_impl(day);
    let mut data: Box<io::Read> = match matches.value_of("input") {
        Some(filename) => { Box::new(fs::File::open(filename).unwrap()) }
        None => { Box::new(io::stdin()) }
    };

    if matches.is_present("part2") {
        println!("{}", implementation.part2(&mut data));
    } else {
        println!("{}", implementation.part1(&mut data));
    }
}

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

    #[test]
    fn test_get_impl() {
        // Verify that we can load all days
        let last_implemented = 1;
        for d in 1..(last_implemented + 1) {
            get_impl(d);
        }
    }
}
<commit_msg>Slightly nicer day parsing.<commit_after>extern crate clap;
extern crate regex;
extern crate itertools;
extern crate permutohedron;
use clap::{Arg, App};
use std::fs;
use std::io;

pub mod common;
pub mod day1;
pub mod day12;
pub mod day15;
pub mod day16;
pub mod day23;
pub mod day24;

fn get_impl(day: &str) -> Box<common::Solution> {
    match day.parse() {
        Ok(1) => { Box::new(day1::Day1::new()) }
        Ok(12) => { Box::new(day12::Day12::new()) }
        Ok(15) => { Box::new(day15::Day15::new()) }
        Ok(16) => { Box::new(day16::Day16::new()) }
        Ok(23) => { Box::new(day23::Day23::new()) }
        Ok(24) => { Box::new(day24::Day24::new()) }
        Ok(val) => {
            panic!("Unimplemented day {}", val)
        },
        Err(_) => {
            panic!("Invalid day");
        }
    }
}

fn main() {
    let matches = App::new("Advent of Code")
        .version("2016")
        .author("Bert Peters <bert@bertptrs.nl>")
        .arg(Arg::with_name("day")
             .value_name("DAY")
             .help("Number of the day to execute")
             .required(true)
             .takes_value(true))
        .arg(Arg::with_name("part2")
             .short("2")
             .help("Whether to run part 2")
             .long("part2"))
        .arg(Arg::with_name("input")
             .short("i")
             .long("input")
             .help("Optional input file, stdin otherwise")
             .takes_value(true))
        .get_matches();

    let mut implementation = get_impl(&matches.value_of("day").unwrap());
    let mut data: Box<io::Read> = match matches.value_of("input") {
        Some(filename) => { Box::new(fs::File::open(filename).unwrap()) }
        None => { Box::new(io::stdin()) }
    };

    if matches.is_present("part2") {
        println!("{}", implementation.part2(&mut data));
    } else {
        println!("{}", implementation.part1(&mut data));
    }
}

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

    #[test]
    fn test_get_impl() {
        // Verify that we can load all days
        let last_implemented = 1;
        for d in 1..(last_implemented + 1) {
            get_impl(&format!("{}", d));
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use appkit::NSRect;

use std::libc;
use std::libc::c_long;
use std::io;
use std::str;

pub type id = libc::intptr_t;
pub type Class = libc::intptr_t;
pub type IMP = *u8;
pub type SEL = libc::intptr_t;
pub type Ivar = libc::intptr_t;

pub static nil : id = 0 as id;

pub extern {
    fn class_addMethod(cls: Class, name: SEL, imp: IMP, types: *libc::c_char) -> bool;
    fn class_addIvar(cls : Class,
                     name : *libc::c_char,
                     size : libc::size_t,
                     alignment: u8,
		             types: *libc::c_char) -> bool;
    fn object_setInstanceVariable(obj : id,
                                  name : *libc::c_char,
                                  value : *libc::c_void);
    fn object_getInstanceVariable(obj : id,
                                  name : *libc::c_char,
                                  outValue : **libc::c_void);
    fn objc_allocateClassPair(superclass : Class,
                              name : *libc::c_char,
                              extraBytes : libc::size_t) -> Class;
    fn objc_getClass(name : *libc::c_char) -> id;
    fn objc_msgSend(theReceiver : id, theSelector : SEL) -> id;
    fn objc_registerClassPair(cls : Class);
    fn sel_registerName(name : *libc::c_char) -> SEL;
}

#[test]
pub fn test_nsapp() {
    let klass = str::as_c_str(~"NSApplication", |s|
        unsafe {
            objc_getClass(s)
        }
    );

    let sel = str::as_c_str(~"sharedApplication", |s|
        unsafe {
            sel_registerName(s)
        }
    );

    unsafe {
        let nsapp = objc_msgSend(klass, sel);
        io::println(fmt!("nsapp: %d", (nsapp as int)));
    }
}

#[test]
pub fn test_custom_obj() {
    extern fn MyObject_doSomething(this : id, _sel : SEL) -> id {
        io::println(~"doSomething");
        return this;
    }

    let NSObject = str::as_c_str(~"NSObject", |s|
        unsafe {
            objc_getClass(s)
        }
    );
    let MyObject = str::as_c_str(~"MyObject", |s|
        unsafe {
            objc_allocateClassPair(NSObject, s, 0 as libc::size_t)
        }
    );
    let doSomething = str::as_c_str(~"doSomething", |s|
        unsafe {
            sel_registerName(s)
        }
    );
    let _ = str::as_c_str(~"@@:", |types|
        unsafe {
            class_addMethod(MyObject,
                                  doSomething,
                                  MyObject_doSomething,
                                  types)
        }
    );

    unsafe {
        objc_registerClassPair(MyObject);
    }

    let alloc = str::as_c_str(~"alloc", |s| unsafe { sel_registerName(s) });
    let init = str::as_c_str(~"init", |s| unsafe { sel_registerName(s) });

    unsafe {
        let mut obj = objc_msgSend(MyObject, alloc);
        obj = objc_msgSend(obj, init);
        objc_msgSend(obj, doSomething);
    }
}

/// Invokes the given selector, which must have the signature:
///
///     double f();
pub fn msg_send_double(theReceiver: id, theSelector: SEL) -> f64 {
    unsafe {
        invoke_msg_double(theReceiver, theSelector)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     id f();
pub fn msg_send_id(theReceiver: id, theSelector: SEL) -> id {
    unsafe {
        invoke_msg_id(theReceiver, theSelector)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     id f(NSRect a);
pub fn msg_send_id_NSRect(theReceiver: id, theSelector: SEL, a: NSRect) -> id {
    unsafe {
        invoke_msg_id_NSRect(theReceiver, theSelector, &a)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     id f(id a, id b, id c, id e, id f);
pub fn msg_send_id_id_id_id_id_id(theReceiver: id,
                                  theSelector: SEL,
                                  a: id,
                                  b: id,
                                  c: id,
                                  d: id,
                                  e: id)
                                  -> id {
    unsafe {
        invoke_msg_id_id_id_id_id_id(theReceiver, theSelector, a, b, c, d, e)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     long f();
pub fn msg_send_long(theReceiver: id, theSelector: SEL) -> c_long {
    unsafe {
        invoke_msg_long(theReceiver, theSelector)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     void f();
pub fn msg_send_void(theReceiver: id, theSelector: SEL) {
    unsafe {
        invoke_msg_void(theReceiver, theSelector)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     void f(BOOL a);
pub fn msg_send_void_bool(theReceiver: id, theSelector: SEL, a: bool) {
    unsafe {
        invoke_msg_void_bool(theReceiver, theSelector, a)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     void f(id a);
pub fn msg_send_void_id(theReceiver: id, theSelector: SEL, a: id) {
    unsafe {
        invoke_msg_void_id(theReceiver, theSelector, a)
    }
}

#[link_args = "-L. -lmsgsend"]
#[nolink]
extern {
    fn invoke_msg_double(theReceiver: id, theSelector: SEL) -> f64;
    fn invoke_msg_id(theReceiver: id, theSelector: SEL) -> id;
    fn invoke_msg_id_id_id_id_id_id(theReceiver: id,
                                    theSelector: SEL,
                                    a: id,
                                    b: id,
                                    c: id,
                                    d: id,
                                    e: id)
                                    -> id;
    fn invoke_msg_id_NSRect(theReceiver: id, theSelector: SEL, a: &NSRect) -> id;
    fn invoke_msg_long(theReceiver: id, theSelector: SEL) -> c_long;
    fn invoke_msg_void(theReceiver: id, theSelector: SEL);
    fn invoke_msg_void_bool(theReceiver: id, theSelector: SEL, a: bool);
    fn invoke_msg_void_id(theReceiver: id, theSelector: SEL, a: id);
}

<commit_msg>Warning police.<commit_after>// Copyright 2013 The Servo Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use appkit::NSRect;

use std::libc;
use std::libc::c_long;

pub type id = libc::intptr_t;
pub type Class = libc::intptr_t;
pub type IMP = *u8;
pub type SEL = libc::intptr_t;
pub type Ivar = libc::intptr_t;

pub static NIL : id = 0 as id;

pub extern {
    fn class_addMethod(cls: Class, name: SEL, imp: IMP, types: *libc::c_char) -> bool;
    fn class_addIvar(cls : Class,
                     name : *libc::c_char,
                     size : libc::size_t,
                     alignment: u8,
		             types: *libc::c_char) -> bool;
    fn object_setInstanceVariable(obj : id,
                                  name : *libc::c_char,
                                  value : *libc::c_void);
    fn object_getInstanceVariable(obj : id,
                                  name : *libc::c_char,
                                  outValue : **libc::c_void);
    fn objc_allocateClassPair(superclass : Class,
                              name : *libc::c_char,
                              extraBytes : libc::size_t) -> Class;
    fn objc_getClass(name : *libc::c_char) -> id;
    fn objc_msgSend(theReceiver : id, theSelector : SEL) -> id;
    fn objc_registerClassPair(cls : Class);
    fn sel_registerName(name : *libc::c_char) -> SEL;
}

#[test]
pub fn test_nsapp() {
    let klass = str::as_c_str(~"NSApplication", |s|
        unsafe {
            objc_getClass(s)
        }
    );

    let sel = str::as_c_str(~"sharedApplication", |s|
        unsafe {
            sel_registerName(s)
        }
    );

    unsafe {
        let nsapp = objc_msgSend(klass, sel);
        io::println(fmt!("nsapp: %d", (nsapp as int)));
    }
}

#[test]
pub fn test_custom_obj() {
    extern fn MyObject_doSomething(this : id, _sel : SEL) -> id {
        io::println(~"doSomething");
        return this;
    }

    let NSObject = str::as_c_str(~"NSObject", |s|
        unsafe {
            objc_getClass(s)
        }
    );
    let MyObject = str::as_c_str(~"MyObject", |s|
        unsafe {
            objc_allocateClassPair(NSObject, s, 0 as libc::size_t)
        }
    );
    let doSomething = str::as_c_str(~"doSomething", |s|
        unsafe {
            sel_registerName(s)
        }
    );
    let _ = str::as_c_str(~"@@:", |types|
        unsafe {
            class_addMethod(MyObject,
                                  doSomething,
                                  MyObject_doSomething,
                                  types)
        }
    );

    unsafe {
        objc_registerClassPair(MyObject);
    }

    let alloc = str::as_c_str(~"alloc", |s| unsafe { sel_registerName(s) });
    let init = str::as_c_str(~"init", |s| unsafe { sel_registerName(s) });

    unsafe {
        let mut obj = objc_msgSend(MyObject, alloc);
        obj = objc_msgSend(obj, init);
        objc_msgSend(obj, doSomething);
    }
}

/// Invokes the given selector, which must have the signature:
///
///     double f();
pub fn msg_send_double(theReceiver: id, theSelector: SEL) -> f64 {
    unsafe {
        invoke_msg_double(theReceiver, theSelector)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     id f();
pub fn msg_send_id(theReceiver: id, theSelector: SEL) -> id {
    unsafe {
        invoke_msg_id(theReceiver, theSelector)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     id f(NSRect a);
pub fn msg_send_id_NSRect(theReceiver: id, theSelector: SEL, a: NSRect) -> id {
    unsafe {
        invoke_msg_id_NSRect(theReceiver, theSelector, &a)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     id f(id a, id b, id c, id e, id f);
pub fn msg_send_id_id_id_id_id_id(theReceiver: id,
                                  theSelector: SEL,
                                  a: id,
                                  b: id,
                                  c: id,
                                  d: id,
                                  e: id)
                                  -> id {
    unsafe {
        invoke_msg_id_id_id_id_id_id(theReceiver, theSelector, a, b, c, d, e)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     long f();
pub fn msg_send_long(theReceiver: id, theSelector: SEL) -> c_long {
    unsafe {
        invoke_msg_long(theReceiver, theSelector)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     void f();
pub fn msg_send_void(theReceiver: id, theSelector: SEL) {
    unsafe {
        invoke_msg_void(theReceiver, theSelector)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     void f(BOOL a);
pub fn msg_send_void_bool(theReceiver: id, theSelector: SEL, a: bool) {
    unsafe {
        invoke_msg_void_bool(theReceiver, theSelector, a)
    }
}

/// Invokes the given selector, which must have the signature:
///
///     void f(id a);
pub fn msg_send_void_id(theReceiver: id, theSelector: SEL, a: id) {
    unsafe {
        invoke_msg_void_id(theReceiver, theSelector, a)
    }
}

#[link_args = "-L. -lmsgsend"]
#[nolink]
extern {
    fn invoke_msg_double(theReceiver: id, theSelector: SEL) -> f64;
    fn invoke_msg_id(theReceiver: id, theSelector: SEL) -> id;
    fn invoke_msg_id_id_id_id_id_id(theReceiver: id,
                                    theSelector: SEL,
                                    a: id,
                                    b: id,
                                    c: id,
                                    d: id,
                                    e: id)
                                    -> id;
    fn invoke_msg_id_NSRect(theReceiver: id, theSelector: SEL, a: &NSRect) -> id;
    fn invoke_msg_long(theReceiver: id, theSelector: SEL) -> c_long;
    fn invoke_msg_void(theReceiver: id, theSelector: SEL);
    fn invoke_msg_void_bool(theReceiver: id, theSelector: SEL, a: bool);
    fn invoke_msg_void_id(theReceiver: id, theSelector: SEL, a: id);
}

<|endoftext|>


<commit_before>#![allow(deprecated)]
use std::path::PathBuf;
use std::{io, fs, slice};
use curl::easy::Easy;
use memmap::{Mmap, Protection};
use flate2::{Flush, Decompress};
use std::io::Write;
use utils::data::{Dataset, DatasetIntf};
use std::rc::Rc;
use tensor::{Tensor, TensorKind};
use torch;

type Sample = (TensorKind, i64);

static URLS: [&str; 4] = ["http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz"];
static RAW_FOLDER: &str = "raw";
static PROCESSED_FOLDER: &str = "processed";
static TRAINING_FILE: &str = "training.pt";
static TEST_FILE: &str = "test.pt";

fn create_dir_f(arg: PathBuf) -> io::Result<()> {
    let result = fs::create_dir(arg);
    if let Some(err) = result.err() {
        match err.kind() {
            io::ErrorKind::AlreadyExists => Ok(()),
            _ => Err(err),
        }
    } else {
        Ok(())
    }
}

fn check_exists(arg: &PathBuf) -> io::Result<bool> {
    let train = fs::metadata(arg.join(TRAINING_FILE))?;
    let test = fs::metadata(arg.join(TEST_FILE))?;

    return Ok(train.is_file() && test.is_file());
}

fn download(root: &String) -> io::Result<()> {
    let raw_path = PathBuf::from(root).join(RAW_FOLDER);
    let processed_path = PathBuf::from(root).join(PROCESSED_FOLDER);
    create_dir_f(raw_path.clone())?;
    create_dir_f(processed_path.clone())?;

    for url in URLS.iter() {
        println!("downloading {}", url);
        let paths: Vec<_> = url.split("/").collect();
        let fname: Vec<_> = paths[5].split("/").collect();
        let file_path = raw_path.join(fname[0]);
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .open(file_path)?;
        let mut data = Vec::new();
        let mut handle = Easy::new();
        handle.url(url).unwrap();
        {
            let mut transfer = handle.transfer();
            transfer
                .write_function(|new_data| {
                                    data.extend_from_slice(new_data);
                                    Ok(new_data.len())
                                })
                .unwrap();
            transfer.perform().unwrap();
        }
        let mut output = Vec::with_capacity(data.len());
        // do we have a header?
        Decompress::new(false)
            .decompress_vec(data.as_slice(), &mut output, Flush::Finish)?;
        file.write_all(output.as_slice())?;
    }
    println!("Proceeding");

    let training_set = vec![read_image_file(raw_path.join("train-images-idx3-ubyte"))?,
                            read_label_file(raw_path.join("train-labels-idx1-ubyte"))?];
    let test_set = vec![read_image_file(raw_path.join("t10k-images-idx3-ubyte"))?,
                        read_label_file(raw_path.join("t10k-labels-idx1-ubyte"))?];

    torch::save(processed_path.join(TRAINING_FILE), &training_set)?;
    torch::save(processed_path.join(TEST_FILE), &test_set)?;

    Ok(())
}

fn read_image_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let (length, nrows, ncols);
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const u32, fp.len()) };
        assert_eq!(data[0], 2051);
        length = data[1] as usize;
        nrows = data[2] as usize;
        ncols = data[3] as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut images: Vec<Vec<u8>> = Vec::with_capacity(length);
    let mut idx = 16;
    for _ in 0..length {
        let mut img = Vec::with_capacity(nrows * ncols);
        for _ in 0..nrows * ncols {
            img.push(data[idx]);
            idx += 1;
        }
        images.push(img)
    }
    Ok(torch::byte_tensor(images).view(&[-1, 28, 28]))
}

fn read_label_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let length;
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const u32, fp.len()) };
        assert_eq!(data[0], 20549);
        length = data[1] as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut labels: Vec<u8> = Vec::with_capacity(length);
    for i in 0..length {
        labels.push(data[8 + i])
    }
    Ok(torch::byte_tensor(labels))
}

pub struct MNIST {
    pub root: String,
    pub train: bool,
    pub data: Tensor<u8>,
    pub labels: Tensor<u8>,
    pub transform: Option<Xfrm>,
}

#[builder(pattern="owned")]
#[derive(Builder)]
pub struct MNISTArgs {
    root: String,
    #[builder(default="true")]
    train: bool,
    #[builder(default="false")]
    download: bool,
}

type Xfrm = Box<fn(&TensorKind) -> TensorKind>;
impl MNISTArgsBuilder {
    pub fn done(self, xfrm: Option<Xfrm>) -> Dataset<Sample> {
        let args = self.build().unwrap();
        Dataset::new(Rc::new(MNIST::new(&args, xfrm)))
    }
}

impl MNIST {
    pub fn build(root: &str) -> MNISTArgsBuilder {
        MNISTArgsBuilder::default().root(root.into())
    }
    pub fn new(args: &MNISTArgs, xfrm: Option<Xfrm>) -> Self {
        if args.download {
            download(&args.root).expect("download failed");
        }
        let processed_path = PathBuf::from(&args.root).join(PROCESSED_FOLDER);
        if !check_exists(&processed_path).expect("Dataset not found, try downloading") {
            panic!("Dataset not found, use download=true to download it");
        }
        let mut v: Vec<TensorKind> = if args.train {
            torch::load(processed_path.join(TRAINING_FILE)).expect("torch load failed")
        } else {
            torch::load(processed_path.join(TEST_FILE)).expect("torch load failed")
        };
        let (data, labels) = (v.remove(0), v.remove(0));

        MNIST {
            root: args.root.clone(),
            train: args.train,
            data: data.into(),
            labels: labels.into(),
            transform: xfrm,
        }
    }
}

impl DatasetIntf<Sample> for MNIST {
    fn len(&self) -> usize {
        if self.train { 60000 } else { 10000 }
    }
    fn iter(&self) -> Box<Iterator<Item = Sample>> {
        unimplemented!();
    }
    fn index(&mut self, idx: usize) -> (TensorKind, i64) {
        let img = self.data.s(idx);
        let img = if let Some(ref transform) = self.transform {
            transform(&img.into())
        } else {
            img.copy().into()
        };
        (img, self.labels[idx] as i64)
    }
}
<commit_msg>fix assert<commit_after>#![allow(deprecated)]
use std::path::PathBuf;
use std::{io, fs, slice};
use curl::easy::Easy;
use memmap::{Mmap, Protection};
use flate2::{Flush, Decompress};
use std::io::Write;
use utils::data::{Dataset, DatasetIntf};
use std::rc::Rc;
use tensor::{Tensor, TensorKind};
use torch;

type Sample = (TensorKind, i64);

static URLS: [&str; 4] = ["http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz",
                          "http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz"];
static RAW_FOLDER: &str = "raw";
static PROCESSED_FOLDER: &str = "processed";
static TRAINING_FILE: &str = "training.pt";
static TEST_FILE: &str = "test.pt";

fn create_dir_f(arg: PathBuf) -> io::Result<()> {
    let result = fs::create_dir(arg);
    if let Some(err) = result.err() {
        match err.kind() {
            io::ErrorKind::AlreadyExists => Ok(()),
            _ => Err(err),
        }
    } else {
        Ok(())
    }
}

fn check_exists(arg: &PathBuf) -> io::Result<bool> {
    let train = fs::metadata(arg.join(TRAINING_FILE))?;
    let test = fs::metadata(arg.join(TEST_FILE))?;

    return Ok(train.is_file() && test.is_file());
}

fn download(root: &String) -> io::Result<()> {
    let raw_path = PathBuf::from(root).join(RAW_FOLDER);
    let processed_path = PathBuf::from(root).join(PROCESSED_FOLDER);
    create_dir_f(raw_path.clone())?;
    create_dir_f(processed_path.clone())?;

    for url in URLS.iter() {
        println!("downloading {}", url);
        let paths: Vec<_> = url.split("/").collect();
        let fname: Vec<_> = paths[5].split("/").collect();
        let file_path = raw_path.join(fname[0]);
        let mut file = fs::OpenOptions::new()
            .write(true)
            .create(true)
            .open(file_path)?;
        let mut data = Vec::new();
        let mut handle = Easy::new();
        handle.url(url).unwrap();
        {
            let mut transfer = handle.transfer();
            transfer
                .write_function(|new_data| {
                                    data.extend_from_slice(new_data);
                                    Ok(new_data.len())
                                })
                .unwrap();
            transfer.perform().unwrap();
        }
        let mut output = Vec::with_capacity(data.len());
        // do we have a header?
        Decompress::new(false)
            .decompress_vec(data.as_slice(), &mut output, Flush::Finish)?;
        file.write_all(output.as_slice())?;
    }
    println!("Proceeding");

    let training_set =
        vec![read_image_file(raw_path.join("train-images-idx3-ubyte"))?,
                            read_label_file(raw_path.join("train-labels-idx1-ubyte"))?];
    let test_set = vec![read_image_file(raw_path.join("t10k-images-idx3-ubyte"))?,
                        read_label_file(raw_path.join("t10k-labels-idx1-ubyte"))?];

    torch::save(processed_path.join(TRAINING_FILE), &training_set)?;
    torch::save(processed_path.join(TEST_FILE), &test_set)?;

    Ok(())
}

fn read_image_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let (length, nrows, ncols);
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const u32, fp.len()) };
        assert_eq!(data[0], 2051);
        length = data[1] as usize;
        nrows = data[2] as usize;
        ncols = data[3] as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut images: Vec<Vec<u8>> = Vec::with_capacity(length);
    let mut idx = 16;
    for _ in 0..length {
        let mut img = Vec::with_capacity(nrows * ncols);
        for _ in 0..nrows * ncols {
            img.push(data[idx]);
            idx += 1;
        }
        images.push(img)
    }
    Ok(torch::byte_tensor(images).view(&[-1, 28, 28]))
}

fn read_label_file(path: PathBuf) -> io::Result<TensorKind> {
    let fp = Mmap::open_path(path, Protection::Read)?;
    let length;
    {
        let data = unsafe { slice::from_raw_parts(fp.ptr() as *const u32, fp.len()) };
        assert_eq!(data[0], 2049);
        length = data[1] as usize;
    }
    let data = unsafe { fp.as_slice() };
    let mut labels: Vec<u8> = Vec::with_capacity(length);
    for i in 0..length {
        labels.push(data[8 + i])
    }
    Ok(torch::byte_tensor(labels))
}

pub struct MNIST {
    pub root: String,
    pub train: bool,
    pub data: Tensor<u8>,
    pub labels: Tensor<u8>,
    pub transform: Option<Xfrm>,
}

#[builder(pattern="owned")]
#[derive(Builder)]
pub struct MNISTArgs {
    root: String,
    #[builder(default="true")]
    train: bool,
    #[builder(default="false")]
    download: bool,
}

type Xfrm = Box<fn(&TensorKind) -> TensorKind>;
impl MNISTArgsBuilder {
    pub fn done(self, xfrm: Option<Xfrm>) -> Dataset<Sample> {
        let args = self.build().unwrap();
        Dataset::new(Rc::new(MNIST::new(&args, xfrm)))
    }
}

impl MNIST {
    pub fn build(root: &str) -> MNISTArgsBuilder {
        MNISTArgsBuilder::default().root(root.into())
    }
    pub fn new(args: &MNISTArgs, xfrm: Option<Xfrm>) -> Self {
        if args.download {
            download(&args.root).expect("download failed");
        }
        let processed_path = PathBuf::from(&args.root).join(PROCESSED_FOLDER);
        if !check_exists(&processed_path).expect("Dataset not found, try downloading") {
            panic!("Dataset not found, use download=true to download it");
        }
        let mut v: Vec<TensorKind> = if args.train {
            torch::load(processed_path.join(TRAINING_FILE)).expect("torch load failed")
        } else {
            torch::load(processed_path.join(TEST_FILE)).expect("torch load failed")
        };
        let (data, labels) = (v.remove(0), v.remove(0));

        MNIST {
            root: args.root.clone(),
            train: args.train,
            data: data.into(),
            labels: labels.into(),
            transform: xfrm,
        }
    }
}

impl DatasetIntf<Sample> for MNIST {
    fn len(&self) -> usize {
        if self.train { 60000 } else { 10000 }
    }
    fn iter(&self) -> Box<Iterator<Item = Sample>> {
        unimplemented!();
    }
    fn index(&mut self, idx: usize) -> (TensorKind, i64) {
        let img = self.data.s(idx);
        let img = if let Some(ref transform) = self.transform {
            transform(&img.into())
        } else {
            img.copy().into()
        };
        (img, self.labels[idx] as i64)
    }
}
<|endoftext|>


<commit_before>use bifrost::vector_clock::{VectorClock, StandardVectorClock, ServerVectorClock};
use bifrost::utils::time::get_time;
use chashmap::{CHashMap, WriteGuard};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use ram::types::{Id};
use ram::cell::{Cell, ReadError, WriteError};
use server::NebServer;
use parking_lot::Mutex;
use std::sync::mpsc::{channel, Sender, Receiver};
use futures;
use super::*;

type TxnAwaits = Arc<Mutex<HashMap<u64, Arc<AwaitingServer>>>>;
type TxnGuard<'a> = WriteGuard<'a, TransactionId, Transaction>;
type ChangedObjs = BTreeMap<u64, Vec<(Id, DataObject)>>;
type DataSiteClients = HashMap<u64, Arc<data_site::AsyncServiceClient>>;

pub static DEFAULT_SERVICE_ID: u64 = hash_ident!(TNX_MANAGER_RPC_SERVICE) as u64;

#[derive(Debug, Serialize, Deserialize)]
pub enum TMError {
    TransactionNotFound,
    CannotLocateCellServer,
    NoResponseFromCellServer,
    AssertionError
}

#[derive(Clone)]
struct DataObject {
    server: u64,
    cell: Option<Cell>,
    new: bool
}

struct Transaction {
    start_time: i64, // use for timeout detecting
    id: TransactionId,
    data: BTreeMap<Id, DataObject>,
    writes: BTreeSet<Id>,
}

service! {
    rpc begin() -> TransactionId;
    rpc read(tid: TransactionId, id: Id) -> TransactionExecResult<Cell, ReadError> | TMError;
    rpc write(tid: TransactionId, id: Id, cell: Cell) -> TransactionExecResult<(), WriteError> | TMError;
    rpc update(tid: TransactionId, cell: Cell) -> TransactionExecResult<(), WriteError> | TMError;
    rpc remove(tid: TransactionId, id: Id) -> TransactionExecResult<(), WriteError> | TMError;

    rpc prepare(tid: TransactionId) -> PrepareResult | TMError;
    rpc commit(tid: TransactionId);
    rpc abort(tid: TransactionId);

    rpc go_ahead(tid: BTreeSet<TransactionId>); // invoked by data site to continue on it's transaction in case of waiting
}

pub struct TransactionManager {
    server: Arc<NebServer>,
    transactions: CHashMap<TransactionId, Transaction>,
    data_sites: CHashMap<u64, Arc<data_site::AsyncServiceClient>>,
    await_manager: AwaitManager
}
dispatch_rpc_service_functions!(TransactionManager);

impl TransactionManager {
    pub fn new(server: &Arc<NebServer>) -> Arc<TransactionManager> {
        Arc::new(TransactionManager {
            server: server.clone(),
            transactions: CHashMap::new(),
            data_sites: CHashMap::new(),
            await_manager: AwaitManager::new()
        })
    }
}

impl TransactionManager {
    fn server_id(&self) -> u64 {
        self.server.server_id
    }
    fn get_data_site(&self, server_id: u64) -> io::Result<Arc<data_site::AsyncServiceClient>> {
        if !self.data_sites.contains_key(&server_id) {
            let client = self.server.get_member_by_server_id(server_id)?;
            self.data_sites.upsert(server_id, || {
                data_site::AsyncServiceClient::new(DEFAULT_SERVICE_ID, &client)
            }, |_| {});
        }
        Ok(self.data_sites .get(&server_id).unwrap().clone())
    }
    fn get_data_site_by_id(&self, id: &Id) -> io::Result<(u64, Arc<data_site::AsyncServiceClient>)> {
        match self.server.get_server_id_by_id(id) {
            Some(id) => match self.get_data_site(id) {
                Ok(site) => Ok((id, site.clone())),
                Err(e) => Err(e)
            },
            _ => Err(io::Error::new(io::ErrorKind::NotFound, "cannot find data site for this id"))
        }
    }
    fn get_clock(&self) -> StandardVectorClock {
        self.server.txn_peer.clock.to_clock()
    }
    fn merge_clock(&self, clock: &StandardVectorClock) {
        self.server.txn_peer.clock.merge_with(clock)
    }
    fn get_transaction(&self, tid: &TransactionId) -> Result<TxnGuard, TMError> {
        match self.transactions.get_mut(tid) {
            Some(tnx) => Ok(tnx),
            _ => {Err(TMError::TransactionNotFound)}
        }
    }
    fn changed_objs(&self, txn: &TxnGuard) -> ChangedObjs {
        let mut changed_objs: BTreeMap<u64, Vec<(Id, DataObject)>> = BTreeMap::new();
        for (id, data_obj) in &txn.data {
            changed_objs.entry(data_obj.server).or_insert_with(|| Vec::new()).push((*id, data_obj.clone()));
        }
        changed_objs
    }
    fn data_sites(&self, changed_objs: &ChangedObjs) -> Result<DataSiteClients, TMError>  {
        let data_sites: HashMap<u64, _> = changed_objs.iter().map(|(server_id, _)| {
            (*server_id, self.get_data_site(*server_id))
        }).collect();
        if data_sites.iter().any(|(_, data_site)| { data_site.is_err() }) {
            return Err(TMError::CannotLocateCellServer)
        }
        Ok(data_sites.into_iter().map(
            |(server_id, data_site)| {
                (server_id, data_site.unwrap())
            }
        ).collect())
    }
    fn site_prepare(
        server: Arc<NebServer>, awaits: TxnAwaits, tid: TransactionId, objs: Vec<(Id, DataObject)>,
        data_site: Arc<data_site::AsyncServiceClient>
    ) -> Box<futures::Future<Item = PrepareResult, Error = TMError>> {
        let self_server_id = server.server_id;
        let cell_ids: Vec<_> = objs.iter().map(|&(id, _)| id).collect();
        let server_for_clock = server.clone();
        let res_future = data_site
            .prepare(&self_server_id, &server.txn_peer.clock.to_clock(), &tid, &cell_ids)
            .map_err(|_| -> TMError {
                TMError::NoResponseFromCellServer
            })
            .map(move |prepare_res| -> PrepareResult {
                let prepare_res = prepare_res.unwrap();
                server_for_clock.txn_peer.clock.merge_with(&prepare_res.clock);
                prepare_res.payload
            })
            .then(move |prepare_payload| {
                match prepare_payload {
                    Ok(payload) => {
                        match payload {
                            PrepareResult::Wait => {
                                AwaitManager::txn_wait(&awaits, data_site.server_id);
                                return TransactionManager::site_prepare(
                                    server, awaits, tid, objs, data_site
                                ).wait() // after waiting, retry
                            },
                            _ => {
                                return Ok(payload)
                            }
                        }
                    },
                    Err(e) => Err(e)
                }
            });
        Box::new(res_future)
    }
    fn sites_prepare(&self, tid: &TransactionId, changed_objs: ChangedObjs, data_sites: DataSiteClients)
        -> Result<PrepareResult, TMError> {
        let prepare_futures: Vec<_> = changed_objs.into_iter().map(|(server, objs)| {
            let data_site = data_sites.get(&server).unwrap().clone();
            TransactionManager::site_prepare(
                self.server.clone(),
                self.await_manager.get_txn(tid),
                tid.clone(), objs, data_site
            )
        }).collect();
        let prepare_results = future::join_all(prepare_futures).wait()?;
        for result in prepare_results {
            match result {
                PrepareResult::Success => {},
                _ => {return Ok(result)}
            }
        }
        Ok(PrepareResult::Success)
    }
}

impl Service for TransactionManager {
    fn begin(&self) -> Result<TransactionId, ()> {
        let id = self.server.txn_peer.clock.inc();
        self.transactions.insert(id.clone(), Transaction {
            start_time: get_time(),
            id: id.clone(),
            data: BTreeMap::new(),
            writes: BTreeSet::new()
        });
        Ok(id)
    }
    fn read(&self, tid: &TransactionId, id: &Id) -> Result<TransactionExecResult<Cell, ReadError>, TMError> {
        let mut txn = self.get_transaction(tid)?;
        if let Some(dataObj) = txn.data.get(id) {
            match dataObj.cell {
                Some(ref cell) => {
                    return Ok(TransactionExecResult::Accepted(cell.clone())) // read from cache
                },
                None => {
                   return Ok(TransactionExecResult::Error(ReadError::CellDoesNotExisted))
                }
            }
        }
        let server = self.get_data_site_by_id(&id);
        match server {
            Ok((server_id, server)) => {
                let self_server_id = self.server.server_id;
                let read_response = server.read(
                    &self_server_id,
                    &self.get_clock(),
                    tid, id
                ).wait();
                match read_response {
                    Ok(dsr) => {
                        let dsr = dsr.unwrap();
                        self.merge_clock(&dsr.clock);
                        match dsr.payload {
                            TransactionExecResult::Accepted(ref cell) => {
                                txn.data.insert(*id, DataObject {
                                    server: server_id,
                                    cell: Some(cell.clone()),
                                    new: false,
                                });
                            },
                            TransactionExecResult::Error(ReadError::CellDoesNotExisted) => {
                                txn.data.insert(*id, DataObject {
                                    server: server_id,
                                    cell: None,
                                    new: false,
                                });
                            }
                            TransactionExecResult::Wait => {} // TODO: deal with wait
                            _ => {}
                        }
                        Ok(dsr.payload)
                    },
                    Err(e) => {
                        error!("{:?}", e);
                        Err(TMError::NoResponseFromCellServer)
                    }
                }
            },
            Err(e) => {
                error!("{:?}", e);
                Err(TMError::CannotLocateCellServer)
            }
        }
    }
    fn write(&self, tid: &TransactionId, id: &Id, cell: &Cell) -> Result<TransactionExecResult<(), WriteError>, TMError> {
        let mut txn = self.get_transaction(tid)?;
        match self.server.get_server_id_by_id(id) {
            Some(server_id) => {
                let have_cached_cell = txn.data.contains_key(id);
                if !have_cached_cell {
                    txn.data.insert(*id, DataObject {
                        server: server_id,
                        cell: Some(cell.clone()),
                        new: true,
                    });
                    Ok(TransactionExecResult::Accepted(()))
                } else {
                    let mut data_obj = txn.data.get_mut(id).unwrap();
                    if !data_obj.cell.is_none() {
                        return Ok(TransactionExecResult::Error(WriteError::CellAlreadyExisted))
                    }
                    data_obj.cell = Some(cell.clone());
                    Ok(TransactionExecResult::Accepted(()))
                }
            },
            None => Err(TMError::CannotLocateCellServer)
        }
    }
    fn update(&self, tid: &TransactionId, cell: &Cell) -> Result<TransactionExecResult<(), WriteError>, TMError> {
        let mut txn = self.get_transaction(tid)?;
        let id = cell.id();
        match self.server.get_server_id_by_id(&id) {
            Some(server_id) => {
                let cell = cell.clone();
                if txn.data.contains_key(&id) {
                    txn.data.get_mut(&id).unwrap().cell = Some(cell)
                } else {
                    txn.data.insert(id, DataObject {
                        server: server_id,
                        cell: Some(cell),
                        new: false,
                    });
                }
                Ok(TransactionExecResult::Accepted(()))
            },
            None => Err(TMError::CannotLocateCellServer)
        }
    }
    fn remove(&self, tid: &TransactionId, id: &Id) -> Result<TransactionExecResult<(), WriteError>, TMError> {
        let mut txn = self.get_transaction(tid)?;
        match self.server.get_server_id_by_id(&id) {
            Some(server_id) => {
                if txn.data.contains_key(&id) {
                    let data_obj = txn.data.get_mut(&id).unwrap();
                    if data_obj.cell.is_none() {
                        return Ok(TransactionExecResult::Error(WriteError::CellDoesNotExisted))
                    }
                    data_obj.cell = None;
                } else {
                    txn.data.insert(*id, DataObject {
                        server: server_id,
                        cell: None,
                        new: false,
                    });
                }
                Ok(TransactionExecResult::Accepted(()))
            },
            None => Err(TMError::CannotLocateCellServer)
        }
    }
    fn prepare(&self, tid: &TransactionId) -> Result<PrepareResult, TMError> {
        let mut txn = self.get_transaction(tid)?;
        let mut changed_objs = self.changed_objs(&txn);
        let data_sites = self.data_sites(&changed_objs)?;
        self.sites_prepare(tid, changed_objs, data_sites)
    }
    fn commit(&self, tid: &TransactionId) -> Result<(), ()> {
        Err(())
    }
    fn abort(&self, tid: &TransactionId) -> Result<(), ()> {
        Err(())
    }
    fn go_ahead(&self, tid: &BTreeSet<TransactionId>) -> Result<(), ()> {
        Err(())
    }
}

struct AwaitingServer {
    sender: Mutex<Sender<()>>,
    receiver: Mutex<Receiver<()>>,
}

impl AwaitingServer {
    pub fn new() -> AwaitingServer {
        let (sender, receiver) = channel();
        AwaitingServer {
            sender: Mutex::new(sender),
            receiver: Mutex::new(receiver)
        }
    }
    pub fn send(&self) {
        self.sender.lock().send(());
    }
    pub fn wait(&self) {
        self.receiver.lock().recv();
    }
}

struct AwaitManager {
    channels: Mutex<HashMap<TransactionId, TxnAwaits>>
}

impl AwaitManager {
    pub fn new() -> AwaitManager {
        AwaitManager {
            channels: Mutex::new(HashMap::new())
        }
    }
    pub fn get_txn(&self, tid: &TransactionId) -> TxnAwaits {
        self.channels.lock()
            .entry(tid.clone())
            .or_insert_with(|| Arc::new(Mutex::new(HashMap::new())))
            .clone()
    }
    pub fn server_from_txn_awaits(awaits: &TxnAwaits, server_id: u64) -> Arc<AwaitingServer> {
        awaits.lock()
            .entry(server_id)
            .or_insert_with(|| Arc::new(AwaitingServer::new()))
            .clone()
    }
    pub fn txn_send(awaits: &TxnAwaits, server_id: u64) {
        AwaitManager::server_from_txn_awaits(awaits, server_id).send();
    }
    pub fn txn_wait(awaits: &TxnAwaits, server_id: u64) {
        AwaitManager::server_from_txn_awaits(awaits, server_id).wait();
    }
}<commit_msg>read with await<commit_after>use bifrost::vector_clock::{VectorClock, StandardVectorClock, ServerVectorClock};
use bifrost::utils::time::get_time;
use chashmap::{CHashMap, WriteGuard};
use std::collections::{BTreeMap, BTreeSet, HashMap};
use ram::types::{Id};
use ram::cell::{Cell, ReadError, WriteError};
use server::NebServer;
use parking_lot::Mutex;
use std::sync::mpsc::{channel, Sender, Receiver};
use futures;
use super::*;

type TxnAwaits = Arc<Mutex<HashMap<u64, Arc<AwaitingServer>>>>;
type TxnGuard<'a> = WriteGuard<'a, TransactionId, Transaction>;
type ChangedObjs = BTreeMap<u64, Vec<(Id, DataObject)>>;
type DataSiteClients = HashMap<u64, Arc<data_site::AsyncServiceClient>>;

pub static DEFAULT_SERVICE_ID: u64 = hash_ident!(TNX_MANAGER_RPC_SERVICE) as u64;

#[derive(Debug, Serialize, Deserialize)]
pub enum TMError {
    TransactionNotFound,
    CannotLocateCellServer,
    NoResponseFromCellServer,
    AssertionError
}

#[derive(Clone)]
struct DataObject {
    server: u64,
    cell: Option<Cell>,
    new: bool
}

struct Transaction {
    start_time: i64, // use for timeout detecting
    id: TransactionId,
    data: BTreeMap<Id, DataObject>,
    writes: BTreeSet<Id>,
}

service! {
    rpc begin() -> TransactionId;
    rpc read(tid: TransactionId, id: Id) -> TransactionExecResult<Cell, ReadError> | TMError;
    rpc write(tid: TransactionId, id: Id, cell: Cell) -> TransactionExecResult<(), WriteError> | TMError;
    rpc update(tid: TransactionId, cell: Cell) -> TransactionExecResult<(), WriteError> | TMError;
    rpc remove(tid: TransactionId, id: Id) -> TransactionExecResult<(), WriteError> | TMError;

    rpc prepare(tid: TransactionId) -> PrepareResult | TMError;
    rpc commit(tid: TransactionId);
    rpc abort(tid: TransactionId);

    rpc go_ahead(tid: BTreeSet<TransactionId>); // invoked by data site to continue on it's transaction in case of waiting
}

pub struct TransactionManager {
    server: Arc<NebServer>,
    transactions: CHashMap<TransactionId, Transaction>,
    data_sites: CHashMap<u64, Arc<data_site::AsyncServiceClient>>,
    await_manager: AwaitManager
}
dispatch_rpc_service_functions!(TransactionManager);

impl TransactionManager {
    pub fn new(server: &Arc<NebServer>) -> Arc<TransactionManager> {
        Arc::new(TransactionManager {
            server: server.clone(),
            transactions: CHashMap::new(),
            data_sites: CHashMap::new(),
            await_manager: AwaitManager::new()
        })
    }
}

impl TransactionManager {
    fn server_id(&self) -> u64 {
        self.server.server_id
    }
    fn get_data_site(&self, server_id: u64) -> io::Result<Arc<data_site::AsyncServiceClient>> {
        if !self.data_sites.contains_key(&server_id) {
            let client = self.server.get_member_by_server_id(server_id)?;
            self.data_sites.upsert(server_id, || {
                data_site::AsyncServiceClient::new(DEFAULT_SERVICE_ID, &client)
            }, |_| {});
        }
        Ok(self.data_sites .get(&server_id).unwrap().clone())
    }
    fn get_data_site_by_id(&self, id: &Id) -> io::Result<(u64, Arc<data_site::AsyncServiceClient>)> {
        match self.server.get_server_id_by_id(id) {
            Some(id) => match self.get_data_site(id) {
                Ok(site) => Ok((id, site.clone())),
                Err(e) => Err(e)
            },
            _ => Err(io::Error::new(io::ErrorKind::NotFound, "cannot find data site for this id"))
        }
    }
    fn get_clock(&self) -> StandardVectorClock {
        self.server.txn_peer.clock.to_clock()
    }
    fn merge_clock(&self, clock: &StandardVectorClock) {
        self.server.txn_peer.clock.merge_with(clock)
    }
    fn get_transaction(&self, tid: &TransactionId) -> Result<TxnGuard, TMError> {
        match self.transactions.get_mut(tid) {
            Some(tnx) => Ok(tnx),
            _ => {Err(TMError::TransactionNotFound)}
        }
    }
    fn changed_objs(&self, txn: &TxnGuard) -> ChangedObjs {
        let mut changed_objs: BTreeMap<u64, Vec<(Id, DataObject)>> = BTreeMap::new();
        for (id, data_obj) in &txn.data {
            changed_objs.entry(data_obj.server).or_insert_with(|| Vec::new()).push((*id, data_obj.clone()));
        }
        changed_objs
    }
    fn data_sites(&self, changed_objs: &ChangedObjs) -> Result<DataSiteClients, TMError>  {
        let data_sites: HashMap<u64, _> = changed_objs.iter().map(|(server_id, _)| {
            (*server_id, self.get_data_site(*server_id))
        }).collect();
        if data_sites.iter().any(|(_, data_site)| { data_site.is_err() }) {
            return Err(TMError::CannotLocateCellServer)
        }
        Ok(data_sites.into_iter().map(
            |(server_id, data_site)| {
                (server_id, data_site.unwrap())
            }
        ).collect())
    }
    fn site_prepare(
        server: Arc<NebServer>, awaits: TxnAwaits, tid: TransactionId, objs: Vec<(Id, DataObject)>,
        data_site: Arc<data_site::AsyncServiceClient>
    ) -> Box<futures::Future<Item = PrepareResult, Error = TMError>> {
        let self_server_id = server.server_id;
        let cell_ids: Vec<_> = objs.iter().map(|&(id, _)| id).collect();
        let server_for_clock = server.clone();
        let res_future = data_site
            .prepare(&self_server_id, &server.txn_peer.clock.to_clock(), &tid, &cell_ids)
            .map_err(|_| -> TMError {
                TMError::NoResponseFromCellServer
            })
            .map(move |prepare_res| -> PrepareResult {
                let prepare_res = prepare_res.unwrap();
                server_for_clock.txn_peer.clock.merge_with(&prepare_res.clock);
                prepare_res.payload
            })
            .then(move |prepare_payload| {
                match prepare_payload {
                    Ok(payload) => {
                        match payload {
                            PrepareResult::Wait => {
                                AwaitManager::txn_wait(&awaits, data_site.server_id);
                                return TransactionManager::site_prepare(
                                    server, awaits, tid, objs, data_site
                                ).wait() // after waiting, retry
                            },
                            _ => {
                                return Ok(payload)
                            }
                        }
                    },
                    Err(e) => Err(e)
                }
            });
        Box::new(res_future)
    }
    fn sites_prepare(&self, tid: &TransactionId, changed_objs: ChangedObjs, data_sites: DataSiteClients)
        -> Result<PrepareResult, TMError> {
        let prepare_futures: Vec<_> = changed_objs.into_iter().map(|(server, objs)| {
            let data_site = data_sites.get(&server).unwrap().clone();
            TransactionManager::site_prepare(
                self.server.clone(),
                self.await_manager.get_txn(tid),
                tid.clone(), objs, data_site
            )
        }).collect();
        let prepare_results = future::join_all(prepare_futures).wait()?;
        for result in prepare_results {
            match result {
                PrepareResult::Success => {},
                _ => {return Ok(result)}
            }
        }
        Ok(PrepareResult::Success)
    }
    fn read_from_site(
        &self, server_id: u64, server: Arc<data_site::AsyncServiceClient>,
        tid: &TransactionId, id: &Id, mut txn: TxnGuard, await: TxnAwaits
    ) -> Result<TransactionExecResult<Cell, ReadError>, TMError> {
        let self_server_id = self.server.server_id;
        let read_response = server.read(
            &self_server_id,
            &self.get_clock(),
            tid, id
        ).wait();
        match read_response {
            Ok(dsr) => {
                let dsr = dsr.unwrap();
                self.merge_clock(&dsr.clock);
                match dsr.payload {
                    TransactionExecResult::Accepted(ref cell) => {
                        txn.data.insert(*id, DataObject {
                            server: server_id,
                            cell: Some(cell.clone()),
                            new: false,
                        });
                    },
                    TransactionExecResult::Error(ReadError::CellDoesNotExisted) => {
                        txn.data.insert(*id, DataObject {
                            server: server_id,
                            cell: None,
                            new: false,
                        });
                    }
                    TransactionExecResult::Wait => {
                        AwaitManager::txn_wait(&await, server_id);
                        self.read_from_site(server_id, server, tid, id, txn, await);
                    }
                    _ => {}
                }
                Ok(dsr.payload)
            },
            Err(e) => {
                error!("{:?}", e);
                Err(TMError::NoResponseFromCellServer)
            }
        }
    }
}

impl Service for TransactionManager {
    fn begin(&self) -> Result<TransactionId, ()> {
        let id = self.server.txn_peer.clock.inc();
        self.transactions.insert(id.clone(), Transaction {
            start_time: get_time(),
            id: id.clone(),
            data: BTreeMap::new(),
            writes: BTreeSet::new()
        });
        Ok(id)
    }
    fn read(&self, tid: &TransactionId, id: &Id) -> Result<TransactionExecResult<Cell, ReadError>, TMError> {
        let mut txn = self.get_transaction(tid)?;
        if let Some(dataObj) = txn.data.get(id) {
            match dataObj.cell {
                Some(ref cell) => {
                    return Ok(TransactionExecResult::Accepted(cell.clone())) // read from cache
                },
                None => {
                   return Ok(TransactionExecResult::Error(ReadError::CellDoesNotExisted))
                }
            }
        }
        let server = self.get_data_site_by_id(&id);
        match server {
            Ok((server_id, server)) => {
                let await = self.await_manager.get_txn(tid);
                self.read_from_site(server_id, server, tid, id, txn, await)
            },
            Err(e) => {
                error!("{:?}", e);
                Err(TMError::CannotLocateCellServer)
            }
        }
    }
    fn write(&self, tid: &TransactionId, id: &Id, cell: &Cell) -> Result<TransactionExecResult<(), WriteError>, TMError> {
        let mut txn = self.get_transaction(tid)?;
        match self.server.get_server_id_by_id(id) {
            Some(server_id) => {
                let have_cached_cell = txn.data.contains_key(id);
                if !have_cached_cell {
                    txn.data.insert(*id, DataObject {
                        server: server_id,
                        cell: Some(cell.clone()),
                        new: true,
                    });
                    Ok(TransactionExecResult::Accepted(()))
                } else {
                    let mut data_obj = txn.data.get_mut(id).unwrap();
                    if !data_obj.cell.is_none() {
                        return Ok(TransactionExecResult::Error(WriteError::CellAlreadyExisted))
                    }
                    data_obj.cell = Some(cell.clone());
                    Ok(TransactionExecResult::Accepted(()))
                }
            },
            None => Err(TMError::CannotLocateCellServer)
        }
    }
    fn update(&self, tid: &TransactionId, cell: &Cell) -> Result<TransactionExecResult<(), WriteError>, TMError> {
        let mut txn = self.get_transaction(tid)?;
        let id = cell.id();
        match self.server.get_server_id_by_id(&id) {
            Some(server_id) => {
                let cell = cell.clone();
                if txn.data.contains_key(&id) {
                    txn.data.get_mut(&id).unwrap().cell = Some(cell)
                } else {
                    txn.data.insert(id, DataObject {
                        server: server_id,
                        cell: Some(cell),
                        new: false,
                    });
                }
                Ok(TransactionExecResult::Accepted(()))
            },
            None => Err(TMError::CannotLocateCellServer)
        }
    }
    fn remove(&self, tid: &TransactionId, id: &Id) -> Result<TransactionExecResult<(), WriteError>, TMError> {
        let mut txn = self.get_transaction(tid)?;
        match self.server.get_server_id_by_id(&id) {
            Some(server_id) => {
                if txn.data.contains_key(&id) {
                    let data_obj = txn.data.get_mut(&id).unwrap();
                    if data_obj.cell.is_none() {
                        return Ok(TransactionExecResult::Error(WriteError::CellDoesNotExisted))
                    }
                    data_obj.cell = None;
                } else {
                    txn.data.insert(*id, DataObject {
                        server: server_id,
                        cell: None,
                        new: false,
                    });
                }
                Ok(TransactionExecResult::Accepted(()))
            },
            None => Err(TMError::CannotLocateCellServer)
        }
    }
    fn prepare(&self, tid: &TransactionId) -> Result<PrepareResult, TMError> {
        let mut txn = self.get_transaction(tid)?;
        let mut changed_objs = self.changed_objs(&txn);
        let data_sites = self.data_sites(&changed_objs)?;
        self.sites_prepare(tid, changed_objs, data_sites)
    }
    fn commit(&self, tid: &TransactionId) -> Result<(), ()> {
        Err(())
    }
    fn abort(&self, tid: &TransactionId) -> Result<(), ()> {
        Err(())
    }
    fn go_ahead(&self, tid: &BTreeSet<TransactionId>) -> Result<(), ()> {
        Err(())
    }
}

struct AwaitingServer {
    sender: Mutex<Sender<()>>,
    receiver: Mutex<Receiver<()>>,
}

impl AwaitingServer {
    pub fn new() -> AwaitingServer {
        let (sender, receiver) = channel();
        AwaitingServer {
            sender: Mutex::new(sender),
            receiver: Mutex::new(receiver)
        }
    }
    pub fn send(&self) {
        self.sender.lock().send(());
    }
    pub fn wait(&self) {
        self.receiver.lock().recv();
    }
}

struct AwaitManager {
    channels: Mutex<HashMap<TransactionId, TxnAwaits>>
}

impl AwaitManager {
    pub fn new() -> AwaitManager {
        AwaitManager {
            channels: Mutex::new(HashMap::new())
        }
    }
    pub fn get_txn(&self, tid: &TransactionId) -> TxnAwaits {
        self.channels.lock()
            .entry(tid.clone())
            .or_insert_with(|| Arc::new(Mutex::new(HashMap::new())))
            .clone()
    }
    pub fn server_from_txn_awaits(awaits: &TxnAwaits, server_id: u64) -> Arc<AwaitingServer> {
        awaits.lock()
            .entry(server_id)
            .or_insert_with(|| Arc::new(AwaitingServer::new()))
            .clone()
    }
    pub fn txn_send(awaits: &TxnAwaits, server_id: u64) {
        AwaitManager::server_from_txn_awaits(awaits, server_id).send();
    }
    pub fn txn_wait(awaits: &TxnAwaits, server_id: u64) {
        AwaitManager::server_from_txn_awaits(awaits, server_id).wait();
    }
}<|endoftext|>


<commit_before>use azure_sdk_core::errors::{check_status_extract_body, AzureError};
use hyper::{self, header, StatusCode};
use hyper_rustls::HttpsConnector;
use ring::hmac;
use std::ops::Add;
use time::Duration;
use url::form_urlencoded;

mod client;
pub use self::client::Client;

type HttpClient = hyper::Client<HttpsConnector<hyper::client::HttpConnector>>;

#[inline]
fn send_event_prepare<B: Into<String>>(
    http_client: &HttpClient,
    namespace: &str,
    event_hub: &str,
    policy_name: &str,
    signing_key: &hmac::Key,
    event_body: B,
    duration: Duration,
) -> Result<hyper::client::ResponseFuture, AzureError> {
    // prepare the url to call
    let url = format!(
        "https://{}.servicebus.windows.net/{}/messages",
        namespace, event_hub
    );
    debug!("url == {:?}", url);

    // generate sas signature based on key name, key value, url and duration.
    let sas = generate_signature(policy_name, signing_key, &url, duration);
    debug!("sas == {}", sas);

    let event_body = event_body.into();
    let request = hyper::Request::post(url)
        .header(header::AUTHORIZATION, sas)
        .body(event_body.into())?;

    Ok(http_client.request(request))
}

async fn send_event(
    http_client: &HttpClient,
    namespace: &str,
    event_hub: &str,
    policy_name: &str,
    hmac: &hmac::Key,
    event_body: &str,
    duration: Duration,
) -> Result<(), AzureError> {
    let req = send_event_prepare(
        http_client,
        namespace,
        event_hub,
        policy_name,
        hmac,
        event_body,
        duration,
    );

    check_status_extract_body(req?, StatusCode::CREATED).await?;
    Ok(())
}

fn generate_signature(
    policy_name: &str,
    signing_key: &hmac::Key,
    url: &str,
    ttl: Duration,
) -> String {
    use url::form_urlencoded::Serializer;

    let expiry = ::chrono::Utc::now().add(ttl).timestamp();
    debug!("expiry == {:?}", expiry);

    let url_encoded: String = form_urlencoded::byte_serialize(url.as_bytes()).collect();
    debug!("url_encoded == {:?}", url_encoded);

    let str_to_sign = format!("{}\n{}", url_encoded, expiry);
    debug!("str_to_sign == {:?}", str_to_sign);

    let sig = hmac::sign(signing_key, str_to_sign.as_bytes());
    let sig = {
        let sig = ::base64::encode(sig.as_ref());
        debug!("sig == {}", sig);
        let mut ser = Serializer::new(String::new());
        ser.append_pair("sig", &sig);
        let sig = ser.finish();
        debug!("sig == {}", sig);
        sig
    };

    debug!("sig == {:?}", sig);

    format!(
        "SharedAccessSignature sr={}&{}&se={}&skn={}",
        &url_encoded, sig, expiry, policy_name
    )
}
<commit_msg>Prepare read message from service bus queue request<commit_after>use azure_sdk_core::errors::{check_status_extract_body, AzureError};
use hyper::{self, header, Body, StatusCode};
use hyper_rustls::HttpsConnector;
use ring::hmac;
use std::ops::Add;
use time::Duration;
use url::form_urlencoded;

mod client;
pub use self::client::Client;

type HttpClient = hyper::Client<HttpsConnector<hyper::client::HttpConnector>>;

#[inline]
fn send_event_prepare<B: Into<String>>(
    http_client: &HttpClient,
    namespace: &str,
    event_hub: &str,
    policy_name: &str,
    signing_key: &hmac::Key,
    event_body: B,
    duration: Duration,
) -> Result<hyper::client::ResponseFuture, AzureError> {
    // prepare the url to call
    let url = format!(
        "https://{}.servicebus.windows.net/{}/messages",
        namespace, event_hub
    );
    debug!("url == {:?}", url);

    // generate sas signature based on key name, key value, url and duration.
    let sas = generate_signature(policy_name, signing_key, &url, duration);
    debug!("sas == {}", sas);

    let event_body = event_body.into();
    let request = hyper::Request::post(url)
        .header(header::AUTHORIZATION, sas)
        .body(event_body.into())?;

    Ok(http_client.request(request))
}

fn peek_lock_event_prepare(
    http_client: &HttpClient,
    namespace: &str,
    event_hub: &str,
    policy_name: &str,
    signing_key: &hmac::Key,
    duration: Duration,
) -> Result<hyper::client::ResponseFuture, AzureError> {
    // prepare the url to call
    let url = format!(
        "https://{}.servicebus.windows.net/{}/messages/head",
        namespace, event_hub
    );
    debug!("url == {:?}", url);

    // generate sas signature based on key name, key value, url and duration.
    let sas = generate_signature(policy_name, signing_key, &url, duration);
    debug!("sas == {}", sas);

    let request = hyper::Request::post(url)
        .header(header::AUTHORIZATION, sas)
        .body(Body::empty())?;

    Ok(http_client.request(request))
}

fn receive_delete_event_prepare(
    http_client: &HttpClient,
    namespace: &str,
    event_hub: &str,
    policy_name: &str,
    signing_key: &hmac::Key,
    duration: Duration,
) -> Result<hyper::client::ResponseFuture, AzureError> {
    // prepare the url to call
    let url = format!(
        "https://{}.servicebus.windows.net/{}/messages/head",
        namespace, event_hub
    );
    debug!("url == {:?}", url);

    // generate sas signature based on key name, key value, url and duration.
    let sas = generate_signature(policy_name, signing_key, &url, duration);
    debug!("sas == {}", sas);

    let request = hyper::Request::delete(url)
        .header(header::AUTHORIZATION, sas)
        .body(Body::empty())?;

    Ok(http_client.request(request))
}

async fn send_event(
    http_client: &HttpClient,
    namespace: &str,
    event_hub: &str,
    policy_name: &str,
    hmac: &hmac::Key,
    event_body: &str,
    duration: Duration,
) -> Result<(), AzureError> {
    let req = send_event_prepare(
        http_client,
        namespace,
        event_hub,
        policy_name,
        hmac,
        event_body,
        duration,
    );

    check_status_extract_body(req?, StatusCode::CREATED).await?;
    Ok(())
}

fn generate_signature(
    policy_name: &str,
    signing_key: &hmac::Key,
    url: &str,
    ttl: Duration,
) -> String {
    use url::form_urlencoded::Serializer;

    let expiry = ::chrono::Utc::now().add(ttl).timestamp();
    debug!("expiry == {:?}", expiry);

    let url_encoded: String = form_urlencoded::byte_serialize(url.as_bytes()).collect();
    debug!("url_encoded == {:?}", url_encoded);

    let str_to_sign = format!("{}\n{}", url_encoded, expiry);
    debug!("str_to_sign == {:?}", str_to_sign);

    let sig = hmac::sign(signing_key, str_to_sign.as_bytes());
    let sig = {
        let sig = ::base64::encode(sig.as_ref());
        debug!("sig == {}", sig);
        let mut ser = Serializer::new(String::new());
        ser.append_pair("sig", &sig);
        let sig = ser.finish();
        debug!("sig == {}", sig);
        sig
    };

    debug!("sig == {:?}", sig);

    format!(
        "SharedAccessSignature sr={}&{}&se={}&skn={}",
        &url_encoded, sig, expiry, policy_name
    )
}
<|endoftext|>


<commit_before>//! SeaHash: A blazingly fast, portable hash function with proven statistical guarantees.
//!
//! SeaHash is a hash function with performance similar (within around ±7% difference) to XXHash
//! and MetroHash, but with stronger guarantees.
//!
//! SeaHash is a portable hash function, meaning that the output is not dependent on the hosting
//! architecture, and makes no assumptions on endianness or the alike. This stable layout allows it
//! to be used for on-disk/permanent storage (e.g. checksums).
//!
//! SeaHash beats xxHash by 3-10% depending on the architecture in use. SeaHash has better quality
//! guarantees as well.
//!
//! # Achieving the performance
//!
//! Like any good general-purpose hash function, SeaHash reads 8 bytes at once effectively reducing
//! the running time by an order of ~5.
//!
//! Secondly, SeaHash achieves the performance by heavily exploiting Instruction-Level Parallelism.
//! In particular, it fetches 4 integers in every round and independently diffuses them. This
//! yields four different states, which are finally combined.
//!
//! # Advantages over other hash functions
//!
//! - Portability: SeaHash always gives the same hash across any platforms, and can thus be used
//!   for e.g. on-disk structures.
//! - Performance: SeaHash beats every high-quality (grading 10/10 in smhasher) hash function that
//!   I know off.
//! - Hardware accelerateable: SeaHash is designed such that ASICs can implement it with really
//!   high performance.
//! - Provable quality guarantees: Contrary to most other non-cryptographic hash function, SeaHash
//!   can be proved to satisfy the avalanche criterion as well as BIC.
//!
//! # Statistical guarantees
//!
//! SeaHash comes with certain proven guarantees about the statistical properties of the output:
//!
//! 1. Pick some _n_-byte sequence, _s_. The number of _n_-byte sequence colliding with _s_ is
//!    independent of the choice of _s_ (all equivalence class have equal size).
//! 2. If you flip any bit in the input, the probability for any bit in the output to be flipped is
//!    0.5.
//!
//! The first guarantee can be derived through deduction, by proving that the diffusion function is
//! bijective (reverse the XORs and find the congruence inverses to the primes).
//!
//! The second guarantee requires more complex calculations: Construct a matrix of probabilities
//! and set one to certain (1), then apply transformations through the respective operations. The
//! proof is a bit long, but relatively simple.

#![no_std]
#![warn(missing_docs)]

use core::num::Wrapping as W;

pub use buffer::hash;
pub use stream::SeaHasher;

pub mod reference;
mod buffer;
mod stream;

/// The diffusion function.
///
/// This is a bijective function emitting chaotic behavior. Such functions are used as building
/// blocks for hash functions.
fn diffuse(mut x: W<u64>) -> W<u64> {
    // Move entropy down by XOR with shifting.
    x ^= x >> 32;
    // Move entropy up by scattering through multiplication.
    x *= W(0x7ed0e9fa0d94a33);
    // We still need more entropy downwards. Flipping higher bits won't flip lower ones, so far.
    // For example, if you flip the most significant bit, the 32'th bit will flip per the XOR-shift
    // subdiffusion, but this flip will only be scattered by the multiplication to flipping bits
    // higher than the 32'th, meaning that the ones lower will be unaffected. As such, we need to
    // get some entropy down.
    x ^= x >> 32;
    // So far, the avalanche diagram looks pretty good, but it still emits stripe patterns. For
    // example, flipping the 5'th lowest bit won't flip the least significant bit because of the
    // choice of scalar (in particular, observe how it leaves the 32'th bit unflipped after the
    // multiplication, which means that the XOR-shift never affects the lowest bit). No choice of
    // scalar will make this go away, it will merely change the unaffected bits. Instead, we need
    // to make the behavior more undeterministic by scattering bits through multiplication.
    x *= W(0x7ed0e9fa0d94a33);
    // This is the final stage of the diffusion function. There are still issues with the lowest
    // bits, which are still unaffected by the multiplication above. However, the multiplication
    // solved the higher bits' dependence, so lending entropy from the higher half will fix the
    // issues with the lower half.
    x ^= x >> 32;

    x
}
<commit_msg>Update documentation on the construction of SeaHash.<commit_after>//! SeaHash: A blazingly fast, portable hash function with proven statistical guarantees.
//!
//! SeaHash is a hash function with performance similar (within around ±7% difference) to XXHash
//! and MetroHash, but with stronger guarantees.
//!
//! SeaHash is a portable hash function, meaning that the output is not dependent on the hosting
//! architecture, and makes no assumptions on endianness or the alike. This stable layout allows it
//! to be used for on-disk/permanent storage (e.g. checksums).
//!
//! SeaHash beats xxHash by 3-10% depending on the architecture in use. SeaHash has better quality
//! guarantees as well.
//!
//! # Achieving the performance
//!
//! Like any good general-purpose hash function, SeaHash reads 8 bytes at once effectively reducing
//! the running time by an order of ~5.
//!
//! Secondly, SeaHash achieves the performance by heavily exploiting Instruction-Level Parallelism.
//! In particular, it fetches 4 integers in every round and independently diffuses them. This
//! yields four different states, which are finally combined.
//!
//! # Advantages over other hash functions
//!
//! - Portability: SeaHash always gives the same hash across any platforms, and can thus be used
//!   for e.g. on-disk structures.
//! - Performance: SeaHash beats every high-quality (grading 10/10 in smhasher) hash function that
//!   I know off.
//! - Hardware accelerateable: SeaHash is designed such that ASICs can implement it with really
//!   high performance.
//! - Provable quality guarantees: Contrary to most other non-cryptographic hash function, SeaHash
//!   can be proved to satisfy the avalanche criterion as well as BIC.
//!
//! # Statistical guarantees
//!
//! SeaHash comes with certain proven guarantees about the statistical properties of the output:
//!
//! 1. Pick some _n_-byte sequence, _s_. The number of _n_-byte sequence colliding with _s_ is
//!    independent of the choice of _s_ (all equivalence class have equal size).
//! 2. If you flip any bit in the input, the probability for any bit in the output to be flipped is
//!    0.5.
//!
//! The first guarantee can be derived through deduction, by proving that the diffusion function is
//! bijective (reverse the XORs and find the congruence inverses to the primes).
//!
//! The second guarantee requires more complex calculations: Construct a matrix of probabilities
//! and set one to certain (1), then apply transformations through the respective operations. The
//! proof is a bit long, but relatively simple.
//!
//! # Inner workings
//!
//! In technical terms, SeaHash follows a alternating 4-state length-padded Merkle–Damgård
//! construction with an add-diffuse compression function:
//!
//! ![A diagram.](http://ticki.github.io/img/seahash_construction_diagram.svg)
//!
//! It starts with 4 initial states, then it alternates between them (increment, wrap on 4) and
//! does modular addition with the respective block. When a state has been visited the diffusion
//! function (f) is applied. The very last block is padded with zeros.
//!
//! After all the blocks have been gone over, all the states are added to the number of bytes
//! written. The sum is then passed through the diffusion function, which produces the final hash
//! value.
//!
//! The diffusion function is drawn below.
//!
//! ```notest
//! x ← x ≫ 32
//! x ← px
//! x ← x ≫ 32
//! x ← px
//! x ← x ≫ 32
//! ```
//!
//! The advantage of having four completely segregated (note that there is no mix round, so they're
//! entirely independent) states is that fast parallelism is possible. For example, if I were to
//! hash 1 TB, I can spawn up four threads which can run independently without _any_
//! intercommunication or syncronization before the last round.
//!
//! If the diffusion function (f) was cryptographically secure, it would pass cryptoanalysis
//! trivially. This might seem irrelavant, as it clearly isn't cryptographically secure, but it
//! tells us something about the inner semantics. In particular, any diffusion function with
//! sufficient statistical quality will make up a good hash function in this construction.

#![no_std]
#![warn(missing_docs)]

use core::num::Wrapping as W;

pub use buffer::hash;
pub use stream::SeaHasher;

pub mod reference;
mod buffer;
mod stream;

/// The diffusion function.
///
/// This is a bijective function emitting chaotic behavior. Such functions are used as building
/// blocks for hash functions.
fn diffuse(mut x: W<u64>) -> W<u64> {
    // Move entropy down by XOR with shifting.
    x ^= x >> 32;
    // Move entropy up by scattering through multiplication.
    x *= W(0x7ed0e9fa0d94a33);
    // We still need more entropy downwards. Flipping higher bits won't flip lower ones, so far.
    // For example, if you flip the most significant bit, the 32'th bit will flip per the XOR-shift
    // subdiffusion, but this flip will only be scattered by the multiplication to flipping bits
    // higher than the 32'th, meaning that the ones lower will be unaffected. As such, we need to
    // get some entropy down.
    x ^= x >> 32;
    // So far, the avalanche diagram looks pretty good, but it still emits stripe patterns. For
    // example, flipping the 5'th lowest bit won't flip the least significant bit because of the
    // choice of scalar (in particular, observe how it leaves the 32'th bit unflipped after the
    // multiplication, which means that the XOR-shift never affects the lowest bit). No choice of
    // scalar will make this go away, it will merely change the unaffected bits. Instead, we need
    // to make the behavior more undeterministic by scattering bits through multiplication.
    x *= W(0x7ed0e9fa0d94a33);
    // This is the final stage of the diffusion function. There are still issues with the lowest
    // bits, which are still unaffected by the multiplication above. However, the multiplication
    // solved the higher bits' dependence, so lending entropy from the higher half will fix the
    // issues with the lower half.
    x ^= x >> 32;

    x
}
<|endoftext|>


<commit_before>use crate::api::{github, ServerResult};
use crate::load::{Config, InputData, TryCommit};
use futures::compat::Future01CompatExt;
use serde::Deserialize;

use regex::Regex;
use reqwest::header::USER_AGENT;

lazy_static::lazy_static! {
    static ref BODY_TRY_COMMIT: Regex =
        Regex::new(r#"(?:\W|^)@rust-timer\s+build\s+(\w+)(?:\W|$)"#).unwrap();
    static ref BODY_QUEUE: Regex =
        Regex::new(r#"(?:\W|^)@rust-timer\s+queue(?:\W|$)"#).unwrap();
}

async fn get_authorized_users() -> ServerResult<Vec<usize>> {
    let url = format!("{}/permissions/perf.json", ::rust_team_data::v1::BASE_URL);
    let client = reqwest::r#async::Client::new();
    client
        .get(&url)
        .send()
        .compat()
        .await
        .map_err(|err| format!("failed to fetch authorized users: {}", err))?
        .error_for_status()
        .map_err(|err| format!("failed to fetch authorized users: {}", err))?
        .json::<rust_team_data::v1::Permission>()
        .compat()
        .await
        .map_err(|err| format!("failed to fetch authorized users: {}", err))
        .map(|perms| perms.github_ids)
}

pub async fn handle_github(
    request: github::Request,
    data: &InputData,
) -> ServerResult<github::Response> {
    if request.comment.body.contains(" homu: ") {
        if let Some(sha) = handle_homu_res(&request, data).await {
            return enqueue_sha(request, data, &sha).await;
        }
    }

    if !request.comment.body.contains("@rust-timer ") {
        return Ok(github::Response);
    }

    if request.comment.author_association != github::Association::Owner
        && !get_authorized_users()
            .await?
            .contains(&request.comment.user.id)
    {
        post_comment(
            &data.config,
            &request.issue,
            "Insufficient permissions to issue commands to rust-timer.",
        )
        .await;
        return Ok(github::Response);
    }

    let body = &request.comment.body;

    if BODY_QUEUE.is_match(&body) {
        {
            let mut persistent = data.persistent.lock();
            persistent.pending_try_builds.insert(request.issue.number);
            persistent.write().expect("successful encode");
        }
        post_comment(
            &data.config,
            &request.issue,
            "Successfully queued build, kicking off bors: @bors try",
        )
        .await;
        return Ok(github::Response);
    }

    if let Some(captures) = BODY_TRY_COMMIT.captures(&body) {
        if let Some(commit) = captures.get(1).map(|c| c.as_str().to_owned()) {
            let commit = commit.trim_start_matches("https://github.com/rust-lang/rust/commit/");
            return enqueue_sha(request, data, &commit).await;
        }
    }

    Ok(github::Response)
}

async fn enqueue_sha(
    request: github::Request,
    data: &InputData,
    commit: &str,
) -> ServerResult<github::Response> {
    let client = reqwest::r#async::Client::new();
    let mut commit_response = client
        .get(&format!(
            "{}/commits/{}",
            request.issue.repository_url, commit
        ))
        .send()
        .compat()
        .await
        .map_err(|_| String::from("cannot get commit"))?;
    let commit_response = commit_response
        .json::<github::Commit>()
        .compat()
        .await
        .map_err(|e| format!("cannot deserialize commit: {:?}", e))?;
    if commit_response.parents.len() != 2 {
        post_comment(
            &data.config,
            &request.issue,
            format!(
                "Bors try commit {} unexpectedly has {} parents.",
                commit_response.sha,
                commit_response.parents.len()
            ),
        )
        .await;
        return Ok(github::Response);
    }
    let try_commit = TryCommit {
        sha: commit_response.sha.clone(),
        parent_sha: commit_response.parents[0].sha.clone(),
        issue: request.issue.clone(),
    };
    {
        let mut persistent = data.persistent.lock();
        if !persistent
            .try_commits
            .iter()
            .any(|c| c.sha() == &commit_response.sha)
        {
            persistent.try_commits.push(try_commit.clone());
        }
        persistent.write().expect("successful encode");
    }
    post_comment(
        &data.config,
        &request.issue,
        format!(
            "Success: Queued {} with parent {}, [comparison URL]({}).",
            commit_response.sha,
            commit_response.parents[0].sha,
            try_commit.comparison_url(),
        ),
    )
    .await;
    Ok(github::Response)
}

#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
enum HomuComment {
    TryBuildCompleted { merge_sha: String },
}

async fn handle_homu_res(request: &github::Request, data: &InputData) -> Option<String> {
    if !request.comment.body.contains("Try build successful") {
        return None;
    }

    {
        let mut persistent = data.persistent.lock();
        if persistent.pending_try_builds.remove(&request.issue.number) {
            persistent.write().expect("successful encode");
        } else {
            log::debug!(
                "Skipping successful try build for pr {}, not in pending",
                request.issue.number
            );
            return None;
        }
    }

    let start = "<!-- homu: ";
    let start_idx = request.comment.body.find(start).expect("found homu") + start.len();
    let end_idx = start_idx + request.comment.body[start_idx..].find(" -->").unwrap();

    let sha = match serde_json::from_str(&request.comment.body[start_idx..end_idx]) {
        Ok(HomuComment::TryBuildCompleted { merge_sha }) => merge_sha,
        Err(err) => {
            log::warn!(
                "failed to parse try build result; comment: {:?}, part: {:?}, err: {:?}",
                request.comment.body,
                &request.comment.body[start_idx..end_idx],
                err
            );
            return None;
        }
    };

    Some(sha)
}

pub async fn post_comment<B>(cfg: &Config, issue: &github::Issue, body: B)
where
    B: Into<String>,
{
    let body = body.into();
    let timer_token = cfg.keys.github.clone().expect("needs rust-timer token");
    let client = reqwest::r#async::Client::new();
    let req = client
        .post(&issue.comments_url)
        .json(&github::PostComment {
            body: body.to_owned(),
        })
        .header(USER_AGENT, "perf-rust-lang-org-server")
        .basic_auth("rust-timer", Some(timer_token));

    if let Err(e) = req.send().compat().await {
        eprintln!("failed to post comment: {:?}", e);
    }
}
<commit_msg>Remove bors try invocation<commit_after>use crate::api::{github, ServerResult};
use crate::load::{Config, InputData, TryCommit};
use futures::compat::Future01CompatExt;
use serde::Deserialize;

use regex::Regex;
use reqwest::header::USER_AGENT;

lazy_static::lazy_static! {
    static ref BODY_TRY_COMMIT: Regex =
        Regex::new(r#"(?:\W|^)@rust-timer\s+build\s+(\w+)(?:\W|$)"#).unwrap();
    static ref BODY_QUEUE: Regex =
        Regex::new(r#"(?:\W|^)@rust-timer\s+queue(?:\W|$)"#).unwrap();
}

async fn get_authorized_users() -> ServerResult<Vec<usize>> {
    let url = format!("{}/permissions/perf.json", ::rust_team_data::v1::BASE_URL);
    let client = reqwest::r#async::Client::new();
    client
        .get(&url)
        .send()
        .compat()
        .await
        .map_err(|err| format!("failed to fetch authorized users: {}", err))?
        .error_for_status()
        .map_err(|err| format!("failed to fetch authorized users: {}", err))?
        .json::<rust_team_data::v1::Permission>()
        .compat()
        .await
        .map_err(|err| format!("failed to fetch authorized users: {}", err))
        .map(|perms| perms.github_ids)
}

pub async fn handle_github(
    request: github::Request,
    data: &InputData,
) -> ServerResult<github::Response> {
    if request.comment.body.contains(" homu: ") {
        if let Some(sha) = handle_homu_res(&request, data).await {
            return enqueue_sha(request, data, &sha).await;
        }
    }

    if !request.comment.body.contains("@rust-timer ") {
        return Ok(github::Response);
    }

    if request.comment.author_association != github::Association::Owner
        && !get_authorized_users()
            .await?
            .contains(&request.comment.user.id)
    {
        post_comment(
            &data.config,
            &request.issue,
            "Insufficient permissions to issue commands to rust-timer.",
        )
        .await;
        return Ok(github::Response);
    }

    let body = &request.comment.body;

    if BODY_QUEUE.is_match(&body) {
        {
            let mut persistent = data.persistent.lock();
            persistent.pending_try_builds.insert(request.issue.number);
            persistent.write().expect("successful encode");
        }
        post_comment(
            &data.config,
            &request.issue,
            "Awaiting bors try build completion",
        )
        .await;
        return Ok(github::Response);
    }

    if let Some(captures) = BODY_TRY_COMMIT.captures(&body) {
        if let Some(commit) = captures.get(1).map(|c| c.as_str().to_owned()) {
            let commit = commit.trim_start_matches("https://github.com/rust-lang/rust/commit/");
            return enqueue_sha(request, data, &commit).await;
        }
    }

    Ok(github::Response)
}

async fn enqueue_sha(
    request: github::Request,
    data: &InputData,
    commit: &str,
) -> ServerResult<github::Response> {
    let client = reqwest::r#async::Client::new();
    let mut commit_response = client
        .get(&format!(
            "{}/commits/{}",
            request.issue.repository_url, commit
        ))
        .send()
        .compat()
        .await
        .map_err(|_| String::from("cannot get commit"))?;
    let commit_response = commit_response
        .json::<github::Commit>()
        .compat()
        .await
        .map_err(|e| format!("cannot deserialize commit: {:?}", e))?;
    if commit_response.parents.len() != 2 {
        post_comment(
            &data.config,
            &request.issue,
            format!(
                "Bors try commit {} unexpectedly has {} parents.",
                commit_response.sha,
                commit_response.parents.len()
            ),
        )
        .await;
        return Ok(github::Response);
    }
    let try_commit = TryCommit {
        sha: commit_response.sha.clone(),
        parent_sha: commit_response.parents[0].sha.clone(),
        issue: request.issue.clone(),
    };
    {
        let mut persistent = data.persistent.lock();
        if !persistent
            .try_commits
            .iter()
            .any(|c| c.sha() == &commit_response.sha)
        {
            persistent.try_commits.push(try_commit.clone());
        }
        persistent.write().expect("successful encode");
    }
    post_comment(
        &data.config,
        &request.issue,
        format!(
            "Queued {} with parent {}, future [comparison URL]({}).",
            commit_response.sha,
            commit_response.parents[0].sha,
            try_commit.comparison_url(),
        ),
    )
    .await;
    Ok(github::Response)
}

#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
enum HomuComment {
    TryBuildCompleted { merge_sha: String },
}

async fn handle_homu_res(request: &github::Request, data: &InputData) -> Option<String> {
    if !request.comment.body.contains("Try build successful") {
        return None;
    }

    {
        let mut persistent = data.persistent.lock();
        if persistent.pending_try_builds.remove(&request.issue.number) {
            persistent.write().expect("successful encode");
        } else {
            log::debug!(
                "Skipping successful try build for pr {}, not in pending",
                request.issue.number
            );
            return None;
        }
    }

    let start = "<!-- homu: ";
    let start_idx = request.comment.body.find(start).expect("found homu") + start.len();
    let end_idx = start_idx + request.comment.body[start_idx..].find(" -->").unwrap();

    let sha = match serde_json::from_str(&request.comment.body[start_idx..end_idx]) {
        Ok(HomuComment::TryBuildCompleted { merge_sha }) => merge_sha,
        Err(err) => {
            log::warn!(
                "failed to parse try build result; comment: {:?}, part: {:?}, err: {:?}",
                request.comment.body,
                &request.comment.body[start_idx..end_idx],
                err
            );
            return None;
        }
    };

    Some(sha)
}

pub async fn post_comment<B>(cfg: &Config, issue: &github::Issue, body: B)
where
    B: Into<String>,
{
    let body = body.into();
    let timer_token = cfg.keys.github.clone().expect("needs rust-timer token");
    let client = reqwest::r#async::Client::new();
    let req = client
        .post(&issue.comments_url)
        .json(&github::PostComment {
            body: body.to_owned(),
        })
        .header(USER_AGENT, "perf-rust-lang-org-server")
        .basic_auth("rust-timer", Some(timer_token));

    if let Err(e) = req.send().compat().await {
        eprintln!("failed to post comment: {:?}", e);
    }
}
<|endoftext|>