<commit_before>use nannou::prelude::*;

struct Model {
    bind_group: wgpu::BindGroup,
    render_pipeline: wgpu::RenderPipeline,
    vertex_buffer: wgpu::Buffer,
}

// The vertex type that we will use to represent a point on our triangle.
#[derive(Clone, Copy)]
struct Vertex {
    position: [f32; 2],
}

// The vertices that make up the rectangle to which the image will be drawn.
const VERTICES: [Vertex; 4] = [
    Vertex {
        position: [-1.0, -1.0],
    },
    Vertex {
        position: [-1.0, 1.0],
    },
    Vertex {
        position: [1.0, -1.0],
    },
    Vertex {
        position: [1.0, 1.0],
    },
];

fn main() {
    nannou::app(model).run();
}

fn model(app: &App) -> Model {
    let w_id = app.new_window().size(512, 512).view(view).build().unwrap();

    // The gpu device associated with the window's swapchain
    let window = app.window(w_id).unwrap();
    let device = window.swap_chain_device();
    let format = Frame::TEXTURE_FORMAT;
    let msaa_samples = window.msaa_samples();

    // Load shader modules.
    let vs = include_bytes!("shaders/vert.spv");
    let vs_spirv =
        wgpu::read_spirv(std::io::Cursor::new(&vs[..])).expect("failed to read hard-coded SPIRV");
    let vs_mod = device.create_shader_module(&vs_spirv);
    let fs = include_bytes!("shaders/frag.spv");
    let fs_spirv =
        wgpu::read_spirv(std::io::Cursor::new(&fs[..])).expect("failed to read hard-coded SPIRV");
    let fs_mod = device.create_shader_module(&fs_spirv);

    // Load an image from disk, then load it to the GPU as a `Texture`.
    let logo_path = app.assets_path().unwrap().join("images").join("Nannou.png");
    let image = image::open(logo_path).unwrap().to_rgba();
    let texture = {
        let mut queue = window.swap_chain_queue().lock().unwrap();
        create_texture_from_rgba_image(device, &mut *queue, image)
    };
    let texture_view = texture.create_default_view();

    // Create the sampler for sampling from the source texture.
    let sampler = wgpu::SamplerBuilder::new().build(device);

    // Create the render pipeline.
    let bind_group_layout = create_bind_group_layout(device);
    let bind_group = create_bind_group(device, &bind_group_layout, &texture_view, &sampler);
    let pipeline_layout = create_pipeline_layout(device, &bind_group_layout);
    let render_pipeline = create_render_pipeline(
        device,
        &pipeline_layout,
        &vs_mod,
        &fs_mod,
        format,
        msaa_samples,
    );

    // Create the vertex buffer.
    let vertex_buffer = device
        .create_buffer_mapped(VERTICES.len(), wgpu::BufferUsage::VERTEX)
        .fill_from_slice(&VERTICES[..]);

    Model {
        bind_group,
        vertex_buffer,
        render_pipeline,
    }
}

// Draw the state of your `Model` into the given `Frame` here.
fn view(_app: &App, model: &Model, frame: &Frame) {
    // Using this we will encode commands that will be submitted to the GPU.
    let mut encoder = frame.command_encoder();

    // A render pass describes how to draw to an output "attachment".
    let render_pass_desc = wgpu::RenderPassDescriptor {
        color_attachments: &[wgpu::RenderPassColorAttachmentDescriptor {
            attachment: frame.texture(),
            resolve_target: None,
            load_op: wgpu::LoadOp::Clear,
            store_op: wgpu::StoreOp::Store,
            clear_color: wgpu::Color::TRANSPARENT,
        }],
        depth_stencil_attachment: None,
    };

    // The render pass can be thought of a single large command consisting of sub commands.
    // Here we begin the render pass and add sub-commands for setting the bind group, render
    // pipeline, vertex buffers and then finally drawing.
    let mut render_pass = encoder.begin_render_pass(&render_pass_desc);
    render_pass.set_bind_group(0, &model.bind_group, &[]);
    render_pass.set_pipeline(&model.render_pipeline);
    render_pass.set_vertex_buffers(0, &[(&model.vertex_buffer, 0)]);

    // We want to draw the whole range of vertices, and we're only drawing one instance of them.
    let vertex_range = 0..VERTICES.len() as u32;
    let instance_range = 0..1;
    render_pass.draw(vertex_range, instance_range);

    // Now we're done! The commands we added will be submitted after `view` completes.
}

// TODO: Generalise this into a nice API exposed by nannou (e.g. `Texture::from_image`).
fn create_texture_from_rgba_image(
    device: &wgpu::Device,
    queue: &mut wgpu::Queue,
    image: image::RgbaImage,
) -> wgpu::Texture {
    const LOGO_TEXTURE_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Bgra8UnormSrgb;

    // Initialise the texture.
    let (width, height) = image.dimensions();
    let logo_tex = wgpu::TextureBuilder::new()
        .size([width, height])
        .format(LOGO_TEXTURE_FORMAT)
        .usage(wgpu::TextureUsage::SAMPLED | wgpu::TextureUsage::COPY_DST)
        .build(device);

    // Upload the pixel data.
    let data = &image.into_raw()[..];
    let buffer = device
        .create_buffer_mapped(data.len(), wgpu::BufferUsage::COPY_SRC)
        .fill_from_slice(data);

    // Submit command for copying pixel data to the texture.
    let pixel_size_bytes = 4; // Rgba8, as above.
    let buffer_copy_view = wgpu::BufferCopyView {
        buffer: &buffer,
        offset: 0,
        row_pitch: width * pixel_size_bytes,
        image_height: height,
    };
    let texture_copy_view = wgpu::TextureCopyView {
        texture: &logo_tex,
        mip_level: 0,
        array_layer: 0,
        origin: wgpu::Origin3d::ZERO,
    };
    let extent = wgpu::Extent3d {
        width: width,
        height: height,
        depth: 1,
    };
    let cmd_encoder_desc = wgpu::CommandEncoderDescriptor { todo: 0 };
    let mut encoder = device.create_command_encoder(&cmd_encoder_desc);
    encoder.copy_buffer_to_texture(buffer_copy_view, texture_copy_view, extent);
    queue.submit(&[encoder.finish()]);

    logo_tex
}

fn create_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    let texture_binding = wgpu::BindGroupLayoutBinding {
        binding: 0,
        visibility: wgpu::ShaderStage::FRAGMENT,
        ty: wgpu::BindingType::SampledTexture {
            multisampled: false,
            dimension: wgpu::TextureViewDimension::D2,
        },
    };
    let sampler_binding = wgpu::BindGroupLayoutBinding {
        binding: 1,
        visibility: wgpu::ShaderStage::FRAGMENT,
        ty: wgpu::BindingType::Sampler,
    };
    let bindings = &[texture_binding, sampler_binding];
    let desc = wgpu::BindGroupLayoutDescriptor { bindings };
    device.create_bind_group_layout(&desc)
}

fn create_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    texture: &wgpu::TextureView,
    sampler: &wgpu::Sampler,
) -> wgpu::BindGroup {
    let texture_binding = wgpu::Binding {
        binding: 0,
        resource: wgpu::BindingResource::TextureView(&texture),
    };
    let sampler_binding = wgpu::Binding {
        binding: 1,
        resource: wgpu::BindingResource::Sampler(&sampler),
    };
    let bindings = &[texture_binding, sampler_binding];
    let desc = wgpu::BindGroupDescriptor { layout, bindings };
    device.create_bind_group(&desc)
}

fn create_pipeline_layout(
    device: &wgpu::Device,
    bind_group_layout: &wgpu::BindGroupLayout,
) -> wgpu::PipelineLayout {
    let desc = wgpu::PipelineLayoutDescriptor {
        bind_group_layouts: &[&bind_group_layout],
    };
    device.create_pipeline_layout(&desc)
}

fn vertex_attrs() -> [wgpu::VertexAttributeDescriptor; 1] {
    [wgpu::VertexAttributeDescriptor {
        format: wgpu::VertexFormat::Float2,
        offset: 0,
        shader_location: 0,
    }]
}

fn create_render_pipeline(
    device: &wgpu::Device,
    layout: &wgpu::PipelineLayout,
    vs_mod: &wgpu::ShaderModule,
    fs_mod: &wgpu::ShaderModule,
    dst_format: wgpu::TextureFormat,
    sample_count: u32,
) -> wgpu::RenderPipeline {
    let vs_desc = wgpu::ProgrammableStageDescriptor {
        module: &vs_mod,
        entry_point: "main",
    };
    let fs_desc = wgpu::ProgrammableStageDescriptor {
        module: &fs_mod,
        entry_point: "main",
    };
    let raster_desc = wgpu::RasterizationStateDescriptor {
        front_face: wgpu::FrontFace::Ccw,
        cull_mode: wgpu::CullMode::None,
        depth_bias: 0,
        depth_bias_slope_scale: 0.0,
        depth_bias_clamp: 0.0,
    };
    let color_state_desc = wgpu::ColorStateDescriptor {
        format: dst_format,
        color_blend: wgpu::BlendDescriptor::REPLACE,
        alpha_blend: wgpu::BlendDescriptor::REPLACE,
        write_mask: wgpu::ColorWrite::ALL,
    };
    let vertex_attrs = vertex_attrs();
    let vertex_buffer_desc = wgpu::VertexBufferDescriptor {
        stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
        step_mode: wgpu::InputStepMode::Vertex,
        attributes: &vertex_attrs[..],
    };
    let desc = wgpu::RenderPipelineDescriptor {
        layout,
        vertex_stage: vs_desc,
        fragment_stage: Some(fs_desc),
        rasterization_state: Some(raster_desc),
        primitive_topology: wgpu::PrimitiveTopology::TriangleStrip,
        color_states: &[color_state_desc],
        depth_stencil_state: None,
        index_format: wgpu::IndexFormat::Uint16,
        vertex_buffers: &[vertex_buffer_desc],
        sample_count,
        sample_mask: !0,
        alpha_to_coverage_enabled: false,
    };
    device.create_render_pipeline(&desc)
}
<commit_msg>Simplify wgpu_image example with new constructor. Fix formatting.<commit_after>use nannou::prelude::*;

struct Model {
    bind_group: wgpu::BindGroup,
    render_pipeline: wgpu::RenderPipeline,
    vertex_buffer: wgpu::Buffer,
}

// The vertex type that we will use to represent a point on our triangle.
#[derive(Clone, Copy)]
struct Vertex {
    position: [f32; 2],
}

// The vertices that make up the rectangle to which the image will be drawn.
const VERTICES: [Vertex; 4] = [
    Vertex {
        position: [-1.0, -1.0],
    },
    Vertex {
        position: [-1.0, 1.0],
    },
    Vertex {
        position: [1.0, -1.0],
    },
    Vertex {
        position: [1.0, 1.0],
    },
];

fn main() {
    nannou::app(model).run();
}

