<commit_before>#![feature(plugin, rustc_private)]
#![plugin(string_cache_plugin)]
extern crate iron;
extern crate router;
extern crate urlencoded;
extern crate serialize;

use std::net::SocketAddrV4;
use std::net::Ipv4Addr;
use iron::prelude::*;
use iron::status;
use router::{Router};
use urlencoded::UrlEncodedQuery;

extern crate html5ever;
extern crate regex;

#[macro_use]
extern crate string_cache;

use pink_spider::scraper::extract_tracks;
use pink_spider::model::{Track, Entry, create_tables, drop_tables};
use serialize::json::{ToJson, Json};

extern crate pink_spider;

fn playlistify(req: &mut Request) -> IronResult<Response> {
    match req.get_ref::<UrlEncodedQuery>() {
        Ok(ref params) => {
            match params.get("url") {
                Some(url) => {
                    let entry = find_or_create_entry(&url[0]);
                    let json_obj: Json = entry.to_json();
                    let json_str: String = json_obj.to_string();
                    return Ok(Response::with((status::Ok, json_str)))
                },
                None => println!("no url")
            }
        }
        Err(ref e) =>
            println!("{:?}", e)
    };

    Ok(Response::with((status::Ok, "{}")))
}

fn find_or_create_entry(url: &str) -> Entry {
    match Entry::find_by_url(url) {
        Some(entry) => {
            println!("Get entry from database cache");
            entry
        },
        None => {
            let tracks = extract_tracks(url);
            match Entry::create_by_url(url.to_string()) {
                Some(mut entry) => {
                    println!("Create new entry to database cache");
                    for t in tracks {
                        match Track::create(t.provider, t.title, t.url, t.identifier) {
                            Some(track) => entry.add_track(track),
                            None        => ()
                        }
                    }
                    entry
                },
                None => {
                    println!("Failed to create entry database cache");
                    Entry {
                            id: 0,
                           url: url.to_string(),
                        tracks: tracks
                    }
                }
            }
        }
    }
}

fn main() {
    drop_tables();
    create_tables();

    let mut router = Router::new();
    router.get("/playlistify", playlistify);
    router.get("/:query", handler);

    fn handler(req: &mut Request) -> IronResult<Response> {
        let ref query = req.extensions.get::<Router>()
            .unwrap().find("query").unwrap_or("/");
        Ok(Response::with((status::Ok, *query)))
    }

    let opt_port = std::env::var("PORT");
    let port_str = match opt_port {
        Ok(n)    => n,
        Err(_) => "8080".to_string()
    };
    let port: u16 = match port_str.trim().parse() {
        Ok(n) => n,
        Err(_) => {
            println!("Faild to parse port");
            return;
        }
    };
    println!("PORT {}", port_str);
    let ip = Ipv4Addr::new(0, 0, 0, 0);
    Iron::new(router).http(SocketAddrV4::new(ip, port)).unwrap();
}
<commit_msg>Implement GET /tracks/:id api: Show track<commit_after>#![feature(plugin, rustc_private)]
#![plugin(string_cache_plugin)]
extern crate iron;
extern crate router;
extern crate urlencoded;
extern crate serialize;
extern crate html5ever;
extern crate regex;
extern crate core;

use std::net::SocketAddrV4;
use std::net::Ipv4Addr;
use iron::prelude::*;
use iron::status;
use iron::headers::{ContentType};
use iron::modifiers::Header;
use iron::mime::Mime;
use router::{Router};
use urlencoded::UrlEncodedQuery;
use core::str::FromStr;

#[macro_use]
extern crate string_cache;

use pink_spider::scraper::extract_tracks;
use pink_spider::model::{Track, Entry, create_tables, drop_tables};
use serialize::json::{ToJson, Json};

extern crate pink_spider;

fn playlistify(req: &mut Request) -> IronResult<Response> {
    match req.get_ref::<UrlEncodedQuery>() {
        Ok(ref params) => {
            match params.get("url") {
                Some(url) => {
                    let entry = find_or_create_entry(&url[0]);
                    let json_obj: Json = entry.to_json();
                    let json_str: String = json_obj.to_string();
                    return Ok(Response::with((status::Ok, json_str)))
                },
                None => println!("no url")
            }
        }
        Err(ref e) =>
            println!("{:?}", e)
    };

    Ok(Response::with((status::Ok, "{}")))
}

fn find_or_create_entry(url: &str) -> Entry {
    match Entry::find_by_url(url) {
        Some(entry) => {
            println!("Get entry from database cache");
            entry
        },
        None => {
            let tracks = extract_tracks(url);
            match Entry::create_by_url(url.to_string()) {
                Some(mut entry) => {
                    println!("Create new entry to database cache");
                    for t in tracks {
                        match Track::create(t.provider, t.title, t.url, t.identifier) {
                            Some(track) => entry.add_track(track),
                            None        => ()
                        }
                    }
                    entry
                },
                None => {
                    println!("Failed to create entry database cache");
                    Entry {
                            id: 0,
                           url: url.to_string(),
                        tracks: tracks
                    }
                }
            }
        }
    }
}

fn show_track(req: &mut Request) -> IronResult<Response> {
    let ref track_id = req.extensions.get::<Router>().unwrap()
                          .find("track_id").unwrap();
    let tid: Option<i32> = track_id.parse().ok();
    let json_type = Header(ContentType(Mime::from_str("application/json").ok().unwrap()));
    match tid {
        Some(n) =>
            match Track::find_by_id(n) {
                Some(track) => {
                    let res = Response::with((status::Ok,
                                              json_type,
                                              track.to_json().to_string()));
                    Ok(res)
                },
                None =>
                    Ok(Response::with((status::Ok, json_type, "{}")))
            },
        None => Ok(Response::with((status::Ok, json_type, "{}")))
    }
}

fn main() {
    drop_tables();
    create_tables();

    let mut router = Router::new();
//    router.post(  "auth",                    signin);
    router.get(   "/playlistify",            playlistify);
//    router.get(   "/entry/:entry_id",        show_entry);
//    router.post(  "/entry/:entry_id",        update_entry);
//    router.delete("/entry/:entry_id",        destroy_entry);
//    router.put(   "/tracks/:track_id",       create_track);
    router.get(   "/tracks/:track_id",       show_track);
//    router.post(  "/tracks/:track_id",       update_track);
//    router.delete("/tracks/:track_id",       destroy_track);
//    router.post(  "/entry/:entry_id/tracks", add_track);*/

    let opt_port = std::env::var("PORT");
    let port_str = match opt_port {
        Ok(n)    => n,
        Err(_) => "8080".to_string()
    };
    let port: u16 = match port_str.trim().parse() {
        Ok(n) => n,
        Err(_) => {
            println!("Faild to parse port");
            return;
        }
    };
    println!("PORT {}", port_str);
    let ip = Ipv4Addr::new(0, 0, 0, 0);
    Iron::new(router).http(SocketAddrV4::new(ip, port)).unwrap();
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate yaml_rust;

extern crate lxd;

use std::env;
use std::fs::File;
use std::io::Read;

use yaml_rust::YamlLoader;

use lxd::{Container,LxdServer};

fn create_dividing_line(widths: &Vec<usize>) -> String {
    let mut dividing_line = String::new();
    dividing_line.push_str("+");
    for width in widths {
        dividing_line.push_str(&format!("{:-^1$}", "", width + 2));
        dividing_line.push_str("+");
    }
    dividing_line.push_str("\n");
    dividing_line
}

fn create_header_line(headers: &Vec<&str>, widths: &Vec<usize>) -> String {
    let mut header_line = String::new();
    header_line.push_str("|");
    for (n, header) in headers.iter().enumerate() {
        header_line.push_str(&format!("{:^1$}", &header, widths[n] + 2));
        header_line.push_str("|");
    }
    header_line.push_str("\n");
    header_line
}

fn create_content_line(item: &Vec<String>, widths: &Vec<usize>) -> String {
    let mut content_line = String::new();
    content_line.push_str("|");
    for (n, column_content) in item.iter().enumerate() {
        content_line.push_str(" ");
        content_line.push_str(&format!("{:1$}", &column_content, widths[n] + 1));
        content_line.push_str("|");
    }
    content_line.push_str("\n");
    content_line
}

fn format_output(headers: &Vec<&str>, items: &Vec<Vec<String>>) -> String {
    let mut widths = Vec::new();
    for header in headers {
        widths.push(header.len());
    }
    for item in items {
        for (n, column) in item.iter().enumerate() {
            if column.len() > widths[n] {
                widths[n] = column.len();
            }
        }
    }
    let dividing_line = create_dividing_line(&widths);
    let mut output_string = String::new();
    output_string.push_str(&dividing_line);
    output_string.push_str(&create_header_line(headers, &widths));
    output_string.push_str(&dividing_line);
    for item in items {
        output_string.push_str(&create_content_line(item, &widths));
    }
    output_string.push_str(&dividing_line);
    output_string
}

fn prepare_container_line(c: &Container) -> Vec<String> {
    let mut ipv4_address = String::new();
    let mut ipv6_address = String::new();
    for ip in &c.status.ips {
        if ip.protocol == "IPV4" && ip.address != "127.0.0.1" {
            ipv4_address = ip.address.clone();
        }
        if ip.protocol == "IPV6" && ip.address != "::1" {
            ipv6_address = ip.address.clone();
        }
    }
    let ephemeral = if c.ephemeral { "YES" } else { "NO" };
    vec![c.name.clone(), c.status.status.clone().to_uppercase(), ipv4_address.to_string(), ipv6_address.to_string(), ephemeral.to_string(), c.snapshot_urls.len().to_string()]
}

fn list() {
    let home_dir = env::var("HOME").unwrap();
    let mut config_file = File::open(home_dir.clone() + "/.config/lxc/config.yml").unwrap();
    let mut file_contents = String::new();
    config_file.read_to_string(&mut file_contents).unwrap();
    let lxd_config = YamlLoader::load_from_str(&file_contents).unwrap();
    let default_remote = lxd_config[0]["default-remote"].as_str().unwrap();
    let lxd_server = lxd_config[0]["remotes"][default_remote]["addr"].as_str().unwrap();
    let server = LxdServer::new(
        lxd_server,
        &(home_dir.clone() + "/.config/lxc/client.crt"),
        &(home_dir.clone() + "/.config/lxc/client.key")
    );
    let headers = vec!["NAME", "STATE", "IPV4", "IPV6", "EPHEMERAL", "SNAPSHOTS"];
    let container_items = server.list_containers().iter().map(prepare_container_line).collect();
    print!("{}", format_output(&headers, &container_items));
}

fn main() {
    let matches = clap_app!(
        myapp => (
            @subcommand list =>
                (about: "lists containers")
        )
    ).get_matches();

    match matches.subcommand_name() {
        Some("list") => list(),
        _ => println!("{}", matches.usage()),
    }
}
<commit_msg>Take a remote to query as an argument to list.<commit_after>extern crate clap;
extern crate yaml_rust;

extern crate lxd;

use std::env;
use std::fs::File;
use std::io::Read;

use clap::{App, Arg, ArgMatches, SubCommand};
use yaml_rust::YamlLoader;

use lxd::{Container,LxdServer};

fn create_dividing_line(widths: &Vec<usize>) -> String {
    let mut dividing_line = String::new();
    dividing_line.push_str("+");
    for width in widths {
        dividing_line.push_str(&format!("{:-^1$}", "", width + 2));
        dividing_line.push_str("+");
    }
    dividing_line.push_str("\n");
    dividing_line
}

fn create_header_line(headers: &Vec<&str>, widths: &Vec<usize>) -> String {
    let mut header_line = String::new();
    header_line.push_str("|");
    for (n, header) in headers.iter().enumerate() {
        header_line.push_str(&format!("{:^1$}", &header, widths[n] + 2));
        header_line.push_str("|");
    }
    header_line.push_str("\n");
    header_line
}

fn create_content_line(item: &Vec<String>, widths: &Vec<usize>) -> String {
    let mut content_line = String::new();
    content_line.push_str("|");
    for (n, column_content) in item.iter().enumerate() {
        content_line.push_str(" ");
        content_line.push_str(&format!("{:1$}", &column_content, widths[n] + 1));
        content_line.push_str("|");
    }
    content_line.push_str("\n");
    content_line
}

fn format_output(headers: &Vec<&str>, items: &Vec<Vec<String>>) -> String {
    let mut widths = Vec::new();
    for header in headers {
        widths.push(header.len());
    }
    for item in items {
        for (n, column) in item.iter().enumerate() {
            if column.len() > widths[n] {
                widths[n] = column.len();
            }
        }
    }
    let dividing_line = create_dividing_line(&widths);
    let mut output_string = String::new();
    output_string.push_str(&dividing_line);
    output_string.push_str(&create_header_line(headers, &widths));
    output_string.push_str(&dividing_line);
    for item in items {
        output_string.push_str(&create_content_line(item, &widths));
    }
    output_string.push_str(&dividing_line);
    output_string
}

fn prepare_container_line(c: &Container) -> Vec<String> {
    let mut ipv4_address = String::new();
    let mut ipv6_address = String::new();
    for ip in &c.status.ips {
        if ip.protocol == "IPV4" && ip.address != "127.0.0.1" {
            ipv4_address = ip.address.clone();
        }
        if ip.protocol == "IPV6" && ip.address != "::1" {
            ipv6_address = ip.address.clone();
        }
    }
    let ephemeral = if c.ephemeral { "YES" } else { "NO" };
    vec![c.name.clone(), c.status.status.clone().to_uppercase(), ipv4_address.to_string(), ipv6_address.to_string(), ephemeral.to_string(), c.snapshot_urls.len().to_string()]
}

fn list(matches: &ArgMatches) {
    let home_dir = env::var("HOME").unwrap();
    let mut config_file = File::open(home_dir.clone() + "/.config/lxc/config.yml").unwrap();
    let mut file_contents = String::new();
    config_file.read_to_string(&mut file_contents).unwrap();
    let lxd_config = YamlLoader::load_from_str(&file_contents).unwrap();
    let default_remote = lxd_config[0]["default-remote"].as_str().unwrap();
    let remote = matches.value_of("resource").unwrap_or(default_remote);
    let lxd_server = match lxd_config[0]["remotes"][remote]["addr"].as_str() {
        Some(remote_addr) => remote_addr,
        None => panic!("No remote named {} configured", remote)
    };
    let server = LxdServer::new(
        lxd_server,
        &(home_dir.clone() + "/.config/lxc/client.crt"),
        &(home_dir.clone() + "/.config/lxc/client.key")
    );
    let headers = vec!["NAME", "STATE", "IPV4", "IPV6", "EPHEMERAL", "SNAPSHOTS"];
    let container_items = server.list_containers().iter().map(prepare_container_line).collect();
    print!("{}", format_output(&headers, &container_items));
}

fn main() {
    let matches = App::new("lxd")
        .subcommand(SubCommand::with_name("list")
                    .arg(Arg::with_name("resource")
                         .help("the resource to use")
                         .required(true)
                         .index(1)))
        .get_matches();

    match matches.subcommand_name() {
        Some("list") => list(matches.subcommand_matches("list").unwrap()),
        _ => println!("{}", matches.usage()),
    }
}
<|endoftext|>


<commit_before>extern crate libusb;
use std::time::Duration;
use std::thread;
use std::process;
use std::io::{self, Write};

fn main() {
    let context = libusb::Context::new().unwrap();
    let device = find_device(&context);

    match device {
        Ok(handle) => {
            if let Err(e) = run(handle) {
                println!("{}", e);
                process::exit(2);
            }
        },
        Err(libusb::Error::NotFound) => {
            println!("Device not found");
            process::exit(1);
        }
        Err(e) => {
            println!("{}", e);
            process::exit(1);
        }
    }
}

fn find_device<'a>(context: &'a libusb::Context) -> Result<libusb::DeviceHandle<'a>, libusb::Error> {
    let devices = try!(context.devices());
    devices.iter().find(|device| {
        device.device_descriptor().ok().map_or(false, |desc| desc.vendor_id() == 0x05c6)
    }).map_or(Err(libusb::Error::NotFound), |device| device.open())
}

