<commit_before>#![crate_id = "benchmark#0.1"]
#![crate_type = "lib"]

extern crate getopts;
extern crate timer;
use std::{result, os};
use std::rand::random;
use std::iter::AdditiveIterator;
use getopts::{getopts, optflag, optflagmulti, optopt, usage};
use timer::Timer;

pub struct Benchmark {
    num_trials: uint,
    trial_size: uint,
    quiet: u8,
    parse_args: bool,
    verify: bool
}

impl Benchmark {
    pub fn new () -> Benchmark {
        Benchmark {
            num_trials: 1,
            trial_size: 10,
            quiet: 0,
            parse_args: true,
            verify: false
        }
    }

    fn parse_opts(&mut self) {
        if self.parse_args {
            let args = os::args();
            let opts = ~[
                optflagmulti("q", "quiet", "Provide less output"),
                optopt("", "trialsize", "Number elements to sort in each trial", ""),
                optopt("", "numtrials", "Number of trials to perform", ""),
                optflag("", "verify", "Verify that the sort was correct"),
                optflag("h", "help", "Show this help")
                ];
            let matches = match getopts(args.tail(), opts) {
                result::Ok(m) => { m }
                result::Err(f) => { fail!(f.to_str()) }
            };
            if matches.opt_present("h") || matches.opt_present("help") {
                let brief = format!("Usage: {} [options]", args.as_slice().head().map(|x| x.as_slice()).unwrap_or(""));
                print!("{}", usage(brief.as_slice(), opts));
                self.num_trials = 0;
                self.parse_args = false;
                return;
            }
            if matches.opt_present("q") {
                self.quiet = matches.opt_count("q") as u8;
                if self.quiet < 1{
                    self.quiet = 1;
                }
            }
            if matches.opt_present("verify") {
                self.verify = true;
            }

            match matches.opt_str("trialsize") {
                Some(size) => {
                    match from_str::<uint>(size.as_slice()) {
                        Some(ts) => { self.trial_size = ts }
                        None => { fail!("Trial size must be an integer") }
                    }
                }
                None => {}
            }

            match matches.opt_str("numtrials") {
                Some(trials) => {
                    match from_str::<uint>(trials.as_slice()) {
                        Some(t) => { self.num_trials = t }
                        None => { fail!("Number of trials must be an integer") }
                    }
                }
                None => {}
            }

            self.parse_args = false;
        }
    }

    pub fn run(&mut self, sort: fn(Vec<uint>) -> Vec<uint>) {
        self.parse_opts();
        let mut timer = Timer::new();
        let mut sort_times = Vec::with_capacity(self.num_trials);

        for trial_number in range(0, self.num_trials) {
            let vals = generate_random_array(self.trial_size);
            /* Run the sort and record the timing */
            match self.quiet {
                0 => { println!("Starting sort ..."); }
                1 => { println!("Trial {}", trial_number); }
                _ => {}
            }

            timer.start();
            let sorted = sort(vals);
            timer.end();

            match self.quiet {
                0 => { println!("Sort finished."); }
                _ => {}
            }

            if self.verify {
                /* Check that it actually is sorted */
                match self.quiet {
                    0 => { println!("Verifying sort ..."); }
                    _ => {}
                }
                if !ensure_sorted(sorted.as_slice()) {
                    /* Print the values so we can see what they actually look like.
                       Note: Should probably only do this if the array is small */
                    for v in sorted.iter() {
                        println!("{}", *v as uint);
                    }
                    fail!(format!("Trial {}: Array was not sorted correctly", trial_number));
                }
                match self.quiet {
                    0 => { println!("Sort was correct."); }
                    _ => {}
                }
            }

            /* Show the time it took */
            match self.quiet {
                0 => { timer.show_time(); }
                _ => {}
            }
            /* Record the time it took */
            sort_times.push(timer.get_total_time());
        }

        if self.num_trials > 0 {
            /* Print out the average time at the end */
            let total_time = sort_times.iter().map(|&x| x).sum();
            let average_time = total_time / (self.num_trials as u64);
            println!("Average time: {}", timer::format_as_time(average_time));
        }
    }
}

pub fn generate_random_array(size: uint) -> Vec<uint> {
    let mut ret = Vec::with_capacity(size);
    for _ in range(0, size) {
        ret.push(random());
    }
    ret
}

fn ensure_sorted(arr: &[uint]) -> bool {
    let mut previous_value = 0;
    for v in arr.iter() {
        if *v < previous_value {
            return false
        }
        previous_value = *v;
    }
    true
}

<commit_msg>Remove final use of ~[]<commit_after>#![crate_id = "benchmark#0.1"]
#![crate_type = "lib"]

extern crate getopts;
extern crate timer;
use std::{result, os};
use std::rand::random;
use std::iter::AdditiveIterator;
use getopts::{getopts, optflag, optflagmulti, optopt, usage};
use timer::Timer;

pub struct Benchmark {
    num_trials: uint,
    trial_size: uint,
    quiet: u8,
    parse_args: bool,
    verify: bool
}

impl Benchmark {
    pub fn new () -> Benchmark {
        Benchmark {
            num_trials: 1,
            trial_size: 10,
            quiet: 0,
            parse_args: true,
            verify: false
        }
    }

    fn parse_opts(&mut self) {
        if self.parse_args {
            let args = os::args();
            let opts = &[
                optflagmulti("q", "quiet", "Provide less output"),
                optopt("", "trialsize", "Number elements to sort in each trial", ""),
                optopt("", "numtrials", "Number of trials to perform", ""),
                optflag("", "verify", "Verify that the sort was correct"),
                optflag("h", "help", "Show this help")
                ];
            let matches = match getopts(args.tail(), opts) {
                result::Ok(m) => { m }
                result::Err(f) => { fail!(f.to_str()) }
            };
            if matches.opt_present("h") || matches.opt_present("help") {
                let brief = format!("Usage: {} [options]", args.as_slice().head().map(|x| x.as_slice()).unwrap_or(""));
                print!("{}", usage(brief.as_slice(), opts));
                self.num_trials = 0;
                self.parse_args = false;
                return;
            }
            if matches.opt_present("q") {
                self.quiet = matches.opt_count("q") as u8;
                if self.quiet < 1{
                    self.quiet = 1;
                }
            }
            if matches.opt_present("verify") {
                self.verify = true;
            }

            match matches.opt_str("trialsize") {
                Some(size) => {
                    match from_str::<uint>(size.as_slice()) {
                        Some(ts) => { self.trial_size = ts }
                        None => { fail!("Trial size must be an integer") }
                    }
                }
                None => {}
            }

            match matches.opt_str("numtrials") {
                Some(trials) => {
                    match from_str::<uint>(trials.as_slice()) {
                        Some(t) => { self.num_trials = t }
                        None => { fail!("Number of trials must be an integer") }
                    }
                }
                None => {}
            }

            self.parse_args = false;
        }
    }

    pub fn run(&mut self, sort: fn(Vec<uint>) -> Vec<uint>) {
        self.parse_opts();
        let mut timer = Timer::new();
        let mut sort_times = Vec::with_capacity(self.num_trials);

        for trial_number in range(0, self.num_trials) {
            let vals = generate_random_array(self.trial_size);
            /* Run the sort and record the timing */
            match self.quiet {
                0 => { println!("Starting sort ..."); }
                1 => { println!("Trial {}", trial_number); }
                _ => {}
            }

            timer.start();
            let sorted = sort(vals);
            timer.end();

            match self.quiet {
                0 => { println!("Sort finished."); }
                _ => {}
            }

            if self.verify {
                /* Check that it actually is sorted */
                match self.quiet {
                    0 => { println!("Verifying sort ..."); }
                    _ => {}
                }
                if !ensure_sorted(sorted.as_slice()) {
                    /* Print the values so we can see what they actually look like.
                       Note: Should probably only do this if the array is small */
                    for v in sorted.iter() {
                        println!("{}", *v as uint);
                    }
                    fail!(format!("Trial {}: Array was not sorted correctly", trial_number));
                }
                match self.quiet {
                    0 => { println!("Sort was correct."); }
                    _ => {}
                }
            }

            /* Show the time it took */
            match self.quiet {
                0 => { timer.show_time(); }
                _ => {}
            }
            /* Record the time it took */
            sort_times.push(timer.get_total_time());
        }

        if self.num_trials > 0 {
            /* Print out the average time at the end */
            let total_time = sort_times.iter().map(|&x| x).sum();
            let average_time = total_time / (self.num_trials as u64);
            println!("Average time: {}", timer::format_as_time(average_time));
        }
    }
}

pub fn generate_random_array(size: uint) -> Vec<uint> {
    let mut ret = Vec::with_capacity(size);
    for _ in range(0, size) {
        ret.push(random());
    }
    ret
}

fn ensure_sorted(arr: &[uint]) -> bool {
    let mut previous_value = 0;
    for v in arr.iter() {
        if *v < previous_value {
            return false
        }
        previous_value = *v;
    }
    true
}

<|endoftext|>


<commit_before>// Higher level OpenCL wrappers.

use CL::*;
use CL::ll::*;

struct Platform {
    id: cl_platform_id
}

pub fn get_platforms() -> ~[Platform] {
    let mut num_platforms = 0;

    // TODO: Check result status
    clGetPlatformIDs(0, ptr::null(), ptr::addr_of(&num_platforms));

    let ids = vec::to_mut(vec::from_elem(num_platforms as uint,
                                         0 as cl_platform_id));

    do vec::as_imm_buf(ids) |ids, len| {
        clGetPlatformIDs(len as cl_uint,
                         ids, ptr::addr_of(&num_platforms))
    };

    do ids.map |id| { Platform { id: *id } }
}

struct Device {
    id: cl_device_id
}

pub fn get_devices(platform: Platform) -> ~[Device] {
    let mut num_devices = 0;

    clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, 0, ptr::null(),
                   ptr::addr_of(&num_devices));

    let ids = vec::to_mut(vec::from_elem(num_devices as uint,
                                         0 as cl_device_id));
    do vec::as_imm_buf(ids) |ids, len| {
        clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, len as cl_uint,
                       ids, ptr::addr_of(&num_devices));
    };

    do ids.map |id| { Device {id: *id }}
}

struct Context {
    ctx: cl_context,

    drop {
        clReleaseContext(self.ctx);
    }
}

pub fn create_context(device: Device) -> Context {
    // TODO: Support for multiple devices
    let mut errcode = 0;

    // TODO: Proper error messages
    let ctx = clCreateContext(ptr::null(), 1, ptr::addr_of(&device.id),
                              ptr::null(), ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create opencl context!";
    }

    Context { ctx: ctx };
}

struct CommandQueue {
    cqueue: cl_command_queue,

    drop {
        clReleaseCommandQueue(self.cqueue);
    }
}

pub fn create_commandqueue(ctx: Context, device: Device) -> CommandQueue {
    let mut errcode = 0;

    let cqueue = clCreateCommandQueue(ctx.ctx, device.id, 0, ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create command queue!";
    }

    CommandQueue { cqueue: cqueue };
}
<commit_msg>Fixed clCreateContext function call issue.<commit_after>// Higher level OpenCL wrappers.

