<commit_before>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;
use extra::priority_queue;
use extra::md4;
use extra::sort;
use std::path;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

static MAX_CACHE_SIZE_BYTES: u64 = 50000000;
static CACHE_MANAGER_A_RATE: u64 = 2000;
static CACHE_MANAGER_B_RATE: u64 = 6000;

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    topPriority: int,
    fileSize: u64
}

impl Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {

        let mut retVal:bool = false;

        //If self has a lower geographic priority, then it is less than
        if (self.topPriority < other.topPriority) {
            retVal = true;
        }
        else if (self.topPriority == other.topPriority) {

            //If the two have equal priorities, then we will sort by file size
            //Smaller files get a higher priority
            if (self.fileSize > other.fileSize) {
                retVal = true;
            }
        }

        retVal
    }
}

struct cache_item {
    name: ~str,
    in_use_flag: bool,
    ssi_flag: bool,
    hash: ~str,
    data: ~[u8],
    count: uint,
    size: u64
}

fn le(this: &cache_item, other: &cache_item) -> bool {
        let mut retVal:bool = false;

        if(this.count <= other.count) {
            retVal = true;
        }

        retVal
    }

fn main() {
    let req_vec: ~priority_queue::PriorityQueue<sched_msg> = ~priority_queue::PriorityQueue::new();
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    //Variables pertaining to safe visitor counting
    let safe_count: ~[uint] = ~[0];
    let shared_visit = arc::RWArc::new(safe_count);

    //Variables pertaining to cache and cache manager
    let cache_list: ~[cache_item] = ~[];
    let shared_cache_list = arc::RWArc::new(cache_list);
    let cache_manager_a = shared_cache_list.clone();
    let cache_manager_b = shared_cache_list.clone();

    //CACHE UPDATE MANAGER (Manager B)
    //This will handle making sure that items in the cache are up-to-date in case they are changed
    do spawn {

        loop {

            do cache_manager_b.write |vec| {

                for i in range(0, (*vec).len()) {

                    //If it is in the cache and in use, we will check to see if the file has been updated
                    if((*vec)[i].in_use_flag) {

                        let curr_path = ~path::Path((*vec)[i].name);
                        if os::path_exists(curr_path) {

                            match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        let temp_md4 = md4::md4_str(file_data.to_owned());

                                        if(temp_md4 != (*vec)[i].hash)
                                        {
                                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                                Some(s) => s,
                                                None => fail!("Could not access file stats for cache")
                                            };

                                            println(fmt!("===== UPDATING FILE: %?", (*vec)[i].name));

                                            (*vec)[i].data = file_data;
                                            (*vec)[i].size = fileInfo.size;
                                            (*vec)[i].hash = temp_md4;
                                        }
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                        }
                    }
                }
            }

            timer::sleep(CACHE_MANAGER_B_RATE);
        }
    }

    //MAIN CACHE MANAGER (Manager A)
    do spawn {

        loop {
            do cache_manager_a.write |vec| {

                //Quick sort sorts in-place, so we don't need to worry about memory overhead
                //Just time overhead
                sort::quick_sort((*vec), le);

                let mut cache_remaining = MAX_CACHE_SIZE_BYTES;

                for i in range(0, (*vec).len()) {

                    if((*vec)[i].size <= cache_remaining && !(*vec)[i].in_use_flag) {

                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {

                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats for cache")
                            };

                            if(fileInfo.size <= cache_remaining) {

                                match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        //NEED TO ADD CHECK IF IT HAS SSI
                                        //PERHAPS EXCLUDE ALL SSI FROM CACHE?
                                        //ADD TO IF: size < remaining, !in use, !ssi?

                                        (*vec)[i].data = file_data.to_owned();
                                        (*vec)[i].size = fileInfo.size;
                                        (*vec)[i].hash = md4::md4_str(file_data);
                                        (*vec)[i].in_use_flag = true;

                                        cache_remaining = cache_remaining - fileInfo.size;
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                            }
                        }
                    }
                    else if((*vec)[i].size <= cache_remaining && (*vec)[i].in_use_flag) {
                        cache_remaining = cache_remaining - (*vec)[i].size;
                    }
                    else if((*vec)[i].size > cache_remaining) {
                        (*vec)[i].in_use_flag = false;
                        (*vec)[i].data = ~[];
                    }

                }
            }

            timer::sleep(CACHE_MANAGER_A_RATE);
        }
    }

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle

                //Check if file is in cache
                //Will do so using RWArc
                let mut serve_from_cache: bool = false;

                do shared_cache_list.write |vec| {

                    let mut found: bool = false;

                    for i in range(0, (*vec).len()) {

                        if( (*vec)[i].name == tf.filepath.to_str() && (*vec)[i].in_use_flag) {
                            serve_from_cache = true;
                            found = true;

                            println(fmt!("===== SERVING FROM CACHE: %?", tf.filepath.to_str()));

                            tf.stream.write((*vec)[i].data);

                            (*vec)[i].count += 1;
                        }
                        else if( (*vec)[i].name == tf.filepath.to_str() && !(*vec)[i].in_use_flag) {
                            (*vec)[i].count += 1;
                            found = true;
                        }
                    }

                    //If it isn't found, we will create a blank entry and enter basic info, and add the
                    //data and md4 later
                    if(!found) {
                        println(fmt!("===== ADDING ITEM %?", tf.filepath.to_str()));

                        let new_cache_item: cache_item = cache_item{name: tf.filepath.to_str(), in_use_flag: false,
                            ssi_flag: false, hash: ~"", data: ~[], count: 1, size: tf.fileSize};

                        (*vec).push(new_cache_item);
                    }
                }

                if(!serve_from_cache) {

                    match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                        Ok(file_data) => {
                            tf.stream.write(file_data);
                            println(fmt!("===== SERVING FROM DISK: %?", tf.filepath.to_str()));

                        }
                        Err(err) => {
                            println("ERROR IN SEND");
                            println(err);
                        }
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {

                    //Since we are using a priority queue, we will use pop()
                    let tf = (*vec).pop();
                    //println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    loop {
        let stream = acceptor.accept();
        let stream = Cell::new(stream);

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();
        let child_arc = shared_visit.clone();

        do spawn {

            let mut actual_count:uint = 0;

            //We will use the cloned RWARC
            do child_arc.write |count_vec| {

                let prev_count = (*count_vec).pop();
                let new_count: uint = prev_count + 1;

                actual_count = new_count;

                (*count_vec).push(new_count);
            }

            let stream = stream.take();

            match stream {
                Some(s) => {

                    let mut stream = s;

                    let mut buf = [0, ..500];
                    stream.read(buf);
                    let request_str = str::from_utf8(buf);

                    let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
                    if req_group.len() > 2 {
                        let path = req_group[1];
                        //println(fmt!("Request for path: \n%?", path));

                        let file_path = ~os::getcwd().push(path.replace("/../", ""));

                        //TO DO: Maybe replace this check to later on since this reduces our response time by
                        //having to check the disk.

                        if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                            //println(fmt!("Request received:\n%s", request_str));
                            let response: ~str = fmt!(
                                "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                                 <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                                 \r\n", actual_count);

                            stream.write(response.as_bytes());
                        }
                        else {
                            // Requests scheduling

                            let mut streamPriority: int = 0;

                            //Retrieving the requesting IP address
                            let ipStr: ~str = match (stream).peer_name() {
                                Some(pr) => pr.ip.to_str(),
                                None => ~"0.0.0.0"
                            };

                            //Split the IP address by '.' so that we can compare
                            let ipSplit: ~[~str] = ipStr.split_iter('.').filter(|&x| x != "")
                                 .map(|x| x.to_owned()).collect();

                            //Assign priority based on geography or if localhost
                            if ( (ipSplit[0] == ~"127" && ipSplit[1] == ~"0") || (ipSplit[0] == ~"128" && ipSplit[1] == ~"143")
                                || (ipSplit[0] == ~"137" && ipSplit[1] == ~"54") ) {

                                streamPriority = 1;
                            }

                            let fileName: ~str = file_path.filename().unwrap().to_owned();
                            let fileNameSplit: ~[~str] = fileName.split_iter('.').filter(|&x| x != "").map(|x| x.to_owned()).collect();

                            //In order to optimize for the benchmark, we will send the HTTP header quickly before adding to the queue
                            let httpHeader: ~str = match fileNameSplit[fileNameSplit.len()-1] {
                                ~"html" | ~"htm" | ~"php" => ~"Content-Type: text/html; charset=UTF-8\r\n\r\n",
                                _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                            };

                            stream.write(httpHeader.as_bytes());
                            stream.flush();

                            //Retrieve file info for additional latency fixes
                            let fileInfo = match std::rt::io::file::stat(file_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats")
                            };

                            let msg: sched_msg = sched_msg{stream: Some(stream), filepath: file_path.clone(), topPriority: streamPriority, fileSize: fileInfo.size};
                            let (sm_port, sm_chan) = std::comm::stream();
                            sm_chan.send(msg);

                            do child_add_vec.write |vec| {
                                let msg = sm_port.recv();
                                (*vec).push(msg); // enqueue new request.
                                //println("add to queue");
                            }
                            child_chan.send(""); //notify the new arriving request.
                            //println(fmt!("get file request: %?", file_path));
                        }
                    }
                    println!("connection terminates")

                },
                None => ()
            }
        }
    }
}<commit_msg>Cleaned up more code<commit_after>//
// zhtta.rs
//
// Running on Rust 0.8
//
// Starting code for PS3
//
// Note: it would be very unwise to run this server on a machine that is
// on the Internet and contains any sensitive files!
//
// University of Virginia - cs4414 Fall 2013
// Weilin Xu and David Evans
// Version 0.3

extern mod extra;

use std::rt::io::*;
use std::rt::io::net::ip::SocketAddr;
use std::io::println;
use std::cell::Cell;
use std::{os, str, io};
use extra::arc;
use std::comm::*;
use extra::priority_queue;
use extra::md4;
use extra::sort;
use std::path;

static PORT:    int = 4414;
static IP: &'static str = "127.0.0.1";

