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

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

use gnuplot::*;
use dft::{Operation, Plan};

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

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

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

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

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

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

struct RTLSDR {
    rtlsdr: RTLSDRDevice,
}

unsafe impl Send for RTLSDR {}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

}

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

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

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

#[inline]
fn moving_average(ntaps: usize) -> Vec<f32> {
    vec![1.0 / ntaps as f32; ntaps]
}
<commit_msg>Wideband FM is working!<commit_after>extern crate num;
extern crate portaudio;
extern crate rtlsdr;
extern crate dft;
extern crate gnuplot;

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

use gnuplot::*;
use dft::{Operation, Plan};

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

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

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

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

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

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

struct RTLSDR {
    rtlsdr: RTLSDRDevice,
}

unsafe impl Send for RTLSDR {}

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

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

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

    // We were decimating way too much. We need to decimate in two stages: one
    // to remove nearby stations and another to match the sound card. For
    // wideband signals, 8 and 3 seems to work well and also hits that sweet
    // spot of no buffer underflows!
    //
    // TODO: need to still calculate this on the fly. In addition, we need to
    // find apporpriate decimation values for narrowband. I'm still confident
    // that the decimation value depends on the bandwitdh, but I need to
    // figure out that appropriate formula.
    //
    let dec_rate_filt = 8;
    let dec_rate_audio = 3;
    let mut sdr =
        RTLSDR { rtlsdr: init_sdr(sdr_index, fm_freq, bandwidth).unwrap() };

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

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

        // After decimation, demodulate the signal and send out of the
        // thread to the receiver.
        let res = demod_fm(iq_vec, prev);
        let mut demod_iq = filter_real(res.0, taps.clone());
        demod_iq = decimate(demod_iq.as_slice(), dec_rate_audio);
        prev = res.1;

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

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

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

fn plot_spectrum<'a>(signal: &Vec<Complex<f32>>, figure: &'a mut Figure) -> &'a mut Figure {
    let mut sig = signal.clone();
    let radix_2: u32 = (sig.len() as f32).log2().ceil() as u32;
    let new_len = 2u32.pow(radix_2) as usize;
    println!("signal len: {}, new len: {}", sig.len(), new_len);
    sig.resize(new_len, Complex::zero());
    let plan = Plan::new(Operation::Forward, new_len);
    dft::transform(&mut sig, &plan);

    // Now generate a plot.
    figure.clear_axes();
    figure.axes2d()
        .lines(0..new_len, sig.iter().map(|x| x.norm()), &[]);
    figure.show();
    figure
}

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

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

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

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

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

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

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

    for samp in iq.iter() {
        let conj = p.conj() * samp;
        demod_queue.push(conj.arg() * gain);
        p = *samp;
    }
    (demod_queue, p)
}

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

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

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

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

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

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


<commit_before>#![feature(conservative_impl_trait)]

extern crate curl;
extern crate env_logger;
extern crate futures;
extern crate tokio_core;
extern crate tokio_curl;
extern crate threadpool;
#[macro_use]
extern crate clap;
extern crate chrono;
extern crate histogram;

mod args;

use args::{Config, parse_args};

use curl::easy::Easy;
use futures::{Future, stream};
use futures::future::{ok, loop_fn, Loop};
use futures::{Stream};
use tokio_core::reactor::{Core};
use tokio_curl::{PerformError, Session};
use threadpool::ThreadPool;
use chrono::*;
use std::sync::mpsc::channel;
use histogram::*;

struct Worker {
    /// Number of succesful requests this worker has made
    num_requests: usize,
    histogram: Histogram,
    /// Number of failed requests this worker has made
    num_failed_requests: usize,
}

impl Worker {
    pub fn new() -> Self {
        Worker { num_requests: 0, histogram: Histogram::new(), num_failed_requests: 0}
    }

    fn send_request(mut self, easy_request: Easy, session: &Session) -> impl Future<Item=(Self, Easy), Error=PerformError> {
        session.perform(easy_request)
            .then(move |res| {
                match res {
                    Ok(mut easy) => {
                        (&mut self).num_requests += 1;
                        let latency = Duration::from_std(easy.total_time().unwrap()).unwrap();
                        self.histogram.increment(latency.num_nanoseconds().unwrap() as u64).unwrap();
                        Ok((self, easy))
                    },
                    Err(error) => {
                        (&mut self).num_failed_requests += 1;
                        Err(error)
                    }
                }
            })
    }
}

struct RunInfo {
    /// Number of requests successfully completed
    pub requests_completed: usize,
    pub duration: Duration,
    histogram: Histogram,
    /// Number of failed requests
    num_failed_requests: usize,
}

impl RunInfo {
    pub fn requests_per_second(&self) -> f32 {
        (self.requests_completed as f32) / (self.duration.num_seconds() as f32)
    }
}

/// Delegates and manages all the work.
struct Boss {
    /// Number of requests successfully completed
    pub requests_completed: usize,
    /// Number of active connections
    pub connections: usize,

    thread_pool: ThreadPool,
    num_threads: usize,
}

impl Boss {
    pub fn new(num_threads: usize) -> Self {
        Boss {
            requests_completed: 0,
            connections: 0,
            thread_pool: ThreadPool::new(num_threads),
            num_threads: num_threads,
        }
    }

    pub fn start_workforce(&self, desired_connections: usize, url: String, duration: Duration) -> RunInfo {
        let (tx, rx) = channel();
        for _ in 0..self.num_threads {
            let tx = tx.clone();
            let url = url.clone();
            self.thread_pool.execute(move || {

                let mut lp = Core::new().unwrap();
                let start_time = Local::now();
                let wanted_end_time = start_time + duration;
                let session = Session::new(lp.handle());

                let iterator = (0..desired_connections).map(|_| {
                    let mut easy_request = Easy::new();
                    easy_request.get(true).unwrap();
                    easy_request.url(&url).unwrap();
                    loop_fn((Worker::new(), easy_request), |(worker, easy)| {
                        worker.send_request(easy, &session)
                            .and_then(|state| {
                                let now_time = Local::now();
                                if now_time < wanted_end_time {
                                    Ok(Loop::Continue(state))
                                } else {
                                    Ok(Loop::Break(state))
                                }
                            })
                    })
                });
                let future = stream::futures_unordered(iterator)
                    .fold((0, 0, Histogram::new()), |(request_acc, request_failed_acc, _), (res, _)|  {
                        let requests = request_acc + res.num_requests;
                        let requests_failed = request_failed_acc + res.num_failed_requests;
                        ok((requests, requests_failed, res.histogram))
                    });
                let res = lp.run(future).unwrap();
                tx.send(res).unwrap();
            });
        }
        // collect information from all threads
        let (total_num_requests, requests_failed, histogram): (usize, usize, Histogram) = rx.iter().take(self.num_threads).fold((0, 0, Histogram::new()), |(request_acc, request_failed_acc, mut histogram_acc), (request, request_failed, histogram)|  {
                        let requests = request_acc + request;
                        let requests_failed = request_failed_acc + request_failed;
                        histogram_acc.merge(&histogram);
                        (requests, requests_failed, histogram_acc)
                    });
        RunInfo {requests_completed: total_num_requests, num_failed_requests: requests_failed, duration: duration, histogram: histogram}
    }
}

fn start(config: Config) -> RunInfo {
    let boss = Boss::new(config.num_threads);
    boss.start_workforce(config.num_connections, config.url, config.duration)
}

fn nanoseconds_to_milliseconds(nanoseconds: u64) -> f64 {
    nanoseconds as f64 / 1_000_000_f64
}

