<commit_before>extern crate base64;

extern crate byteorder;

extern crate hyper;
extern crate hyper_native_tls;

extern crate rayon;

#[macro_use]
extern crate serde_derive;
extern crate serde_json;

extern crate pem;

#[macro_use]
extern crate prettytable;

use std::{env, process};
use std::fs::File;
use std::io::{Read, Write};

use byteorder::{BigEndian, WriteBytesExt};

use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

#[derive(Deserialize)]
struct LogsResponseLogs {
    description: String,
    url: String,
    operated_by: Vec<u32>,
    disqualified_at: Option<u64>,
}

#[derive(Deserialize)]
struct LogsResponseOperators {
    name: String,
    id: u32,
}

#[derive(Deserialize)]
struct LogsResponse {
    logs: Vec<LogsResponseLogs>,
    operators: Vec<LogsResponseOperators>,
}

#[derive(Debug, Clone)]
struct Log {
    description: String,
    url: String,
    is_google: bool,
}

const LOG_LIST_URL: &'static str = "https://www.gstatic.com/ct/log_list/log_list.json";

fn fetch_trusted_ct_logs(http_client: &hyper::Client) -> Vec<Log> {
    let response = http_client.get(LOG_LIST_URL).send().unwrap();
    let logs_response: LogsResponse = serde_json::from_reader(response).unwrap();

    let google_id = logs_response
        .operators
        .iter()
        .find(|o| o.name == "Google")
        .map(|o| o.id);

    return logs_response
               .logs
               .into_iter()
               .filter(|log| log.disqualified_at.is_none())
               .map(|log| {
                        Log {
                            url: log.url,
                            description: log.description,
                            is_google: log.operated_by.contains(&google_id.unwrap()),
                        }
                    })
               .collect();
}

#[derive(Debug, Deserialize)]
struct SignedCertificateTimestamp {
    sct_version: u8,
    id: String,
    timestamp: u64,
    extensions: String,
    signature: String,
}

impl SignedCertificateTimestamp {
    fn to_raw_bytes(&self) -> Vec<u8> {
        let mut b = Vec::new();
        b.write_u8(self.sct_version).unwrap();

        let log_id = base64::decode(&self.id).unwrap();
        b.write(&log_id).unwrap();

        b.write_u64::<BigEndian>(self.timestamp).unwrap();

        let extensions = base64::decode(&self.extensions).unwrap();
        assert!(extensions.len() <= 65535);
        b.write_u16::<BigEndian>(extensions.len() as u16)
            .unwrap();
        b.write(&extensions).unwrap();

        let signature = base64::decode(&self.signature).unwrap();
        b.write(&signature).unwrap();

        return b;
    }
}

fn submit_to_log(http_client: &hyper::Client,
                 url: &str,
                 payload: &[u8])
                 -> Option<SignedCertificateTimestamp> {
    let mut url = "https://".to_string() + url;
    if !url.ends_with("/") {
        url += "/";
    }
    url += "ct/v1/add-chain";
    let response = http_client
        .post(&url)
        .body(hyper::client::Body::BufBody(payload, payload.len()))
        .header(hyper::header::ContentType::json())
        .send();
    let response = match response {
        Ok(r) => r,
        // TODO: maybe not all of these should be silently ignored.
        Err(_) => return None,
    };

    // 400, 403, and probably some others generally indicate a log doesn't accept certs from this
    // root, or that the log isn't accepting new submissions.
    if response.status.is_client_error() {
        return None;
    }

    return Some(serde_json::from_reader(response).unwrap());
}

#[derive(Serialize)]
struct AddChainRequest {
    chain: Vec<String>,
}

fn submit_cert_to_logs<'a>(http_client: &hyper::Client,
                           logs: &'a [Log],
                           cert: Vec<Vec<u8>>)
                           -> Vec<(&'a Log, SignedCertificateTimestamp)> {
    let payload = serde_json::to_vec(&AddChainRequest {
                                          chain: cert.iter().map(|r| base64::encode(r)).collect(),
                                      })
            .unwrap();

    return logs.par_iter()
               .filter_map(|ref log| {
                               let sct = submit_to_log(http_client, &log.url, &payload);
                               return sct.map(|s| (log.clone(), s));
                           })
               .collect();
}

fn main() {
    let path = match env::args().skip(1).next() {
        Some(p) => p,
        None => {
            println!("Usage: {} <path/to/chain.pem>", env::args().next().unwrap());
            process::exit(1);
        }
    };

    let mut contents = String::new();
    File::open(path)
        .unwrap()
        .read_to_string(&mut contents)
        .unwrap();

    let chain = pem::parse_many(contents)
        .into_iter()
        .filter(|p| p.tag == "CERTIFICATE")
        .map(|p| p.contents)
        .collect();

    let http_client = hyper::Client::with_connector(
        hyper::net::HttpsConnector::new(hyper_native_tls::NativeTlsClient::new().unwrap())
    );
    let logs = fetch_trusted_ct_logs(&http_client);

    let scts = submit_cert_to_logs(&http_client, &logs, chain);

    let mut table = prettytable::Table::new();
    table.add_row(row!["Log", "SCT"]);
    for (log, sct) in scts {
        table.add_row(row![log.description, base64::encode(&sct.to_raw_bytes())]);
    }
    table.printstd();
}
<commit_msg>Limit the network reads for safety<commit_after>extern crate base64;

extern crate byteorder;

extern crate hyper;
extern crate hyper_native_tls;

extern crate rayon;

#[macro_use]
extern crate serde_derive;
extern crate serde_json;

extern crate pem;

#[macro_use]
extern crate prettytable;

use std::{env, process};
use std::fs::File;
use std::io::{Read, Write};

use byteorder::{BigEndian, WriteBytesExt};

use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

#[derive(Deserialize)]
struct LogsResponseLogs {
    description: String,
    url: String,
    operated_by: Vec<u32>,
    disqualified_at: Option<u64>,
}

#[derive(Deserialize)]
struct LogsResponseOperators {
    name: String,
    id: u32,
}

#[derive(Deserialize)]
struct LogsResponse {
    logs: Vec<LogsResponseLogs>,
    operators: Vec<LogsResponseOperators>,
}

#[derive(Debug, Clone)]
struct Log {
    description: String,
    url: String,
    is_google: bool,
}

const LOG_LIST_URL: &'static str = "https://www.gstatic.com/ct/log_list/log_list.json";

fn fetch_trusted_ct_logs(http_client: &hyper::Client) -> Vec<Log> {
    let response = http_client.get(LOG_LIST_URL).send().unwrap();
    // Limit the response to 10MB at most, to be resillient to DoS.
    let logs_response: LogsResponse = serde_json::from_reader(response.take(10 * 1024 * 1024))
        .unwrap();

    let google_id = logs_response
        .operators
        .iter()
        .find(|o| o.name == "Google")
        .map(|o| o.id);

    return logs_response
               .logs
               .into_iter()
               .filter(|log| log.disqualified_at.is_none())
               .map(|log| {
                        Log {
                            url: log.url,
                            description: log.description,
                            is_google: log.operated_by.contains(&google_id.unwrap()),
                        }
                    })
               .collect();
}

#[derive(Debug, Deserialize)]
struct SignedCertificateTimestamp {
    sct_version: u8,
    id: String,
    timestamp: u64,
    extensions: String,
    signature: String,
}

impl SignedCertificateTimestamp {
    fn to_raw_bytes(&self) -> Vec<u8> {
        let mut b = Vec::new();
        b.write_u8(self.sct_version).unwrap();

        let log_id = base64::decode(&self.id).unwrap();
        b.write(&log_id).unwrap();

        b.write_u64::<BigEndian>(self.timestamp).unwrap();

        let extensions = base64::decode(&self.extensions).unwrap();
        assert!(extensions.len() <= 65535);
        b.write_u16::<BigEndian>(extensions.len() as u16)
            .unwrap();
        b.write(&extensions).unwrap();

        let signature = base64::decode(&self.signature).unwrap();
        b.write(&signature).unwrap();

        return b;
    }
}

fn submit_to_log(http_client: &hyper::Client,
                 url: &str,
                 payload: &[u8])
                 -> Option<SignedCertificateTimestamp> {
    let mut url = "https://".to_string() + url;
    if !url.ends_with("/") {
        url += "/";
    }
    url += "ct/v1/add-chain";
    let response = http_client
        .post(&url)
        .body(hyper::client::Body::BufBody(payload, payload.len()))
        .header(hyper::header::ContentType::json())
        .send();
    let response = match response {
        Ok(r) => r,
        // TODO: maybe not all of these should be silently ignored.
        Err(_) => return None,
    };

    // 400, 403, and probably some others generally indicate a log doesn't accept certs from this
    // root, or that the log isn't accepting new submissions.
    if response.status.is_client_error() {
        return None;
    }

    // Limt the response to 10MB (well above what would ever be needed) to be resilient to DoS in
    // the face of a dumb or malicious log.
    return Some(serde_json::from_reader(response.take(10 * 1024 * 1024)).unwrap());
}

#[derive(Serialize)]
struct AddChainRequest {
    chain: Vec<String>,
}

fn submit_cert_to_logs<'a>(http_client: &hyper::Client,
                           logs: &'a [Log],
                           cert: Vec<Vec<u8>>)
                           -> Vec<(&'a Log, SignedCertificateTimestamp)> {
    let payload = serde_json::to_vec(&AddChainRequest {
                                          chain: cert.iter().map(|r| base64::encode(r)).collect(),
                                      })
            .unwrap();

    return logs.par_iter()
               .filter_map(|ref log| {
                               let sct = submit_to_log(http_client, &log.url, &payload);
                               return sct.map(|s| (log.clone(), s));
                           })
               .collect();
}