fn model(app: &App) -> Model {
    let w_id = app.new_window().size(512, 512).view(view).build().unwrap();

    let window = app.window(w_id).unwrap();
    let device = window.swap_chain_device();
    let format = Frame::TEXTURE_FORMAT;
    let msaa_samples = window.msaa_samples();

    let vs = include_bytes!("shaders/vert.spv");
    let vs_spirv =
        wgpu::read_spirv(std::io::Cursor::new(&vs[..])).expect("failed to read hard-coded SPIRV");
    let vs_mod = device.create_shader_module(&vs_spirv);
    let fs = include_bytes!("shaders/frag.spv");
    let fs_spirv =
        wgpu::read_spirv(std::io::Cursor::new(&fs[..])).expect("failed to read hard-coded SPIRV");
    let fs_mod = device.create_shader_module(&fs_spirv);

    let logo_path = app.assets_path().unwrap().join("images").join("Nannou.png");
    let image = image::open(logo_path).unwrap().to_rgba();
    let texture = {
        // The wgpu device queue used to load the image data.
        let mut queue = window.swap_chain_queue().lock().unwrap();
        // Describe how we will use the texture so that the GPU may handle it efficiently.
        let usage = wgpu::TextureUsage::SAMPLED;
        wgpu::Texture::load_from_image_buffer(device, &mut *queue, &image, usage)
    };
    let texture_view = texture.create_default_view();

    // Create the sampler for sampling from the source texture.
    let sampler = wgpu::SamplerBuilder::new().build(device);

    let bind_group_layout = create_bind_group_layout(device);
    let bind_group = create_bind_group(device, &bind_group_layout, &texture_view, &sampler);
    let pipeline_layout = create_pipeline_layout(device, &bind_group_layout);
    let render_pipeline = create_render_pipeline(
        device,
        &pipeline_layout,
        &vs_mod,
        &fs_mod,
        format,
        msaa_samples,
    );

    // Create the vertex buffer.
    let vertex_buffer = device
        .create_buffer_mapped(VERTICES.len(), wgpu::BufferUsage::VERTEX)
        .fill_from_slice(&VERTICES[..]);

    Model {
        bind_group,
        vertex_buffer,
        render_pipeline,
    }
}

fn view(_app: &App, model: &Model, frame: &Frame) {
    let mut encoder = frame.command_encoder();

    let render_pass_desc = wgpu::RenderPassDescriptor {
        color_attachments: &[wgpu::RenderPassColorAttachmentDescriptor {
            attachment: frame.texture(),
            resolve_target: None,
            load_op: wgpu::LoadOp::Clear,
            store_op: wgpu::StoreOp::Store,
            clear_color: wgpu::Color::TRANSPARENT,
        }],
        depth_stencil_attachment: None,
    };

    let mut render_pass = encoder.begin_render_pass(&render_pass_desc);
    render_pass.set_bind_group(0, &model.bind_group, &[]);
    render_pass.set_pipeline(&model.render_pipeline);
    render_pass.set_vertex_buffers(0, &[(&model.vertex_buffer, 0)]);

    let vertex_range = 0..VERTICES.len() as u32;
    let instance_range = 0..1;
    render_pass.draw(vertex_range, instance_range);
}

fn create_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    let texture_binding = wgpu::BindGroupLayoutBinding {
        binding: 0,
        visibility: wgpu::ShaderStage::FRAGMENT,
        ty: wgpu::BindingType::SampledTexture {
            multisampled: false,
            dimension: wgpu::TextureViewDimension::D2,
        },
    };
    let sampler_binding = wgpu::BindGroupLayoutBinding {
        binding: 1,
        visibility: wgpu::ShaderStage::FRAGMENT,
        ty: wgpu::BindingType::Sampler,
    };
    let bindings = &[texture_binding, sampler_binding];
    let desc = wgpu::BindGroupLayoutDescriptor { bindings };
    device.create_bind_group_layout(&desc)
}

fn create_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    texture: &wgpu::TextureView,
    sampler: &wgpu::Sampler,
) -> wgpu::BindGroup {
    let texture_binding = wgpu::Binding {
        binding: 0,
        resource: wgpu::BindingResource::TextureView(&texture),
    };
    let sampler_binding = wgpu::Binding {
        binding: 1,
        resource: wgpu::BindingResource::Sampler(&sampler),
    };
    let bindings = &[texture_binding, sampler_binding];
    let desc = wgpu::BindGroupDescriptor { layout, bindings };
    device.create_bind_group(&desc)
}

fn create_pipeline_layout(
    device: &wgpu::Device,
    bind_group_layout: &wgpu::BindGroupLayout,
) -> wgpu::PipelineLayout {
    let desc = wgpu::PipelineLayoutDescriptor {
        bind_group_layouts: &[&bind_group_layout],
    };
    device.create_pipeline_layout(&desc)
}

fn vertex_attrs() -> [wgpu::VertexAttributeDescriptor; 1] {
    [wgpu::VertexAttributeDescriptor {
        format: wgpu::VertexFormat::Float2,
        offset: 0,
        shader_location: 0,
    }]
}

fn create_render_pipeline(
    device: &wgpu::Device,
    layout: &wgpu::PipelineLayout,
    vs_mod: &wgpu::ShaderModule,
    fs_mod: &wgpu::ShaderModule,
    dst_format: wgpu::TextureFormat,
    sample_count: u32,
) -> wgpu::RenderPipeline {
    let vs_desc = wgpu::ProgrammableStageDescriptor {
        module: &vs_mod,
        entry_point: "main",
    };
    let fs_desc = wgpu::ProgrammableStageDescriptor {
        module: &fs_mod,
        entry_point: "main",
    };
    let raster_desc = wgpu::RasterizationStateDescriptor {
        front_face: wgpu::FrontFace::Ccw,
        cull_mode: wgpu::CullMode::None,
        depth_bias: 0,
        depth_bias_slope_scale: 0.0,
        depth_bias_clamp: 0.0,
    };
    let color_state_desc = wgpu::ColorStateDescriptor {
        format: dst_format,
        color_blend: wgpu::BlendDescriptor::REPLACE,
        alpha_blend: wgpu::BlendDescriptor::REPLACE,
        write_mask: wgpu::ColorWrite::ALL,
    };
    let vertex_attrs = vertex_attrs();
    let vertex_buffer_desc = wgpu::VertexBufferDescriptor {
        stride: std::mem::size_of::<Vertex>() as wgpu::BufferAddress,
        step_mode: wgpu::InputStepMode::Vertex,
        attributes: &vertex_attrs[..],
    };
    let desc = wgpu::RenderPipelineDescriptor {
        layout,
        vertex_stage: vs_desc,
        fragment_stage: Some(fs_desc),
        rasterization_state: Some(raster_desc),
        primitive_topology: wgpu::PrimitiveTopology::TriangleStrip,
        color_states: &[color_state_desc],
        depth_stencil_state: None,
        index_format: wgpu::IndexFormat::Uint16,
        vertex_buffers: &[vertex_buffer_desc],
        sample_count,
        sample_mask: !0,
        alpha_to_coverage_enabled: false,
    };
    device.create_render_pipeline(&desc)
}
<|endoftext|>


<commit_before>extern crate kafka;
extern crate getopts;
extern crate env_logger;

use std::{env, fmt, process};
use std::fs::File;
use std::io::{self, stdin, stderr, Write, BufRead, BufReader};
use std::ops::{Deref, DerefMut};

use kafka::client::{KafkaClient, Compression, RequiredAcks};
use kafka::producer::{AsBytes, Producer, Record};

// how many brokers do we require to acknowledge the send messages
const REQUIRED_ACKS: RequiredAcks = RequiredAcks::One;
// how long do we allow wainting for the acknowledgement
const ACK_TIMEOUT: i32 = 100;

/// This is a very simple command line application sending every
/// non-empty line of standard input to a specified kafka topic; one
/// by one (not the most efficient way, of course).
///
/// Alternatively, messages can be read from an input file and sent do
/// kafka in batches (the typical use-case).
fn main() {
    env_logger::init().unwrap();

    let cfg = match Config::from_cmdline() {
        Ok(cfg) => cfg,
        Err(e) => {
            println!("{}", e);
            process::exit(1);
        }
    };
    if let Err(e) = produce(&cfg) {
        println!("{}", e);
        process::exit(1);
    }
}

fn produce(cfg: &Config) -> Result<(), Error> {
    let mut client = KafkaClient::new(cfg.brokers.clone());
    client.set_compression(cfg.compression);
    try!(client.load_metadata_all());

    // ~ verify that the remote brokers do know about the target topic
    if !client.topics().contains(&cfg.topic) {
        return Err(Error::Literal(format!("No such topic at {:?}: {}", cfg.brokers, cfg.topic)));
    }
    match cfg.input_file {
        None => {
            let stdin = stdin();
            let mut stdin = stdin.lock();
            produce_impl(&mut stdin, client, &cfg)
        }
        Some(ref file) => {
            let mut r = BufReader::new(try!(File::open(file)));
            produce_impl(&mut r, client, &cfg)
        }
    }
}

fn produce_impl(src: &mut BufRead, client: KafkaClient, cfg: &Config) -> Result<(), Error> {
    let mut producer = try!(Producer::from_client(client)
                            .with_ack_timeout(ACK_TIMEOUT)
                            .with_required_acks(REQUIRED_ACKS)
                            .create());
    if cfg.batch_size < 2 {
        produce_impl_nobatch(&mut producer, src, cfg)
    } else {
        produce_impl_inbatches(&mut producer, src, cfg)
    }
}

struct Trimmed(String);

impl AsBytes for Trimmed {
    fn as_bytes(&self) -> &[u8] {
        self.0.trim().as_bytes()
    }
}