use CL::*;
use CL::ll::*;

struct Platform {
    id: cl_platform_id
}

pub fn get_platforms() -> ~[Platform] {
    let mut num_platforms = 0;

    // TODO: Check result status
    clGetPlatformIDs(0, ptr::null(), ptr::addr_of(&num_platforms));

    let ids = vec::to_mut(vec::from_elem(num_platforms as uint,
                                         0 as cl_platform_id));

    do vec::as_imm_buf(ids) |ids, len| {
        clGetPlatformIDs(len as cl_uint,
                         ids, ptr::addr_of(&num_platforms))
    };

    do ids.map |id| { Platform { id: *id } }
}

struct Device {
    id: cl_device_id
}

pub fn get_devices(platform: Platform) -> ~[Device] {
    let mut num_devices = 0;

    clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, 0, ptr::null(),
                   ptr::addr_of(&num_devices));

    let ids = vec::to_mut(vec::from_elem(num_devices as uint,
                                         0 as cl_device_id));
    do vec::as_imm_buf(ids) |ids, len| {
        clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, len as cl_uint,
                       ids, ptr::addr_of(&num_devices));
    };

    do ids.map |id| { Device {id: *id }}
}

struct Context {
    ctx: cl_context,

    drop {
        clReleaseContext(self.ctx);
    }
}

pub fn create_context(device: Device) -> Context {
    // TODO: Support for multiple devices
    let mut errcode = 0;

    // TODO: Proper error messages
    let ctx = clCreateContext(ptr::null(), 1, ptr::addr_of(&device.id),
                              ptr::null(), ptr::null(), ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create opencl context!";
    }

    Context { ctx: ctx };
}

struct CommandQueue {
    cqueue: cl_command_queue,

    drop {
        clReleaseCommandQueue(self.cqueue);
    }
}

pub fn create_commandqueue(ctx: Context, device: Device) -> CommandQueue {
    let mut errcode = 0;

    let cqueue = clCreateCommandQueue(ctx.ctx, device.id, 0, ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create command queue!";
    }

    CommandQueue { cqueue: cqueue };
}
<|endoftext|>


<commit_before>// Higher level OpenCL wrappers.

use CL::*;
use CL::ll::*;

struct Platform {
    id: cl_platform_id
}

pub fn get_platforms() -> ~[Platform] {
    let mut num_platforms = 0;

    // TODO: Check result status
    clGetPlatformIDs(0, ptr::null(), ptr::addr_of(&num_platforms));

    let ids = vec::to_mut(vec::from_elem(num_platforms as uint,
                                         0 as cl_platform_id));

    do vec::as_imm_buf(ids) |ids, len| {
        clGetPlatformIDs(len as cl_uint,
                         ids, ptr::addr_of(&num_platforms))
    };

    do ids.map |id| { Platform { id: *id } }
}

struct Device {
    id: cl_device_id
}

pub fn get_devices(platform: Platform) -> ~[Device] {
    let mut num_devices = 0;

    clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, 0, ptr::null(),
                   ptr::addr_of(&num_devices));

    let ids = vec::to_mut(vec::from_elem(num_devices as uint,
                                         0 as cl_device_id));
    do vec::as_imm_buf(ids) |ids, len| {
        clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, len as cl_uint,
                       ids, ptr::addr_of(&num_devices));
    };

    do ids.map |id| { Device {id: *id }}
}

struct Context {
    ctx: cl_context,

    drop {
        clReleaseContext(self.ctx);
    }
}

pub fn create_context(device: Device) -> Context {
    // TODO: Support for multiple devices
    let mut errcode = 0;

    // TODO: Proper error messages
    let ctx = clCreateContext(ptr::null(), 1, ptr::addr_of(&device.id),
                              ptr::null(), ptr::null(), ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create opencl context!"
    }

    Context { ctx: ctx }
}

struct CommandQueue {
    cqueue: cl_command_queue,

    drop {
        clReleaseCommandQueue(self.cqueue);
    }
}

pub fn create_commandqueue(ctx: & Context, device: Device) -> CommandQueue {
    let mut errcode = 0;

    let cqueue = clCreateCommandQueue(ctx.ctx, device.id, 0, ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create command queue!"
    }

    CommandQueue { cqueue: cqueue }
}

struct Buffer {
    buffer: cl_mem,
    size: int,

    drop {
        clReleaseMemObject(self.buffer);
    }
}

// TODO: How to make this function cleaner and nice
pub fn create_buffer(ctx: & Context, size: int, flags: cl_mem_flags) -> Buffer {
    let mut errcode = 0;

    let buffer = clCreateBuffer(ctx.ctx, flags, size as libc::size_t, ptr::null(),
                                ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create buffer!"
    }

    Buffer { buffer: buffer, size: size }
}

pub fn enqueue_write_buffer(cqueue: & CommandQueue, buf: & Buffer, host_vector: & ~[float]) unsafe {
    let ret = clEnqueueWriteBuffer(cqueue.cqueue, buf.buffer, CL_TRUE, 0,
                                   buf.size as libc::size_t,
                                   vec::raw::to_ptr(*host_vector) as *libc::c_void,
                                   0, ptr::null(), ptr::null());
    if ret != CL_SUCCESS {
        fail ~"Failed to enqueue write buffer!"
    }
}

pub fn enqueue_read_buffer(cqueue: & CommandQueue, buf: & Buffer, host_vector: & ~[mut float]) unsafe {
    let mut ret = 0;
     do vec::as_imm_buf(*host_vector) |elements, _len| {
        ret = clEnqueueReadBuffer(cqueue.cqueue, buf.buffer, CL_TRUE, 0,
                            buf.size as libc::size_t,
                            elements as *libc::c_void, 0, ptr::null(), ptr::null());
    };

    if ret != CL_SUCCESS {
        fail ~"Failed to enqueue read buffer!"
    }
}

struct Program {
    prg: cl_program,

    drop {
        clReleaseProgram(self.prg);
    }
}

// TODO: Support multiple devices
pub fn create_program_with_binary(ctx: & Context, device: Device, binary_path: & Path) -> Program{
    let mut errcode = 0;
    let binary = match move io::read_whole_file_str(binary_path) {
        result::Ok(move binary) => move binary,
        Err(e) => fail fmt!("%?", e)
    };
    let program = do str::as_c_str(binary) |kernel_binary| {
        clCreateProgramWithBinary(ctx.ctx, 1, ptr::addr_of(&device.id),
                                  ptr::addr_of(&(binary.len() + 1)) as *libc::size_t,
                                  ptr::addr_of(&kernel_binary) as **libc::c_uchar,
                                  ptr::null(),
                                  ptr::addr_of(&errcode))
    };

    if errcode != CL_SUCCESS {
        fail ~"Failed to create open cl program with binary!"
    }

    Program { prg: program }
}

pub fn build_program(program: & Program, device: Device){
    let ret = clBuildProgram(program.prg, 1, ptr::addr_of(&device.id),
                             ptr::null(), ptr::null(), ptr::null());
    if ret != CL_SUCCESS {
        let mut logv = ~"";
        for uint::range(0,1024*1204) |_i| {
            str::push_char(& mut logv, ' ');
        }
        do str::as_buf(logv) |logs, l| {
            let r = clGetProgramBuildInfo(program.prg, device.id, CL_PROGRAM_BUILD_LOG, l as libc::size_t, logs as *libc::c_void, ptr::null());
            if r != CL_SUCCESS{
                io::println(~"failed to get build info!");
            }
        }
        //io::println(logv);
        fail ~"Failure during program building!"
    }
}

struct Kernel {
    kernel: cl_kernel,

    drop {
        clReleaseKernel(self.kernel);
    }
}

impl Kernel {
    fn set_arg<T: KernelArg>(i: uint, x: &T) {
        set_kernel_arg(&self, i as CL::cl_uint, x)
    }
}

pub fn create_kernel(program: & Program, kernel: & str) -> Kernel unsafe{
    let mut errcode = 0;
    let bytes = str::to_bytes(kernel);
    let kernel = clCreateKernel(program.prg, vec::raw::to_ptr(bytes) as *libc::c_char, ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create kernel!"
    }

    Kernel { kernel: kernel }
}

pub trait KernelArg{
    fn get_value() -> *libc::c_void;
}

pub fn set_kernel_arg<T: KernelArg>(kernel: & Kernel, position: cl_uint, arg: & T) unsafe{
    // TODO: How to set different argument types. Currently only support cl_mem
    let ret = clSetKernelArg(kernel.kernel, position,
                             sys::size_of::<cl_mem>() as libc::size_t,
                             arg.get_value());

    if ret != CL_SUCCESS {
        fail ~"Failed to set kernel arg!"
    }
}

pub fn enqueue_nd_range_kernel(cqueue: & CommandQueue, kernel: & Kernel, work_dim: cl_uint,
                               _global_work_offset: int, global_work_size: int,
                               local_work_size: int) unsafe{
    let ret = clEnqueueNDRangeKernel(cqueue.cqueue, kernel.kernel, work_dim,
                                     // ptr::addr_of(&global_work_offset) as *libc::size_t,
                                     ptr::null(),
                                     ptr::addr_of(&global_work_size) as *libc::size_t,
                                     ptr::addr_of(&local_work_size) as *libc::size_t,
                                     0, ptr::null(), ptr::null());
    if ret != CL_SUCCESS {
        io::println((ret as int).to_str());
        fail ~"Failed to enqueue nd range kernel!"
    }
}

pub impl Buffer: KernelArg{
    fn get_value() -> *libc::c_void{
        ptr::addr_of(&self.buffer) as *libc::c_void
    }
}

pub impl *libc::c_void: KernelArg{
    fn get_value() -> *libc::c_void {
        self
    }
}<commit_msg>Added an impl for Platform.<commit_after>// Higher level OpenCL wrappers.

use CL::*;
use CL::ll::*;

struct Platform {
    id: cl_platform_id
}

impl Platform {
    fn get_devices() -> ~[Device] {
        get_devices(self)
    }
}

pub fn get_platforms() -> ~[Platform] {
    let mut num_platforms = 0;

    // TODO: Check result status
    clGetPlatformIDs(0, ptr::null(), ptr::addr_of(&num_platforms));

    let ids = vec::to_mut(vec::from_elem(num_platforms as uint,
                                         0 as cl_platform_id));

    do vec::as_imm_buf(ids) |ids, len| {
        clGetPlatformIDs(len as cl_uint,
                         ids, ptr::addr_of(&num_platforms))
    };

    do ids.map |id| { Platform { id: *id } }
}

struct Device {
    id: cl_device_id
}

pub fn get_devices(platform: Platform) -> ~[Device] {
    let mut num_devices = 0;

    clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, 0, ptr::null(),
                   ptr::addr_of(&num_devices));

    let ids = vec::to_mut(vec::from_elem(num_devices as uint,
                                         0 as cl_device_id));
    do vec::as_imm_buf(ids) |ids, len| {
        clGetDeviceIDs(platform.id, CL_DEVICE_TYPE_GPU, len as cl_uint,
                       ids, ptr::addr_of(&num_devices));
    };

    do ids.map |id| { Device {id: *id }}
}