fn main() {
    let path = match env::args().skip(1).next() {
        Some(p) => p,
        None => {
            println!("Usage: {} <path/to/chain.pem>", env::args().next().unwrap());
            process::exit(1);
        }
    };

    let mut contents = String::new();
    File::open(path)
        .unwrap()
        .read_to_string(&mut contents)
        .unwrap();

    let chain = pem::parse_many(contents)
        .into_iter()
        .filter(|p| p.tag == "CERTIFICATE")
        .map(|p| p.contents)
        .collect();

    let http_client = hyper::Client::with_connector(
        hyper::net::HttpsConnector::new(hyper_native_tls::NativeTlsClient::new().unwrap())
    );
    let logs = fetch_trusted_ct_logs(&http_client);

    let scts = submit_cert_to_logs(&http_client, &logs, chain);

    let mut table = prettytable::Table::new();
    table.add_row(row!["Log", "SCT"]);
    for (log, sct) in scts {
        table.add_row(row![log.description, base64::encode(&sct.to_raw_bytes())]);
    }
    table.printstd();
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate clap;
extern crate toml;
extern crate shogi;

mod config;
mod environment;
mod error;
mod game;
mod stats;
mod usi;

use std::io::Write;
use std::sync::atomic::Ordering;
use std::sync::mpsc::{channel, Receiver};
use std::thread;
use clap::{App, Arg};
use shogi::Color;

use config::*;
use environment::*;
use error::*;
use stats::*;
use usi::*;

fn main() {
    let matches = App::new("usirun")
        .version(crate_version!())
        .about("A command line utility for running games between USI compliant Shogi engines.")
        .arg(Arg::with_name("config")
            .short("c")
            .long("config")
            .value_name("TOML")
            .help("Loads a configuration file for setting up match rules")
            .required(true)
            .takes_value(true))
        .arg(Arg::with_name("display")
            .short("d")
            .long("display")
            .value_name("MODE")
            .help("Displays ")
            .takes_value(true)
            .possible_values(&["board", "command", "simple"]))
        .get_matches();

    let mut match_config = MatchConfig::default();
    if let Some(config_path) = matches.value_of("config") {
        match_config.load(config_path)
            .expect(&format!("failed to open the config file at {}", config_path));
    }

    if let Some(display) = matches.value_of("display") {
        match_config.display = match display {
            "board" => DisplayMode::Board,
            "command" => DisplayMode::Command,
            _ => DisplayMode::Simple,
        }
    }

    match run_match(&match_config) {
        Ok(ref stats) => {
            println!("");
            println!("Match statistics");
            println!("Black: {} ({:.2}%)",
                     stats.black_wins(),
                     stats.black_win_rate() * 100.0);
            println!("White: {} ({:.2}%)",
                     stats.white_wins(),
                     stats.white_win_rate() * 100.0);
            println!("Draw : {} ({:.2}%)",
                     stats.draw_games(),
                     stats.draw_game_rate() * 100.0);

        }
        Err(e) => {
            println!("an error occurred during the match: {}", e);
        }
    }
}

fn run_match(config: &MatchConfig) -> Result<MatchStatistics, Error> {
    let mut env = Environment::new(&config.time.to_time_control())
        .initial_sfen(&config.initial_pos)
        .max_ply(config.max_ply);

    let (black_tx, black_rx) = channel();
    let (white_tx, white_rx) = channel();
    let (monitor_tx, monitor_rx) = channel();

    let mut black_engine = try!(UsiEngine::launch(Color::Black, &config.black_engine));
    let mut white_engine = try!(UsiEngine::launch(Color::White, &config.white_engine));

    if config.display == DisplayMode::Command {
        set_command_logger(Color::Black, &mut black_engine);
        set_command_logger(Color::White, &mut white_engine);
    }

    black_engine.listen(env.new_sender(), black_rx);
    white_engine.listen(env.new_sender(), white_rx);

    let monitor_handle = start_monitor_thread(&config, &black_engine, &white_engine, monitor_rx);

    for _ in 0..config.num_games {
        try!(env.start_game(&[&black_tx, &white_tx, &monitor_tx]));

        if config.display == DisplayMode::Simple {
            print!(".");
            try!(std::io::stdout().flush());
        }
    }

    try!(black_engine.kill());
    try!(white_engine.kill());

    let stats = monitor_handle.join().expect("unexpected error occurred in the monitoring thread");

    Ok(stats)
}

fn start_monitor_thread(config: &MatchConfig,
                        black_engine: &UsiEngine,
                        white_engine: &UsiEngine,
                        rx: Receiver<Event>)
                        -> thread::JoinHandle<MatchStatistics> {
    let num_games = config.num_games;
    let display = config.display;
    let black_score = black_engine.score.clone();
    let white_score = white_engine.score.clone();

    thread::spawn(move || {
        let mut results = MatchStatistics::default();

        while let Some(event) = rx.recv().ok() {
            match event {
                Event::NewTurn(game) => {
                    if display == DisplayMode::Board {
                        if let Some(game) = game.read().ok() {
                            println!("{}", game.pos);
                            println!("Time: (Black) {}s, (White) {}s",
                                     game.time.black_time().as_secs(),
                                     game.time.white_time().as_secs());
                            println!("Score (Black) {}, (White) {}",
                                     black_score.load(Ordering::Relaxed),
                                     white_score.load(Ordering::Relaxed));
                            println!("");
                        }
                    }
                }
                Event::GameOver(winner, reason) => {
                    if display == DisplayMode::Board {
                        match winner {
                            Some(c) => {
                                let name = if c == Color::Black { "Black" } else { "White" };
                                println!("Game #{}: {} wins({:?})",
                                         results.total_games() + 1,
                                         name,
                                         reason);
                            }
                            None => {
                                println!("Game #{}: Draw({:?})", results.total_games() + 1, reason);
                            }
                        }
                    }

                    results.record_game(winner);
                    if num_games == results.total_games() {
                        break;
                    }
                }
                _ => {}
            }
        }

        results
    })
}

fn set_command_logger(color: Color, engine: &mut UsiEngine) {
    let prefix = if color == Color::Black { "B" } else { "W" };

    engine.set_read_hook(Some(Box::new(move |output| {
        print!("{}> {}", prefix, output.raw_str());
    })));

    engine.set_write_hook(Some(Box::new(move |_, raw_str| {
        print!("{}< {}", prefix, raw_str);
    })));
}<commit_msg>refactor: use shogi-rs:0.6.0 for better performance (#2)<commit_after>#[macro_use]
extern crate clap;
extern crate toml;
extern crate shogi;

mod config;
mod environment;
mod error;
mod game;
mod stats;
mod usi;

use std::io::Write;
use std::sync::atomic::Ordering;
use std::sync::mpsc::{channel, Receiver};
use std::thread;
use clap::{App, Arg};
use shogi::Color;
use shogi::bitboard::Factory;

use config::*;
use environment::*;
use error::*;
use stats::*;
use usi::*;

fn main() {
    let matches = App::new("usirun")
        .version(crate_version!())
        .about("A command line utility for running games between USI compliant Shogi engines.")
        .arg(Arg::with_name("config")
            .short("c")
            .long("config")
            .value_name("TOML")
            .help("Loads a configuration file for setting up match rules")
            .required(true)
            .takes_value(true))
        .arg(Arg::with_name("display")
            .short("d")
            .long("display")
            .value_name("MODE")
            .help("Displays ")
            .takes_value(true)
            .possible_values(&["board", "command", "simple"]))
        .get_matches();

    let mut match_config = MatchConfig::default();
    if let Some(config_path) = matches.value_of("config") {
        match_config.load(config_path)
            .expect(&format!("failed to open the config file at {}", config_path));
    }

    if let Some(display) = matches.value_of("display") {
        match_config.display = match display {
            "board" => DisplayMode::Board,
            "command" => DisplayMode::Command,
            _ => DisplayMode::Simple,
        }
    }

    Factory::init();

    match run_match(&match_config) {
        Ok(ref stats) => {
            println!("");
            println!("Match statistics");
            println!("Black: {} ({:.2}%)",
                     stats.black_wins(),
                     stats.black_win_rate() * 100.0);
            println!("White: {} ({:.2}%)",
                     stats.white_wins(),
                     stats.white_win_rate() * 100.0);
            println!("Draw : {} ({:.2}%)",
                     stats.draw_games(),
                     stats.draw_game_rate() * 100.0);

        }
        Err(e) => {
            println!("an error occurred during the match: {}", e);
        }
    }
}

fn run_match(config: &MatchConfig) -> Result<MatchStatistics, Error> {
    let mut env = Environment::new(&config.time.to_time_control())
        .initial_sfen(&config.initial_pos)
        .max_ply(config.max_ply);

    let (black_tx, black_rx) = channel();
    let (white_tx, white_rx) = channel();
    let (monitor_tx, monitor_rx) = channel();

    let mut black_engine = try!(UsiEngine::launch(Color::Black, &config.black_engine));
    let mut white_engine = try!(UsiEngine::launch(Color::White, &config.white_engine));

    if config.display == DisplayMode::Command {
        set_command_logger(Color::Black, &mut black_engine);
        set_command_logger(Color::White, &mut white_engine);
    }

    black_engine.listen(env.new_sender(), black_rx);
    white_engine.listen(env.new_sender(), white_rx);

    let monitor_handle = start_monitor_thread(&config, &black_engine, &white_engine, monitor_rx);

    for _ in 0..config.num_games {
        try!(env.start_game(&[&black_tx, &white_tx, &monitor_tx]));

        if config.display == DisplayMode::Simple {
            print!(".");
            try!(std::io::stdout().flush());
        }
    }

    try!(black_engine.kill());
    try!(white_engine.kill());

    let stats = monitor_handle.join().expect("unexpected error occurred in the monitoring thread");

    Ok(stats)
}

fn start_monitor_thread(config: &MatchConfig,
                        black_engine: &UsiEngine,
                        white_engine: &UsiEngine,
                        rx: Receiver<Event>)
                        -> thread::JoinHandle<MatchStatistics> {
    let num_games = config.num_games;
    let display = config.display;
    let black_score = black_engine.score.clone();
    let white_score = white_engine.score.clone();

    thread::spawn(move || {
        let mut results = MatchStatistics::default();

        while let Some(event) = rx.recv().ok() {
            match event {
                Event::NewTurn(game) => {
                    if display == DisplayMode::Board {
                        if let Some(game) = game.read().ok() {
                            println!("{}", game.pos);
                            println!("Time: (Black) {}s, (White) {}s",
                                     game.time.black_time().as_secs(),
                                     game.time.white_time().as_secs());
                            println!("Score (Black) {}, (White) {}",
                                     black_score.load(Ordering::Relaxed),
                                     white_score.load(Ordering::Relaxed));
                            println!("");
                        }
                    }
                }
                Event::GameOver(winner, reason) => {
                    if display == DisplayMode::Board {
                        match winner {
                            Some(c) => {
                                let name = if c == Color::Black { "Black" } else { "White" };
                                println!("Game #{}: {} wins({:?})",
                                         results.total_games() + 1,
                                         name,
                                         reason);
                            }
                            None => {
                                println!("Game #{}: Draw({:?})", results.total_games() + 1, reason);
                            }
                        }
                    }

                    results.record_game(winner);
                    if num_games == results.total_games() {
                        break;
                    }
                }
                _ => {}
            }
        }

        results
    })
}

fn set_command_logger(color: Color, engine: &mut UsiEngine) {
    let prefix = if color == Color::Black { "B" } else { "W" };

    engine.set_read_hook(Some(Box::new(move |output| {
        print!("{}> {}", prefix, output.raw_str());
    })));

    engine.set_write_hook(Some(Box::new(move |_, raw_str| {
        print!("{}< {}", prefix, raw_str);
    })));
}<|endoftext|>


<commit_before>extern crate rand;

use std::string::String;
use std::collections::HashMap;
use std::vec::Vec;
use std::fs::{File};
use std::env;
use std::path::{Path};
use std::io::{self, BufReader, BufRead, Write};
use std::iter;

use rand::{Closed01, Rng};
use rand::distributions::{IndependentSample, Sample, Gamma, LogNormal, RandSample, Range};

type Bag = HashMap<usize, usize>;

fn load_bags<P: AsRef<Path>>(path: P) -> io::Result<(Vec<usize>, Vec<Bag>, usize)> {
    let mut bags = Vec::new();
    let mut labels = Vec::new();
    let mut vocab_size = 0;
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        let mut iter = line.split_whitespace();
        let label = iter.next().unwrap().parse::<usize>().unwrap();
        let mut bag = Bag::new();
        for elm in iter {
            let mut iter = elm.split(':');
            let index = iter.next().unwrap().parse::<usize>().unwrap();
            let value = iter.next().unwrap().parse::<usize>().unwrap();
            bag.insert(index - 1, value);    // We assume one-based indexing
            if vocab_size < index {
                vocab_size = index;
            }
        }
        labels.push(label);
        bags.push(bag);
    }
    Ok((labels, bags, vocab_size))
}

fn load_text_vocabulary<P: AsRef<Path>>(path: P) -> io::Result<Vec<String>> {
    let mut vocab = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        vocab.push(line);
    }
    Ok(vocab)
}