/// Presents the results to the user
fn present(run_info: RunInfo) {
    let requests_completed = run_info.requests_completed;
    println!("{} requests in {}s", requests_completed, run_info.duration.num_seconds());
    println!("{} failed requests", run_info.num_failed_requests);
    println!("Requests per second: {}", run_info.requests_per_second());
    println!("Latency distribution: \n50%: {} ms\n75%: {} ms\n90%: {} ms\n95%: {} ms\n99%: {} ms",
        run_info.histogram.percentile(50.0).map(nanoseconds_to_milliseconds).unwrap(),
        run_info.histogram.percentile(75.0).map(nanoseconds_to_milliseconds).unwrap(),
        run_info.histogram.percentile(90.0).map(nanoseconds_to_milliseconds).unwrap(),
        run_info.histogram.percentile(95.0).map(nanoseconds_to_milliseconds).unwrap(),
        run_info.histogram.percentile(99.0).map(nanoseconds_to_milliseconds).unwrap(),
    );
}

fn main() {
    env_logger::init().unwrap();
    let config = parse_args();
    let run_info = start(config);
    present(run_info)
}
<commit_msg>less use of unwrap, just map values around<commit_after>#![feature(conservative_impl_trait)]

extern crate curl;
extern crate env_logger;
extern crate futures;
extern crate tokio_core;
extern crate tokio_curl;
extern crate threadpool;
#[macro_use]
extern crate clap;
extern crate chrono;
extern crate histogram;

mod args;

use args::{Config, parse_args};

use curl::easy::Easy;
use futures::{Future, stream};
use futures::future::{ok, loop_fn, Loop};
use futures::{Stream};
use tokio_core::reactor::{Core};
use tokio_curl::{PerformError, Session};
use threadpool::ThreadPool;
use chrono::*;
use std::sync::mpsc::channel;
use histogram::*;

struct Worker {
    /// Number of succesful requests this worker has made
    num_requests: usize,
    histogram: Histogram,
    /// Number of failed requests this worker has made
    num_failed_requests: usize,
}

impl Worker {
    pub fn new() -> Self {
        Worker { num_requests: 0, histogram: Histogram::new(), num_failed_requests: 0}
    }

    fn send_request(mut self, easy_request: Easy, session: &Session) -> impl Future<Item=(Self, Easy), Error=PerformError> {
        session.perform(easy_request)
            .then(|res| {
                match res {
                    Ok(mut easy) => {
                        self.num_requests += 1;
                        easy.total_time()
                            .ok()
                            .and_then(|x|
                                Duration::from_std(x).ok()
                            )
                            .and_then(|latency| {
                                latency.num_nanoseconds()
                                    .map(|x| self.histogram.increment(x as u64).ok())
                            });

                        Ok((self, easy))
                    },
                    Err(error) => {
                        self.num_failed_requests += 1;
                        Err(error)
                    }
                }
            })
    }
}

struct RunInfo {
    /// Number of requests successfully completed
    pub requests_completed: usize,
    pub duration: Duration,
    histogram: Histogram,
    /// Number of failed requests
    num_failed_requests: usize,
}

impl RunInfo {
    pub fn requests_per_second(&self) -> f32 {
        (self.requests_completed as f32) / (self.duration.num_seconds() as f32)
    }
}

/// Delegates and manages all the work.
struct Boss {
    /// Number of requests successfully completed
    pub requests_completed: usize,
    /// Number of active connections
    pub connections: usize,

    thread_pool: ThreadPool,
    num_threads: usize,
}

impl Boss {
    pub fn new(num_threads: usize) -> Self {
        Boss {
            requests_completed: 0,
            connections: 0,
            thread_pool: ThreadPool::new(num_threads),
            num_threads: num_threads,
        }
    }

    pub fn start_workforce(&self, desired_connections: usize, url: String, duration: Duration) -> RunInfo {
        let (tx, rx) = channel();
        for _ in 0..self.num_threads {
            let tx = tx.clone();
            let url = url.clone();
            self.thread_pool.execute(move || {

                let mut lp = Core::new().unwrap();
                let start_time = Local::now();
                let wanted_end_time = start_time + duration;
                let session = Session::new(lp.handle());

                let iterator = (0..desired_connections).map(|_| {
                    let mut easy_request = Easy::new();
                    easy_request.get(true).unwrap();
                    easy_request.url(&url).unwrap();
                    loop_fn((Worker::new(), easy_request), |(worker, easy)| {
                        worker.send_request(easy, &session)
                            .and_then(|state| {
                                let now_time = Local::now();
                                if now_time < wanted_end_time {
                                    Ok(Loop::Continue(state))
                                } else {
                                    Ok(Loop::Break(state))
                                }
                            })
                    })
                });
                let future = stream::futures_unordered(iterator)
                    .fold((0, 0, Histogram::new()), |(request_acc, request_failed_acc, _), (res, _)|  {
                        let requests = request_acc + res.num_requests;
                        let requests_failed = request_failed_acc + res.num_failed_requests;
                        ok((requests, requests_failed, res.histogram))
                    });
                let res = lp.run(future).unwrap();
                tx.send(res).unwrap();
            });
        }
        // collect information from all threads
        let (total_num_requests, requests_failed, histogram): (usize, usize, Histogram) = rx.iter().take(self.num_threads).fold((0, 0, Histogram::new()), |(request_acc, request_failed_acc, mut histogram_acc), (request, request_failed, histogram)|  {
                        let requests = request_acc + request;
                        let requests_failed = request_failed_acc + request_failed;
                        histogram_acc.merge(&histogram);
                        (requests, requests_failed, histogram_acc)
                    });
        RunInfo {requests_completed: total_num_requests, num_failed_requests: requests_failed, duration: duration, histogram: histogram}
    }
}

fn start(config: Config) -> RunInfo {
    let boss = Boss::new(config.num_threads);
    boss.start_workforce(config.num_connections, config.url, config.duration)
}

fn nanoseconds_to_milliseconds(nanoseconds: u64) -> f64 {
    nanoseconds as f64 / 1_000_000_f64
}

/// Presents the results to the user
fn present(run_info: RunInfo) {
    let requests_completed = run_info.requests_completed;
    println!("{} requests in {}s", requests_completed, run_info.duration.num_seconds());
    println!("{} failed requests", run_info.num_failed_requests);
    println!("Requests per second: {}", run_info.requests_per_second());
    println!("Latency distribution: \n50%: {} ms\n75%: {} ms\n90%: {} ms\n95%: {} ms\n99%: {} ms",
        run_info.histogram.percentile(50.0).map(nanoseconds_to_milliseconds).unwrap(),
        run_info.histogram.percentile(75.0).map(nanoseconds_to_milliseconds).unwrap(),
        run_info.histogram.percentile(90.0).map(nanoseconds_to_milliseconds).unwrap(),
        run_info.histogram.percentile(95.0).map(nanoseconds_to_milliseconds).unwrap(),
        run_info.histogram.percentile(99.0).map(nanoseconds_to_milliseconds).unwrap(),
    );
}

fn main() {
    env_logger::init().unwrap();
    let config = parse_args();
    let run_info = start(config);
    present(run_info)
}
<|endoftext|>


<commit_before>extern crate card_engine;
extern crate rand;

use card_engine::{Round, GameState, Action, ActionError};
use card_engine::cards::BasicCard;
use rand::{thread_rng};
use rand::distributions::{IndependentSample, Range};

fn format_hand(hand: &[BasicCard], gs: &GameState) -> String {
    let mut cards: Vec<_> = hand.iter().collect();
    cards.sort_by(|a, b| gs.display_order(a, b));

    let sc: Vec<_> = cards.iter().map(|x| format!("{}", x)).collect();
    sc.join(" ")
}

fn format_state(gs: &GameState) -> String {
    format!("Player 1: Score {}, Hand: {}\nPlayer 2: Score {}, Hand: {}\nTrump: {}",
            gs.score[0], &format_hand(&gs.hands[0], gs),
            gs.score[1], &format_hand(&gs.hands[1], gs),
            gs.trump)
}

