<commit_before>extern crate libc;

use libc::{c_int, size_t};
use std::io;
use std::ptr;
use std::mem;
use std::io::process::{Command, StdioContainer};

use termios::*;
use signal::*;
use constants::*;

mod termios;
mod signal;
mod constants;

type Stdr = io::stdio::StdinReader;
type Stdw = io::LineBufferedWriter<io::stdio::StdWriter>;

// start off as null pointer
static mut instate_location:size_t = 0;

struct InputState {
    line: Vec<String>,
    word: String,
    part: String,
    bpart: String,
    stdin: Stdr,
    stdout: Stdw,
    stderr: Stdw
}

impl InputState {
    fn new() -> InputState {
        InputState {
            line: Vec::<String>::new(),
            word: String::new(),
            part: String::new(),
            bpart: String::new(),
            stdin: io::stdin(),
            stdout: io::stdout(),
            stderr: io::stderr()
        }
    }

    fn clear(&mut self) {
        self.word.clear();
        self.line.clear();
        self.part.clear();
        self.bpart.clear();
    }

    fn flush(&mut self) {
        self.stdout.flush().unwrap();
    }

    fn outw(&mut self, msg:&str) {
        self.stdout.write_str(msg).unwrap();
    }

    fn outc(&mut self, ch:char) {
        self.stdout.write_char(ch).unwrap();
    }

    fn err(&mut self, msg:&str) {
        self.stderr.write_str(msg).unwrap();
    }
}

#[allow(unused_variables)]
unsafe extern fn handle_sigint(signum:c_int, siginfo:*const SigInfo, context:size_t) {
    if instate_location == 0 {
        panic!("Shell state location uninitialized");
    }
    print!("^C\n");
    // Below: completely disregarding everything Rust stands for
    let state:&mut InputState = mem::transmute(instate_location);
    state.clear();
    io::stdio::flush();
}

fn set_instate_location(state:&mut InputState) {
    unsafe {
        if instate_location != 0 {
            panic!("Tried to set shell state location twice");
        }
        instate_location = mem::transmute::<&mut InputState, size_t>(state);
    }
}

fn prepare_terminal(tios:&mut Termios) {
    tios.ldisable(ICANON);
    tios.ldisable(ECHO);
}

fn update_terminal(tios:Termios, state:&mut InputState) -> bool {
    if !tios.set() {
        state.err("Warning: Could not set terminal mode\n");
        return false;
    }
    return true;
}

fn handle_escape(state:&mut InputState) {
    // Handle an ANSI escape sequence
    if state.stdin.read_char() != Ok(ANSI) {
        return;
    }
    match state.stdin.read_char() {
        Err(_) => return,
        Ok('D') => {
            // left
            match state.word.pop() {
                Some(c) => {
                    state.bpart.clear();
                    state.part.push(c);
                    cursor_left(state);
                },
                None => match state.line.pop() {
                    Some(s) => {
                        state.bpart.clear();
                        state.part.push(' ');
                        state.word.clear();
                        state.word.push_str(s.as_slice());
                        cursor_left(state);
                    }
                    None => return
                }
            }
        },
        Ok('C') => {
            // right
            match state.part.pop() {
                Some(' ') => {
                    state.bpart.clear();
                    state.line.push(state.word.clone());
                    state.word.clear();
                    cursor_right(state);
                },
                Some(c) => {
                    state.bpart.clear();
                    state.word.push(c);
                    cursor_right(state);
                },
                None => return
            }
        },
        Ok(_) => return
    }
}

fn cursor_left(state:&mut InputState) {
    state.outc(DEL);
}

fn cursor_right(state:&mut InputState) {
    state.outw(CRSR_RIGHT);
}

fn draw_part(state:&mut InputState) {
    // quick out if part is empty
    if state.part.is_empty() {
        return;
    }
    if state.bpart.is_empty() {
        // only calculate bpart when it needs to be recalculated
        let mut cpart = state.part.clone();
        loop {
            match cpart.pop() {
                Some(c) => state.bpart.push(c),
                None => break
            }
        }
    }
    let c = state.bpart.clone();
    state.outw(c.as_slice());
}

// work around lack of DST
fn build_string(ch:char, count:uint) -> String {
    let mut s = String::new();
    let mut i = 0u;
    loop {
        if i == count {
            return s;
        }
        s.push(ch);
        i += 1;
    }
}

fn cursors_left(by:uint, state:&mut InputState) {
    // move back by a given number of characters
    state.outw(build_string(DEL, by).as_slice());
}

fn idraw_part(state:&mut InputState) {
    // in-place draw of the line part
    draw_part(state);
    cursors_left(state.part.len(), state);
}

fn prepare_signals(state:&mut InputState) {
    let mut sa = SigAction {
        handler: handle_sigint,
        mask: [0; SIGSET_NWORDS],
        flags: SA_RESTART | SA_SIGINFO
    };
    unsafe {
        if sigfillset(&mut sa.mask) != 0 {
            state.err("Warning: could not fill mask set for SIGINT handler\n");
        }
        if sigaction(SIGINT, &sa, ptr::null_mut::<SigAction>()) != 0 {
            state.err("Warning: could not set handler for SIGINT\n");
        }
    }
}

fn run_command(state:&mut InputState) {
    let mut process = Command::new(&state.line[0]);
    process.args(state.line.slice_from(1));
    process.stdout(StdioContainer::InheritFd(STDOUT));
    process.stdin(StdioContainer::InheritFd(STDIN));
    process.stderr(StdioContainer::InheritFd(STDERR));
    let mut child = match process.spawn() {
        Err(e) => {
            let s = state.line[0].clone();
            state.err(format!("Couldn't spawn {}: {}\n", s, e).as_slice());
            return;
        },
        Ok(child) => child
    };
    match child.wait() {
        Err(e) => {
            state.err(format!("Couldn't wait for child to exit: {}\n", e.desc).as_slice());
        },
        Ok(_) => {
            // nothing
        }
    };
}

fn main() {
    let mut state = &mut InputState::new();
    prepare_signals(state);
    let mut tios = Termios::new();
    tios.get();
    let old_tios = tios.clone();
    prepare_terminal(&mut tios);
    update_terminal(tios, state);

    set_instate_location(state);
    loop {
        // Note: in non-canonical mode
        match state.stdin.read_char() {
            Ok(EOF) => {
                if state.line.is_empty() && state.word.is_empty() {
                    break;
                }
            },
            Ok(NL) => {
                // start command output on next line
                state.outc(NL);
                // push any remaining word onto the line
                if !state.word.is_empty() {
                    state.line.push(state.word.clone());
                }
                // debug info
                let numargs = state.line.len();
                state.outw(format!("Number of arguments: {}\n", numargs).as_slice());
                // run command if one was specified
                if !state.line.is_empty() {
                    // run command
                    update_terminal(old_tios, state);
                    run_command(state);
                    update_terminal(tios, state);
                }
                // clear the state
                state.clear();
            },
            Ok(DEL) => {
                if state.word.is_empty() {
                    state.word = match state.line.pop() {
                        Some(s) => s,
                        None => continue
                    };
                } else {
                    state.word.pop();
                }
                cursor_left(state);
                draw_part(state);
                state.outc(SPC);
                cursors_left(state.part.len() + 1, state);
            },
            Ok(ESC) => handle_escape(state),
            Ok(SPC) => {
                state.line.push(state.word.clone());
                state.word.clear();
                state.outc(SPC);
                idraw_part(state);
            },
            Ok(c) => {
                state.word.push(c);
                state.outc(c);
                idraw_part(state);
            },
            Err(e) => {
                state.err(format!("Error: {}\n", e).as_slice());
                break;
            }
        }
        // flush output
        state.flush();
    }
    state.outw("Exiting\n");
    // restore old term state
    update_terminal(old_tios, state);
}
<commit_msg>Remove debug output<commit_after>extern crate libc;

use libc::{c_int, size_t};
use std::io;
use std::ptr;
use std::mem;
use std::io::process::{Command, StdioContainer};

use termios::*;
use signal::*;
use constants::*;

mod termios;
mod signal;
mod constants;

type Stdr = io::stdio::StdinReader;
type Stdw = io::LineBufferedWriter<io::stdio::StdWriter>;

// start off as null pointer
static mut instate_location:size_t = 0;

struct InputState {
    line: Vec<String>,
    word: String,
    part: String,
    bpart: String,
    stdin: Stdr,
    stdout: Stdw,
    stderr: Stdw
}

impl InputState {
    fn new() -> InputState {
        InputState {
            line: Vec::<String>::new(),
            word: String::new(),
            part: String::new(),
            bpart: String::new(),
            stdin: io::stdin(),
            stdout: io::stdout(),
            stderr: io::stderr()
        }
    }

    fn clear(&mut self) {
        self.word.clear();
        self.line.clear();
        self.part.clear();
        self.bpart.clear();
    }

    fn flush(&mut self) {
        self.stdout.flush().unwrap();
    }

    fn outw(&mut self, msg:&str) {
        self.stdout.write_str(msg).unwrap();
    }

    fn outc(&mut self, ch:char) {
        self.stdout.write_char(ch).unwrap();
    }

    fn err(&mut self, msg:&str) {
        self.stderr.write_str(msg).unwrap();
    }
}

#[allow(unused_variables)]
unsafe extern fn handle_sigint(signum:c_int, siginfo:*const SigInfo, context:size_t) {
    if instate_location == 0 {
        panic!("Shell state location uninitialized");
    }
    print!("^C\n");
    // Below: completely disregarding everything Rust stands for
    let state:&mut InputState = mem::transmute(instate_location);
    state.clear();
    io::stdio::flush();
}

fn set_instate_location(state:&mut InputState) {
    unsafe {
        if instate_location != 0 {
            panic!("Tried to set shell state location twice");
        }
        instate_location = mem::transmute::<&mut InputState, size_t>(state);
    }
}

fn prepare_terminal(tios:&mut Termios) {
    tios.ldisable(ICANON);
    tios.ldisable(ECHO);
}

fn update_terminal(tios:Termios, state:&mut InputState) -> bool {
    if !tios.set() {
        state.err("Warning: Could not set terminal mode\n");
        return false;
    }
    return true;
}

fn handle_escape(state:&mut InputState) {
    // Handle an ANSI escape sequence
    if state.stdin.read_char() != Ok(ANSI) {
        return;
    }
    match state.stdin.read_char() {
        Err(_) => return,
        Ok('D') => {
            // left
            match state.word.pop() {
                Some(c) => {
                    state.bpart.clear();
                    state.part.push(c);
                    cursor_left(state);
                },
                None => match state.line.pop() {
                    Some(s) => {
                        state.bpart.clear();
                        state.part.push(' ');
                        state.word.clear();
                        state.word.push_str(s.as_slice());
                        cursor_left(state);
                    }
                    None => return
                }
            }
        },
        Ok('C') => {
            // right
            match state.part.pop() {
                Some(' ') => {
                    state.bpart.clear();
                    state.line.push(state.word.clone());
                    state.word.clear();
                    cursor_right(state);
                },
                Some(c) => {
                    state.bpart.clear();
                    state.word.push(c);
                    cursor_right(state);
                },
                None => return
            }
        },
        Ok(_) => return
    }
}