#[derive(Clone)]
struct Categorical {
    prop: Vec<f64>,
}

impl Categorical {
    fn new(weights: Vec<f64>) -> Categorical {
        let sum: f64 = weights.iter().sum();
        let prop: Vec<f64> = weights.iter().map(|w| w / sum).collect();
        Categorical {
            prop: prop,
        }
    }
}

impl Sample<usize> for Categorical {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> usize {
        self.ind_sample(rng)
    }
}

impl IndependentSample<usize> for Categorical {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> usize {
        let closed01 = RandSample::<Closed01<f64>>::new();
        let Closed01(x) = closed01.ind_sample(rng);

        let mut sum = 0.0;
        let mut ret = self.prop.len() - 1;
        for (k, &p) in self.prop.iter().enumerate() {
            sum += p;
            if x <= sum {
                ret = k;
                break;
            }
        }
        ret
    }
}

#[derive(Clone)]
struct Dirichlet {
    alpha: Vec<f64>,
}

impl Dirichlet {
    fn new(alpha: Vec<f64>) -> Dirichlet {
        Dirichlet {
            alpha: alpha,
        }
    }
}

impl Sample<Vec<f64>> for Dirichlet {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Vec<f64> {
        self.ind_sample(rng)
    }
}

impl IndependentSample<Vec<f64>> for Dirichlet {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Vec<f64> {
        let mut sum = 0.0;
        let mut xs = Vec::new();
        for &a in &self.alpha {
            let gamma = Gamma::new(a, 1.0);
            let y = gamma.ind_sample(rng);
            xs.push(y);
            sum += y;
        }
        for x in xs.iter_mut() {
            *x /= sum;
        }
        xs
    }
}

fn lda(dataset: Vec<Bag>, alpha: Vec<f64>, beta: Vec<f64>, burn_in: usize, num_samples: usize) {
    let num_docs: usize = dataset.len();
    let num_topics: usize = alpha.len();
    let vocab_size: usize = beta.len();
    let mut rng = rand::thread_rng();
    println!("K = {}", num_topics);
    println!("M = {}", num_docs);
    println!("V = {}", vocab_size);

    write!(&mut std::io::stderr(), "Initializing...");

    // Construct zero-filled nested arrays
    let mut w:     Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut z:     Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut theta: Vec<Vec<f64>>   = Vec::with_capacity(num_docs);
    let mut phi:   Vec<Vec<f64>>   = Vec::with_capacity(num_topics);
    let mut ndk:   Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut nkv:   Vec<Vec<usize>> = Vec::with_capacity(num_topics);
    let mut z_samples: Vec<Vec<Vec<usize>>> = Vec::with_capacity(num_docs);
    for bag in &dataset {
        // n_d
        let mut n_d = 0;
        for &count in bag.values() {
            n_d += count;
        }
        w.push(vec![0; n_d]);
        z.push(vec![0; n_d]);
        theta.push(vec![0.0; num_topics]);
        ndk.push(vec![0; num_topics]);
        z_samples.push(vec![vec![0; num_topics]; n_d]);
    }
    for k in 0..num_topics {
        phi.push(vec![0.0; vocab_size]);
        nkv.push(vec![0; vocab_size]);
    }

    // Initialize w, z, ndk and nkv
    let among_topics = Range::new(0, num_topics);
    for (d, bag) in dataset.iter().enumerate() {
        let mut i = 0;
        for (&v, &count) in bag {
            for _ in 0..count {
                w[d][i] = v;
                let k = among_topics.ind_sample(&mut rng);
                z[d][i] = k;
                ndk[d][k] += 1;
                nkv[k][v] += 1;
                i += 1;
            }
        }
    }
    // Initialize theta, phi
    for d in 0..num_docs {
        let dir = Dirichlet::new(alpha.clone());
        theta[d] = dir.ind_sample(&mut rng);
    }
    for k in 0..num_topics {
        let dir = Dirichlet::new(beta.clone());
        phi[k] = dir.ind_sample(&mut rng);
    }

    println!("z = {:?}", z);
    println!("theta = {:?}", theta);
    println!("ndk = {:?}", ndk);
    println!("phi = {:?}", phi);
    println!("nkv = {:?}", nkv);

    writeln!(&mut std::io::stderr(), "\rInitialized.");

    // Sampling
    write!(&mut std::io::stderr(), "Sampling...");
    for s in 0..(burn_in + num_samples) {
        for (d, w_d) in w.iter().enumerate() {
            for (i, &w_di) in w_d.iter().enumerate() {
                let v = w_di;
                // Sample z_di
                let mut weights: Vec<f64> = Vec::with_capacity(num_topics);
                for k in 0..num_topics {
                    weights.push(theta[d][k] * phi[k][v]);
                }
                let cat = Categorical::new(weights);
                let old_z_di = z[d][i];
                let new_z_di = cat.ind_sample(&mut rng);
                z[d][i] = new_z_di;
                if s >= burn_in {
                    z_samples[d][i][new_z_di] += 1;
                }
                // Update ndk and nkv
                ndk[d][old_z_di] -= 1;
                ndk[d][new_z_di] += 1;
                nkv[old_z_di][v] -= 1;
                nkv[new_z_di][v] += 1;
            }
            // Sample theta_d
            let mut alpha_d: Vec<f64> = Vec::with_capacity(num_topics);
            for k in 0..num_topics {
                alpha_d.push(ndk[d][k] as f64 + alpha[k]);
            }
            let dir = Dirichlet::new(alpha_d);
            theta[d] = dir.ind_sample(&mut rng);
        }
        for k in 0..num_topics {
            // Sample phi_k
            let mut beta_k: Vec<f64> = Vec::with_capacity(vocab_size);
            for v in 0..vocab_size {
                beta_k.push(nkv[k][v] as f64 + beta[v]);
            }
            let dir = Dirichlet::new(beta_k);
            phi[k] = dir.ind_sample(&mut rng);
        }
        // Evaluate the log-likelihood value for the current parameters
        let mut log_likelihood = 0.0;
        for k in 0..num_topics {
            for v in 0..vocab_size {
                log_likelihood += nkv[k][v] as f64 * f64::ln(phi[k][v]);
            }
        }
        println!("log_likelihood = {}", log_likelihood);
        write!(&mut std::io::stderr(), "\rSampling... {}/{}", s, num_samples);
    }
    writeln!(&mut std::io::stderr(), "\rSampled.");

    // Distribution of z
    for (d, z_samples_d) in z_samples.iter().enumerate() {
        for (i, samples) in z_samples_d.iter().enumerate() {
            let mut sum = 0.0;
            for k in 0..num_topics {
                sum += samples[k] as f64;
            }
            print!("z[{:3}][{:3}] = ", d, i);
            for k in 0..num_topics {
                print!("{:3}:{:.2} ", k, samples[k] as f64 / sum);
            }
            println!("");
        }
    }
}

fn make_dataset(num_docs: usize, mean_nd: f64, std_dev_nd: f64, alpha: Vec<f64>, beta: Vec<f64>) -> Vec<Bag> {
    let num_topics: usize = alpha.len();
    let vocab_size: usize = beta.len();
    let mut rng = rand::thread_rng();
    // phi
    let dir_beta = Dirichlet::new(beta);
    let mut phi: Vec<Vec<f64>> = Vec::with_capacity(num_topics);
    for k in 0..num_topics {
        // Sample phi_k
        phi.push(dir_beta.ind_sample(&mut rng));
    }
    // theta, nd, z, w
    let dir_alpha = Dirichlet::new(alpha);
    let lognorm = LogNormal::new(mean_nd, std_dev_nd);
    let mut theta: Vec<Vec<f64>> = Vec::with_capacity(num_docs);
    let mut nd: Vec<usize> = Vec::with_capacity(num_docs);
    let mut z: Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut w: Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    for d in 0..num_docs {
        theta.push(dir_alpha.ind_sample(&mut rng));
        nd.push(f64::floor(lognorm.ind_sample(&mut rng)) as usize);
        let cat_theta = Categorical::new(theta[d].clone());
        let mut z_d = Vec::with_capacity(nd[d]);
        let mut w_d = Vec::with_capacity(nd[d]);
        for i in 0..nd[d] {
            z_d.push(cat_theta.ind_sample(&mut rng));
            let cat_phi = Categorical::new(phi[z_d[i]].clone());
            w_d.push(cat_phi.ind_sample(&mut rng));
        }
        z.push(z_d);
        w.push(w_d);
    }
    // Make bags
    let mut bags: Vec<Bag> = Vec::new();
    for w_d in w {
        let mut bag: Bag = Bag::new();
        for v in w_d {
            let counter = bag.entry(v).or_insert(0);
            *counter += 1;
        }
        bags.push(bag);
    }
    bags
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let (labels, dataset, vocab_size) = load_bags(&args[1]).unwrap();
    let vocab: Option<Vec<String>> = if args.len() < 3 {
        None
    }
    else {
        Some(load_text_vocabulary(&args[2]).unwrap())
    };
    println!("{:?}", labels);
    println!("{:?}", dataset);
    println!("{:?}", vocab);

    let num_topics = 3;
    let alpha: Vec<f64> = iter::repeat(1.0).take(num_topics).collect();
    let beta: Vec<f64> = iter::repeat(1.0).take(vocab_size).collect();

    lda(dataset, alpha, beta, 1000, 10000);
}
<commit_msg>Improve displaying progress<commit_after>extern crate rand;

use std::string::String;
use std::collections::HashMap;
use std::vec::Vec;
use std::fs::{File};
use std::env;
use std::path::{Path};
use std::io::{self, BufReader, BufRead, Write};
use std::iter;