struct Context {
    ctx: cl_context,

    drop {
        clReleaseContext(self.ctx);
    }
}

pub fn create_context(device: Device) -> Context {
    // TODO: Support for multiple devices
    let mut errcode = 0;

    // TODO: Proper error messages
    let ctx = clCreateContext(ptr::null(), 1, ptr::addr_of(&device.id),
                              ptr::null(), ptr::null(), ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create opencl context!"
    }

    Context { ctx: ctx }
}

struct CommandQueue {
    cqueue: cl_command_queue,
    device: Device

    drop {
        clReleaseCommandQueue(self.cqueue);
    }
}

pub fn create_commandqueue(ctx: & Context, device: Device) -> CommandQueue {
    let mut errcode = 0;

    let cqueue = clCreateCommandQueue(ctx.ctx, device.id, 0, ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create command queue!"
    }

    CommandQueue {
        cqueue: cqueue,
        device: Device
    }
}

struct Buffer {
    buffer: cl_mem,
    size: int,

    drop {
        clReleaseMemObject(self.buffer);
    }
}

// TODO: How to make this function cleaner and nice
pub fn create_buffer(ctx: & Context, size: int, flags: cl_mem_flags) -> Buffer {
    let mut errcode = 0;

    let buffer = clCreateBuffer(ctx.ctx, flags, size as libc::size_t, ptr::null(),
                                ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create buffer!"
    }

    Buffer { buffer: buffer, size: size }
}

pub fn enqueue_write_buffer(cqueue: & CommandQueue, buf: & Buffer, host_vector: & ~[float]) unsafe {
    let ret = clEnqueueWriteBuffer(cqueue.cqueue, buf.buffer, CL_TRUE, 0,
                                   buf.size as libc::size_t,
                                   vec::raw::to_ptr(*host_vector) as *libc::c_void,
                                   0, ptr::null(), ptr::null());
    if ret != CL_SUCCESS {
        fail ~"Failed to enqueue write buffer!"
    }
}

pub fn enqueue_read_buffer(cqueue: & CommandQueue, buf: & Buffer, host_vector: & ~[mut float]) unsafe {
    let mut ret = 0;
     do vec::as_imm_buf(*host_vector) |elements, _len| {
        ret = clEnqueueReadBuffer(cqueue.cqueue, buf.buffer, CL_TRUE, 0,
                            buf.size as libc::size_t,
                            elements as *libc::c_void, 0, ptr::null(), ptr::null());
    };

    if ret != CL_SUCCESS {
        fail ~"Failed to enqueue read buffer!"
    }
}

struct Program {
    prg: cl_program,

    drop {
        clReleaseProgram(self.prg);
    }
}

// TODO: Support multiple devices
pub fn create_program_with_binary(ctx: & Context, device: Device,
                                  binary_path: & Path) -> Program {
    let mut errcode = 0;
    let binary = match move io::read_whole_file_str(binary_path) {
        result::Ok(move binary) => move binary,
        Err(e) => fail fmt!("%?", e)
    };
    let program = do str::as_c_str(binary) |kernel_binary| {
        clCreateProgramWithBinary(ctx.ctx, 1, ptr::addr_of(&device.id),
                                  ptr::addr_of(&(binary.len() + 1)) as *libc::size_t,
                                  ptr::addr_of(&kernel_binary) as **libc::c_uchar,
                                  ptr::null(),
                                  ptr::addr_of(&errcode))
    };

    if errcode != CL_SUCCESS {
        fail ~"Failed to create open cl program with binary!"
    }

    Program { prg: program }
}

pub fn build_program(program: & Program, device: Device){
    let ret = clBuildProgram(program.prg, 1, ptr::addr_of(&device.id),
                             ptr::null(), ptr::null(), ptr::null());
    if ret != CL_SUCCESS {
        let mut logv = ~"";
        for uint::range(0,1024*1204) |_i| {
            str::push_char(& mut logv, ' ');
        }
        do str::as_buf(logv) |logs, l| {
            let r = clGetProgramBuildInfo(program.prg, device.id, CL_PROGRAM_BUILD_LOG, l as libc::size_t, logs as *libc::c_void, ptr::null());
            if r != CL_SUCCESS{
                io::println(~"failed to get build info!");
            }
        }
        //io::println(logv);
        fail ~"Failure during program building!"
    }
}

struct Kernel {
    kernel: cl_kernel,

    drop {
        clReleaseKernel(self.kernel);
    }
}

impl Kernel {
    fn set_arg<T: KernelArg>(i: uint, x: &T) {
        set_kernel_arg(&self, i as CL::cl_uint, x)
    }
}

pub fn create_kernel(program: & Program, kernel: & str) -> Kernel unsafe{
    let mut errcode = 0;
    let bytes = str::to_bytes(kernel);
    let kernel = clCreateKernel(program.prg, vec::raw::to_ptr(bytes) as *libc::c_char, ptr::addr_of(&errcode));

    if errcode != CL_SUCCESS {
        fail ~"Failed to create kernel!"
    }

    Kernel { kernel: kernel }
}

pub trait KernelArg{
    fn get_value() -> *libc::c_void;
}

pub fn set_kernel_arg<T: KernelArg>(kernel: & Kernel, position: cl_uint, arg: & T) unsafe{
    // TODO: How to set different argument types. Currently only support cl_mem
    let ret = clSetKernelArg(kernel.kernel, position,
                             sys::size_of::<cl_mem>() as libc::size_t,
                             arg.get_value());

    if ret != CL_SUCCESS {
        fail ~"Failed to set kernel arg!"
    }
}

pub fn enqueue_nd_range_kernel(cqueue: & CommandQueue, kernel: & Kernel, work_dim: cl_uint,
                               _global_work_offset: int, global_work_size: int,
                               local_work_size: int) unsafe{
    let ret = clEnqueueNDRangeKernel(cqueue.cqueue, kernel.kernel, work_dim,
                                     // ptr::addr_of(&global_work_offset) as *libc::size_t,
                                     ptr::null(),
                                     ptr::addr_of(&global_work_size) as *libc::size_t,
                                     ptr::addr_of(&local_work_size) as *libc::size_t,
                                     0, ptr::null(), ptr::null());
    if ret != CL_SUCCESS {
        io::println((ret as int).to_str());
        fail ~"Failed to enqueue nd range kernel!"
    }
}

pub impl Buffer: KernelArg{
    fn get_value() -> *libc::c_void{
        ptr::addr_of(&self.buffer) as *libc::c_void
    }
}

pub impl *libc::c_void: KernelArg{
    fn get_value() -> *libc::c_void {
        self
    }
}<|endoftext|>


<commit_before>use std::fmt;
use std::hash;
use std::mem;

use runtime::{Message, Object};
use {ToMessage};

pub trait Identifier<T: Message> { }

#[unsafe_no_drop_flag]
pub struct Id<T> {
	ptr: *mut T,
}

impl<T: Message> Id<T> {
	pub unsafe fn from_ptr(ptr: *mut T) -> Id<T> {
		match Id::maybe_from_ptr(ptr) {
			Some(id) => id,
			None => fail!("Attempted to construct an Id from a null pointer"),
		}
	}

	pub unsafe fn from_retained_ptr(ptr: *mut T) -> Id<T> {
		match Id::maybe_from_retained_ptr(ptr) {
			Some(id) => id,
			None => fail!("Attempted to construct an Id from a null pointer"),
		}
	}

	pub unsafe fn maybe_from_ptr(ptr: *mut T) -> Option<Id<T>> {
		// objc_msgSend is a no-op on null pointers
		msg_send![ptr retain];
		Id::maybe_from_retained_ptr(ptr)
	}

	pub unsafe fn maybe_from_retained_ptr(ptr: *mut T) -> Option<Id<T>> {
		if ptr.is_null() {
			None
		} else {
			Some(Id { ptr: ptr })
		}
	}
}

impl<T: Message> Identifier<T> for Id<T> { }

impl<T: Message> ToMessage for Id<T> {
	fn as_ptr(&self) -> *mut Object {
		self.ptr.as_ptr()
	}
}

#[unsafe_destructor]
impl<T: Message> Drop for Id<T> {
	fn drop(&mut self) {
		if !self.ptr.is_null() {
			let ptr = mem::replace(&mut self.ptr, RawPtr::null());
			unsafe {
				msg_send![ptr release];
			}
		}
	}
}

impl<T> Deref<T> for Id<T> {
	fn deref(&self) -> &T {
		unsafe { &*self.ptr }
	}
}

impl<T: PartialEq> PartialEq for Id<T> {
	fn eq(&self, other: &Id<T>) -> bool {
		self.deref() == other.deref()
	}

	fn ne(&self, other: &Id<T>) -> bool {
		self.deref() != other.deref()
	}
}

impl<T: Eq> Eq for Id<T> { }

impl<T: hash::Hash> hash::Hash for Id<T> {
	fn hash(&self, state: &mut hash::sip::SipState) {
		self.deref().hash(state)
	}
}

impl<T: fmt::Show> fmt::Show for Id<T> {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		self.deref().fmt(f)
	}
}

pub trait IdVector<T> {
	fn as_refs_slice(&self) -> &[&T];
}

impl<T, V: Vector<Id<T>>> IdVector<T> for V {
	fn as_refs_slice(&self) -> &[&T] {
		unsafe {
			mem::transmute(self.as_slice())
		}
	}
}

pub trait IntoIdVector<T> {
	unsafe fn into_id_vec(self) -> Vec<Id<T>>;
}

impl<T: Message> IntoIdVector<T> for Vec<*mut T> {
	unsafe fn into_id_vec(self) -> Vec<Id<T>> {
		for &ptr in self.iter() {
			msg_send![ptr retain];
		}
		mem::transmute(self)
	}
}

impl<'a, T: Message> IntoIdVector<T> for Vec<&'a T> {
	unsafe fn into_id_vec(self) -> Vec<Id<T>> {
		for &obj in self.iter() {
			msg_send![obj retain];
		}
		mem::transmute(self)
	}
}
<commit_msg>Use new trait for IdVector traits.<commit_after>use std::fmt;
use std::hash;
use std::mem;

use runtime::{Message, Object};
use {ToMessage};

pub trait Identifier<T: Message> { }

#[unsafe_no_drop_flag]
pub struct Id<T> {
	ptr: *mut T,
}

impl<T: Message> Id<T> {
	pub unsafe fn from_ptr(ptr: *mut T) -> Id<T> {
		match Id::maybe_from_ptr(ptr) {
			Some(id) => id,
			None => fail!("Attempted to construct an Id from a null pointer"),
		}
	}

	pub unsafe fn from_retained_ptr(ptr: *mut T) -> Id<T> {
		match Id::maybe_from_retained_ptr(ptr) {
			Some(id) => id,
			None => fail!("Attempted to construct an Id from a null pointer"),
		}
	}