impl Deref for Trimmed {
    type Target = String;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Trimmed {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

fn produce_impl_nobatch(producer: &mut Producer, src: &mut BufRead, cfg: &Config) -> Result<(), Error> {
    let mut stderr = stderr();
    let mut rec = Record::from_value(&cfg.topic, Trimmed(String::new()));
    loop {
        rec.value.clear();
        if try!(src.read_line(&mut rec.value)) == 0 {
            break; // ~ EOF reached
        }
        if rec.value.trim().is_empty() {
            continue; // ~ skip empty lines
        }
        // ~ directly send to kafka
        try!(producer.send(&rec));
        let _ = write!(stderr, "Sent: {}", *rec.value);
    }
    Ok(())
}

// This implementation wants to be efficient.  It buffers N lines from
// the source and sends these in batches to Kafka.  Line buffers
// across batches are re-used for the sake of avoiding allocations.
fn produce_impl_inbatches(producer: &mut Producer, src: &mut BufRead, cfg: &Config)
                          -> Result<(), Error>
{
    assert!(cfg.batch_size > 1);

    // ~ a buffer of prepared records to be send in a batch to Kafka
    // ~ in the loop following, we'll only modify the 'value' of the
    // cached records
    let mut rec_stash: Vec<Record<(), Trimmed>> =
        (0 .. cfg.batch_size)
        .map(|_| Record::from_value(&cfg.topic, Trimmed(String::new())))
        .collect();
    // ~ points to the next free slot in `rec_stash`.  if it reaches
    // `rec_stash.len()` we'll send `rec_stash` to kafka
    let mut next_rec = 0;
    loop {
        // ~ send out a batch if it's ready
        if next_rec == rec_stash.len() {
            try!(send_batch(producer, &rec_stash));
            next_rec = 0;
        }
        let mut rec = &mut rec_stash[next_rec];
        rec.value.clear();
        if try!(src.read_line(&mut rec.value)) == 0 {
            break; // ~ EOF reached
        }
        if rec.value.trim().is_empty() {
            continue; // ~ skip empty lines
        }
        // ~ ok, we got a line. read the next one in a new buffer
        next_rec += 1;
    }
    // ~ flush pending messages - if any
    if next_rec > 0 {
        try!(send_batch(producer, &rec_stash[..next_rec]));
    }
    Ok(())
}

fn send_batch(producer: &mut Producer, batch: &[Record<(), Trimmed>]) -> Result<(), Error> {
    let rs = try!(producer.send_all(batch));
    for r in rs {
        if let Err(e) = r.offset {
            return Err(From::from(e));
        }
    }
    Ok(())
}

// --------------------------------------------------------------------

enum Error {
    Kafka(kafka::error::Error),
    Io(io::Error),
    Literal(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::Kafka(ref e) => write!(f, "kafka-error: {}", e),
            &Error::Io(ref e) => write!(f, "io-error: {}", e),
            &Error::Literal(ref s) => write!(f, "{}", s),
        }
    }
}

impl From<kafka::error::Error> for Error {
    fn from(e: kafka::error::Error) -> Self { Error::Kafka(e) }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { Error::Io(e) }
}

// --------------------------------------------------------------------

struct Config {
    brokers: Vec<String>,
    topic: String,
    input_file: Option<String>,
    compression: Compression,
    batch_size: usize,
}

impl Config {
    fn from_cmdline() -> Result<Config, Error> {
        let args: Vec<String> = env::args().collect();
        let mut opts = getopts::Options::new();
        opts.optflag("h", "help", "Print this help screen");
        opts.optopt("", "brokers", "Specify kafka brokers (comma separated)", "HOSTS");
        opts.optopt("", "topic", "Specify target topic", "NAME");
        opts.optopt("", "input", "Specify input file", "FILE");
        opts.optopt("", "compression", "Compress messages [NONE, GZIP, SNAPPY]", "TYPE");
        opts.optopt("", "batch-size", "Send N message in one batch.", "N");

        let m = match opts.parse(&args[1..]) {
            Ok(m) => m,
            Err(e) => return Err(Error::Literal(e.to_string())),
        };
        if m.opt_present("help") {
            let brief = format!("{} [options]", args[0]);
            return Err(Error::Literal(opts.usage(&brief)));
        }
        Ok(Config {
            brokers: m.opt_str("brokers")
                .unwrap_or_else(|| "localhost:9092".to_owned())
                .split(',')
                .map(|s| s.trim().to_owned())
                .collect(),
            topic: m.opt_str("topic")
                .unwrap_or_else(|| "my-topic".to_owned()),
            input_file: m.opt_str("input"),
            compression: {
                let s = m.opt_str("compression")
                    .unwrap_or_else(|| "NONE".to_owned());
                match s.trim() {
                    "none" | "NONE" => Compression::NONE,
                    #[cfg(feature = "gzip")]
                    "gzip" | "GZIP" => Compression::GZIP,
                    #[cfg(feature = "snappy")]
                    "snappy" | "SNAPPY" => Compression::SNAPPY,
                    _ => return Err(Error::Literal(format!("Unsupported compression type: {}", s))),
                }
            },
            batch_size: {
                match m.opt_str("batch-size") {
                    None => 1,
                    Some(s) => {
                        match s.parse::<usize>() {
                            Ok(n) => n,
                            Err(_) => return Err(Error::Literal(format!("Not a number: {}", s))),
                        }
                    }
                }
            },
        })
    }
}
<commit_msg>console-producer: allow specifying required-acks<commit_after>extern crate kafka;
extern crate getopts;
extern crate env_logger;

use std::{env, fmt, process};
use std::fs::File;
use std::io::{self, stdin, stderr, Write, BufRead, BufReader};
use std::ops::{Deref, DerefMut};

use kafka::client::{KafkaClient, Compression, RequiredAcks};
use kafka::producer::{AsBytes, Producer, Record};

// how long do we allow wainting for the acknowledgement
const ACK_TIMEOUT: i32 = 100;

/// This is a very simple command line application sending every
/// non-empty line of standard input to a specified kafka topic; one
/// by one (not the most efficient way, of course).
///
/// Alternatively, messages can be read from an input file and sent do
/// kafka in batches (the typical use-case).
fn main() {
    env_logger::init().unwrap();

    let cfg = match Config::from_cmdline() {
        Ok(cfg) => cfg,
        Err(e) => {
            println!("{}", e);
            process::exit(1);
        }
    };
    if let Err(e) = produce(&cfg) {
        println!("{}", e);
        process::exit(1);
    }
}

fn produce(cfg: &Config) -> Result<(), Error> {
    let mut client = KafkaClient::new(cfg.brokers.clone());
    client.set_compression(cfg.compression);
    try!(client.load_metadata_all());

    // ~ verify that the remote brokers do know about the target topic
    if !client.topics().contains(&cfg.topic) {
        return Err(Error::Literal(format!("No such topic at {:?}: {}", cfg.brokers, cfg.topic)));
    }
    match cfg.input_file {
        None => {
            let stdin = stdin();
            let mut stdin = stdin.lock();
            produce_impl(&mut stdin, client, &cfg)
        }
        Some(ref file) => {
            let mut r = BufReader::new(try!(File::open(file)));
            produce_impl(&mut r, client, &cfg)
        }
    }
}

fn produce_impl(src: &mut BufRead, client: KafkaClient, cfg: &Config) -> Result<(), Error> {
    let mut producer = try!(Producer::from_client(client)
                            .with_ack_timeout(ACK_TIMEOUT)
                            .with_required_acks(cfg.required_acks)
                            .create());
    if cfg.batch_size < 2 {
        produce_impl_nobatch(&mut producer, src, cfg)
    } else {
        produce_impl_inbatches(&mut producer, src, cfg)
    }
}

struct Trimmed(String);

impl AsBytes for Trimmed {
    fn as_bytes(&self) -> &[u8] {
        self.0.trim().as_bytes()
    }
}

impl Deref for Trimmed {
    type Target = String;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl DerefMut for Trimmed {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

fn produce_impl_nobatch(producer: &mut Producer, src: &mut BufRead, cfg: &Config) -> Result<(), Error> {
    let mut stderr = stderr();
    let mut rec = Record::from_value(&cfg.topic, Trimmed(String::new()));
    loop {
        rec.value.clear();
        if try!(src.read_line(&mut rec.value)) == 0 {
            break; // ~ EOF reached
        }
        if rec.value.trim().is_empty() {
            continue; // ~ skip empty lines
        }
        // ~ directly send to kafka
        try!(producer.send(&rec));
        let _ = write!(stderr, "Sent: {}", *rec.value);
    }
    Ok(())
}

// This implementation wants to be efficient.  It buffers N lines from
// the source and sends these in batches to Kafka.  Line buffers
// across batches are re-used for the sake of avoiding allocations.
fn produce_impl_inbatches(producer: &mut Producer, src: &mut BufRead, cfg: &Config)
                          -> Result<(), Error>
{
    assert!(cfg.batch_size > 1);

    // ~ a buffer of prepared records to be send in a batch to Kafka
    // ~ in the loop following, we'll only modify the 'value' of the
    // cached records
    let mut rec_stash: Vec<Record<(), Trimmed>> =
        (0 .. cfg.batch_size)
        .map(|_| Record::from_value(&cfg.topic, Trimmed(String::new())))
        .collect();
    // ~ points to the next free slot in `rec_stash`.  if it reaches
    // `rec_stash.len()` we'll send `rec_stash` to kafka
    let mut next_rec = 0;
    loop {
        // ~ send out a batch if it's ready
        if next_rec == rec_stash.len() {
            try!(send_batch(producer, &rec_stash));
            next_rec = 0;
        }
        let mut rec = &mut rec_stash[next_rec];
        rec.value.clear();
        if try!(src.read_line(&mut rec.value)) == 0 {
            break; // ~ EOF reached
        }
        if rec.value.trim().is_empty() {
            continue; // ~ skip empty lines
        }
        // ~ ok, we got a line. read the next one in a new buffer
        next_rec += 1;
    }
    // ~ flush pending messages - if any
    if next_rec > 0 {
        try!(send_batch(producer, &rec_stash[..next_rec]));
    }
    Ok(())
}

fn send_batch(producer: &mut Producer, batch: &[Record<(), Trimmed>]) -> Result<(), Error> {
    let rs = try!(producer.send_all(batch));
    for r in rs {
        if let Err(e) = r.offset {
            return Err(From::from(e));
        }
    }
    Ok(())
}

// --------------------------------------------------------------------

enum Error {
    Kafka(kafka::error::Error),
    Io(io::Error),
    Literal(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::Kafka(ref e) => write!(f, "kafka-error: {}", e),
            &Error::Io(ref e) => write!(f, "io-error: {}", e),
            &Error::Literal(ref s) => write!(f, "{}", s),
        }
    }
}

impl From<kafka::error::Error> for Error {
    fn from(e: kafka::error::Error) -> Self { Error::Kafka(e) }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { Error::Io(e) }
}

// --------------------------------------------------------------------

struct Config {
    brokers: Vec<String>,
    topic: String,
    input_file: Option<String>,
    compression: Compression,
    required_acks: RequiredAcks,
    batch_size: usize,
}

impl Config {
    fn from_cmdline() -> Result<Config, Error> {
        use std::ascii::AsciiExt;

        let args: Vec<String> = env::args().collect();
        let mut opts = getopts::Options::new();
        opts.optflag("h", "help", "Print this help screen");
        opts.optopt("", "brokers", "Specify kafka brokers (comma separated)", "HOSTS");
        opts.optopt("", "topic", "Specify target topic", "NAME");
        opts.optopt("", "input", "Specify input file", "FILE");
        opts.optopt("", "compression", "Compress messages [NONE, GZIP, SNAPPY]", "TYPE");
        opts.optopt("", "required-acks", "Specify amount of required broker acknowledgments \
                                          [NONE, ONE, ALL]", "TYPE");
        opts.optopt("", "batch-size", "Send N message in one batch.", "N");

        let m = match opts.parse(&args[1..]) {
            Ok(m) => m,
            Err(e) => return Err(Error::Literal(e.to_string())),
        };
        if m.opt_present("help") {
            let brief = format!("{} [options]", args[0]);
            return Err(Error::Literal(opts.usage(&brief)));
        }
        Ok(Config {
            brokers: m.opt_str("brokers")
                .unwrap_or_else(|| "localhost:9092".to_owned())
                .split(',')
                .map(|s| s.trim().to_owned())
                .collect(),
            topic: m.opt_str("topic")
                .unwrap_or_else(|| "my-topic".to_owned()),
            input_file: m.opt_str("input"),
            compression: match m.opt_str("compression") {
                None => Compression::NONE,
                Some(ref s) if s.eq_ignore_ascii_case("none") => Compression::NONE,
                #[cfg(feature = "gzip")]
                Some(ref s) if s.eq_ignore_ascii_case("gzip") => Compression::GZIP,
                #[cfg(feature = "snappy")]
                Some(ref s) if s.eq_ignore_ascii_case("snappy") => Compression::SNAPPY,
                Some(s) => return Err(Error::Literal(format!("Unsupported compression type: {}", s))),
            },
            required_acks: match m.opt_str("required-acks") {
                None => RequiredAcks::One,
                Some(ref s) if s.eq_ignore_ascii_case("none") => RequiredAcks::None,
                Some(ref s) if s.eq_ignore_ascii_case("one") => RequiredAcks::One,
                Some(ref s) if s.eq_ignore_ascii_case("all") => RequiredAcks::All,
                Some(s) => return Err(Error::Literal(format!("Unknown --required-acks argument: {}", s))),
            },
            batch_size: {
                match m.opt_str("batch-size") {
                    None => 1,
                    Some(s) => {
                        match s.parse::<usize>() {
                            Ok(n) => n,
                            Err(_) => return Err(Error::Literal(format!("Not a number: {}", s))),
                        }
                    }
                }
            },
        })
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Guillaume Pinot (@TeXitoi) <texitoi@texitoi.eu>,
// Kevin Knapp (@kbknapp) <kbknapp@gmail.com>, and
// Andrew Hobden (@hoverbear) <andrew@hoverbear.org>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// This work was derived from Structopt (https://github.com/TeXitoi/structopt)
// commit#ea76fa1b1b273e65e3b0b1046643715b49bec51f which is licensed under the
// MIT/Apache 2.0 license.

use std::env;

use proc_macro2::{Span, TokenStream};
use proc_macro_error::{abort, abort_call_site};
use quote::{quote, quote_spanned};
use syn::{
    punctuated::Punctuated, spanned::Spanned, token::Comma, Attribute, Data, DataStruct,
    DeriveInput, Field, Fields, Ident, Type,
};

use crate::{
    attrs::{Attrs, GenOutput, Kind, Name, ParserKind, DEFAULT_CASING, DEFAULT_ENV_CASING},
    dummies,
    utils::{sub_type, subty_if_name, Sp, Ty},
};

pub fn derive_into_app(input: &DeriveInput) -> TokenStream {
    let ident = &input.ident;

    dummies::into_app(ident);

    match input.data {
        Data::Struct(DataStruct {
            fields: Fields::Named(ref fields),
            ..
        }) => gen_for_struct(ident, &fields.named, &input.attrs).0,
        Data::Struct(DataStruct {
            fields: Fields::Unit,
            ..
        }) => gen_for_struct(ident, &Punctuated::<Field, Comma>::new(), &input.attrs).0,
        Data::Enum(_) => gen_for_enum(ident, &input.attrs),
        _ => abort_call_site!("`#[derive(IntoApp)]` only supports non-tuple structs and enums"),
    }
}

pub fn gen_for_struct(
    struct_name: &Ident,
    fields: &Punctuated<Field, Comma>,
    attrs: &[Attribute],
) -> GenOutput {
    let (into_app, attrs) = gen_into_app_fn(attrs);
    let augment_clap = gen_augment_clap_fn(fields, &attrs);

    let tokens = quote! {
        #[allow(dead_code, unreachable_code, unused_variables)]
        #[allow(
            clippy::style,
            clippy::complexity,
            clippy::pedantic,
            clippy::restriction,
            clippy::perf,
            clippy::deprecated,
            clippy::nursery,
            clippy::cargo
        )]
        #[deny(clippy::correctness)]
        impl clap::IntoApp for #struct_name {
            #into_app
            #augment_clap
        }
    };