use rand::{Closed01, Rng};
use rand::distributions::{IndependentSample, Sample, Gamma, LogNormal, RandSample, Range};

type Bag = HashMap<usize, usize>;

fn load_bags<P: AsRef<Path>>(path: P) -> io::Result<(Vec<usize>, Vec<Bag>, usize)> {
    let mut bags = Vec::new();
    let mut labels = Vec::new();
    let mut vocab_size = 0;
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        let mut iter = line.split_whitespace();
        let label = iter.next().unwrap().parse::<usize>().unwrap();
        let mut bag = Bag::new();
        for elm in iter {
            let mut iter = elm.split(':');
            let index = iter.next().unwrap().parse::<usize>().unwrap();
            let value = iter.next().unwrap().parse::<usize>().unwrap();
            bag.insert(index - 1, value);    // We assume one-based indexing
            if vocab_size < index {
                vocab_size = index;
            }
        }
        labels.push(label);
        bags.push(bag);
    }
    Ok((labels, bags, vocab_size))
}

fn load_text_vocabulary<P: AsRef<Path>>(path: P) -> io::Result<Vec<String>> {
    let mut vocab = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        vocab.push(line);
    }
    Ok(vocab)
}

#[derive(Clone)]
struct Categorical {
    prop: Vec<f64>,
}

impl Categorical {
    fn new(weights: Vec<f64>) -> Categorical {
        let sum: f64 = weights.iter().sum();
        let prop: Vec<f64> = weights.iter().map(|w| w / sum).collect();
        Categorical {
            prop: prop,
        }
    }
}

impl Sample<usize> for Categorical {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> usize {
        self.ind_sample(rng)
    }
}

impl IndependentSample<usize> for Categorical {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> usize {
        let closed01 = RandSample::<Closed01<f64>>::new();
        let Closed01(x) = closed01.ind_sample(rng);

        let mut sum = 0.0;
        let mut ret = self.prop.len() - 1;
        for (k, &p) in self.prop.iter().enumerate() {
            sum += p;
            if x <= sum {
                ret = k;
                break;
            }
        }
        ret
    }
}

#[derive(Clone)]
struct Dirichlet {
    alpha: Vec<f64>,
}

impl Dirichlet {
    fn new(alpha: Vec<f64>) -> Dirichlet {
        Dirichlet {
            alpha: alpha,
        }
    }
}

impl Sample<Vec<f64>> for Dirichlet {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> Vec<f64> {
        self.ind_sample(rng)
    }
}

impl IndependentSample<Vec<f64>> for Dirichlet {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> Vec<f64> {
        let mut sum = 0.0;
        let mut xs = Vec::new();
        for &a in &self.alpha {
            let gamma = Gamma::new(a, 1.0);
            let y = gamma.ind_sample(rng);
            xs.push(y);
            sum += y;
        }
        for x in xs.iter_mut() {
            *x /= sum;
        }
        xs
    }
}

fn lda(dataset: Vec<Bag>, alpha: Vec<f64>, beta: Vec<f64>, burn_in: usize, num_samples: usize) {
    let num_docs: usize = dataset.len();
    let num_topics: usize = alpha.len();
    let vocab_size: usize = beta.len();
    let mut rng = rand::thread_rng();
    println!("K = {}", num_topics);
    println!("M = {}", num_docs);
    println!("V = {}", vocab_size);

    write!(&mut std::io::stderr(), "Initializing...").unwrap();

    // Construct zero-filled nested arrays
    let mut w:     Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut z:     Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut theta: Vec<Vec<f64>>   = Vec::with_capacity(num_docs);
    let mut phi:   Vec<Vec<f64>>   = Vec::with_capacity(num_topics);
    let mut ndk:   Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut nkv:   Vec<Vec<usize>> = Vec::with_capacity(num_topics);
    let mut z_samples: Vec<Vec<Vec<usize>>> = Vec::with_capacity(num_docs);
    for bag in &dataset {
        // n_d
        let mut n_d = 0;
        for &count in bag.values() {
            n_d += count;
        }
        w.push(vec![0; n_d]);
        z.push(vec![0; n_d]);
        theta.push(vec![0.0; num_topics]);
        ndk.push(vec![0; num_topics]);
        z_samples.push(vec![vec![0; num_topics]; n_d]);
    }
    for k in 0..num_topics {
        phi.push(vec![0.0; vocab_size]);
        nkv.push(vec![0; vocab_size]);
    }

    // Initialize w, z, ndk and nkv
    let among_topics = Range::new(0, num_topics);
    for (d, bag) in dataset.iter().enumerate() {
        let mut i = 0;
        for (&v, &count) in bag {
            for _ in 0..count {
                w[d][i] = v;
                let k = among_topics.ind_sample(&mut rng);
                z[d][i] = k;
                ndk[d][k] += 1;
                nkv[k][v] += 1;
                i += 1;
            }
        }
    }
    // Initialize theta, phi
    for d in 0..num_docs {
        let dir = Dirichlet::new(alpha.clone());
        theta[d] = dir.ind_sample(&mut rng);
    }
    for k in 0..num_topics {
        let dir = Dirichlet::new(beta.clone());
        phi[k] = dir.ind_sample(&mut rng);
    }

    println!("z = {:?}", z);
    println!("theta = {:?}", theta);
    println!("ndk = {:?}", ndk);
    println!("phi = {:?}", phi);
    println!("nkv = {:?}", nkv);

    writeln!(&mut std::io::stderr(), "\rInitialized.").unwrap();

    // Sampling
    write!(&mut std::io::stderr(), "Sampling...").unwrap();
    for s in 0..(burn_in + num_samples) {
        for (d, w_d) in w.iter().enumerate() {
            for (i, &w_di) in w_d.iter().enumerate() {
                let v = w_di;
                // Sample z_di
                let mut weights: Vec<f64> = Vec::with_capacity(num_topics);
                for k in 0..num_topics {
                    weights.push(theta[d][k] * phi[k][v]);
                }
                let cat = Categorical::new(weights);
                let old_z_di = z[d][i];
                let new_z_di = cat.ind_sample(&mut rng);
                z[d][i] = new_z_di;
                if s >= burn_in {
                    z_samples[d][i][new_z_di] += 1;
                }
                // Update ndk and nkv
                ndk[d][old_z_di] -= 1;
                ndk[d][new_z_di] += 1;
                nkv[old_z_di][v] -= 1;
                nkv[new_z_di][v] += 1;
            }
            // Sample theta_d
            let mut alpha_d: Vec<f64> = Vec::with_capacity(num_topics);
            for k in 0..num_topics {
                alpha_d.push(ndk[d][k] as f64 + alpha[k]);
            }
            let dir = Dirichlet::new(alpha_d);
            theta[d] = dir.ind_sample(&mut rng);
        }
        for k in 0..num_topics {
            // Sample phi_k
            let mut beta_k: Vec<f64> = Vec::with_capacity(vocab_size);
            for v in 0..vocab_size {
                beta_k.push(nkv[k][v] as f64 + beta[v]);
            }
            let dir = Dirichlet::new(beta_k);
            phi[k] = dir.ind_sample(&mut rng);
        }
        // Evaluate the log-likelihood value for the current parameters
        let mut log_likelihood = 0.0;
        for k in 0..num_topics {
            for v in 0..vocab_size {
                log_likelihood += nkv[k][v] as f64 * f64::ln(phi[k][v]);
            }
        }
        println!("log_likelihood = {}", log_likelihood);
        if s < burn_in {
            write!(&mut std::io::stderr(), "\rBurn-in... {}/{}", s + 1, burn_in).unwrap();
        }
        else {
            if s == burn_in {
                writeln!(&mut std::io::stderr(), "").unwrap();
            }
            write!(&mut std::io::stderr(), "\rSampling... {}/{}", s - burn_in + 1, num_samples).unwrap();
        }
    }
    writeln!(&mut std::io::stderr(), "\rSampled.").unwrap();

    // Distribution of z
    for (d, z_samples_d) in z_samples.iter().enumerate() {
        for (i, samples) in z_samples_d.iter().enumerate() {
            let mut sum = 0.0;
            for k in 0..num_topics {
                sum += samples[k] as f64;
            }
            print!("z[{:3}][{:3}] = ", d, i);
            for k in 0..num_topics {
                print!("{:3}:{:.2} ", k, samples[k] as f64 / sum);
            }
            println!("");
        }
    }
}

fn make_dataset(num_docs: usize, mean_nd: f64, std_dev_nd: f64, alpha: Vec<f64>, beta: Vec<f64>) -> Vec<Bag> {
    let num_topics: usize = alpha.len();
    let vocab_size: usize = beta.len();
    let mut rng = rand::thread_rng();
    // phi
    let dir_beta = Dirichlet::new(beta);
    let mut phi: Vec<Vec<f64>> = Vec::with_capacity(num_topics);
    for k in 0..num_topics {
        // Sample phi_k
        phi.push(dir_beta.ind_sample(&mut rng));
    }
    // theta, nd, z, w
    let dir_alpha = Dirichlet::new(alpha);
    let lognorm = LogNormal::new(mean_nd, std_dev_nd);
    let mut theta: Vec<Vec<f64>> = Vec::with_capacity(num_docs);
    let mut nd: Vec<usize> = Vec::with_capacity(num_docs);
    let mut z: Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    let mut w: Vec<Vec<usize>> = Vec::with_capacity(num_docs);
    for d in 0..num_docs {
        theta.push(dir_alpha.ind_sample(&mut rng));
        nd.push(f64::floor(lognorm.ind_sample(&mut rng)) as usize);
        let cat_theta = Categorical::new(theta[d].clone());
        let mut z_d = Vec::with_capacity(nd[d]);
        let mut w_d = Vec::with_capacity(nd[d]);
        for i in 0..nd[d] {
            z_d.push(cat_theta.ind_sample(&mut rng));
            let cat_phi = Categorical::new(phi[z_d[i]].clone());
            w_d.push(cat_phi.ind_sample(&mut rng));
        }
        z.push(z_d);
        w.push(w_d);
    }
    // Make bags
    let mut bags: Vec<Bag> = Vec::new();
    for w_d in w {
        let mut bag: Bag = Bag::new();
        for v in w_d {
            let counter = bag.entry(v).or_insert(0);
            *counter += 1;
        }
        bags.push(bag);
    }
    bags
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let (labels, dataset, vocab_size) = load_bags(&args[1]).unwrap();
    let vocab: Option<Vec<String>> = if args.len() < 3 {
        None
    }
    else {
        Some(load_text_vocabulary(&args[2]).unwrap())
    };
    println!("{:?}", labels);
    println!("{:?}", dataset);
    println!("{:?}", vocab);

    let num_topics = 3;
    let alpha: Vec<f64> = iter::repeat(1.0).take(num_topics).collect();
    let beta: Vec<f64> = iter::repeat(1.0).take(vocab_size).collect();

    lda(dataset, alpha, beta, 1000, 10000);
}
<|endoftext|>