fn run(mut handle: libusb::DeviceHandle) -> Result<(), libusb::Error> {
    let mut stdout = io::stdout();
    println!("Found device");
    try!(handle.claim_interface(0));

    let mut read_buf = [0; 4096];

    loop {
        println!("Connecting DMSS");

        try!(handle.write_bulk(0x01, &[0x0c, 0x00, 0x00, 0x00,  0x00, 0x47, 0xb8, 0x7e], Duration::from_secs(10)));
        let len = try!(handle.read_bulk(0x81, &mut read_buf, Duration::from_secs(10)));
        let rx = &read_buf[0..len];

        if rx == &[0x13, 0x0c, 0x00, 0x00,  0x00, 0x00, 0x72, 0xce, 0x7e] {
            break;
        } else {
            thread::sleep(Duration::from_millis(100));
        }
    }

    try!(handle.write_bulk(0x01, &[0x7d, 0x5d, 0x04, 0x34,  0x21, 0x34, 0x21, 0x00,  0x00, 0x1f, 0x00, 0x00,  0x00, 0xab, 0xdf, 0x7e], Duration::from_secs(1)));
    println!("DMSS is connected");

    loop {
        let len = try!(handle.read_bulk(0x81, &mut read_buf, Duration::from_secs(0)));

        for packet in Hdlc::new(read_buf[0..len].iter().map(|&x| x)) {
            if packet.len() > 0 {
                parse_packet(&mut stdout, packet).unwrap();
            }
        }
    }
}

fn split_byte(s: &[u8], sep: u8) -> (&[u8], &[u8]) {
    if let Some(pos) = s.iter().position(|&i| i == sep) {
        let (a, b) = s.split_at(pos);
        (a, &b[1..])
    } else {
        (s, &[])
    }
}

fn le32(s: &[u8]) -> u32 {
    (s[0] as u32) | (s[1] as u32) << 8 | (s[2] as u32) << 16 | (s[3] as u32) << 24
}

fn le16(s: &[u8]) -> u16 {
    (s[0] as u16) | (s[1] as u16) << 8
}

fn parse_packet(stdout: &mut Write, packet: Vec<u8>) -> io::Result<()> {
    if packet.len() < 24 || packet[0] != 0x79 || packet[2] != 0 {
        print!("[unknown packet: ");
        for b in &packet { print!("{:02x}", b); }
        println!("]");
        return Ok(());
    }

    let (header, rest) = packet.split_at(20);
    let timestamp = le32(&header[6..10]);
    let lineno = le16(&header[12..14]);
    let (msg, rest) = split_byte(rest, 0);
    let (file, _) = split_byte(rest, 0);

    //for b in &packet { print!("{:02x} ", b); }

    try!(write!(stdout, "{:10}: ", timestamp));
    try!(stdout.write_all(msg));
    try!(write!(stdout, " ("));
    try!(stdout.write_all(file));
    try!(write!(stdout, ":{})\n", lineno));

    Ok(())
}

/// Iterator over packets in a HDLC-framed sequence of bytes
struct Hdlc<I> {
    input: I,
    end: bool,
}

impl<I> Hdlc<I> where I:Iterator<Item=u8> {
    fn new(input: I) -> Hdlc<I> {
        Hdlc {
            input: input,
            end: false,
        }
    }

    fn input_byte(&mut self) -> Option<u8> {
        let x = self.input.next();
        self.end |= x.is_none();
        x
    }
}

impl<I> Iterator for Hdlc<I> where I:Iterator<Item=u8> {
    type Item = Vec<u8>;
    fn next(&mut self) -> Option<Vec<u8>> {
        if self.end {
            return None;
        }

        let mut out = vec![];

        loop {
            match self.input_byte() {
                Some(0x7e) | None => break,
                Some(0x7d) => {
                    match self.input_byte() {
                        Some(c) => out.push(c ^ 0x20),
                        None => break,
                    }
                }
                Some(c) => out.push(c),
            };
        }

        Some(out)
    }
}<commit_msg>Strip trailing newlines<commit_after>extern crate libusb;
use std::time::Duration;
use std::thread;
use std::process;
use std::io::{self, Write};

fn main() {
    let context = libusb::Context::new().unwrap();
    let device = find_device(&context);

    match device {
        Ok(handle) => {
            if let Err(e) = run(handle) {
                println!("{}", e);
                process::exit(2);
            }
        },
        Err(libusb::Error::NotFound) => {
            println!("Device not found");
            process::exit(1);
        }
        Err(e) => {
            println!("{}", e);
            process::exit(1);
        }
    }
}

fn find_device<'a>(context: &'a libusb::Context) -> Result<libusb::DeviceHandle<'a>, libusb::Error> {
    let devices = try!(context.devices());
    devices.iter().find(|device| {
        device.device_descriptor().ok().map_or(false, |desc| desc.vendor_id() == 0x05c6)
    }).map_or(Err(libusb::Error::NotFound), |device| device.open())
}

fn run(mut handle: libusb::DeviceHandle) -> Result<(), libusb::Error> {
    let mut stdout = io::stdout();
    println!("Found device");
    try!(handle.claim_interface(0));

    let mut read_buf = [0; 4096];

    loop {
        println!("Connecting DMSS");

        try!(handle.write_bulk(0x01, &[0x0c, 0x00, 0x00, 0x00,  0x00, 0x47, 0xb8, 0x7e], Duration::from_secs(10)));
        let len = try!(handle.read_bulk(0x81, &mut read_buf, Duration::from_secs(10)));
        let rx = &read_buf[0..len];

        if rx == &[0x13, 0x0c, 0x00, 0x00,  0x00, 0x00, 0x72, 0xce, 0x7e] {
            break;
        } else {
            thread::sleep(Duration::from_millis(100));
        }
    }

    try!(handle.write_bulk(0x01, &[0x7d, 0x5d, 0x04, 0x34,  0x21, 0x34, 0x21, 0x00,  0x00, 0x1f, 0x00, 0x00,  0x00, 0xab, 0xdf, 0x7e], Duration::from_secs(1)));
    println!("DMSS is connected");

    loop {
        let len = try!(handle.read_bulk(0x81, &mut read_buf, Duration::from_secs(0)));

        for packet in Hdlc::new(read_buf[0..len].iter().map(|&x| x)) {
            if packet.len() > 0 {
                parse_packet(&mut stdout, packet).unwrap();
            }
        }
    }
}

fn split_byte(s: &[u8], sep: u8) -> (&[u8], &[u8]) {
    if let Some(pos) = s.iter().position(|&i| i == sep) {
        let (a, b) = s.split_at(pos);
        (a, &b[1..])
    } else {
        (s, &[])
    }
}

fn strip_trailing_newlines(mut s: &[u8]) -> &[u8] {
    while s.ends_with(b"\n") {
        s = &s[0..s.len()-1];
    }
    s
}

fn le32(s: &[u8]) -> u32 {
    (s[0] as u32) | (s[1] as u32) << 8 | (s[2] as u32) << 16 | (s[3] as u32) << 24
}

fn le16(s: &[u8]) -> u16 {
    (s[0] as u16) | (s[1] as u16) << 8
}

fn parse_packet(stdout: &mut Write, packet: Vec<u8>) -> io::Result<()> {
    if packet.len() < 24 || packet[0] != 0x79 || packet[2] != 0 {
        print!("[unknown packet: ");
        for b in &packet { print!("{:02x}", b); }
        println!("]");
        return Ok(());
    }

    let (header, rest) = packet.split_at(20);
    let timestamp = le32(&header[6..10]);
    let lineno = le16(&header[12..14]);
    let (msg, rest) = split_byte(rest, 0);
    let (file, _) = split_byte(rest, 0);

    //for b in &packet { print!("{:02x} ", b); }

    try!(write!(stdout, "{:10}: ", timestamp));
    try!(stdout.write_all(strip_trailing_newlines(msg)));
    try!(write!(stdout, " ("));
    try!(stdout.write_all(file));
    try!(write!(stdout, ":{})\n", lineno));

    Ok(())
}