static MAX_CACHE_SIZE_BYTES: u64 = 50000000;
static CACHE_MANAGER_A_RATE: u64 = 2000;
static CACHE_MANAGER_B_RATE: u64 = 6000;

struct sched_msg {
    stream: Option<std::rt::io::net::tcp::TcpStream>,
    filepath: ~std::path::PosixPath,
    topPriority: int,
    fileSize: u64
}

impl Ord for sched_msg {
    fn lt(&self, other: &sched_msg) -> bool {

        let mut retVal:bool = false;

        //If self has a lower geographic priority, then it is less than
        if (self.topPriority < other.topPriority) {
            retVal = true;
        }
        else if (self.topPriority == other.topPriority) {

            //If the two have equal priorities, then we will sort by file size
            //Smaller files get a higher priority
            if (self.fileSize > other.fileSize) {
                retVal = true;
            }
        }

        retVal
    }
}

struct cache_item {
    name: ~str,
    in_use_flag: bool,
    ssi_flag: bool,
    hash: ~str,
    data: ~[u8],
    count: uint,
    size: u64
}

fn le(this: &cache_item, other: &cache_item) -> bool {
        let mut retVal:bool = false;

        if(this.count <= other.count) {
            retVal = true;
        }

        retVal
    }

fn main() {
    let req_vec: ~priority_queue::PriorityQueue<sched_msg> = ~priority_queue::PriorityQueue::new();
    let shared_req_vec = arc::RWArc::new(req_vec);
    let add_vec = shared_req_vec.clone();
    let take_vec = shared_req_vec.clone();

    let (port, chan) = stream();
    let chan = SharedChan::new(chan);

    //Variables pertaining to safe visitor counting
    let safe_count: ~[uint] = ~[0];
    let shared_visit = arc::RWArc::new(safe_count);

    //Variables pertaining to cache and cache manager
    let cache_list: ~[cache_item] = ~[];
    let shared_cache_list = arc::RWArc::new(cache_list);
    let cache_manager_a = shared_cache_list.clone();
    let cache_manager_b = shared_cache_list.clone();

    //CACHE UPDATE MANAGER (Manager B)
    //This will handle making sure that items in the cache are up-to-date in case they are changed
    do spawn {
        loop {
            do cache_manager_b.write |vec| {
                for i in range(0, (*vec).len()) {
                    //If it is in the cache and in use, we will check to see if the file has been updated
                    if((*vec)[i].in_use_flag) {
                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {
                            match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {
                                        let temp_md4 = md4::md4_str(file_data.to_owned());

                                        if(temp_md4 != (*vec)[i].hash)
                                        {
                                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                                Some(s) => s,
                                                None => fail!("Could not access file stats for cache")
                                            };

                                            println(fmt!("===== UPDATING FILE: %?", (*vec)[i].name));

                                            (*vec)[i].data = file_data;
                                            (*vec)[i].size = fileInfo.size;
                                            (*vec)[i].hash = temp_md4;
                                        }
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                        }
                    }
                }
            }

            timer::sleep(CACHE_MANAGER_B_RATE);
        }
    }

    //MAIN CACHE MANAGER (Manager A)
    do spawn {
        loop {
            do cache_manager_a.write |vec| {
                //Quick sort sorts in-place, so we don't need to worry about memory overhead
                //Just time overhead
                sort::quick_sort((*vec), le);

                let mut cache_remaining = MAX_CACHE_SIZE_BYTES;

                for i in range(0, (*vec).len()) {
                    if((*vec)[i].size <= cache_remaining && !(*vec)[i].in_use_flag) {
                        let curr_path = ~path::Path((*vec)[i].name);

                        if os::path_exists(curr_path) {
                            let fileInfo = match std::rt::io::file::stat(curr_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats for cache")
                            };

                            if(fileInfo.size <= cache_remaining) {
                                match io::read_whole_file(curr_path) {
                                    Ok(file_data) => {

                                        //NEED TO ADD CHECK IF IT HAS SSI
                                        //PERHAPS EXCLUDE ALL SSI FROM CACHE?
                                        //ADD TO IF: size < remaining, !in use, !ssi?

                                        (*vec)[i].data = file_data.to_owned();
                                        (*vec)[i].size = fileInfo.size;
                                        (*vec)[i].hash = md4::md4_str(file_data);
                                        (*vec)[i].in_use_flag = true;

                                        cache_remaining = cache_remaining - fileInfo.size;
                                    }
                                    Err(err) => {
                                        println("ERROR IN UPDATE CACHE");
                                        println(err);
                                    }
                                }
                            }
                        }
                    }
                    else if((*vec)[i].size <= cache_remaining && (*vec)[i].in_use_flag) {
                        cache_remaining = cache_remaining - (*vec)[i].size;
                    }
                    else if((*vec)[i].size > cache_remaining) {
                        (*vec)[i].in_use_flag = false;
                        (*vec)[i].data = ~[];
                    }

                }
            }

            timer::sleep(CACHE_MANAGER_A_RATE);
        }
    }

    // dequeue file requests, and send responses.
    // FIFO
    do spawn {
        let (sm_port, sm_chan) = stream();

        // a task for sending responses.
        do spawn {
            loop {
                let mut tf: sched_msg = sm_port.recv(); // wait for the dequeued request to handle

                //Check if file is in cache
                //Will do so using RWArc
                let mut serve_from_cache: bool = false;

                do shared_cache_list.write |vec| {
                    let mut found: bool = false;

                    for i in range(0, (*vec).len()) {
                        if( (*vec)[i].name == tf.filepath.to_str() && (*vec)[i].in_use_flag) {
                            serve_from_cache = true;
                            found = true;

                            println(fmt!("===== SERVING FROM CACHE: %?", tf.filepath.to_str()));

                            tf.stream.write((*vec)[i].data);

                            (*vec)[i].count += 1;
                        }
                        else if( (*vec)[i].name == tf.filepath.to_str() && !(*vec)[i].in_use_flag) {
                            (*vec)[i].count += 1;
                            found = true;
                        }
                    }

                    //If it isn't found, we will create a blank entry and enter basic info, and add the
                    //data and md4 later
                    if(!found) {
                        println(fmt!("===== ADDING ITEM %?", tf.filepath.to_str()));

                        let new_cache_item: cache_item = cache_item{name: tf.filepath.to_str(), in_use_flag: false,
                            ssi_flag: false, hash: ~"", data: ~[], count: 1, size: tf.fileSize};

                        (*vec).push(new_cache_item);
                    }
                }

                if(!serve_from_cache) {
                    match io::read_whole_file(tf.filepath) { // killed if file size is larger than memory size.
                        Ok(file_data) => {
                            tf.stream.write(file_data);
                            println(fmt!("===== SERVING FROM DISK: %?", tf.filepath.to_str()));

                        }
                        Err(err) => {
                            println("ERROR IN SEND");
                            println(err);
                        }
                    }
                }
            }
        }

        loop {
            port.recv(); // wait for arrving notification
            do take_vec.write |vec| {
                if ((*vec).len() > 0) {

                    //Since we are using a priority queue, we will use pop()
                    let tf = (*vec).pop();
                    //println(fmt!("shift from queue, size: %ud", (*vec).len()));
                    sm_chan.send(tf); // send the request to send-response-task to serve.
                }
            }
        }
    }

    let ip = match FromStr::from_str(IP) { Some(ip) => ip,
                                           None => { println(fmt!("Error: Invalid IP address <%s>", IP));
                                                     return;},
                                         };

    let socket = net::tcp::TcpListener::bind(SocketAddr {ip: ip, port: PORT as u16});

    println(fmt!("Listening on %s:%d ...", ip.to_str(), PORT));
    let mut acceptor = socket.listen().unwrap();

    loop {
        let stream = acceptor.accept();
        let stream = Cell::new(stream);

        // Start a new task to handle the each connection
        let child_chan = chan.clone();
        let child_add_vec = add_vec.clone();
        let child_arc = shared_visit.clone();

        do spawn {

            let mut actual_count:uint = 0;

            //We will use the cloned RWARC
            do child_arc.write |count_vec| {
                let prev_count = (*count_vec).pop();
                let new_count: uint = prev_count + 1;

                actual_count = new_count;
                (*count_vec).push(new_count);
            }

            let stream = stream.take();

            match stream {
                Some(s) => {

                    let mut stream = s;

                    let mut buf = [0, ..500];
                    stream.read(buf);
                    let request_str = str::from_utf8(buf);

                    let req_group : ~[&str]= request_str.splitn_iter(' ', 3).collect();
                    if req_group.len() > 2 {
                        let path = req_group[1];

                        let file_path = ~os::getcwd().push(path.replace("/../", ""));

                        if !os::path_exists(file_path) || os::path_is_dir(file_path) {
                            let response: ~str = fmt!(
                                "HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=UTF-8\r\n\r\n
                                 <doctype !html><title>Hello, Rust!</title>

Greetings, Krusty!

Visitor count: %u

                                 \r\n", actual_count);

                            stream.write(response.as_bytes());
                        }
                        else {
                            // Requests scheduling

                            let mut streamPriority: int = 0;

                            //Retrieving the requesting IP address
                            let ipStr: ~str = match (stream).peer_name() {
                                Some(pr) => pr.ip.to_str(),
                                None => ~"0.0.0.0"
                            };

                            //Split the IP address by '.' so that we can compare
                            let ipSplit: ~[~str] = ipStr.split_iter('.').filter(|&x| x != "")
                                 .map(|x| x.to_owned()).collect();

                            //Assign priority based on geography or if localhost
                            if ( (ipSplit[0] == ~"127" && ipSplit[1] == ~"0") || (ipSplit[0] == ~"128" && ipSplit[1] == ~"143")
                                || (ipSplit[0] == ~"137" && ipSplit[1] == ~"54") ) {

                                streamPriority = 1;
                            }

                            let fileName: ~str = file_path.filename().unwrap().to_owned();
                            let fileNameSplit: ~[~str] = fileName.split_iter('.').filter(|&x| x != "").map(|x| x.to_owned()).collect();

                            //In order to optimize for the benchmark, we will send the HTTP header quickly before adding to the queue
                            let httpHeader: ~str = match fileNameSplit[fileNameSplit.len()-1] {
                                ~"html" | ~"htm" | ~"php" => ~"Content-Type: text/html; charset=UTF-8\r\n\r\n",
                                _ => ~"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream; charset=UTF-8\r\n\r\n"
                            };

                            stream.write(httpHeader.as_bytes());
                            stream.flush();

                            //Retrieve file info for additional latency fixes
                            let fileInfo = match std::rt::io::file::stat(file_path) {
                                Some(s) => s,
                                None => fail!("Could not access file stats")
                            };

                            let msg: sched_msg = sched_msg{stream: Some(stream), filepath: file_path.clone(), topPriority: streamPriority, fileSize: fileInfo.size};
                            let (sm_port, sm_chan) = std::comm::stream();
                            sm_chan.send(msg);

                            do child_add_vec.write |vec| {
                                let msg = sm_port.recv();
                                (*vec).push(msg); // enqueue new request.
                            }
                            child_chan.send(""); //notify the new arriving request.
                        }
                    }
                    println!("connection terminates")

                },
                None => ()
            }
        }
    }
}<|endoftext|>


<commit_before>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::cmp::min;

use crypto::{Hash, PublicKey, HASH_SIZE};
use super::super::StorageKey;

pub const BRANCH_KEY_PREFIX: u8 = 00;
pub const LEAF_KEY_PREFIX: u8 = 01;

/// Size in bytes of the `ProofMapKey`.
pub const KEY_SIZE: usize = HASH_SIZE;
pub const DB_KEY_SIZE: usize = KEY_SIZE + 2;

/// A trait that defines a subset of storage key types which are suitable for use with
/// `ProofMapIndex`.
///
/// The size of the keys must be exactly 32 bytes and the keys must have a uniform distribution.
pub trait ProofMapKey: StorageKey {}

impl ProofMapKey for Hash {}
impl ProofMapKey for PublicKey {}
impl ProofMapKey for [u8; KEY_SIZE] {}

impl StorageKey for [u8; KEY_SIZE] {
    fn size(&self) -> usize {
        KEY_SIZE
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self.as_ref())
    }

    fn read(buffer: &[u8]) -> Self {
        let mut value = [0; KEY_SIZE];
        value.copy_from_slice(buffer);
        value
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ChildKind {
    Left,
    Right,
}

/// A struct that represents a bit slices of the proof map keys.
#[derive(Clone)]
pub struct DBKey {
    data: [u8; KEY_SIZE],
    from: u16,
    to: u16,
}

impl ::std::ops::Not for ChildKind {
    type Output = ChildKind;

    fn not(self) -> ChildKind {
        match self {
            ChildKind::Left => ChildKind::Right,
            ChildKind::Right => ChildKind::Left,
        }
    }
}

impl DBKey {
    /// Create a new bit slice from the given binary data.
    pub fn leaf<K: ProofMapKey>(key: &K) -> DBKey {
        debug_assert_eq!(key.size(), KEY_SIZE);

        let mut data = [0; KEY_SIZE];
        key.write(&mut data);
        DBKey {
            data: data,
            from: 0,
            to: (KEY_SIZE * 8) as u16,
        }
    }

    #[doc(hidden)]
    pub fn from(&self) -> u16 {
        self.from
    }

    // TODO: terrible hack, try to remove this (ECR-22)
    #[doc(hidden)]
    pub fn set_from(&mut self, from: u16) {
        self.from = from
    }

    #[doc(hidden)]
    pub fn to(&self) -> u16 {
        self.to
    }

    /// Returns length of the `DBKey`.
    pub fn len(&self) -> u16 {
        self.to - self.from
    }

    /// Returns true if `DBKey` has zero length.
    pub fn is_empty(&self) -> bool {
        self.to == self.from
    }

    /// Get bit at position `idx`.
    pub fn get(&self, idx: u16) -> ChildKind {
        debug_assert!(self.from + idx < self.to);

        let pos = self.from + idx;
        let chunk = self.data[(pos / 8) as usize];
        let bit = 7 - pos % 8;
        let value = (1 << bit) & chunk;
        if value != 0 {
            ChildKind::Right
        } else {
            ChildKind::Left
        }
    }

    /// Shortens this DBKey to the specified length.
    pub fn prefix(&self, mid: u16) -> DBKey {
        DBKey {
            data: self.data,
            from: self.from,
            to: self.from + mid,
        }
    }

    /// Return object which represents a view on to this slice (further) offset by `i` bits.
    pub fn suffix(&self, mid: u16) -> DBKey {
        debug_assert!(self.from + mid <= self.to);

        DBKey {
            data: self.data,
            from: self.from + mid,
            to: self.to,
        }
    }

    /// Shortens this DBKey to the specified length.
    pub fn truncate(&self, size: u16) -> DBKey {
        debug_assert!(self.from + size <= self.to);

        DBKey {
            data: self.data,
            from: self.from,
            to: self.from + size,
        }
    }

    /// Returns how many bits at the beginning matches with `other`
    pub fn common_prefix(&self, other: &Self) -> u16 {
        // We assume that all slices created from byte arrays with the same length
        if self.from != other.from {
            0
        } else {
            let from = self.from / 8;
            let to = min((self.to + 7) / 8, (other.to + 7) / 8);
            let max_len = min(self.len(), other.len());

            for i in from..to {
                let x = self.data[i as usize] ^ other.data[i as usize];
                if x != 0 {
                    let tail = x.leading_zeros() as u16;
                    return min(i * 8 + tail - self.from, max_len);
                }
            }
            max_len
        }
    }

    /// Returns true if we starts with the same prefix at the whole of `Other`
    pub fn starts_with(&self, other: &Self) -> bool {
        self.common_prefix(other) == other.len()
    }

    /// Returns true if self.to not changed
    pub fn is_leaf(&self) -> bool {
        self.to == (KEY_SIZE * 8) as u16
    }

    // TODO: terrible hack, try to remove this (ECR-22)
    /// Represents `DBKey` as byte array and returns it
    pub fn as_bytes(&self) -> Box<[u8]> {
        let mut buffer = Box::new([0u8; DB_KEY_SIZE as usize]);
        self.write(buffer.as_mut());
        buffer
    }
}

impl AsRef<[u8]> for DBKey {
    fn as_ref(&self) -> &[u8] {
        &self.data
    }
}

impl StorageKey for DBKey {
    fn size(&self) -> usize {
        DB_KEY_SIZE
    }

    fn write(&self, buffer: &mut [u8]) {
        if self.is_leaf() {
            buffer[0] = LEAF_KEY_PREFIX;
            buffer[1..KEY_SIZE + 1].copy_from_slice(&self.data);
            buffer[KEY_SIZE + 1] = 0;
        } else {
            buffer[0] = BRANCH_KEY_PREFIX;
            let right = (self.to as usize + 7) / 8;
            buffer[1..right + 1].copy_from_slice(&self.data[0..right]);
            if self.to % 8 != 0 {
                buffer[right] &= !(255u8 >> (self.to % 8));
            }
            for i in buffer.iter_mut().take(KEY_SIZE + 1).skip(right + 1) {
                *i = 0
            }
            buffer[KEY_SIZE + 1] = self.to as u8;
        }
    }

    fn read(buffer: &[u8]) -> Self {
        let mut data = [0; KEY_SIZE];
        data[..].copy_from_slice(&buffer[1..KEY_SIZE + 1]);
        let to = match buffer[0] {
            LEAF_KEY_PREFIX => KEY_SIZE as u16 * 8,
            BRANCH_KEY_PREFIX => u16::from(buffer[DB_KEY_SIZE - 1]),
            _ => unreachable!("wrong key prefix"),
        };
        DBKey {
            data: data,
            from: 0,
            to: to,
        }
    }
}

impl PartialEq for DBKey {
    fn eq(&self, other: &Self) -> bool {
        self.len() == other.len() && self.starts_with(other)
    }
}

impl ::std::fmt::Debug for DBKey {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "DBKey(")?;
        for i in 0..self.len() {
            write!(
                f,
                "{}",
                match self.get(i) {
                    ChildKind::Left => '0',
                    ChildKind::Right => '1',
                }
            )?;
        }
        write!(f, ")")
    }
}
<commit_msg>Restore DBKey tests<commit_after>// Copyright 2017 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::cmp::min;