<commit_before>extern crate net2;
extern crate bincode;
extern crate rustc_serialize;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate byteorder;

use std::net::Ipv4Addr;
#[cfg(not(windows))]
use net2::unix::UnixUdpBuilderExt;
use net2::UdpBuilder;
use bincode::rustc_serialize::decode;
use byteorder::{ByteOrder, NetworkEndian};

#[repr(C)]
#[repr(packed)]
#[derive(Default)]
#[derive(Debug)]
#[derive(RustcEncodable, RustcDecodable)]
struct MDNSPacketHeader {
    id: u16, // transaction ID
    flags: u16,
    num_qn: u16,
    num_ans_rr: u16,
    num_auth_rr: u16,
    num_add_rr: u16,
}

#[derive(Debug)]
enum MDNSType {
    A,
    PTR,
    TXT,
    AAAA,
    SRV,
    NSEC,
    ANY,
    UNKNOW(u16),
}

impl MDNSType {
    #[doc(hidden)]
    pub fn from_u16(n: u16) -> MDNSType {
        match n {
            0x01 => MDNSType::A,
            0x0C => MDNSType::PTR,
            0x10 => MDNSType::TXT,
            0x1C => MDNSType::AAAA,
            0x21 => MDNSType::SRV,
            0x2F => MDNSType::NSEC,
            0xFF => MDNSType::ANY,
            _ => MDNSType::UNKNOW(n),
        }
    }

    #[doc(hidden)]
    pub fn to_u16(&self) -> u16 {
        match *self {
            MDNSType::A => 0x01,
            MDNSType::PTR => 0x0C,
            MDNSType::TXT => 0x10,
            MDNSType::AAAA => 0x1C,
            MDNSType::SRV => 0x21,
            MDNSType::NSEC => 0x2F,
            MDNSType::ANY => 0xFF,
            MDNSType::UNKNOW(n) => n,
        }
    }
}

#[derive(Debug)]
struct MDNSData {
    id: u16,
    flags: u16,
    questions: Vec<MDNSQuestion>,
    answers: Vec<MDNSAnswer>,
}

#[derive(Debug)]
struct MDNSQuestion {
    name: String,
    rr_type: MDNSType,
    rr_class: u16,
}

#[derive(Debug)]
struct MDNSAnswer {
    id: u16,
    flags: u16,
}

fn decompress_label(packet: &[u8], mut offset: usize) -> (String, usize) {
    // create the string
    let mut full_string = String::new();
    // second offset to let the caller know how much we moved
    // from the current offset
    let mut label_offset = 0;

    // helper that gets true as soon as we have hit a compress token
    // the label_offset calculation is different when we have have
    // a commpressed label
    let mut compressed = false;

    // Now go until we found a string with ZERO length
    while 0 != packet[offset] {
        // if this is not the first string we have found,
        // add a . to join them
        if full_string.len() > 0 {
            full_string.push('.');
        }

        trace!("size buffer: {:?}", &packet[offset..(offset + 2)]);

        // is this a compressed token?
        if 0xC0 == (packet[offset] & 0xC0) {
            // yes
            if false == compressed {
                compressed = true;
                // we have moved 2 bytes from the current position in memory
                label_offset += 2;
            }

            // now the second byte after the token
            // is the lower size of the uint16 offset
            // from the packet start!!!
            let small_size = packet[offset + 1] as usize;
            trace!("In BIG SIZE: {}", small_size);

            // get the higher size
            offset = (packet[offset] & 0x3F) as usize;
            trace!("In BIG SIZE_1: {}", offset);

            // shift 8 bits to the left
            offset *= 256;
            trace!("In BIG SIZE_2: {}", offset);

            // add the lower part
            offset += small_size;
        }

        // get size of the string
        let string_size = packet[offset] as usize;
        // jump over the size byte
        offset += 1;

        trace!("Before UTF8_2: {} {}", offset, string_size);
        trace!("Before UTF8: {:?}", &packet[offset..(offset + string_size)]);

        // now convert the bytes to a string
        // TODO: UTF8???? Or ASCII?
        full_string += std::str::from_utf8(&packet[offset..(offset + string_size)]).unwrap();

        // move to the next token
        offset += string_size;

        if false == compressed {
            // if we haven't had a compress token yet
            // just a the size byte and the string length
            label_offset += string_size;
            label_offset += 1;
        }
    }

    // skip the last ZERO byte indicating the end of the string
    if false == compressed {
        // but only if we have NO compressed string
        label_offset += 1;
    }

    trace!("{:?} {}", full_string, label_offset);

    // done, return the string and the offset we moved
    // in memory to get the String
    (full_string, label_offset)
}

fn parse_packet(packet: &[u8]) -> Result<MDNSData, String> {
    let decoded: MDNSPacketHeader = try!(decode(&packet[0..12]).map_err(|e| e.to_string()));

    let mut result = MDNSData {
        id: decoded.id,
        flags: decoded.flags,
        questions: vec![],
        answers: vec![],
    };

    let mut decode_position = std::mem::size_of::<MDNSPacketHeader>();
    for _ in 0..decoded.num_qn {
        // first read the label
        let (label_string, label_offset) = decompress_label(packet, decode_position);
        // move the offset
        decode_position += label_offset;

        // now read the uint16 type
        let rr_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;
        // and the class (uint16)
        let class_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;

        // Add the create and push the question struct
        result.questions.push(MDNSQuestion {
            name: label_string,
            rr_type: MDNSType::from_u16(rr_type),
            rr_class: class_type,
        });
    }

    trace!("{:?}", decoded);

    Ok(result)
}

fn main() {
    env_logger::init().unwrap();

    let udp_builder = UdpBuilder::new_v4().unwrap();

    match udp_builder.reuse_address(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't reuse_address: {}", e),
    };

    match udp_builder.reuse_port(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't reuse_address: {}", e),
    };

    let udp_socket = match udp_builder.bind("224.0.0.251:5353") {
        Ok(s) => s,
        Err(e) => panic!("couldn't bind: {}", e),
    };

    let ip_mdns = Ipv4Addr::new(224, 0, 0, 251);
    let ip_any = Ipv4Addr::new(0, 0, 0, 0);

    match udp_socket.join_multicast_v4(&ip_mdns, &ip_any) {
        Ok(s) => s,
        Err(e) => panic!("couldn't join_multicast_v4: {}", e),
    };

    match udp_socket.set_multicast_loop_v4(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't set_multicast_loop_v4: {}", e),
    };

    info!("start reading");

    loop {
        // 	let mut read_packet: mdns_packet = Default::default();
        let mut byte_array = [0u8; 65536];

        let (amt, src) = udp_socket.recv_from(&mut byte_array).unwrap();
        // println!("{:?} {:?}", amt, src);

        if let Ok(mdns_data) = parse_packet(&byte_array[0..amt]) {
            debug!("{:?}", mdns_data);
        }
    }
}
<commit_msg>Answer initial implemenation (incomplete)<commit_after>extern crate net2;
extern crate bincode;
extern crate rustc_serialize;
#[macro_use]
extern crate log;
extern crate env_logger;
extern crate byteorder;

use std::collections::HashMap;
use std::net::Ipv4Addr;
#[cfg(not(windows))]
use net2::unix::UnixUdpBuilderExt;
use net2::UdpBuilder;
use bincode::rustc_serialize::decode;
use byteorder::{ByteOrder, NetworkEndian};

#[repr(C)]
#[repr(packed)]
#[derive(Default)]
#[derive(Debug)]
#[derive(RustcEncodable, RustcDecodable)]
struct MDNSPacketHeader {
    id: u16, // transaction ID
    flags: u16,
    num_qn: u16,
    num_ans_rr: u16,
    num_auth_rr: u16,
    num_add_rr: u16,
}

#[derive(Debug)]
enum MDNSType {
    A,
    PTR,
    TXT,
    AAAA,
    SRV,
    NSEC,
    ANY,
    UNKNOW(u16),
}

impl MDNSType {
    #[doc(hidden)]
    pub fn from_u16(n: u16) -> MDNSType {
        match n {
            0x01 => MDNSType::A,
            0x0C => MDNSType::PTR,
            0x10 => MDNSType::TXT,
            0x1C => MDNSType::AAAA,
            0x21 => MDNSType::SRV,
            0x2F => MDNSType::NSEC,
            0xFF => MDNSType::ANY,
            _ => MDNSType::UNKNOW(n),
        }
    }

    #[doc(hidden)]
    pub fn to_u16(&self) -> u16 {
        match *self {
            MDNSType::A => 0x01,
            MDNSType::PTR => 0x0C,
            MDNSType::TXT => 0x10,
            MDNSType::AAAA => 0x1C,
            MDNSType::SRV => 0x21,
            MDNSType::NSEC => 0x2F,
            MDNSType::ANY => 0xFF,
            MDNSType::UNKNOW(n) => n,
        }
    }
}

#[derive(Debug)]
struct MDNSData {
    id: u16,
    flags: u16,
    questions: Vec<MDNSQuestion>,
    answers: Vec<MDNSAnswer>,
}

#[derive(Debug)]
struct MDNSQuestion {
    name: String,
    rr_type: MDNSType,
    rr_class: u16,
}

#[derive(Debug)]
struct MDNSAnswer {
    name: String,
    rr_type: MDNSType,
    rr_class: u16,
    ttl: u32,
}

fn decompress_label(packet: &[u8], mut offset: usize) -> (String, usize) {
    // create the string
    let mut full_string = String::new();
    // second offset to let the caller know how much we moved
    // from the current offset
    let mut label_offset = 0;

    // helper that gets true as soon as we have hit a compress token
    // the label_offset calculation is different when we have have
    // a commpressed label
    let mut compressed = false;

    // Now go until we found a string with ZERO length
    while 0 != packet[offset] {
        // if this is not the first string we have found,
        // add a . to join them
        if full_string.len() > 0 {
            full_string.push('.');
        }

        trace!("size buffer: {:?}", &packet[offset..(offset + 2)]);

        // is this a compressed token?
        if 0xC0 == (packet[offset] & 0xC0) {
            // yes
            if false == compressed {
                compressed = true;
                // we have moved 2 bytes from the current position in memory
                label_offset += 2;
            }

            // now the second byte after the token
            // is the lower size of the uint16 offset
            // from the packet start!!!
            let small_size = packet[offset + 1] as usize;
            trace!("In BIG SIZE: {}", small_size);

            // get the higher size
            offset = (packet[offset] & 0x3F) as usize;
            trace!("In BIG SIZE_1: {}", offset);

            // shift 8 bits to the left
            offset *= 256;
            trace!("In BIG SIZE_2: {}", offset);

            // add the lower part
            offset += small_size;
        }

        // get size of the string
        let string_size = packet[offset] as usize;
        // jump over the size byte
        offset += 1;

        trace!("Before UTF8_2: {} {}", offset, string_size);
        trace!("Before UTF8: {:?}", &packet[offset..(offset + string_size)]);

        // now convert the bytes to a string
        // TODO: UTF8???? Or ASCII?
        full_string += std::str::from_utf8(&packet[offset..(offset + string_size)]).unwrap();

        // move to the next token
        offset += string_size;

        if false == compressed {
            // if we haven't had a compress token yet
            // just a the size byte and the string length
            label_offset += string_size;
            label_offset += 1;
        }
    }

    // skip the last ZERO byte indicating the end of the string
    if false == compressed {
        // but only if we have NO compressed string
        label_offset += 1;
    }

    trace!("{:?} {}", full_string, label_offset);

    // done, return the string and the offset we moved
    // in memory to get the String
    (full_string, label_offset)
}