    (tokens, attrs)
}

pub fn gen_for_enum(enum_name: &Ident, attrs: &[Attribute]) -> TokenStream {
    let app_name = env::var("CARGO_PKG_NAME").ok().unwrap_or_default();

    let attrs = Attrs::from_struct(
        Span::call_site(),
        attrs,
        Name::Assigned(quote!(#app_name)),
        Sp::call_site(DEFAULT_CASING),
        Sp::call_site(DEFAULT_ENV_CASING),
    );
    let name = attrs.cased_name();

    quote! {
        #[allow(dead_code, unreachable_code, unused_variables)]
        #[allow(
            clippy::style,
            clippy::complexity,
            clippy::pedantic,
            clippy::restriction,
            clippy::perf,
            clippy::deprecated,
            clippy::nursery,
            clippy::cargo
        )]
        #[deny(clippy::correctness)]
        impl clap::IntoApp for #enum_name {
            fn into_app<'b>() -> clap::App<'b> {
                let app = clap::App::new(#name)
                    .setting(clap::AppSettings::SubcommandRequiredElseHelp);
                <#enum_name as clap::IntoApp>::augment_clap(app)
            }

            fn augment_clap<'b>(app: clap::App<'b>) -> clap::App<'b> {
                <#enum_name as clap::Subcommand>::augment_subcommands(app)
            }

            fn into_app_for_update<'b>() -> clap::App<'b> {
                let app = clap::App::new(#name);
                <#enum_name as clap::IntoApp>::augment_clap_for_update(app)
            }

            fn augment_clap_for_update<'b>(app: clap::App<'b>) -> clap::App<'b> {
                <#enum_name as clap::Subcommand>::augment_subcommands_for_update(app)
            }
        }
    }
}

fn gen_into_app_fn(attrs: &[Attribute]) -> GenOutput {
    let app_name = env::var("CARGO_PKG_NAME").ok().unwrap_or_default();

    let attrs = Attrs::from_struct(
        Span::call_site(),
        attrs,
        Name::Assigned(quote!(#app_name)),
        Sp::call_site(DEFAULT_CASING),
        Sp::call_site(DEFAULT_ENV_CASING),
    );
    let name = attrs.cased_name();

    let tokens = quote! {
        fn into_app<'b>() -> clap::App<'b> {
            Self::augment_clap(clap::App::new(#name))
        }
        fn into_app_for_update<'b>() -> clap::App<'b> {
            Self::augment_clap_for_update(clap::App::new(#name))
        }
    };

    (tokens, attrs)
}

fn gen_augment_clap_fn(fields: &Punctuated<Field, Comma>, parent_attribute: &Attrs) -> TokenStream {
    let app_var = Ident::new("app", Span::call_site());
    let augmentation = gen_app_augmentation(fields, &app_var, parent_attribute, false);
    let augmentation_update = gen_app_augmentation(fields, &app_var, parent_attribute, true);
    quote! {
        fn augment_clap<'b>(#app_var: clap::App<'b>) -> clap::App<'b> {
            #augmentation
        }
        fn augment_clap_for_update<'b>(#app_var: clap::App<'b>) -> clap::App<'b> {
            #augmentation_update
        }
    }
}

fn gen_arg_enum_possible_values(ty: &Type) -> TokenStream {
    quote_spanned! { ty.span()=>
        .possible_values(&<#ty as clap::ArgEnum>::VARIANTS)
    }
}

/// Generate a block of code to add arguments/subcommands corresponding to
/// the `fields` to an app.
pub fn gen_app_augmentation(
    fields: &Punctuated<Field, Comma>,
    app_var: &Ident,
    parent_attribute: &Attrs,
    override_required: bool,
) -> TokenStream {
    let mut subcmds = fields.iter().filter_map(|field| {
        let attrs = Attrs::from_field(
            &field,
            parent_attribute.casing(),
            parent_attribute.env_casing(),
        );
        let kind = attrs.kind();
        if let Kind::Subcommand(ty) = &*kind {
            let subcmd_type = match (**ty, sub_type(&field.ty)) {
                (Ty::Option, Some(sub_type)) => sub_type,
                _ => &field.ty,
            };
            let required = if **ty == Ty::Option {
                quote!()
            } else {
                quote_spanned! { kind.span()=>
                    let #app_var = #app_var.setting(
                        clap::AppSettings::SubcommandRequiredElseHelp
                    );
                }
            };

            let span = field.span();
            let ts = if override_required {
                quote! {
                    let #app_var = <#subcmd_type as clap::Subcommand>::augment_subcommands_for_update( #app_var );
                }
            } else{
                quote! {
                    let #app_var = <#subcmd_type as clap::Subcommand>::augment_subcommands( #app_var );
                    #required
                }
            };
            Some((span, ts))
        } else {
            None
        }
    });
    let subcmd = subcmds.next().map(|(_, ts)| ts);
    if let Some((span, _)) = subcmds.next() {
        abort!(
            span,
            "multiple subcommand sets are not allowed, that's the second"
        );
    }

    let args = fields.iter().filter_map(|field| {
        let attrs = Attrs::from_field(
            field,
            parent_attribute.casing(),
            parent_attribute.env_casing(),
        );
        let kind = attrs.kind();
        match &*kind {
            Kind::Subcommand(_) | Kind::Skip(_) | Kind::FromGlobal(_) | Kind::ExternalSubcommand => None,
            Kind::Flatten => {
                let ty = &field.ty;
                Some(quote_spanned! { kind.span()=>
                    let #app_var = <#ty as clap::IntoApp>::augment_clap(#app_var);
                })
            }
            Kind::Arg(ty) => {
                let convert_type = match **ty {
                    Ty::Vec | Ty::Option => sub_type(&field.ty).unwrap_or(&field.ty),
                    Ty::OptionOption | Ty::OptionVec => {
                        sub_type(&field.ty).and_then(sub_type).unwrap_or(&field.ty)
                    }
                    _ => &field.ty,
                };

                let occurrences = *attrs.parser().kind == ParserKind::FromOccurrences;
                let flag = *attrs.parser().kind == ParserKind::FromFlag;

                let parser = attrs.parser();
                let func = &parser.func;

                let validator = match *parser.kind {
                    _ if attrs.is_enum() => quote!(),
                    ParserKind::TryFromStr => quote_spanned! { func.span()=>
                        .validator(|s| {
                            #func(s)
                            .map(|_: #convert_type| ())
                        })
                    },
                    ParserKind::TryFromOsStr => quote_spanned! { func.span()=>
                        .validator_os(|s| #func(s).map(|_: #convert_type| ()))
                    },
                    _ => quote!(),
                };

                let modifier = match **ty {
                    Ty::Bool => quote!(),

                    Ty::Option => {
                        let mut possible_values = quote!();

                        if attrs.is_enum() {
                            if let Some(subty) = subty_if_name(&field.ty, "Option") {
                                possible_values = gen_arg_enum_possible_values(subty);
                            }
                        };

                        quote_spanned! { ty.span()=>
                            .takes_value(true)
                            #possible_values
                            #validator
                        }
                    }

                    Ty::OptionOption => quote_spanned! { ty.span()=>
                        .takes_value(true)
                        .multiple(false)
                        .min_values(0)
                        .max_values(1)
                        #validator
                    },

                    Ty::OptionVec => quote_spanned! { ty.span()=>
                        .takes_value(true)
                        .multiple(true)
                        .min_values(0)
                        #validator
                    },

                    Ty::Vec => {
                        let mut possible_values = quote!();

                        if attrs.is_enum() {
                            if let Some(subty) = subty_if_name(&field.ty, "Vec") {
                                possible_values = gen_arg_enum_possible_values(subty);
                            }
                        };

                        quote_spanned! { ty.span()=>
                            .takes_value(true)
                            .multiple(true)
                            #possible_values
                            #validator
                        }
                    }

                    Ty::Other if occurrences => quote_spanned! { ty.span()=>
                        .multiple_occurrences(true)
                    },

                    Ty::Other if flag => quote_spanned! { ty.span()=>
                        .takes_value(false)
                        .multiple(false)
                    },

                    Ty::Other => {
                        let required = !attrs.has_method("default_value") && !override_required;
                        let mut possible_values = quote!();

                        if attrs.is_enum() {
                            possible_values = gen_arg_enum_possible_values(&field.ty);
                        };

                        quote_spanned! { ty.span()=>
                            .takes_value(true)
                            .required(#required)
                            #possible_values
                            #validator
                        }
                    }
                };

                let name = attrs.cased_name();
                let methods = attrs.field_methods();

                Some(quote_spanned! { field.span()=>
                    let #app_var = #app_var.arg(
                        clap::Arg::new(#name)
                            #modifier
                            #methods
                    );
                })
            }
        }
    });

    let app_methods = parent_attribute.top_level_methods();
    let version = parent_attribute.version();
    quote! {{
        let #app_var = #app_var#app_methods;
        #( #args )*
        #subcmd
        #app_var#version
    }}
}
<commit_msg>Run fmt<commit_after>// Copyright 2018 Guillaume Pinot (@TeXitoi) <texitoi@texitoi.eu>,
// Kevin Knapp (@kbknapp) <kbknapp@gmail.com>, and
// Andrew Hobden (@hoverbear) <andrew@hoverbear.org>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// This work was derived from Structopt (https://github.com/TeXitoi/structopt)
// commit#ea76fa1b1b273e65e3b0b1046643715b49bec51f which is licensed under the
// MIT/Apache 2.0 license.

use std::env;

use proc_macro2::{Span, TokenStream};
use proc_macro_error::{abort, abort_call_site};
use quote::{quote, quote_spanned};
use syn::{
    punctuated::Punctuated, spanned::Spanned, token::Comma, Attribute, Data, DataStruct,
    DeriveInput, Field, Fields, Ident, Type,
};

use crate::{
    attrs::{Attrs, GenOutput, Kind, Name, ParserKind, DEFAULT_CASING, DEFAULT_ENV_CASING},
    dummies,
    utils::{sub_type, subty_if_name, Sp, Ty},
};

pub fn derive_into_app(input: &DeriveInput) -> TokenStream {
    let ident = &input.ident;

    dummies::into_app(ident);

    match input.data {
        Data::Struct(DataStruct {
            fields: Fields::Named(ref fields),
            ..
        }) => gen_for_struct(ident, &fields.named, &input.attrs).0,
        Data::Struct(DataStruct {
            fields: Fields::Unit,
            ..
        }) => gen_for_struct(ident, &Punctuated::<Field, Comma>::new(), &input.attrs).0,
        Data::Enum(_) => gen_for_enum(ident, &input.attrs),
        _ => abort_call_site!("`#[derive(IntoApp)]` only supports non-tuple structs and enums"),
    }
}

pub fn gen_for_struct(
    struct_name: &Ident,
    fields: &Punctuated<Field, Comma>,
    attrs: &[Attribute],
) -> GenOutput {
    let (into_app, attrs) = gen_into_app_fn(attrs);
    let augment_clap = gen_augment_clap_fn(fields, &attrs);

    let tokens = quote! {
        #[allow(dead_code, unreachable_code, unused_variables)]
        #[allow(
            clippy::style,
            clippy::complexity,
            clippy::pedantic,
            clippy::restriction,
            clippy::perf,
            clippy::deprecated,
            clippy::nursery,
            clippy::cargo
        )]
        #[deny(clippy::correctness)]
        impl clap::IntoApp for #struct_name {
            #into_app
            #augment_clap
        }
    };

    (tokens, attrs)
}

pub fn gen_for_enum(enum_name: &Ident, attrs: &[Attribute]) -> TokenStream {
    let app_name = env::var("CARGO_PKG_NAME").ok().unwrap_or_default();

    let attrs = Attrs::from_struct(
        Span::call_site(),
        attrs,
        Name::Assigned(quote!(#app_name)),
        Sp::call_site(DEFAULT_CASING),
        Sp::call_site(DEFAULT_ENV_CASING),
    );
    let name = attrs.cased_name();

    quote! {
        #[allow(dead_code, unreachable_code, unused_variables)]
        #[allow(
            clippy::style,
            clippy::complexity,
            clippy::pedantic,
            clippy::restriction,
            clippy::perf,
            clippy::deprecated,
            clippy::nursery,
            clippy::cargo
        )]
        #[deny(clippy::correctness)]
        impl clap::IntoApp for #enum_name {
            fn into_app<'b>() -> clap::App<'b> {
                let app = clap::App::new(#name)
                    .setting(clap::AppSettings::SubcommandRequiredElseHelp);
                <#enum_name as clap::IntoApp>::augment_clap(app)
            }

            fn augment_clap<'b>(app: clap::App<'b>) -> clap::App<'b> {
                <#enum_name as clap::Subcommand>::augment_subcommands(app)
            }

            fn into_app_for_update<'b>() -> clap::App<'b> {
                let app = clap::App::new(#name);
                <#enum_name as clap::IntoApp>::augment_clap_for_update(app)
            }

            fn augment_clap_for_update<'b>(app: clap::App<'b>) -> clap::App<'b> {
                <#enum_name as clap::Subcommand>::augment_subcommands_for_update(app)
            }
        }
    }
}