use crypto::{Hash, PublicKey, HASH_SIZE};
use super::super::StorageKey;

pub const BRANCH_KEY_PREFIX: u8 = 00;
pub const LEAF_KEY_PREFIX: u8 = 01;

/// Size in bytes of the `ProofMapKey`.
pub const KEY_SIZE: usize = HASH_SIZE;
pub const DB_KEY_SIZE: usize = KEY_SIZE + 2;

/// A trait that defines a subset of storage key types which are suitable for use with
/// `ProofMapIndex`.
///
/// The size of the keys must be exactly 32 bytes and the keys must have a uniform distribution.
pub trait ProofMapKey: StorageKey {}

impl ProofMapKey for Hash {}
impl ProofMapKey for PublicKey {}
impl ProofMapKey for [u8; KEY_SIZE] {}

impl StorageKey for [u8; KEY_SIZE] {
    fn size(&self) -> usize {
        KEY_SIZE
    }

    fn write(&self, buffer: &mut [u8]) {
        buffer.copy_from_slice(self.as_ref())
    }

    fn read(buffer: &[u8]) -> Self {
        let mut value = [0; KEY_SIZE];
        value.copy_from_slice(buffer);
        value
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ChildKind {
    Left,
    Right,
}

/// A struct that represents a bit slices of the proof map keys.
#[derive(Clone)]
pub struct DBKey {
    data: [u8; KEY_SIZE],
    from: u16,
    to: u16,
}

impl ::std::ops::Not for ChildKind {
    type Output = ChildKind;

    fn not(self) -> ChildKind {
        match self {
            ChildKind::Left => ChildKind::Right,
            ChildKind::Right => ChildKind::Left,
        }
    }
}

impl DBKey {
    /// Create a new bit slice from the given binary data.
    pub fn leaf<K: ProofMapKey>(key: &K) -> DBKey {
        debug_assert_eq!(key.size(), KEY_SIZE);

        let mut data = [0; KEY_SIZE];
        key.write(&mut data);
        DBKey {
            data: data,
            from: 0,
            to: (KEY_SIZE * 8) as u16,
        }
    }

    #[doc(hidden)]
    pub fn from(&self) -> u16 {
        self.from
    }

    // TODO: terrible hack, try to remove this (ECR-22)
    #[doc(hidden)]
    pub fn set_from(&mut self, from: u16) {
        self.from = from
    }

    #[doc(hidden)]
    pub fn to(&self) -> u16 {
        self.to
    }

    /// Returns length of the `DBKey`.
    pub fn len(&self) -> u16 {
        self.to - self.from
    }

    /// Returns true if `DBKey` has zero length.
    pub fn is_empty(&self) -> bool {
        self.to == self.from
    }

    /// Get bit at position `idx`.
    pub fn get(&self, idx: u16) -> ChildKind {
        debug_assert!(self.from + idx < self.to);

        let pos = self.from + idx;
        let chunk = self.data[(pos / 8) as usize];
        let bit = 7 - pos % 8;
        let value = (1 << bit) & chunk;
        if value != 0 {
            ChildKind::Right
        } else {
            ChildKind::Left
        }
    }

    /// Shortens this DBKey to the specified length.
    pub fn prefix(&self, suffix: u16) -> DBKey {
        debug_assert!(self.to + suffix <= self.data.len() as u16 * 8);

        DBKey {
            data: self.data,
            from: self.from,
            to: self.from + suffix,
        }
    }

    /// Return object which represents a view on to this slice (further) offset by `i` bits.
    pub fn suffix(&self, suffix: u16) -> DBKey {
        debug_assert!(self.from + suffix <= self.to);

        DBKey {
            data: self.data,
            from: self.from + suffix,
            to: self.to,
        }
    }

    /// Shortens this DBKey to the specified length.
    pub fn truncate(&self, size: u16) -> DBKey {
        debug_assert!(self.from + size <= self.to);

        DBKey {
            data: self.data,
            from: self.from,
            to: self.from + size,
        }
    }

    /// Returns how many bits at the beginning matches with `other`
    pub fn common_prefix(&self, other: &Self) -> u16 {
        // We assume that all slices created from byte arrays with the same length
        if self.from != other.from {
            0
        } else {
            let from = self.from / 8;
            let to = min((self.to + 7) / 8, (other.to + 7) / 8);
            let max_len = min(self.len(), other.len());

            for i in from..to {
                let x = self.data[i as usize] ^ other.data[i as usize];
                if x != 0 {
                    let tail = x.leading_zeros() as u16;
                    return min(i * 8 + tail - self.from, max_len);
                }
            }
            max_len
        }
    }

    /// Returns true if we starts with the same prefix at the whole of `Other`
    pub fn starts_with(&self, other: &Self) -> bool {
        self.common_prefix(other) == other.len()
    }

    /// Returns true if self.to not changed
    pub fn is_leaf(&self) -> bool {
        self.to == (KEY_SIZE * 8) as u16
    }

    // TODO: terrible hack, try to remove this (ECR-22)
    /// Represents `DBKey` as byte array and returns it
    pub fn as_bytes(&self) -> Box<[u8]> {
        let mut buffer = Box::new([0u8; DB_KEY_SIZE as usize]);
        self.write(buffer.as_mut());
        buffer
    }
}

impl AsRef<[u8]> for DBKey {
    fn as_ref(&self) -> &[u8] {
        &self.data
    }
}

impl StorageKey for DBKey {
    fn size(&self) -> usize {
        DB_KEY_SIZE
    }

    fn write(&self, buffer: &mut [u8]) {
        if self.is_leaf() {
            buffer[0] = LEAF_KEY_PREFIX;
            buffer[1..KEY_SIZE + 1].copy_from_slice(&self.data);
            buffer[KEY_SIZE + 1] = 0;
        } else {
            buffer[0] = BRANCH_KEY_PREFIX;
            let right = (self.to as usize + 7) / 8;
            buffer[1..right + 1].copy_from_slice(&self.data[0..right]);
            if self.to % 8 != 0 {
                buffer[right] &= !(255u8 >> (self.to % 8));
            }
            for i in buffer.iter_mut().take(KEY_SIZE + 1).skip(right + 1) {
                *i = 0
            }
            buffer[KEY_SIZE + 1] = self.to as u8;
        }
    }

    fn read(buffer: &[u8]) -> Self {
        debug_assert_eq!(buffer.len(), DB_KEY_SIZE);
        let mut data = [0; KEY_SIZE];
        data[..].copy_from_slice(&buffer[1..KEY_SIZE + 1]);
        let to = match buffer[0] {
            LEAF_KEY_PREFIX => KEY_SIZE as u16 * 8,
            BRANCH_KEY_PREFIX => u16::from(buffer[DB_KEY_SIZE - 1]),
            _ => unreachable!("wrong key prefix"),
        };
        DBKey {
            data: data,
            from: 0,
            to: to,
        }
    }
}

impl PartialEq for DBKey {
    fn eq(&self, other: &Self) -> bool {
        self.len() == other.len() && self.starts_with(other)
    }
}

impl ::std::fmt::Debug for DBKey {
    fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
        write!(f, "DBKey(")?;
        for i in 0..self.len() {
            write!(
                f,
                "{}",
                match self.get(i) {
                    ChildKind::Left => '0',
                    ChildKind::Right => '1',
                }
            )?;
        }
        write!(f, ")")
    }
}

#[test]
fn test_dbkey_suffix() {
    let b = DBKey::read(b"\x00\x01\x02\xFF\x0C0000000000000000000000000000\x20");
    assert_eq!(b.len(), 32);
    assert_eq!(b.get(0), ChildKind::Left);
    assert_eq!(b.get(7), ChildKind::Right);
    assert_eq!(b.get(8), ChildKind::Left);
    assert_eq!(b.get(14), ChildKind::Right);
    assert_eq!(b.get(15), ChildKind::Left);
    assert_eq!(b.get(16), ChildKind::Right);
    assert_eq!(b.get(20), ChildKind::Right);
    assert_eq!(b.get(23), ChildKind::Right);
    assert_eq!(b.get(31), ChildKind::Left);
    let b2 = b.suffix(8);
    assert_eq!(b2.len(), 24);
    assert_eq!(b2.get(0), ChildKind::Left);
    assert_eq!(b2.get(6), ChildKind::Right);
    assert_eq!(b2.get(7), ChildKind::Left);
    assert_eq!(b2.get(12), ChildKind::Right);
    assert_eq!(b2.get(15), ChildKind::Right);
    let b3 = b2.suffix(24);
    assert_eq!(b3.len(), 0);
    let b4 = b.suffix(1);
    assert_eq!(b4.get(6), ChildKind::Right);
    assert_eq!(b4.get(7), ChildKind::Left);
    assert_eq!(b4.get(13), ChildKind::Right);
}

#[test]
fn test_dbkey_truncate() {
    let b = DBKey::read(b"\x00\x80wertyuiopasdfghjklzxcvbnm123456\x08");
    assert_eq!(b.len(), 8);
    assert_eq!(b.truncate(1).get(0), ChildKind::Right);
    assert_eq!(b.truncate(1).len(), 1);
}

#[test]
fn test_dbkey_len() {
    let b = DBKey::read(b"\x01qwertyuiopasdfghjklzxcvbnm123456\x00");
    assert_eq!(b.len(), 256);
}

#[test]
#[should_panic(expected = "self.from + idx < self.to")]
fn test_dbkey_at_overflow() {
    let b = DBKey::read(b"\x00qwertyuiopasdfghjklzxcvbnm123456\x0F");
    b.get(32);
}

#[test]
#[should_panic(expected = "self.from + suffix <= self.to")]
fn test_dbkey_suffix_overflow() {
    let b = DBKey::read(b"\x00qwertyuiopasdfghjklzxcvbnm123456\xFF");
    assert_eq!(b"\x01qwertyuiopasdfghjklzxcvbnm123456\x00".len(), 34);
    b.suffix(255).suffix(2);
}

#[test]
#[should_panic(expected = "self.from + idx < self.to")]
fn test_dbkey_suffix_at_overflow() {
    let b = DBKey::read(b"\x00qwertyuiopasdfghjklzxcvbnm123456\xFF");
    b.suffix(1).get(255);
}

#[test]
fn test_dbkey_common_prefix() {
    let b1 = DBKey::read(b"\x01abcd0000000000000000000000000000\x00");
    let b2 = DBKey::read(b"\x01abde0000000000000000000000000000\x00");
    assert_eq!(b1.common_prefix(&b1), 256);
    let c = b1.common_prefix(&b2);
    assert_eq!(c, 21);
    let c = b2.common_prefix(&b1);
    assert_eq!(c, 21);
    let b1 = b1.suffix(9);
    let b2 = b2.suffix(9);
    let c = b1.common_prefix(&b2);
    assert_eq!(c, 12);
    let b3 = DBKey::read(b"\x01\xFF0000000000000000000000000000000\x00");
    let b4 = DBKey::read(b"\x01\xFE0000000000000000000000000000000\x00");
    assert_eq!(b3.common_prefix(&b4), 7);
    assert_eq!(b4.common_prefix(&b3), 7);
    assert_eq!(b3.common_prefix(&b3), 256);
    let b3 = b3.suffix(30);
    assert_eq!(b3.common_prefix(&b3), 226);
    let b3 = b3.truncate(200);
    assert_eq!(b3.common_prefix(&b3), 200);
    let b5 = DBKey::read(b"\x01\xFF0000000000000000000000000000000\x00");
    assert_eq!(b5.truncate(0).common_prefix(&b3), 0);
}

#[test]
fn test_dbkey_is_leaf() {
    let b = DBKey::read(b"\x01qwertyuiopasdfghjklzxcvbnm123456\x00");
    assert_eq!(b.len(), 256);
    assert_eq!(b.suffix(4).is_leaf(), true);
    assert_eq!(b.suffix(8).is_leaf(), true);
    assert_eq!(b.suffix(250).is_leaf(), true);
    assert_eq!(b.truncate(16).is_leaf(), false);
}

#[test]
fn test_dbkey_is_branch() {
    let b = DBKey::read(b"\x00qwertyuiopasdfghjklzxcvbnm123456\xFF");
    assert_eq!(b.len(), 255);
    assert_eq!(b.is_leaf(), false);
}
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};
use std::fs::File;
use std::os::unix::fs::{PermissionsExt, MetadataExt};