	pub unsafe fn maybe_from_ptr(ptr: *mut T) -> Option<Id<T>> {
		// objc_msgSend is a no-op on null pointers
		msg_send![ptr retain];
		Id::maybe_from_retained_ptr(ptr)
	}

	pub unsafe fn maybe_from_retained_ptr(ptr: *mut T) -> Option<Id<T>> {
		if ptr.is_null() {
			None
		} else {
			Some(Id { ptr: ptr })
		}
	}
}

impl<T: Message> Identifier<T> for Id<T> { }

impl<T: Message> ToMessage for Id<T> {
	fn as_ptr(&self) -> *mut Object {
		self.ptr.as_ptr()
	}
}

#[unsafe_destructor]
impl<T: Message> Drop for Id<T> {
	fn drop(&mut self) {
		if !self.ptr.is_null() {
			let ptr = mem::replace(&mut self.ptr, RawPtr::null());
			unsafe {
				msg_send![ptr release];
			}
		}
	}
}

impl<T> Deref<T> for Id<T> {
	fn deref(&self) -> &T {
		unsafe { &*self.ptr }
	}
}

impl<T: PartialEq> PartialEq for Id<T> {
	fn eq(&self, other: &Id<T>) -> bool {
		self.deref() == other.deref()
	}

	fn ne(&self, other: &Id<T>) -> bool {
		self.deref() != other.deref()
	}
}

impl<T: Eq> Eq for Id<T> { }

impl<T: hash::Hash> hash::Hash for Id<T> {
	fn hash(&self, state: &mut hash::sip::SipState) {
		self.deref().hash(state)
	}
}

impl<T: fmt::Show> fmt::Show for Id<T> {
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
		self.deref().fmt(f)
	}
}

pub trait IdVector<T> {
	fn as_refs_slice(&self) -> &[&T];
}

impl<T, I: Identifier<T>, V: Vector<I>> IdVector<T> for V {
	fn as_refs_slice(&self) -> &[&T] {
		unsafe {
			mem::transmute(self.as_slice())
		}
	}
}

pub trait IntoIdVector<T, I: Identifier<T>> {
	unsafe fn into_id_vec(self) -> Vec<I>;
}

impl<T: Message, I: Identifier<T>> IntoIdVector<T, I> for Vec<*mut T> {
	unsafe fn into_id_vec(self) -> Vec<I> {
		for &ptr in self.iter() {
			msg_send![ptr retain];
		}
		mem::transmute(self)
	}
}

impl<'a, T: Message, I: Identifier<T>> IntoIdVector<T, I> for Vec<&'a T> {
	unsafe fn into_id_vec(self) -> Vec<I> {
		for &obj in self.iter() {
			msg_send![obj retain];
		}
		mem::transmute(self)
	}
}
<|endoftext|>


<commit_before>use slog::{Level, Record, Serializer, KV};

pub(crate) struct Visitor {
    kvs: Vec<(String, String)>,
}

impl Visitor {
    pub fn new() -> Self {
        Self { kvs: vec![] }
    }
}

impl<'kvs, 'a> log::kv::Visitor<'kvs> for Visitor {
    fn visit_pair(
        &mut self,
        key: log::kv::Key<'kvs>,
        val: log::kv::Value<'kvs>,
    ) -> Result<(), log::kv::Error> {
        let key = key.to_string();
        if let Some(val) = val.to_borrowed_str() {
            let val = val.to_string();
            self.kvs.push((key, val));
        }
        Ok(())
    }
}

fn string_to_static_str(s: String) -> &'static str {
    Box::leak(s.into_boxed_str())
}

impl slog::KV for Visitor {
    fn serialize(&self, _record: &Record, serializer: &mut dyn Serializer) -> slog::Result {
        for (key, val) in &self.kvs {
            let key = string_to_static_str(key.to_owned());
            serializer.emit_str(key, val.as_str())?;
        }
        Ok(())
    }
}<commit_msg>Update: add newline at end of file.<commit_after>use slog::{Level, Record, Serializer, KV};

pub(crate) struct Visitor {
    kvs: Vec<(String, String)>,
}

impl Visitor {
    pub fn new() -> Self {
        Self { kvs: vec![] }
    }
}

impl<'kvs, 'a> log::kv::Visitor<'kvs> for Visitor {
    fn visit_pair(
        &mut self,
        key: log::kv::Key<'kvs>,
        val: log::kv::Value<'kvs>,
    ) -> Result<(), log::kv::Error> {
        let key = key.to_string();
        if let Some(val) = val.to_borrowed_str() {
            let val = val.to_string();
            self.kvs.push((key, val));
        }
        Ok(())
    }
}

fn string_to_static_str(s: String) -> &'static str {
    Box::leak(s.into_boxed_str())
}

impl slog::KV for Visitor {
    fn serialize(&self, _record: &Record, serializer: &mut dyn Serializer) -> slog::Result {
        for (key, val) in &self.kvs {
            let key = string_to_static_str(key.to_owned());
            serializer.emit_str(key, val.as_str())?;
        }
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Client Connection Pooling
use std::borrow::ToOwned;
use std::collections::HashMap;
use std::io::{self, Read, Write};
use std::net::{SocketAddr, Shutdown};
use std::sync::{Arc, Mutex};

use std::time::Duration;

use net::{NetworkConnector, NetworkStream, DefaultConnector};

/// The `NetworkConnector` that behaves as a connection pool used by hyper's `Client`.
pub struct Pool<C: NetworkConnector> {
    connector: C,
    inner: Arc<Mutex<PoolImpl<<C as NetworkConnector>::Stream>>>
}

/// Config options for the `Pool`.
#[derive(Debug)]
pub struct Config {
    /// The maximum idle connections *per host*.
    pub max_idle: usize,
}

impl Default for Config {
    #[inline]
    fn default() -> Config {
        Config {
            max_idle: 5,
        }
    }
}

#[derive(Debug)]
struct PoolImpl<S> {
    conns: HashMap<Key, Vec<PooledStreamInner<S>>>,
    config: Config,
}

type Key = (String, u16, Scheme);

fn key<T: Into<Scheme>>(host: &str, port: u16, scheme: T) -> Key {
    (host.to_owned(), port, scheme.into())
}

#[derive(Clone, PartialEq, Eq, Debug, Hash)]
enum Scheme {
    Http,
    Https,
    Other(String)
}

impl<'a> From<&'a str> for Scheme {
    fn from(s: &'a str) -> Scheme {
        match s {
            "http" => Scheme::Http,
            "https" => Scheme::Https,
            s => Scheme::Other(String::from(s))
        }
    }
}

impl Pool<DefaultConnector> {
    /// Creates a `Pool` with a `DefaultConnector`.
    #[inline]
    pub fn new(config: Config) -> Pool<DefaultConnector> {
        Pool::with_connector(config, DefaultConnector::default())
    }
}

impl<C: NetworkConnector> Pool<C> {
    /// Creates a `Pool` with a specified `NetworkConnector`.
    #[inline]
    pub fn with_connector(config: Config, connector: C) -> Pool<C> {
        Pool {
            connector: connector,
            inner: Arc::new(Mutex::new(PoolImpl {
                conns: HashMap::new(),
                config: config,
            }))
        }
    }

    /// Clear all idle connections from the Pool, closing them.
    #[inline]
    pub fn clear_idle(&mut self) {
        self.inner.lock().unwrap().conns.clear();
    }
}

impl<S> PoolImpl<S> {
    fn reuse(&mut self, key: Key, conn: PooledStreamInner<S>) {
        trace!("reuse {:?}", key);
        let conns = self.conns.entry(key).or_insert(vec![]);
        if conns.len() < self.config.max_idle {
            conns.push(conn);
        }
    }
}

impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector for Pool<C> {
    type Stream = PooledStream<S>;
    fn connect(&self, host: &str, port: u16, scheme: &str) -> ::Result<PooledStream<S>> {
        let key = key(host, port, scheme);
        let mut should_remove = false;
        let inner = match self.inner.lock().unwrap().conns.get_mut(&key) {
            Some(ref mut vec) => {
                trace!("Pool had connection, using");
                should_remove = vec.len() == 1;
                vec.pop().unwrap()
            }
            None => PooledStreamInner {
                key: key.clone(),
                stream: try!(self.connector.connect(host, port, scheme)),
                previous_response_expected_no_content: false,
            }
        };
        if should_remove {
            self.inner.lock().unwrap().conns.remove(&key);
        }
        Ok(PooledStream {
            inner: Some(inner),
            is_closed: false,
            pool: self.inner.clone(),
        })
    }
}

/// A Stream that will try to be returned to the Pool when dropped.
#[derive(Debug)]
pub struct PooledStream<S> {
    inner: Option<PooledStreamInner<S>>,
    is_closed: bool,
    pool: Arc<Mutex<PoolImpl<S>>>,
}

impl<S: NetworkStream> PooledStream<S> {
    /// Take the wrapped stream out of the pool completely.
    pub fn into_inner(mut self) -> S {
        self.inner.take().expect("PooledStream lost its inner stream").stream
    }
}

#[derive(Debug)]
struct PooledStreamInner<S> {
    key: Key,
    stream: S,
    previous_response_expected_no_content: bool,
}

impl<S: NetworkStream> Read for PooledStream<S> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.inner.as_mut().unwrap().stream.read(buf) {
            Ok(0) => {
                // if the wrapped stream returns EOF (Ok(0)), that means the
                // server has closed the stream. we must be sure this stream
                // is dropped and not put back into the pool.
                self.is_closed = true;
                Ok(0)
            },
            r => r
        }
    }
}

impl<S: NetworkStream> Write for PooledStream<S> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.as_mut().unwrap().stream.write(buf)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.inner.as_mut().unwrap().stream.flush()
    }
}

impl<S: NetworkStream> NetworkStream for PooledStream<S> {
    #[inline]
    fn peer_addr(&mut self) -> io::Result<SocketAddr> {
        self.inner.as_mut().unwrap().stream.peer_addr()
    }

    #[inline]
    fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
        self.inner.as_ref().unwrap().stream.set_read_timeout(dur)
    }

    #[inline]
    fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
        self.inner.as_ref().unwrap().stream.set_write_timeout(dur)
    }

    #[inline]
    fn close(&mut self, how: Shutdown) -> io::Result<()> {
        self.is_closed = true;
        self.inner.as_mut().unwrap().stream.close(how)
    }

    #[inline]
    fn set_previous_response_expected_no_content(&mut self, expected: bool) {
        trace!("set_previous_response_expected_no_content {}", expected);
        self.inner.as_mut().unwrap().previous_response_expected_no_content = expected;
    }

    #[inline]
    fn previous_response_expected_no_content(&self) -> bool {
        let answer = self.inner.as_ref().unwrap().previous_response_expected_no_content;
        trace!("previous_response_expected_no_content {}", answer);
        answer
    }
}