fn cursor_left(state:&mut InputState) {
    state.outc(DEL);
}

fn cursor_right(state:&mut InputState) {
    state.outw(CRSR_RIGHT);
}

fn draw_part(state:&mut InputState) {
    // quick out if part is empty
    if state.part.is_empty() {
        return;
    }
    if state.bpart.is_empty() {
        // only calculate bpart when it needs to be recalculated
        let mut cpart = state.part.clone();
        loop {
            match cpart.pop() {
                Some(c) => state.bpart.push(c),
                None => break
            }
        }
    }
    let c = state.bpart.clone();
    state.outw(c.as_slice());
}

// work around lack of DST
fn build_string(ch:char, count:uint) -> String {
    let mut s = String::new();
    let mut i = 0u;
    loop {
        if i == count {
            return s;
        }
        s.push(ch);
        i += 1;
    }
}

fn cursors_left(by:uint, state:&mut InputState) {
    // move back by a given number of characters
    state.outw(build_string(DEL, by).as_slice());
}

fn idraw_part(state:&mut InputState) {
    // in-place draw of the line part
    draw_part(state);
    cursors_left(state.part.len(), state);
}

fn prepare_signals(state:&mut InputState) {
    let mut sa = SigAction {
        handler: handle_sigint,
        mask: [0; SIGSET_NWORDS],
        flags: SA_RESTART | SA_SIGINFO
    };
    unsafe {
        if sigfillset(&mut sa.mask) != 0 {
            state.err("Warning: could not fill mask set for SIGINT handler\n");
        }
        if sigaction(SIGINT, &sa, ptr::null_mut::<SigAction>()) != 0 {
            state.err("Warning: could not set handler for SIGINT\n");
        }
    }
}

fn run_command(state:&mut InputState) {
    let mut process = Command::new(&state.line[0]);
    process.args(state.line.slice_from(1));
    process.stdout(StdioContainer::InheritFd(STDOUT));
    process.stdin(StdioContainer::InheritFd(STDIN));
    process.stderr(StdioContainer::InheritFd(STDERR));
    let mut child = match process.spawn() {
        Err(e) => {
            let s = state.line[0].clone();
            state.err(format!("Couldn't spawn {}: {}\n", s, e).as_slice());
            return;
        },
        Ok(child) => child
    };
    match child.wait() {
        Err(e) => {
            state.err(format!("Couldn't wait for child to exit: {}\n", e.desc).as_slice());
        },
        Ok(_) => {
            // nothing
        }
    };
}

fn main() {
    let mut state = &mut InputState::new();
    prepare_signals(state);
    let mut tios = Termios::new();
    tios.get();
    let old_tios = tios.clone();
    prepare_terminal(&mut tios);
    update_terminal(tios, state);

    set_instate_location(state);
    loop {
        // Note: in non-canonical mode
        match state.stdin.read_char() {
            Ok(EOF) => {
                if state.line.is_empty() && state.word.is_empty() {
                    break;
                }
            },
            Ok(NL) => {
                // start command output on next line
                state.outc(NL);
                // push any remaining word onto the line
                if !state.word.is_empty() {
                    state.line.push(state.word.clone());
                }
                // run command if one was specified
                if !state.line.is_empty() {
                    // run command
                    update_terminal(old_tios, state);
                    run_command(state);
                    update_terminal(tios, state);
                }
                // clear the state
                state.clear();
            },
            Ok(DEL) => {
                if state.word.is_empty() {
                    state.word = match state.line.pop() {
                        Some(s) => s,
                        None => continue
                    };
                } else {
                    state.word.pop();
                }
                cursor_left(state);
                draw_part(state);
                state.outc(SPC);
                cursors_left(state.part.len() + 1, state);
            },
            Ok(ESC) => handle_escape(state),
            Ok(SPC) => {
                state.line.push(state.word.clone());
                state.word.clear();
                state.outc(SPC);
                idraw_part(state);
            },
            Ok(c) => {
                state.word.push(c);
                state.outc(c);
                idraw_part(state);
            },
            Err(e) => {
                state.err(format!("Error: {}\n", e).as_slice());
                break;
            }
        }
        // flush output
        state.flush();
    }
    state.outw("Exiting\n");
    // restore old term state
    update_terminal(old_tios, state);
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, All Contributors (see CONTRIBUTORS file)
//
// 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(slice_patterns, advanced_slice_patterns)]
#![cfg_attr(test, feature(test))]

#![cfg_attr(not(target_os = "windows"), feature(alloc, heap_api))]
#![feature(alloc)]

include!("crates.rs");

extern crate num_cpus;
extern crate log4rs;

pub mod script;
pub mod server;
pub mod timestamp;
pub mod pubsub;
pub mod storage;

use std::fs;
use std::fs::{OpenOptions};
use std::path::PathBuf;
use std::thread;
use std::sync::Arc;
use std::sync::mpsc;
use memmap::{Mmap, Protection};

lazy_static! {
 static ref ENVIRONMENT: lmdb::Environment = {
    let _ = config::set_default("storage.path", "pumpkin.db");
    let storage_path = config::get_str("storage.path").unwrap().into_owned();
    fs::create_dir_all(storage_path.as_str()).expect("can't create directory");
    let map_size = config::get_int("storage.mapsize");
    storage::create_environment(storage_path, map_size)
 };
}

/// Accepts storage path, filename and length and prepares the file. It is important that the length
/// is the total length of the memory mapped file, otherwise the application _will segfault_ when
/// trying to read those sections later. There is no way to handle that.
fn prepare_mmap(storage_path: &str, filename: &str, length: u64) -> Mmap {
    let mut scratchpad_pathbuf = PathBuf::from(storage_path);
    scratchpad_pathbuf.push(filename);
    scratchpad_pathbuf.set_extension("dat");
    let scratchpad_path = scratchpad_pathbuf.as_path();
    let scratchpad_file = OpenOptions::new().create_new(false).write(true)
        .open(scratchpad_path).expect("Could not open or create scratchpad");
    let _ = scratchpad_file.set_len(length);
    Mmap::open_path(scratchpad_path, Protection::ReadWrite)
        .expect("Could not open scratchpad")
}

fn main() {
    let _ = config::merge(config::Environment::new("pumpkindb"));
    let _ = config::merge(config::File::new("pumpkindb.toml", config::FileFormat::Toml));
    let _ = config::set_default("server.port", 9981);
    let _ = config::set_default("storage.path", "pumpkin.db");
    let storage_path = config::get_str("storage.path").unwrap().into_owned();
    fs::create_dir_all(storage_path.as_str()).expect("can't create directory");

    // Initialize Mmap
    let scratchpad = prepare_mmap(storage_path.as_str(), "scratchpad", 20);
    let mut hlc_state = scratchpad.into_view_sync();
    hlc_state.restrict(0, 20).expect("Could not prepare HLC state");

    // Initialize logging
    let log_config = config::get_str("logging.config");
    let mut log_configured = false;
    if log_config.is_some() {
        let log_file_path = log_config.unwrap().into_owned();
        if fs::metadata(&log_file_path).is_ok() {
            log4rs::init_file(&log_file_path, Default::default()).unwrap();
            log_configured = true;
        } else {
            println!("{} not found", &log_file_path);
        }
    }

    if !log_configured {
        let appender = log4rs::config::Appender::builder()
            .build("console",
                   Box::new(log4rs::append::console::ConsoleAppender::builder().build()));
        let root = log4rs::config::Root::builder().appender("console").build(log::LogLevelFilter::Info);
        let _ = log4rs::init_config(log4rs::config::Config::builder().appender(appender).build(root).unwrap());
        warn!("No logging configuration specified, switching to console logging");
    }

    info!("Starting up");

    let mut senders = Vec::new();

    let mut publisher = pubsub::Publisher::new();
    let publisher_accessor = publisher.accessor();
    let _ = thread::spawn(move || publisher.run());
    let storage = Arc::new(storage::Storage::new(&ENVIRONMENT));
    let timestamp = Arc::new(timestamp::Timestamp::new(Some(hlc_state)));

    for i in 0..num_cpus::get() {
        info!("Starting scheduler on core {}.", i);
        let (sender, receiver) = mpsc::sync_channel(0);
        let publisher_clone = publisher_accessor.clone();
        let storage_clone = storage.clone();
        let timestamp_clone = timestamp.clone();
        thread::spawn(move || {
            let mut scheduler = script::Scheduler::new(
                &storage_clone,
                publisher_clone,
                timestamp_clone,
                sender,
            );
            scheduler.run()
        });
        let sender = receiver.recv().expect("Could not receive sender");
        senders.push(sender)
    }

    server::run(config::get_int("server.port").unwrap(), senders, publisher_accessor.clone());

}
<commit_msg>Problem: PDB crashes if scratchpad doesn’t exist<commit_after>// Copyright (c) 2017, All Contributors (see CONTRIBUTORS file)
//
// 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(slice_patterns, advanced_slice_patterns)]
#![cfg_attr(test, feature(test))]

#![cfg_attr(not(target_os = "windows"), feature(alloc, heap_api))]
#![feature(alloc)]

include!("crates.rs");

extern crate num_cpus;
extern crate log4rs;

pub mod script;
pub mod server;
pub mod timestamp;
pub mod pubsub;
pub mod storage;

use std::fs;
use std::fs::{OpenOptions};
use std::path::PathBuf;
use std::thread;
use std::sync::Arc;
use std::sync::mpsc;
use memmap::{Mmap, Protection};

lazy_static! {
 static ref ENVIRONMENT: lmdb::Environment = {
    let _ = config::set_default("storage.path", "pumpkin.db");
    let storage_path = config::get_str("storage.path").unwrap().into_owned();
    fs::create_dir_all(storage_path.as_str()).expect("can't create directory");
    let map_size = config::get_int("storage.mapsize");
    storage::create_environment(storage_path, map_size)
 };
}

/// Accepts storage path, filename and length and prepares the file. It is important that the length
/// is the total length of the memory mapped file, otherwise the application _will segfault_ when
/// trying to read those sections later. There is no way to handle that.
fn prepare_mmap(storage_path: &str, filename: &str, length: u64) -> Mmap {
    let mut scratchpad_pathbuf = PathBuf::from(storage_path);
    scratchpad_pathbuf.push(filename);
    scratchpad_pathbuf.set_extension("dat");
    let scratchpad_path = scratchpad_pathbuf.as_path();
    let scratchpad_file = OpenOptions::new().create(true).write(true)
        .open(scratchpad_path).expect("Could not open or create scratchpad");
    let _ = scratchpad_file.set_len(length);
    Mmap::open_path(scratchpad_path, Protection::ReadWrite)
        .expect("Could not open scratchpad")
}

fn main() {
    let _ = config::merge(config::Environment::new("pumpkindb"));
    let _ = config::merge(config::File::new("pumpkindb.toml", config::FileFormat::Toml));
    let _ = config::set_default("server.port", 9981);
    let _ = config::set_default("storage.path", "pumpkin.db");
    let storage_path = config::get_str("storage.path").unwrap().into_owned();
    fs::create_dir_all(storage_path.as_str()).expect("can't create directory");

    // Initialize Mmap
    let scratchpad = prepare_mmap(storage_path.as_str(), "scratchpad", 20);
    let mut hlc_state = scratchpad.into_view_sync();
    hlc_state.restrict(0, 20).expect("Could not prepare HLC state");

    // Initialize logging
    let log_config = config::get_str("logging.config");
    let mut log_configured = false;
    if log_config.is_some() {
        let log_file_path = log_config.unwrap().into_owned();
        if fs::metadata(&log_file_path).is_ok() {
            log4rs::init_file(&log_file_path, Default::default()).unwrap();
            log_configured = true;
        } else {
            println!("{} not found", &log_file_path);
        }
    }

    if !log_configured {
        let appender = log4rs::config::Appender::builder()
            .build("console",
                   Box::new(log4rs::append::console::ConsoleAppender::builder().build()));
        let root = log4rs::config::Root::builder().appender("console").build(log::LogLevelFilter::Info);
        let _ = log4rs::init_config(log4rs::config::Config::builder().appender(appender).build(root).unwrap());
        warn!("No logging configuration specified, switching to console logging");
    }

    info!("Starting up");

    let mut senders = Vec::new();

    let mut publisher = pubsub::Publisher::new();
    let publisher_accessor = publisher.accessor();
    let _ = thread::spawn(move || publisher.run());
    let storage = Arc::new(storage::Storage::new(&ENVIRONMENT));
    let timestamp = Arc::new(timestamp::Timestamp::new(Some(hlc_state)));

    for i in 0..num_cpus::get() {
        info!("Starting scheduler on core {}.", i);
        let (sender, receiver) = mpsc::sync_channel(0);
        let publisher_clone = publisher_accessor.clone();
        let storage_clone = storage.clone();
        let timestamp_clone = timestamp.clone();
        thread::spawn(move || {
            let mut scheduler = script::Scheduler::new(
                &storage_clone,
                publisher_clone,
                timestamp_clone,
                sender,
            );
            scheduler.run()
        });
        let sender = receiver.recv().expect("Could not receive sender");
        senders.push(sender)
    }

    server::run(config::get_int("server.port").unwrap(), senders, publisher_accessor.clone());

}
<|endoftext|>


<commit_before>extern crate bytes;
extern crate env_logger;
extern crate mio;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;
extern crate failure;
extern crate slab;

use clap::{App, Arg};
use failure::Error;
use failure::ResultExt;
use mio::net::{TcpListener, TcpStream};
use slab::Slab;
use std::collections::VecDeque;
use std::io::{Read, Write};
use std::ops;

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

trait RuleHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready);
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error>;
    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error>;
    fn should_close(&self) -> bool;
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::Poll, listener: TcpListener) -> Result<(), Error> {
        let token = {
            let e = self.connections.vacant_entry();
            let token = mio::Token(e.key());
            let l = EventHandler::Listener(Listener::new(listener, token));
            e.insert(l);
            token
        };
        &self.connections[token.into()].register(event_loop, token)?;
        Ok(())
    }

    fn process_action(&mut self, src: usize, msg: MiChatCommand, event_loop: &mut mio::Poll) {
        trace!("{:p}; from {:?}; got {:?}", self, src, msg);
        match msg {
            MiChatCommand::Broadcast(s) => self.process_broadcast(s),
            MiChatCommand::NewConnection(socket) => self.process_new_connection(event_loop, socket),
        }
    }

    fn process_broadcast(&mut self, s: String) {
        for c in self.client_connections_mut() {
            c.enqueue(&s);
        }
    }

    fn process_new_connection(&mut self, event_loop: &mut mio::Poll, socket: mio::tcp::TcpStream) {
        let token = {
            let e = self.connections.vacant_entry();
            let token = mio::Token(e.key());
            let l = EventHandler::Conn(Connection::new(socket, token));
            e.insert(l);
            token
        };
        &self.connections[token.into()].register(event_loop, token);
    }

    fn client_connections_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut Connection> + 'a {
        self.connections.iter_mut().filter_map(|(_, it)| match it {
            &mut EventHandler::Conn(ref mut c) => Some(c),
            _ => None,
        })
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::Ready,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    write_buf: Vec<u8>,
}