use libmount::{BindMount, Remount};

use quire::validate as V;
use config::read_config;
use config::containers::Container as Cont;
use version::short_version;
use container::util::{copy_dir};
use file_util::{Dir, ShallowCopy};
use build_step::{BuildStep, VersionError, StepError, Digest, Config, Guard};

use builder::error::StepError as E;
use builder::dns::revert_name_files;
use builder::commands::copy::{create_path_filter, hash_path};
use builder::commands::copy::{hash_file_content};

// Build Steps
#[derive(Debug, Deserialize)]
pub struct Container(String);

impl Container {
    pub fn config() -> V::Scalar {
        V::Scalar::new()
    }
}

#[derive(Deserialize, Debug)]
pub struct Build {
    pub container: String,
    pub source: PathBuf,
    pub path: Option<PathBuf>,
    pub temporary_mount: Option<PathBuf>,
    pub content_hash: bool,
}

impl Build {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("container", V::Scalar::new())
        .member("source".to_string(),
            V::Directory::new().absolute(true).default("/"))
        .member("path".to_string(),
            V::Directory::new().absolute(true).optional())
        .member("temporary_mount".to_string(),
            V::Directory::new().absolute(true).optional())
        .member("content_hash", V::Scalar::new().default(false))
    }
}

#[derive(Deserialize, Debug)]
pub struct GitSource {
    pub url: String,
    pub revision: Option<String>,
    pub branch: Option<String>,
}