/// Iterator over packets in a HDLC-framed sequence of bytes
struct Hdlc<I> {
    input: I,
    end: bool,
}

impl<I> Hdlc<I> where I:Iterator<Item=u8> {
    fn new(input: I) -> Hdlc<I> {
        Hdlc {
            input: input,
            end: false,
        }
    }

    fn input_byte(&mut self) -> Option<u8> {
        let x = self.input.next();
        self.end |= x.is_none();
        x
    }
}

impl<I> Iterator for Hdlc<I> where I:Iterator<Item=u8> {
    type Item = Vec<u8>;
    fn next(&mut self) -> Option<Vec<u8>> {
        if self.end {
            return None;
        }

        let mut out = vec![];

        loop {
            match self.input_byte() {
                Some(0x7e) | None => break,
                Some(0x7d) => {
                    match self.input_byte() {
                        Some(c) => out.push(c ^ 0x20),
                        None => break,
                    }
                }
                Some(c) => out.push(c),
            };
        }

        Some(out)
    }
}<|endoftext|>


<commit_before>extern crate base64;
extern crate docopt;
extern crate rustc_serialize;

use std::error::Error;
use std::fs::File;
use std::io::{self, Write};
use std::path::Path;
use std::str;

use docopt::Docopt;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

const BUF_SIZE: usize = 4 * 1024;

const USAGE: &'static str = "
Base64 encode or decode FILE, or standard input, to standard output.

Usage: base64 [options] [<file>]
       base64 (--help | --version)

Options:
  -d --decode   decode data
  -h --help     display this help and exit
  --version     output version information and exit
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_file: String,
    flag_decode: bool,
    flag_help: bool,
    flag_version: bool,
}

fn process_file<F>(func: F, file: &mut io::Read)
where F: Fn(&mut [u8; BUF_SIZE], usize) {
    let mut buf: [u8; BUF_SIZE] = [0; BUF_SIZE];

    loop {
        match file.read(&mut buf[..]) {
            Ok(count) => func(&mut buf, count),
            Err(err) => panic!(err),
        }
    }
}

fn encode(file: &mut io::Read) {
    let mut buf: [u8; BUF_SIZE] = [0; BUF_SIZE];
    let mut stdout = io::stdout();

    loop {
        match file.read(&mut buf[..]) {
            Ok(count) => {
                if count == 0 {
                    print!("\n");
                    break;
                }
                let res = base64::encode(&buf);
                stdout.write(&res[..count + 1]).unwrap();
            },
            Err(err) => panic!(err),
        }
    }
}

fn decode(file: &mut io::Read) {
    let mut buf: [u8; BUF_SIZE] = [0; BUF_SIZE];
    let mut stdout = io::stdout();
    let mut stderr = io::stderr();

    loop {
        match file.read(&mut buf[..]) {
            Ok(count) => {
                if count == 0 {
                     break;
                }
                println!("{:?}", str::from_utf8(&buf[..count])); // DEBUG
                match base64::decode(&buf[..count]) {
                    Ok(res) => {
                        stdout.write(&res).unwrap();
                    },
                    Err(msg) => {
                        stderr.write(msg.as_bytes()).unwrap();

                    },
                }
            },
            Err(err) => panic!(err),
        }
    }
}

fn process(file: &mut io::Read, is_decode: bool) {
    if is_decode {
        decode(file);
    } else {
        encode(file);
    }
}

fn main() {
    /*
    match base64::decode(b"cXdlCg==") {
        Err(msg) => println!("{}", msg),
        Ok(data) => println!("{}", str::from_utf8(&data).unwrap()),
    }
    */
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("base64 v{}", VERSION);
        return;
    }

    if args.arg_file != "" {
        let path = Path::new(&args.arg_file);
        let mut file = match File::open(&path) {
            Err(why) => panic!("couldn't open {}: {}",
                               path.display(), why.description()),
            Ok(file) => file,
        };
        process(&mut file, args.flag_decode);
    } else {
        process(&mut io::stdin(), args.flag_decode);
    }
}
<commit_msg>Cleanup<commit_after>extern crate base64;
extern crate docopt;
extern crate rustc_serialize;

use std::error::Error;
use std::fs::File;
use std::io::{self, Write};
use std::path::Path;
use std::str;

use docopt::Docopt;

const VERSION: &'static str = env!("CARGO_PKG_VERSION");

const BUF_SIZE: usize = 4 * 1024;

const USAGE: &'static str = "
Base64 encode or decode FILE, or standard input, to standard output.

Usage: base64 [options] [<file>]
       base64 (--help | --version)

Options:
  -d --decode   decode data
  -h --help     display this help and exit
  --version     output version information and exit
";

#[derive(Debug, RustcDecodable)]
struct Args {
    arg_file: String,
    flag_decode: bool,
    flag_help: bool,
    flag_version: bool,
}

fn encode(file: &mut io::Read) {
    let mut buf: [u8; BUF_SIZE] = [0; BUF_SIZE];
    let mut stdout = io::stdout();

    loop {
        match file.read(&mut buf[..]) {
            Ok(count) => {
                if count == 0 {
                    print!("\n");
                    break;
                }
                let res = base64::encode(&buf);
                stdout.write(&res[..count + 1]).unwrap();
            },
            Err(err) => panic!(err),
        }
    }
}

fn decode(file: &mut io::Read) {
    let mut buf: [u8; BUF_SIZE] = [0; BUF_SIZE];
    let mut stdout = io::stdout();
    let mut stderr = io::stderr();

    loop {
        match file.read(&mut buf[..]) {
            Ok(count) => {
                if count == 0 {
                     break;
                }
                match base64::decode(&buf[..count]) {
                    Ok(res) => {
                        stdout.write(&res).unwrap();
                    },
                    Err(msg) => {
                        stderr.write(msg.as_bytes()).unwrap();

                    },
                }
            },
            Err(err) => panic!(err),
        }
    }
}

fn process(file: &mut io::Read, is_decode: bool) {
    if is_decode {
        decode(file);
    } else {
        encode(file);
    }
}

fn main() {
    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.decode())
        .unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("base64 v{}", VERSION);
        return;
    }

    if args.arg_file != "" {
        let path = Path::new(&args.arg_file);
        let mut file = match File::open(&path) {
            Err(why) => panic!("couldn't open {}: {}",
                               path.display(), why.description()),
            Ok(file) => file,
        };
        process(&mut file, args.flag_decode);
    } else {
        process(&mut io::stdin(), args.flag_decode);
    }
}
<|endoftext|>


<commit_before>extern crate num;
extern crate portaudio;
extern crate rtlsdr;

use portaudio as pa;
use rtlsdr::{RTLSDRDevice, RTLSDRError};
use num::complex::Complex;
use std::env;
use std::collections::VecDeque;
use std::f32::consts::PI;
use std::error::Error;

// Make sure that the buffer size is radix-2, otherwise the read_sync function
// will fail with an error code of -8.
const SDR_BUF_SIZE: usize = 524288;

// Separating out the buffer sizes seems to help some.
const BUF_SIZE: usize = 262144;

// Most other sample rates fail, but this one works for my particular device.
// I will investigate exactly what's happening here and generate a list of
// possible sample rates (if there are others besides this one).
const SAMPLE_RATE: u32 = 2800000;

// Here is the sample rate of the output waveform we'll try to use.
const SAMPLE_RATE_AUDIO: f64 = 48000.0;

// The number of channels being used for audio.
const CHANNELS: i32 = 2;

fn main() {
    match run() {
        Ok(()) => (),
        Err(e) => println!("ERROR: {}", e),
    }
}

struct RTLSDR {
    rtlsdr: RTLSDRDevice,
}

unsafe impl Send for RTLSDR {}

fn run() -> Result<(), Box<Error>> {
    let sdr_index: i32 = env::args().nth(1)
        .expect("Please specify the SDR index")
        .parse::<i32>()?;
    let fm_freq_mhz: f32 = env::args().nth(2)
        .expect("Please specify a center frequency (MHz).")
        .parse::<f32>()?;
    let bandwidth: u32 = env::args().nth(3)
        .expect("Please specify the bandwidth.")
        .parse::<u32>()?;

    // Initialize the PortAudio class.
    let audio = try!(pa::PortAudio::new());
    let settings =
        try!(audio.default_output_stream_settings(CHANNELS,
                                                  SAMPLE_RATE_AUDIO,
                                                  BUF_SIZE as u32));
    let mut stream = try!(audio.open_blocking_stream(settings));
    try!(stream.start());

    let fm_freq: u32 = (fm_freq_mhz * 1e6) as u32;

    // TODO: we need to figure out the appropriate decimation here. I've seen
    // talk that we decimate based on the bandwidth but after decimation the
    // audio signal comes out much slower. Changing the decimation value
    // seems to affect the speed of the audio, so I need to establish the
    // relationship here. Just playing with the decimation value, 30 seems
    // to work well but I need to figure out _why_.
    // let dec_rate = (SAMPLE_RATE as f64 / bandwidth as f64) as usize;
    let dec_rate = 30;
    println!("dec rate is {}", dec_rate);
    let mut sdr =
        RTLSDR { rtlsdr: init_sdr(sdr_index, fm_freq, bandwidth).unwrap() };
    let mut prev = Complex::new(0.0, 0.0);

    loop {
        // Read the samples and decimate down to match the sample rate of
        // the audio that'll be going out.
        let bytes = sdr.rtlsdr.read_sync(SDR_BUF_SIZE).unwrap();
        let mut iq_vec = decimate(read_samples(bytes).unwrap().as_slice(),
                                  dec_rate);
        iq_vec = filter(iq_vec, moving_average(10));

        // After decimation, demodulate the signal and send out of the
        // thread to the receiver.
        let res = demod_fm(iq_vec, prev);
        let mut demod_iq = res.0; // filter_queue(res.0, moving_average(1));
        prev = res.1;

        // Get the write stream and write our samples to the stream.
        let out_frames = match stream.write_available() {
            Ok(available) => {
                match available {
                    pa::StreamAvailable::Frames(frames) => frames as u32,
                    _ => return Err(Box::new(pa::error::Error::NullCallback)),
                }
            }
            Err(e) => return Err(Box::new(e)),
        };

        let buffer_frames = (demod_iq.len() / CHANNELS as usize) as u32;
        let write_frames = std::cmp::min(buffer_frames, out_frames);
        if write_frames == 0 {
            println!("write frames is 0, are we not generating fast enough?");
        }
        let n_write_samples = write_frames as usize * CHANNELS as usize;

        stream.write(write_frames,
                   |output| for ix in 0..n_write_samples as usize {
                       output[ix] = 0.005 * demod_iq.pop_front().unwrap();
                   })?;
    }
}

fn init_sdr(sdr_index: i32,
            fm_freq: u32,
            bandwidth: u32)
            -> Result<RTLSDRDevice, RTLSDRError> {
    let mut sdr = try!(rtlsdr::open(sdr_index));
    try!(sdr.set_center_freq(fm_freq));
    try!(sdr.set_sample_rate(SAMPLE_RATE));
    try!(sdr.set_tuner_bandwidth(bandwidth));
    try!(sdr.set_agc_mode(false));
    try!(sdr.reset_buffer());
    Ok(sdr)
}

fn decimate(signal: &[Complex<f32>], dec_rate: usize) -> Vec<Complex<f32>> {
    let mut ix = 0;
    let new_size = (signal.len() / dec_rate + 1) as usize;
    let mut signal_dec = Vec::<Complex<f32>>::with_capacity(new_size);
    while ix < signal.len() {
        signal_dec.push(signal[ix]);
        ix += dec_rate;
    }
    signal_dec
}

fn read_samples(bytes: Vec<u8>) -> Option<Vec<Complex<f32>>> {
    // First, check that we've been given an even number of bytes. Not sure
    // what to do if we don't get I and Q.
    let bytes_len = bytes.len();
    if bytes_len % 2 != 0 {
        println!("Received an odd number of samples.");
        return None;

    }
    let mut iq_vec: Vec<Complex<f32>> = Vec::with_capacity(bytes_len / 2);

    // Write the values to the complex value and normalize from [0, 255] to
    // [-1, 1]. If we don't normalize, it seems we don't get anything.
    for iq in bytes.chunks(2) {
        let iq_cmplx = Complex::new((iq[0] as f32 - 127.0) / 127.0,
                                    (iq[1] as f32 - 127.0) / 127.0);
        iq_vec.push(iq_cmplx);
    }
    Some(iq_vec)
}