fn gen_into_app_fn(attrs: &[Attribute]) -> GenOutput {
    let app_name = env::var("CARGO_PKG_NAME").ok().unwrap_or_default();

    let attrs = Attrs::from_struct(
        Span::call_site(),
        attrs,
        Name::Assigned(quote!(#app_name)),
        Sp::call_site(DEFAULT_CASING),
        Sp::call_site(DEFAULT_ENV_CASING),
    );
    let name = attrs.cased_name();

    let tokens = quote! {
        fn into_app<'b>() -> clap::App<'b> {
            Self::augment_clap(clap::App::new(#name))
        }
        fn into_app_for_update<'b>() -> clap::App<'b> {
            Self::augment_clap_for_update(clap::App::new(#name))
        }
    };

    (tokens, attrs)
}

fn gen_augment_clap_fn(fields: &Punctuated<Field, Comma>, parent_attribute: &Attrs) -> TokenStream {
    let app_var = Ident::new("app", Span::call_site());
    let augmentation = gen_app_augmentation(fields, &app_var, parent_attribute, false);
    let augmentation_update = gen_app_augmentation(fields, &app_var, parent_attribute, true);
    quote! {
        fn augment_clap<'b>(#app_var: clap::App<'b>) -> clap::App<'b> {
            #augmentation
        }
        fn augment_clap_for_update<'b>(#app_var: clap::App<'b>) -> clap::App<'b> {
            #augmentation_update
        }
    }
}

fn gen_arg_enum_possible_values(ty: &Type) -> TokenStream {
    quote_spanned! { ty.span()=>
        .possible_values(&<#ty as clap::ArgEnum>::VARIANTS)
    }
}

/// Generate a block of code to add arguments/subcommands corresponding to
/// the `fields` to an app.
pub fn gen_app_augmentation(
    fields: &Punctuated<Field, Comma>,
    app_var: &Ident,
    parent_attribute: &Attrs,
    override_required: bool,
) -> TokenStream {
    let mut subcmds = fields.iter().filter_map(|field| {
        let attrs = Attrs::from_field(
            &field,
            parent_attribute.casing(),
            parent_attribute.env_casing(),
        );
        let kind = attrs.kind();
        if let Kind::Subcommand(ty) = &*kind {
            let subcmd_type = match (**ty, sub_type(&field.ty)) {
                (Ty::Option, Some(sub_type)) => sub_type,
                _ => &field.ty,
            };
            let required = if **ty == Ty::Option {
                quote!()
            } else {
                quote_spanned! { kind.span()=>
                    let #app_var = #app_var.setting(
                        clap::AppSettings::SubcommandRequiredElseHelp
                    );
                }
            };

            let span = field.span();
            let ts = if override_required {
                quote! {
                    let #app_var = <#subcmd_type as clap::Subcommand>::augment_subcommands_for_update( #app_var );
                }
            } else{
                quote! {
                    let #app_var = <#subcmd_type as clap::Subcommand>::augment_subcommands( #app_var );
                    #required
                }
            };
            Some((span, ts))
        } else {
            None
        }
    });
    let subcmd = subcmds.next().map(|(_, ts)| ts);
    if let Some((span, _)) = subcmds.next() {
        abort!(
            span,
            "multiple subcommand sets are not allowed, that's the second"
        );
    }

    let args = fields.iter().filter_map(|field| {
        let attrs = Attrs::from_field(
            field,
            parent_attribute.casing(),
            parent_attribute.env_casing(),
        );
        let kind = attrs.kind();
        match &*kind {
            Kind::Subcommand(_)
            | Kind::Skip(_)
            | Kind::FromGlobal(_)
            | Kind::ExternalSubcommand => None,
            Kind::Flatten => {
                let ty = &field.ty;
                Some(quote_spanned! { kind.span()=>
                    let #app_var = <#ty as clap::IntoApp>::augment_clap(#app_var);
                })
            }
            Kind::Arg(ty) => {
                let convert_type = match **ty {
                    Ty::Vec | Ty::Option => sub_type(&field.ty).unwrap_or(&field.ty),
                    Ty::OptionOption | Ty::OptionVec => {
                        sub_type(&field.ty).and_then(sub_type).unwrap_or(&field.ty)
                    }
                    _ => &field.ty,
                };

                let occurrences = *attrs.parser().kind == ParserKind::FromOccurrences;
                let flag = *attrs.parser().kind == ParserKind::FromFlag;

                let parser = attrs.parser();
                let func = &parser.func;

                let validator = match *parser.kind {
                    _ if attrs.is_enum() => quote!(),
                    ParserKind::TryFromStr => quote_spanned! { func.span()=>
                        .validator(|s| {
                            #func(s)
                            .map(|_: #convert_type| ())
                        })
                    },
                    ParserKind::TryFromOsStr => quote_spanned! { func.span()=>
                        .validator_os(|s| #func(s).map(|_: #convert_type| ()))
                    },
                    _ => quote!(),
                };

                let modifier = match **ty {
                    Ty::Bool => quote!(),

                    Ty::Option => {
                        let mut possible_values = quote!();

                        if attrs.is_enum() {
                            if let Some(subty) = subty_if_name(&field.ty, "Option") {
                                possible_values = gen_arg_enum_possible_values(subty);
                            }
                        };

                        quote_spanned! { ty.span()=>
                            .takes_value(true)
                            #possible_values
                            #validator
                        }
                    }

                    Ty::OptionOption => quote_spanned! { ty.span()=>
                        .takes_value(true)
                        .multiple(false)
                        .min_values(0)
                        .max_values(1)
                        #validator
                    },

                    Ty::OptionVec => quote_spanned! { ty.span()=>
                        .takes_value(true)
                        .multiple(true)
                        .min_values(0)
                        #validator
                    },

                    Ty::Vec => {
                        let mut possible_values = quote!();

                        if attrs.is_enum() {
                            if let Some(subty) = subty_if_name(&field.ty, "Vec") {
                                possible_values = gen_arg_enum_possible_values(subty);
                            }
                        };

                        quote_spanned! { ty.span()=>
                            .takes_value(true)
                            .multiple(true)
                            #possible_values
                            #validator
                        }
                    }

                    Ty::Other if occurrences => quote_spanned! { ty.span()=>
                        .multiple_occurrences(true)
                    },

                    Ty::Other if flag => quote_spanned! { ty.span()=>
                        .takes_value(false)
                        .multiple(false)
                    },

                    Ty::Other => {
                        let required = !attrs.has_method("default_value") && !override_required;
                        let mut possible_values = quote!();

                        if attrs.is_enum() {
                            possible_values = gen_arg_enum_possible_values(&field.ty);
                        };

                        quote_spanned! { ty.span()=>
                            .takes_value(true)
                            .required(#required)
                            #possible_values
                            #validator
                        }
                    }
                };

                let name = attrs.cased_name();
                let methods = attrs.field_methods();

                Some(quote_spanned! { field.span()=>
                    let #app_var = #app_var.arg(
                        clap::Arg::new(#name)
                            #modifier
                            #methods
                    );
                })
            }
        }
    });

    let app_methods = parent_attribute.top_level_methods();
    let version = parent_attribute.version();
    quote! {{
        let #app_var = #app_var#app_methods;
        #( #args )*
        #subcmd
        #app_var#version
    }}
}
<|endoftext|>


<commit_before>//! Symbol frequency model implemented by Binary Indexed Tree.

use std::boxed::Box;
use std::vec::Vec;
use super::Model;
use super::Parameters;
use super::super::Result;
use super::super::Error;

/// Adaptive model that uses a Binary Indexed Tree for storing cumulative frequencies.
pub struct AdaptiveTreeModel {
    /// Tree of cumulative frequencies.
    tree: Vec<u64>,
    /// Arithmetic parameters.
    params: Parameters,
}

trait LastOne<T> {
    fn last_one(self) -> T;
}

impl LastOne<usize> for usize {
    fn last_one(self) -> usize {
        self & self.wrapping_neg()
    }
}

impl AdaptiveTreeModel {
    /// Initializes the model with the given parameters.
    pub fn new(p: Parameters) -> Box<AdaptiveTreeModel> {
        let mut m = AdaptiveTreeModel {
            tree: vec![0; p.symbol_count],
            params: p,
        };

        m.tree[0] = 1u64;
        for i in 1..m.tree.len() {
            m.tree[i] = i.last_one() as u64;
        }

        return Box::new(m);
    }

    /// Returns the cumulated frequency of the symbol
    fn get_frequency_single(&self, symbol: usize) -> u64 {
        let mut i = symbol;
        let mut sum = self.tree[0];
        while i > 0 {
            sum += self.tree[i];
            i -= i.last_one();
        }
        return sum;
    }

    /// Updates the cumulative frequencies for the given symbol.
    fn update(&mut self, symbol: usize) {
        if symbol == 0 {
            self.tree[0] += 1;
        } else {
            let mut i = symbol;
            while i <= self.params.symbol_eof {
                self.tree[i] += 1;
                i += i.last_one();
            }
        }
    }
}

impl Model for AdaptiveTreeModel {
    fn parameters<'a>(&'a self) -> &'a Parameters {
        &self.params
    }

    fn total_frequency(&self) -> u64 {
        self.get_frequency_single(self.params.symbol_eof)
    }

    fn get_frequency(&mut self, symbol: usize) -> Result<(u64, u64)> {
        if symbol > self.params.symbol_eof {
            return Err(Error::InvalidInput);
        }

        let result = if symbol > 0 {
            let mut sumh = 0u64;
            let mut suml = 0u64;
            let mut h = symbol;
            let mut l = symbol - 1;
            while h != l {
                if h > l {
                    sumh += self.tree[h];
                    h -= h.last_one();
                } else {
                    suml += self.tree[l];
                    l -= l.last_one();
                }
            }

            let sumr = self.get_frequency_single(h);
            (suml + sumr, sumh + sumr)
        } else {
            (0u64, self.tree[0])
        };

        self.update(symbol);
        return Ok(result);
    }

    fn get_symbol(&mut self, value: u64) -> Result<(usize, u64, u64)> {
        let mut m = self.params.symbol_eof;
        let mut i = 0usize;
        let mut v = value;
        while (m > 0) && (i < self.params.symbol_eof) {
            let t = i + m;
            if v >= self.tree[t] {
                i = t;
                v -= self.tree[t];
            }
            m >>= 1;
        }

        if v > 0 {
            Err(Error::InvalidInput)
        } else {
            match self.get_frequency(i) {
                Ok((l, h)) => Ok((i, l, h)),
                Err(e) => Err(e)
            }
        }
    }
}

<commit_msg>Try to fix decoding error<commit_after>//! Symbol frequency model implemented by Binary Indexed Tree.

use std::boxed::Box;
use std::vec::Vec;
use super::Model;
use super::Parameters;
use super::super::Result;
use super::super::Error;

/// Adaptive model that uses a Binary Indexed Tree for storing cumulative frequencies.
pub struct AdaptiveTreeModel {
    /// Tree of cumulative frequencies.
    tree: Vec<u64>,
    /// Arithmetic parameters.
    params: Parameters,
}

trait LastOne<T> {
    fn last_one(self) -> T;
}

impl LastOne<usize> for usize {
    fn last_one(self) -> usize {
        self & self.wrapping_neg()
    }
}

impl AdaptiveTreeModel {
    /// Initializes the model with the given parameters.
    pub fn new(p: Parameters) -> Box<AdaptiveTreeModel> {
        let mut m = AdaptiveTreeModel {
            tree: vec![0; p.symbol_count],
            params: p,
        };

        m.tree[0] = 1u64;
        for i in 1..m.tree.len() {
            m.tree[i] = i.last_one() as u64;
        }

        return Box::new(m);
    }

    /// Returns the cumulated frequency of the symbol
    fn get_frequency_single(&self, symbol: usize) -> u64 {
        let mut i = symbol;
        let mut sum = self.tree[0];
        while i > 0 {
            sum += self.tree[i];
            i -= i.last_one();
        }
        return sum;
    }

    /// Updates the cumulative frequencies for the given symbol.
    fn update(&mut self, symbol: usize) {
        if symbol == 0 {
            self.tree[0] += 1;
        } else {
            let mut i = symbol;
            while i <= self.params.symbol_eof {
                self.tree[i] += 1;
                i += i.last_one();
            }
        }
    }
}

impl Model for AdaptiveTreeModel {
    fn parameters<'a>(&'a self) -> &'a Parameters {
        &self.params
    }

    fn total_frequency(&self) -> u64 {
        self.get_frequency_single(self.params.symbol_eof)
    }

    fn get_frequency(&mut self, symbol: usize) -> Result<(u64, u64)> {
        if symbol > self.params.symbol_eof {
            return Err(Error::InvalidInput);
        }

        let result = if symbol > 0 {
            let mut sumh = 0u64;
            let mut suml = 0u64;
            let mut h = symbol;
            let mut l = symbol - 1;
            while h != l {
                if h > l {
                    sumh += self.tree[h];
                    h -= h.last_one();
                } else {
                    suml += self.tree[l];
                    l -= l.last_one();
                }
            }

            let sumr = self.get_frequency_single(h);
            (suml + sumr, sumh + sumr)
        } else {
            (0u64, self.tree[0])
        };

        self.update(symbol);
        return Ok(result);
    }

    fn get_symbol(&mut self, value: u64) -> Result<(usize, u64, u64)> {
        let mut m = self.params.symbol_eof;
        let mut i = 0usize;
        let mut v = value;
        while (m > 0) && (i < self.params.symbol_eof) {
            let ti = i + m;
            let tv = self.tree[ti];

            if v >= tv {
                i = ti;
                v -= tv;
            }
            m >>= 1;
        }

        let (l, h) = try!(self.get_frequency(i));
        self.update(i);
        return Ok((i, l, h));
    }
}

<|endoftext|>


<commit_before>use traits::Grid;
use traits::Engine;
use traits::Consumer;

pub struct Sequential<G: Grid, C: Consumer> {
    grid: G,
    consumer: C
}

impl<G: Grid, C: Consumer> Sequential<G, C> {

    pub fn new(grid: G, consumer: C) -> Self {
        Sequential {
            grid: grid,
            consumer: consumer
        }
    }
}

impl<G: Grid, C: Consumer> Engine for Sequential<G, C> {
    fn run_times(&mut self, times: i64) {
        for _ in 0..times {
            self.grid.step();
            self.consumer.consume(&self.grid);
        }
    }
}<commit_msg>Support for the new constraint.<commit_after>use traits::Cell;
use traits::Consumer;
use traits::Grid;
use traits::Engine;

pub struct Sequential<C: Cell, G: Grid<Cell=C>, Con: Consumer<Cell=C>> {
    grid: G,
    consumer: Con
}

impl<C: Cell, G: Grid<Cell=C>, Con: Consumer<Cell=C>> Sequential<C, G, Con> {

    pub fn new(grid: G, consumer: Con) -> Self {
        Sequential {
            grid: grid,
            consumer: consumer
        }
    }
}

impl<C: Cell, G: Grid<Cell=C>, Con: Consumer<Cell=C>> Engine for Sequential<C, G, Con> {
    fn run_times(&mut self, times: i64) {
        for _ in 0..times {
            self.grid.step();
            self.consumer.consume(&self.grid);
        }
    }
}<|endoftext|>


<commit_before>use std::time::Duration;

use wlroots_sys::{wlr_event_keyboard_key, wlr_key_state, xkb_keysym_t, xkb_state,
                  xkb_state_key_get_syms};

pub type Key = xkb_keysym_t;

#[derive(Debug)]
pub struct KeyEvent {
    key: *mut wlr_event_keyboard_key,
    xkb_state: *mut xkb_state
}

impl KeyEvent {
    /// Constructs a KeyEvent from the raw key event pointer information.
    pub(crate) unsafe fn new(key: *mut wlr_event_keyboard_key, xkb_state: *mut xkb_state) -> Self {
        KeyEvent { key, xkb_state }
    }

    /// Gets the raw keycode from the device.
    ///
    /// Usually you want to use `KeyEvent::input_keys` since you care about what
    /// value XKB says this is.
    pub fn keycode(&self) -> u32 {
        unsafe { (*self.key).keycode + 8 }
    }

    /// Get how long the key has been pressed down, in milliseconds.
    pub fn time_msec(&self) -> Duration {
        Duration::from_millis(unsafe { (*self.key).time_msec } as u64)
    }

    /// TODO What is this?
    pub fn update_state(&self) -> bool {
        unsafe { (*self.key).update_state }
    }

    /// Get the pressed/released state of the key.
    pub fn key_state(&self) -> wlr_key_state {
        unsafe { (*self.key).state }
    }

    /// Gets the keys that are pressed using XKB to convert them to a more
    /// programmer friendly form.
    pub fn pressed_keys(&self) -> Vec<Key> {
        unsafe {
            let mut syms = 0 as *const xkb_keysym_t;
            let key_length = xkb_state_key_get_syms(self.xkb_state, self.keycode(), &mut syms);
            (0..key_length).map(|index| *syms.offset(index as isize))
                           .collect()
        }
    }
}
<commit_msg>Don't increment keycode automatically<commit_after>use std::time::Duration;

use wlroots_sys::{wlr_event_keyboard_key, wlr_key_state, xkb_keysym_t, xkb_state,
                  xkb_state_key_get_syms};

pub type Key = xkb_keysym_t;

#[derive(Debug)]
pub struct KeyEvent {
    key: *mut wlr_event_keyboard_key,
    xkb_state: *mut xkb_state
}

impl KeyEvent {
    /// Constructs a KeyEvent from the raw key event pointer information.
    pub(crate) unsafe fn new(key: *mut wlr_event_keyboard_key, xkb_state: *mut xkb_state) -> Self {
        KeyEvent { key, xkb_state }
    }

    /// Gets the raw keycode from the device.
    ///
    /// Usually you want to use `KeyEvent::input_keys` since you care about what
    /// value XKB says this is.
    pub fn keycode(&self) -> u32 {
        unsafe { (*self.key).keycode }
    }

    /// Get how long the key has been pressed down, in milliseconds.
    pub fn time_msec(&self) -> Duration {
        Duration::from_millis(unsafe { (*self.key).time_msec } as u64)
    }

    /// TODO What is this?
    pub fn update_state(&self) -> bool {
        unsafe { (*self.key).update_state }
    }

    /// Get the pressed/released state of the key.
    pub fn key_state(&self) -> wlr_key_state {
        unsafe { (*self.key).state }
    }

    /// Gets the keys that are pressed using XKB to convert them to a more
    /// programmer friendly form.
    pub fn pressed_keys(&self) -> Vec<Key> {
        unsafe {
            let mut syms = 0 as *const xkb_keysym_t;
            let key_length = xkb_state_key_get_syms(self.xkb_state, self.keycode() + 8, &mut syms);
            (0..key_length).map(|index| *syms.offset(index as isize))
                           .collect()
        }
    }
}
<|endoftext|>


<commit_before>use Bytes;
use write::{OutputFileWriter, OutputChunkWriter};
use client::GetResponse;
use hyper::client::Client;
use hyper::error::Error;
use hyper::header::{ByteRangeSpec, Headers, Range};
use pbr::{MultiBar, Pipe, ProgressBar, Units};
use response::CheckResponseStatus;
use std::cmp::min;
use std::io::Read;
use std::thread;
use std::time::{Instant, Duration};

/// Constant to represent the length of the buffer to download
/// the remote content
const DOWNLOAD_BUFFER_BYTES: usize = 1024 * 64;

/// Constant to represent the refresh interval (in milliseconds)
/// for the CLI
const PROGRESS_UPDATE_INTERVAL_MILLIS: u64 = 500;

/// Represents a range between two Bytes types
#[derive(Debug, PartialEq)]
struct RangeBytes(Bytes, Bytes);

/// Function to get the current chunk length, based on the chunk index.
fn get_chunk_length(chunk_index: u64,
                    content_length: Bytes,
                    global_chunk_length: Bytes)
                    -> Option<RangeBytes> {

    if content_length == 0 || global_chunk_length == 0 {
        return None;
    }

    let b_range: Bytes = chunk_index * global_chunk_length;

    if b_range >= (content_length - 1) {
        return None;
    }

    let e_range: Bytes = min(content_length - 1,
                             ((chunk_index + 1) * global_chunk_length) - 1);

    Some(RangeBytes(b_range, e_range))

}

/// Function to get the HTTP header to send to the file server, for a chunk (specified by its index)
fn get_header_from_index(chunk_index: u64,
                         content_length: Bytes,
                         global_chunk_length: Bytes)
                         -> Option<(Headers, RangeBytes)> {

    match get_chunk_length(chunk_index, content_length, global_chunk_length) {
        Some(range) => {
            let mut header = Headers::new();
            header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(range.0, range.1)]));
            Some((header, RangeBytes(range.0, range.1 - range.0)))
        }
        None => None,
    }

}