impl RuleHandler for Connection {
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error> {
        event_loop
            .register(
                &self.socket,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).context("event loop register")?;
        Ok(())
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready) {
        self.sock_status.insert(events);
        info!(
            "Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.socket.peer_addr(),
            events,
            self.sock_status
        );
    }

    // actions are processed here on down.

    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error> {
        if self.sock_status.is_readable() {
            self.sock_status.remove(mio::Ready::readable());
            self.read()?;
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.sock_status.remove(mio::Ready::writable());
            self.write()?;
        }

        if !self.should_close() {
            self.reregister(event_loop)?;
        }

        Ok(())
    }

    fn should_close(&self) -> bool {
        self.read_eof && self.write_buf.is_empty()
    }
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::Ready::empty(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
        }
    }
    fn process_buffer(&mut self, to_parent: &mut FnMut(MiChatCommand)) {
        let mut prev = 0;
        info!(
            "{:?}: Read buffer: {:?}",
            self.socket.peer_addr(),
            self.read_buf
        );
        for n in self.read_buf.iter().enumerate().filter_map(|(i, e)| {
            if *e == '\n' as u8 {
                Some(i)
            } else {
                None
            }
        }) {
            info!(
                "{:?}: Pos: {:?}; chunk: {:?}",
                self.socket.peer_addr(),
                n,
                &self.read_buf[prev..n]
            );
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent(cmd);
            prev = n + 1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!(
            "{:?}: read Remainder: {}",
            self.socket.peer_addr(),
            remainder.len()
        );
        self.read_buf = remainder;
    }

    fn read(&mut self) -> Result<(), Error> {
        let mut abuf = vec![0; 1024];
        match self.socket.read(&mut abuf).context("Reading socket")? {
            0 => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true;
                Ok(())
            }
            n => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(&abuf[0..n]);
                Ok(())
            }
        }
    }

    fn write(&mut self) -> Result<(), Error> {
        let n = self
            .socket
            .write(&mut self.write_buf)
            .context("writing buffer")?;
        info!(
            "{:?}: Wrote {} of {} in buffer",
            self.socket.peer_addr(),
            n,
            self.write_buf.len()
        );
        self.write_buf = self.write_buf[n..].to_vec();
        info!(
            "{:?}: Now {:?}b",
            self.socket.peer_addr(),
            self.write_buf.len()
        );
        Ok(())
    }

    fn reregister(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        let mut flags = mio::Ready::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::Ready::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop
            .reregister(&self.socket, self.token, flags, mio::PollOpt::oneshot())
            .context("EventLoop#reregister")?;

        Ok(())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::Ready,
    token: mio::Token,
}

impl Listener {
    fn new(listener: TcpListener, token: mio::Token) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::Ready::empty(),
            token: token,
        }
    }
    fn reregister(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error> {
        event_loop
            .reregister(
                &self.listener,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).context("Register listener")?;
        Ok(())
    }
}

impl RuleHandler for Listener {
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error> {
        event_loop
            .register(
                &self.listener,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).context("Register listener")?;
        Ok(())
    }

    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!(
            "Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.listener.local_addr(),
            events,
            self.sock_status
        );
    }
    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error> {
        if self.sock_status.is_readable() {
            info!("the listener socket is ready to accept a connection");
            self.sock_status.remove(mio::Ready::readable());
            let (socket, _client) = self.listener.accept().context("Accept listening socket")?;
            let cmd = MiChatCommand::NewConnection(socket);
            to_parent(cmd);

            self.reregister(event_loop, self.token)?;
        }
        Ok(())
    }

    fn should_close(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener(Listener),
    Conn(Connection),
}

impl ops::Deref for EventHandler {
    type Target = RuleHandler;
    fn deref(&self) -> &Self::Target {
        match self {
            &EventHandler::Conn(ref conn) => conn,
            &EventHandler::Listener(ref listener) => listener,
        }
    }
}
impl ops::DerefMut for EventHandler {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn,
            &mut EventHandler::Listener(ref mut listener) => listener,
        }
    }
}

impl MiChat {
    fn ready(&mut self, event_loop: &mut mio::Poll, token: mio::Token, events: mio::Ready) {
        info!("{:?}: {:?}", token, events);
        self.connections[token.into()].handle_event(event_loop, events);
    }

    fn process(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        let mut parent_actions = VecDeque::new();
        let mut failed = Vec::new();
        loop {
            for (idx, conn) in self.connections.iter_mut() {
                if let Err(e) = conn.process_rules(event_loop, &mut |action| {
                    parent_actions.push_back((idx, action))
                }) {
                    failed.push(idx);
                    error!("Error on connection {:?}; Dropping: {:?}", idx, e);
                }

                if conn.should_close() {
                    failed.push(idx);
                    info!("Removing; {:?}", idx);
                }
            }
            for idx in failed.drain(..) {
                self.connections.remove(idx);
            }
            // Anything left to process?
            if parent_actions.is_empty() {
                break;
            }

            for (src, action) in parent_actions.drain(..) {
                self.process_action(src, action, event_loop);
            }
        }

        Ok(())
    }

    fn run(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        loop {
            self.run_once(event_loop)?;
        }
    }
    fn run_once(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        let mut events = mio::Events::with_capacity(1024);
        event_loop.poll(&mut events, None)?;
        for ev in &events {
            self.ready(event_loop, ev.token(), ev.readiness())
        }

        self.process(event_loop)
    }
}

fn main() {
    env_logger::init();
    let matches = App::new("michat")
        .arg(
            Arg::with_name("bind")
                .short("l")
                .takes_value(true)
                .required(true),
        ).get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::Poll::new().expect("Create event loop");
    let mut service = MiChat::new();

    service
        .listen(&mut event_loop, listener)
        .expect("Starting listener");
    info!("running michat listener at: {:?}", address);
    service.run(&mut event_loop).expect("run");
}
<commit_msg>Deregister from poller.<commit_after>extern crate bytes;
extern crate env_logger;
extern crate mio;
#[macro_use]
extern crate log;
#[macro_use]
extern crate clap;
extern crate failure;
extern crate slab;

use clap::{App, Arg};
use failure::Error;
use failure::ResultExt;
use mio::net::{TcpListener, TcpStream};
use slab::Slab;
use std::collections::VecDeque;
use std::io::{Read, Write};
use std::ops;