fn demod_fm(iq: Vec<Complex<f32>>,
            prev: Complex<f32>)
            -> (VecDeque<f32>, Complex<f32>) {

    let mut p = prev.clone();
    let mut demod_queue: VecDeque<f32> = VecDeque::with_capacity(iq.len());
    let gain = SAMPLE_RATE as f32 / (2.0 * PI * 75e3 / 8.0);

    for samp in iq.iter() {
        let conj = p.conj() * samp;
        let fm_val = conj.im.atan2(conj.re);
        demod_queue.push_back(fm_val * gain);
        p = *samp;
    }
    (demod_queue, p)
}

fn filter(samples: Vec<Complex<f32>>, taps: Vec<f32>) -> Vec<Complex<f32>> {
    let mut filt_samps: Vec<Complex<f32>> = Vec::new();
    for window in samples.as_slice().windows(taps.len()) {
        let iter = window.iter().zip(taps.iter());
        let filt_samp = iter.map(|(x, y)| x * y)
            .fold(Complex::new(0.0, 0.0), |acc, x| acc + x);
        filt_samps.push(filt_samp);
    }
    filt_samps
}

fn filter_queue(samples: VecDeque<f32>, taps: Vec<f32>) -> VecDeque<f32> {
    let mut filt_samps: VecDeque<f32> = VecDeque::new();
    for window in samples.as_slices().0.windows(taps.len()) {
        let iter = window.iter().zip(taps.iter());
        let filt_samp = iter.map(|(x, y)| x * y).sum();
        filt_samps.push_back(filt_samp);
    }
    filt_samps

}

fn moving_average(ntaps: usize) -> Vec<f32> {
    vec![1.0 / ntaps as f32; ntaps]
}
<commit_msg>Added low pass filter and windowing function support<commit_after>extern crate num;
extern crate portaudio;
extern crate rtlsdr;

use portaudio as pa;
use rtlsdr::{RTLSDRDevice, RTLSDRError};
use num::complex::Complex;
use std::env;
use std::collections::VecDeque;
use std::f32::consts::PI;
use std::error::Error;

// Make sure that the buffer size is radix-2, otherwise the read_sync function
// will fail with an error code of -8.
const SDR_BUF_SIZE: usize = 524288;

// Separating out the buffer sizes seems to help some.
const BUF_SIZE: usize = 262144;

// Most other sample rates fail, but this one works for my particular device.
// I will investigate exactly what's happening here and generate a list of
// possible sample rates (if there are others besides this one).
const SAMPLE_RATE: u32 = 2800000;

// Here is the sample rate of the output waveform we'll try to use.
const SAMPLE_RATE_AUDIO: f32 = 44100.0;

// The number of channels being used for audio.
const CHANNELS: i32 = 2;

fn main() {
    match run() {
        Ok(()) => (),
        Err(e) => println!("ERROR: {}", e),
    }
}

struct RTLSDR {
    rtlsdr: RTLSDRDevice,
}

unsafe impl Send for RTLSDR {}

fn run() -> Result<(), Box<Error>> {
    let sdr_index: i32 = env::args().nth(1)
        .expect("Please specify the SDR index")
        .parse::<i32>()?;
    let fm_freq_mhz: f32 = env::args().nth(2)
        .expect("Please specify a center frequency (MHz).")
        .parse::<f32>()?;
    let bandwidth: u32 = env::args().nth(3)
        .expect("Please specify the bandwidth.")
        .parse::<u32>()?;

    // Initialize the PortAudio class.
    let audio = try!(pa::PortAudio::new());
    let settings =
        try!(audio.default_output_stream_settings(CHANNELS,
                                                  SAMPLE_RATE_AUDIO as f64,
                                                  BUF_SIZE as u32));
    let mut stream = try!(audio.open_blocking_stream(settings));
    try!(stream.start());

    let fm_freq: u32 = (fm_freq_mhz * 1e6) as u32;

    // TODO: we need to figure out the appropriate decimation here. I've seen
    // talk that we decimate based on the bandwidth but after decimation the
    // audio signal comes out much slower. Changing the decimation value
    // seems to affect the speed of the audio, so I need to establish the
    // relationship here. Just playing with the decimation value, 30 seems
    // to work well but I need to figure out _why_.
    let dec_rate = (SAMPLE_RATE as f64 / bandwidth as f64 * 2.0) as usize;
    println!("dec rate is {}", dec_rate);
    let mut sdr =
        RTLSDR { rtlsdr: init_sdr(sdr_index, fm_freq, bandwidth).unwrap() };

    // This holds the previous value for the moving average filter so we can
    // keep track in between iterations of reading from the SDR.
    let mut prev = Complex::new(0.0, 0.0);

    loop {
        // Read the samples and decimate down to match the sample rate of
        // the audio that'll be going out.
        let bytes = sdr.rtlsdr.read_sync(SDR_BUF_SIZE).unwrap();
        let mut iq_vec = read_samples(bytes).unwrap();
        iq_vec = decimate(iq_vec.as_slice(), dec_rate);

        // After decimation, demodulate the signal and send out of the
        // thread to the receiver.
        let res = demod_fm(iq_vec, prev);
        let mut demod_iq =
            filter_queue(res.0, low_pass(50000.0, SAMPLE_RATE as f32, 64));
        prev = res.1;

        // Get the write stream and write our samples to the stream.
        let out_frames = match stream.write_available() {
            Ok(available) => {
                match available {
                    pa::StreamAvailable::Frames(frames) => frames as u32,
                    _ => return Err(Box::new(pa::error::Error::NullCallback)),
                }
            }
            Err(e) => return Err(Box::new(e)),
        };

        let buffer_frames = (demod_iq.len() / CHANNELS as usize) as u32;
        let write_frames = std::cmp::min(buffer_frames, out_frames);
        if write_frames == 0 {
            println!("write frames is 0, are we not generating fast enough?");
        }
        let n_write_samples = write_frames as usize * CHANNELS as usize;

        stream.write(write_frames,
                   |output| for ix in 0..n_write_samples as usize {
                       output[ix] = 0.005 * demod_iq.pop_front().unwrap();
                   })?;
    }
}

fn init_sdr(sdr_index: i32,
            fm_freq: u32,
            bandwidth: u32)
            -> Result<RTLSDRDevice, RTLSDRError> {
    let mut sdr = try!(rtlsdr::open(sdr_index));
    try!(sdr.set_center_freq(fm_freq));
    try!(sdr.set_sample_rate(SAMPLE_RATE));
    try!(sdr.set_tuner_bandwidth(bandwidth));
    try!(sdr.set_agc_mode(false));
    try!(sdr.reset_buffer());
    Ok(sdr)
}

fn decimate(signal: &[Complex<f32>], dec_rate: usize) -> Vec<Complex<f32>> {
    let mut ix = 0;
    let new_size = (signal.len() / dec_rate + 1) as usize;
    let mut signal_dec = Vec::<Complex<f32>>::with_capacity(new_size);
    while ix < signal.len() {
        signal_dec.push(signal[ix]);
        ix += dec_rate;
    }
    signal_dec
}

fn decimate_queue(signal: &[f32], dec_rate: usize) -> VecDeque<f32> {
    let mut ix = 0;
    let new_size = (signal.len() / dec_rate + 1) as usize;
    let mut signal_dec = VecDeque::<f32>::with_capacity(new_size);
    while ix < signal.len() {
        signal_dec.push_back(signal[ix]);
        ix += dec_rate;
    }
    signal_dec
}

fn read_samples(bytes: Vec<u8>) -> Option<Vec<Complex<f32>>> {
    // First, check that we've been given an even number of bytes. Not sure
    // what to do if we don't get I and Q.
    let bytes_len = bytes.len();
    if bytes_len % 2 != 0 {
        println!("Received an odd number of samples.");
        return None;

    }
    let mut iq_vec: Vec<Complex<f32>> = Vec::with_capacity(bytes_len / 2);

    // Write the values to the complex value and normalize from [0, 255] to
    // [-1, 1]. If we don't normalize, it seems we don't get anything.
    for iq in bytes.chunks(2) {
        let iq_cmplx = Complex::new((iq[0] as f32 - 127.0) / 127.0,
                                    (iq[1] as f32 - 127.0) / 127.0);
        iq_vec.push(iq_cmplx);
    }
    Some(iq_vec)
}

fn demod_fm(iq: Vec<Complex<f32>>,
            prev: Complex<f32>)
            -> (VecDeque<f32>, Complex<f32>) {

    let mut p = prev.clone();
    let mut demod_queue: VecDeque<f32> = VecDeque::with_capacity(iq.len());
    let gain = SAMPLE_RATE as f32 / (2.0 * PI * 75e3 / 8.0);

    for samp in iq.iter() {
        let conj = p.conj() * samp;
        let fm_val = conj.im.atan2(conj.re);
        demod_queue.push_back(fm_val * gain);
        p = *samp;
    }
    (demod_queue, p)
}

fn filter(samples: Vec<Complex<f32>>, taps: Vec<f32>) -> Vec<Complex<f32>> {
    let mut filt_samps: Vec<Complex<f32>> = Vec::new();
    for window in samples.as_slice().windows(taps.len()) {
        let iter = window.iter().zip(taps.iter());
        let filt_samp = iter.map(|(x, y)| x * y)
            .fold(Complex::new(0.0, 0.0), |acc, x| acc + x);
        filt_samps.push(filt_samp);
    }
    filt_samps
}

fn filter_queue(samples: VecDeque<f32>, taps: Vec<f32>) -> VecDeque<f32> {
    let mut filt_samps: VecDeque<f32> = VecDeque::new();
    for window in samples.as_slices().0.windows(taps.len()) {
        let iter = window.iter().zip(taps.iter());
        let filt_samp = iter.map(|(x, y)| x * y).sum();
        filt_samps.push_back(filt_samp);
    }
    filt_samps

}

fn hamming(ntaps: usize) -> Vec<f32> {
    let m: f32 = ntaps as f32 - 1.0;
    let mut taps: Vec<f32> = Vec::with_capacity(ntaps);
    for ix in 0..ntaps {
        taps.push(0.54 - 0.46 * (2.0 * PI * ix as f32 / m).cos());
    }
    taps
}

fn low_pass(cutoff_freq: f32, sample_rate: f32, ntaps: usize) -> Vec<f32> {
    let wc: f32 = 2.0 * PI * cutoff_freq / sample_rate;
    let mut taps: Vec<f32> = Vec::with_capacity(ntaps);
    let win_taps: Vec<f32> = hamming(ntaps);
    let ntaps_i: i32 = ntaps as i32;
    let m: i32 = (ntaps_i - 1) / 2;

    for ix in -m..m + 1 {
        let ix_f: f32 = ix as f32;
        if ix == 0 {
            let hi = wc / PI * win_taps[(ix + m) as usize];
            taps.push(hi);
        } else {
            let hi = (ix_f * wc).sin() / (ix_f * PI) *
                     win_taps[(ix + m) as usize];
            taps.push(hi);
        }
    }
    taps
}

#[inline]
fn moving_average(ntaps: usize) -> Vec<f32> {
    vec![1.0 / ntaps as f32; ntaps]
}
<|endoftext|>


<commit_before>extern crate glob;
extern crate http;
extern crate hyper;
extern crate hyper_tls;
#[macro_use]
extern crate lazy_static;
extern crate libc;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate structopt;
extern crate toml;

mod error;
mod ghclient;
mod statics;
mod structs;
#[macro_use]
mod syslog;