/// Function to get from the server the content of a chunk.
/// This function returns a Result type - Bytes if the content of the header is accessible, an Error type otherwise.
fn download_a_chunk(http_client: &Client,
                    http_header: Headers,
                    mut chunk_writer: OutputChunkWriter,
                    url: &str,
                    mpb: &mut ProgressBar<Pipe>)
                    -> Result<Bytes, Error> {

    let mut body = http_client.get_http_response_using_headers(url, http_header).unwrap();
    if !body.check_partialcontent_status() {
        return Err(Error::Status);
    }
    let mut bytes_buffer = [0; DOWNLOAD_BUFFER_BYTES];
    let mut sum_bytes = 0;

    let progress_update_interval = Duration::from_millis(PROGRESS_UPDATE_INTERVAL_MILLIS);
    let mut last_progress_bytes = 0;
    let mut last_progress_time = Instant::now() - progress_update_interval;

    while let Ok(n) = body.read(&mut bytes_buffer) {
        if n == 0 {
            return Ok(sum_bytes);
        }

        chunk_writer.write(sum_bytes, &bytes_buffer[0..n]);

        sum_bytes += n as u64;

        // Update the CLI
        if Instant::now().duration_since(last_progress_time) > progress_update_interval {
            last_progress_time = Instant::now();
            let progress_bytes_delta = sum_bytes - last_progress_bytes;
            last_progress_bytes = sum_bytes;
            mpb.add(progress_bytes_delta);
        }
    }
    mpb.add(sum_bytes - last_progress_bytes);
    return Ok(0u64);
}