impl<S> Drop for PooledStream<S> {
    fn drop(&mut self) {
        trace!("PooledStream.drop, is_closed={}", self.is_closed);
        if !self.is_closed {
            self.inner.take().map(|inner| {
                if let Ok(mut pool) = self.pool.lock() {
                    pool.reuse(inner.key.clone(), inner);
                }
                // else poisoned, give up
            });
        }
    }
}

#[cfg(test)]
mod tests {
    use std::net::Shutdown;
    use std::io::Read;
    use mock::{MockConnector};
    use net::{NetworkConnector, NetworkStream};

    use super::{Pool, key};

    macro_rules! mocked {
        () => ({
            Pool::with_connector(Default::default(), MockConnector)
        })
    }

    #[test]
    fn test_connect_and_drop() {
        let pool = mocked!();
        let key = key("127.0.0.1", 3000, "http");
        pool.connect("127.0.0.1", 3000, "http").unwrap();
        {
            let locked = pool.inner.lock().unwrap();
            assert_eq!(locked.conns.len(), 1);
            assert_eq!(locked.conns.get(&key).unwrap().len(), 1);
        }
        pool.connect("127.0.0.1", 3000, "http").unwrap(); //reused
        {
            let locked = pool.inner.lock().unwrap();
            assert_eq!(locked.conns.len(), 1);
            assert_eq!(locked.conns.get(&key).unwrap().len(), 1);
        }
    }

    #[test]
    fn test_closed() {
        let pool = mocked!();
        let mut stream = pool.connect("127.0.0.1", 3000, "http").unwrap();
        stream.close(Shutdown::Both).unwrap();
        drop(stream);
        let locked = pool.inner.lock().unwrap();
        assert_eq!(locked.conns.len(), 0);
    }

    #[test]
    fn test_eof_closes() {
        let pool = mocked!();

        let mut stream = pool.connect("127.0.0.1", 3000, "http").unwrap();
        assert_eq!(stream.read(&mut [0]).unwrap(), 0);
        drop(stream);
        let locked = pool.inner.lock().unwrap();
        assert_eq!(locked.conns.len(), 0);
    }
}
<commit_msg>fix(client): fix panic in Pool::connect<commit_after>//! Client Connection Pooling
use std::borrow::ToOwned;
use std::collections::HashMap;
use std::io::{self, Read, Write};
use std::net::{SocketAddr, Shutdown};
use std::sync::{Arc, Mutex};

use std::time::Duration;

use net::{NetworkConnector, NetworkStream, DefaultConnector};

/// The `NetworkConnector` that behaves as a connection pool used by hyper's `Client`.
pub struct Pool<C: NetworkConnector> {
    connector: C,
    inner: Arc<Mutex<PoolImpl<<C as NetworkConnector>::Stream>>>
}

/// Config options for the `Pool`.
#[derive(Debug)]
pub struct Config {
    /// The maximum idle connections *per host*.
    pub max_idle: usize,
}

impl Default for Config {
    #[inline]
    fn default() -> Config {
        Config {
            max_idle: 5,
        }
    }
}

#[derive(Debug)]
struct PoolImpl<S> {
    conns: HashMap<Key, Vec<PooledStreamInner<S>>>,
    config: Config,
}

type Key = (String, u16, Scheme);

fn key<T: Into<Scheme>>(host: &str, port: u16, scheme: T) -> Key {
    (host.to_owned(), port, scheme.into())
}

#[derive(Clone, PartialEq, Eq, Debug, Hash)]
enum Scheme {
    Http,
    Https,
    Other(String)
}

impl<'a> From<&'a str> for Scheme {
    fn from(s: &'a str) -> Scheme {
        match s {
            "http" => Scheme::Http,
            "https" => Scheme::Https,
            s => Scheme::Other(String::from(s))
        }
    }
}

impl Pool<DefaultConnector> {
    /// Creates a `Pool` with a `DefaultConnector`.
    #[inline]
    pub fn new(config: Config) -> Pool<DefaultConnector> {
        Pool::with_connector(config, DefaultConnector::default())
    }
}

impl<C: NetworkConnector> Pool<C> {
    /// Creates a `Pool` with a specified `NetworkConnector`.
    #[inline]
    pub fn with_connector(config: Config, connector: C) -> Pool<C> {
        Pool {
            connector: connector,
            inner: Arc::new(Mutex::new(PoolImpl {
                conns: HashMap::new(),
                config: config,
            }))
        }
    }

    /// Clear all idle connections from the Pool, closing them.
    #[inline]
    pub fn clear_idle(&mut self) {
        self.inner.lock().unwrap().conns.clear();
    }
}

impl<S> PoolImpl<S> {
    fn reuse(&mut self, key: Key, conn: PooledStreamInner<S>) {
        trace!("reuse {:?}", key);
        let conns = self.conns.entry(key).or_insert(vec![]);
        if conns.len() < self.config.max_idle {
            conns.push(conn);
        }
    }
}

impl<C: NetworkConnector<Stream=S>, S: NetworkStream + Send> NetworkConnector for Pool<C> {
    type Stream = PooledStream<S>;
    fn connect(&self, host: &str, port: u16, scheme: &str) -> ::Result<PooledStream<S>> {
        let key = key(host, port, scheme);

        let inner = {
            // keep the mutex locked only in this block
            let mut locked = self.inner.lock().unwrap();
            let mut should_remove = false;
            let inner = locked.conns.get_mut(&key).map(|vec| {
                trace!("Pool had connection, using");
                should_remove = vec.len() == 1;
                vec.pop().unwrap()
            });
            if should_remove {
                locked.conns.remove(&key);
            }
            inner
        };

        let inner = match inner {
            Some(inner) => inner,
            None => PooledStreamInner {
                key: key.clone(),
                stream: try!(self.connector.connect(host, port, scheme)),
                previous_response_expected_no_content: false,
            }

        };
        Ok(PooledStream {
            inner: Some(inner),
            is_closed: false,
            pool: self.inner.clone(),
        })
    }
}

/// A Stream that will try to be returned to the Pool when dropped.
#[derive(Debug)]
pub struct PooledStream<S> {
    inner: Option<PooledStreamInner<S>>,
    is_closed: bool,
    pool: Arc<Mutex<PoolImpl<S>>>,
}

impl<S: NetworkStream> PooledStream<S> {
    /// Take the wrapped stream out of the pool completely.
    pub fn into_inner(mut self) -> S {
        self.inner.take().expect("PooledStream lost its inner stream").stream
    }
}

#[derive(Debug)]
struct PooledStreamInner<S> {
    key: Key,
    stream: S,
    previous_response_expected_no_content: bool,
}

impl<S: NetworkStream> Read for PooledStream<S> {
    #[inline]
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self.inner.as_mut().unwrap().stream.read(buf) {
            Ok(0) => {
                // if the wrapped stream returns EOF (Ok(0)), that means the
                // server has closed the stream. we must be sure this stream
                // is dropped and not put back into the pool.
                self.is_closed = true;
                Ok(0)
            },
            r => r
        }
    }
}

impl<S: NetworkStream> Write for PooledStream<S> {
    #[inline]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.as_mut().unwrap().stream.write(buf)
    }

    #[inline]
    fn flush(&mut self) -> io::Result<()> {
        self.inner.as_mut().unwrap().stream.flush()
    }
}

impl<S: NetworkStream> NetworkStream for PooledStream<S> {
    #[inline]
    fn peer_addr(&mut self) -> io::Result<SocketAddr> {
        self.inner.as_mut().unwrap().stream.peer_addr()
    }

    #[inline]
    fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
        self.inner.as_ref().unwrap().stream.set_read_timeout(dur)
    }

    #[inline]
    fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
        self.inner.as_ref().unwrap().stream.set_write_timeout(dur)
    }

    #[inline]
    fn close(&mut self, how: Shutdown) -> io::Result<()> {
        self.is_closed = true;
        self.inner.as_mut().unwrap().stream.close(how)
    }

    #[inline]
    fn set_previous_response_expected_no_content(&mut self, expected: bool) {
        trace!("set_previous_response_expected_no_content {}", expected);
        self.inner.as_mut().unwrap().previous_response_expected_no_content = expected;
    }

    #[inline]
    fn previous_response_expected_no_content(&self) -> bool {
        let answer = self.inner.as_ref().unwrap().previous_response_expected_no_content;
        trace!("previous_response_expected_no_content {}", answer);
        answer
    }
}

impl<S> Drop for PooledStream<S> {
    fn drop(&mut self) {
        trace!("PooledStream.drop, is_closed={}", self.is_closed);
        if !self.is_closed {
            self.inner.take().map(|inner| {
                if let Ok(mut pool) = self.pool.lock() {
                    pool.reuse(inner.key.clone(), inner);
                }
                // else poisoned, give up
            });
        }
    }
}

#[cfg(test)]
mod tests {
    use std::net::Shutdown;
    use std::io::Read;
    use mock::{MockConnector};
    use net::{NetworkConnector, NetworkStream};

    use super::{Pool, key};

    macro_rules! mocked {
        () => ({
            Pool::with_connector(Default::default(), MockConnector)
        })
    }

    #[test]
    fn test_connect_and_drop() {
        let pool = mocked!();
        let key = key("127.0.0.1", 3000, "http");
        pool.connect("127.0.0.1", 3000, "http").unwrap();
        {
            let locked = pool.inner.lock().unwrap();
            assert_eq!(locked.conns.len(), 1);
            assert_eq!(locked.conns.get(&key).unwrap().len(), 1);
        }
        pool.connect("127.0.0.1", 3000, "http").unwrap(); //reused
        {
            let locked = pool.inner.lock().unwrap();
            assert_eq!(locked.conns.len(), 1);
            assert_eq!(locked.conns.get(&key).unwrap().len(), 1);
        }
    }

    #[test]
    fn test_closed() {
        let pool = mocked!();
        let mut stream = pool.connect("127.0.0.1", 3000, "http").unwrap();
        stream.close(Shutdown::Both).unwrap();
        drop(stream);
        let locked = pool.inner.lock().unwrap();
        assert_eq!(locked.conns.len(), 0);
    }

    #[test]
    fn test_eof_closes() {
        let pool = mocked!();

        let mut stream = pool.connect("127.0.0.1", 3000, "http").unwrap();
        assert_eq!(stream.read(&mut [0]).unwrap(), 0);
        drop(stream);
        let locked = pool.inner.lock().unwrap();
        assert_eq!(locked.conns.len(), 0);
    }
}
<|endoftext|>


<commit_before>//! [Standard terms](https://en.wikipedia.org/wiki/Lambda_calculus#Standard_terms)
//!
//! * [SKI](https://en.wikipedia.org/wiki/SKI_combinator_calculus)
//! * [BCKW](https://en.wikipedia.org/wiki/B,_C,_K,_W_system)
// //! * the recursion combinator U - needs more research
//! * the looping combinator ω
//! * the divergent combinator Ω
//! * [the fixed-point combinator Y](https://en.wikipedia.org/wiki/Fixed-point_combinator)

use term::*;
use term::Term::*;

/// I - the identity combinator.
///
/// i := λx.x = λ 1
pub fn i() -> Term { abs(Var(1)) }