fn print_round(round: &Round) {
    println!("{}", &format_state(round.get_state()));
    println!("{}\n", round.get_phase());
}

fn print_action(action: &Action) {
    println!("Player {} plays {}.", action.player + 1, action.card);
}

/// Randomly choose actions at each play
fn play_random_game(verbose: bool) -> Result<[usize; 2], ActionError> {
    let mut round = Round::new();
    let mut rng = thread_rng();

    if verbose {
        print_round(&round);
    }

    let mut actions = round.possible_actions();

    while actions.len() > 0 {

        let card_range = Range::new(0, actions.len());
        let ri = card_range.ind_sample(&mut rng) as usize;

        let action = actions[ri].clone();
        if verbose {
            print_action(&action);
        }
        round.play_action(action)?;

        if verbose {
            println!("**************");
            print_round(&round);
        }

        actions = round.possible_actions();
    }

    Ok(round.get_state().score.clone())
}

fn main() {
    let mut games_won = [0, 0];
    for _ in 0..100 {
        let score = play_random_game(false).expect("bad game");
        let winner: usize = if score[0] > score[1]  { 0 } else { 1 };
        games_won[winner] += 1;
    }

    println!("Player 1 Record: {}-{}", games_won[0], games_won[1]);
}
<commit_msg>player API<commit_after>extern crate card_engine;
extern crate rand;

use card_engine::{Round, Action, ActionError};
use card_engine::cards::{BasicCard, Rank, Suit};
use card_engine::germanwhist::util::*;
use card_engine::germanwhist::PlayerView;
use rand::{thread_rng};
use rand::distributions::{IndependentSample, Range};

trait Player {
    /// Return a card to play, based on the current view of the world.
    fn play_card(&self, view: &PlayerView) -> BasicCard;
}

pub struct RandomPlayer;

impl RandomPlayer {
    fn new() -> RandomPlayer {
        RandomPlayer { }
    }
}

impl Player for RandomPlayer {
    fn play_card(&self, view: &PlayerView) -> BasicCard {
        let mut rng = thread_rng();
        let cards = view.playable_cards();
        let card_range = Range::new(0, cards.len());
        let ri = card_range.ind_sample(&mut rng) as usize;
        cards[ri].clone()
    }
}

pub struct BasicPlayer {
    // try to win any non-trump with rank above this
    min_nontrump_rank_to_win: Option<u8>
}

impl BasicPlayer {
    fn new(mntr: Option<Rank>) -> BasicPlayer {
        BasicPlayer { min_nontrump_rank_to_win: mntr.map(|c| c.ord_ace_high()) }
    }

    fn try_to_win(&self, card: &BasicCard, trump: Suit) -> bool {
        card.suit == trump || self.min_nontrump_rank_to_win.map(|c| card.rank.ord_ace_high() >= c).unwrap_or(false)
    }
}

impl Player for BasicPlayer {
    fn play_card(&self, view: &PlayerView) -> BasicCard {
        let mut cards = view.playable_cards();

        cards.sort_by_key(|c| (c.rank.ord_ace_high(), view.ord_suit(c.suit)));

        match &view.revealed {
            // playing for cards
            &Some(ref c)  => {

                // go all-out for trumps, kings, or better
                if self.try_to_win(c, view.trump) {
                    // Play the highest non-trump, otherwise play the lowest trump
                    cards.iter().rev().find(|p| p.suit != view.trump)
                        .unwrap_or(&cards[0]).clone()
                } else {
                    // try to ditch
                    cards.iter().find(|p| p.suit != view.trump).unwrap_or(&cards[0]).clone()
                }
            },

            // playing for points
            &None => {
                match &view.leading_card {
                    &Some(ref lc) => {
                        // play the lowest card to beat it, otherwise ditch
                        cards.iter().find(|p| view.wins_against(lc, p))
                            .unwrap_or(&cards[0]).clone()
                    },
                    &None => { cards[0].clone() }
                }
            }
        }
    }
}

/// Randomly choose actions at each play
fn play_random_game(start: usize, r: Option<Rank>, verbose: bool) -> Result<[usize; 2], ActionError> {
    let mut round = Round::new(start, (1, 2));

    if verbose {
        println!("{}", format_round(&round));
    }

    let mut actions = round.possible_actions();

    let players: [Box<Player>; 2] = [Box::new(BasicPlayer::new(r)), Box::new(RandomPlayer::new())];

    while actions.len() > 0 {
        let action = {
            let player_view = round.active_player_view();
            let card = players[player_view.player].play_card(&player_view);

            Action { card, player: player_view.player }
        };

        if verbose {
            println!("{}", format_action(&action));
        }

        round.play_action(action)?;

        if verbose {
            println!("**************");
            println!("{}", format_round(&round));
        }

        actions = round.possible_actions();
    }

    Ok(round.get_state().score.clone())
}

fn test_basic_player(r: Option<Rank>) -> [usize; 2] {
    let mut games_won = [0, 0];
    let mut starting_player: usize = 0;

    for _ in 0..10000 {
        let score = play_random_game(starting_player, r, false).expect("bad game");
        let winner: usize = if score[0] > score[1]  { 0 } else { 1 };
        games_won[winner] += 1;
        starting_player = 1 - winner;
    }

    games_won
}

fn main() {

    for r in Rank::iterator() {
        let games_won = test_basic_player(Some(r.clone()));

        println!("Player 1 Record, Rank {}: {}-{}", r, games_won[0], games_won[1]);
    }

    let games_won = test_basic_player(None);

    println!("Player 1 Record, Never: {}-{}", games_won[0], games_won[1]);

}

<|endoftext|>


<commit_before>#![cfg_attr(test, allow(dead_code))]
#![feature(collections, os, env, io, std_misc, libc, core)]
#![cfg_attr(not(windows), feature(path))]

extern crate libc;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::env;
use std::old_io;
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            env::set_exit_status(1);
            return;
        },
    };

    if args.help { return; }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        event_loop(&choices, &initial_search);
    }
}

fn event_loop(choices: &[&str], initial_search: &str) {
    let mut search = Search::new(choices, initial_search.to_string());
    let mut screen = Screen::open_screen();

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(&mut screen, &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in keys.iter() {
            handle_key(&mut search, key, screen.visible_choices);
        }
    }

    screen.move_cursor_to_bottom();
    println!("{}", search.get_selection());
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace => search.backspace(),
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') => search.cancel(),
        Control('g') => search.cancel(),
        Control('n') => search.down(visible_choices),
        Control('p') => search.up(),
        Enter => search.done(),
        _ => {}
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    index: usize,
    state: SearchState,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            index: 0,
            state: InProgress,
        }
    }

    fn up(&mut self) {
        self.index = if self.index == 0 { 0 } else { self.index - 1 };
    }

    fn down(&mut self, visible_choices: u16) {
        let limit = min(visible_choices as usize - 1, self.matches.len() - 1);
        self.index = min(self.index + 1, limit);
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.index = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.index = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query);
            self.stale = false;
        }
    }

    fn get_selection(&mut self) -> String {
        if self.state == Canceled {
            "".to_string()
        } else {
            self.recompute_matches();
            self.matches.get(self.index).unwrap_or(&"").to_string()
        }
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices)\n", search.query, search.matches.len(), search.choices.len()));

    print_matches(screen, &search.matches, search.index);

    let start = screen.start_line;
    screen.move_cursor(start, 2 + search.query.len() as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut Screen, matches: &Vec<&str>, index: usize) {
    let mut i = 1;
    let max_width = screen.width as usize;
    for choice in matches.iter() {
        let choice = choice.slice_chars(0, min(max_width, choice.width(false)));
        if i == index + 1 {
            screen.write_inverted(&choice);
        } else {
            screen.write(&choice);
        }
        if i >= screen.visible_choices as usize {
            return;
        } else {
            screen.write("\n");
        }
        i += 1;
    }
}