/// Function to download each chunk of a remote content (given by its URL).
/// This function takes as parameters:
/// * the remote content length,
/// * a mutable reference to share between threads, which contains each chunk,
/// * the number of chunks that contains the remote content,
/// * the URL of the remote content server.
pub fn download_chunks(content_length: u64,
                       mut out_file: OutputFileWriter,
                       nb_chunks: u64,
                       url: &str) {

    // let mut downloaded_chunks: Arc<Mutex<Chunks>> =
    //     Arc::new(Mutex::new(Chunks::with_capacity(nb_chunks as usize)));
    let global_chunk_length: u64 = (content_length / nb_chunks) + 1;
    let mut jobs = vec![];

    let mut mpb = MultiBar::new();
    mpb.println(&format!("Downloading {} chunks: ", nb_chunks));

    for chunk_index in 0..nb_chunks {

        let (http_header, RangeBytes(chunk_offset, chunk_length)) =
            get_header_from_index(chunk_index, content_length, global_chunk_length).unwrap();
        let hyper_client = Client::new();
        let url_clone = String::from(url);

        // Progress bar customization
        let mut mp = mpb.create_bar(chunk_length);
        mp.tick_format("▏▎▍▌▋▊▉██▉▊▋▌▍▎▏");
        mp.format("|#--|");
        mp.show_tick = true;
        mp.show_speed = true;
        mp.show_percent = true;
        mp.show_counter = false;
        mp.show_time_left = true;
        mp.set_units(Units::Bytes);
        mp.message(&format!("Chunk {} ", chunk_index));

        let chunk_writer = out_file.get_chunk_writer(chunk_offset);
        jobs.push(thread::spawn(move || match download_a_chunk(&hyper_client,
                                                                    http_header,
                                                                    chunk_writer,
                                                                    &url_clone,
                                                                    &mut mp) {
            Ok(bytes_written) => {
                if bytes_written > 0 {
                    mp.finish();
                } else {
                    panic!("The downloaded chunk {} is empty", chunk_index);
                }
            }
            Err(error) => {
                panic!("Cannot download the chunk {}, due to error {}",
                       chunk_index,
                       error);
            }
        }));
    }

    mpb.listen();

    for child in jobs {
        let _ = child.join();
    }

}

#[cfg(test)]
mod test_chunk_length {

    use super::get_chunk_length;
    use super::RangeBytes;

    #[test]
    fn wrong_content_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 15, 0));
    }

    #[test]
    fn wrong_global_chunk_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 15));
    }

    #[test]
    fn wrong_length_parameters_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 0));
    }

    #[test]
    fn get_the_first_range_in_chunk() {
        assert_eq!(Some(RangeBytes(0, 249)), get_chunk_length(0, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_chunk() {
        assert_eq!(Some(RangeBytes(750, 999)), get_chunk_length(3, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_shorten_chunk() {
        assert_eq!(Some(RangeBytes(750, 997)), get_chunk_length(3, 998, 250));
    }

    #[test]
    fn wrong_index_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(4, 1000, 250));
    }

}

#[cfg(test)]
mod test_header {

    use super::{get_header_from_index, RangeBytes};
    use hyper::header::{ByteRangeSpec, Headers, Range};

    #[test]
    fn wrong_chunk_length_should_return_none() {
        assert_eq!(None, get_header_from_index(0, 0, 0));
    }

    #[test]
    fn good_chunk_length_should_return_a_good_header() {
        let mut test_header = Headers::new();
        test_header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(750, 997)]));
        assert_eq!(Some((test_header, RangeBytes(750, 247))),
                   get_header_from_index(3, 998, 250));
    }

}
<commit_msg>rewrite get_header_from_index<commit_after>use Bytes;
use write::{OutputFileWriter, OutputChunkWriter};
use client::GetResponse;
use hyper::client::Client;
use hyper::error::Error;
use hyper::header::{ByteRangeSpec, Headers, Range};
use pbr::{MultiBar, Pipe, ProgressBar, Units};
use response::CheckResponseStatus;
use std::cmp::min;
use std::io::Read;
use std::thread;
use std::time::{Instant, Duration};

/// Constant to represent the length of the buffer to download
/// the remote content
const DOWNLOAD_BUFFER_BYTES: usize = 1024 * 64;

/// Constant to represent the refresh interval (in milliseconds)
/// for the CLI
const PROGRESS_UPDATE_INTERVAL_MILLIS: u64 = 500;

/// Represents a range between two Bytes types
#[derive(Debug, PartialEq)]
struct RangeBytes(Bytes, Bytes);

/// Function to get the current chunk length, based on the chunk index.
fn get_chunk_length(chunk_index: u64,
                    content_length: Bytes,
                    global_chunk_length: Bytes)
                    -> Option<RangeBytes> {

    if content_length == 0 || global_chunk_length == 0 {
        return None;
    }

    let b_range: Bytes = chunk_index * global_chunk_length;

    if b_range >= (content_length - 1) {
        return None;
    }

    let e_range: Bytes = min(content_length - 1,
                             ((chunk_index + 1) * global_chunk_length) - 1);

    Some(RangeBytes(b_range, e_range))

}

/// Function to get the HTTP header to send to the file server, for a chunk (specified by its index)
fn get_header_from_index(chunk_index: u64,
                         content_length: Bytes,
                         global_chunk_length: Bytes)
                         -> Option<(Headers, RangeBytes)> {

    get_chunk_length(chunk_index, content_length, global_chunk_length).map(
        |range| {
            let mut header = Headers::new();
            header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(range.0, range.1)]));
            (header, RangeBytes(range.0, range.1 - range.0))
        }
    )
}