#[derive(Debug)]
enum MiChatCommand {
    Broadcast(String),
    NewConnection(TcpStream),
}

struct MiChat {
    connections: Slab<EventHandler>,
}

trait RuleHandler {
    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready);
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error>;
    fn deregister(&self, event_loop: &mut mio::Poll) -> Result<(), Error>;
    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error>;
    fn should_close(&self) -> bool;
}

impl MiChat {
    fn new() -> MiChat {
        MiChat {
            connections: Slab::new(),
        }
    }

    fn listen(&mut self, event_loop: &mut mio::Poll, listener: TcpListener) -> Result<(), Error> {
        let token = {
            let e = self.connections.vacant_entry();
            let token = mio::Token(e.key());
            let l = EventHandler::Listener(Listener::new(listener, token));
            e.insert(l);
            token
        };
        &self.connections[token.into()].register(event_loop, token)?;
        Ok(())
    }

    fn process_action(&mut self, src: usize, msg: MiChatCommand, event_loop: &mut mio::Poll) {
        trace!("{:p}; from {:?}; got {:?}", self, src, msg);
        match msg {
            MiChatCommand::Broadcast(s) => self.process_broadcast(s),
            MiChatCommand::NewConnection(socket) => self.process_new_connection(event_loop, socket),
        }
    }

    fn process_broadcast(&mut self, s: String) {
        for c in self.client_connections_mut() {
            c.enqueue(&s);
        }
    }

    fn process_new_connection(&mut self, event_loop: &mut mio::Poll, socket: mio::tcp::TcpStream) {
        let token = {
            let e = self.connections.vacant_entry();
            let token = mio::Token(e.key());
            let l = EventHandler::Conn(Connection::new(socket, token));
            e.insert(l);
            token
        };
        &self.connections[token.into()].register(event_loop, token);
    }

    fn client_connections_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut Connection> + 'a {
        self.connections.iter_mut().filter_map(|(_, it)| match it {
            &mut EventHandler::Conn(ref mut c) => Some(c),
            _ => None,
        })
    }
}

#[derive(Debug)]
struct Connection {
    socket: TcpStream,
    sock_status: mio::Ready,
    token: mio::Token,
    read_buf: Vec<u8>,
    read_eof: bool,
    write_buf: Vec<u8>,
}

impl RuleHandler for Connection {
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error> {
        event_loop
            .register(
                &self.socket,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).context("event loop register")?;
        Ok(())
    }
    fn deregister(&self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        event_loop.deregister(&self.socket)?;
        Ok(())
    }

    // Event updates arrive here
    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready) {
        self.sock_status.insert(events);
        info!(
            "Connection::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.socket.peer_addr(),
            events,
            self.sock_status
        );
    }

    // actions are processed here on down.

    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error> {
        if self.sock_status.is_readable() {
            self.sock_status.remove(mio::Ready::readable());
            self.read()?;
        }

        self.process_buffer(to_parent);

        if self.sock_status.is_writable() {
            self.sock_status.remove(mio::Ready::writable());
            self.write()?;
        }

        if !self.should_close() {
            self.reregister(event_loop)?;
        }

        Ok(())
    }

    fn should_close(&self) -> bool {
        self.read_eof && self.write_buf.is_empty()
    }
}

impl Connection {
    fn new(socket: TcpStream, token: mio::Token) -> Connection {
        Connection {
            socket: socket,
            sock_status: mio::Ready::empty(),
            token: token,
            read_buf: Vec::with_capacity(1024),
            write_buf: Vec::new(),
            read_eof: false,
        }
    }
    fn process_buffer(&mut self, to_parent: &mut FnMut(MiChatCommand)) {
        let mut prev = 0;
        info!(
            "{:?}: Read buffer: {:?}",
            self.socket.peer_addr(),
            self.read_buf
        );
        for n in self.read_buf.iter().enumerate().filter_map(|(i, e)| {
            if *e == '\n' as u8 {
                Some(i)
            } else {
                None
            }
        }) {
            info!(
                "{:?}: Pos: {:?}; chunk: {:?}",
                self.socket.peer_addr(),
                n,
                &self.read_buf[prev..n]
            );
            let s = String::from_utf8_lossy(&self.read_buf[prev..n]).to_string();
            let cmd = MiChatCommand::Broadcast(s);
            info!("Send! {:?}", cmd);
            to_parent(cmd);
            prev = n + 1;
        }
        let remainder = self.read_buf[prev..].to_vec();
        info!(
            "{:?}: read Remainder: {}",
            self.socket.peer_addr(),
            remainder.len()
        );
        self.read_buf = remainder;
    }

    fn read(&mut self) -> Result<(), Error> {
        let mut abuf = vec![0; 1024];
        match self.socket.read(&mut abuf).context("Reading socket")? {
            0 => {
                info!("{:?}: EOF!", self.socket.peer_addr());
                self.read_eof = true;
                Ok(())
            }
            n => {
                info!("{:?}: Read {}bytes", self.socket.peer_addr(), n);
                self.read_buf.extend(&abuf[0..n]);
                Ok(())
            }
        }
    }

    fn write(&mut self) -> Result<(), Error> {
        let n = self
            .socket
            .write(&mut self.write_buf)
            .context("writing buffer")?;
        info!(
            "{:?}: Wrote {} of {} in buffer",
            self.socket.peer_addr(),
            n,
            self.write_buf.len()
        );
        self.write_buf = self.write_buf[n..].to_vec();
        info!(
            "{:?}: Now {:?}b",
            self.socket.peer_addr(),
            self.write_buf.len()
        );
        Ok(())
    }

    fn reregister(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        let mut flags = mio::Ready::readable();
        if !self.write_buf.is_empty() {
            flags.insert(mio::Ready::writable());
        }
        info!("Registering {:?} with {:?}", self, flags);

        event_loop
            .reregister(&self.socket, self.token, flags, mio::PollOpt::oneshot())
            .context("EventLoop#reregister")?;

        Ok(())
    }

    fn enqueue(&mut self, s: &str) {
        self.write_buf.extend(s.as_bytes());
        self.write_buf.push('\n' as u8);
    }
}

#[derive(Debug)]
struct Listener {
    listener: TcpListener,
    sock_status: mio::Ready,
    token: mio::Token,
}

impl Listener {
    fn new(listener: TcpListener, token: mio::Token) -> Listener {
        Listener {
            listener: listener,
            sock_status: mio::Ready::empty(),
            token: token,
        }
    }
    fn reregister(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error> {
        event_loop
            .reregister(
                &self.listener,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).context("Register listener")?;
        Ok(())
    }
}

impl RuleHandler for Listener {
    fn register(&self, event_loop: &mut mio::Poll, token: mio::Token) -> Result<(), Error> {
        event_loop
            .register(
                &self.listener,
                token,
                mio::Ready::readable(),
                mio::PollOpt::edge() | mio::PollOpt::oneshot(),
            ).context("Register listener")?;
        Ok(())
    }
    fn deregister(&self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        event_loop.deregister(&self.listener)?;
        Ok(())
    }

    fn handle_event(&mut self, _event_loop: &mut mio::Poll, events: mio::Ready) {
        assert!(events.is_readable());
        self.sock_status.insert(events);
        info!(
            "Listener::handle_event: {:?}; this time: {:?}; now: {:?}",
            self.listener.local_addr(),
            events,
            self.sock_status
        );
    }
    fn process_rules(
        &mut self,
        event_loop: &mut mio::Poll,
        to_parent: &mut FnMut(MiChatCommand),
    ) -> Result<(), Error> {
        if self.sock_status.is_readable() {
            info!("the listener socket is ready to accept a connection");
            self.sock_status.remove(mio::Ready::readable());
            let (socket, _client) = self.listener.accept().context("Accept listening socket")?;
            let cmd = MiChatCommand::NewConnection(socket);
            to_parent(cmd);

            self.reregister(event_loop, self.token)?;
        }
        Ok(())
    }

    fn should_close(&self) -> bool {
        false
    }
}

#[derive(Debug)]
enum EventHandler {
    Listener(Listener),
    Conn(Connection),
}

impl ops::Deref for EventHandler {
    type Target = RuleHandler;
    fn deref(&self) -> &Self::Target {
        match self {
            &EventHandler::Conn(ref conn) => conn,
            &EventHandler::Listener(ref listener) => listener,
        }
    }
}
impl ops::DerefMut for EventHandler {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self {
            &mut EventHandler::Conn(ref mut conn) => conn,
            &mut EventHandler::Listener(ref mut listener) => listener,
        }
    }
}

impl MiChat {
    fn ready(&mut self, event_loop: &mut mio::Poll, token: mio::Token, events: mio::Ready) {
        info!("{:?}: {:?}", token, events);
        self.connections[token.into()].handle_event(event_loop, events);
    }

    fn process(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        let mut parent_actions = VecDeque::new();
        let mut failed = Vec::new();
        loop {
            for (idx, conn) in self.connections.iter_mut() {
                if let Err(e) = conn.process_rules(event_loop, &mut |action| {
                    parent_actions.push_back((idx, action))
                }) {
                    failed.push(idx);
                    error!("Error on connection {:?}; Dropping: {:?}", idx, e);
                }

                if conn.should_close() {
                    failed.push(idx);
                    info!("Removing; {:?}", idx);
                }
            }
            for idx in failed.drain(..) {
                let conn = self.connections.remove(idx);
                conn.deregister(event_loop)?;
            }
            // Anything left to process?
            if parent_actions.is_empty() {
                break;
            }

            for (src, action) in parent_actions.drain(..) {
                self.process_action(src, action, event_loop);
            }
        }

        Ok(())
    }

    fn run(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        loop {
            self.run_once(event_loop)?;
        }
    }
    fn run_once(&mut self, event_loop: &mut mio::Poll) -> Result<(), Error> {
        let mut events = mio::Events::with_capacity(1024);
        event_loop.poll(&mut events, None)?;
        for ev in &events {
            self.ready(event_loop, ev.token(), ev.readiness())
        }

        self.process(event_loop)
    }
}

fn main() {
    env_logger::init();
    let matches = App::new("michat")
        .arg(
            Arg::with_name("bind")
                .short("l")
                .takes_value(true)
                .required(true),
        ).get_matches();

    let address = value_t_or_exit!(matches.value_of("bind"), std::net::SocketAddr);
    let listener = TcpListener::bind(&address).expect("bind");

    let mut event_loop = mio::Poll::new().expect("Create event loop");
    let mut service = MiChat::new();

    service
        .listen(&mut event_loop, listener)
        .expect("Starting listener");
    info!("running michat listener at: {:?}", address);
    service.run(&mut event_loop).expect("run");
}
<|endoftext|>


<commit_before>extern crate docopt;
extern crate i2cdev;
#[macro_use]
extern crate serde_derive;

use docopt::Docopt;

use std::{thread, time};

#[cfg(target_os = "linux")]
use i2cdev::linux::*;