fn read_choices() -> Vec<String> {
    let mut stdin = old_io::stdio::stdin();
    let mut lines = Vec::new();

    while let Ok(mut s) = stdin.read_line() {
        trim(&mut s);
        lines.push(s);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

#[test]
fn trim_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        trim(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("\n", "");
    should_become("\n\n", "");
    should_become("asdf", "asdf");
    should_become("asdf\n", "asdf");
    should_become("asdf\nasdf\n", "asdf\nasdf");
}

#[test]
fn delete_word_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        delete_last_word(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("a", "");
    should_become("asdf", "");
    should_become("asdf asdf asdf", "asdf asdf ");
    should_become("asdf asdf asdf ", "asdf asdf ");
    should_become("asdf asdf asdf  ", "asdf asdf ");
}
<commit_msg>Fix Unicode truncation panic<commit_after>#![cfg_attr(test, allow(dead_code))]
#![feature(collections, os, env, io, std_misc, libc, core)]
#![cfg_attr(not(windows), feature(path))]

extern crate libc;

mod args;
mod matching;
mod screen;
#[cfg(not(windows))] mod ansi;

use std::env;
use std::old_io;
use std::cmp::min;
use screen::Screen;
use screen::Key;
use screen::Key::*;
use self::SearchState::*;

fn main() {
    let args = match args::parse_args() {
        Some(args) => args,
        None => {
            env::set_exit_status(1);
            return;
        },
    };

    if args.help { return; }

    let choices = read_choices();
    let initial_search = args.initial_search.clone();
    let choices = choices.iter().map(|x| &x[]).collect::<Vec<&str>>();
    if args.use_first {
        let matches = matching::compute_matches(&choices, &initial_search);
        println!("{}", matches.get(0).unwrap_or(&""));
        return;
    } else {
        event_loop(&choices, &initial_search);
    }
}

fn event_loop(choices: &[&str], initial_search: &str) {
    let mut search = Search::new(choices, initial_search.to_string());
    let mut screen = Screen::open_screen();

    loop {
        search.recompute_matches();

        match search.state {
            InProgress => draw_screen(&mut screen, &search),
            _ => break,
        }

        let keys = screen.get_buffered_keys();
        for key in keys.iter() {
            handle_key(&mut search, key, screen.visible_choices);
        }
    }

    screen.move_cursor_to_bottom();
    println!("{}", search.get_selection());
}

fn handle_key(search: &mut Search, key: &Key, visible_choices: u16) {
    match *key {
        Char(x) => search.append(x),
        Backspace => search.backspace(),
        Control('h') => search.backspace(),
        Control('w') => search.delete_word(),
        Control('u') => search.clear_query(),
        Control('c') => search.cancel(),
        Control('g') => search.cancel(),
        Control('n') => search.down(visible_choices),
        Control('p') => search.up(),
        Enter => search.done(),
        _ => {}
    }
}

struct Search<'a> {
    choices: &'a [&'a str],
    query: String,
    matches: Vec<&'a str>,
    stale: bool,
    index: usize,
    state: SearchState,
}

#[derive(PartialEq, Eq)]
enum SearchState {
    InProgress,
    Done,
    Canceled,
}

impl<'a> Search<'a> {
    fn new(choices: &'a [&'a str], initial_search: String) -> Search<'a> {
        let matches = choices.to_vec();
        Search {
            choices: choices,
            query: initial_search,
            matches: matches,
            stale: true,
            index: 0,
            state: InProgress,
        }
    }

    fn up(&mut self) {
        self.index = if self.index == 0 { 0 } else { self.index - 1 };
    }

    fn down(&mut self, visible_choices: u16) {
        let limit = min(visible_choices as usize - 1, self.matches.len() - 1);
        self.index = min(self.index + 1, limit);
    }

    fn backspace(&mut self) {
        self.query.pop();
        self.stale = true;
        self.index = 0;
        self.matches = self.choices.to_vec();
    }

    fn delete_word(&mut self) {
        self.stale = true;
        delete_last_word(&mut self.query);
        self.matches = self.choices.to_vec();
    }

    fn append(&mut self, c: char) {
        self.query.push(c);
        self.stale = true;
        self.index = 0;
    }

    fn clear_query(&mut self) {
        self.query.clear();
        self.matches = self.choices.to_vec();
    }

    fn recompute_matches(&mut self) {
        if self.stale {
            self.matches = matching::compute_matches(&self.matches, &self.query);
            self.stale = false;
        }
    }

    fn get_selection(&mut self) -> String {
        if self.state == Canceled {
            "".to_string()
        } else {
            self.recompute_matches();
            self.matches.get(self.index).unwrap_or(&"").to_string()
        }
    }

    fn cancel(&mut self) {
        self.state = Canceled
    }

    fn done(&mut self) {
        self.state = Done
    }
}

fn draw_screen(screen: &mut Screen, search: &Search) {
    screen.hide_cursor();
    screen.blank_screen();
    screen.write(&format!("> {} ({}/{} choices)\n", search.query, search.matches.len(), search.choices.len()));

    print_matches(screen, &search.matches, search.index);

    let start = screen.start_line;
    screen.move_cursor(start, 2 + search.query.len() as u16);
    screen.show_cursor();
}

fn print_matches(screen: &mut Screen, matches: &Vec<&str>, index: usize) {
    let mut i = 1;
    let max_width = screen.width as usize;
    for choice in matches.iter() {
        let choice = choice.slice_chars(0, min(max_width, choice.chars().count()));
        if i == index + 1 {
            screen.write_inverted(&choice);
        } else {
            screen.write(&choice);
        }
        if i >= screen.visible_choices as usize {
            return;
        } else {
            screen.write("\n");
        }
        i += 1;
    }
}

fn read_choices() -> Vec<String> {
    let mut stdin = old_io::stdio::stdin();
    let mut lines = Vec::new();

    while let Ok(mut s) = stdin.read_line() {
        trim(&mut s);
        lines.push(s);
    }

    lines
}

pub fn trim(s: &mut String) {
    while let Some(x) = s.pop() {
        if x != '\n' {
            s.push(x);
            return;
        }
    }
}

fn delete_last_word(s: &mut String) {
    let mut deleted_something = false;
    while let Some(x) = s.pop() {
        if x == ' ' {
            if deleted_something {
                s.push(x);
                return;
            }
        } else {
            deleted_something = true;
        }
    }
}

#[test]
fn trim_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        trim(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("\n", "");
    should_become("\n\n", "");
    should_become("asdf", "asdf");
    should_become("asdf\n", "asdf");
    should_become("asdf\nasdf\n", "asdf\nasdf");
}

#[test]
fn delete_word_test() {
    fn should_become(before: &str, after: &str) {
        let mut x = before.to_string();
        delete_last_word(&mut x);
        assert_eq!(after.to_string(), x);
    }
    should_become("", "");
    should_become("a", "");
    should_become("asdf", "");
    should_become("asdf asdf asdf", "asdf asdf ");
    should_become("asdf asdf asdf ", "asdf asdf ");
    should_become("asdf asdf asdf  ", "asdf asdf ");
}
<|endoftext|>