use statics::CONF_PATH;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
enum Command {
    /// Gets user public key
    #[structopt(name = "key")]
    Key {
        #[structopt(parse(from_str))]
        user: String,
    },
    /// Executes pam check
    #[structopt(name = "pam")]
    Pam,
    /// Check configuration
    #[structopt(name = "check")]
    Check {
        #[structopt(parse(from_os_str))]
        confpath: std::path::PathBuf,
    },
    /// Cleans caches up
    #[structopt(name = "cleanup")]
    CleanUp,
    /// Get rate limit for github api
    #[structopt(name = "ratelimit")]
    RateLimit,
    /// Displays version details
    #[structopt(name = "version")]
    Version,
    /// Displays completion
    #[structopt(name = "completion")]
    Completion {
        #[structopt(subcommand)]
        shell: Shell,
    },
}

#[derive(Debug, StructOpt)]
enum Shell {
    #[structopt(name = "bash")]
    Bash,
    #[structopt(name = "fish")]
    Fish,
    #[structopt(name = "zsh")]
    Zsh,
    #[structopt(name = "powershell")]
    PowerShell,
    #[structopt(name = "elvish")]
    Elvish,
}

fn main() {
    let command = Command::from_args();

    use ghclient::GithubClient;
    use std::env;
    use std::process;
    use structs::Config;

    match command {
        Command::Check { confpath } => match Config::from_path(&confpath) {
            Ok(_) => process::exit(0),
            Err(_) => process::exit(11),
        },
        Command::Key { user } => {
            match Config::from_path(&CONF_PATH).and_then(|conf| Ok(GithubClient::new(&conf)))
                                               .and_then(|client| client.print_user_public_key(&user))
            {
                Ok(_) => process::exit(0),
                Err(_) => process::exit(21),
            }
        }
        Command::Pam => match env::var("PAM_USER") {
            Ok(user) => match Config::from_path(&CONF_PATH).and_then(|conf| Ok(GithubClient::new(&conf)))
                                                           .and_then(|client| client.check_pam(&user))
            {
                Ok(true) => process::exit(0),
                Ok(false) => process::exit(1),
                Err(_) => process::exit(31),
            },
            Err(_) => process::exit(41),
        },
        Command::CleanUp => match Config::from_path(&CONF_PATH).and_then(|conf| Ok(GithubClient::new(&conf)))
                                                               .and_then(|client| client.clear_all_caches())
        {
            Ok(_) => process::exit(0),
            Err(_) => process::exit(51),
        },
        Command::RateLimit => match Config::from_path(&CONF_PATH).and_then(|conf| Ok(GithubClient::new(&conf)))
                                                                 .and_then(|client| client.print_rate_limit())
        {
            Ok(_) => process::exit(0),
            Err(_) => process::exit(61),
        },
        Command::Version => {
            println!("{}",
                     concat!(env!("CARGO_PKG_VERSION"),
                             include_str!(concat!(env!("OUT_DIR"), "/commit-info.txt"))));
        }
        Command::Completion { shell } => {
            let shell = match shell {
                Shell::Bash => structopt::clap::Shell::Bash,
                Shell::Fish => structopt::clap::Shell::Fish,
                Shell::Zsh => structopt::clap::Shell::Zsh,
                Shell::PowerShell => structopt::clap::Shell::PowerShell,
                Shell::Elvish => structopt::clap::Shell::Elvish,
            };
            Command::clap().gen_completions_to("sectora", shell, &mut std::io::stdout());
        }
    };
}
<commit_msg>Simplify using rename_all<commit_after>extern crate glob;
extern crate http;
extern crate hyper;
extern crate hyper_tls;
#[macro_use]
extern crate lazy_static;
extern crate libc;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;
extern crate structopt;
extern crate toml;

mod error;
mod ghclient;
mod statics;
mod structs;
#[macro_use]
mod syslog;

use statics::CONF_PATH;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
enum Command {
    /// Gets user public key
    Key {
        #[structopt(parse(from_str))]
        user: String,
    },
    /// Executes pam check
    Pam,
    /// Check configuration
    Check {
        #[structopt(parse(from_os_str))]
        confpath: std::path::PathBuf,
    },
    /// Cleans caches up
    CleanUp,
    /// Get rate limit for github api
    RateLimit,
    /// Displays version details
    Version,
    /// Displays completion
    Completion {
        #[structopt(subcommand)]
        shell: Shell,
    },
}

#[derive(Debug, StructOpt)]
#[structopt(rename_all = "kebab-case")]
enum Shell {
    Bash,
    Fish,
    Zsh,
    PowerShell,
    Elvish,
}

fn main() {
    let command = Command::from_args();

    use ghclient::GithubClient;
    use std::env;
    use std::process;
    use structs::Config;

    match command {
        Command::Check { confpath } => match Config::from_path(&confpath) {
            Ok(_) => process::exit(0),
            Err(_) => process::exit(11),
        },
        Command::Key { user } => {
            match Config::from_path(&CONF_PATH).and_then(|conf| Ok(GithubClient::new(&conf)))
                                               .and_then(|client| client.print_user_public_key(&user))
            {
                Ok(_) => process::exit(0),
                Err(_) => process::exit(21),
            }
        }
        Command::Pam => match env::var("PAM_USER") {
            Ok(user) => match Config::from_path(&CONF_PATH).and_then(|conf| Ok(GithubClient::new(&conf)))
                                                           .and_then(|client| client.check_pam(&user))
            {
                Ok(true) => process::exit(0),
                Ok(false) => process::exit(1),
                Err(_) => process::exit(31),
            },
            Err(_) => process::exit(41),
        },
        Command::CleanUp => match Config::from_path(&CONF_PATH).and_then(|conf| Ok(GithubClient::new(&conf)))
                                                               .and_then(|client| client.clear_all_caches())
        {
            Ok(_) => process::exit(0),
            Err(_) => process::exit(51),
        },
        Command::RateLimit => match Config::from_path(&CONF_PATH).and_then(|conf| Ok(GithubClient::new(&conf)))
                                                                 .and_then(|client| client.print_rate_limit())
        {
            Ok(_) => process::exit(0),
            Err(_) => process::exit(61),
        },
        Command::Version => {
            println!("{}",
                     concat!(env!("CARGO_PKG_VERSION"),
                             include_str!(concat!(env!("OUT_DIR"), "/commit-info.txt"))));
        }
        Command::Completion { shell } => {
            let shell = match shell {
                Shell::Bash => structopt::clap::Shell::Bash,
                Shell::Fish => structopt::clap::Shell::Fish,
                Shell::Zsh => structopt::clap::Shell::Zsh,
                Shell::PowerShell => structopt::clap::Shell::PowerShell,
                Shell::Elvish => structopt::clap::Shell::Elvish,
            };
            Command::clap().gen_completions_to("sectora", shell, &mut std::io::stdout());
        }
    };
}
<|endoftext|>


<commit_before>extern crate bit_vector;
extern crate num;

use std::cmp;
use std::env;
use std::fmt::Display;
use std::str::FromStr;

use bit_vector::{BitVector,BitSliceMut};

use num::Bounded;

macro_rules! err_exit {
    ($($arg:tt)*) => (
        {
            use std::io::prelude::*;
            use std::process;
            if let Err(e) = write!(&mut ::std::io::stderr(), "{}\n", format_args!($($arg)*)) {
                panic!("Failed to write to stderr.\
                    \nOriginal error output: {}\
                    \nSecondary error writing to stderr: {}", format!($($arg)*), e);
            }
            process::exit(1)
        }
    )
}

type SieveStorage = usize;

fn main() {
    let requested_threads: usize = parse_command_line_argument(1, "threads");
    let max_prime: usize = parse_command_line_argument(2, "max_prime");

    let threads = calculate_actual_threads(requested_threads, max_prime);

    let sieve = Sieve {
        threads: threads,
        max_prime: max_prime
    };
    println!("{:?}", sieve);
    let result = sieve.find_primes();
    println!("Primes: {:?}", result.primes());
    println!("Found number of primes: {}", result.number_of_primes());
}

fn parse_command_line_argument<T: FromStr + Bounded + Display>(position: usize, name: &str) -> T {
    match env::args().nth(position) {
        None => err_exit!("Usage: sieve_multithreaded_shared <threads> <max_prime>"),
        Some(val) => match val.parse::<T>() {
            Err(_) => err_exit!("Expected <{}>: {} to be >= 0 and <= {}", name, val, T::max_value()),
            Ok(val) => val
        }
    }
}

fn calculate_actual_threads(requested_threads: usize, max_prime: usize) -> usize {
    let sqrt_prime = (max_prime as f64).sqrt().ceil() as usize;
    cmp::min(sqrt_prime, requested_threads)
}

#[derive(Debug)]
struct Sieve {
    threads: usize,
    max_prime: usize
}

impl Sieve {
    fn find_primes(&self) -> SieveResult {
        let mut bit_vector = BitVector::with_capacity(self.max_prime + 1, true);

        {
            let indices = self.calculate_indices();
            let bit_slices = self.split_into_bit_slices(&mut bit_vector, &indices);

            println!("{:?}", bit_slices);
        }

        SieveResult {
            threads: self.threads,
            max_prime: self.max_prime,
            bit_vector: bit_vector
        }
    }

    fn calculate_indices(&self) -> Vec<usize> {
        let storage_size = std::mem::size_of::<SieveStorage>() * 8;
        let numbers_per_segment = (self.max_prime + 1) / self.threads;

        let mut indices = vec![];
        let mut current_index = self.max_prime + 1;
        loop {
            if indices.len() == self.threads - 1 {
                break;
            }

            current_index = match current_index.checked_sub(numbers_per_segment) {
                None => break,
                Some(val) => val
            };
            indices.insert(0, current_index);
        }

        for index in &mut indices {
            *index = ((*index / storage_size) + 1) * storage_size;
        }

        indices
    }

    fn split_into_bit_slices<'a>(&self, bit_vector: &'a mut BitVector<SieveStorage>, indices: &[usize]) -> Vec<BitSliceMut<'a, SieveStorage>> {
        let mut bit_slices = vec![];

        //TODO request as_bit_slice and as_bit_slice_mut methods?
        bit_slices.push(bit_vector.split_at_mut(0).1);
        let mut split_indices = 0;
        for index in indices {
            let last_slice = bit_slices.pop().unwrap();
            let (new_slice, remainder) = last_slice.split_at_mut(index - split_indices);
            split_indices = *index;
            bit_slices.push(new_slice);
            bit_slices.push(remainder);
        }

        bit_slices
    }
}

#[derive(Debug)]
struct SieveResult {
    threads: usize,
    max_prime: usize,
    bit_vector: BitVector<SieveStorage>
}

impl SieveResult {
    fn number_of_primes(&self) -> usize {
        self.bit_vector.iter().filter(|x| *x).count()
    }

    fn primes(&self) -> Vec<usize> {
        self.bit_vector.iter().enumerate().filter(|x| x.1).map(|x| x.0).collect()
    }
}<commit_msg>Created the PrimeSlice struct to represent the pair of a BitSliceMut and a starting number such that indexes in the bit slice can be related to numbers of the whole bit vector.<commit_after>extern crate bit_vector;
extern crate num;

use std::cmp;
use std::env;
use std::fmt::Display;
use std::str::FromStr;

use bit_vector::{BitVector,BitSliceMut};

use num::Bounded;

macro_rules! err_exit {
    ($($arg:tt)*) => (
        {
            use std::io::prelude::*;
            use std::process;
            if let Err(e) = write!(&mut ::std::io::stderr(), "{}\n", format_args!($($arg)*)) {
                panic!("Failed to write to stderr.\
                    \nOriginal error output: {}\
                    \nSecondary error writing to stderr: {}", format!($($arg)*), e);
            }
            process::exit(1)
        }
    )
}

type SieveStorage = usize;

fn main() {
    let requested_threads: usize = parse_command_line_argument(1, "threads");
    let max_prime: usize = parse_command_line_argument(2, "max_prime");

    let threads = calculate_actual_threads(requested_threads, max_prime);

    let sieve = Sieve {
        threads: threads,
        max_prime: max_prime
    };
    println!("{:?}", sieve);
    let result = sieve.find_primes();
    println!("Primes: {:?}", result.primes());
    println!("Found number of primes: {}", result.number_of_primes());
}