/// K - the constant / discarding combinator.
///
/// k := λxy.x = λ λ 2 = true
pub fn k() -> Term { abs(abs(Var(2))) }

/// S - the substitution combinator.
///
/// s := λxyz.x z (y z) = λ λ λ 3 1 (2 1)
pub fn s() -> Term { abs(abs(abs(Var(3).app(Var(1)).app(Var(2).app(Var(1)))))) }

/// B - the composition combinator.
///
/// b := λxyz.x (y z) = λ λ λ 3 (2 1)
pub fn b() -> Term { abs(abs(abs(Var(3).app(Var(2).app(Var(1)))))) }

/// C - the swapping combinator.
///
/// c := λxyz.x z y = λ λ λ 3 1 2
pub fn c() -> Term { abs(abs(abs(Var(3).app(Var(1)).app(Var(2))))) }

/// W - the duplicating combinator.
///
/// w := λxy.x y y = λ λ 2 1 1
pub fn w() -> Term { abs(abs(Var(2).app(Var(1)).app(Var(1)))) }
/*
/// U - the recursion combinator.
///
/// u := λxy.y (x x y) = λ λ 1 (2 2 1)
pub fn u() -> Term { abs(abs(Var(1).app(Var(2).app(Var(2)).app(Var(1))))) }
*/
/// ω - the looping combinator.
///
/// ω := λx.x x
pub fn om() -> Term { abs(Var(1).app(Var(1))) }

/// Ω - the divergent combinator.
///
/// Ω := ω ω
pub fn omm() -> Term { om().app(om()) }

/// Y - the fixed-point combinator.
///
/// Y := λg.(λx.g (x x)) (λx.g (x x)) = λ (λ 2 (1 1)) (λ 2 (1 1))
pub fn y() -> Term { abs(app(abs(Var(2).app(Var(1).app(Var(1)))), abs(Var(2).app(Var(1).app(Var(1)))))) }

<commit_msg>Add Iota, capitalize names in descriptions<commit_after>//! [Standard terms](https://en.wikipedia.org/wiki/Lambda_calculus#Standard_terms)
//!
//! * [SKI](https://en.wikipedia.org/wiki/SKI_combinator_calculus)
//! * [Iota](https://en.wikipedia.org/wiki/Iota_and_Jot)
//! * [BCKW](https://en.wikipedia.org/wiki/B,_C,_K,_W_system)
// //! * the recursion combinator U - needs more research
//! * the looping combinator ω
//! * the divergent combinator Ω
//! * [the fixed-point combinator Y](https://en.wikipedia.org/wiki/Fixed-point_combinator)

use term::*;
use term::Term::*;

/// I - the identity combinator.
///
/// I := λx.x = λ 1
pub fn i() -> Term { abs(Var(1)) }

/// K - the constant / discarding combinator.
///
/// K := λxy.x = λ λ 2 = true
pub fn k() -> Term { abs(abs(Var(2))) }

/// S - the substitution combinator.
///
/// S := λxyz.x z (y z) = λ λ λ 3 1 (2 1)
pub fn s() -> Term { abs(abs(abs(Var(3).app(Var(1)).app(Var(2).app(Var(1)))))) }

/// X - the Iota combinator.
///
/// X := λx.x S K = λ 1 S K
pub fn x() -> Term { abs(Var(1).app(k()).app(s())) }

/// B - the composition combinator.
///
/// B := λxyz.x (y z) = λ λ λ 3 (2 1)
pub fn b() -> Term { abs(abs(abs(Var(3).app(Var(2).app(Var(1)))))) }

/// C - the swapping combinator.
///
/// C := λxyz.x z y = λ λ λ 3 1 2
pub fn c() -> Term { abs(abs(abs(Var(3).app(Var(1)).app(Var(2))))) }

/// W - the duplicating combinator.
///
/// W := λxy.x y y = λ λ 2 1 1
pub fn w() -> Term { abs(abs(Var(2).app(Var(1)).app(Var(1)))) }
/*
/// U - the recursion combinator.
///
/// U := λxy.y (x x y) = λ λ 1 (2 2 1)
pub fn u() -> Term { abs(abs(Var(1).app(Var(2).app(Var(2)).app(Var(1))))) }
*/
/// ω - the looping combinator.
///
/// ω := λx.x x
pub fn om() -> Term { abs(Var(1).app(Var(1))) }

/// Ω - the divergent combinator.
///
/// Ω := ω ω
pub fn omm() -> Term { om().app(om()) }

/// Y - the fixed-point combinator.
///
/// Y := λg.(λx.g (x x)) (λx.g (x x)) = λ (λ 2 (1 1)) (λ 2 (1 1))
pub fn y() -> Term { abs(app(abs(Var(2).app(Var(1).app(Var(1)))), abs(Var(2).app(Var(1).app(Var(1)))))) }

<|endoftext|>


<commit_before>#![feature(rustc_private, i128_type)]

extern crate getopts;
extern crate miri;
extern crate rustc;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate env_logger;
extern crate log_settings;
extern crate syntax;
extern crate log;

use rustc::session::Session;
use rustc::middle::cstore::CrateStore;
use rustc_driver::{Compilation, CompilerCalls, RustcDefaultCalls};
use rustc_driver::driver::{CompileState, CompileController};
use rustc::session::config::{self, Input, ErrorOutputType};
use rustc::hir::{self, itemlikevisit};
use rustc::ty::TyCtxt;
use syntax::ast::{self, MetaItemKind, NestedMetaItemKind};
use std::path::PathBuf;

struct MiriCompilerCalls {
    default: RustcDefaultCalls,
}

impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
    fn early_callback(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        descriptions: &rustc_errors::registry::Registry,
        output: ErrorOutputType,
    ) -> Compilation {
        self.default.early_callback(
            matches,
            sopts,
            cfg,
            descriptions,
            output,
        )
    }
    fn no_input(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
        descriptions: &rustc_errors::registry::Registry,
    ) -> Option<(Input, Option<PathBuf>)> {
        self.default.no_input(
            matches,
            sopts,
            cfg,
            odir,
            ofile,
            descriptions,
        )
    }
    fn late_callback(
        &mut self,
        matches: &getopts::Matches,
        sess: &Session,
        cstore: &CrateStore,
        input: &Input,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
    ) -> Compilation {
        self.default.late_callback(matches, sess, cstore, input, odir, ofile)
    }
    fn build_controller(
        &mut self,
        sess: &Session,
        matches: &getopts::Matches,
    ) -> CompileController<'a> {
        let mut control = self.default.build_controller(sess, matches);
        control.after_hir_lowering.callback = Box::new(after_hir_lowering);
        control.after_analysis.callback = Box::new(after_analysis);
        if sess.target.target != sess.host {
            // only fully compile targets on the host. linking will fail for cross-compilation.
            control.after_analysis.stop = Compilation::Stop;
        }
        control
    }
}

fn after_hir_lowering(state: &mut CompileState) {
    let attr = (
        String::from("miri"),
        syntax::feature_gate::AttributeType::Whitelisted,
    );
    state.session.plugin_attributes.borrow_mut().push(attr);
}

fn after_analysis<'a, 'tcx>(state: &mut CompileState<'a, 'tcx>) {
    state.session.abort_if_errors();

    let tcx = state.tcx.unwrap();
    let limits = resource_limits_from_attributes(state);

    if std::env::args().any(|arg| arg == "--test") {
        struct Visitor<'a, 'tcx: 'a>(
            miri::ResourceLimits,
            TyCtxt<'a, 'tcx, 'tcx>,
            &'a CompileState<'a, 'tcx>
        );
        impl<'a, 'tcx: 'a, 'hir> itemlikevisit::ItemLikeVisitor<'hir> for Visitor<'a, 'tcx> {
            fn visit_item(&mut self, i: &'hir hir::Item) {
                if let hir::Item_::ItemFn(_, _, _, _, _, body_id) = i.node {
                    if i.attrs.iter().any(|attr| {
                        attr.name().map_or(false, |n| n == "test")
                    })
                    {
                        let did = self.1.hir.body_owner_def_id(body_id);
                        println!(
                            "running test: {}",
                            self.1.hir.def_path(did).to_string(self.1)
                        );
                        miri::eval_main(self.1, did, None, self.0);
                        self.2.session.abort_if_errors();
                    }
                }
            }
            fn visit_trait_item(&mut self, _trait_item: &'hir hir::TraitItem) {}
            fn visit_impl_item(&mut self, _impl_item: &'hir hir::ImplItem) {}
        }
        state.hir_crate.unwrap().visit_all_item_likes(
            &mut Visitor(limits, tcx, state),
        );
    } else if let Some((entry_node_id, _)) = *state.session.entry_fn.borrow() {
        let entry_def_id = tcx.hir.local_def_id(entry_node_id);
        let start_wrapper = tcx.lang_items().start_fn().and_then(|start_fn| {
            if tcx.is_mir_available(start_fn) {
                Some(start_fn)
            } else {
                None
            }
        });
        miri::eval_main(tcx, entry_def_id, start_wrapper, limits);

        state.session.abort_if_errors();
    } else {
        println!("no main function found, assuming auxiliary build");
    }
}

fn resource_limits_from_attributes(state: &CompileState) -> miri::ResourceLimits {
    let mut limits = miri::ResourceLimits::default();
    let krate = state.hir_crate.as_ref().unwrap();
    let err_msg = "miri attributes need to be in the form `miri(key = value)`";
    let extract_int = |lit: &syntax::ast::Lit| -> u128 {
        match lit.node {
            syntax::ast::LitKind::Int(i, _) => i,
            _ => {
                state.session.span_fatal(
                    lit.span,
                    "expected an integer literal",
                )
            }
        }
    };

    for attr in krate.attrs.iter().filter(|a| {
        a.name().map_or(false, |n| n == "miri")
    })
    {
        if let Some(items) = attr.meta_item_list() {
            for item in items {
                if let NestedMetaItemKind::MetaItem(ref inner) = item.node {
                    if let MetaItemKind::NameValue(ref value) = inner.node {
                        match &inner.name().as_str()[..] {
                            "memory_size" => limits.memory_size = extract_int(value) as u64,
                            "step_limit" => limits.step_limit = extract_int(value) as u64,
                            "stack_limit" => limits.stack_limit = extract_int(value) as usize,
                            _ => state.session.span_err(item.span, "unknown miri attribute"),
                        }
                    } else {
                        state.session.span_err(inner.span, err_msg);
                    }
                } else {
                    state.session.span_err(item.span, err_msg);
                }
            }
        } else {
            state.session.span_err(attr.span, err_msg);
        }
    }
    limits
}