<commit_before>use csv;
use indicatif;
use lazy_static;
use prettytable;
use regex;
use rusqlite;
use rustyline;
use std::env;
use std::error::Error;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).unwrap();
        while let Some(row) = records.next() {
            stmt.execute(&row?).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.description().to_owned())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = prettytable::Table::new();
    let mut title_row = prettytable::Row::new(vec![]);
    for col in stmt.column_names() {
        title_row.add_cell(prettytable::Cell::new(col));
    }
    table.set_titles(title_row);
    table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = prettytable::Row::new(vec![]);
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(prettytable::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    table.printstd();
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap())?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path)?;
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}
<commit_msg>Fix for deprecation of Error::description (#253)<commit_after>use csv;
use indicatif;
use lazy_static;
use prettytable;
use regex;
use rusqlite;
use rustyline;
use std::env;
use std::fs::File;

fn _normalize_col(col: &str) -> String {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"\(.*?\)").unwrap();
    }
    RE.replace_all(col, "")
        .to_lowercase()
        .trim()
        .replace(" ", "_")
        .replace(".", "_")
        .replace("-", "_")
        .replace("/", "_")
        .replace("?", "")
}

fn _create_table(db: &mut rusqlite::Connection, table_name: &str, cols: &[String]) {
    let create_columns = cols
        .iter()
        .map(|c| format!("{} varchar", c))
        .collect::<Vec<String>>()
        .join(", ");
    db.execute(
        &format!("CREATE TABLE {} ({})", table_name, create_columns),
        &[] as &[&dyn rusqlite::types::ToSql],
    )
    .unwrap();
}

fn _load_table_from_path(
    db: &mut rusqlite::Connection,
    table_name: &str,
    path: String,
) -> Result<Vec<String>, Box<dyn std::error::Error>> {
    let mut num_rows = 0;
    let f = File::open(path)?;
    let file_size = f.metadata().unwrap().len();
    let mut reader = csv::Reader::from_reader(f);

    let normalized_cols =
        reader
            .headers()?
            .iter()
            .map(_normalize_col)
            .fold(vec![], |mut v, orig_col| {
                let mut col = orig_col.clone();
                let mut i = 1;
                while v.contains(&col) {
                    col = format!("{}_{}", orig_col, i);
                    i += 1
                }
                v.push(col);
                v
            });
    _create_table(db, table_name, &normalized_cols);

    let insert_query = format!(
        "INSERT INTO {} VALUES ({})",
        table_name,
        normalized_cols
            .iter()
            .map(|_| "?")
            .collect::<Vec<_>>()
            .join(", ")
    );
    let pb = indicatif::ProgressBar::new(file_size);
    pb.set_style(
        indicatif::ProgressStyle::default_bar()
            .template("[{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})")
            .progress_chars("#>-"),
    );
    let mut records = reader.records();
    let tx = db.transaction().unwrap();
    {
        let mut stmt = tx.prepare(&insert_query).unwrap();
        while let Some(row) = records.next() {
            stmt.execute(&row?).unwrap();

            num_rows += 1;
            if num_rows % 10000 == 0 {
                pb.set_position(records.reader().position().byte())
            }
        }
    }
    tx.commit().unwrap();
    pb.finish();

    println!(
        "Loaded {} rows into {}({})",
        num_rows,
        table_name,
        normalized_cols.join(", "),
    );
    Ok(normalized_cols)
}

struct FromAnySqlType {
    value: String,
}

impl rusqlite::types::FromSql for FromAnySqlType {
    fn column_result(
        value: rusqlite::types::ValueRef<'_>,
    ) -> Result<FromAnySqlType, rusqlite::types::FromSqlError> {
        let result = match value {
            rusqlite::types::ValueRef::Null => "null".to_string(),
            rusqlite::types::ValueRef::Integer(v) => v.to_string(),
            rusqlite::types::ValueRef::Real(v) => v.to_string(),
            rusqlite::types::ValueRef::Blob(v) | rusqlite::types::ValueRef::Text(v) => {
                String::from_utf8(v.to_vec()).unwrap()
            }
        };
        Ok(FromAnySqlType { value: result })
    }
}

fn _prepare_query<'a>(
    conn: &'a mut rusqlite::Connection,
    query: &str,
) -> Result<rusqlite::Statement<'a>, String> {
    conn.prepare(&query).map_err(|e| e.to_string())
}

fn _handle_query(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    let mut stmt = _prepare_query(conn, line)?;

    let mut table = prettytable::Table::new();
    let mut title_row = prettytable::Row::new(vec![]);
    for col in stmt.column_names() {
        title_row.add_cell(prettytable::Cell::new(col));
    }
    table.set_titles(title_row);
    table.set_format(*prettytable::format::consts::FORMAT_NO_LINESEP_WITH_TITLE);

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        let mut row = prettytable::Row::new(vec![]);
        for i in 0..r.column_count() {
            let cell: FromAnySqlType = r.get(i).unwrap();
            row.add_cell(prettytable::Cell::new(&cell.value));
        }
        table.add_row(row);
    }
    table.printstd();
    Ok(())
}

fn _handle_export(conn: &mut rusqlite::Connection, line: &str) -> Result<(), String> {
    lazy_static::lazy_static! {
        static ref RE: regex::Regex = regex::Regex::new(r"^\.export\(([\w_\-\./]+)\) (.*)").unwrap();
    }
    let caps = RE
        .captures(line)
        .ok_or_else(|| "Must match `.export(file-name) SQL`".to_owned())?;
    let destination_path = &caps[1];
    let query = &caps[2];

    let mut stmt = _prepare_query(conn, query)?;

    let mut writer = csv::Writer::from_path(destination_path).unwrap();
    writer.write_record(stmt.column_names()).unwrap();

    let mut results = stmt.query(&[] as &[&dyn rusqlite::types::ToSql]).unwrap();
    while let Ok(Some(r)) = results.next() {
        writer
            .write_record((0..r.column_count()).map(|i| {
                let cell: FromAnySqlType = r.get(i).unwrap();
                cell.value
            }))
            .unwrap();
    }

    Ok(())
}

fn _process_query(conn: &mut rusqlite::Connection, line: &str) {
    let result = if line.starts_with(".export") {
        _handle_export(conn, line)
    } else {
        _handle_query(conn, line)
    };
    if let Err(e) = result {
        println!("{}", e);
    }
}

struct SimpleWordCompleter {
    words: Vec<String>,
}

static BREAK_CHARS: [u8; 4] = [b' ', b'(', b')', b','];
impl SimpleWordCompleter {
    fn new(words: Vec<String>) -> SimpleWordCompleter {
        SimpleWordCompleter { words }
    }
}

impl rustyline::Helper for SimpleWordCompleter {}

impl rustyline::hint::Hinter for SimpleWordCompleter {
    fn hint(&self, _line: &str, _pos: usize, _ctx: &rustyline::Context<'_>) -> Option<String> {
        None
    }
}

impl rustyline::highlight::Highlighter for SimpleWordCompleter {}

impl rustyline::completion::Completer for SimpleWordCompleter {
    type Candidate = String;