fn parse_command_line_argument<T: FromStr + Bounded + Display>(position: usize, name: &str) -> T {
    match env::args().nth(position) {
        None => err_exit!("Usage: sieve_multithreaded_shared <threads> <max_prime>"),
        Some(val) => match val.parse::<T>() {
            Err(_) => err_exit!("Expected <{}>: {} to be >= 0 and <= {}", name, val, T::max_value()),
            Ok(val) => val
        }
    }
}

fn calculate_actual_threads(requested_threads: usize, max_prime: usize) -> usize {
    let sqrt_prime = (max_prime as f64).sqrt().ceil() as usize;
    cmp::min(sqrt_prime, requested_threads)
}

#[derive(Debug)]
struct Sieve {
    threads: usize,
    max_prime: usize
}

impl Sieve {
    fn find_primes(&self) -> SieveResult {
        let mut bit_vector = BitVector::with_capacity(self.max_prime + 1, true);

        {
            let indices = self.calculate_indices();
            let prime_slices = self.split_into_prime_slices(&mut bit_vector, &indices);

            println!("{:?}", prime_slices);
        }

        SieveResult {
            threads: self.threads,
            max_prime: self.max_prime,
            bit_vector: bit_vector
        }
    }

    fn calculate_indices(&self) -> Vec<usize> {
        let storage_size = std::mem::size_of::<SieveStorage>() * 8;
        let numbers_per_segment = (self.max_prime + 1) / self.threads;

        let mut indices = vec![];
        let mut current_index = self.max_prime + 1;
        loop {
            if indices.len() == self.threads - 1 {
                break;
            }

            current_index = match current_index.checked_sub(numbers_per_segment) {
                None => break,
                Some(val) => val
            };
            indices.insert(0, current_index);
        }

        for index in &mut indices {
            *index = ((*index / storage_size) + 1) * storage_size;
        }

        indices
    }

    fn split_into_prime_slices<'a>(&self, bit_vector: &'a mut BitVector<SieveStorage>, indices: &[usize]) -> Vec<PrimeSlice<'a>> {
        let mut bit_slices = vec![];

        //TODO request as_bit_slice and as_bit_slice_mut methods?
        bit_slices.push(bit_vector.split_at_mut(0).1);
        let mut split_indices = 0;
        for index in indices {
            let last_slice = bit_slices.pop().unwrap();
            let (new_slice, remainder) = last_slice.split_at_mut(index - split_indices);
            split_indices = *index;
            bit_slices.push(new_slice);
            bit_slices.push(remainder);
        }

        let mut prime_slices = vec![];
        let mut rolling_index = 0;
        for bit_slice in bit_slices {
            let bit_slice_capacity = bit_slice.capacity();
            prime_slices.push(PrimeSlice::new(bit_slice, rolling_index));
            rolling_index += bit_slice_capacity;
        }

        prime_slices
    }
}

#[derive(Debug)]
struct PrimeSlice<'a> {
    bit_slice: BitSliceMut<'a, SieveStorage>,
    start_number: usize
}

impl<'a> PrimeSlice<'a> {
    fn new(bit_slice: BitSliceMut<'a, SieveStorage>, start_number: usize) -> PrimeSlice<'a> {
        PrimeSlice {
            bit_slice: bit_slice,
            start_number: start_number
        }
    }
}

#[derive(Debug)]
struct SieveResult {
    threads: usize,
    max_prime: usize,
    bit_vector: BitVector<SieveStorage>
}

impl SieveResult {
    fn number_of_primes(&self) -> usize {
        self.bit_vector.iter().filter(|x| *x).count()
    }

    fn primes(&self) -> Vec<usize> {
        self.bit_vector.iter().enumerate().filter(|x| x.1).map(|x| x.0).collect()
    }
}<|endoftext|>


<commit_before>extern crate chrono;
#[macro_use]
extern crate clap;
#[macro_use]
extern crate lazy_static;
extern crate hyper;
extern crate mozprofile;
extern crate mozrunner;
extern crate regex;
extern crate rustc_serialize;
#[macro_use]
extern crate slog;
extern crate slog_atomic;
extern crate slog_stdlog;
extern crate slog_stream;
extern crate zip;
#[macro_use]
extern crate webdriver;

#[macro_use]
extern crate log;

use std::borrow::ToOwned;
use std::io::Write;
use std::net::{SocketAddr, IpAddr};
use std::path::PathBuf;
use std::str::FromStr;

use clap::{App, Arg};

macro_rules! try_opt {
    ($expr:expr, $err_type:expr, $err_msg:expr) => ({
        match $expr {
            Some(x) => x,
            None => return Err(WebDriverError::new($err_type, $err_msg))
        }
    })
}

mod logging;
mod prefs;
mod marionette;
mod capabilities;

use logging::LogLevel;
use marionette::{MarionetteHandler, MarionetteSettings, extension_routes};

type ProgramResult = std::result::Result<(), (ExitCode, String)>;

enum ExitCode {
    Ok = 0,
    Usage = 64,
}

fn app<'a, 'b>() -> App<'a, 'b> {
    App::new(format!("geckodriver {}", crate_version!()))
        .about("WebDriver implementation for Firefox.")
        .arg(Arg::with_name("webdriver_host")
            .long("host")
            .value_name("HOST")
            .help("Host ip to use for WebDriver server (default: 127.0.0.1)")
            .takes_value(true))
        .arg(Arg::with_name("webdriver_port")
            .short("p")
            .long("port")
            .value_name("PORT")
            .help("Port to use for WebDriver server (default: 4444)")
            .takes_value(true)
            .alias("webdriver-port"))
        .arg(Arg::with_name("binary")
            .short("b")
            .long("binary")
            .value_name("BINARY")
            .help("Path to the Firefox binary")
            .takes_value(true))
        .arg(Arg::with_name("marionette_port")
            .long("marionette-port")
            .value_name("PORT")
            .help("Port to use to connect to Gecko (default: random free port)")
            .takes_value(true))
        .arg(Arg::with_name("connect_existing")
            .long("connect-existing")
            .requires("marionette_port")
            .help("Connect to an existing Firefox instance"))
        .arg(Arg::with_name("verbosity")
            .short("v")
            .multiple(true)
            .conflicts_with("log_level")
            .help("Log level verbosity (-v for debug and -vv for trace level)"))
        .arg(Arg::with_name("log_level")
            .long("log")
            .takes_value(true)
            .value_name("LEVEL")
            .possible_values(&["fatal", "error", "warn", "info", "config", "debug", "trace"])
            .help("Set Gecko log level"))
        .arg(Arg::with_name("version")
            .short("V")
            .long("version")
            .help("Prints version and copying information"))
}

fn run() -> ProgramResult {
    let matches = app().get_matches();

    if matches.is_present("version") {
        println!("geckodriver {}\n\n{}", crate_version!(),
"The source code of this program is available at
https://github.com/mozilla/geckodriver.

This program is subject to the terms of the Mozilla Public License 2.0.
You can obtain a copy of the license at https://mozilla.org/MPL/2.0/.");
        return Ok(())
    }

    let host = matches.value_of("webdriver_host").unwrap_or("127.0.0.1");
    let port = match u16::from_str(matches.value_of("webdriver_port")
        .or(matches.value_of("webdriver_port_alias"))
        .unwrap_or("4444")) {
        Ok(x) => x,
        Err(_) => return Err((ExitCode::Usage, "invalid WebDriver port".to_owned())),
    };
    let addr = match IpAddr::from_str(host) {
        Ok(addr) => SocketAddr::new(addr, port),
        Err(_) => return Err((ExitCode::Usage, "invalid host address".to_owned())),
    };

    let binary = matches.value_of("binary").map(|x| PathBuf::from(x));

    let marionette_port = match matches.value_of("marionette_port") {
        Some(x) => match u16::from_str(x) {
            Ok(x) => Some(x),
            Err(_) => return Err((ExitCode::Usage, "invalid Marionette port".to_owned())),
        },
        None => None
    };

    // overrides defaults in Gecko
    // which are info for optimised builds
    // and debug for debug builds
    let log_level = if matches.is_present("log_level") {
        LogLevel::from_str(matches.value_of("log_level").unwrap()).ok()
    } else {
        match matches.occurrences_of("verbosity") {
            0 => Some(LogLevel::Info),
            1 => Some(LogLevel::Debug),
            _ => Some(LogLevel::Trace),
        }
    };
    logging::init(&log_level);

    let settings = MarionetteSettings {
        port: marionette_port,
        binary: binary,
        connect_existing: matches.is_present("connect_existing"),
        log_level: log_level,
    };

    let handler = MarionetteHandler::new(settings);
    let listening = try!(webdriver::server::start(addr, handler, &extension_routes())
        .map_err(|err| (ExitCode::Usage, err.to_string())));
    info!("Listening on {}", listening.socket);
    Ok(())
}

fn main() {
    let exit_code = match run() {
        Ok(_) => ExitCode::Ok,
        Err((exit_code, reason)) => {
            error!("{}", reason);
            exit_code
        },
    };

    std::io::stdout().flush().unwrap();
    std::process::exit(exit_code as i32);
}
<commit_msg>main: use exit code 69 to indicate that port is unavailable<commit_after>extern crate chrono;
#[macro_use]
extern crate clap;
#[macro_use]
extern crate lazy_static;
extern crate hyper;
extern crate mozprofile;
extern crate mozrunner;
extern crate regex;
extern crate rustc_serialize;
#[macro_use]
extern crate slog;
extern crate slog_atomic;
extern crate slog_stdlog;
extern crate slog_stream;
extern crate zip;
#[macro_use]
extern crate webdriver;

#[macro_use]
extern crate log;

use std::borrow::ToOwned;
use std::io::Write;
use std::net::{SocketAddr, IpAddr};
use std::path::PathBuf;
use std::str::FromStr;

use clap::{App, Arg};

macro_rules! try_opt {
    ($expr:expr, $err_type:expr, $err_msg:expr) => ({
        match $expr {
            Some(x) => x,
            None => return Err(WebDriverError::new($err_type, $err_msg))
        }
    })
}

mod logging;
mod prefs;
mod marionette;
mod capabilities;

use logging::LogLevel;
use marionette::{MarionetteHandler, MarionetteSettings, extension_routes};

type ProgramResult = std::result::Result<(), (ExitCode, String)>;

enum ExitCode {
    Ok = 0,
    Usage = 64,
    Unavailable = 69,
}

fn app<'a, 'b>() -> App<'a, 'b> {
    App::new(format!("geckodriver {}", crate_version!()))
        .about("WebDriver implementation for Firefox.")
        .arg(Arg::with_name("webdriver_host")
            .long("host")
            .value_name("HOST")
            .help("Host ip to use for WebDriver server (default: 127.0.0.1)")
            .takes_value(true))
        .arg(Arg::with_name("webdriver_port")
            .short("p")
            .long("port")
            .value_name("PORT")
            .help("Port to use for WebDriver server (default: 4444)")
            .takes_value(true)
            .alias("webdriver-port"))
        .arg(Arg::with_name("binary")
            .short("b")
            .long("binary")
            .value_name("BINARY")
            .help("Path to the Firefox binary")
            .takes_value(true))
        .arg(Arg::with_name("marionette_port")
            .long("marionette-port")
            .value_name("PORT")
            .help("Port to use to connect to Gecko (default: random free port)")
            .takes_value(true))
        .arg(Arg::with_name("connect_existing")
            .long("connect-existing")
            .requires("marionette_port")
            .help("Connect to an existing Firefox instance"))
        .arg(Arg::with_name("verbosity")
            .short("v")
            .multiple(true)
            .conflicts_with("log_level")
            .help("Log level verbosity (-v for debug and -vv for trace level)"))
        .arg(Arg::with_name("log_level")
            .long("log")
            .takes_value(true)
            .value_name("LEVEL")
            .possible_values(&["fatal", "error", "warn", "info", "config", "debug", "trace"])
            .help("Set Gecko log level"))
        .arg(Arg::with_name("version")
            .short("V")
            .long("version")
            .help("Prints version and copying information"))
}

fn run() -> ProgramResult {
    let matches = app().get_matches();

    if matches.is_present("version") {
        println!("geckodriver {}\n\n{}", crate_version!(),
"The source code of this program is available at
https://github.com/mozilla/geckodriver.

This program is subject to the terms of the Mozilla Public License 2.0.
You can obtain a copy of the license at https://mozilla.org/MPL/2.0/.");
        return Ok(())
    }

    let host = matches.value_of("webdriver_host").unwrap_or("127.0.0.1");
    let port = match u16::from_str(matches.value_of("webdriver_port")
        .or(matches.value_of("webdriver_port_alias"))
        .unwrap_or("4444")) {
        Ok(x) => x,
        Err(_) => return Err((ExitCode::Usage, "invalid WebDriver port".to_owned())),
    };
    let addr = match IpAddr::from_str(host) {
        Ok(addr) => SocketAddr::new(addr, port),
        Err(_) => return Err((ExitCode::Usage, "invalid host address".to_owned())),
    };

    let binary = matches.value_of("binary").map(|x| PathBuf::from(x));

    let marionette_port = match matches.value_of("marionette_port") {
        Some(x) => match u16::from_str(x) {
            Ok(x) => Some(x),
            Err(_) => return Err((ExitCode::Usage, "invalid Marionette port".to_owned())),
        },
        None => None
    };

    // overrides defaults in Gecko
    // which are info for optimised builds
    // and debug for debug builds
    let log_level = if matches.is_present("log_level") {
        LogLevel::from_str(matches.value_of("log_level").unwrap()).ok()
    } else {
        match matches.occurrences_of("verbosity") {
            0 => Some(LogLevel::Info),
            1 => Some(LogLevel::Debug),
            _ => Some(LogLevel::Trace),
        }
    };
    logging::init(&log_level);

    let settings = MarionetteSettings {
        port: marionette_port,
        binary: binary,
        connect_existing: matches.is_present("connect_existing"),
        log_level: log_level,
    };

    let handler = MarionetteHandler::new(settings);
    let listening = try!(webdriver::server::start(addr, handler, &extension_routes())
        .map_err(|err| (ExitCode::Unavailable, err.to_string())));
    info!("Listening on {}", listening.socket);
    Ok(())
}

fn main() {
    let exit_code = match run() {
        Ok(_) => ExitCode::Ok,
        Err((exit_code, reason)) => {
            error!("{}", reason);
            exit_code
        },
    };

    std::io::stdout().flush().unwrap();
    std::process::exit(exit_code as i32);
}
<|endoftext|>


<commit_before>// Based on smallpt, http://www.kevinbeason.com/smallpt/ which is also licensed under
// the MIT license.
use std::fs::File;
use std::io::prelude::*;

extern crate rand;
use rand::{Rng, XorShiftRng, SeedableRng};

extern crate threadpool;

mod path_tracer;
use path_tracer::*;

enum Material {
    Diffuse,
    Specular,
    Refractive
}

struct Sphere {
    material: Material,
    radius: f64,
    position: Vec3d,
    emission: Vec3d,
    colour: Vec3d
}

impl Sphere {
    fn new(material: Material, radius: f64, position: Vec3d, emission: Vec3d, colour: Vec3d) -> Sphere {
        Sphere {
          material: material,
          radius: radius,
          position: position,
          emission: emission,
          colour: colour
        }
    }
    fn intersect(&self, ray: &Ray) -> Option<f64> {
        let op = self.position - ray.origin;
        let b = op.dot(ray.direction);
        let determinant = b * b - op.dot(op) + self.radius * self.radius;
        if determinant < 0.0 { return None; }
        let determinant = determinant.sqrt();
        let t1 = b - determinant;
        let t2 = b + determinant;
        const EPSILON : f64 = 0.0001;
        if t1 > EPSILON {
            Some(t1)
        } else if t2 > EPSILON {
            Some(t2)
        } else {
            None
        }
    }
}

#[test]
fn intersection() {
    let sphere = Sphere::new(
        Material::Diffuse,
        100.0,
        Vec3d::new(0.0, 0.0, 200.0),
        Vec3d::zero(),
        Vec3d::zero());
    let ray = Ray::new(Vec3d::zero(), Vec3d::new(0.0, 0.0, 1.0));
    match sphere.intersect(&ray) {
        Some(x) => assert_eq!(x, 100.0),
        None => panic!("unexpected")
    }
    let ray = Ray::new(Vec3d::zero(), Vec3d::new(0.0, 1.0, 0.0));
    match sphere.intersect(&ray) {
        Some(_) => panic!("unexpected"),
        None => {}
    }
}

struct HitRecord<'a> {
    sphere: &'a Sphere,
    dist: f64
}

fn intersect<'a>(scene: &'a Vec<Sphere>, ray: &Ray) -> Option<HitRecord<'a>> {
    let mut result : Option<HitRecord<'a>> = None;
    for sph in scene.iter() {
        if let Some(dist) = sph.intersect(&ray) {
            if match result { None => true, Some(ref x) => dist < x.dist } {
                result = Some(HitRecord { sphere: &sph, dist: dist });
            }
        }
    }
    result
}