fn parse_packet(packet: &[u8]) -> Result<MDNSData, String> {
    let decoded: MDNSPacketHeader = try!(decode(&packet[0..12]).map_err(|e| e.to_string()));

    let mut result = MDNSData {
        id: decoded.id,
        flags: decoded.flags,
        questions: vec![],
        answers: vec![],
    };

    let mut decode_position = std::mem::size_of::<MDNSPacketHeader>();

    // parse the questions
    for _ in 0..decoded.num_qn {
        // first read the label
        let (label_string, label_offset) = decompress_label(packet, decode_position);
        // move the offset
        decode_position += label_offset;

        // now read the uint16 type
        let rr_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;
        // and the class (uint16)
        let class_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;

        // Add the create and push the question struct
        result.questions.push(MDNSQuestion {
            name: label_string,
            rr_type: MDNSType::from_u16(rr_type),
            rr_class: class_type,
        });
    }

    for _ in 0..decoded.num_ans_rr {
        // first read the label
        let (label_string, label_offset) = decompress_label(packet, decode_position);
        // move the offset
        decode_position += label_offset;

        // now read the uint16 type
        let rr_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;
        // and the class (uint16)
        let class_type = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;

        let ttl = NetworkEndian::read_u32(&packet[decode_position..(decode_position + 4)]);
        decode_position += 4;

        // RR data
        let data_len = NetworkEndian::read_u16(&packet[decode_position..(decode_position + 2)]);
        decode_position += 2;

        trace!("Data len: {}, type: {:?}",
               data_len,
               MDNSType::from_u16(rr_type));

        match MDNSType::from_u16(rr_type) {
            MDNSType::TXT => {
                let mut txt_map: Vec<String> = vec![];

                txt_map.push("Oli".to_string());
            }
            _ => {}
        }
        decode_position += data_len as usize;

        // Add the create and push the question struct
        result.answers.push(MDNSAnswer {
            name: label_string,
            rr_type: MDNSType::from_u16(rr_type),
            rr_class: class_type,
            ttl: ttl,
        });
    }

    trace!("{:?}", decoded);

    Ok(result)
}