    fn complete(
        &self,
        line: &str,
        pos: usize,
        _ctx: &rustyline::Context<'_>,
    ) -> rustyline::Result<(usize, Vec<String>)> {
        let (start, word) = rustyline::completion::extract_word(line, pos, None, &BREAK_CHARS);

        let matches = self
            .words
            .iter()
            .filter(|w| w.starts_with(word))
            .cloned()
            .collect();
        Ok((start, matches))
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut paths = env::args().skip(1);

    let mut conn = rusqlite::Connection::open_in_memory().unwrap();

    let mut base_words = [
        "distinct", "select", "from", "group", "by", "order", "where", "count", "limit", "offset",
        ".export",
    ]
    .iter()
    .map(|&s| s.to_string())
    .collect::<Vec<String>>();

    if paths.len() == 1 {
        let mut col_names = _load_table_from_path(&mut conn, "t", paths.next().unwrap())?;
        base_words.append(&mut col_names);
    } else {
        for (idx, path) in paths.enumerate() {
            let mut col_names = _load_table_from_path(&mut conn, &format!("t{}", idx + 1), path)?;
            base_words.append(&mut col_names);
        }
    }

    let completer = SimpleWordCompleter::new(base_words);
    let mut rl = rustyline::Editor::new();
    rl.set_helper(Some(completer));
    loop {
        match rl.readline("> ") {
            Ok(line) => {
                if line.trim().is_empty() {
                    continue;
                }
                _process_query(&mut conn, &line);
                rl.add_history_entry(line);
            }
            Err(rustyline::error::ReadlineError::Interrupted) => {
                println!("Interrupted");
                continue;
            }
            Err(rustyline::error::ReadlineError::Eof) => {
                break;
            }
            Err(err) => {
                println!("Error: {}", err);
                break;
            }
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate hyperlocal;

use hyperlocal::{DomainUrl, SocketConnector};

fn main() {
    let client = hyper::Client::with_connector(hyperlocal::SocketConnector);
    let mut res = client.get(DomainUrl::new("/var/run/docker.sock", "info")).send().unwrap();
    std::io::copy(&mut res, &mut std::io::stdout()).unwrap();
}
<commit_msg>expect a slash<commit_after>extern crate hyper;
extern crate hyperlocal;

use hyperlocal::{DomainUrl, SocketConnector};

fn main() {
    let client = hyper::Client::with_connector(hyperlocal::SocketConnector);
    let mut res = client.get(DomainUrl::new("/var/run/docker.sock", "/info")).send().unwrap();
    std::io::copy(&mut res, &mut std::io::stdout()).unwrap();
}
<|endoftext|>


<commit_before>//! Ruma is a server for Matrix.org's client-server API.

#![feature(custom_attribute, custom_derive, plugin)]
#![plugin(diesel_codegen)]
#![plugin(serde_macros)]
#![cfg_attr(feature="clippy", plugin(clippy))]

extern crate argon2rs;
extern crate base64;
extern crate bodyparser;
extern crate chrono;
extern crate clap;
extern crate env_logger;
#[macro_use] extern crate diesel;
#[macro_use] extern crate iron;
#[cfg(test)] extern crate iron_test;
#[macro_use] extern crate log;
extern crate macaroons;
extern crate mount;
extern crate plugin;
extern crate persistent;
extern crate r2d2;
extern crate r2d2_diesel;
extern crate rand;
extern crate router;
extern crate serde;
extern crate serde_json;
extern crate unicase;

use clap::{App, AppSettings, SubCommand};

use config::load;
use crypto::generate_macaroon_secret_key;
use server::Server;

pub mod access_token;
/// API endpoints as Iron handlers.
pub mod api {
    /// API endpoints for the 0.x.x version of the Matrix spec.
    pub mod r0 {
        pub mod login;
        pub mod registration;
        pub mod versions;
    }
}
pub mod authentication;
pub mod config;
pub mod crypto;
pub mod db;
pub mod error;
pub mod middleware;
pub mod modifier;
pub mod schema;
pub mod server;
pub mod swagger;
#[cfg(test)] pub mod test;
pub mod user;

fn main() {
    env_logger::init().expect("Failed to initialize logger.");

    let matches = App::new("ruma")
        .version(env!("CARGO_PKG_VERSION"))
        .about("A Matrix client-server API")
        .setting(AppSettings::GlobalVersion)
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .subcommand(
            SubCommand::with_name("run")
                .about("Runs the Ruma server")
        )
        .subcommand(
            SubCommand::with_name("secret")
                .about("Generates a random value to be used as a macaroon secret key")
        )
        .get_matches();

    match matches.subcommand() {
        ("run", Some(_)) => {
            let config = match load("ruma.json") {
                Ok(config) => config,
                Err(error) => {
                    println!("Failed to load configuration file: {}", error);

                    return;
                }
            };

            match Server::new(&config) {
                Ok(server) => {
                    if let Err(error) = server.run() {
                        println!("{}", error);
                    }
                },
                Err(error) => {
                    println!("Failed to create server: {}", error);

                    return;
                }
            }
        }
        ("secret", Some(_)) => match generate_macaroon_secret_key() {
            Ok(key) => println!("{}", key),
            Err(error) => println!("Failed to generate macaroon secret key: {}", error),
        },
        _ => println!("{}", matches.usage()),
    };
}
<commit_msg>Remove plugin attribute for clippy.<commit_after>//! Ruma is a server for Matrix.org's client-server API.

#![feature(custom_attribute, custom_derive, plugin)]
#![plugin(diesel_codegen)]
#![plugin(serde_macros)]

extern crate argon2rs;
extern crate base64;
extern crate bodyparser;
extern crate chrono;
extern crate clap;
extern crate env_logger;
#[macro_use] extern crate diesel;
#[macro_use] extern crate iron;
#[cfg(test)] extern crate iron_test;
#[macro_use] extern crate log;
extern crate macaroons;
extern crate mount;
extern crate plugin;
extern crate persistent;
extern crate r2d2;
extern crate r2d2_diesel;
extern crate rand;
extern crate router;
extern crate serde;
extern crate serde_json;
extern crate unicase;

use clap::{App, AppSettings, SubCommand};

use config::load;
use crypto::generate_macaroon_secret_key;
use server::Server;

pub mod access_token;
/// API endpoints as Iron handlers.
pub mod api {
    /// API endpoints for the 0.x.x version of the Matrix spec.
    pub mod r0 {
        pub mod login;
        pub mod registration;
        pub mod versions;
    }
}
pub mod authentication;
pub mod config;
pub mod crypto;
pub mod db;
pub mod error;
pub mod middleware;
pub mod modifier;
pub mod schema;
pub mod server;
pub mod swagger;
#[cfg(test)] pub mod test;
pub mod user;

fn main() {
    env_logger::init().expect("Failed to initialize logger.");

    let matches = App::new("ruma")
        .version(env!("CARGO_PKG_VERSION"))
        .about("A Matrix client-server API")
        .setting(AppSettings::GlobalVersion)
        .setting(AppSettings::SubcommandRequiredElseHelp)
        .subcommand(
            SubCommand::with_name("run")
                .about("Runs the Ruma server")
        )
        .subcommand(
            SubCommand::with_name("secret")
                .about("Generates a random value to be used as a macaroon secret key")
        )
        .get_matches();

    match matches.subcommand() {
        ("run", Some(_)) => {
            let config = match load("ruma.json") {
                Ok(config) => config,
                Err(error) => {
                    println!("Failed to load configuration file: {}", error);

                    return;
                }
            };

            match Server::new(&config) {
                Ok(server) => {
                    if let Err(error) = server.run() {
                        println!("{}", error);
                    }
                },
                Err(error) => {
                    println!("Failed to create server: {}", error);

                    return;
                }
            }
        }
        ("secret", Some(_)) => match generate_macaroon_secret_key() {
            Ok(key) => println!("{}", key),
            Err(error) => println!("Failed to generate macaroon secret key: {}", error),
        },
        _ => println!("{}", matches.usage()),
    };
}
<|endoftext|>


<commit_before>#![feature(custom_derive, plugin)]
#![plugin(serde_macros)]

extern crate serde;
extern crate serde_json;

extern crate jsonrpc_core;

#[macro_use]
extern crate log;
extern crate env_logger;
extern crate gluon;

pub mod language_server;

use jsonrpc_core::{Error, ErrorCode, IoHandler, MethodCommand, NotificationCommand, Params, Value};
use serde_json::value::{from_value, to_value};

use gluon::base::ast;
use gluon::base::metadata::Metadata;
use gluon::base::symbol::Symbol;
use gluon::check::completion;
use gluon::import::{CheckImporter, Import};
use gluon::vm::internal::Value as GluonValue;
use gluon::vm::thread::ThreadInternal;
use gluon::{Compiler, Error as GluonError, Result as GluonResult, RootedThread, new_vm,
            filename_to_module};

use std::error::Error as StdError;
use std::io;
use std::io::{Read, Write};
use std::str;

use language_server::*;

struct Initialize(RootedThread);
impl MethodCommand for Initialize {
    fn execute(&self, _params: Params) -> Result<Value, Error> {
        let result = InitializeResult {
            capabilities: ServerCapabilities {
                text_document_sync: Some(TextDocumentSyncKind::Full),
                completion_provider: Some(CompletionOptions {
                    resolve_provider: Some(true),
                    trigger_characters: vec![".".into()],
                }),
            },
        };
        Ok(to_value(&result))
    }
}

struct Completion(RootedThread);
impl MethodCommand for Completion {
    fn execute(&self, params: Params) -> Result<Value, Error> {
        let change: TextDocumentPositionParams = match params {
            Params::Map(map) => {
                match from_value(Value::Object(map)) {
                    Ok(change) => change,
                    Err(_) => return Err(Error::invalid_params()),
                }
            }
            _ => return Err(Error::invalid_params()),
        };
        let thread = &self.0;
        let module = change.text_document.uri;
        let import = thread.get_macros().get("import").expect("Import macro");
        let import = import.downcast_ref::<Import<CheckImporter>>().expect("Check importer");
        let importer = import.importer.0.lock().unwrap();
        let expr = try!(importer.get(&module).ok_or_else(|| {
            Error {
                code: ErrorCode::InternalError,
                message: format!("Module `{}` is not defined", module),
                data: None,
            }
        }));
        let suggestions = completion::suggest(&ast::EmptyEnv::new(),
                                              expr,
                                              ast::Location {
                                                  row: change.position.line as i32,
                                                  column: change.position.character as i32,
                                                  absolute: 0,
                                              });
        let items: Vec<_> = suggestions.into_iter()
                                       .map(|symbol| {
                                           CompletionItem {
                                               label:
                                                   String::from(symbol.as_ref()
                                                                      .split(':')
                                                                      .next()
                                                                      .unwrap_or(symbol.as_ref())),
                                               detail: Some("comment".into()),
                                               kind: Some(CompletionItemKind::Variable),
                                               ..CompletionItem::default()
                                           }
                                       })
                                       .collect();
        Ok(to_value(&items))
    }
}

fn location_to_position(loc: &ast::Location) -> Position {
    Position {
        line: loc.row as u64,
        character: loc.column as u64,
    }
}
fn span_to_range(span: &ast::Span) -> Range {
    Range {
        start: location_to_position(&span.start),
        end: location_to_position(&span.end),
    }
}

struct TextDocumentDidChange(RootedThread);
impl NotificationCommand for TextDocumentDidChange {
    fn execute(&self, params: Params) {
        let change: DidChangeTextDocumentParams = match params {
            Params::Map(map) => {
                match from_value(Value::Object(map)) {
                    Ok(change) => change,
                    Err(_) => return,
                }
            }
            _ => return,
        };
        let diagnostics = match self.execute2(&change) {
            Ok(_) => vec![],
            Err(err) => {
                match err {
                    GluonError::Typecheck(err) => {
                        err.errors()
                           .errors
                           .into_iter()
                           .map(|err| {
                               Diagnostic {
                                   message: format!("{}", err.value),
                                   severity: Some(DiagnosticSeverity::Error),
                                   range: span_to_range(&err.span),
                                   ..Diagnostic::default()
                               }
                           })
                           .collect()
                    }
                    err => {
                        vec![Diagnostic {
                                 message: format!("{}", err),
                                 severity: Some(DiagnosticSeverity::Error),
                                 ..Diagnostic::default()
                             }]
                    }
                }
            }
        };
        let filename = &change.text_document.uri;
        let r = format!(r#"{{
                            "jsonrpc": "2.0",
                            "method": "textDocument/publishDiagnostics",
                            "params": {}
                        }}"#,
                        to_value(&PublishDiagnosticsParams {
            uri: filename.clone(),
            diagnostics: diagnostics,
        }));
        print!("Content-Length: {}\r\n\r\n{}", r.len(), r);
    }
}

impl TextDocumentDidChange {
    fn execute2(&self, change: &DidChangeTextDocumentParams) -> GluonResult<()> {
        use gluon::compiler_pipeline::*;

        let thread = &self.0;
        let filename = &change.text_document.uri;
        let fileinput = &change.content_changes[0].text;
        let name = filename_to_module(filename);
        let mut compiler = Compiler::new();
        let MacroValue(mut expr) = try!(fileinput.expand_macro(&mut compiler, thread, &name));
        let result = match compiler.typecheck_expr(thread, &name, fileinput, &mut expr) {
            Ok(typ) => {
                let metadata = Metadata::default();
                try!(thread.global_env()
                           .set_global(Symbol::new(filename), typ, metadata, GluonValue::Int(0)));
                Ok(())
            }
            Err(err) => Err(err),
        };
        let import = thread.get_macros().get("import").expect("Import macro");
        let import = import.downcast_ref::<Import<CheckImporter>>()
                           .expect("Check importer");
        let mut importer = import.importer.0.lock().unwrap();
        importer.insert(filename.clone(), expr);
        result
    }
}

fn log_message(message: String) {
    let r = format!(r#"{{"jsonrpc": "2.0", "method": "window/logMessage", "params": {} }}"#,
                    to_value(&LogMessageParams {
                        typ: MessageType::Log,
                        message: message,
                    }));
    print!("Content-Length: {}\r\n\r\n{}", r.len(), r);
}

fn main_loop(io: &mut IoHandler) -> Result<(), Box<StdError>> {
    let stdin = io::stdin();
    loop {
        let mut header = String::new();
        let n = try!(stdin.read_line(&mut header));
        if n == 0 {
            // EOF
            return Ok(());
        }
        debug!("{}", header);
        if header.starts_with("Content-Length: ") {
            let content_length = {
                let len = header["Content-Length:".len()..].trim();
                debug!("{}", len);
                try!(len.parse::<usize>())
            };
            while header != "\r\n" {
                header.clear();
                try!(io::stdin().read_line(&mut header));
            }
            let mut content = vec![0; content_length];
            try!(stdin.lock().read_exact(&mut content));
            let json = try!(str::from_utf8(&content));
            if let Some(response) = io.handle_request(json) {
                print!("Content-Length: {}\r\n\r\n{}", response.len(), response);
                try!(io::stdout().flush());
            }
        }
    }
}

fn main() {
    ::env_logger::init().unwrap();
    let handle = ::std::thread::spawn(|| {
        let thread = new_vm();
        let import = Import::new(CheckImporter::new());
        thread.get_macros().insert("import".into(), import);

        let mut io = IoHandler::new();
        io.add_method("initialize", Initialize(thread.clone()));
        io.add_method("textDocument/completion", Completion(thread.clone()));
        io.add_notification("textDocument/didChange", TextDocumentDidChange(thread));

        main_loop(&mut io).unwrap();
    });
    if let Err(err) = handle.join() {
        let msg = err.downcast_ref::<&'static str>()
                     .cloned()
                     .or_else(|| err.downcast_ref::<String>().map(|s| &s[..]))
                     .unwrap_or("Any");
        log_message(format!("Panic: `{}`", msg));
    }
}
<commit_msg>fix: Line numbers from visual code start from 0 while gluon starts from 1<commit_after>#![feature(custom_derive, plugin)]
#![plugin(serde_macros)]

extern crate serde;
extern crate serde_json;

extern crate jsonrpc_core;

#[macro_use]
extern crate log;
extern crate env_logger;
extern crate gluon;

pub mod language_server;

use jsonrpc_core::{Error, ErrorCode, IoHandler, MethodCommand, NotificationCommand, Params, Value};
use serde_json::value::{from_value, to_value};

use gluon::base::ast;
use gluon::base::metadata::Metadata;
use gluon::base::symbol::Symbol;
use gluon::check::completion;
use gluon::import::{CheckImporter, Import};
use gluon::vm::internal::Value as GluonValue;
use gluon::vm::thread::ThreadInternal;
use gluon::{Compiler, Error as GluonError, Result as GluonResult, RootedThread, new_vm,
            filename_to_module};

use std::error::Error as StdError;
use std::io;
use std::io::{Read, Write};
use std::str;

use language_server::*;

struct Initialize(RootedThread);
impl MethodCommand for Initialize {
    fn execute(&self, _params: Params) -> Result<Value, Error> {
        let result = InitializeResult {
            capabilities: ServerCapabilities {
                text_document_sync: Some(TextDocumentSyncKind::Full),
                completion_provider: Some(CompletionOptions {
                    resolve_provider: Some(true),
                    trigger_characters: vec![".".into()],
                }),
            },
        };
        Ok(to_value(&result))
    }
}

struct Completion(RootedThread);
impl MethodCommand for Completion {
    fn execute(&self, params: Params) -> Result<Value, Error> {
        let change: TextDocumentPositionParams = match params {
            Params::Map(map) => {
                match from_value(Value::Object(map)) {
                    Ok(change) => change,
                    Err(_) => return Err(Error::invalid_params()),
                }
            }
            _ => return Err(Error::invalid_params()),
        };
        let thread = &self.0;
        let module = change.text_document.uri;
        let import = thread.get_macros().get("import").expect("Import macro");
        let import = import.downcast_ref::<Import<CheckImporter>>().expect("Check importer");
        let importer = import.importer.0.lock().unwrap();
        let expr = try!(importer.get(&module).ok_or_else(|| {
            Error {
                code: ErrorCode::InternalError,
                message: format!("Module `{}` is not defined", module),
                data: None,
            }
        }));
        let suggestions = completion::suggest(&ast::EmptyEnv::new(),
                                              expr,
                                              ast::Location {
                                                  row: (change.position.line + 1) as i32,
                                                  column: change.position.character as i32,
                                                  absolute: 0,
                                              });
        let items: Vec<_> = suggestions.into_iter()
                                       .map(|symbol| {
                                           CompletionItem {
                                               label:
                                                   String::from(symbol.as_ref()
                                                                      .split(':')
                                                                      .next()
                                                                      .unwrap_or(symbol.as_ref())),
                                               detail: Some("comment".into()),
                                               kind: Some(CompletionItemKind::Variable),
                                               ..CompletionItem::default()
                                           }
                                       })
                                       .collect();
        Ok(to_value(&items))
    }
}

fn location_to_position(loc: &ast::Location) -> Position {
    Position {
        line: loc.row as u64 + 1,
        character: loc.column as u64,
    }
}
fn span_to_range(span: &ast::Span) -> Range {
    Range {
        start: location_to_position(&span.start),
        end: location_to_position(&span.end),
    }
}

struct TextDocumentDidChange(RootedThread);
impl NotificationCommand for TextDocumentDidChange {
    fn execute(&self, params: Params) {
        let change: DidChangeTextDocumentParams = match params {
            Params::Map(map) => {
                match from_value(Value::Object(map)) {
                    Ok(change) => change,
                    Err(_) => return,
                }
            }
            _ => return,
        };
        let diagnostics = match self.execute2(&change) {
            Ok(_) => vec![],
            Err(err) => {
                match err {
                    GluonError::Typecheck(err) => {
                        err.errors()
                           .errors
                           .into_iter()
                           .map(|err| {
                               Diagnostic {
                                   message: format!("{}", err.value),
                                   severity: Some(DiagnosticSeverity::Error),
                                   range: span_to_range(&err.span),
                                   ..Diagnostic::default()
                               }
                           })
                           .collect()
                    }
                    err => {
                        vec![Diagnostic {
                                 message: format!("{}", err),
                                 severity: Some(DiagnosticSeverity::Error),
                                 ..Diagnostic::default()
                             }]
                    }
                }
            }
        };
        let filename = &change.text_document.uri;
        let r = format!(r#"{{
                            "jsonrpc": "2.0",
                            "method": "textDocument/publishDiagnostics",
                            "params": {}
                        }}"#,
                        to_value(&PublishDiagnosticsParams {
            uri: filename.clone(),
            diagnostics: diagnostics,
        }));
        print!("Content-Length: {}\r\n\r\n{}", r.len(), r);
    }
}