fn radiance<R: Rng>(scene: &Vec<Sphere>, ray: &Ray, depth: i32, rng: &mut R) -> Vec3d {
    if let Some(hit) = intersect(&scene, &ray) {
        let hit_pos = ray.origin + ray.direction * hit.dist;
        let hit_normal = (hit_pos - hit.sphere.position).normalized();
        let n1 = if hit_normal.dot(ray.direction) < 0.0 { hit_normal } else { hit_normal.neg() };
        let mut colour = hit.sphere.colour;
        let max_reflectance = colour.max_component();
        let depth = depth + 1;
        if depth > 5 {
            let rand = rng.gen::<f64>();
            if rand < max_reflectance && depth < 100 {
               colour = colour * (1.0 / max_reflectance);
            } else {
                return hit.sphere.emission;
            }
        }
        match hit.sphere.material {
            Material::Diffuse => {
                // Get a random polar coordinate
                let r1 = rng.gen::<f64>() * 2.0 * std::f64::consts::PI;
                let r2 = rng.gen::<f64>();
                let r2s = r2.sqrt();
                // Create a coordinate system u,v,w local to the point, where the w is the normal
                // pointing out of the sphere and the u and v are orthonormal to w.
                let w = n1;
                // Pick an arbitrary non-zero preferred axis for u
                let u = if n1.x.abs() > 0.1 { Vec3d::new(0.0, 1.0, 0.0) } else { Vec3d::new(1.0, 0.0, 0.0) }.cross(w);
                let v = w.cross(u);
                // construct the new direction
                let new_dir = u * r1.cos() * r2s + v * r1.sin() * r2s + w * (1.0 - r2).sqrt();
                colour = colour * radiance(scene, &Ray::new(hit_pos, new_dir.normalized()), depth, rng);
            },
            Material::Specular => {
                let reflection = ray.direction - hit_normal * 2.0 * hit_normal.dot(ray.direction);
                let reflected_ray = Ray::new(hit_pos, reflection);
                colour = colour * radiance(scene, &reflected_ray, depth, rng);
            },
            Material::Refractive => {
                let reflection = ray.direction - hit_normal * 2.0 * hit_normal.dot(ray.direction);
                let reflected_ray = Ray::new(hit_pos, reflection);
                let into = if hit_normal.dot(n1) > 0.0 { true } else { false };
                let nc = 1.0;
                let nt = 1.5;
                let nnt = if into { nc/nt } else { nt/nc };
                let ddn = ray.direction.dot(n1);
                let cos2t = 1.0 - nnt * nnt * (1.0 - ddn * ddn);
                if cos2t < 0.0 {
                    // Total internal reflection
                    colour = colour * radiance(scene, &reflected_ray, depth, rng);
                } else {
                    let tbd = ddn * nnt + cos2t.sqrt();
                    let tbd = if into { tbd } else { -tbd };
                    let tdir = (ray.direction * nnt - hit_normal * tbd).normalized();
                    let transmitted_ray = Ray::new(hit_pos, tdir);
                    let a = nt - nc;
                    let b = nt + nc;
                    let r0 = (a * a) / (b * b);
                    let c = 1.0 - if into { -ddn } else { tdir.dot(hit_normal) };
                    let re = r0 + (1.0 - r0) * c * c * c * c * c;
                    let tr = 1.0 - re;
                    let p = 0.25 + 0.5 * re;
                    let rp = re / p;
                    let tp = tr / (1.0 - p);
                    colour = colour * if depth > 2 {
                        if rng.gen::<f64>() < p {
                            radiance(scene, &reflected_ray, depth, rng) * rp
                        } else {
                            radiance(scene, &transmitted_ray, depth, rng) * tp
                        }
                    } else {
                        radiance(scene, &reflected_ray, depth, rng) * re +
                            radiance(scene, &transmitted_ray, depth, rng) * tr
                    }
                }
            }
        }
        hit.sphere.emission + colour
    } else {
        Vec3d::zero()
    }
}

fn random_samp<T: Rng>(rng: &mut T) -> f64 {
    let r = 2.0 * rng.gen::<f64>();
    if r < 1.0 { r.sqrt() - 1.0 } else { 1.0 - (2.0 - r).sqrt() }
}

fn to_int(v: f64) -> u8 {
    let ch = (v.powf(1.0/2.2) * 255.0 + 0.5) as i64;
    if ch < 0 { 0 } else if ch > 255 { 255 } else { ch as u8 }
}

fn main() {
    const BLACK : Vec3d = Vec3d { x: 0.0, y: 0.0, z: 0.0 };
    const RED : Vec3d = Vec3d { x: 0.75, y: 0.25, z: 0.25 };
    const BLUE : Vec3d = Vec3d { x: 0.25, y: 0.25, z: 0.75 };
    const GREY : Vec3d = Vec3d { x: 0.75, y: 0.75, z: 0.75 };
    const WHITE : Vec3d = Vec3d { x: 0.999, y: 0.999, z: 0.999 };
    let scene = vec!{
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(1e5+1.0, 40.8, 81.6),
            BLACK, RED),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(-1e5+99.0, 40.8, 81.6),
            BLACK, BLUE),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 40.8, 1e5),
            BLACK, GREY),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 40.8, -1e5 + 170.0),
            BLACK, BLACK),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 1e5, 81.6),
            BLACK, GREY),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, -1e5 + 81.6, 81.6),
            BLACK, RED),
        Sphere::new(Material::Specular, 16.5,
            Vec3d::new(27.0, 16.5, 47.0),
            BLACK, WHITE),
        Sphere::new(Material::Refractive, 16.5,
            Vec3d::new(73.0, 16.5, 78.0),
            BLACK, WHITE),
        Sphere::new(Material::Diffuse, 600.0,
            Vec3d::new(50.0, 681.6 - 0.27, 81.6),
            Vec3d::new(12.0, 12.0, 12.0), BLACK),
    };

    const WIDTH: usize = 256;//1024;
    const HEIGHT: usize = 192;//768;
    let samps = 16;
    let camera_pos = Vec3d::new(50.0, 52.0, 295.6);
    let camera_dir = Vec3d::new(0.0, -0.042612, -1.0);
    let camera_x = Vec3d::new(WIDTH as f64 * 0.5135 / HEIGHT as f64, 0.0, 0.0);
    let camera_y = camera_x.cross(camera_dir).normalized() * 0.5135;

    let mut screen: Vec<Vec<Vec3d>> = Vec::with_capacity(HEIGHT);
    for _y in 0..HEIGHT {
        screen.push(Vec::with_capacity(WIDTH));
    }

    for y in 0..HEIGHT {
        println!("Rendering ({} spp) {:.4}%...", samps * 4, 100.0 * y as f64 / HEIGHT as f64);
        for x in 0..WIDTH {
            let mut rng = XorShiftRng::from_seed([10, 200, 999999, y as u32]);
            let mut sum = Vec3d::zero();
            for sx in 0..2 {
                for sy in 0..2 {
                    for _samp in 0..samps {
                        let dx = random_samp(&mut rng);
                        let dy = random_samp(&mut rng);
                        let dir = camera_x * (((sx as f64 + 0.5 + dx)/2.0 + x as f64) / WIDTH as f64 - 0.5) +
                            camera_y * (((sy as f64 + 0.5 + dy)/2.0 + (HEIGHT - y - 1) as f64) / HEIGHT as f64  - 0.5) + camera_dir;
                        let jittered_ray = Ray::new(camera_pos + dir * 140.0, dir.normalized());
                        let sample = radiance(&scene, &jittered_ray, 0, &mut rng);
                        sum = sum + sample
                    }
                }
            }
            screen[y].push(sum / (samps * 4) as f64);
        }
    }
    let mut output_file = File::create("image.ppm").unwrap();
    write!(&mut output_file, "P3\n{} {}\n255\n", WIDTH, HEIGHT).unwrap();
    for y in 0..HEIGHT {
        for x in 0..WIDTH {
            let sum = screen[y][x];
            write!(&mut output_file, "{} {} {} ", to_int(sum.x), to_int(sum.y), to_int(sum.z)).unwrap();
        }
    }
}
<commit_msg>Add threading<commit_after>// Based on smallpt, http://www.kevinbeason.com/smallpt/ which is also licensed under
// the MIT license.
use std::fs::File;
use std::io::prelude::*;