fn main() {
    env_logger::init().unwrap();

    let udp_builder = UdpBuilder::new_v4().unwrap();

    match udp_builder.reuse_address(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't reuse_address: {}", e),
    };

    match udp_builder.reuse_port(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't reuse_address: {}", e),
    };

    let udp_socket = match udp_builder.bind("224.0.0.251:5353") {
        Ok(s) => s,
        Err(e) => panic!("couldn't bind: {}", e),
    };

    let ip_mdns = Ipv4Addr::new(224, 0, 0, 251);
    let ip_any = Ipv4Addr::new(0, 0, 0, 0);

    match udp_socket.join_multicast_v4(&ip_mdns, &ip_any) {
        Ok(s) => s,
        Err(e) => panic!("couldn't join_multicast_v4: {}", e),
    };

    match udp_socket.set_multicast_loop_v4(true) {
        Ok(s) => s,
        Err(e) => panic!("couldn't set_multicast_loop_v4: {}", e),
    };

    info!("start reading");

    loop {
        // 	let mut read_packet: mdns_packet = Default::default();
        let mut byte_array = [0u8; 65536];

        let (amt, src) = udp_socket.recv_from(&mut byte_array).unwrap();
        // println!("{:?} {:?}", amt, src);

        if let Ok(mdns_data) = parse_packet(&byte_array[0..amt]) {
            debug!("{:?}", mdns_data);
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![feature(box_syntax)]
#![feature(link_args)]

#[macro_use]
extern crate log;

#[cfg(all(not(feature = "force-glutin"), target_os = "macos"))]
extern crate libc;
#[cfg(all(not(feature = "force-glutin"), target_os = "macos"))]
extern crate cocoa;
#[cfg(all(not(feature = "force-glutin"), target_os = "macos"))]
#[macro_use]
extern crate objc;

#[cfg(any(feature = "force-glutin", not(target_os = "macos")))]
extern crate glutin;
#[cfg(any(feature = "force-glutin", not(target_os = "macos")))]
extern crate tinyfiledialogs;

#[cfg(target_os = "windows")]
extern crate winapi;
#[cfg(target_os = "windows")]
extern crate user32;
#[cfg(target_os = "windows")]
extern crate gdi32;

extern crate open;

mod traits;
mod platform;
mod servo;
mod state;
mod logs;

use platform::App;
use servo::{Servo, ServoEvent, ServoUrl, WebRenderDebugOption};
use std::env::args;
use state::{AppState, WindowState};
use traits::app::{AppEvent, AppCommand, AppMethods};
use traits::view::*;
use traits::window::{WindowEvent, WindowCommand};

const PKG_VERSION: &'static str = env!("CARGO_PKG_VERSION");
const PKG_NAME: &'static str = env!("CARGO_PKG_NAME");

fn main() {

    let logs = logs::Logger::init();

    info!("starting");

    args().find(|arg| arg == "--version").map(|_| {
        println!("{} {}", PKG_NAME, PKG_VERSION);
        std::process::exit(0);
    });

    let resources_path = App::get_resources_path().expect("Can't find resources path");

    let app = App::new().expect("Can't create application");
    let window = app.new_window().expect("Can't create application");

    let view = window.new_view().unwrap();

    Servo::configure(resources_path.clone());

    let servo = {
        let geometry = view.get_geometry();
        let waker = window.new_event_loop_waker();
        Servo::new(geometry, view.clone(), waker)
    };

    let home_url = resources_path.parent().unwrap().join("shell_resources").join("home.html");
    let home_url = ServoUrl::from_file_path(&home_url).unwrap().into_string();

    // Skip first argument (executable), and find the first
    // argument that doesn't start with `-`
    let url = args().skip(1).find(|arg| {
        !arg.starts_with("-")
    }).unwrap_or(home_url);

    let browser = servo.new_browser(&url);
    servo.select_browser(browser.id);

    let mut app_state = AppState::new();
    app_state.current_window_index = Some(0);

    let mut win_state = WindowState::new();
    win_state.current_browser_index = Some(0);
    win_state.browsers.push(browser);

    app.render(&app_state);
    window.render(&win_state);

    info!("Servo version: {}", servo.version());

    let handle_events = || {

        // Loop until no events are available anymore.
        loop {

            let before_app_state = app_state.clone();
            let before_win_state = win_state.clone();

            let app_events = app.get_events();
            let win_events = window.get_events();
            let view_events = view.get_events();
            let servo_events = servo.get_events();

            if app_events.is_empty() &&
               win_events.is_empty() &&
               view_events.is_empty() &&
               servo_events.is_empty() {
                   break
            }

            // FIXME: it's really annoying we need this
            let mut force_sync = false;

            for event in app_events {
                match event {
                    AppEvent::DidFinishLaunching => {
                        // FIXME: does this work?
                    }
                    AppEvent::WillTerminate => {
                        // FIXME: does this work?
                    }
                    AppEvent::DidChangeScreenParameters => {
                        // FIXME: does this work?
                        servo.update_geometry(view.get_geometry());
                        view.update_drawable();
                    }
                    AppEvent::DoCommand(cmd) => {
                        match cmd {
                            AppCommand::ClearHistory => {
                                // FIXME
                            }
                            AppCommand::ToggleOptionDarkTheme => {
                                app_state.dark_theme = !app_state.dark_theme;
                            }
                        }
                    }
                }
            }

            for event in win_events {
                match event {
                    WindowEvent::EventLoopAwaken => {
                        force_sync = true;
                    }
                    WindowEvent::GeometryDidChange => {
                        servo.update_geometry(view.get_geometry());
                        view.update_drawable();
                    }
                    WindowEvent::DidEnterFullScreen => {
                        // FIXME
                    }
                    WindowEvent::DidExitFullScreen => {
                        // FIXME
                    }
                    WindowEvent::WillClose => {
                        // FIXME
                    }
                    WindowEvent::OptionsClosed => {
                        win_state.options_open = false;
                    }
                    WindowEvent::UrlbarFocusChanged(focused) => {
                        win_state.urlbar_focused = focused;
                    }
                    WindowEvent::DoCommand(cmd) => {
                        let idx = win_state.current_browser_index.unwrap();
                        let bid = win_state.browsers[idx].id;
                        match cmd {
                            WindowCommand::Stop => {
                                // FIXME
                            }
                            WindowCommand::Reload => {
                                servo.reload(bid);
                            }
                            WindowCommand::NavigateBack => {
                                servo.go_back(bid);
                            }
                            WindowCommand::NavigateForward => {
                                servo.go_forward(bid);
                            }
                            WindowCommand::OpenLocation => {
                                win_state.urlbar_focused = true;
                            }
                            WindowCommand::OpenInDefaultBrowser => {
                                if let Some(ref url) = win_state.browsers[idx].url {
                                    open::that(url.clone()).ok();
                                }
                            }
                            WindowCommand::ZoomIn => {
                                win_state.browsers[idx].zoom *= 1.1;
                                servo.zoom(win_state.browsers[idx].zoom);
                            }
                            WindowCommand::ZoomOut => {
                                win_state.browsers[idx].zoom /= 1.1;
                                servo.zoom(win_state.browsers[idx].zoom);
                            }
                            WindowCommand::ZoomToActualSize => {
                                win_state.browsers[idx].zoom = 1.0;
                                servo.reset_zoom();
                            }

                            WindowCommand::ToggleSidebar => {
                                win_state.sidebar_is_open = !win_state.sidebar_is_open;
                            }

                            WindowCommand::ShowOptions => {
                                win_state.options_open = !win_state.options_open;
                            }

                            WindowCommand::Load(request) => {
                                win_state.browsers[idx].user_input = Some(request.clone());
                                win_state.urlbar_focused = false;
                                let url = ServoUrl::parse(&request).or_else(|error| {
                                    // FIXME: weak
                                    if request.ends_with(".com") || request.ends_with(".org") || request.ends_with(".net") {
                                        ServoUrl::parse(&format!("http://{}", request))
                                    } else {
                                        Err(error)
                                    }
                                }).or_else(|_| {
                                    ServoUrl::parse(&format!("https://duckduckgo.com/html/?q={}", request))
                                });
                                match url {
                                    Ok(url) => {
                                        servo.load_url(bid, url)
                                    },
                                    Err(err) => warn!("Can't parse url: {}", err),
                                }
                            }
                            WindowCommand::ToggleOptionShowLogs => {
                                win_state.logs_visible = !win_state.logs_visible;
                            },
                            WindowCommand::NewTab => {
                                let browser = servo.new_browser("about:blank");
                                servo.select_browser(browser.id);
                                servo.update_geometry(view.get_geometry());
                                win_state.current_browser_index = Some(idx + 1);
                                win_state.browsers.push(browser);
                                if cfg!(all(not(feature = "force-glutin"), target_os = "macos")) {
                                    // Focus urlbar, but only on cocoa
                                    win_state.urlbar_focused = true;
                                }
                            },
                            WindowCommand::CloseTab => {
                                if win_state.browsers.len() > 1 {
                                    let id = win_state.browsers[idx].id;
                                    let new_id = if idx == win_state.browsers.len() - 1 {
                                        win_state.current_browser_index = Some(idx - 1);
                                        win_state.browsers[idx - 1].id
                                    } else {
                                        win_state.browsers[idx + 1].id
                                    };
                                    servo.select_browser(new_id);
                                    servo.close_browser(id);
                                    win_state.browsers.remove(idx);
                                }
                            },
                            WindowCommand::PrevTab => {
                                let new_idx = if idx == 0 {
                                    win_state.browsers.len() - 1
                                } else {
                                    idx - 1
                                };
                                win_state.current_browser_index = Some(new_idx);
                                let id = win_state.browsers[new_idx].id;
                                servo.select_browser(id);
                            },
                            WindowCommand::NextTab => {
                                let new_idx = if idx == win_state.browsers.len() - 1 {
                                    0
                                } else {
                                    idx + 1
                                };
                                win_state.current_browser_index = Some(new_idx);
                                let id = win_state.browsers[new_idx].id;
                                servo.select_browser(id);
                            },
                            WindowCommand::SelectTab(idx) => {
                                if win_state.current_browser_index != Some(idx) {
                                    win_state.current_browser_index = Some(idx);
                                    let id = win_state.browsers[idx].id;
                                    servo.select_browser(id);
                                }
                            },
                            WindowCommand::ToggleOptionFragmentBorders => { },
                            WindowCommand::ToggleOptionParallelDisplayListBuidling => { },
                            WindowCommand::ToggleOptionShowParallelLayout => { },
                            WindowCommand::ToggleOptionConvertMouseToTouch => { },
                            WindowCommand::ToggleOptionTileBorders => { },

                            WindowCommand::ToggleOptionWRProfiler => {
                                win_state.debug_options.wr_profiler = !win_state.debug_options.wr_profiler;
                                servo.toggle_webrender_debug_option(WebRenderDebugOption::Profiler);
                            },

                            WindowCommand::ToggleOptionWRTextureCacheDebug => {
                                win_state.debug_options.wr_texture_cache_debug = !win_state.debug_options.wr_texture_cache_debug;
                                servo.toggle_webrender_debug_option(WebRenderDebugOption::TextureCacheDebug);
                            },

                            WindowCommand::ToggleOptionWRTargetDebug => {
                                win_state.debug_options.wr_render_target_debug = !win_state.debug_options.wr_render_target_debug;
                                servo.toggle_webrender_debug_option(WebRenderDebugOption::RenderTargetDebug);
                            },
                        }
                    }
                }
            }

            for event in view_events {
                let idx = win_state.current_browser_index.unwrap();
                let state = &mut win_state.browsers[idx];
                match event {
                    ViewEvent::GeometryDidChange => {
                        servo.update_geometry(view.get_geometry());
                        view.update_drawable();
                    }
                    ViewEvent::MouseWheel(delta, phase) => {
                        // FIXME: magic value
                        static LINE_HEIGHT: f32 = 38.0;
                        let (mut x, mut y) = match delta {
                            MouseScrollDelta::PixelDelta(x, y) => (x, y),
                            MouseScrollDelta::LineDelta(x, y) => (x, y * LINE_HEIGHT),
                        };
                        if y.abs() >= x.abs() { x = 0.0; } else { y = 0.0; }
                        servo.perform_scroll(0, 0, x, y, phase);
                    }
                    ViewEvent::MouseMoved(x, y) => {
                        servo.perform_mouse_move(x, y);
                    }
                    ViewEvent::MouseInput(element_state, button, x, y) => {
                        servo.perform_click(x, y, element_state, button);
                    }
                    ViewEvent::KeyEvent(c, key, keystate, modifiers) => {
                        servo.send_key(state.id, c, key, keystate, modifiers);
                    }
                }
            }

            for event in servo_events {
                match event {
                    ServoEvent::SetWindowInnerSize(..) => {
                        // ignore
                    }
                    ServoEvent::SetWindowPosition(..) => {
                        // ignore
                    }
                    ServoEvent::SetFullScreenState(fullscreen) => {
                        if fullscreen {
                            view.enter_fullscreen();
                        } else {
                            view.exit_fullscreen();
                        }
                    }
                    ServoEvent::TitleChanged(id, title) => {
                        match win_state.browsers.iter_mut().find(|b| b.id == id) {
                            Some(browser) => {
                                browser.title = title;
                            }
                            None => warn!("Got message for unkown browser:  {:?}", id)
                        }
                    }
                    ServoEvent::StatusChanged(status) => {
                        win_state.status = status;
                    }
                    ServoEvent::LoadStart(id) => {
                        match win_state.browsers.iter_mut().find(|b| b.id == id) {
                            Some(browser) => {
                                browser.is_loading = true;
                            }
                            None => warn!("Got message for unkown browser:  {:?}", id)
                        }
                    }
                    ServoEvent::LoadEnd(id) => {
                        match win_state.browsers.iter_mut().find(|b| b.id == id) {
                            Some(browser) => {
                                browser.is_loading = false;
                            }
                            None => warn!("Got message for unkown browser:  {:?}", id)
                        }
                    }
                    ServoEvent::HeadParsed(..) => {
                        // FIXME
                    }
                    ServoEvent::HistoryChanged(id, entries, current) => {
                        match win_state.browsers.iter_mut().find(|b| b.id == id) {
                            Some(browser) => {
                                let url = entries[current].url.to_string();
                                browser.url = Some(url);
                                browser.can_go_back = current > 0;
                                browser.can_go_forward = current < entries.len() - 1;
                            }
                            None => warn!("Got message for unkown browser:  {:?}", id)
                        }
                    }
                    ServoEvent::CursorChanged(cursor) => {
                        app_state.cursor = cursor;
                    }
                    ServoEvent::FaviconChanged(..) => {
                        // FIXME
                    }
                    ServoEvent::Key(..) => {
                        // FIXME
                    }
                    ServoEvent::OpenInDefaultBrowser(url) => {
                        open::that(url).ok();
                    }
                }
            }

            let app_has_changed = before_app_state == app_state;
            let win_has_changed = before_win_state == win_state;

            if app_has_changed || win_has_changed {
                app.render(&app_state);
                window.render(&win_state);
            }

            servo.sync(force_sync);
        }

        // Here, only stuff that we know for sure won't trigger any
        // new events

        // FIXME: logs will grow until pulled
        if win_state.logs_visible {
            window.append_logs(&logs.get_logs());
        }
    };

    view.set_live_resize_callback(&handle_events);

    app.run(handle_events);

}
<commit_msg>window -> win<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#![feature(box_syntax)]
#![feature(link_args)]

#[macro_use]
extern crate log;

#[cfg(all(not(feature = "force-glutin"), target_os = "macos"))]
extern crate libc;
#[cfg(all(not(feature = "force-glutin"), target_os = "macos"))]
extern crate cocoa;
#[cfg(all(not(feature = "force-glutin"), target_os = "macos"))]
#[macro_use]
extern crate objc;

#[cfg(any(feature = "force-glutin", not(target_os = "macos")))]
extern crate glutin;
#[cfg(any(feature = "force-glutin", not(target_os = "macos")))]
extern crate tinyfiledialogs;

#[cfg(target_os = "windows")]
extern crate winapi;
#[cfg(target_os = "windows")]
extern crate user32;
#[cfg(target_os = "windows")]
extern crate gdi32;

extern crate open;

mod traits;
mod platform;
mod servo;
mod state;
mod logs;

use platform::App;
use servo::{Servo, ServoEvent, ServoUrl, WebRenderDebugOption};
use std::env::args;
use state::{AppState, WindowState};
use traits::app::{AppEvent, AppCommand, AppMethods};
use traits::view::*;
use traits::window::{WindowEvent, WindowCommand};

const PKG_VERSION: &'static str = env!("CARGO_PKG_VERSION");
const PKG_NAME: &'static str = env!("CARGO_PKG_NAME");

fn main() {

    let logs = logs::Logger::init();

    info!("starting");

    args().find(|arg| arg == "--version").map(|_| {
        println!("{} {}", PKG_NAME, PKG_VERSION);
        std::process::exit(0);
    });

    let resources_path = App::get_resources_path().expect("Can't find resources path");

    let app = App::new().expect("Can't create application");
    let win = app.new_window().expect("Can't create application");

    let view = win.new_view().unwrap();

    Servo::configure(resources_path.clone());

    let servo = {
        let geometry = view.get_geometry();
        let waker = win.new_event_loop_waker();
        Servo::new(geometry, view.clone(), waker)
    };

    let home_url = resources_path.parent().unwrap().join("shell_resources").join("home.html");
    let home_url = ServoUrl::from_file_path(&home_url).unwrap().into_string();

    // Skip first argument (executable), and find the first
    // argument that doesn't start with `-`
    let url = args().skip(1).find(|arg| {
        !arg.starts_with("-")
    }).unwrap_or(home_url);

    let browser = servo.new_browser(&url);
    servo.select_browser(browser.id);

    let mut app_state = AppState::new();
    app_state.current_window_index = Some(0);

    let mut win_state = WindowState::new();
    win_state.current_browser_index = Some(0);
    win_state.browsers.push(browser);

    app.render(&app_state);
    win.render(&win_state);

    info!("Servo version: {}", servo.version());

    let handle_events = || {

        // Loop until no events are available anymore.
        loop {

            let before_app_state = app_state.clone();
            let before_win_state = win_state.clone();

            let app_events = app.get_events();
            let win_events = win.get_events();
            let view_events = view.get_events();
            let servo_events = servo.get_events();

            if app_events.is_empty() &&
               win_events.is_empty() &&
               view_events.is_empty() &&
               servo_events.is_empty() {
                   break
            }

            // FIXME: it's really annoying we need this
            let mut force_sync = false;

            for event in app_events {
                match event {
                    AppEvent::DidFinishLaunching => {
                        // FIXME: does this work?
                    }
                    AppEvent::WillTerminate => {
                        // FIXME: does this work?
                    }
                    AppEvent::DidChangeScreenParameters => {
                        // FIXME: does this work?
                        servo.update_geometry(view.get_geometry());
                        view.update_drawable();
                    }
                    AppEvent::DoCommand(cmd) => {
                        match cmd {
                            AppCommand::ClearHistory => {
                                // FIXME
                            }
                            AppCommand::ToggleOptionDarkTheme => {
                                app_state.dark_theme = !app_state.dark_theme;
                            }
                        }
                    }
                }
            }

            for event in win_events {
                match event {
                    WindowEvent::EventLoopAwaken => {
                        force_sync = true;
                    }
                    WindowEvent::GeometryDidChange => {
                        servo.update_geometry(view.get_geometry());
                        view.update_drawable();
                    }
                    WindowEvent::DidEnterFullScreen => {
                        // FIXME
                    }
                    WindowEvent::DidExitFullScreen => {
                        // FIXME
                    }
                    WindowEvent::WillClose => {
                        // FIXME
                    }
                    WindowEvent::OptionsClosed => {
                        win_state.options_open = false;
                    }
                    WindowEvent::UrlbarFocusChanged(focused) => {
                        win_state.urlbar_focused = focused;
                    }
                    WindowEvent::DoCommand(cmd) => {
                        let idx = win_state.current_browser_index.unwrap();
                        let bid = win_state.browsers[idx].id;
                        match cmd {
                            WindowCommand::Stop => {
                                // FIXME
                            }
                            WindowCommand::Reload => {
                                servo.reload(bid);
                            }
                            WindowCommand::NavigateBack => {
                                servo.go_back(bid);
                            }
                            WindowCommand::NavigateForward => {
                                servo.go_forward(bid);
                            }
                            WindowCommand::OpenLocation => {
                                win_state.urlbar_focused = true;
                            }
                            WindowCommand::OpenInDefaultBrowser => {
                                if let Some(ref url) = win_state.browsers[idx].url {
                                    open::that(url.clone()).ok();
                                }
                            }
                            WindowCommand::ZoomIn => {
                                win_state.browsers[idx].zoom *= 1.1;
                                servo.zoom(win_state.browsers[idx].zoom);
                            }
                            WindowCommand::ZoomOut => {
                                win_state.browsers[idx].zoom /= 1.1;
                                servo.zoom(win_state.browsers[idx].zoom);
                            }
                            WindowCommand::ZoomToActualSize => {
                                win_state.browsers[idx].zoom = 1.0;
                                servo.reset_zoom();
                            }

                            WindowCommand::ToggleSidebar => {
                                win_state.sidebar_is_open = !win_state.sidebar_is_open;
                            }

                            WindowCommand::ShowOptions => {
                                win_state.options_open = !win_state.options_open;
                            }

                            WindowCommand::Load(request) => {
                                win_state.browsers[idx].user_input = Some(request.clone());
                                win_state.urlbar_focused = false;
                                let url = ServoUrl::parse(&request).or_else(|error| {
                                    // FIXME: weak
                                    if request.ends_with(".com") || request.ends_with(".org") || request.ends_with(".net") {
                                        ServoUrl::parse(&format!("http://{}", request))
                                    } else {
                                        Err(error)
                                    }
                                }).or_else(|_| {
                                    ServoUrl::parse(&format!("https://duckduckgo.com/html/?q={}", request))
                                });
                                match url {
                                    Ok(url) => {
                                        servo.load_url(bid, url)
                                    },
                                    Err(err) => warn!("Can't parse url: {}", err),
                                }
                            }
                            WindowCommand::ToggleOptionShowLogs => {
                                win_state.logs_visible = !win_state.logs_visible;
                            },
                            WindowCommand::NewTab => {
                                let browser = servo.new_browser("about:blank");
                                servo.select_browser(browser.id);
                                servo.update_geometry(view.get_geometry());
                                win_state.current_browser_index = Some(idx + 1);
                                win_state.browsers.push(browser);
                                if cfg!(all(not(feature = "force-glutin"), target_os = "macos")) {
                                    // Focus urlbar, but only on cocoa
                                    win_state.urlbar_focused = true;
                                }
                            },
                            WindowCommand::CloseTab => {
                                if win_state.browsers.len() > 1 {
                                    let id = win_state.browsers[idx].id;
                                    let new_id = if idx == win_state.browsers.len() - 1 {
                                        win_state.current_browser_index = Some(idx - 1);
                                        win_state.browsers[idx - 1].id
                                    } else {
                                        win_state.browsers[idx + 1].id
                                    };
                                    servo.select_browser(new_id);
                                    servo.close_browser(id);
                                    win_state.browsers.remove(idx);
                                }
                            },
                            WindowCommand::PrevTab => {
                                let new_idx = if idx == 0 {
                                    win_state.browsers.len() - 1
                                } else {
                                    idx - 1
                                };
                                win_state.current_browser_index = Some(new_idx);
                                let id = win_state.browsers[new_idx].id;
                                servo.select_browser(id);
                            },
                            WindowCommand::NextTab => {
                                let new_idx = if idx == win_state.browsers.len() - 1 {
                                    0
                                } else {
                                    idx + 1
                                };
                                win_state.current_browser_index = Some(new_idx);
                                let id = win_state.browsers[new_idx].id;
                                servo.select_browser(id);
                            },
                            WindowCommand::SelectTab(idx) => {
                                if win_state.current_browser_index != Some(idx) {
                                    win_state.current_browser_index = Some(idx);
                                    let id = win_state.browsers[idx].id;
                                    servo.select_browser(id);
                                }
                            },
                            WindowCommand::ToggleOptionFragmentBorders => { },
                            WindowCommand::ToggleOptionParallelDisplayListBuidling => { },
                            WindowCommand::ToggleOptionShowParallelLayout => { },
                            WindowCommand::ToggleOptionConvertMouseToTouch => { },
                            WindowCommand::ToggleOptionTileBorders => { },

                            WindowCommand::ToggleOptionWRProfiler => {
                                win_state.debug_options.wr_profiler = !win_state.debug_options.wr_profiler;
                                servo.toggle_webrender_debug_option(WebRenderDebugOption::Profiler);
                            },

                            WindowCommand::ToggleOptionWRTextureCacheDebug => {
                                win_state.debug_options.wr_texture_cache_debug = !win_state.debug_options.wr_texture_cache_debug;
                                servo.toggle_webrender_debug_option(WebRenderDebugOption::TextureCacheDebug);
                            },

                            WindowCommand::ToggleOptionWRTargetDebug => {
                                win_state.debug_options.wr_render_target_debug = !win_state.debug_options.wr_render_target_debug;
                                servo.toggle_webrender_debug_option(WebRenderDebugOption::RenderTargetDebug);
                            },
                        }
                    }
                }
            }

            for event in view_events {
                let idx = win_state.current_browser_index.unwrap();
                let state = &mut win_state.browsers[idx];
                match event {
                    ViewEvent::GeometryDidChange => {
                        servo.update_geometry(view.get_geometry());
                        view.update_drawable();
                    }
                    ViewEvent::MouseWheel(delta, phase) => {
                        // FIXME: magic value
                        static LINE_HEIGHT: f32 = 38.0;
                        let (mut x, mut y) = match delta {
                            MouseScrollDelta::PixelDelta(x, y) => (x, y),
                            MouseScrollDelta::LineDelta(x, y) => (x, y * LINE_HEIGHT),
                        };
                        if y.abs() >= x.abs() { x = 0.0; } else { y = 0.0; }
                        servo.perform_scroll(0, 0, x, y, phase);
                    }
                    ViewEvent::MouseMoved(x, y) => {
                        servo.perform_mouse_move(x, y);
                    }
                    ViewEvent::MouseInput(element_state, button, x, y) => {
                        servo.perform_click(x, y, element_state, button);
                    }
                    ViewEvent::KeyEvent(c, key, keystate, modifiers) => {
                        servo.send_key(state.id, c, key, keystate, modifiers);
                    }
                }
            }

            for event in servo_events {
                match event {
                    ServoEvent::SetWindowInnerSize(..) => {
                        // ignore
                    }
                    ServoEvent::SetWindowPosition(..) => {
                        // ignore
                    }
                    ServoEvent::SetFullScreenState(fullscreen) => {
                        if fullscreen {
                            view.enter_fullscreen();
                        } else {
                            view.exit_fullscreen();
                        }
                    }
                    ServoEvent::TitleChanged(id, title) => {
                        match win_state.browsers.iter_mut().find(|b| b.id == id) {
                            Some(browser) => {
                                browser.title = title;
                            }
                            None => warn!("Got message for unkown browser:  {:?}", id)
                        }
                    }
                    ServoEvent::StatusChanged(status) => {
                        win_state.status = status;
                    }
                    ServoEvent::LoadStart(id) => {
                        match win_state.browsers.iter_mut().find(|b| b.id == id) {
                            Some(browser) => {
                                browser.is_loading = true;
                            }
                            None => warn!("Got message for unkown browser:  {:?}", id)
                        }
                    }
                    ServoEvent::LoadEnd(id) => {
                        match win_state.browsers.iter_mut().find(|b| b.id == id) {
                            Some(browser) => {
                                browser.is_loading = false;
                            }
                            None => warn!("Got message for unkown browser:  {:?}", id)
                        }
                    }
                    ServoEvent::HeadParsed(..) => {
                        // FIXME
                    }
                    ServoEvent::HistoryChanged(id, entries, current) => {
                        match win_state.browsers.iter_mut().find(|b| b.id == id) {
                            Some(browser) => {
                                let url = entries[current].url.to_string();
                                browser.url = Some(url);
                                browser.can_go_back = current > 0;
                                browser.can_go_forward = current < entries.len() - 1;
                            }
                            None => warn!("Got message for unkown browser:  {:?}", id)
                        }
                    }
                    ServoEvent::CursorChanged(cursor) => {
                        app_state.cursor = cursor;
                    }
                    ServoEvent::FaviconChanged(..) => {
                        // FIXME
                    }
                    ServoEvent::Key(..) => {
                        // FIXME
                    }
                    ServoEvent::OpenInDefaultBrowser(url) => {
                        open::that(url).ok();
                    }
                }
            }

            let app_has_changed = before_app_state == app_state;
            let win_has_changed = before_win_state == win_state;
            if app_has_changed || win_has_changed {
                app.render(&app_state);
                win.render(&win_state);
            }

            servo.sync(force_sync);
        }

        // Here, only stuff that we know for sure won't trigger any
        // new events

        // FIXME: logs will grow until pulled
        if win_state.logs_visible {
            win.append_logs(&logs.get_logs());
        }
    };

    view.set_live_resize_callback(&handle_events);

    app.run(handle_events);

}
<|endoftext|>