fn init_logger() {
    let format = |record: &log::LogRecord| {
        if record.level() == log::LogLevel::Trace {
            // prepend spaces to indent the final string
            let indentation = log_settings::settings().indentation;
            format!(
                "{lvl}:{module}:{indent:<indentation$} {text}",
                lvl = record.level(),
                module = record.location().module_path(),
                indentation = indentation,
                indent = "",
                text = record.args(),
            )
        } else {
            format!(
                "{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.location().module_path(),
                text = record.args(),
            )
        }
    };

    let mut builder = env_logger::LogBuilder::new();
    builder.format(format).filter(
        None,
        log::LogLevelFilter::Info,
    );

    if std::env::var("MIRI_LOG").is_ok() {
        builder.parse(&std::env::var("MIRI_LOG").unwrap());
    }

    builder.init().unwrap();
}

fn find_sysroot() -> String {
    if let Ok(sysroot) = std::env::var("MIRI_SYSROOT") {
        return sysroot;
    }

    // Taken from https://github.com/Manishearth/rust-clippy/pull/911.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => {
            option_env!("RUST_SYSROOT")
                .expect(
                    "need to specify RUST_SYSROOT env var or use rustup or multirust",
                )
                .to_owned()
        }
    }
}

fn main() {
    init_logger();
    let mut args: Vec<String> = std::env::args().collect();

    let sysroot_flag = String::from("--sysroot");
    if !args.contains(&sysroot_flag) {
        args.push(sysroot_flag);
        args.push(find_sysroot());
    }

    // Make sure we always have all the MIR (e.g. for auxilary builds in unit tests).
    args.push("-Zalways-encode-mir".to_owned());

    rustc_driver::run_compiler(&args, &mut MiriCompilerCalls {
        default: RustcDefaultCalls,
    }, None, None);
}
<commit_msg>Use numbers instead of indentations<commit_after>#![feature(rustc_private, i128_type)]

extern crate getopts;
extern crate miri;
extern crate rustc;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate env_logger;
extern crate log_settings;
extern crate syntax;
extern crate log;

use rustc::session::Session;
use rustc::middle::cstore::CrateStore;
use rustc_driver::{Compilation, CompilerCalls, RustcDefaultCalls};
use rustc_driver::driver::{CompileState, CompileController};
use rustc::session::config::{self, Input, ErrorOutputType};
use rustc::hir::{self, itemlikevisit};
use rustc::ty::TyCtxt;
use syntax::ast::{self, MetaItemKind, NestedMetaItemKind};
use std::path::PathBuf;

struct MiriCompilerCalls {
    default: RustcDefaultCalls,
}

impl<'a> CompilerCalls<'a> for MiriCompilerCalls {
    fn early_callback(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        descriptions: &rustc_errors::registry::Registry,
        output: ErrorOutputType,
    ) -> Compilation {
        self.default.early_callback(
            matches,
            sopts,
            cfg,
            descriptions,
            output,
        )
    }
    fn no_input(
        &mut self,
        matches: &getopts::Matches,
        sopts: &config::Options,
        cfg: &ast::CrateConfig,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
        descriptions: &rustc_errors::registry::Registry,
    ) -> Option<(Input, Option<PathBuf>)> {
        self.default.no_input(
            matches,
            sopts,
            cfg,
            odir,
            ofile,
            descriptions,
        )
    }
    fn late_callback(
        &mut self,
        matches: &getopts::Matches,
        sess: &Session,
        cstore: &CrateStore,
        input: &Input,
        odir: &Option<PathBuf>,
        ofile: &Option<PathBuf>,
    ) -> Compilation {
        self.default.late_callback(matches, sess, cstore, input, odir, ofile)
    }
    fn build_controller(
        &mut self,
        sess: &Session,
        matches: &getopts::Matches,
    ) -> CompileController<'a> {
        let mut control = self.default.build_controller(sess, matches);
        control.after_hir_lowering.callback = Box::new(after_hir_lowering);
        control.after_analysis.callback = Box::new(after_analysis);
        if sess.target.target != sess.host {
            // only fully compile targets on the host. linking will fail for cross-compilation.
            control.after_analysis.stop = Compilation::Stop;
        }
        control
    }
}

fn after_hir_lowering(state: &mut CompileState) {
    let attr = (
        String::from("miri"),
        syntax::feature_gate::AttributeType::Whitelisted,
    );
    state.session.plugin_attributes.borrow_mut().push(attr);
}

fn after_analysis<'a, 'tcx>(state: &mut CompileState<'a, 'tcx>) {
    state.session.abort_if_errors();

    let tcx = state.tcx.unwrap();
    let limits = resource_limits_from_attributes(state);

    if std::env::args().any(|arg| arg == "--test") {
        struct Visitor<'a, 'tcx: 'a>(
            miri::ResourceLimits,
            TyCtxt<'a, 'tcx, 'tcx>,
            &'a CompileState<'a, 'tcx>
        );
        impl<'a, 'tcx: 'a, 'hir> itemlikevisit::ItemLikeVisitor<'hir> for Visitor<'a, 'tcx> {
            fn visit_item(&mut self, i: &'hir hir::Item) {
                if let hir::Item_::ItemFn(_, _, _, _, _, body_id) = i.node {
                    if i.attrs.iter().any(|attr| {
                        attr.name().map_or(false, |n| n == "test")
                    })
                    {
                        let did = self.1.hir.body_owner_def_id(body_id);
                        println!(
                            "running test: {}",
                            self.1.hir.def_path(did).to_string(self.1)
                        );
                        miri::eval_main(self.1, did, None, self.0);
                        self.2.session.abort_if_errors();
                    }
                }
            }
            fn visit_trait_item(&mut self, _trait_item: &'hir hir::TraitItem) {}
            fn visit_impl_item(&mut self, _impl_item: &'hir hir::ImplItem) {}
        }
        state.hir_crate.unwrap().visit_all_item_likes(
            &mut Visitor(limits, tcx, state),
        );
    } else if let Some((entry_node_id, _)) = *state.session.entry_fn.borrow() {
        let entry_def_id = tcx.hir.local_def_id(entry_node_id);
        let start_wrapper = tcx.lang_items().start_fn().and_then(|start_fn| {
            if tcx.is_mir_available(start_fn) {
                Some(start_fn)
            } else {
                None
            }
        });
        miri::eval_main(tcx, entry_def_id, start_wrapper, limits);

        state.session.abort_if_errors();
    } else {
        println!("no main function found, assuming auxiliary build");
    }
}

fn resource_limits_from_attributes(state: &CompileState) -> miri::ResourceLimits {
    let mut limits = miri::ResourceLimits::default();
    let krate = state.hir_crate.as_ref().unwrap();
    let err_msg = "miri attributes need to be in the form `miri(key = value)`";
    let extract_int = |lit: &syntax::ast::Lit| -> u128 {
        match lit.node {
            syntax::ast::LitKind::Int(i, _) => i,
            _ => {
                state.session.span_fatal(
                    lit.span,
                    "expected an integer literal",
                )
            }
        }
    };

    for attr in krate.attrs.iter().filter(|a| {
        a.name().map_or(false, |n| n == "miri")
    })
    {
        if let Some(items) = attr.meta_item_list() {
            for item in items {
                if let NestedMetaItemKind::MetaItem(ref inner) = item.node {
                    if let MetaItemKind::NameValue(ref value) = inner.node {
                        match &inner.name().as_str()[..] {
                            "memory_size" => limits.memory_size = extract_int(value) as u64,
                            "step_limit" => limits.step_limit = extract_int(value) as u64,
                            "stack_limit" => limits.stack_limit = extract_int(value) as usize,
                            _ => state.session.span_err(item.span, "unknown miri attribute"),
                        }
                    } else {
                        state.session.span_err(inner.span, err_msg);
                    }
                } else {
                    state.session.span_err(item.span, err_msg);
                }
            }
        } else {
            state.session.span_err(attr.span, err_msg);
        }
    }
    limits
}

fn init_logger() {
    let format = |record: &log::LogRecord| {
        if record.level() == log::LogLevel::Trace {
            // prepend frame number
            let indentation = log_settings::settings().indentation;
            format!(
                "{indentation}:{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.location().module_path(),
                indentation = indentation,
                text = record.args(),
            )
        } else {
            format!(
                "{lvl}:{module}: {text}",
                lvl = record.level(),
                module = record.location().module_path(),
                text = record.args(),
            )
        }
    };

    let mut builder = env_logger::LogBuilder::new();
    builder.format(format).filter(
        None,
        log::LogLevelFilter::Info,
    );

    if std::env::var("MIRI_LOG").is_ok() {
        builder.parse(&std::env::var("MIRI_LOG").unwrap());
    }

    builder.init().unwrap();
}

fn find_sysroot() -> String {
    if let Ok(sysroot) = std::env::var("MIRI_SYSROOT") {
        return sysroot;
    }

    // Taken from https://github.com/Manishearth/rust-clippy/pull/911.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => {
            option_env!("RUST_SYSROOT")
                .expect(
                    "need to specify RUST_SYSROOT env var or use rustup or multirust",
                )
                .to_owned()
        }
    }
}

fn main() {
    init_logger();
    let mut args: Vec<String> = std::env::args().collect();

    let sysroot_flag = String::from("--sysroot");
    if !args.contains(&sysroot_flag) {
        args.push(sysroot_flag);
        args.push(find_sysroot());
    }

    // Make sure we always have all the MIR (e.g. for auxilary builds in unit tests).
    args.push("-Zalways-encode-mir".to_owned());

    rustc_driver::run_compiler(&args, &mut MiriCompilerCalls {
        default: RustcDefaultCalls,
    }, None, None);
}
<|endoftext|>


<commit_before>use dirs;
use super::error::{S3Error, Result};
use ini::Ini;
use std::collections::HashMap;
use std::env;

/// AWS access credentials: access key, secret key, and optional token.
///
/// # Example
///
/// Loads from the standard AWS credentials file with the given profile name,
/// defaults to "default".
///
/// ```no_run
/// # // Do not execute this as it would cause unit tests to attempt to access
/// # // real user credentials.
/// use s3::credentials::Credentials;
///
/// // Load credentials from `[default]` profile
/// let credentials = Credentials::default();
///
/// // Also loads credentials from `[default]` profile
/// let credentials = Credentials::new(None, None, None, None);
///
/// // Load credentials from `[my-profile]` profile
/// let credentials = Credentials::new(None, None, None, Some("my-profile".into()));
/// ```
///
/// Credentials may also be initialized directly or by the following environment variables:
///
///   - `AWS_ACCESS_KEY_ID`,
///   - `AWS_SECRET_ACCESS_KEY`
///   - `AWS_SESSION_TOKEN`
///
/// The order of preference is arguments, then environment, and finally AWS
/// credentials file.
///
/// ```
/// use s3::credentials::Credentials;
///
/// // Load credentials directly
/// let access_key = String::from("AKIAIOSFODNN7EXAMPLE");
/// let secret_key = String::from("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
/// let credentials = Credentials::new(Some(access_key), Some(secret_key), None, None);
///
/// // Load credentials from the environment
/// use std::env;
/// env::set_var("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE");
/// env::set_var("AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
/// let credentials = Credentials::new(None, None, None, None);
/// ```
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Credentials {
    /// AWS public access key.
    pub access_key: String,
    /// AWS secret key.
    pub secret_key: String,
    /// Temporary token issued by AWS service.
    pub token: Option<String>,
    _private: (),
}