mod utils;

mod bme280;
use bme280::*;

const USAGE: &'static str = "
Reading BME280 sensor value

Usage:
  bme280 <device> [--address <addr>]
  bme280 (-h | --help)
  bme280 (-v | --version)

Options:
  -h --help    Show this help text.
  --address <addr>     I2C device address [default: 119] (=0x77)
  -v --version    Show version.
";

const BME280_DEFAULT_ADDRESS: u16 = 0x76;

#[derive(Debug, Deserialize)]
struct Args {
    arg_device: String,
    flag_address: Option<u16>,
    flag_version: bool,
}

#[cfg(not(target_os = "linux"))]
fn main() {
    println!("This program can run only on Linux")
}

#[cfg(target_os = "linux")]
fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.deserialize()).unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("bme280 {}", env!("CARGO_PKG_VERSION"));
        return;
    }

    let address = args.flag_address.unwrap_or(BME280_DEFAULT_ADDRESS);
    let config: Config = Config {
        mode: Mode::Force,
        oversampling_temperature: 1,
        oversampling_pressure: 1,
        oversampling_humidity: 1,
        standby_time: 5,
        iir_filter: 0,
        spi3w_enabled: false,
    };

    let dev = LinuxI2CDevice::new(args.arg_device, address).unwrap();
    let mut bme280 = BME280::new(dev, config).unwrap();

    // Sleep while measureing finishes.
    let ten_millis = time::Duration::from_millis(10);
    thread::sleep(ten_millis);

    let raw_pressure = bme280.raw_pressure().unwrap();
    let raw_temperature = bme280.raw_temperature().unwrap();
    let raw_humidity = bme280.raw_humidity().unwrap();
    println!("raw_temp = {}, raw_press = {}, raw_humid = {}",
             raw_temperature,
             raw_pressure,
             raw_humidity);

    println!("temperature = {}", bme280.temperature().unwrap());
    println!("humidity = {}", bme280.humidity().unwrap());
}
<commit_msg>Add options to change output values.<commit_after>extern crate docopt;
extern crate i2cdev;
#[macro_use]
extern crate serde_derive;

use docopt::Docopt;

use std::{thread, time};

#[cfg(target_os = "linux")]
use i2cdev::linux::*;

mod utils;

mod bme280;
use bme280::*;

const USAGE: &'static str = "
Reading BME280 sensor value

Usage:
  bme280 <device> [--address <addr>] [--temperature] [--humidity]
  bme280 (-h | --help)
  bme280 (-v | --version)

Options:
  -h --help    Show this help text.
  --address <addr>     I2C device address [default: 119] (=0x77)
  --temperature    Show temperature.
  --humidity    Show humidity
  -v --version    Show version.
";

const BME280_DEFAULT_ADDRESS: u16 = 0x76;

#[derive(Debug, Deserialize)]
struct Args {
    arg_device: String,
    flag_address: Option<u16>,
    flag_version: bool,
    flag_temperature: bool,
    flag_humidity: bool,
}

#[cfg(not(target_os = "linux"))]
fn main() {
    println!("This program can run only on Linux")
}

#[cfg(target_os = "linux")]
fn main() {
    let args: Args = Docopt::new(USAGE).and_then(|d| d.deserialize()).unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("bme280 {}", env!("CARGO_PKG_VERSION"));
        return;
    }

    let address = args.flag_address.unwrap_or(BME280_DEFAULT_ADDRESS);
    let config: Config = Config {
        mode: Mode::Force,
        oversampling_temperature: 1,
        oversampling_pressure: 1,
        oversampling_humidity: 2,
        standby_time: 5,
        iir_filter: 0,
        spi3w_enabled: false,
    };

    let dev = LinuxI2CDevice::new(args.arg_device, address).unwrap();
    let mut bme280 = BME280::new(dev, config).unwrap();

    // Sleep while measureing finishes.
    let millis = time::Duration::from_millis(100);
    thread::sleep(millis);

    if args.flag_temperature {
        println!("{}", bme280.temperature().unwrap());
    }
    if args.flag_humidity {
        println!("{:.2}", bme280.humidity().unwrap());
    }
}
<|endoftext|>


<commit_before>use std;
use env_logger;
use time;
use portaudio;
use rgsl;
#[cfg(feature="gnuplot")]
use gnuplot;

use thread;
use note;
use mvar;

fn consume<T>(_: T) {}

fn to_fft(mut samples: Vec<f64>) -> Result<Vec<f64>, String> {
  let len = samples.len();
  let r = rgsl::fft::real_radix2::transform(&mut samples, 1, len);
  if r != rgsl::Value::Success {
    return Err(format!("rgsl returned {:?}", r));
  }

  Ok(samples)
}

fn of_fft(mut fft: Vec<f64>) -> Result<Vec<f64>, String> {
  let len = fft.len();
  let r = rgsl::fft::real_radix2::inverse(&mut fft, 1, len);
  if r != rgsl::Value::Success {
    return Err(format!("rgsl returned {:?}", r));
  }

  Ok(fft)
}

fn string_err<Ok, Err: std::fmt::Display>(r: Result<Ok, Err>) -> Result<Ok, String> {
  r.map_err(|err| format!("{}", err))
}

fn sine_wave(note: note::T, sample_frequency: f64, start: u32, end: u32) -> Vec<f32> {
  let mut buf = Vec::new();

  let f = note.to_frequency();

  for t in start .. end {
    let t = t as f64 / sample_frequency;
    let sample = (t * f * 2.0 * std::f64::consts::PI).sin();
    let amplitude = 1.0;
    buf.push((sample * amplitude) as f32);
  }

  buf
}

fn with_play_channel<T, F>(sample_rate: f64, buf_size: u32, f: F) -> Result<T, String> where
  F: FnOnce(&mut portaudio::pa::Stream<f32, f32>) -> Result<T, String>,
{
  let mut stream: portaudio::pa::Stream<f32, f32> = portaudio::pa::Stream::new();
  try!(string_err(
    stream.open_default_blocking(
      sample_rate,
      buf_size,
      0,
      1,
      portaudio::pa::SampleFormat::Float32,
    )));
  try!(string_err(stream.start()));

  let r = try!(f(&mut stream));

  try!(string_err(stream.stop()));

  Ok(r)
}

fn play_note(note: note::T, secs: u64) -> Result<(), String> {
  info!("Writing..");

  let sample_rate = 44100 as f64;
  let buf_size = 1 << 10;

  try!(with_play_channel(
    sample_rate,
    buf_size,
    |stream| {
      let start_time = time::precise_time_ns();
      let mut i = 0;
      while time::precise_time_ns() <= start_time + secs*1_000_000_000 {
        let buf = sine_wave(note, sample_rate, i*buf_size, (i+1)*buf_size);
        assert!(buf.len() == buf_size as usize);
        try!(string_err(stream.write(buf, buf_size)));

        i = i + 1;
      }

      Ok(())
    }
  ));

  Ok(())
}

fn record(sample_frequency: f64, delta_t_ns: u64) -> Result<Vec<f64>, String> {
  let buf_size = 1 << 10;
  assert!(buf_size & (buf_size - 1) == 0);

  info!("Reading..");

  let mut stream: portaudio::pa::Stream<f32, f32> = portaudio::pa::Stream::new();
  try!(string_err(
    stream.open_default_blocking(
      sample_frequency,
      buf_size,
      1,
      0,
      portaudio::pa::SampleFormat::Float32,
    )));
  try!(string_err(stream.start()));

  let mut buf = Vec::new();
  let start_time = time::precise_time_ns();
  while time::precise_time_ns() <= start_time + delta_t_ns || (buf.len() & (buf.len() - 1)) != 0 {
    let new = try!(string_err(stream.read(buf_size)));
    assert!(new.len () == buf_size as usize);
    buf.extend(new.into_iter().map(|f| f as f64));
  }

  try!(string_err(stream.stop()));

  Ok(buf)
}

fn detect_frequency(samples: Vec<f64>, timestep: f64) -> Result<Option<f64>, String> {
  let num_samples = samples.len();
  let samples = try!(to_fft(samples));

  let buckets: Vec<_> =
    (0 .. 1 + num_samples / 2)
    .filter_map(|i| {
      let real = samples[i];
      let imag =
        if i % num_samples / 2 == 0 {
          0.0
        } else {
          samples[num_samples - i]
        };

      let noise_threshold = 0.0;
      let x = real*real + imag*imag - noise_threshold;
      if x > noise_threshold {
        let f = i as f64 / num_samples as f64 / timestep;
        Some((f, x))
      } else {
        None
      }
    })
    .collect();
  consume(samples);

  if buckets.is_empty() {
    return Ok(None)
  }

  #[cfg(feature="gnuplot")]
  {
    let x: Vec<_> = buckets.iter().map(|&(f, _)| f).collect();
    let y: Vec<_> = buckets.iter().map(|&(_, y)| y).collect();
    let mut fg = gnuplot::Figure::new();
    fg
      .axes2d()
      .boxes(
        &x,
        &y,
        &[],
      );
    fg.echo_to_file("/tmp/gnuplot.txt");
  }

  let (max_f, max_val) =
    buckets.iter()
      .fold((0.0, 0.0), |(max_f, max_val), &(f, amp)| {
        let val = amp;
        if val > max_val {
          (f, val)
        } else {
          (max_f, max_val)
        }
      });

  debug!("Max index is {}", max_f);
  debug!("Max value is {}", max_val);

  Ok(Some(max_f))
}

fn harmony(fft: &[f64], semitones_up: f64) -> Vec<f64> {
  let num_samples = fft.len();

  let shift_factor: f64 = (2.0 as f64).powf(-semitones_up / 12.0);

  let mut shifted_fft: Vec<_> = std::iter::repeat(0.0).take(num_samples).collect();
  for i in 1..(num_samples / 2) {
    let source_i = (i as f64 * shift_factor).round() as usize;
    shifted_fft[i] = shifted_fft[i] + fft[source_i];
    shifted_fft[num_samples - i] = shifted_fft[num_samples - i] + fft[num_samples - source_i];
  }

  // This might be broken for nonzero shifts?
  shifted_fft[0] = fft[0];
  shifted_fft[num_samples / 2] = fft[num_samples / 2];

  shifted_fft
}

fn detect_pitch_main() -> Result<(), String> {
  let sample_frequency = 44100.0;
  let samples = mvar::new();

  let _record_thread =
    unsafe {
      thread::scoped(|| {
        loop {
          let new_samples = record(sample_frequency, 100_000_000).unwrap();
          info!("Collected {} samples", new_samples.len());
          samples.overwrite(new_samples);
        }
      })
    };

  let _detect_thread =
    unsafe {
      thread::scoped(|| {
        info!("Detecting pitch..");

        loop {
          let samples = samples.take();
          match detect_frequency(samples, 1.0 / sample_frequency).unwrap() {
            None => {},
            Some(f) => {
              debug!("Estimated frequency is {}", f);
              let note = note::of_frequency(f);
              println!("Estimated note is {}", note.to_string_human());
            },
          }
        }
      })
    };

  Ok(())
}