#[derive(Deserialize, Debug)]
pub enum Source {
    Git(GitSource),
    Container(String),
    Directory,
}

#[derive(Deserialize, Debug)]
pub struct SubConfig {
    pub source: Source,
    pub path: PathBuf,
    pub container: String,
    pub cache: Option<bool>,
}

impl SubConfig {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("source", V::Enum::new()
            .option("Directory", V::Nothing)
            .option("Container", V::Scalar::new())
            .option("Git", V::Structure::new()
                .member("url", V::Scalar::new())
                .member("revision", V::Scalar::new().optional())
                .member("branch", V::Scalar::new().optional()))
            .optional()
            .default_tag("Directory"))
        .member("path".to_string(), V::Directory::new()
            .absolute(false)
            .default("vagga.yaml"))
        .member("container", V::Scalar::new())
        .member("cache", V::Scalar::new().optional())
    }
}

pub fn build(binfo: &Build, guard: &mut Guard, build: bool)
    -> Result<(), StepError>
{
    let ref name = binfo.container;
    let cont = guard.ctx.config.containers.get(name)
        .expect("Subcontainer not found");  // TODO
    if build {
        let version = short_version(&cont, &guard.ctx.config)
            .map_err(|(s, e)| format!("step {}: {}", s, e))?;
        let container = Path::new("/vagga/base/.roots")
            .join(format!("{}.{}", name, version));
        let path = container.join("root")
            .join(binfo.source.strip_prefix("/").unwrap());

        // Update container use when using it as subcontainer (fixes #267)
        File::create(Path::new(&container).join("last_use"))
            .map_err(|e| warn!("Can't write image usage info: {}", e)).ok();

        if let Some(ref dest_rel) = binfo.path {
            let dest = Path::new("/vagga/root")
                .join(dest_rel.strip_prefix("/").unwrap());
            if path.is_dir() {
                try_msg!(copy_dir(&path, &dest, None, None),
                    "Error copying dir {p:?}: {err}", p=path);
            } else {
                try_msg!(ShallowCopy::new(&path, &dest).copy(),
                    "Error copying file {p:?}: {err}", p=path);
            }
        } else if let Some(ref dest_rel) = binfo.temporary_mount {
            let dest = Path::new("/vagga/root")
                .join(dest_rel.strip_prefix("/").unwrap());
            try_msg!(Dir::new(&dest).create(),
                "Error creating destination dir: {err}");
            BindMount::new(&path, &dest).mount()?;
            Remount::new(&dest).bind(true).readonly(true).remount()?;
            guard.ctx.mounted.push(dest);
        }
    }
    Ok(())
}

fn real_copy(name: &String, cont: &Cont, guard: &mut Guard)
    -> Result<(), StepError>
{
    let version = short_version(&cont, &guard.ctx.config)
        .map_err(|(s, e)| format!("step {}: {}", s, e))?;
    let container = format!("/vagga/base/.roots/{}.{}", name, version);

    // Update container use when using it as subcontainer (fixes #267)
    File::create(Path::new(&container).join("last_use"))
        .map_err(|e| warn!("Can't write image usage info: {}", e)).ok();

    let root = Path::new(&container).join("root");
    try_msg!(copy_dir(&root, &Path::new("/vagga/root"),
                      None, None),
        "Error copying dir {p:?}: {err}", p=root);
    Ok(())
}

pub fn clone(name: &String, guard: &mut Guard, build: bool)
    -> Result<(), StepError>
{
    let cont = guard.ctx.config.containers.get(name)
        .expect("Subcontainer not found");  // TODO
    for b in cont.setup.iter() {
        b.build(guard, false)
            .map_err(|e| E::SubStep(b.0.clone(), Box::new(e)))?;
    }
    if build {
        real_copy(name, cont, guard)?;
    }
    Ok(())
}

fn find_config(cfg: &SubConfig, guard: &mut Guard)
    -> Result<Config, StepError>
{
    let path = match cfg.source {
        Source::Container(ref container) => {
            let cont = guard.ctx.config.containers.get(container)
                .expect("Subcontainer not found");  // TODO
            let version = short_version(&cont, &guard.ctx.config)
                .map_err(|(s, e)| format!("step {}: {}", s, e))?;
            let container = Path::new("/vagga/base/.roots")
                .join(format!("{}.{}", container, version));

            // Update container use when using it as subcontainer (fixes #267)
            File::create(Path::new(&container).join("last_use"))
                .map_err(|e| warn!("Can't write image usage info: {}", e))
                .ok();

            container.join("root").join(&cfg.path)
        }
        Source::Git(ref _git) => {
            unimplemented!();
        }
        Source::Directory => {
            Path::new("/work").join(&cfg.path)
        }
    };
    Ok(read_config(&path, true)?)
}

pub fn subconfig(cfg: &SubConfig, guard: &mut Guard, build: bool)
    -> Result<(), StepError>
{
    let subcfg = find_config(cfg, guard)?;
    let cont = subcfg.containers.get(&cfg.container)
        .expect("Subcontainer not found");  // TODO
    for b in cont.setup.iter() {
        b.build(guard, build)
            .map_err(|e| E::SubStep(b.0.clone(), Box::new(e)))?;
    }
    Ok(())
}

impl BuildStep for Container {
    fn name(&self) -> &'static str { "Container" }
    fn hash(&self, cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        let cont = cfg.containers.get(&self.0)
            .ok_or(VersionError::ContainerNotFound(self.0.to_string()))?;
        for b in cont.setup.iter() {
            debug!("Versioning setup: {:?}", b);
            hash.command(b.name());
            b.hash(cfg, hash)?;
        }
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        clone(&self.0, guard, build)?;
        revert_name_files()?;
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        Some(&self.0)
    }
}
impl BuildStep for Build {
    fn name(&self) -> &'static str { "Build" }
    fn hash(&self, cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        let cinfo = cfg.containers.get(&self.container)
            .ok_or(VersionError::ContainerNotFound(self.container.clone()))?;
        if self.content_hash {
            let version = short_version(&cinfo, cfg)?;
            let root = Path::new("/vagga/base/.roots")
                .join(format!("{}.{}", self.container, version))
                .join("root");
            println!("CONTAINER PATH {:?}", root);
            if !root.exists() {
                return Err(VersionError::New);
            }
            if let Some(ref dest_rel) = self.path {
                let filter = create_path_filter(&Vec::new(), Some(true),
                    &None, &None, false)?;
                let spath = self.source.strip_prefix("/")
                    .expect("absolute_source_path");
                hash_path(hash, &root.join(&spath), &filter, |h, p, st| {
                    h.field("filename", p);
                    h.field("mode", st.permissions().mode() & 0o7777);
                    h.field("uid", st.uid());
                    h.field("gid", st.gid());
                    hash_file_content(h, p, st)
                        .map_err(|e| VersionError::Io(e, PathBuf::from(p)))?;
                    Ok(())
                })?;
                hash.field("path", dest_rel);
            } else if let Some(_) = self.temporary_mount {
                unimplemented!("Build: combination of \
                    content-hash and temporary-mount are not supported yet");
            }
        } else {
            for b in cinfo.setup.iter() {
                debug!("Versioning setup: {:?}", b);
                hash.command(b.name());
                b.hash(cfg, hash)?;
            }
            // TODO(tailhook) should we hash our params?!?!
        }
        Ok(())
    }
    fn build(&self, guard: &mut Guard, do_build: bool)
        -> Result<(), StepError>
    {
        build(&self, guard, do_build)
    }
    fn is_dependent_on(&self) -> Option<&str> {
        Some(&self.container)
    }
}
impl BuildStep for SubConfig {
    fn name(&self) -> &'static str { "SubConfig" }
    fn hash(&self, cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        let path = match self.source {
            Source::Container(ref container) => {
                let cinfo = cfg.containers.get(container)
                    .ok_or(VersionError::ContainerNotFound(container.clone()))?;
                let version = short_version(&cinfo, cfg)?;
                Path::new("/vagga/base/.roots")
                    .join(format!("{}.{}", container, version))
                    .join("root").join(&self.path)
            }
            Source::Git(ref _git) => {
                unimplemented!();
            }
            Source::Directory => {
                Path::new("/work").join(&self.path)
            }
        };
        if !path.exists() {
            return Err(VersionError::New);
        }
        let subcfg = read_config(&path, true)?;
        let cont = subcfg.containers.get(&self.container)
            .ok_or(VersionError::ContainerNotFound(self.container.to_string()))?;
        for b in cont.setup.iter() {
            debug!("Versioning setup: {:?}", b);
            hash.command(b.name());
            b.hash(cfg, hash)?;
        }
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        subconfig(self, guard, build)?;
        revert_name_files()?;
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        match self.source {
            Source::Directory => None,
            Source::Container(ref name) => Some(name),
            Source::Git(ref _git) => None,
        }
    }
}
<commit_msg>Remove debugging print<commit_after>use std::path::{Path, PathBuf};
use std::fs::File;
use std::os::unix::fs::{PermissionsExt, MetadataExt};