impl Credentials {
    /// Initialize Credentials directly with key ID, secret key, and optional
    /// token.
    pub fn new(
        access_key: Option<String>,
        secret_key: Option<String>,
        token: Option<String>,
        profile: Option<String>,
    ) -> Credentials {
        let credentials = match access_key {
            Some(key) => match secret_key {
                Some(secret) => match token {
                    Some(t) => Some(Credentials {
                        access_key: key,
                        secret_key: secret,
                        token: Some(t),
                        _private: (),
                    }),
                    None => Some(Credentials {
                        access_key: key,
                        secret_key: secret,
                        token: None,
                        _private: (),
                    }),
                },
                None => None,
            },
            None => None,
        };
        match credentials {
            Some(c) => c,
            None => match Credentials::from_env() {
                Ok(c) => c,
                Err(_) => match Credentials::from_profile(profile) {
                    Ok(c) => c,
                    Err(_) => match futures::executor::block_on(Credentials::from_instance_metadata()) {
                        Ok(c) => c,
                        Err(e) => panic!("No credentials provided as arguments, in the environment or in the profile file. \n {}", e)
                    }
                }
            }
        }
    }

    fn from_env() -> Result<Credentials> {
        let access_key = env::var("AWS_ACCESS_KEY_ID")?;
        let secret_key = env::var("AWS_SECRET_ACCESS_KEY")?;
        let token = match env::var("AWS_SESSION_TOKEN") {
            Ok(x) => Some(x),
            Err(_) => None,
        };
        Ok(Credentials {
            access_key,
            secret_key,
            token,
            _private: (),
        })
    }

    async fn from_instance_metadata() -> Result<Credentials> {
        if !Credentials::is_ec2() {
            return Err(S3Error::from("Not an EC2 instance"));
        }

        let resp: HashMap<String, String> =
            reqwest::get("http://169.254.169.254/latest/meta-data/iam/info").await?.json().await?;
        let credentials = if let Some(arn) = resp.get("InstanceProfileArn") {
            if let Some(role) = arn.split('/').last() {
                let resp: HashMap<String, String> = reqwest::get(&format!(
                    "http://169.254.169.254/latest/meta-data/iam/security-credentials/{}",
                    role
                )).await?
                .json().await?;
                let access_key = resp.get("AccessKeyId").unwrap().clone();
                let secret_key = resp.get("SecretAccessKey").unwrap().clone();
                let token = Some(resp.get("Token").unwrap().clone());
                Some(Credentials {
                    access_key,
                    secret_key,
                    token,
                    _private: (),
                })
            } else {
                None
            }
        } else {
            None
        };

        Ok(credentials.unwrap())
    }

    fn is_ec2() -> bool {
        if let Ok(uuid) = std::fs::read_to_string("/sys/hypervisor/uuid") {
            if uuid.len() >= 3 && &uuid[..3] == "ec2" {
                return true;
            }
        }
        if let Ok(uuid) = std::fs::read_to_string("/sys/class/dmi/id/board_vendor") {
            if uuid.len() >= 10 && &uuid[..10] == "Amazon EC2" {
                return true;
            }
        }
        false
    }

    pub fn from_profile(section: Option<String>) -> Result<Credentials> {
        let home_dir = match dirs::home_dir() {
            Some(path) => Ok(path),
            None => Err(S3Error::from("Invalid home dir")),
        };
        let profile = format!("{}/.aws/credentials", home_dir?.display());
        let conf = Ini::load_from_file(&profile)?;
        let section = match section {
            Some(s) => s,
            None => String::from("default"),
        };
        let mut access_key = Err(S3Error::from("Missing aws_access_key_id section"));
        let mut secret_key = Err(S3Error::from("Missing aws_secret_access_key section"));
        let mut token = None;
        if let Some(data) = conf.section(Some(section)) {
            access_key = match data.get("aws_access_key_id") {
                Some(x) => Ok(x.to_owned()),
                None => Err(S3Error::from("Missing aws_access_key_id section")),
            };
            secret_key = match data.get("aws_secret_access_key") {
                Some(x) => Ok(x.to_owned()),
                None => Err(S3Error::from("Missing aws_secret_access_key section")),
            };
            token = match data.get("aws_security_token") {
                Some(x) => Some(x.to_owned()),
                None => None,
            }
        }

        Ok(Credentials {
            access_key: access_key?,
            secret_key: secret_key?,
            token,
            _private: (),
        })
    }
}

impl Default for Credentials {
    fn default() -> Self {
        Credentials::new(None, None, None, None)
    }
}
<commit_msg>Fix credentials<commit_after>use dirs;
use super::error::{S3Error, Result};
use ini::Ini;
use std::collections::HashMap;
use std::env;
use tokio::runtime::Runtime;

/// AWS access credentials: access key, secret key, and optional token.
///
/// # Example
///
/// Loads from the standard AWS credentials file with the given profile name,
/// defaults to "default".
///
/// ```no_run
/// # // Do not execute this as it would cause unit tests to attempt to access
/// # // real user credentials.
/// use s3::credentials::Credentials;
///
/// // Load credentials from `[default]` profile
/// let credentials = Credentials::default();
///
/// // Also loads credentials from `[default]` profile
/// let credentials = Credentials::new(None, None, None, None);
///
/// // Load credentials from `[my-profile]` profile
/// let credentials = Credentials::new(None, None, None, Some("my-profile".into()));
/// ```
///
/// Credentials may also be initialized directly or by the following environment variables:
///
///   - `AWS_ACCESS_KEY_ID`,
///   - `AWS_SECRET_ACCESS_KEY`
///   - `AWS_SESSION_TOKEN`
///
/// The order of preference is arguments, then environment, and finally AWS
/// credentials file.
///
/// ```
/// use s3::credentials::Credentials;
///
/// // Load credentials directly
/// let access_key = String::from("AKIAIOSFODNN7EXAMPLE");
/// let secret_key = String::from("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
/// let credentials = Credentials::new(Some(access_key), Some(secret_key), None, None);
///
/// // Load credentials from the environment
/// use std::env;
/// env::set_var("AWS_ACCESS_KEY_ID", "AKIAIOSFODNN7EXAMPLE");
/// env::set_var("AWS_SECRET_ACCESS_KEY", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY");
/// let credentials = Credentials::new(None, None, None, None);
/// ```
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Credentials {
    /// AWS public access key.
    pub access_key: String,
    /// AWS secret key.
    pub secret_key: String,
    /// Temporary token issued by AWS service.
    pub token: Option<String>,
    _private: (),
}

impl Credentials {
    /// Initialize Credentials directly with key ID, secret key, and optional
    /// token.
    pub fn new(
        access_key: Option<String>,
        secret_key: Option<String>,
        token: Option<String>,
        profile: Option<String>,
    ) -> Credentials {
        let credentials = match access_key {
            Some(key) => match secret_key {
                Some(secret) => match token {
                    Some(t) => Some(Credentials {
                        access_key: key,
                        secret_key: secret,
                        token: Some(t),
                        _private: (),
                    }),
                    None => Some(Credentials {
                        access_key: key,
                        secret_key: secret,
                        token: None,
                        _private: (),
                    }),
                },
                None => None,
            },
            None => None,
        };
        let mut rt = Runtime::new().unwrap();
        match credentials {
            Some(c) => c,
            None => match Credentials::from_env() {
                Ok(c) => c,
                Err(_) => match Credentials::from_profile(profile) {
                    Ok(c) => c,
                    Err(_) => match rt.block_on(Credentials::from_instance_metadata()) {
                        Ok(c) => c,
                        Err(e) => panic!("No credentials provided as arguments, in the environment or in the profile file. \n {}", e)
                    }
                }
            }
        }
    }

    fn from_env() -> Result<Credentials> {
        let access_key = env::var("AWS_ACCESS_KEY_ID")?;
        let secret_key = env::var("AWS_SECRET_ACCESS_KEY")?;
        let token = match env::var("AWS_SESSION_TOKEN") {
            Ok(x) => Some(x),
            Err(_) => None,
        };
        Ok(Credentials {
            access_key,
            secret_key,
            token,
            _private: (),
        })
    }

    async fn from_instance_metadata() -> Result<Credentials> {
        if !Credentials::is_ec2() {
            return Err(S3Error::from("Not an EC2 instance"));
        }

        let resp: HashMap<String, String> =
            reqwest::get("http://169.254.169.254/latest/meta-data/iam/info").await?.json().await?;
        let credentials = if let Some(arn) = resp.get("InstanceProfileArn") {
            if let Some(role) = arn.split('/').last() {
                let resp: HashMap<String, String> = reqwest::get(&format!(
                    "http://169.254.169.254/latest/meta-data/iam/security-credentials/{}",
                    role
                )).await?
                .json().await?;
                let access_key = resp.get("AccessKeyId").unwrap().clone();
                let secret_key = resp.get("SecretAccessKey").unwrap().clone();
                let token = Some(resp.get("Token").unwrap().clone());
                Some(Credentials {
                    access_key,
                    secret_key,
                    token,
                    _private: (),
                })
            } else {
                None
            }
        } else {
            None
        };

        Ok(credentials.unwrap())
    }

    fn is_ec2() -> bool {
        if let Ok(uuid) = std::fs::read_to_string("/sys/hypervisor/uuid") {
            if uuid.len() >= 3 && &uuid[..3] == "ec2" {
                return true;
            }
        }
        if let Ok(uuid) = std::fs::read_to_string("/sys/class/dmi/id/board_vendor") {
            if uuid.len() >= 10 && &uuid[..10] == "Amazon EC2" {
                return true;
            }
        }
        false
    }

    pub fn from_profile(section: Option<String>) -> Result<Credentials> {
        let home_dir = match dirs::home_dir() {
            Some(path) => Ok(path),
            None => Err(S3Error::from("Invalid home dir")),
        };
        let profile = format!("{}/.aws/credentials", home_dir?.display());
        let conf = Ini::load_from_file(&profile)?;
        let section = match section {
            Some(s) => s,
            None => String::from("default"),
        };
        let mut access_key = Err(S3Error::from("Missing aws_access_key_id section"));
        let mut secret_key = Err(S3Error::from("Missing aws_secret_access_key section"));
        let mut token = None;
        if let Some(data) = conf.section(Some(section)) {
            access_key = match data.get("aws_access_key_id") {
                Some(x) => Ok(x.to_owned()),
                None => Err(S3Error::from("Missing aws_access_key_id section")),
            };
            secret_key = match data.get("aws_secret_access_key") {
                Some(x) => Ok(x.to_owned()),
                None => Err(S3Error::from("Missing aws_secret_access_key section")),
            };
            token = match data.get("aws_security_token") {
                Some(x) => Some(x.to_owned()),
                None => None,
            }
        }

        Ok(Credentials {
            access_key: access_key?,
            secret_key: secret_key?,
            token,
            _private: (),
        })
    }
}

impl Default for Credentials {
    fn default() -> Self {
        Credentials::new(None, None, None, None)
    }
}
<|endoftext|>