fn add_harmonies(samples: Vec<f64>, semitones: &[f64]) -> Result<Vec<f64>, String> {
  let len = samples.len();
  let original_fft = try!(to_fft(samples));
  let mut fft = original_fft.clone();

  for &semitones in semitones {
    let harmony = harmony(&original_fft, semitones);
    for i in 0 .. len {
      fft[i] += harmony[i];
    }
  }

  let fft = fft.into_iter().map(|f| f / semitones.len() as f64).collect();

  let samples = try!(of_fft(fft));
  Ok(samples)
}

fn harmony_main(semitones: Vec<f64>) -> Result<(), String> {
  let sample_frequency = 44100.0;

  let samples = record(sample_frequency, 1_000_000_000).unwrap();
  info!("Collected {} samples", samples.len());

  let samples = try!(add_harmonies(samples, &semitones));
  let mut samples = samples.into_iter().map(|f| f as f32);

  let sample_rate = 44100 as f64;
  let buf_size = 1 << 10;

  try!(with_play_channel(
    sample_rate,
    buf_size,
    |stream| {
      loop {
        let mut snip = Vec::new();
        for _ in 0..buf_size {
          match samples.next() {
            None => break,
            Some(x) => snip.push(x),
          }
        }

        if snip.is_empty() {
          break
        }

        let len = snip.len();
        try!(string_err(stream.write(snip, len as u32)));
      }

      Ok(())
    }
  ));

  Ok(())
}

fn errorful_main() -> Result<(), String> {
  try!(string_err(env_logger::init()));
  try!(string_err(portaudio::pa::initialize()));

  let matches =
    clap_app!(pitchdetect =>
      (@subcommand play =>
        (about: "Play a note")
        (@arg note: +required +takes_value --note)
        (@arg time: +required +takes_value --time)
      )
      (@subcommand detect =>
        (about: "Detect pitch from the microphone")
      )
      (@subcommand harmony =>
        (about: "Add harmony to a recorded clip")
      )
    ).get_matches();

  if let Some(matches) = matches.subcommand_matches("play") {
    let note = try!(note::from_str(matches.value_of("note").unwrap()));
    let duration = try!(string_err(matches.value_of("time").unwrap().parse()));
    try!(play_note(note, duration));
  } else if let Some(_matches) = matches.subcommand_matches("detect") {
    try!(detect_pitch_main());
  } else if let Some(_harmony) = matches.subcommand_matches("harmony") {
    //let semitones = vec!(4, 7, 12);
    let semitones = vec!(4.0);
    try!(harmony_main(semitones));
  }

  Ok(())
}

#[main]
fn main() {
  errorful_main().unwrap();
}
<commit_msg>Add harmony in snippets<commit_after>use std;
use env_logger;
use time;
use portaudio;
use rgsl;
#[cfg(feature="gnuplot")]
use gnuplot;

use thread;
use note;
use mvar;

fn consume<T>(_: T) {}

fn to_fft(mut samples: Vec<f64>) -> Result<Vec<f64>, String> {
  let len = samples.len();
  let r = rgsl::fft::real_radix2::transform(&mut samples, 1, len);
  if r != rgsl::Value::Success {
    return Err(format!("rgsl returned {:?}", r));
  }

  Ok(samples)
}

fn of_fft(mut fft: Vec<f64>) -> Result<Vec<f64>, String> {
  let len = fft.len();
  let r = rgsl::fft::real_radix2::inverse(&mut fft, 1, len);
  if r != rgsl::Value::Success {
    return Err(format!("rgsl returned {:?}", r));
  }

  Ok(fft)
}

fn string_err<Ok, Err: std::fmt::Display>(r: Result<Ok, Err>) -> Result<Ok, String> {
  r.map_err(|err| format!("{}", err))
}

fn sine_wave(note: note::T, sample_frequency: f64, start: u32, end: u32) -> Vec<f32> {
  let mut buf = Vec::new();

  let f = note.to_frequency();

  for t in start .. end {
    let t = t as f64 / sample_frequency;
    let sample = (t * f * 2.0 * std::f64::consts::PI).sin();
    let amplitude = 1.0;
    buf.push((sample * amplitude) as f32);
  }

  buf
}

fn with_play_channel<T, F>(sample_rate: f64, buf_size: u32, f: F) -> Result<T, String> where
  F: FnOnce(&mut portaudio::pa::Stream<f32, f32>) -> Result<T, String>,
{
  let mut stream: portaudio::pa::Stream<f32, f32> = portaudio::pa::Stream::new();
  try!(string_err(
    stream.open_default_blocking(
      sample_rate,
      buf_size,
      0,
      1,
      portaudio::pa::SampleFormat::Float32,
    )));
  try!(string_err(stream.start()));

  let r = try!(f(&mut stream));

  try!(string_err(stream.stop()));

  Ok(r)
}

fn play_note(note: note::T, secs: u64) -> Result<(), String> {
  info!("Writing..");

  let sample_rate = 44100 as f64;
  let buf_size = 1 << 10;

  try!(with_play_channel(
    sample_rate,
    buf_size,
    |stream| {
      let start_time = time::precise_time_ns();
      let mut i = 0;
      while time::precise_time_ns() <= start_time + secs*1_000_000_000 {
        let buf = sine_wave(note, sample_rate, i*buf_size, (i+1)*buf_size);
        assert!(buf.len() == buf_size as usize);
        try!(string_err(stream.write(buf, buf_size)));

        i = i + 1;
      }

      Ok(())
    }
  ));

  Ok(())
}

fn record(sample_frequency: f64, delta_t_ns: u64) -> Result<Vec<f64>, String> {
  let buf_size = 1 << 10;
  assert!(buf_size & (buf_size - 1) == 0);

  info!("Reading..");

  let mut stream: portaudio::pa::Stream<f32, f32> = portaudio::pa::Stream::new();
  try!(string_err(
    stream.open_default_blocking(
      sample_frequency,
      buf_size,
      1,
      0,
      portaudio::pa::SampleFormat::Float32,
    )));
  try!(string_err(stream.start()));

  let mut buf = Vec::new();
  let start_time = time::precise_time_ns();
  while time::precise_time_ns() <= start_time + delta_t_ns || (buf.len() & (buf.len() - 1)) != 0 {
    let new = try!(string_err(stream.read(buf_size)));
    assert!(new.len () == buf_size as usize);
    buf.extend(new.into_iter().map(|f| f as f64));
  }

  try!(string_err(stream.stop()));

  Ok(buf)
}

fn detect_frequency(samples: Vec<f64>, timestep: f64) -> Result<Option<f64>, String> {
  let num_samples = samples.len();
  let samples = try!(to_fft(samples));

  let buckets: Vec<_> =
    (0 .. 1 + num_samples / 2)
    .filter_map(|i| {
      let real = samples[i];
      let imag =
        if i % num_samples / 2 == 0 {
          0.0
        } else {
          samples[num_samples - i]
        };

      let noise_threshold = 0.0;
      let x = real*real + imag*imag - noise_threshold;
      if x > noise_threshold {
        let f = i as f64 / num_samples as f64 / timestep;
        Some((f, x))
      } else {
        None
      }
    })
    .collect();
  consume(samples);

  if buckets.is_empty() {
    return Ok(None)
  }

  #[cfg(feature="gnuplot")]
  {
    let x: Vec<_> = buckets.iter().map(|&(f, _)| f).collect();
    let y: Vec<_> = buckets.iter().map(|&(_, y)| y).collect();
    let mut fg = gnuplot::Figure::new();
    fg
      .axes2d()
      .boxes(
        &x,
        &y,
        &[],
      );
    fg.echo_to_file("/tmp/gnuplot.txt");
  }

  let (max_f, max_val) =
    buckets.iter()
      .fold((0.0, 0.0), |(max_f, max_val), &(f, amp)| {
        let val = amp;
        if val > max_val {
          (f, val)
        } else {
          (max_f, max_val)
        }
      });

  debug!("Max index is {}", max_f);
  debug!("Max value is {}", max_val);

  Ok(Some(max_f))
}

fn harmony(fft: &[f64], semitones_up: f64) -> Vec<f64> {
  let len = fft.len();
  let shift_factor: f64 = (2.0 as f64).powf(semitones_up / 12.0);

  let mut shifted_fft: Vec<_> = std::iter::repeat(0.0).take(len).collect();
  for i in 1..(len / 2) {
    let source_i = (i as f64 / shift_factor).round() as usize;
    if source_i >= len / 2 {
      break
    }
    shifted_fft[i] += fft[source_i];
    shifted_fft[len - i] += fft[len - source_i];
  }

  shifted_fft[0] = fft[0];

  shifted_fft
}

fn detect_pitch_main() -> Result<(), String> {
  let sample_frequency = 44100.0;
  let samples = mvar::new();

  let _record_thread =
    unsafe {
      thread::scoped(|| {
        loop {
          let new_samples = record(sample_frequency, 100_000_000).unwrap();
          info!("Collected {} samples", new_samples.len());
          samples.overwrite(new_samples);
        }
      })
    };

  let _detect_thread =
    unsafe {
      thread::scoped(|| {
        info!("Detecting pitch..");

        loop {
          let samples = samples.take();
          match detect_frequency(samples, 1.0 / sample_frequency).unwrap() {
            None => {},
            Some(f) => {
              debug!("Estimated frequency is {}", f);
              let note = note::of_frequency(f);
              println!("Estimated note is {}", note.to_string_human());
            },
          }
        }
      })
    };

  Ok(())
}

fn add_harmonies(samples: Vec<f64>, semitones: &[f64]) -> Result<Vec<f64>, String> {
  let len = samples.len();

  let original_fft = try!(to_fft(samples));
  let mut fft = original_fft.clone();

  for &semitones in semitones {
    let harmony = harmony(&original_fft, semitones);
    assert!(fft.len() == harmony.len());
    for i in 0 .. len {
      fft[i] += harmony[i];
    }
  }

  let fft = fft.into_iter().map(|f| f / semitones.len() as f64).collect();

  let samples = try!(of_fft(fft));

  Ok(samples)
}

fn harmony_main(semitones: Vec<f64>) -> Result<(), String> {
  let sample_rate = 44100.0;
  let buf_size = 1 << 14;

  let samples = record(sample_rate, 2_000_000_000).unwrap();
  let len = samples.len() as u32;
  info!("Collected {} samples", len);
  let mut samples = samples.into_iter();

  let mut play_buffer: Vec<Vec<f32>> = Vec::new();
  loop {
    let mut snip = Vec::new();
    for _ in 0..buf_size {
      match samples.next() {
        None => break,
        Some(x) => snip.push(x),
      }
    }

    if snip.len() < buf_size as usize {
      break
    }

    let snip = try!(add_harmonies(snip, &semitones));
    let snip = snip.into_iter().map(|f| f as f32).collect();
    play_buffer.push(snip);
  }

  try!(with_play_channel(
    sample_rate,
    buf_size,
    |stream| {
      for snip in play_buffer {
        try!(string_err(stream.write(snip, buf_size)));
      }
      Ok(())
    },
  ));

  Ok(())
}