use libmount::{BindMount, Remount};

use quire::validate as V;
use config::read_config;
use config::containers::Container as Cont;
use version::short_version;
use container::util::{copy_dir};
use file_util::{Dir, ShallowCopy};
use build_step::{BuildStep, VersionError, StepError, Digest, Config, Guard};

use builder::error::StepError as E;
use builder::dns::revert_name_files;
use builder::commands::copy::{create_path_filter, hash_path};
use builder::commands::copy::{hash_file_content};

// Build Steps
#[derive(Debug, Deserialize)]
pub struct Container(String);

impl Container {
    pub fn config() -> V::Scalar {
        V::Scalar::new()
    }
}

#[derive(Deserialize, Debug)]
pub struct Build {
    pub container: String,
    pub source: PathBuf,
    pub path: Option<PathBuf>,
    pub temporary_mount: Option<PathBuf>,
    pub content_hash: bool,
}

impl Build {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("container", V::Scalar::new())
        .member("source".to_string(),
            V::Directory::new().absolute(true).default("/"))
        .member("path".to_string(),
            V::Directory::new().absolute(true).optional())
        .member("temporary_mount".to_string(),
            V::Directory::new().absolute(true).optional())
        .member("content_hash", V::Scalar::new().default(false))
    }
}

#[derive(Deserialize, Debug)]
pub struct GitSource {
    pub url: String,
    pub revision: Option<String>,
    pub branch: Option<String>,
}

#[derive(Deserialize, Debug)]
pub enum Source {
    Git(GitSource),
    Container(String),
    Directory,
}

#[derive(Deserialize, Debug)]
pub struct SubConfig {
    pub source: Source,
    pub path: PathBuf,
    pub container: String,
    pub cache: Option<bool>,
}

impl SubConfig {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("source", V::Enum::new()
            .option("Directory", V::Nothing)
            .option("Container", V::Scalar::new())
            .option("Git", V::Structure::new()
                .member("url", V::Scalar::new())
                .member("revision", V::Scalar::new().optional())
                .member("branch", V::Scalar::new().optional()))
            .optional()
            .default_tag("Directory"))
        .member("path".to_string(), V::Directory::new()
            .absolute(false)
            .default("vagga.yaml"))
        .member("container", V::Scalar::new())
        .member("cache", V::Scalar::new().optional())
    }
}

pub fn build(binfo: &Build, guard: &mut Guard, build: bool)
    -> Result<(), StepError>
{
    let ref name = binfo.container;
    let cont = guard.ctx.config.containers.get(name)
        .expect("Subcontainer not found");  // TODO
    if build {
        let version = short_version(&cont, &guard.ctx.config)
            .map_err(|(s, e)| format!("step {}: {}", s, e))?;
        let container = Path::new("/vagga/base/.roots")
            .join(format!("{}.{}", name, version));
        let path = container.join("root")
            .join(binfo.source.strip_prefix("/").unwrap());

        // Update container use when using it as subcontainer (fixes #267)
        File::create(Path::new(&container).join("last_use"))
            .map_err(|e| warn!("Can't write image usage info: {}", e)).ok();

        if let Some(ref dest_rel) = binfo.path {
            let dest = Path::new("/vagga/root")
                .join(dest_rel.strip_prefix("/").unwrap());
            if path.is_dir() {
                try_msg!(copy_dir(&path, &dest, None, None),
                    "Error copying dir {p:?}: {err}", p=path);
            } else {
                try_msg!(ShallowCopy::new(&path, &dest).copy(),
                    "Error copying file {p:?}: {err}", p=path);
            }
        } else if let Some(ref dest_rel) = binfo.temporary_mount {
            let dest = Path::new("/vagga/root")
                .join(dest_rel.strip_prefix("/").unwrap());
            try_msg!(Dir::new(&dest).create(),
                "Error creating destination dir: {err}");
            BindMount::new(&path, &dest).mount()?;
            Remount::new(&dest).bind(true).readonly(true).remount()?;
            guard.ctx.mounted.push(dest);
        }
    }
    Ok(())
}

fn real_copy(name: &String, cont: &Cont, guard: &mut Guard)
    -> Result<(), StepError>
{
    let version = short_version(&cont, &guard.ctx.config)
        .map_err(|(s, e)| format!("step {}: {}", s, e))?;
    let container = format!("/vagga/base/.roots/{}.{}", name, version);

    // Update container use when using it as subcontainer (fixes #267)
    File::create(Path::new(&container).join("last_use"))
        .map_err(|e| warn!("Can't write image usage info: {}", e)).ok();

    let root = Path::new(&container).join("root");
    try_msg!(copy_dir(&root, &Path::new("/vagga/root"),
                      None, None),
        "Error copying dir {p:?}: {err}", p=root);
    Ok(())
}

pub fn clone(name: &String, guard: &mut Guard, build: bool)
    -> Result<(), StepError>
{
    let cont = guard.ctx.config.containers.get(name)
        .expect("Subcontainer not found");  // TODO
    for b in cont.setup.iter() {
        b.build(guard, false)
            .map_err(|e| E::SubStep(b.0.clone(), Box::new(e)))?;
    }
    if build {
        real_copy(name, cont, guard)?;
    }
    Ok(())
}

fn find_config(cfg: &SubConfig, guard: &mut Guard)
    -> Result<Config, StepError>
{
    let path = match cfg.source {
        Source::Container(ref container) => {
            let cont = guard.ctx.config.containers.get(container)
                .expect("Subcontainer not found");  // TODO
            let version = short_version(&cont, &guard.ctx.config)
                .map_err(|(s, e)| format!("step {}: {}", s, e))?;
            let container = Path::new("/vagga/base/.roots")
                .join(format!("{}.{}", container, version));

            // Update container use when using it as subcontainer (fixes #267)
            File::create(Path::new(&container).join("last_use"))
                .map_err(|e| warn!("Can't write image usage info: {}", e))
                .ok();

            container.join("root").join(&cfg.path)
        }
        Source::Git(ref _git) => {
            unimplemented!();
        }
        Source::Directory => {
            Path::new("/work").join(&cfg.path)
        }
    };
    Ok(read_config(&path, true)?)
}

pub fn subconfig(cfg: &SubConfig, guard: &mut Guard, build: bool)
    -> Result<(), StepError>
{
    let subcfg = find_config(cfg, guard)?;
    let cont = subcfg.containers.get(&cfg.container)
        .expect("Subcontainer not found");  // TODO
    for b in cont.setup.iter() {
        b.build(guard, build)
            .map_err(|e| E::SubStep(b.0.clone(), Box::new(e)))?;
    }
    Ok(())
}

impl BuildStep for Container {
    fn name(&self) -> &'static str { "Container" }
    fn hash(&self, cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        let cont = cfg.containers.get(&self.0)
            .ok_or(VersionError::ContainerNotFound(self.0.to_string()))?;
        for b in cont.setup.iter() {
            debug!("Versioning setup: {:?}", b);
            hash.command(b.name());
            b.hash(cfg, hash)?;
        }
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        clone(&self.0, guard, build)?;
        revert_name_files()?;
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        Some(&self.0)
    }
}
impl BuildStep for Build {
    fn name(&self) -> &'static str { "Build" }
    fn hash(&self, cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        let cinfo = cfg.containers.get(&self.container)
            .ok_or(VersionError::ContainerNotFound(self.container.clone()))?;
        if self.content_hash {
            let version = short_version(&cinfo, cfg)?;
            let root = Path::new("/vagga/base/.roots")
                .join(format!("{}.{}", self.container, version))
                .join("root");
            if !root.exists() {
                return Err(VersionError::New);
            }
            if let Some(ref dest_rel) = self.path {
                let filter = create_path_filter(&Vec::new(), Some(true),
                    &None, &None, false)?;
                let spath = self.source.strip_prefix("/")
                    .expect("absolute_source_path");
                hash_path(hash, &root.join(&spath), &filter, |h, p, st| {
                    h.field("filename", p);
                    h.field("mode", st.permissions().mode() & 0o7777);
                    h.field("uid", st.uid());
                    h.field("gid", st.gid());
                    hash_file_content(h, p, st)
                        .map_err(|e| VersionError::Io(e, PathBuf::from(p)))?;
                    Ok(())
                })?;
                hash.field("path", dest_rel);
            } else if let Some(_) = self.temporary_mount {
                unimplemented!("Build: combination of \
                    content-hash and temporary-mount are not supported yet");
            }
        } else {
            for b in cinfo.setup.iter() {
                debug!("Versioning setup: {:?}", b);
                hash.command(b.name());
                b.hash(cfg, hash)?;
            }
            // TODO(tailhook) should we hash our params?!?!
        }
        Ok(())
    }
    fn build(&self, guard: &mut Guard, do_build: bool)
        -> Result<(), StepError>
    {
        build(&self, guard, do_build)
    }
    fn is_dependent_on(&self) -> Option<&str> {
        Some(&self.container)
    }
}
impl BuildStep for SubConfig {
    fn name(&self) -> &'static str { "SubConfig" }
    fn hash(&self, cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        let path = match self.source {
            Source::Container(ref container) => {
                let cinfo = cfg.containers.get(container)
                    .ok_or(VersionError::ContainerNotFound(container.clone()))?;
                let version = short_version(&cinfo, cfg)?;
                Path::new("/vagga/base/.roots")
                    .join(format!("{}.{}", container, version))
                    .join("root").join(&self.path)
            }
            Source::Git(ref _git) => {
                unimplemented!();
            }
            Source::Directory => {
                Path::new("/work").join(&self.path)
            }
        };
        if !path.exists() {
            return Err(VersionError::New);
        }
        let subcfg = read_config(&path, true)?;
        let cont = subcfg.containers.get(&self.container)
            .ok_or(VersionError::ContainerNotFound(self.container.to_string()))?;
        for b in cont.setup.iter() {
            debug!("Versioning setup: {:?}", b);
            hash.command(b.name());
            b.hash(cfg, hash)?;
        }
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        subconfig(self, guard, build)?;
        revert_name_files()?;
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        match self.source {
            Source::Directory => None,
            Source::Container(ref name) => Some(name),
            Source::Git(ref _git) => None,
        }
    }
}
<|endoftext|>