impl TextDocumentDidChange {
    fn execute2(&self, change: &DidChangeTextDocumentParams) -> GluonResult<()> {
        use gluon::compiler_pipeline::*;

        let thread = &self.0;
        let filename = &change.text_document.uri;
        let fileinput = &change.content_changes[0].text;
        let name = filename_to_module(filename);
        let mut compiler = Compiler::new();
        let MacroValue(mut expr) = try!(fileinput.expand_macro(&mut compiler, thread, &name));
        let result = match compiler.typecheck_expr(thread, &name, fileinput, &mut expr) {
            Ok(typ) => {
                let metadata = Metadata::default();
                try!(thread.global_env()
                           .set_global(Symbol::new(filename), typ, metadata, GluonValue::Int(0)));
                Ok(())
            }
            Err(err) => Err(err),
        };
        let import = thread.get_macros().get("import").expect("Import macro");
        let import = import.downcast_ref::<Import<CheckImporter>>()
                           .expect("Check importer");
        let mut importer = import.importer.0.lock().unwrap();
        importer.insert(filename.clone(), expr);
        result
    }
}

fn log_message(message: String) {
    let r = format!(r#"{{"jsonrpc": "2.0", "method": "window/logMessage", "params": {} }}"#,
                    to_value(&LogMessageParams {
                        typ: MessageType::Log,
                        message: message,
                    }));
    print!("Content-Length: {}\r\n\r\n{}", r.len(), r);
}