extern crate rand;
use rand::{Rng, XorShiftRng, SeedableRng};

extern crate threadpool;
use threadpool::ThreadPool;
use std::sync::mpsc::channel;
use std::sync::Arc;

mod path_tracer;
use path_tracer::*;

enum Material {
    Diffuse,
    Specular,
    Refractive
}

struct Sphere {
    material: Material,
    radius: f64,
    position: Vec3d,
    emission: Vec3d,
    colour: Vec3d
}

impl Sphere {
    fn new(material: Material, radius: f64, position: Vec3d, emission: Vec3d, colour: Vec3d) -> Sphere {
        Sphere {
          material: material,
          radius: radius,
          position: position,
          emission: emission,
          colour: colour
        }
    }
    fn intersect(&self, ray: &Ray) -> Option<f64> {
        let op = self.position - ray.origin;
        let b = op.dot(ray.direction);
        let determinant = b * b - op.dot(op) + self.radius * self.radius;
        if determinant < 0.0 { return None; }
        let determinant = determinant.sqrt();
        let t1 = b - determinant;
        let t2 = b + determinant;
        const EPSILON : f64 = 0.0001;
        if t1 > EPSILON {
            Some(t1)
        } else if t2 > EPSILON {
            Some(t2)
        } else {
            None
        }
    }
}

#[test]
fn intersection() {
    let sphere = Sphere::new(
        Material::Diffuse,
        100.0,
        Vec3d::new(0.0, 0.0, 200.0),
        Vec3d::zero(),
        Vec3d::zero());
    let ray = Ray::new(Vec3d::zero(), Vec3d::new(0.0, 0.0, 1.0));
    match sphere.intersect(&ray) {
        Some(x) => assert_eq!(x, 100.0),
        None => panic!("unexpected")
    }
    let ray = Ray::new(Vec3d::zero(), Vec3d::new(0.0, 1.0, 0.0));
    match sphere.intersect(&ray) {
        Some(_) => panic!("unexpected"),
        None => {}
    }
}

struct HitRecord<'a> {
    sphere: &'a Sphere,
    dist: f64
}

fn intersect<'a>(scene: &'a Vec<Sphere>, ray: &Ray) -> Option<HitRecord<'a>> {
    let mut result : Option<HitRecord<'a>> = None;
    for sph in scene.iter() {
        if let Some(dist) = sph.intersect(&ray) {
            if match result { None => true, Some(ref x) => dist < x.dist } {
                result = Some(HitRecord { sphere: &sph, dist: dist });
            }
        }
    }
    result
}

fn radiance<R: Rng>(scene: &Vec<Sphere>, ray: &Ray, depth: i32, rng: &mut R) -> Vec3d {
    if let Some(hit) = intersect(&scene, &ray) {
        let hit_pos = ray.origin + ray.direction * hit.dist;
        let hit_normal = (hit_pos - hit.sphere.position).normalized();
        let n1 = if hit_normal.dot(ray.direction) < 0.0 { hit_normal } else { hit_normal.neg() };
        let mut colour = hit.sphere.colour;
        let max_reflectance = colour.max_component();
        let depth = depth + 1;
        if depth > 5 {
            let rand = rng.gen::<f64>();
            if rand < max_reflectance && depth < 100 {
               colour = colour * (1.0 / max_reflectance);
            } else {
                return hit.sphere.emission;
            }
        }
        match hit.sphere.material {
            Material::Diffuse => {
                // Get a random polar coordinate
                let r1 = rng.gen::<f64>() * 2.0 * std::f64::consts::PI;
                let r2 = rng.gen::<f64>();
                let r2s = r2.sqrt();
                // Create a coordinate system u,v,w local to the point, where the w is the normal
                // pointing out of the sphere and the u and v are orthonormal to w.
                let w = n1;
                // Pick an arbitrary non-zero preferred axis for u
                let u = if n1.x.abs() > 0.1 { Vec3d::new(0.0, 1.0, 0.0) } else { Vec3d::new(1.0, 0.0, 0.0) }.cross(w);
                let v = w.cross(u);
                // construct the new direction
                let new_dir = u * r1.cos() * r2s + v * r1.sin() * r2s + w * (1.0 - r2).sqrt();
                colour = colour * radiance(scene, &Ray::new(hit_pos, new_dir.normalized()), depth, rng);
            },
            Material::Specular => {
                let reflection = ray.direction - hit_normal * 2.0 * hit_normal.dot(ray.direction);
                let reflected_ray = Ray::new(hit_pos, reflection);
                colour = colour * radiance(scene, &reflected_ray, depth, rng);
            },
            Material::Refractive => {
                let reflection = ray.direction - hit_normal * 2.0 * hit_normal.dot(ray.direction);
                let reflected_ray = Ray::new(hit_pos, reflection);
                let into = if hit_normal.dot(n1) > 0.0 { true } else { false };
                let nc = 1.0;
                let nt = 1.5;
                let nnt = if into { nc/nt } else { nt/nc };
                let ddn = ray.direction.dot(n1);
                let cos2t = 1.0 - nnt * nnt * (1.0 - ddn * ddn);
                if cos2t < 0.0 {
                    // Total internal reflection
                    colour = colour * radiance(scene, &reflected_ray, depth, rng);
                } else {
                    let tbd = ddn * nnt + cos2t.sqrt();
                    let tbd = if into { tbd } else { -tbd };
                    let tdir = (ray.direction * nnt - hit_normal * tbd).normalized();
                    let transmitted_ray = Ray::new(hit_pos, tdir);
                    let a = nt - nc;
                    let b = nt + nc;
                    let r0 = (a * a) / (b * b);
                    let c = 1.0 - if into { -ddn } else { tdir.dot(hit_normal) };
                    let re = r0 + (1.0 - r0) * c * c * c * c * c;
                    let tr = 1.0 - re;
                    let p = 0.25 + 0.5 * re;
                    let rp = re / p;
                    let tp = tr / (1.0 - p);
                    colour = colour * if depth > 2 {
                        if rng.gen::<f64>() < p {
                            radiance(scene, &reflected_ray, depth, rng) * rp
                        } else {
                            radiance(scene, &transmitted_ray, depth, rng) * tp
                        }
                    } else {
                        radiance(scene, &reflected_ray, depth, rng) * re +
                            radiance(scene, &transmitted_ray, depth, rng) * tr
                    }
                }
            }
        }
        hit.sphere.emission + colour
    } else {
        Vec3d::zero()
    }
}

fn random_samp<T: Rng>(rng: &mut T) -> f64 {
    let r = 2.0 * rng.gen::<f64>();
    if r < 1.0 { r.sqrt() - 1.0 } else { 1.0 - (2.0 - r).sqrt() }
}

fn to_int(v: f64) -> u8 {
    let ch = (v.powf(1.0/2.2) * 255.0 + 0.5) as i64;
    if ch < 0 { 0 } else if ch > 255 { 255 } else { ch as u8 }
}

fn main() {
    const BLACK : Vec3d = Vec3d { x: 0.0, y: 0.0, z: 0.0 };
    const RED : Vec3d = Vec3d { x: 0.75, y: 0.25, z: 0.25 };
    const BLUE : Vec3d = Vec3d { x: 0.25, y: 0.25, z: 0.75 };
    const GREY : Vec3d = Vec3d { x: 0.75, y: 0.75, z: 0.75 };
    const WHITE : Vec3d = Vec3d { x: 0.999, y: 0.999, z: 0.999 };
    let scene = Arc::new(vec!{
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(1e5+1.0, 40.8, 81.6),
            BLACK, RED),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(-1e5+99.0, 40.8, 81.6),
            BLACK, BLUE),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 40.8, 1e5),
            BLACK, GREY),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 40.8, -1e5 + 170.0),
            BLACK, BLACK),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, 1e5, 81.6),
            BLACK, GREY),
        Sphere::new(Material::Diffuse, 1e5,
            Vec3d::new(50.0, -1e5 + 81.6, 81.6),
            BLACK, RED),
        Sphere::new(Material::Specular, 16.5,
            Vec3d::new(27.0, 16.5, 47.0),
            BLACK, WHITE),
        Sphere::new(Material::Refractive, 16.5,
            Vec3d::new(73.0, 16.5, 78.0),
            BLACK, WHITE),
        Sphere::new(Material::Diffuse, 600.0,
            Vec3d::new(50.0, 681.6 - 0.27, 81.6),
            Vec3d::new(12.0, 12.0, 12.0), BLACK),
    });

    const WIDTH: usize = 256;//1024;
    const HEIGHT: usize = 192;//768;
    let samps = 1024;
    let camera_pos = Vec3d::new(50.0, 52.0, 295.6);
    let camera_dir = Vec3d::new(0.0, -0.042612, -1.0);
    let camera_x = Vec3d::new(WIDTH as f64 * 0.5135 / HEIGHT as f64, 0.0, 0.0);
    let camera_y = camera_x.cross(camera_dir).normalized() * 0.5135;

    let pool = ThreadPool::new(4);
    let (tx, rx) = channel();

    for y in 0..HEIGHT {
        let tx = tx.clone();
        let scene = scene.clone();
        pool.execute(move || {
            let mut line = Vec::with_capacity(WIDTH);
            for x in 0..WIDTH {
                let mut rng = XorShiftRng::from_seed([10, 200, 999999, y as u32]);
                let mut sum = Vec3d::zero();
                for sx in 0..2 {
                    for sy in 0..2 {
                        for _samp in 0..samps {
                            let dx = random_samp(&mut rng);
                            let dy = random_samp(&mut rng);
                            let dir = camera_x * (((sx as f64 + 0.5 + dx)/2.0 + x as f64) / WIDTH as f64 - 0.5) +
            camera_y * (((sy as f64 + 0.5 + dy)/2.0 + (HEIGHT - y - 1) as f64) / HEIGHT as f64  - 0.5) + camera_dir;
        let jittered_ray = Ray::new(camera_pos + dir * 140.0, dir.normalized());
        let sample = radiance(&scene, &jittered_ray, 0, &mut rng);
        sum = sum + sample;
                        }
                    }
                }
                line.push(sum / (samps * 4) as f64);
            }
            tx.send((y, line)).unwrap();
        });
    }
    let mut left = HEIGHT;
    let mut screen : Vec<Vec<Vec3d>> = Vec::new();
    for _y in 0..HEIGHT {
        screen.push(Vec::new());
    }
    while left > 0 {
        println!("Rendering ({} spp) {:.4}%...", samps * 4, 100.0 * (HEIGHT - left) as f64 / HEIGHT as f64);
        let (y, line) = rx.recv().unwrap();
        screen[y] = line;
        left -= 1;
    }
    println!("Writing output");
    let mut output_file = File::create("image.ppm").unwrap();
    write!(&mut output_file, "P3\n{} {}\n255\n", WIDTH, HEIGHT).unwrap();
    for y in 0..HEIGHT {
        for x in 0..WIDTH {
            let sum = screen[y][x];
            write!(&mut output_file, "{} {} {} ", to_int(sum.x), to_int(sum.y), to_int(sum.z)).unwrap();
        }
    }
}
<|endoftext|>