<commit_before>const USAGE: &'static str = "
Usage: sieve bench
       sieve --help

Sieve of Eratosthenes

A sieve finds prime numbers by first assuming all candidates are prime,
then progressively using small primes to mark their multiples composite.

Note that a given prime p only needs to start marking its multiples >=p²,
as anything less will already have been marked by an even smaller prime
factor.  This also means that sieving up to N only requires primes up to
sqrt(N) - a useful fact to prepare for parallel sieving.

Here we have three forms of the sieve:

- `sieve_serial`: direct iteration in serial
- `sieve_chunks`: chunked iteration, still serial
- `sieve_parallel`: chunked iteration in parallel

The performance improvement of `sieve_chunks` is largely due to cache
locality.  Since the sieve has to sweep repeatedly over the same memory,
it's helpful that the chunk size can fit entirely in the CPU cache.  Then
`sieve_parallel` also benefits from this as long as there's cache room for
multiple chunks, for the separate jobs in each thread.

Commands:
    bench              Run the benchmark in different modes and print the timings.

Options:
    -h, --help         Show this message.
";

#[derive(RustcDecodable)]
pub struct Args {
    cmd_bench: bool,
}

#[cfg(test)]
mod bench;

use docopt::Docopt;
use itertools::StrideMut;
use rayon::prelude::*;
use time;

const CHUNK_SIZE: usize = 100_000;

// Number of Primes < 10^n
// https://oeis.org/A006880
const NUM_PRIMES: &'static [usize] = &[
    0, // primes in 0..10^0
    4, // primes in 0..10^1
    25, // etc
    168,
    1229,
    9592,
    78498,
    664579,
    5761455,
    50847534,
    455052511,
    4118054813,
    37607912018,
    346065536839,
    3204941750802,
    29844570422669,
    279238341033925,
    2623557157654233,
    24739954287740860,
    234057667276344607,
    2220819602560918840,
];

// For all of these sieves, sieve[i]==true -> 2*i+1 is prime

fn max(magnitude: usize) -> usize {
    10_usize.pow(magnitude as u32)
}

/// Sieve odd integers for primes < max.
fn sieve_serial(max: usize) -> Vec<bool> {
    let mut sieve = vec![true; max / 2];
    sieve[0] = false; // 1 is not prime
    for i in 1 .. {
        if sieve[i] {
            let p = 2 * i + 1;
            let pp = p * p;
            if pp >= max { break }
            clear_stride(&mut sieve, pp / 2, p);
        }
    }
    sieve
}

/// Sieve odd integers for primes < max using chunks.
fn sieve_chunks(max: usize) -> Vec<bool> {
    // first compute the small primes, up to sqrt(max).
    let small_max = (max as f64).sqrt().ceil() as usize;
    let mut sieve = sieve_serial(small_max);
    sieve.resize(max / 2, true);

    {
        let (low, high) = sieve.split_at_mut(small_max / 2);
        for (chunk_index, chunk) in high.chunks_mut(CHUNK_SIZE).enumerate() {
            let i = small_max / 2 + chunk_index * CHUNK_SIZE;
            let base = i * 2 + 1;
            update_chunk(low, chunk, base);
        }
    }

    sieve
}

/// Sieve odd integers for primes < max, in parallel!
fn sieve_parallel(max: usize) -> Vec<bool> {
    // first compute the small primes, up to sqrt(max).
    let small_max = (max as f64).sqrt().ceil() as usize;
    let mut sieve = sieve_serial(small_max);
    sieve.resize(max / 2, true);

    {
        let (low, high) = sieve.split_at_mut(small_max / 2);
        high.par_chunks_mut(CHUNK_SIZE)
            .enumerate() // to figure out where this chunk came from
            // FIXME needed? .weight_max() // ensure every single chunk is a separate rayon job
            .for_each(|(chunk_index, chunk)| {
                let i = small_max / 2 + chunk_index * CHUNK_SIZE;
                let base = i * 2 + 1;
                update_chunk(low, chunk, base);
            });
    }

    sieve
}

/// Update a chunk with low primes
fn update_chunk(low: &[bool], chunk: &mut [bool], base: usize) {
    let max = base + chunk.len() * 2;
    for (i, &is_prime) in low.iter().enumerate() {
        if is_prime {
            let p = 2 * i + 1;
            let pp = p * p;
            if pp >= max { break }

            let pm = if pp < base {
                // p² is too small - find the first odd multiple that's in range
                ((base + p - 1) / p | 1) * p
            } else { pp };

            if pm < max {
                clear_stride(chunk, (pm - base) / 2, p);
            }
        }
    }
}

fn clear_stride(slice: &mut [bool], from: usize, stride: usize) {
    let slice = &mut slice[from..];
    for x in StrideMut::from_slice(slice, stride as isize) {
        *x = false;
    }
}

fn measure(f: fn(usize) -> Vec<bool>) -> u64 {
    const MAGNITUDE: usize = 9;

    let start = time::precise_time_ns();
    let sieve = f(max(MAGNITUDE));
    let duration = time::precise_time_ns() - start;

    // sanity check the number of primes found
    let num_primes = 1 + sieve.into_iter().filter(|&b| b).count();
    assert_eq!(num_primes, NUM_PRIMES[MAGNITUDE]);

    duration
}

pub fn main(args: &[String]) {
    let args: Args =
        Docopt::new(USAGE)
            .and_then(|d| d.argv(args).decode())
            .unwrap_or_else(|e| e.exit());

    if args.cmd_bench {
        let serial = measure(sieve_serial);
        println!("  serial: {:10} ns", serial);

        let chunks = measure(sieve_chunks);
        println!("  chunks: {:10} ns -> {:.2}x speedup", chunks,
                 serial as f64 / chunks as f64);

        let parallel = measure(sieve_parallel);
        println!("parallel: {:10} ns -> {:.2}x speedup", parallel,
                 chunks as f64 / parallel as f64);
    }
}
<commit_msg>Use set_max_len(1) in the sieve demo<commit_after>const USAGE: &'static str = "
Usage: sieve bench
       sieve --help

Sieve of Eratosthenes

A sieve finds prime numbers by first assuming all candidates are prime,
then progressively using small primes to mark their multiples composite.

Note that a given prime p only needs to start marking its multiples >=p²,
as anything less will already have been marked by an even smaller prime
factor.  This also means that sieving up to N only requires primes up to
sqrt(N) - a useful fact to prepare for parallel sieving.

Here we have three forms of the sieve:

- `sieve_serial`: direct iteration in serial
- `sieve_chunks`: chunked iteration, still serial
- `sieve_parallel`: chunked iteration in parallel

The performance improvement of `sieve_chunks` is largely due to cache
locality.  Since the sieve has to sweep repeatedly over the same memory,
it's helpful that the chunk size can fit entirely in the CPU cache.  Then
`sieve_parallel` also benefits from this as long as there's cache room for
multiple chunks, for the separate jobs in each thread.

Commands:
    bench              Run the benchmark in different modes and print the timings.

Options:
    -h, --help         Show this message.
";

#[derive(RustcDecodable)]
pub struct Args {
    cmd_bench: bool,
}

#[cfg(test)]
mod bench;

use docopt::Docopt;
use itertools::StrideMut;
use rayon::prelude::*;
use time;

const CHUNK_SIZE: usize = 100_000;

// Number of Primes < 10^n
// https://oeis.org/A006880
const NUM_PRIMES: &'static [usize] = &[
    0, // primes in 0..10^0
    4, // primes in 0..10^1
    25, // etc
    168,
    1229,
    9592,
    78498,
    664579,
    5761455,
    50847534,
    455052511,
    4118054813,
    37607912018,
    346065536839,
    3204941750802,
    29844570422669,
    279238341033925,
    2623557157654233,
    24739954287740860,
    234057667276344607,
    2220819602560918840,
];

// For all of these sieves, sieve[i]==true -> 2*i+1 is prime

fn max(magnitude: usize) -> usize {
    10_usize.pow(magnitude as u32)
}

/// Sieve odd integers for primes < max.
fn sieve_serial(max: usize) -> Vec<bool> {
    let mut sieve = vec![true; max / 2];
    sieve[0] = false; // 1 is not prime
    for i in 1 .. {
        if sieve[i] {
            let p = 2 * i + 1;
            let pp = p * p;
            if pp >= max { break }
            clear_stride(&mut sieve, pp / 2, p);
        }
    }
    sieve
}

/// Sieve odd integers for primes < max using chunks.
fn sieve_chunks(max: usize) -> Vec<bool> {
    // first compute the small primes, up to sqrt(max).
    let small_max = (max as f64).sqrt().ceil() as usize;
    let mut sieve = sieve_serial(small_max);
    sieve.resize(max / 2, true);

    {
        let (low, high) = sieve.split_at_mut(small_max / 2);
        for (chunk_index, chunk) in high.chunks_mut(CHUNK_SIZE).enumerate() {
            let i = small_max / 2 + chunk_index * CHUNK_SIZE;
            let base = i * 2 + 1;
            update_chunk(low, chunk, base);
        }
    }

    sieve
}