fn errorful_main() -> Result<(), String> {
  try!(string_err(env_logger::init()));
  try!(string_err(portaudio::pa::initialize()));

  let matches =
    clap_app!(pitchdetect =>
      (@subcommand play =>
        (about: "Play a note")
        (@arg note: +required +takes_value --note)
        (@arg time: +required +takes_value --time)
      )
      (@subcommand detect =>
        (about: "Detect pitch from the microphone")
      )
      (@subcommand harmony =>
        (about: "Add harmony to a recorded clip")
      )
    ).get_matches();

  if let Some(matches) = matches.subcommand_matches("play") {
    let note = try!(note::from_str(matches.value_of("note").unwrap()));
    let duration = try!(string_err(matches.value_of("time").unwrap().parse()));
    try!(play_note(note, duration));
  } else if let Some(_matches) = matches.subcommand_matches("detect") {
    try!(detect_pitch_main());
  } else if let Some(_harmony) = matches.subcommand_matches("harmony") {
    let semitones = vec!(4.0, 7.0);
    try!(harmony_main(semitones));
  }

  Ok(())
}

#[main]
fn main() {
  errorful_main().unwrap();
}
<|endoftext|>


<commit_before> #![feature(std_misc)]
extern crate time;

mod models;
mod parser;
mod mayberef;

use std::fs::{File, read_dir};
use std::io::{BufReader, BufRead, Read, Write, stdin};
use std::path::{Path};
use std::collections::HashMap;
use std::collections::hash_map::Entry::*;
use std::ascii::OwnedAsciiExt;
use models::*;

fn get_line() -> String {
    let mut stdin = stdin();
    let mut buffer = String::new();
    stdin.read_line(&mut buffer).unwrap();
    buffer
}

fn find_most_common_words(corpus_loc: &str, outfile: &str) {
    let path = Path::new(corpus_loc);
    let mut word_counts = HashMap::new();
    let words = files(path).flat_map(|file| read_words(file));

    for word in words {
        match word_counts.entry(word) {
            Vacant(e) => { e.insert(1); },
            Occupied(mut e) => { *e.get_mut() += 1; },
        }
    }

    let mut counts: Vec<_> = word_counts.into_iter().collect();
    counts.sort_by(|a, b| b.1.cmp(&a.1));

    let mut out = File::create(&Path::new(outfile)).unwrap();
    for (word, count) in counts {
        out.write_all(format!("{}, {}\n", word, count).as_bytes()).unwrap();
    }

}

fn load_most_common_words(filename: &str, num: usize) -> Vec<String> {
    let file = File::open(&Path::new(filename)).unwrap();
    let reader = BufReader::new(file);
    reader.lines().map(|line|
        line.unwrap().split(',').next().unwrap().to_string()
    )
    .take(num)
    .collect()
}

fn sentences<T: Read + 'static>(reader: BufReader<T>) -> Box<Iterator<Item = String>> {
    Box::new(reader.split('.' as u8)
                   .filter_map(|v| {
                        let lowercase = v.unwrap().into_ascii_lowercase();
                        String::from_utf8(lowercase).ok()
                    }))
}

fn read_words<R: BufRead + 'static>(reader: R) -> Box<Iterator<Item = String>> {
    Box::new(reader.lines()
                   .filter_map(|line| line.ok())
                   .flat_map(|line| {
                       line.split(|c| match c {
                                'a'...'z' | 'A'...'Z' => false,
                                _ => true,
                            })
                           .filter(|word| !word.is_empty())
                           .map(|word| word.to_string().into_ascii_lowercase())
                           .collect::<Vec<_>>()
                           .into_iter()
                   }))
}

fn files(path: &Path) -> Box<Iterator<Item = BufReader<File>>> {
    Box::new(read_dir(path)
            .unwrap()
            .map(|path| path.unwrap().path())
            .filter(|path| path.to_str().unwrap().ends_with(".txt"))
            .map(|path| {
                let file = File::open(&path).unwrap();
                BufReader::new(file)
            }))
}

fn create_model(path: &Path) {
    const CORPUS_DIR: &'static str = "/home/jamougha/corpus/pg";
    const WORDS: &'static str = "/home/jamougha/corpus/pg/word_counts.csv";
    let start_time = time::get_time();
    // find_most_common_words(CORPUS_DIR, WORDS);
    let words = load_most_common_words(WORDS, 30000);
    let mut builder = LanguageModelBuilder::new(10, words);

    let corpus_path = Path::new(CORPUS_DIR);

    let mut num_words = 0;

    for sentence in files(corpus_path).flat_map(sentences) {
        let mut acc = builder.new_sentence();
        for word in sentence.split(|c| match c {
            'a'...'z' => false,
            _ => true,
        }).filter(|w| !w.is_empty())
        {
            num_words += 1;
            acc.add_word(word);
            if num_words % 10_000_000 == 0 {
                println!("Loaded {} million words in {} seconds",
                         num_words / 1_000_000,
                         time::get_time().sec - start_time.sec);
            }
        }
    }

    builder.save(path);
    let end_time = time::get_time();
    println!("Model built in {}s", end_time.sec - start_time.sec);
}

fn main() {
    let path = Path::new("/home/jamougha/corpus/pg/model.data");
    // create_model(&path);
    let start_time = time::get_time();
    let builder = LanguageModelBuilder::load(&path);
    let model = builder.build();
    println!("Model loaded in {}s", time::get_time().sec - start_time.sec);

    loop {
        println!("");
        let input = get_line();
        if input.starts_with(":q") {
            break;
        }

        match parser::parse(input.trim_matches(|c: char| c.is_whitespace()), &model) {
            Ok(word_vec) => {
                let nearest = model.nearest_words(&word_vec);
                println!(" = {:?}", word_vec);
                println!("-------------");
                for word in nearest.iter().take(20) {
                    println!("{:?}, {}", word, word_vec.distance(word));
                }
            }
            Err(e) => println!("{:?}", e)
        }

    }

}
<commit_msg>Update to rust 1.4<commit_after> #![feature(std_misc)]
extern crate time;

mod models;
mod parser;
mod mayberef;

use std::fs::{File, read_dir};
use std::io::{BufReader, BufRead, Read, Write, stdin};
use std::path::{Path};
use std::collections::HashMap;
use std::collections::hash_map::Entry::*;
use models::*;

fn get_line() -> String {
    let mut stdin = stdin();
    let mut buffer = String::new();
    stdin.read_line(&mut buffer).unwrap();
    buffer
}

fn find_most_common_words(corpus_loc: &str, outfile: &str) {
    let path = Path::new(corpus_loc);
    let mut word_counts = HashMap::new();
    let words = files(path).flat_map(|file| read_words(file));

    for word in words {
        match word_counts.entry(word) {
            Vacant(e) => { e.insert(1); },
            Occupied(mut e) => { *e.get_mut() += 1; },
        }
    }

    let mut counts: Vec<_> = word_counts.into_iter().collect();
    counts.sort_by(|a, b| b.1.cmp(&a.1));

    let mut out = File::create(&Path::new(outfile)).unwrap();
    for (word, count) in counts {
        out.write_all(format!("{}, {}\n", word, count).as_bytes()).unwrap();
    }

}

fn load_most_common_words(filename: &str, num: usize) -> Vec<String> {
    let file = File::open(&Path::new(filename)).unwrap();
    let reader = BufReader::new(file);
    reader.lines().map(|line|
        line.unwrap().split(',').next().unwrap().to_string()
    )
    .take(num)
    .collect()
}

fn sentences<T: Read + 'static>(reader: BufReader<T>) -> Box<Iterator<Item = String>> {
    Box::new(reader.split('.' as u8)
                   .filter_map(|v| {
                        let lowercase = v.unwrap();
                        String::from_utf8(lowercase).ok().map(|s| s.to_uppercase())
                    }))
}

fn read_words<R: BufRead + 'static>(reader: R) -> Box<Iterator<Item = String>> {
    Box::new(reader.lines()
                   .filter_map(|line| line.ok())
                   .flat_map(|line| {
                       line.split(|c| match c {
                                'a'...'z' | 'A'...'Z' => false,
                                _ => true,
                            })
                           .filter(|word| !word.is_empty())
                           .map(|word| word.to_string().to_uppercase())
                           .collect::<Vec<_>>()
                           .into_iter()
                   }))
}

fn files(path: &Path) -> Box<Iterator<Item = BufReader<File>>> {
    Box::new(read_dir(path)
            .unwrap()
            .map(|path| path.unwrap().path())
            .filter(|path| path.to_str().unwrap().ends_with(".txt"))
            .map(|path| {
                let file = File::open(&path).unwrap();
                BufReader::new(file)
            }))
}

fn create_model(path: &Path) {
    const CORPUS_DIR: &'static str = "/home/jamougha/corpus/pg";
    const WORDS: &'static str = "/home/jamougha/corpus/pg/word_counts.csv";
    let start_time = time::get_time();
    // find_most_common_words(CORPUS_DIR, WORDS);
    let words = load_most_common_words(WORDS, 30000);
    let mut builder = LanguageModelBuilder::new(10, words);

    let corpus_path = Path::new(CORPUS_DIR);

    let mut num_words = 0;

    for sentence in files(corpus_path).flat_map(sentences) {
        let mut acc = builder.new_sentence();
        for word in sentence.split(|c| match c {
            'a'...'z' => false,
            _ => true,
        }).filter(|w| !w.is_empty())
        {
            num_words += 1;
            acc.add_word(word);
            if num_words % 10_000_000 == 0 {
                println!("Loaded {} million words in {} seconds",
                         num_words / 1_000_000,
                         time::get_time().sec - start_time.sec);
            }
        }
    }

    builder.save(path);
    let end_time = time::get_time();
    println!("Model built in {}s", end_time.sec - start_time.sec);
}

fn main() {
    let path = Path::new("/home/jamougha/corpus/pg/model.data");
    create_model(&path);
    let start_time = time::get_time();
    let builder = LanguageModelBuilder::load(&path);
    let model = builder.build();
    println!("Model loaded in {}s", time::get_time().sec - start_time.sec);

    loop {
        println!("");
        let input = get_line();
        if input.starts_with(":q") {
            break;
        }

        match parser::parse(input.trim_matches(|c: char| c.is_whitespace()), &model) {
            Ok(word_vec) => {
                let nearest = model.nearest_words(&word_vec);
                println!(" = {:?}", word_vec);
                println!("-------------");
                for word in nearest.iter().take(20) {
                    println!("{:?}, {}", word, word_vec.distance(word));
                }
            }
            Err(e) => println!("{:?}", e)
        }

    }

}
<|endoftext|>