/// Function to get from the server the content of a chunk.
/// This function returns a Result type - Bytes if the content of the header is accessible, an Error type otherwise.
fn download_a_chunk(http_client: &Client,
                    http_header: Headers,
                    mut chunk_writer: OutputChunkWriter,
                    url: &str,
                    mpb: &mut ProgressBar<Pipe>)
                    -> Result<Bytes, Error> {

    let mut body = http_client.get_http_response_using_headers(url, http_header).unwrap();
    if !body.check_partialcontent_status() {
        return Err(Error::Status);
    }
    let mut bytes_buffer = [0; DOWNLOAD_BUFFER_BYTES];
    let mut sum_bytes = 0;

    let progress_update_interval = Duration::from_millis(PROGRESS_UPDATE_INTERVAL_MILLIS);
    let mut last_progress_bytes = 0;
    let mut last_progress_time = Instant::now() - progress_update_interval;

    while let Ok(n) = body.read(&mut bytes_buffer) {
        if n == 0 {
            return Ok(sum_bytes);
        }

        chunk_writer.write(sum_bytes, &bytes_buffer[0..n]);

        sum_bytes += n as u64;

        // Update the CLI
        if Instant::now().duration_since(last_progress_time) > progress_update_interval {
            last_progress_time = Instant::now();
            let progress_bytes_delta = sum_bytes - last_progress_bytes;
            last_progress_bytes = sum_bytes;
            mpb.add(progress_bytes_delta);
        }
    }
    mpb.add(sum_bytes - last_progress_bytes);
    return Ok(0u64);
}

/// Function to download each chunk of a remote content (given by its URL).
/// This function takes as parameters:
/// * the remote content length,
/// * a mutable reference to share between threads, which contains each chunk,
/// * the number of chunks that contains the remote content,
/// * the URL of the remote content server.
pub fn download_chunks(content_length: u64,
                       mut out_file: OutputFileWriter,
                       nb_chunks: u64,
                       url: &str) {

    // let mut downloaded_chunks: Arc<Mutex<Chunks>> =
    //     Arc::new(Mutex::new(Chunks::with_capacity(nb_chunks as usize)));
    let global_chunk_length: u64 = (content_length / nb_chunks) + 1;
    let mut jobs = vec![];

    let mut mpb = MultiBar::new();
    mpb.println(&format!("Downloading {} chunks: ", nb_chunks));

    for chunk_index in 0..nb_chunks {

        let (http_header, RangeBytes(chunk_offset, chunk_length)) =
            get_header_from_index(chunk_index, content_length, global_chunk_length).unwrap();
        let hyper_client = Client::new();
        let url_clone = String::from(url);

        // Progress bar customization
        let mut mp = mpb.create_bar(chunk_length);
        mp.tick_format("▏▎▍▌▋▊▉██▉▊▋▌▍▎▏");
        mp.format("|#--|");
        mp.show_tick = true;
        mp.show_speed = true;
        mp.show_percent = true;
        mp.show_counter = false;
        mp.show_time_left = true;
        mp.set_units(Units::Bytes);
        mp.message(&format!("Chunk {} ", chunk_index));

        let chunk_writer = out_file.get_chunk_writer(chunk_offset);
        jobs.push(thread::spawn(move || match download_a_chunk(&hyper_client,
                                                                    http_header,
                                                                    chunk_writer,
                                                                    &url_clone,
                                                                    &mut mp) {
            Ok(bytes_written) => {
                if bytes_written > 0 {
                    mp.finish();
                } else {
                    panic!("The downloaded chunk {} is empty", chunk_index);
                }
            }
            Err(error) => {
                panic!("Cannot download the chunk {}, due to error {}",
                       chunk_index,
                       error);
            }
        }));
    }

    mpb.listen();

    for child in jobs {
        let _ = child.join();
    }

}

#[cfg(test)]
mod test_chunk_length {

    use super::get_chunk_length;
    use super::RangeBytes;

    #[test]
    fn wrong_content_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 15, 0));
    }

    #[test]
    fn wrong_global_chunk_length_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 15));
    }

    #[test]
    fn wrong_length_parameters_should_return_none() {
        assert_eq!(None, get_chunk_length(0, 0, 0));
    }

    #[test]
    fn get_the_first_range_in_chunk() {
        assert_eq!(Some(RangeBytes(0, 249)), get_chunk_length(0, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_chunk() {
        assert_eq!(Some(RangeBytes(750, 999)), get_chunk_length(3, 1000, 250));
    }

    #[test]
    fn get_the_last_range_in_shorten_chunk() {
        assert_eq!(Some(RangeBytes(750, 997)), get_chunk_length(3, 998, 250));
    }

    #[test]
    fn wrong_index_parameter_should_return_none() {
        assert_eq!(None, get_chunk_length(4, 1000, 250));
    }

}

#[cfg(test)]
mod test_header {

    use super::{get_header_from_index, RangeBytes};
    use hyper::header::{ByteRangeSpec, Headers, Range};

    #[test]
    fn wrong_chunk_length_should_return_none() {
        assert_eq!(None, get_header_from_index(0, 0, 0));
    }

    #[test]
    fn good_chunk_length_should_return_a_good_header() {
        let mut test_header = Headers::new();
        test_header.set(Range::Bytes(vec![ByteRangeSpec::FromTo(750, 997)]));
        assert_eq!(Some((test_header, RangeBytes(750, 247))),
                   get_header_from_index(3, 998, 250));
    }

}
<|endoftext|>


<commit_before>// See LICENSE file for copyright and license details.

use std::f32::consts::{PI};
use std::num::{Float};
use cgmath::{Vector3, Vector, Rad, Angle, rad};
use common::types::{ZInt, ZFloat, MapPos};
use zgl::types::{VertexCoord, WorldPos};
use core::geom;

pub use core::geom::{HEX_IN_RADIUS, HEX_EX_RADIUS};

pub const MINIMAL_LIFT_HEIGHT: ZFloat = 0.1;

pub fn map_pos_to_world_pos(i: &MapPos) -> WorldPos {
    WorldPos{v: geom::map_pos_to_world_pos(i).extend(0.0)}
}

pub fn lift(v: Vector3<ZFloat>) -> Vector3<ZFloat> {
    let mut v = v;
    v.z += MINIMAL_LIFT_HEIGHT;
    v
}

pub fn index_to_circle_vertex(count: ZInt, i: ZInt) -> VertexCoord {
    let n = PI / 2.0 + 2.0 * PI * (i as ZFloat) / (count as ZFloat);
    VertexCoord {
        v: Vector3{x: n.cos(), y: n.sin(), z: 0.0}.mul_s(HEX_EX_RADIUS)
    }
}

pub fn index_to_hex_vertex(i: ZInt) -> VertexCoord {
    index_to_circle_vertex(6, i)
}

pub fn index_to_hex_vertex_s(scale: ZFloat, i: ZInt) -> VertexCoord {
    let v = index_to_hex_vertex(i).v.mul_s(scale);
    VertexCoord{v: v}
}

// TODO: ZFloat -> WorldDistance
pub fn dist(a: &WorldPos, b: &WorldPos) -> ZFloat {
    let dx = (b.v.x - a.v.x).abs();
    let dy = (b.v.y - a.v.y).abs();
    let dz = (b.v.z - a.v.z).abs();
    ((dx.powi(2) + dy.powi(2) + dz.powi(2)) as ZFloat).sqrt()
}

pub fn get_rot_angle(a: &WorldPos, b: &WorldPos) -> Rad<ZFloat> {
    let diff = b.v - a.v;
    let angle = Float::atan2(diff.x, diff.y);
    rad(-angle).normalize()
}

#[cfg(test)]
mod tests {
    use std::f32::consts::{PI};
    use std::num::{Float};
    use cgmath::{Vector3};
    use common::types::{ZFloat};
    use zgl::types::{WorldPos};
    use super::{get_rot_angle, index_to_circle_vertex};

    const EPS: ZFloat = 0.001;

    #[test]
    fn test_get_rot_angle_30_deg() {
        let count = 12;
        for i in 0 .. count {
            let a = WorldPos{v: Vector3{x: 0.0, y: 0.0, z: 0.0}};
            let b = WorldPos{v: index_to_circle_vertex(count, i).v};
            let expected_angle = i as ZFloat * (PI * 2.0) / (count as ZFloat);
            let angle = get_rot_angle(&a, &b);
            let diff = (expected_angle - angle.s).abs();
            assert!(diff < EPS, "{} != {}", expected_angle, angle.s);
        }
    }
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<commit_msg>geom.rs: 'MINIMAL_' -> 'MIN_'<commit_after>// See LICENSE file for copyright and license details.

use std::f32::consts::{PI};
use std::num::{Float};
use cgmath::{Vector3, Vector, Rad, Angle, rad};
use common::types::{ZInt, ZFloat, MapPos};
use zgl::types::{VertexCoord, WorldPos};
use core::geom;

pub use core::geom::{HEX_IN_RADIUS, HEX_EX_RADIUS};

pub const MIN_LIFT_HEIGHT: ZFloat = 0.1;

pub fn map_pos_to_world_pos(i: &MapPos) -> WorldPos {
    WorldPos{v: geom::map_pos_to_world_pos(i).extend(0.0)}
}

pub fn lift(v: Vector3<ZFloat>) -> Vector3<ZFloat> {
    let mut v = v;
    v.z += MIN_LIFT_HEIGHT;
    v
}

pub fn index_to_circle_vertex(count: ZInt, i: ZInt) -> VertexCoord {
    let n = PI / 2.0 + 2.0 * PI * (i as ZFloat) / (count as ZFloat);
    VertexCoord {
        v: Vector3{x: n.cos(), y: n.sin(), z: 0.0}.mul_s(HEX_EX_RADIUS)
    }
}

pub fn index_to_hex_vertex(i: ZInt) -> VertexCoord {
    index_to_circle_vertex(6, i)
}

pub fn index_to_hex_vertex_s(scale: ZFloat, i: ZInt) -> VertexCoord {
    let v = index_to_hex_vertex(i).v.mul_s(scale);
    VertexCoord{v: v}
}

// TODO: ZFloat -> WorldDistance
pub fn dist(a: &WorldPos, b: &WorldPos) -> ZFloat {
    let dx = (b.v.x - a.v.x).abs();
    let dy = (b.v.y - a.v.y).abs();
    let dz = (b.v.z - a.v.z).abs();
    ((dx.powi(2) + dy.powi(2) + dz.powi(2)) as ZFloat).sqrt()
}

pub fn get_rot_angle(a: &WorldPos, b: &WorldPos) -> Rad<ZFloat> {
    let diff = b.v - a.v;
    let angle = Float::atan2(diff.x, diff.y);
    rad(-angle).normalize()
}

#[cfg(test)]
mod tests {
    use std::f32::consts::{PI};
    use std::num::{Float};
    use cgmath::{Vector3};
    use common::types::{ZFloat};
    use zgl::types::{WorldPos};
    use super::{get_rot_angle, index_to_circle_vertex};

    const EPS: ZFloat = 0.001;

    #[test]
    fn test_get_rot_angle_30_deg() {
        let count = 12;
        for i in 0 .. count {
            let a = WorldPos{v: Vector3{x: 0.0, y: 0.0, z: 0.0}};
            let b = WorldPos{v: index_to_circle_vertex(count, i).v};
            let expected_angle = i as ZFloat * (PI * 2.0) / (count as ZFloat);
            let angle = get_rot_angle(&a, &b);
            let diff = (expected_angle - angle.s).abs();
            assert!(diff < EPS, "{} != {}", expected_angle, angle.s);
        }
    }
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<|endoftext|>