/// Sieve odd integers for primes < max, in parallel!
fn sieve_parallel(max: usize) -> Vec<bool> {
    // first compute the small primes, up to sqrt(max).
    let small_max = (max as f64).sqrt().ceil() as usize;
    let mut sieve = sieve_serial(small_max);
    sieve.resize(max / 2, true);

    {
        let (low, high) = sieve.split_at_mut(small_max / 2);
        high.par_chunks_mut(CHUNK_SIZE)
            .enumerate() // to figure out where this chunk came from
            .set_max_len(1) // ensure every single chunk is a separate rayon job
            .for_each(|(chunk_index, chunk)| {
                let i = small_max / 2 + chunk_index * CHUNK_SIZE;
                let base = i * 2 + 1;
                update_chunk(low, chunk, base);
            });
    }

    sieve
}

/// Update a chunk with low primes
fn update_chunk(low: &[bool], chunk: &mut [bool], base: usize) {
    let max = base + chunk.len() * 2;
    for (i, &is_prime) in low.iter().enumerate() {
        if is_prime {
            let p = 2 * i + 1;
            let pp = p * p;
            if pp >= max { break }

            let pm = if pp < base {
                // p² is too small - find the first odd multiple that's in range
                ((base + p - 1) / p | 1) * p
            } else { pp };

            if pm < max {
                clear_stride(chunk, (pm - base) / 2, p);
            }
        }
    }
}

fn clear_stride(slice: &mut [bool], from: usize, stride: usize) {
    let slice = &mut slice[from..];
    for x in StrideMut::from_slice(slice, stride as isize) {
        *x = false;
    }
}

fn measure(f: fn(usize) -> Vec<bool>) -> u64 {
    const MAGNITUDE: usize = 9;

    let start = time::precise_time_ns();
    let sieve = f(max(MAGNITUDE));
    let duration = time::precise_time_ns() - start;

    // sanity check the number of primes found
    let num_primes = 1 + sieve.into_iter().filter(|&b| b).count();
    assert_eq!(num_primes, NUM_PRIMES[MAGNITUDE]);

    duration
}

pub fn main(args: &[String]) {
    let args: Args =
        Docopt::new(USAGE)
            .and_then(|d| d.argv(args).decode())
            .unwrap_or_else(|e| e.exit());

    if args.cmd_bench {
        let serial = measure(sieve_serial);
        println!("  serial: {:10} ns", serial);

        let chunks = measure(sieve_chunks);
        println!("  chunks: {:10} ns -> {:.2}x speedup", chunks,
                 serial as f64 / chunks as f64);

        let parallel = measure(sieve_parallel);
        println!("parallel: {:10} ns -> {:.2}x speedup", parallel,
                 chunks as f64 / parallel as f64);
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use model::Model;

pub struct Classifier {
    models: HashMap<String, Model>
}

impl Classifier {
    pub fn new() -> Classifier {
        Classifier { models: HashMap::new() }
    }

    pub fn train(&mut self, text: &str, name: &str) {
        let model = Model::build_from_text(text);
        self.models.insert(name.to_string(), model);
    }

    pub fn classify(&self, text: &str) -> &String {
        let model = Model::build_from_text(text);
        let (name, matched_model) = self.models
            .iter()
            .min_by_key(|&(name, other_model)| model.compare(other_model))
            .unwrap();
        name
    }
}
<commit_msg>Silence an unused variable warning<commit_after>use std::collections::HashMap;

use model::Model;

pub struct Classifier {
    models: HashMap<String, Model>
}

impl Classifier {
    pub fn new() -> Classifier {
        Classifier { models: HashMap::new() }
    }

    pub fn train(&mut self, text: &str, name: &str) {
        let model = Model::build_from_text(text);
        self.models.insert(name.to_string(), model);
    }

    pub fn classify(&self, text: &str) -> &String {
        let model = Model::build_from_text(text);
        let (name, _) = self.models
            .iter()
            .min_by_key(|&(_, other_model)| model.compare(other_model))
            .unwrap();
        name
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

#![feature(std_misc)]

extern crate rustfmt;

use std::collections::HashMap;
use std::fs;
use std::io::Read;
use rustfmt::*;

// For now, the only supported regression tests are idempotent tests - the input and
// output must match exactly.
// FIXME(#28) would be good to check for error messages and fail on them, or at least report.
#[test]
fn idempotent_tests() {
    println!("Idempotent tests:");

    // Get all files in the tests/idem directory
    let files = fs::read_dir("tests/idem").unwrap();
    let files2 = fs::read_dir("tests").unwrap();
    let files3 = fs::read_dir("src/bin").unwrap();
    // For each file, run rustfmt and collect the output

    let mut count = 0;
    let mut fails = 0;
    for entry in files.chain(files2).chain(files3) {
        let path = entry.unwrap().path();
        let file_name = path.to_str().unwrap();
        if !file_name.ends_with(".rs") {
            continue;
        }
        println!("Testing '{}'...", file_name);
        match idempotent_check(file_name.to_owned()) {
            Ok(()) => {},
            Err(m) => {
                print_mismatches(m);
                fails += 1;
            },
        }
        count += 1;
    }
    // And also dogfood rustfmt!
    println!("Testing 'src/lib.rs'...");
    match idempotent_check("src/lib.rs".to_owned()) {
        Ok(()) => {},
        Err(m) => {
            print_mismatches(m);
            fails += 1;
        },
    }
    count += 1;

    // Display results
    println!("Ran {} idempotent tests; {} failures.", count, fails);
    assert!(fails == 0, "{} idempotent tests failed", fails);
}

// Compare output to input.
fn print_mismatches(result: HashMap<String, String>) {
    for (file_name, fmt_text) in result {
        println!("Mismatch in {}.", file_name);
        println!("{}", fmt_text);
    }
}

// Ick, just needed to get a &'static to handle_result.
static HANDLE_RESULT: &'static Fn(HashMap<String, String>) = &handle_result;

pub fn idempotent_check(filename: String) -> Result<(), HashMap<String, String>> {
    use std::thread;
    use std::fs;
    use std::io::Read;
	let args = vec!["rustfmt".to_owned(), filename];
    thread::spawn(move || {
        run(args, WriteMode::Return(HANDLE_RESULT));
    }).join().map_err(|mut any|
        any.downcast_mut::<HashMap<String, String>>()
           .unwrap() // i know it is a hashmap
           .drain() // i only get a reference :(
           .collect() // so i need to turn it into an iter and then back
    )
}

// Compare output to input.
fn handle_result(result: HashMap<String, String>) {
    let mut failures = HashMap::new();

    for (file_name, fmt_text) in result {
        let mut f = fs::File::open(&file_name).unwrap();
        let mut text = String::new();
        // TODO: speedup by running through bytes iterator
        f.read_to_string(&mut text).unwrap();
        if fmt_text != text {
            failures.insert(file_name, fmt_text);
        }
    }
    if !failures.is_empty() {
        panic!(failures);
    }
}<commit_msg>forgot newline at eof<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// 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.

#![feature(std_misc)]

extern crate rustfmt;

use std::collections::HashMap;
use std::fs;
use std::io::Read;
use rustfmt::*;

// For now, the only supported regression tests are idempotent tests - the input and
// output must match exactly.
// FIXME(#28) would be good to check for error messages and fail on them, or at least report.
#[test]
fn idempotent_tests() {
    println!("Idempotent tests:");

    // Get all files in the tests/idem directory
    let files = fs::read_dir("tests/idem").unwrap();
    let files2 = fs::read_dir("tests").unwrap();
    let files3 = fs::read_dir("src/bin").unwrap();
    // For each file, run rustfmt and collect the output

    let mut count = 0;
    let mut fails = 0;
    for entry in files.chain(files2).chain(files3) {
        let path = entry.unwrap().path();
        let file_name = path.to_str().unwrap();
        if !file_name.ends_with(".rs") {
            continue;
        }
        println!("Testing '{}'...", file_name);
        match idempotent_check(file_name.to_owned()) {
            Ok(()) => {},
            Err(m) => {
                print_mismatches(m);
                fails += 1;
            },
        }
        count += 1;
    }
    // And also dogfood rustfmt!
    println!("Testing 'src/lib.rs'...");
    match idempotent_check("src/lib.rs".to_owned()) {
        Ok(()) => {},
        Err(m) => {
            print_mismatches(m);
            fails += 1;
        },
    }
    count += 1;

    // Display results
    println!("Ran {} idempotent tests; {} failures.", count, fails);
    assert!(fails == 0, "{} idempotent tests failed", fails);
}

// Compare output to input.
fn print_mismatches(result: HashMap<String, String>) {
    for (file_name, fmt_text) in result {
        println!("Mismatch in {}.", file_name);
        println!("{}", fmt_text);
    }
}

// Ick, just needed to get a &'static to handle_result.
static HANDLE_RESULT: &'static Fn(HashMap<String, String>) = &handle_result;

pub fn idempotent_check(filename: String) -> Result<(), HashMap<String, String>> {
    use std::thread;
    use std::fs;
    use std::io::Read;
	let args = vec!["rustfmt".to_owned(), filename];
    thread::spawn(move || {
        run(args, WriteMode::Return(HANDLE_RESULT));
    }).join().map_err(|mut any|
        any.downcast_mut::<HashMap<String, String>>()
           .unwrap() // i know it is a hashmap
           .drain() // i only get a reference :(
           .collect() // so i need to turn it into an iter and then back
    )
}

// Compare output to input.
fn handle_result(result: HashMap<String, String>) {
    let mut failures = HashMap::new();

    for (file_name, fmt_text) in result {
        let mut f = fs::File::open(&file_name).unwrap();
        let mut text = String::new();
        // TODO: speedup by running through bytes iterator
        f.read_to_string(&mut text).unwrap();
        if fmt_text != text {
            failures.insert(file_name, fmt_text);
        }
    }
    if !failures.is_empty() {
        panic!(failures);
    }
}
<|endoftext|>