fn main_loop(io: &mut IoHandler) -> Result<(), Box<StdError>> {
    let stdin = io::stdin();
    loop {
        let mut header = String::new();
        let n = try!(stdin.read_line(&mut header));
        if n == 0 {
            // EOF
            return Ok(());
        }
        debug!("{}", header);
        if header.starts_with("Content-Length: ") {
            let content_length = {
                let len = header["Content-Length:".len()..].trim();
                debug!("{}", len);
                try!(len.parse::<usize>())
            };
            while header != "\r\n" {
                header.clear();
                try!(io::stdin().read_line(&mut header));
            }
            let mut content = vec![0; content_length];
            try!(stdin.lock().read_exact(&mut content));
            let json = try!(str::from_utf8(&content));
            if let Some(response) = io.handle_request(json) {
                print!("Content-Length: {}\r\n\r\n{}", response.len(), response);
                try!(io::stdout().flush());
            }
        }
    }
}

fn main() {
    ::env_logger::init().unwrap();
    let handle = ::std::thread::spawn(|| {
        let thread = new_vm();
        let import = Import::new(CheckImporter::new());
        thread.get_macros().insert("import".into(), import);

        let mut io = IoHandler::new();
        io.add_method("initialize", Initialize(thread.clone()));
        io.add_method("textDocument/completion", Completion(thread.clone()));
        io.add_notification("textDocument/didChange", TextDocumentDidChange(thread));

        main_loop(&mut io).unwrap();
    });
    if let Err(err) = handle.join() {
        let msg = err.downcast_ref::<&'static str>()
                     .cloned()
                     .or_else(|| err.downcast_ref::<String>().map(|s| &s[..]))
                     .unwrap_or("Any");
        log_message(format!("Panic: `{}`", msg));
    }
}
<|endoftext|>