<commit_before>use serde_json::{json, Value};
use std::io::stdin;
use std::string::String;
mod parser;

fn main() {
    loop {
        let mut buf = String::new();
        stdin().read_line(&mut buf).unwrap();
        if buf == "" {
            break;
        }
        let v: Value = serde_json::from_str(&buf).unwrap();
        let f = v.as_object().unwrap().get("f").unwrap().as_str().unwrap();
        let args = v.as_object().unwrap().get("a").unwrap().as_array().unwrap();
        if f == "echo" {
            let s = args[0].as_str().unwrap();
            println!(
                "{}",
                json!({
                    "error": "",
                    "data": echo(s)
                })
            )
        } else if f == "is_op" {
            let token_type = args[0].as_str().unwrap();
            println!(
                "{}",
                json!({
                    "error": "",
                    "data": parser::is_op(String::from(token_type))
                })
            )
        } else if f == "operator_strength" {
            let op = args[0].as_str().unwrap();
            println!(
                "{}",
                json!({
                    "error": "",
                    "data": parser::operator_strength(String::from(op))
                })
            )
        } else if f == "is_postfix_op" {
            let op = args[0].as_str().unwrap();
            println!(
                "{}",
                json!({
                    "error": "",
                    "data": parser::is_postfix_op(String::from(op))
                })
            )
        } else {
            println!(
                "{}",
                json!({
                    "error": "unknown function",
                    "data": null
                })
            )
        }
    }
}

fn echo(s: &str) -> String {
    return format!("echo: {}", s);
}
<commit_msg>Clean up function dispatch<commit_after>use serde_json::{json, Value};
use std::io::stdin;
use std::string::String;
mod parser;

fn main() {
    loop {
        let mut buf = String::new();
        stdin().read_line(&mut buf).unwrap();
        if buf == "" {
            break;
        }
        let v: Value = serde_json::from_str(&buf).unwrap();
        let f = v.as_object().unwrap().get("f").unwrap().as_str().unwrap();
        let args = v.as_object().unwrap().get("a").unwrap().as_array().unwrap();
        route(f, args);
    }
}

fn route(f: &str, args: &Vec<serde_json::Value>) {
    let result = match f {
        "echo" => {
            let s = args[0].as_str().unwrap();
            json!({
                "error": "",
                "data": echo(s)
            })
        }
        "is_op" => {
            let token_type = args[0].as_str().unwrap();
            json!({
                "error": "",
                "data": parser::is_op(String::from(token_type))
            })
        }
        "operator_strength" => {
            let op = args[0].as_str().unwrap();
            json!({
                "error": "",
                "data": parser::operator_strength(String::from(op))
            })
        }
        "is_postfix_op" => {
            let op = args[0].as_str().unwrap();
            json!({
                "error": "",
                "data": parser::is_postfix_op(String::from(op))
            })
        }
        _ => {
            json!({
                "error": "unknown function",
                "data": null
            })
        }
    };
    println!("{}", result);
}

fn echo(s: &str) -> String {
    return format!("echo: {}", s);
}
<|endoftext|>


<commit_before>/// An example of a client/server agnostic WebSocket that takes input from stdin and sends that
/// input to all other peers.
///
/// For example, to create a network like this:
///
/// 3013 ---- 3012 ---- 3014
///   \        |
///    \       |
///     \      |
///      \     |
///       \    |
///        \   |
///         \  |
///          3015
///
/// Run these commands in separate processes
/// ./peer2peer
/// ./peer2peer --server localhost:3013 ws://localhost:3012
/// ./peer2peer --server localhost:3014 ws://localhost:3012
/// ./peer2peer --server localhost:3015 ws://localhost:3012 ws://localhost:3013
///
/// Stdin on 3012 will be sent to all other peers
/// Stdin on 3013 will be sent to 3012 and 3015
/// Stdin on 3014 will be sent to 3012 only
/// Stdin on 3015 will be sent to 3012 and 2013

extern crate ws;
extern crate url;
extern crate clap;
extern crate env_logger;
extern crate rustc_serialize;
extern crate bincode;
#[macro_use] extern crate log;

use std::clone::Clone;
use std::io;
use std::io::prelude::*;
use std::thread;

use bincode::rustc_serialize::{encode, decode};

use clap::{App, Arg};

fn main() {
    // Setup logging
    env_logger::init().unwrap();

    // Parse command line arguments
    let matches = App::new("Simple Peer 2 Peer")
        .version("1.0")
        .author("Jonathan Almeida <hello@jonalmeida.com")
        .about("Connect to other peers and listen for incoming connections.")
        .arg(Arg::with_name("server")
             .short("s")
             .long("server")
             .value_name("SERVER")
             .help("Set the address to listen for new connections. (default: localhost:3012"))
        .arg(Arg::with_name("PEER")
             .help("A WebSocket URL to attempt to connect to at start.")
             .multiple(true))
        .get_matches();

    // Get address of this peer
    let my_addr = matches.value_of("server").unwrap_or("localhost:3012");
    let sending_addr = String::new() + my_addr;

    // Create simple websocket that just prints out messages
    let mut me = ws::WebSocket::new(|sender| {
        move |msg| {
            match msg {
                ws::Message::Binary(vector) => {
                    let encoded_msg = &*vector.into_boxed_slice();
                    let message: PeerMessage = decode(encoded_msg).unwrap();
                    Ok(info!("Peer {} got message: {}", message.sender, message.message))
                },
                ws::Message::Text(string) => {
                    Ok(warn!("We received a string, but we don't want to handle it: {}", string))
                },
            }
        }
    }).unwrap();

    // Get a sender for ALL connections to the websocket
    let broacaster = me.broadcaster();

    // Setup thread for listening to stdin and sending messages to connections
    let input = thread::spawn(move || {
        let stdin = io::stdin();
        for line in stdin.lock().lines() {
            // Send a message to all connections regardless of
            // how those connections were established
            let clocks = vec![0u32, 1u32, 1u32];
            let message = PeerMessage { sender: sending_addr.clone(), clocks: clocks, message: line.unwrap() };
            let encoded: Vec<u8> = encode(&message, bincode::SizeLimit::Infinite).unwrap();
            //broacaster.send(line.unwrap()).unwrap();
            broacaster.send(&*encoded.into_boxed_slice()).unwrap();
        }
    });

    // Connect to any existing peers specified on the cli
    if let Some(peers) = matches.values_of("PEER") {
        for peer in peers {
            me.connect(url::Url::parse(peer).unwrap()).unwrap();
        }
    }

    // Run the websocket
    me.listen(my_addr).unwrap();
    input.join().unwrap();

}

#[derive(RustcEncodable, RustcDecodable, PartialEq)]
struct PeerMessage {
    // Host address passed in via CLI. Used for identification.
    sender: String,
    // Vector clocks for each process.
    clocks: Vec<u32>,
    // The message that we want to send. For now, we're testing with a string.
    message: String,
}
<commit_msg>Clocks are getting sent as well<commit_after>/// An example of a client/server agnostic WebSocket that takes input from stdin and sends that
/// input to all other peers.
///
/// For example, to create a network like this:
///
/// 3013 ---- 3012 ---- 3014
///   \        |
///    \       |
///     \      |
///      \     |
///       \    |
///        \   |
///         \  |
///          3015
///
/// Run these commands in separate processes
/// ./peer2peer
/// ./peer2peer --server localhost:3013 ws://localhost:3012
/// ./peer2peer --server localhost:3014 ws://localhost:3012
/// ./peer2peer --server localhost:3015 ws://localhost:3012 ws://localhost:3013
///
/// Stdin on 3012 will be sent to all other peers
/// Stdin on 3013 will be sent to 3012 and 3015
/// Stdin on 3014 will be sent to 3012 only
/// Stdin on 3015 will be sent to 3012 and 2013

extern crate ws;
extern crate url;
extern crate clap;
extern crate env_logger;
extern crate rustc_serialize;
extern crate bincode;
#[macro_use] extern crate log;

use std::clone::Clone;
use std::io;
use std::io::prelude::*;
use std::thread;

use bincode::rustc_serialize::{encode, decode};

use clap::{App, Arg};

fn main() {
    // Setup logging
    env_logger::init().unwrap();

    // Parse command line arguments
    let matches = App::new("Simple Peer 2 Peer")
        .version("1.0")
        .author("Jonathan Almeida <hello@jonalmeida.com")
        .about("Connect to other peers and listen for incoming connections.")
        .arg(Arg::with_name("server")
             .short("s")
             .long("server")
             .value_name("SERVER")
             .help("Set the address to listen for new connections. (default: localhost:3012"))
        .arg(Arg::with_name("PEER")
             .help("A WebSocket URL to attempt to connect to at start.")
             .multiple(true))
        .get_matches();

    // Get address of this peer
    let my_addr = matches.value_of("server").unwrap_or("localhost:3012");
    let sending_addr = String::new() + my_addr;

    // Create simple websocket that just prints out messages
    let mut me = ws::WebSocket::new(|sender| {
        move |msg| {
            match msg {
                ws::Message::Binary(vector) => {
                    let encoded_msg = &*vector.into_boxed_slice();
                    let message: PeerMessage = decode(encoded_msg).unwrap();
                    Ok(info!("Peer {} with clocks: {:?} got message: {}",
                             message.sender,
                             message.clocks,
                             message.message))
                },
                ws::Message::Text(string) => {
                    Ok(warn!("We received a string, but we don't want to handle it: {}", string))
                },
            }
        }
    }).unwrap();

    // Get a sender for ALL connections to the websocket
    let broacaster = me.broadcaster();

    // Setup thread for listening to stdin and sending messages to connections
    let input = thread::spawn(move || {
        let stdin = io::stdin();
        for line in stdin.lock().lines() {
            // Send a message to all connections regardless of
            // how those connections were established
            let clocks = vec![0, 1, 1];
            let message = PeerMessage { sender: sending_addr.clone(), clocks: clocks, message: line.unwrap() };
            let encoded: Vec<u8> = encode(&message, bincode::SizeLimit::Infinite).unwrap();
            //broacaster.send(line.unwrap()).unwrap();
            broacaster.send(&*encoded.into_boxed_slice()).unwrap();
        }
    });

    // Connect to any existing peers specified on the cli
    if let Some(peers) = matches.values_of("PEER") {
        for peer in peers {
            me.connect(url::Url::parse(peer).unwrap()).unwrap();
        }
    }

    // Run the websocket
    me.listen(my_addr).unwrap();
    input.join().unwrap();

}

#[derive(RustcEncodable, RustcDecodable, PartialEq)]
struct PeerMessage {
    // Host address passed in via CLI. Used for identification.
    sender: String,
    // Vector clocks for each process.
    clocks: Vec<u32>,
    // The message that we want to send. For now, we're testing with a string.
    message: String,
}
<|endoftext|>