<commit_before>//! # OpenTelemetry Zipkin Exporter
//!
//! Collects OpenTelemetry spans and reports them to a given Zipkin
//! `collector` endpoint. See the [Zipkin Docs](https://zipkin.io/) for details
//! and deployment information.
//!
//! ### Zipkin collector example
//!
//! This example expects a Zipkin collector running on `localhost:9411`.
//!
//! ```rust,no_run
//! use opentelemetry::{api::Key, global, sdk};
//! use opentelemetry_zipkin::ExporterConfig;
//! use std::net::{SocketAddr, IpAddr, Ipv4Addr};
//!
//! fn init_tracer() {
//!     let exporter = opentelemetry_zipkin::Exporter::from_config(
//!        ExporterConfig::builder()
//!            .with_service_name("opentelemetry-backend".to_owned())
//!            .with_service_endpoint(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080))
//!            .build());
//!     let provider = sdk::Provider::builder()
//!         .with_simple_exporter(exporter)
//!         .with_config(sdk::Config {
//!             default_sampler: Box::new(sdk::Sampler::AlwaysOn),
//!             ..Default::default()
//!         })
//!         .build();
//!
//!     global::set_provider(provider);
//! }
//! ```
//!
#![deny(missing_docs, unreachable_pub, missing_debug_implementations)]
#![cfg_attr(test, deny(warnings))]

#[macro_use]
extern crate typed_builder;

mod model;
mod uploader;

use model::{annotation, endpoint, span};
use opentelemetry::api;
use opentelemetry::exporter::trace;
use std::collections::HashMap;
use std::net;
use std::sync::Arc;
use std::time::{Duration, SystemTime};

/// Default Zipkin collector endpoint if none specified
static DEFAULT_COLLECTOR_ENDPOINT: &str = "127.0.0.1:9411";

/// Zipkin span exporter
#[derive(Debug)]
pub struct Exporter {
    config: ExporterConfig,
    uploader: uploader::Uploader,
}

/// Zipkin-specific configuration used to initialize the `Exporter`.
#[derive(Clone, Debug)]
pub struct ExporterConfig {
    local_endpoint: endpoint::Endpoint,
    collector_endpoint: String,
}

/// Builder for `ExporterConfig` struct.
#[derive(Debug)]
pub struct ExporterConfigBuilder {
    service_name: Option<String>,
    service_endpoint: Option<net::SocketAddr>,
    collector_endpoint: Option<String>,
}

impl Default for ExporterConfigBuilder {
    fn default() -> Self {
        ExporterConfigBuilder {
            collector_endpoint: None,
            service_name: None,
            service_endpoint: None,
        }
    }
}

impl ExporterConfig {
    /// Create an export config builder
    pub fn builder() -> ExporterConfigBuilder {
        ExporterConfigBuilder::default()
    }
}

impl ExporterConfigBuilder {
    /// Create `ExporterConfig` struct from current `ExporterConfigBuilder`
    pub fn build(&self) -> ExporterConfig {
        let local_endpoint: endpoint::Endpoint;
        let service_name = self
            .service_name
            .clone()
            .unwrap_or_else(|| "DEFAULT".to_owned());
        match self.service_endpoint {
            Some(socket_addr) => {
                match socket_addr.ip() {
                    net::IpAddr::V4(addr) => {
                        local_endpoint = endpoint::Endpoint::builder()
                            .service_name(service_name)
                            .ipv4(addr)
                            .port(socket_addr.port())
                            .build()
                    }
                    net::IpAddr::V6(addr) => {
                        local_endpoint = endpoint::Endpoint::builder()
                            .service_name(service_name)
                            .ipv6(addr)
                            .port(socket_addr.port())
                            .build()
                    }
                };
            }
            None => {
                local_endpoint = endpoint::Endpoint::builder()
                    .service_name(service_name)
                    .build()
            }
        };
        ExporterConfig {
            collector_endpoint: self
                .collector_endpoint
                .clone()
                .unwrap_or_else(|| DEFAULT_COLLECTOR_ENDPOINT.parse().unwrap()),
            local_endpoint,
        }
    }

    /// Assign the service name for `ConfigBuilder`
    pub fn with_service_name(&mut self, name: String) -> &mut Self {
        self.service_name = Some(name);
        self
    }

    /// Assign the service endpoint for `ConfigBuilder`
    pub fn with_service_endpoint(&mut self, endpoint: net::SocketAddr) -> &mut Self {
        self.service_endpoint = Some(endpoint);
        self
    }
}

impl Exporter {
    /// Creates new `Exporter` from a given `ExporterConfig`.
    pub fn from_config(config: ExporterConfig) -> Self {
        Exporter {
            config: config.clone(),
            uploader: uploader::Uploader::new(
                config.collector_endpoint,
                uploader::UploaderFormat::HTTP,
            ),
        }
    }
}

impl trace::SpanExporter for Exporter {
    /// Export spans to Zipkin collector.
    fn export(&self, batch: Vec<Arc<trace::SpanData>>) -> trace::ExportResult {
        let zipkin_spans: Vec<span::Span> = batch
            .into_iter()
            .map(|span| into_zipkin_span(&self.config, span))
            .collect();
        self.uploader.upload(span::ListOfSpans(zipkin_spans))
    }

    fn shutdown(&self) {}
}

/// Converts `api::Event` into an `annotation::Annotation`
impl Into<annotation::Annotation> for api::Event {
    fn into(self) -> annotation::Annotation {
        let timestamp = self
            .timestamp
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap_or_else(|_| Duration::from_secs(0))
            .as_micros() as u64;

        annotation::Annotation::builder()
            .timestamp(timestamp)
            .value(self.name)
            .build()
    }
}

/// Converts `api::SpanKind` into an `Option<span::Kind>`
fn into_zipkin_span_kind(kind: api::SpanKind) -> Option<span::Kind> {
    match kind {
        api::SpanKind::Client => Some(span::Kind::Client),
        api::SpanKind::Server => Some(span::Kind::Server),
        api::SpanKind::Producer => Some(span::Kind::Producer),
        api::SpanKind::Consumer => Some(span::Kind::Consumer),
        api::SpanKind::Internal => None,
    }
}

/// Converts a `trace::SpanData` to a `span::SpanData` for a given `ExporterConfig`, which can then
/// be ingested into a Zipkin collector.
fn into_zipkin_span(config: &ExporterConfig, span_data: Arc<trace::SpanData>) -> span::Span {
    let mut user_defined_span_kind = false;
    let tags = map_from_kvs(
        span_data
            .attributes
            .iter()
            .map(|(k, v)| {
                if k == &api::Key::new("span.kind") {
                    user_defined_span_kind = true;
                }
                api::KeyValue::new(k.clone(), v.clone())
            })
            .chain(
                span_data
                    .resource
                    .iter()
                    .map(|(k, v)| api::KeyValue::new(k.clone(), v.clone())),
            ),
    );

    span::Span::builder()
        .trace_id(format!(
            "{:032x}",
            span_data.span_context.trace_id().to_u128()
        ))
        .parent_id(format!("{:016x}", span_data.parent_span_id.to_u64()))
        .id(format!(
            "{:016x}",
            span_data.span_context.span_id().to_u64()
        ))
        .name(span_data.name.clone())
        .kind(if user_defined_span_kind {
            None
        } else {
            into_zipkin_span_kind(span_data.span_kind.clone())
        })
        .timestamp(
            span_data
                .start_time
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap_or_else(|_| Duration::from_secs(0))
                .as_micros() as u64,
        )
        .duration(
            span_data
                .end_time
                .duration_since(span_data.start_time)
                .unwrap_or_else(|_| Duration::from_secs(0))
                .as_micros() as u64,
        )
        .local_endpoint(config.local_endpoint.clone())
        .annotations(
            span_data
                .message_events
                .iter()
                .cloned()
                .map(Into::into)
                .collect(),
        )
        .tags(tags)
        .build()
}

fn map_from_kvs<T>(kvs: T) -> HashMap<String, String>
where
    T: IntoIterator<Item = api::KeyValue>,
{
    let mut map: HashMap<String, String> = HashMap::new();
    for kv in kvs {
        map.insert(kv.key.into(), kv.value.into());
    }
    map
}
<commit_msg>fix(zipkin): add with_collector_endpoint into config builder (#155)<commit_after>//! # OpenTelemetry Zipkin Exporter
//!
//! Collects OpenTelemetry spans and reports them to a given Zipkin
//! `collector` endpoint. See the [Zipkin Docs](https://zipkin.io/) for details
//! and deployment information.
//!
//! ### Zipkin collector example
//!
//! This example expects a Zipkin collector running on `localhost:9411`.
//!
//! ```rust,no_run
//! use opentelemetry::{api::Key, global, sdk};
//! use opentelemetry_zipkin::ExporterConfig;
//! use std::net::{SocketAddr, IpAddr, Ipv4Addr};
//!
//! fn init_tracer() {
//!     let exporter = opentelemetry_zipkin::Exporter::from_config(
//!        ExporterConfig::builder()
//!            .with_service_name("opentelemetry-backend".to_owned())
//!            .with_service_endpoint(SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080))
//!            .build());
//!     let provider = sdk::Provider::builder()
//!         .with_simple_exporter(exporter)
//!         .with_config(sdk::Config {
//!             default_sampler: Box::new(sdk::Sampler::AlwaysOn),
//!             ..Default::default()
//!         })
//!         .build();
//!
//!     global::set_provider(provider);
//! }
//! ```
//!
#![deny(missing_docs, unreachable_pub, missing_debug_implementations)]
#![cfg_attr(test, deny(warnings))]

#[macro_use]
extern crate typed_builder;

mod model;
mod uploader;

use model::{annotation, endpoint, span};
use opentelemetry::api;
use opentelemetry::exporter::trace;
use std::collections::HashMap;
use std::net;
use std::sync::Arc;
use std::time::{Duration, SystemTime};

/// Default Zipkin collector endpoint if none specified
static DEFAULT_COLLECTOR_ENDPOINT: &str = "127.0.0.1:9411";

/// Zipkin span exporter
#[derive(Debug)]
pub struct Exporter {
    config: ExporterConfig,
    uploader: uploader::Uploader,
}

/// Zipkin-specific configuration used to initialize the `Exporter`.
#[derive(Clone, Debug)]
pub struct ExporterConfig {
    local_endpoint: endpoint::Endpoint,
    collector_endpoint: String,
}

/// Builder for `ExporterConfig` struct.
#[derive(Debug)]
pub struct ExporterConfigBuilder {
    service_name: Option<String>,
    service_endpoint: Option<net::SocketAddr>,
    collector_endpoint: Option<String>,
}

impl Default for ExporterConfigBuilder {
    fn default() -> Self {
        ExporterConfigBuilder {
            collector_endpoint: None,
            service_name: None,
            service_endpoint: None,
        }
    }
}

impl ExporterConfig {
    /// Create an export config builder
    pub fn builder() -> ExporterConfigBuilder {
        ExporterConfigBuilder::default()
    }
}

impl ExporterConfigBuilder {
    /// Create `ExporterConfig` struct from current `ExporterConfigBuilder`
    pub fn build(&self) -> ExporterConfig {
        let local_endpoint: endpoint::Endpoint;
        let service_name = self
            .service_name
            .clone()
            .unwrap_or_else(|| "DEFAULT".to_owned());
        match self.service_endpoint {
            Some(socket_addr) => {
                match socket_addr.ip() {
                    net::IpAddr::V4(addr) => {
                        local_endpoint = endpoint::Endpoint::builder()
                            .service_name(service_name)
                            .ipv4(addr)
                            .port(socket_addr.port())
                            .build()
                    }
                    net::IpAddr::V6(addr) => {
                        local_endpoint = endpoint::Endpoint::builder()
                            .service_name(service_name)
                            .ipv6(addr)
                            .port(socket_addr.port())
                            .build()
                    }
                };
            }
            None => {
                local_endpoint = endpoint::Endpoint::builder()
                    .service_name(service_name)
                    .build()
            }
        };
        ExporterConfig {
            collector_endpoint: self
                .collector_endpoint
                .clone()
                .unwrap_or_else(|| DEFAULT_COLLECTOR_ENDPOINT.parse().unwrap()),
            local_endpoint,
        }
    }

    /// Assign the service name for `ConfigBuilder`
    pub fn with_service_name(&mut self, name: String) -> &mut Self {
        self.service_name = Some(name);
        self
    }

    /// Assign the service endpoint for `ConfigBuilder`
    pub fn with_service_endpoint(&mut self, endpoint: net::SocketAddr) -> &mut Self {
        self.service_endpoint = Some(endpoint);
        self
    }

    /// Assign the collector endpoint for `ConfigBuilder`
    pub fn with_collector_endpoint(&mut self, endpoint: String) -> &mut Self {
        self.collector_endpoint = Some(endpoint);
        self
    }
}

impl Exporter {
    /// Creates new `Exporter` from a given `ExporterConfig`.
    pub fn from_config(config: ExporterConfig) -> Self {
        Exporter {
            config: config.clone(),
            uploader: uploader::Uploader::new(
                config.collector_endpoint,
                uploader::UploaderFormat::HTTP,
            ),
        }
    }
}

impl trace::SpanExporter for Exporter {
    /// Export spans to Zipkin collector.
    fn export(&self, batch: Vec<Arc<trace::SpanData>>) -> trace::ExportResult {
        let zipkin_spans: Vec<span::Span> = batch
            .into_iter()
            .map(|span| into_zipkin_span(&self.config, span))
            .collect();
        self.uploader.upload(span::ListOfSpans(zipkin_spans))
    }

    fn shutdown(&self) {}
}

/// Converts `api::Event` into an `annotation::Annotation`
impl Into<annotation::Annotation> for api::Event {
    fn into(self) -> annotation::Annotation {
        let timestamp = self
            .timestamp
            .duration_since(SystemTime::UNIX_EPOCH)
            .unwrap_or_else(|_| Duration::from_secs(0))
            .as_micros() as u64;

        annotation::Annotation::builder()
            .timestamp(timestamp)
            .value(self.name)
            .build()
    }
}

/// Converts `api::SpanKind` into an `Option<span::Kind>`
fn into_zipkin_span_kind(kind: api::SpanKind) -> Option<span::Kind> {
    match kind {
        api::SpanKind::Client => Some(span::Kind::Client),
        api::SpanKind::Server => Some(span::Kind::Server),
        api::SpanKind::Producer => Some(span::Kind::Producer),
        api::SpanKind::Consumer => Some(span::Kind::Consumer),
        api::SpanKind::Internal => None,
    }
}

/// Converts a `trace::SpanData` to a `span::SpanData` for a given `ExporterConfig`, which can then
/// be ingested into a Zipkin collector.
fn into_zipkin_span(config: &ExporterConfig, span_data: Arc<trace::SpanData>) -> span::Span {
    let mut user_defined_span_kind = false;
    let tags = map_from_kvs(
        span_data
            .attributes
            .iter()
            .map(|(k, v)| {
                if k == &api::Key::new("span.kind") {
                    user_defined_span_kind = true;
                }
                api::KeyValue::new(k.clone(), v.clone())
            })
            .chain(
                span_data
                    .resource
                    .iter()
                    .map(|(k, v)| api::KeyValue::new(k.clone(), v.clone())),
            ),
    );

    span::Span::builder()
        .trace_id(format!(
            "{:032x}",
            span_data.span_context.trace_id().to_u128()
        ))
        .parent_id(format!("{:016x}", span_data.parent_span_id.to_u64()))
        .id(format!(
            "{:016x}",
            span_data.span_context.span_id().to_u64()
        ))
        .name(span_data.name.clone())
        .kind(if user_defined_span_kind {
            None
        } else {
            into_zipkin_span_kind(span_data.span_kind.clone())
        })
        .timestamp(
            span_data
                .start_time
                .duration_since(SystemTime::UNIX_EPOCH)
                .unwrap_or_else(|_| Duration::from_secs(0))
                .as_micros() as u64,
        )
        .duration(
            span_data
                .end_time
                .duration_since(span_data.start_time)
                .unwrap_or_else(|_| Duration::from_secs(0))
                .as_micros() as u64,
        )
        .local_endpoint(config.local_endpoint.clone())
        .annotations(
            span_data
                .message_events
                .iter()
                .cloned()
                .map(Into::into)
                .collect(),
        )
        .tags(tags)
        .build()
}

fn map_from_kvs<T>(kvs: T) -> HashMap<String, String>
where
    T: IntoIterator<Item = api::KeyValue>,
{
    let mut map: HashMap<String, String> = HashMap::new();
    for kv in kvs {
        map.insert(kv.key.into(), kv.value.into());
    }
    map
}
<|endoftext|>


<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#![crate_name="calculator"]
#![crate_type="bin"]

extern crate thrift;

use std::io::net::ip;
use thrift::protocol::{MessageType, Type};
use thrift::transport::Transport;
use thrift::protocol::Protocol;
use thrift::protocol::binary_protocol::BinaryProtocol;

// TODO: shall be generated
//#[allow(dead_code)]
struct CalculatorPingArgs;

impl CalculatorPingArgs {

    #[allow(unused_variables)]
    fn write(&self, oprot: &Protocol, transport: & mut Transport) {
        oprot.write_struct_begin(transport, "Calculator_ping_args");
        oprot.write_field_stop(transport);
        oprot.write_struct_end(transport);
    }
}

struct CalculatorPingResult;

impl CalculatorPingResult {

    #[allow(unused_variables)]
    fn read(&self, iprot: &Protocol, transport: & mut Transport) {
        let fname = iprot.read_struct_begin(transport);

        loop {
            let (fname, ftype, fid) = iprot.read_field_begin(transport);
            if ftype == Type::TStop {
              break;
            }
            iprot.skip(transport, ftype);
            iprot.read_field_end(transport);
        }

        iprot.read_struct_end(transport);
    }
}

struct CalculatorAddArgs {
    num1: i32,
    num2: i32,
}

impl CalculatorAddArgs {

    #[allow(unused_variables)]
    fn write(&self, oprot: &Protocol, transport: & mut Transport) {
        oprot.write_struct_begin(transport, "Calculator_add_args");

        oprot.write_field_begin(transport, "num1", Type::TI32, 1);
        oprot.write_i32(transport, self.num1);
        oprot.write_field_end(transport);

        oprot.write_field_begin(transport, "num2", Type::TI32, 2);
        oprot.write_i32(transport, self.num2);
        oprot.write_field_end(transport);

        oprot.write_field_stop(transport);
        oprot.write_struct_end(transport);
    }
}

struct CalculatorAddResult {
    success: Option<i32>,
}

impl CalculatorAddResult {

    fn read(& mut self, iprot: &Protocol, transport: & mut Transport) {
        let fname = iprot.read_struct_begin(transport);

        loop {
            let (fname, ftype, fid) = iprot.read_field_begin(transport);
            if ftype == Type::TStop {
                break;
            }
            match (fid, ftype) {
                (0, Type::TI32) => {
                    self.success = Some(iprot.read_i32(transport));
                }
                _ => {
                    iprot.skip(transport, ftype);
                }
            }
            iprot.read_field_end(transport);
        }
        iprot.read_struct_end(transport);
    }
}

struct CalculatorClient<T: Transport, P: Protocol> {
    transport: T,
    protocol: P,
}

impl <T: Transport, P: Protocol> CalculatorClient<T, P> {

    fn ping(& mut self) {
        self.send_ping();
        self.receive_ping();
    }

    fn send_ping(& mut self) {
        let cseqid: i32 = 0;
        self.protocol.write_message_begin(& mut self.transport, "ping", MessageType::MtCall, cseqid);

        let args = CalculatorPingArgs;
        args.write(&self.protocol, & mut self.transport);

        self.protocol.write_message_end(& mut self.transport);

        //self.transport.write_end();
        self.transport.flush();
    }

    fn receive_ping(& mut self) {
      let (fname, mtype, rseqid) = self.protocol.read_message_begin(& mut self.transport);
      match mtype {
        MessageType::MtException => {
            // TODO
            //let x = ApplicationException;
            //x.read(& mut self protocol)
            //self.protocol.read_message_end();
            //transport.read_end();
            //throw x
        }
        MessageType::MtReply => {
            match fname.as_slice() {
                "ping" => {
                    let result = CalculatorPingResult;
                    result.read(&self.protocol, & mut self.transport);
                }
                _ => {
                    self.protocol.skip(& mut self.transport, Type::TStruct);
                }
            }
            self.protocol.read_message_end(& mut self.transport);
            //self.transport.read_end();
        }
        _ => {
            self.protocol.skip(& mut self.transport, Type::TStruct);
            self.protocol.read_message_end(& mut self.transport);
            //self.transport.read_end();
        }
      }
    }

    fn add(& mut self, num1: i32, num2: i32) -> Option<i32> {
        self.send_add(num1, num2);
        self.receive_add()
    }

    fn send_add(& mut self, num1: i32, num2: i32) {
        let cseqid: i32 = 0;
        self.protocol.write_message_begin(& mut self.transport, "add", MessageType::MtCall, cseqid);

        let args = CalculatorAddArgs { num1: num1, num2: num2 };
        args.write(&self.protocol, & mut self.transport);

        self.protocol.write_message_end(& mut self.transport);

        //self.transport.write_end();
        self.transport.flush();
    }

    fn receive_add(& mut self) -> Option<i32> {
      let mut result = CalculatorAddResult { success: None };
      let (fname, mtype, rseqid) = self.protocol.read_message_begin(& mut self.transport);
      match mtype {
        MessageType::MtException => {
            // TODO
            //let x = ApplicationException;
            //x.read(& mut self protocol)
            //self.protocol.read_message_end();
            //transport.read_end();
            //throw x
        }
        MessageType::MtReply => {
            match fname.as_slice() {
                "add" => {
                    result.read(&self.protocol, & mut self.transport);
                }
                _ => {
                    self.protocol.skip(& mut self.transport, Type::TStruct);
                }
            }
            self.protocol.read_message_end(& mut self.transport);
            //self.transport.read_end();
        }
        _ => {
            self.protocol.skip(& mut self.transport, Type::TStruct);
            self.protocol.read_message_end(& mut self.transport);
            //self.transport.read_end();
        }
      }
      result.success
    }
}

pub fn main() {
    let server_address = "127.0.0.1:9090";
    let addr: ip::SocketAddr = std::str::FromStr::from_str(server_address)
        .expect("bad server address");
    let tcp = std::io::TcpStream::connect(addr).unwrap();

    let mut client = CalculatorClient{ protocol: BinaryProtocol, transport: tcp };

    client.ping();
    println!("ping()");

    println!("1 + 1 = {}", client.add(1, 1).unwrap());

    println!("PASS");
}
<commit_msg>Fixed warnings<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#![crate_name="calculator"]
#![crate_type="bin"]

extern crate thrift;

use std::io::net::ip;
use thrift::protocol::{MessageType, Type};
use thrift::transport::Transport;
use thrift::protocol::Protocol;
use thrift::protocol::binary_protocol::BinaryProtocol;

// TODO: shall be generated
//#[allow(dead_code)]
struct CalculatorPingArgs;

impl CalculatorPingArgs {

    #[allow(unused_variables)]
    fn write(&self, oprot: &Protocol, transport: & mut Transport) {
        oprot.write_struct_begin(transport, "Calculator_ping_args");
        oprot.write_field_stop(transport);
        oprot.write_struct_end(transport);
    }
}

struct CalculatorPingResult;

impl CalculatorPingResult {

    fn read(&self, iprot: &Protocol, transport: & mut Transport) {
        iprot.skip(transport, Type::TStruct);
    }
}

struct CalculatorAddArgs {
    num1: i32,
    num2: i32,
}

impl CalculatorAddArgs {

    #[allow(unused_variables)]
    fn write(&self, oprot: &Protocol, transport: & mut Transport) {
        oprot.write_struct_begin(transport, "Calculator_add_args");

        oprot.write_field_begin(transport, "num1", Type::TI32, 1);
        oprot.write_i32(transport, self.num1);
        oprot.write_field_end(transport);

        oprot.write_field_begin(transport, "num2", Type::TI32, 2);
        oprot.write_i32(transport, self.num2);
        oprot.write_field_end(transport);

        oprot.write_field_stop(transport);
        oprot.write_struct_end(transport);
    }
}

struct CalculatorAddResult {
    success: Option<i32>,
}

impl CalculatorAddResult {

    #[allow(unused_variables)]
    fn read(& mut self, iprot: &Protocol, transport: & mut Transport) {
        iprot.read_struct_begin(transport);

        loop {
            let (fname, ftype, fid) = iprot.read_field_begin(transport);
            if ftype == Type::TStop {
                break;
            }
            match (fid, ftype) {
                (0, Type::TI32) => {
                    self.success = Some(iprot.read_i32(transport));
                }
                _ => {
                    iprot.skip(transport, ftype);
                }
            }
            iprot.read_field_end(transport);
        }
        iprot.read_struct_end(transport);
    }
}

struct CalculatorClient<T: Transport, P: Protocol> {
    transport: T,
    protocol: P,
}

impl <T: Transport, P: Protocol> CalculatorClient<T, P> {

    fn ping(& mut self) {
        self.send_ping();
        self.receive_ping();
    }

    #[allow(unused_variables)]
    fn send_ping(& mut self) {
        let cseqid: i32 = 0;
        self.protocol.write_message_begin(& mut self.transport, "ping", MessageType::MtCall, cseqid);

        let args = CalculatorPingArgs;
        args.write(&self.protocol, & mut self.transport);

        self.protocol.write_message_end(& mut self.transport);

        //self.transport.write_end();
        self.transport.flush().unwrap();
    }

    #[allow(unused_variables)]
    fn receive_ping(& mut self) {
      let (fname, mtype, rseqid) = self.protocol.read_message_begin(& mut self.transport);
      match mtype {
        MessageType::MtException => {
            // TODO
            //let x = ApplicationException;
            //x.read(& mut self protocol)
            //self.protocol.read_message_end();
            //transport.read_end();
            //throw x
        }
        MessageType::MtReply => {
            match fname.as_slice() {
                "ping" => {
                    let result = CalculatorPingResult;
                    result.read(&self.protocol, & mut self.transport);
                }
                _ => {
                    self.protocol.skip(& mut self.transport, Type::TStruct);
                }
            }
            self.protocol.read_message_end(& mut self.transport);
            //self.transport.read_end();
        }
        _ => {
            self.protocol.skip(& mut self.transport, Type::TStruct);
            self.protocol.read_message_end(& mut self.transport);
            //self.transport.read_end();
        }
      }
    }

    fn add(& mut self, num1: i32, num2: i32) -> Option<i32> {
        self.send_add(num1, num2);
        self.receive_add()
    }

    fn send_add(& mut self, num1: i32, num2: i32) {
        let cseqid: i32 = 0;
        self.protocol.write_message_begin(& mut self.transport, "add", MessageType::MtCall, cseqid);

        let args = CalculatorAddArgs { num1: num1, num2: num2 };
        args.write(&self.protocol, & mut self.transport);

        self.protocol.write_message_end(& mut self.transport);

        //self.transport.write_end();
        self.transport.flush().unwrap();
    }

    #[allow(unused_variables)]
    fn receive_add(& mut self) -> Option<i32> {
      let mut result = CalculatorAddResult { success: None };
      let (fname, mtype, rseqid) = self.protocol.read_message_begin(& mut self.transport);
      match mtype {
        MessageType::MtException => {
            // TODO
            //let x = ApplicationException;
            //x.read(& mut self protocol)
            //self.protocol.read_message_end();
            //transport.read_end();
            //throw x
        }
        MessageType::MtReply => {
            match fname.as_slice() {
                "add" => {
                    result.read(&self.protocol, & mut self.transport);
                }
                _ => {
                    self.protocol.skip(& mut self.transport, Type::TStruct);
                }
            }
            self.protocol.read_message_end(& mut self.transport);
            //self.transport.read_end();
        }
        _ => {
            self.protocol.skip(& mut self.transport, Type::TStruct);
            self.protocol.read_message_end(& mut self.transport);
            //self.transport.read_end();
        }
      }
      result.success
    }
}

pub fn main() {
    let server_address = "127.0.0.1:9090";
    let addr: ip::SocketAddr = std::str::FromStr::from_str(server_address)
        .expect("bad server address");
    let tcp = std::io::TcpStream::connect(addr).unwrap();

    let mut client = CalculatorClient{ protocol: BinaryProtocol, transport: tcp };

    client.ping();
    println!("ping()");

    println!("1 + 1 = {}", client.add(1, 1).unwrap());

    println!("PASS");
}
<|endoftext|>


<commit_before>//! NVPTX intrinsics (experimental)
//!
//! These intrinsics form the foundation of the CUDA
//! programming model.
//!
//! The reference is the [CUDA C Programming Guide][cuda_c]. Relevant is also
//! the [LLVM NVPTX Backend documentation][llvm_docs].
//!
//! [cuda_c]:
//! http://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
//! [llvm_docs]:
//! https://llvm.org/docs/NVPTXUsage.html

#[allow(improper_ctypes)]
extern "C" {
    #[link_name = "llvm.cuda.syncthreads"]
    fn syncthreads() -> ();
    #[link_name = "llvm.nvvm.read.ptx.sreg.ntid.x"]
    fn block_dim_x() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ntid.y"]
    fn block_dim_y() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ntid.z"]
    fn block_dim_z() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ctaid.x"]
    fn block_idx_x() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ctaid.y"]
    fn block_idx_y() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ctaid.z"]
    fn block_idx_z() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.nctaid.x"]
    fn grid_dim_x() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.nctaid.y"]
    fn grid_dim_y() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.nctaid.z"]
    fn grid_dim_z() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.tid.x"]
    fn thread_idx_x() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.tid.y"]
    fn thread_idx_y() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.tid.z"]
    fn thread_idx_z() -> i32;
}

/// Synchronizes all threads in the block.
#[inline]
pub unsafe fn _syncthreads() -> () {
    syncthreads()
}

/// x-th thread-block dimension.
#[inline]
pub unsafe fn _block_dim_x() -> i32 {
    block_dim_x()
}

/// y-th thread-block dimension.
#[inline]
pub unsafe fn _block_dim_y() -> i32 {
    block_dim_y()
}

/// z-th thread-block dimension.
#[inline]
pub unsafe fn _block_dim_z() -> i32 {
    block_dim_z()
}

/// x-th thread-block index.
#[inline]
pub unsafe fn _block_idx_x() -> i32 {
    block_idx_x()
}

/// y-th thread-block index.
#[inline]
pub unsafe fn _block_idx_y() -> i32 {
    block_idx_y()
}

/// z-th thread-block index.
#[inline]
pub unsafe fn _block_idx_z() -> i32 {
    block_idx_z()
}

/// x-th block-grid dimension.
#[inline]
pub unsafe fn _grid_dim_x() -> i32 {
    grid_dim_x()
}

/// y-th block-grid dimension.
#[inline]
pub unsafe fn _grid_dim_y() -> i32 {
    grid_dim_y()
}

/// z-th block-grid dimension.
#[inline]
pub unsafe fn _grid_dim_z() -> i32 {
    grid_dim_z()
}

/// x-th thread index.
#[inline]
pub unsafe fn _thread_idx_x() -> i32 {
    thread_idx_x()
}

/// y-th thread index.
#[inline]
pub unsafe fn _thread_idx_y() -> i32 {
    thread_idx_y()
}

/// z-th thread index.
#[inline]
pub unsafe fn _thread_idx_z() -> i32 {
    thread_idx_z()
}
<commit_msg>Rename `llvm.cuda.syncthreads` -> `llvm.nvvm.barrier0` to keep up with (#589)<commit_after>//! NVPTX intrinsics (experimental)
//!
//! These intrinsics form the foundation of the CUDA
//! programming model.
//!
//! The reference is the [CUDA C Programming Guide][cuda_c]. Relevant is also
//! the [LLVM NVPTX Backend documentation][llvm_docs].
//!
//! [cuda_c]:
//! http://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html
//! [llvm_docs]:
//! https://llvm.org/docs/NVPTXUsage.html

#[allow(improper_ctypes)]
extern "C" {
    #[link_name = "llvm.nvvm.barrier0"]
    fn syncthreads() -> ();
    #[link_name = "llvm.nvvm.read.ptx.sreg.ntid.x"]
    fn block_dim_x() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ntid.y"]
    fn block_dim_y() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ntid.z"]
    fn block_dim_z() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ctaid.x"]
    fn block_idx_x() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ctaid.y"]
    fn block_idx_y() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.ctaid.z"]
    fn block_idx_z() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.nctaid.x"]
    fn grid_dim_x() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.nctaid.y"]
    fn grid_dim_y() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.nctaid.z"]
    fn grid_dim_z() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.tid.x"]
    fn thread_idx_x() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.tid.y"]
    fn thread_idx_y() -> i32;
    #[link_name = "llvm.nvvm.read.ptx.sreg.tid.z"]
    fn thread_idx_z() -> i32;
}

/// Synchronizes all threads in the block.
#[inline]
pub unsafe fn _syncthreads() -> () {
    syncthreads()
}

/// x-th thread-block dimension.
#[inline]
pub unsafe fn _block_dim_x() -> i32 {
    block_dim_x()
}

/// y-th thread-block dimension.
#[inline]
pub unsafe fn _block_dim_y() -> i32 {
    block_dim_y()
}

/// z-th thread-block dimension.
#[inline]
pub unsafe fn _block_dim_z() -> i32 {
    block_dim_z()
}

/// x-th thread-block index.
#[inline]
pub unsafe fn _block_idx_x() -> i32 {
    block_idx_x()
}

/// y-th thread-block index.
#[inline]
pub unsafe fn _block_idx_y() -> i32 {
    block_idx_y()
}

/// z-th thread-block index.
#[inline]
pub unsafe fn _block_idx_z() -> i32 {
    block_idx_z()
}

/// x-th block-grid dimension.
#[inline]
pub unsafe fn _grid_dim_x() -> i32 {
    grid_dim_x()
}

/// y-th block-grid dimension.
#[inline]
pub unsafe fn _grid_dim_y() -> i32 {
    grid_dim_y()
}

/// z-th block-grid dimension.
#[inline]
pub unsafe fn _grid_dim_z() -> i32 {
    grid_dim_z()
}

/// x-th thread index.
#[inline]
pub unsafe fn _thread_idx_x() -> i32 {
    thread_idx_x()
}

/// y-th thread index.
#[inline]
pub unsafe fn _thread_idx_y() -> i32 {
    thread_idx_y()
}

/// z-th thread index.
#[inline]
pub unsafe fn _thread_idx_z() -> i32 {
    thread_idx_z()
}
<|endoftext|>


<commit_before>#[crate_id(name="base64", vers="1.0.0", author="Jordy Dickinson")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Jordy Dickinson <jordy.dickinson@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE file
 * that was distributed with this source code.
 */

extern mod extra;

use std::char;
use std::io::{File, stdin, stdout};
use std::os;
use std::str;

use extra::getopts::groups::{
    getopts,
    optflag,
    optopt,
    usage
};
use extra::base64;
use extra::base64::{FromBase64, ToBase64};

fn main() {
    let args = ~os::args();
    let opts = ~[
        optflag("d", "decode", "decode data"),
        optflag("i", "ignore-garbage", "when decoding, ignore non-alphabetic characters"),
        optopt("w", "wrap",
            "wrap encoded lines after COLS character (default 76, 0 to disable wrapping)", "COLS"
        ),
        optflag("h", "help", "display this help text and exit"),
        optflag("V", "version", "output version information and exit")
    ];
    let matches = match getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(e) => {
            error!("error: {:s}", e.to_err_msg());
            fail!()
        }
    };

    let progname = args[0].clone();
    let usage = usage("Base64 encode or decode FILE, or standard input, to standard output.", opts);
    let mode = if matches.opt_present("help") {
        Help
    } else if matches.opt_present("version") {
        Version
    } else if matches.opt_present("decode") {
        Decode
    } else {
        Encode
    };
    let ignore_garbage = matches.opt_present("ignore-garbage");
    let line_wrap = match matches.opt_str("wrap") {
        Some(s) => match from_str(s) {
            Some(s) => s,
            None => {
                error!("error: {:s}", "Argument to option 'wrap' improperly formatted.");
                fail!()
            }
        },
        None => 76
    };
    let mut input = if matches.free.is_empty() || matches.free[0] == ~"-" {
        ~stdin() as ~Reader
    } else {
        let path = Path::new(matches.free[0]);
        ~File::open(&path) as ~Reader
    };

    match mode {
        Decode  => decode(input, ignore_garbage),
        Encode  => encode(input, line_wrap),
        Help    => help(progname, usage),
        Version => version()
    }
}

fn decode(input: &mut Reader, ignore_garbage: bool) {
    let mut to_decode = str::from_utf8_owned(input.read_to_end());

    to_decode = to_decode.replace("\n", "");

    if ignore_garbage {
        let standard_chars =
            bytes!("ABCDEFGHIJKLMNOPQRSTUVWXYZ",
            "abcdefghijklmnopqrstuvwxyz",
            "0123456789+/").map(|b| char::from_u32(*b as u32).unwrap());

        to_decode = to_decode
            .trim_chars(&|c| !standard_chars.contains(&c))
            .to_owned();
    }

    match to_decode.from_base64() {
        Ok(bytes) => {
            let mut out = stdout();

            out.write(bytes);
            out.flush();
        }
        Err(s) => {
            error!("error: {:s}", s);
            fail!()
        }
    }
}

fn encode(input: &mut Reader, line_wrap: uint) {
    let b64_conf = base64::Config {
        char_set: base64::Standard,
        pad: true,
        line_length: match line_wrap {
            0 => None,
            _ => Some(line_wrap)
        }
    };
    let to_encode = input.read_to_end();
    let mut encoded = to_encode.to_base64(b64_conf);

    // To my knowledge, RFC 3548 does not specify which line endings to use. It
    // seems that rust's base64 algorithm uses CRLF as prescribed by RFC 2045.
    // However, since GNU base64 outputs only LF (presumably because that is
    // the standard UNIX line ending), we strip CRs from the output to maintain
    // compatibility.
    encoded = encoded.replace("\r", "");

    println(encoded);
}

fn help(progname: &str, usage: &str) {
    println!("Usage: {:s} [OPTION]... [FILE]", progname);
    println("");
    println(usage);

    let msg = ~"With no FILE, or when FILE is -, read standard input.\n\n\
        The data are encoded as described for the base64 alphabet in RFC \
        3548. When\ndecoding, the input may contain newlines in addition \
        to the bytes of the formal\nbase64 alphabet. Use --ignore-garbage \
        to attempt to recover from any other\nnon-alphabet bytes in the \
        encoded stream.";

    println(msg);
}

fn version() {
    println("base64 1.0.0");
}

enum Mode {
    Decode,
    Encode,
    Help,
    Version
}

<commit_msg>trying to update base64 for print changes<commit_after>#[crate_id(name="base64", vers="1.0.0", author="Jordy Dickinson")];

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Jordy Dickinson <jordy.dickinson@gmail.com>
 *
 * For the full copyright and license information, please view the LICENSE file
 * that was distributed with this source code.
 */

extern mod extra;

use std::char;
use std::io::{println, File, stdin, stdout};
use std::os;
use std::str;

use extra::getopts::groups::{
    getopts,
    optflag,
    optopt,
    usage
};
use extra::base64;
use extra::base64::{FromBase64, ToBase64};

fn main() {
    let args = ~os::args();
    let opts = ~[
        optflag("d", "decode", "decode data"),
        optflag("i", "ignore-garbage", "when decoding, ignore non-alphabetic characters"),
        optopt("w", "wrap",
            "wrap encoded lines after COLS character (default 76, 0 to disable wrapping)", "COLS"
        ),
        optflag("h", "help", "display this help text and exit"),
        optflag("V", "version", "output version information and exit")
    ];
    let matches = match getopts(args.tail(), opts) {
        Ok(m) => m,
        Err(e) => {
            error!("error: {:s}", e.to_err_msg());
            fail!()
        }
    };

    let progname = args[0].clone();
    let usage = usage("Base64 encode or decode FILE, or standard input, to standard output.", opts);
    let mode = if matches.opt_present("help") {
        Help
    } else if matches.opt_present("version") {
        Version
    } else if matches.opt_present("decode") {
        Decode
    } else {
        Encode
    };
    let ignore_garbage = matches.opt_present("ignore-garbage");
    let line_wrap = match matches.opt_str("wrap") {
        Some(s) => match from_str(s) {
            Some(s) => s,
            None => {
                error!("error: {:s}", "Argument to option 'wrap' improperly formatted.");
                fail!()
            }
        },
        None => 76
    };
    let mut input = if matches.free.is_empty() || matches.free[0] == ~"-" {
        ~stdin() as ~Reader
    } else {
        let path = Path::new(matches.free[0]);
        ~File::open(&path) as ~Reader
    };

    match mode {
        Decode  => decode(input, ignore_garbage),
        Encode  => encode(input, line_wrap),
        Help    => help(progname, usage),
        Version => version()
    }
}

fn decode(input: &mut Reader, ignore_garbage: bool) {
    let mut to_decode = str::from_utf8_owned(input.read_to_end());

    to_decode = to_decode.replace("\n", "");

    if ignore_garbage {
        let standard_chars =
            bytes!("ABCDEFGHIJKLMNOPQRSTUVWXYZ",
            "abcdefghijklmnopqrstuvwxyz",
            "0123456789+/").map(|b| char::from_u32(*b as u32).unwrap());

        to_decode = to_decode
            .trim_chars(&|c| !standard_chars.contains(&c))
            .to_owned();
    }

    match to_decode.from_base64() {
        Ok(bytes) => {
            let mut out = stdout();

            out.write(bytes);
            out.flush();
        }
        Err(s) => {
            error!("error: {:s}", s);
            fail!()
        }
    }
}

fn encode(input: &mut Reader, line_wrap: uint) {
    let b64_conf = base64::Config {
        char_set: base64::Standard,
        pad: true,
        line_length: match line_wrap {
            0 => None,
            _ => Some(line_wrap)
        }
    };
    let to_encode = input.read_to_end();
    let mut encoded = to_encode.to_base64(b64_conf);

    // To my knowledge, RFC 3548 does not specify which line endings to use. It
    // seems that rust's base64 algorithm uses CRLF as prescribed by RFC 2045.
    // However, since GNU base64 outputs only LF (presumably because that is
    // the standard UNIX line ending), we strip CRs from the output to maintain
    // compatibility.
    encoded = encoded.replace("\r", "");

    println(encoded);
}

fn help(progname: &str, usage: &str) {
    println!("Usage: {:s} [OPTION]... [FILE]", progname);
    println!("");
    println(usage);

    let msg = ~"With no FILE, or when FILE is -, read standard input.\n\n\
        The data are encoded as described for the base64 alphabet in RFC \
        3548. When\ndecoding, the input may contain newlines in addition \
        to the bytes of the formal\nbase64 alphabet. Use --ignore-garbage \
        to attempt to recover from any other\nnon-alphabet bytes in the \
        encoded stream.";

    println(msg);
}

fn version() {
    println!("base64 1.0.0");
}

enum Mode {
    Decode,
    Encode,
    Help,
    Version
}

<|endoftext|>


<commit_before>extern crate rustty;

use rustty::{
    Terminal,
    Event,
};
use rustty::ui::{
    Painter,
    Dialog,
    ButtonResult,
    Alignable,
    HorizontalAlign,
    VerticalAlign,
    StdButton,
    Widget,
};

fn create_maindlg() -> Dialog {
    let mut maindlg = Dialog::new(60, 10);
    maindlg.window_mut().draw_box();

    let mut b1 = StdButton::new("Quit", 'q', ButtonResult::Ok);
    b1.pack(&maindlg, HorizontalAlign::Left, VerticalAlign::Top, 6);
    let mut b2 = StdButton::new("Foo", 'f', ButtonResult::Custom(1));
    b2.pack(&maindlg, HorizontalAlign::Middle, VerticalAlign::Top, 6);
    let mut b3 = StdButton::new("Bar", 'b', ButtonResult::Custom(2));
    b3.pack(&maindlg, HorizontalAlign::Right, VerticalAlign::Top, 6);

    maindlg.add_button(b1);
    maindlg.add_button(b2);
    maindlg.add_button(b3);
    maindlg
}

fn main() {
    let mut term = Terminal::new().unwrap();
    let mut maindlg = create_maindlg();
    maindlg.window_mut().align(&term, HorizontalAlign::Middle, VerticalAlign::Middle, 0);
    'main: loop {
        while let Some(Event::Key(ch)) = term.get_event(0).unwrap() {
            match maindlg.result_for_key(ch) {
                Some(ButtonResult::Ok)          => break 'main,
                Some(ButtonResult::Custom(i))   => {
                    let msg = if i == 1 { "Foo!" } else { "Bar!" };
                    let w = maindlg.window_mut();
                    let x = w.halign_line(msg, HorizontalAlign::Middle, 1);
                    let y = w.valign_line(msg, VerticalAlign::Middle, 1);
                    w.printline(x, y, msg);
                },
                _ => {},
            }
        }

        maindlg.window().draw_into(&mut term);
        term.swap_buffers().unwrap();
    }
}

<commit_msg>fix include errors in ui.rs<commit_after>extern crate rustty;

use rustty::{
    Terminal,
    Event,
};

use rustty::ui::core::{
    Painter,
    Base,
    Widget,
    ButtonResult,
    Alignable,
    HorizontalAlign,
    VerticalAlign,
};

use rustty::ui::{
    StdButton,
    Dialog
};

fn create_maindlg() -> Dialog {
    let mut maindlg = Dialog::new(60, 10);
    maindlg.window_mut().draw_box();

    let mut b1 = StdButton::new("Quit", 'q', ButtonResult::Ok);
    b1.pack(&maindlg, HorizontalAlign::Left, VerticalAlign::Top, 6);
    let mut b2 = StdButton::new("Foo", 'f', ButtonResult::Custom(1));
    b2.pack(&maindlg, HorizontalAlign::Middle, VerticalAlign::Top, 6);
    let mut b3 = StdButton::new("Bar", 'b', ButtonResult::Custom(2));
    b3.pack(&maindlg, HorizontalAlign::Right, VerticalAlign::Top, 6);

    maindlg.add_button(b1);
    maindlg.add_button(b2);
    maindlg.add_button(b3);
    maindlg
}

fn main() {
    let mut term = Terminal::new().unwrap();
    let mut maindlg = create_maindlg();
    maindlg.window_mut().align(&term, HorizontalAlign::Middle, VerticalAlign::Middle, 0);
    'main: loop {
        while let Some(Event::Key(ch)) = term.get_event(0).unwrap() {
            match maindlg.result_for_key(ch) {
                Some(ButtonResult::Ok)          => break 'main,
                Some(ButtonResult::Custom(i))   => {
                    let msg = if i == 1 { "Foo!" } else { "Bar!" };
                    let w = maindlg.window_mut();
                    let x = w.halign_line(msg, HorizontalAlign::Middle, 1);
                    let y = w.valign_line(msg, VerticalAlign::Middle, 1);
                    w.printline(x, y, msg);
                },
                _ => {},
            }
        }

        maindlg.window().draw_into(&mut term);
        term.swap_buffers().unwrap();
    }
}

<|endoftext|>


<commit_before>//! Utilities for working with the WTF-8 encoding.
//!
//! WTF-8 is a hack that allows to represent potentially ill-formed UTF-16 byte
//! strings in a way somewhat compatible with UTF-8. Such bugs UTF-16 strings
//! can be found for example in Windows paths for legacy reasons.
//!
//! See the official [WTF-8][wtf8] specification for more information.
//!
//! [wtf8]: https://simonsapin.github.io/wtf-8

/// Decodes the given potentially ill-formed UTF-16 into a WTF-8 byte sequence.
///
/// If the given byte sequence is a well-formed UTF-16 string, then the result
/// is guaranteed to be a valid UTF-8string.
pub fn decode_ill_formed_utf16(units: impl Iterator<Item=u16>) -> Vec<u8> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let unit = match iter.next() {
            Some(unit) => unit,
            None => return res,
        };

        // Potentially ill-formed UTF-16 to code point conversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-ill-formed-utf-16
        let mut point = unit as u32;
        if is_lead_surrogate(&unit) {
            if let Some(next) = iter.next_if(is_trail_surrogate) {
                let lead = (unit as u32 - 0xD800) << 10;
                let trail = next as u32 - 0xDC00;
                point = 0x10000 + lead + trail;
            }
        }

        // Code point to WTF-8 encoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-wtf-8
        match point {
            0x0000..=0x007F => {
                res.push(point as u8);
            }
            0x0080..=0x07FF => {
                res.push(0xC0 | (point >> 6) as u8);
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x0800..=0xFFFF => {
                res.push(0xE0 | (point >> 12) as u8);
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x10000..=0x10FFFF => {
                res.push(0xF0 | (point >> 18) as u8);
                res.push(0x80 | ((point >> 12) as u8 & 0x3F));
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            _ => panic!(), // Not possible by construction.
        }
    }
}

/// Encodes the given WTF-8 byte sequence into potentially ill-formed UTF-16.
///
/// If the given byte sequence is a valid UTF-8 string, then the result is
/// guaranteed to be a well-formed UTF-16 string.
pub fn encode_ill_formed_utf16(units: impl Iterator<Item=u8>) -> Vec<u16> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let byte1 = match iter.next() {
            Some(byte1) => byte1,
            None => return res,
        };

        // TODO: Add support for proper error handling.

        // WTF-8 to a code point decoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-wtf-8
        let mut point = 0;
        match byte1 {
            0x00..=0x7F => {
                point += byte1 as u32;
            }
            0xC2..=0xDF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x1F) as u32) << 6;
                point += ((byte2 & 0x3F) as u32) << 0;
            }
            0xE0..=0xEF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x0F) as u32) << 12;
                point += ((byte2 & 0x3F) as u32) << 6;
                point += ((byte3 & 0x3F) as u32) << 0;
            }
            0xF0..=0xF4 => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                let byte4 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x07) as u32) << 18;
                point += ((byte2 & 0x3F) as u32) << 12;
                point += ((byte3 & 0x3F) as u32) << 6;
                point += ((byte4 & 0x3F) as u32) << 0;
            }
            _ => todo!(),
        }

        // Code point to potentially ill-formed UTF-16 coversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-ill-formed-utf-16
        if is_supplementary(&point) {
            res.push(((point - 0x10000) >> 10) as u16 + 0xD800);
            res.push(((point - 0x10000) as u16 & 0x3FF) + 0xDC00);
        } else {
            res.push(point as u16);
        }
    }
}

/// Determines whether the given UTF-16 code unit is a lead surrogate.
#[inline]
fn is_lead_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xD800..=0xDBFF)
}

/// Determines whether the given UTF-16 code unit is a trail surrogate.
#[inline]
fn is_trail_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xDC00..=0xDFFF)
}

/// Determines whether the given code point is in the supplementary.
#[inline]
fn is_supplementary(point: &u32) -> bool {
    matches!(point, 0x10000..=0x10FFFF)
}

#[cfg(test)]
mod tests {

    use crate::wtf8;

    #[test]
    fn decode_empty() {
        assert_eq!(wtf8::decode_ill_formed_utf16(std::iter::empty()), b"");
    }

    #[test]
    fn decode_small_bytes() {
        let bytes_utf16 = vec![0x04, 0x08, 0x16, 0x23, 0x42];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\x04\x08\x16\x23\x42");
    }

    #[test]
    fn decode_bmp_bytes() {
        let bytes_utf16 = vec![0xABCD, 0xBEEF, 0xFFFF];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xEA\xAF\x8D\xEB\xBB\xAF\xEF\xBF\xBF");
    }

    #[test]
    fn decode_sup_bytes() {
        let bytes_utf16 = vec![0xD800, 0xDC00, 0xDBFF, 0xDFFF];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xF0\x90\x80\x80\xF4\x8F\xBF\xBF");
    }

    #[test]
    fn decode_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = wtf8::decode_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn decode_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = wtf8::decode_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn encode_empty() {
        assert_eq!(wtf8::encode_ill_formed_utf16(std::iter::empty()), vec![]);
    }

    #[test]
    fn encode_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = wtf8::encode_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }

    #[test]
    fn encode_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = wtf8::encode_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }
}
<commit_msg>Use less confusing names in the WTF-8 module.<commit_after>//! Utilities for working with the WTF-8 encoding.
//!
//! WTF-8 is a hack that allows to represent potentially ill-formed UTF-16 byte
//! strings in a way somewhat compatible with UTF-8. Such bugs UTF-16 strings
//! can be found for example in Windows paths for legacy reasons.
//!
//! See the official [WTF-8][wtf8] specification for more information.
//!
//! [wtf8]: https://simonsapin.github.io/wtf-8

/// Converts the given potentially ill-formed UTF-16 into a WTF-8 byte sequence.
///
/// If the given byte sequence is a well-formed UTF-16 string, then the result
/// is guaranteed to be a valid UTF-8string.
pub fn from_ill_formed_utf16(units: impl Iterator<Item=u16>) -> Vec<u8> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let unit = match iter.next() {
            Some(unit) => unit,
            None => return res,
        };

        // Potentially ill-formed UTF-16 to code point conversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-ill-formed-utf-16
        let mut point = unit as u32;
        if is_lead_surrogate(&unit) {
            if let Some(next) = iter.next_if(is_trail_surrogate) {
                let lead = (unit as u32 - 0xD800) << 10;
                let trail = next as u32 - 0xDC00;
                point = 0x10000 + lead + trail;
            }
        }

        // Code point to WTF-8 encoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-wtf-8
        match point {
            0x0000..=0x007F => {
                res.push(point as u8);
            }
            0x0080..=0x07FF => {
                res.push(0xC0 | (point >> 6) as u8);
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x0800..=0xFFFF => {
                res.push(0xE0 | (point >> 12) as u8);
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x10000..=0x10FFFF => {
                res.push(0xF0 | (point >> 18) as u8);
                res.push(0x80 | ((point >> 12) as u8 & 0x3F));
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            _ => panic!(), // Not possible by construction.
        }
    }
}

/// Converts the given WTF-8 byte sequence into potentially ill-formed UTF-16.
///
/// If the given byte sequence is a valid UTF-8 string, then the result is
/// guaranteed to be a well-formed UTF-16 string.
pub fn into_ill_formed_utf16(units: impl Iterator<Item=u8>) -> Vec<u16> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let byte1 = match iter.next() {
            Some(byte1) => byte1,
            None => return res,
        };

        // TODO: Add support for proper error handling.

        // WTF-8 to a code point decoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-wtf-8
        let mut point = 0;
        match byte1 {
            0x00..=0x7F => {
                point += byte1 as u32;
            }
            0xC2..=0xDF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x1F) as u32) << 6;
                point += ((byte2 & 0x3F) as u32) << 0;
            }
            0xE0..=0xEF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x0F) as u32) << 12;
                point += ((byte2 & 0x3F) as u32) << 6;
                point += ((byte3 & 0x3F) as u32) << 0;
            }
            0xF0..=0xF4 => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                let byte4 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x07) as u32) << 18;
                point += ((byte2 & 0x3F) as u32) << 12;
                point += ((byte3 & 0x3F) as u32) << 6;
                point += ((byte4 & 0x3F) as u32) << 0;
            }
            _ => todo!(),
        }

        // Code point to potentially ill-formed UTF-16 coversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-ill-formed-utf-16
        if is_supplementary(&point) {
            res.push(((point - 0x10000) >> 10) as u16 + 0xD800);
            res.push(((point - 0x10000) as u16 & 0x3FF) + 0xDC00);
        } else {
            res.push(point as u16);
        }
    }
}

/// Determines whether the given UTF-16 code unit is a lead surrogate.
#[inline]
fn is_lead_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xD800..=0xDBFF)
}

/// Determines whether the given UTF-16 code unit is a trail surrogate.
#[inline]
fn is_trail_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xDC00..=0xDFFF)
}

/// Determines whether the given code point is in the supplementary.
#[inline]
fn is_supplementary(point: &u32) -> bool {
    matches!(point, 0x10000..=0x10FFFF)
}

#[cfg(test)]
mod tests {

    use crate::wtf8;

    #[test]
    fn from_empty() {
        assert_eq!(wtf8::from_ill_formed_utf16(std::iter::empty()), b"");
    }

    #[test]
    fn from_small_bytes() {
        let bytes_utf16 = vec![0x04, 0x08, 0x16, 0x23, 0x42];
        let bytes_wtf8 = wtf8::from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\x04\x08\x16\x23\x42");
    }

    #[test]
    fn from_bmp_bytes() {
        let bytes_utf16 = vec![0xABCD, 0xBEEF, 0xFFFF];
        let bytes_wtf8 = wtf8::from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xEA\xAF\x8D\xEB\xBB\xAF\xEF\xBF\xBF");
    }

    #[test]
    fn from_sup_bytes() {
        let bytes_utf16 = vec![0xD800, 0xDC00, 0xDBFF, 0xDFFF];
        let bytes_wtf8 = wtf8::from_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xF0\x90\x80\x80\xF4\x8F\xBF\xBF");
    }

    #[test]
    fn from_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = wtf8::from_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn from_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = wtf8::from_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn into_empty() {
        assert_eq!(wtf8::into_ill_formed_utf16(std::iter::empty()), vec![]);
    }

    #[test]
    fn into_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = wtf8::into_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }

    #[test]
    fn into_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = wtf8::into_ill_formed_utf16(string.bytes());
        let string_utf16 = string.encode_utf16().collect::<Vec<_>>();
        assert_eq!(string_wtf8, string_utf16);
    }
}
<|endoftext|>


<commit_before>//! The decompression algorithm.

use std::io;
use byteorder::{LittleEndian, ByteOrder};

/// A LZ4 decoder.
///
/// This will decode in accordance to the LZ4 format. It represents a particular state of the
/// decompressor.
struct Decoder<'a> {
    /// The compressed input.
    input: &'a [u8],
    /// The decompressed output.
    output: &'a mut Vec<u8>,
    /// The current block's "token".
    ///
    /// This token contains to 4-bit "fields", a higher and a lower, representing the literals'
    /// length and the back reference's length, respectively. LSIC is used if either are their
    /// maximal values.
    token: u8,
}

impl<'a> Decoder<'a> {
    /// Internal (partial) function for `take`.
    #[inline]
    fn take_imp(input: &mut &'a [u8], n: usize) -> io::Result<&'a [u8]> {
        // Check if we have enough bytes left.
        if input.len() < n {
            // No extra bytes. This is clearly not expected, so we return an error.
            Err(io::Error::new(io::ErrorKind::InvalidData, "Expected another byte, found none."))
        } else {
            // Take the first n bytes.
            let res = Ok(&input[..n]);
            // Shift the stream to left, so that it is no longer the first byte.
            *input = &input[n..];

            // Return the former first byte.
            res
        }
    }

    /// Pop n bytes from the start of the input stream.
    fn take(&mut self, n: usize) -> io::Result<&[u8]> {
        Self::take_imp(&mut self.input, n)
    }

    /// Write a buffer to the output stream.
    ///
    /// The reason this doesn't take `&mut self` is that we need partial borrowing due to the rules
    /// of the borrow checker. For this reason, we instead take some number of segregated
    /// references so we can read and write them independently.
    fn output(output: &mut Vec<u8>, buf: &[u8]) {
        // We use simple memcpy to extend the vector.
        output.extend_from_slice(&buf[..buf.len()]);
    }

    /// Write an already decompressed match to the output stream.
    ///
    /// This is used for the essential part of the algorithm: deduplication. We start at some
    /// position `start` and then keep pushing the following element until we've added
    /// `match_length` elements.
    fn duplicate(&mut self, start: usize, match_length: usize) {
        // We cannot simply use memcpy or `extend_from_slice`, because these do not allow
        // self-referential copies: http://ticki.github.io/img/lz4_runs_encoding_diagram.svg
        for i in start..start + match_length {
            let b = self.output[i];
            self.output.push(b);
        }
    }

    /// Read an integer LSIC (linear small integer code) encoded.
    ///
    /// In LZ4, we encode small integers in a way that we can have an arbitrary number of bytes. In
    /// particular, we add the bytes repeatedly until we hit a non-0xFF byte. When we do, we add
    /// this byte to our sum and terminate the loop.
    ///
    /// # Example
    ///
    /// ```notest
    ///     255, 255, 255, 4, 2, 3, 4, 6, 7
    /// ```
    ///
    /// is encoded to _255 + 255 + 255 + 4 = 769_. The bytes after the first 4 is ignored, because
    /// 4 is the first non-0xFF byte.
    #[inline]
    fn read_integer(&mut self) -> io::Result<usize> {
        // We start at zero and count upwards.
        let mut n = 0;
        // If this byte takes value 255 (the maximum value it can take), another byte is read
        // and added to the sum. This repeats until a byte lower than 255 is read.
        while {
            // We add the next byte until we get a byte which we add to the counting variable.
            let extra = self.take(1)?[0];
            n += extra as usize;

            // We continue if we got 255.
            extra == 0xFF
        } {}

        Ok(n)
    }

    /// Read a little-endian 16-bit integer from the input stream.
    #[inline]
    fn read_u16(&mut self) -> io::Result<u16> {
        // We use byteorder to read an u16 in little endian.
        Ok(LittleEndian::read_u16(self.take(2)?))
    }

    /// Read the literals section of a block.
    ///
    /// The literals section encodes some bytes which are to be copied to the output without any
    /// modification.
    ///
    /// It consists of two parts:
    ///
    /// 1. An LSIC integer extension to the literals length as defined by the first part of the
    ///    token, if it takes the highest value (15).
    /// 2. The literals themself.
    fn read_literal_section(&mut self) -> io::Result<()> {
        // The higher token is the literals part of the token. It takes a value from 0 to 15.
        let mut literal = (self.token >> 4) as usize;
        // If the initial value is 15, it is indicated that another byte will be read and added to
        // it.
        if literal == 15 {
            // The literal length took the maximal value, indicating that there is more than 15
            // literal bytes. We read the extra integer.
            literal += self.read_integer()?;
        }

        // Now we know the literal length. The number will be used to indicate how long the
        // following literal copied to the output buffer is.

        // Read the literals segment and output them without processing.
        Self::output(&mut self.output, Self::take_imp(&mut self.input, literal)?);

        Ok(())
    }

    /// Read the duplicates section of the block.
    ///
    /// The duplicates section serves to reference an already decoded segment. This consists of two
    /// parts:
    ///
    /// 1. A 16-bit little-endian integer defining the "offset", i.e. how long back we need to go
    ///    in the decoded buffer and copy.
    /// 2. An LSIC integer extension to the duplicate length as defined by the first part of the
    ///    token, if it takes the highest value (15).
    fn read_duplicate_section(&mut self) -> io::Result<()> {
        // Now, we will obtain the offset which we will use to copy from the output. It is an
        // 16-bit integer.
        let offset = self.read_u16()?;

        // Obtain the initial match length. The match length is the length of the duplicate segment
        // which will later be copied from data previously decompressed into the output buffer. The
        // initial length is derived from the second part of the token (the lower nibble), we read
        // earlier. Since having a match length of less than 4 would mean negative compression
        // ratio, we start at 4.
        let mut match_length = (4 + self.token & 0xF) as usize;

        // The intial match length can maximally be 19. As with the literal length, this indicates
        // that there are more bytes to read.
        if match_length == 4 + 15 {
            // The match length took the maximal value, indicating that there is more bytes. We
            // read the extra integer.
            match_length += self.read_integer()?;
        }

        // We now copy from the already decompressed buffer. This allows us for storing duplicates
        // by simply referencing the other location.

        // Calculate the start of this duplicate segment. We use wrapping subtraction to avoid
        // overflow checks, which we will catch later.
        let start = self.output.len().wrapping_sub(offset as usize);

        // We'll do a bound check to avoid panicking.
        if start < self.output.len() {
            // Write the duplicate segment to the output buffer.
            self.duplicate(start, match_length);

            Ok(())
        } else {
            Err(io::Error::new(io::ErrorKind::InvalidData, "The offset to copy is not contained in \
                               the decompressed buffer."))
        }
    }

    /// Complete the decompression by reading all the blocks.
    ///
    /// # Decompressing a block
    ///
    /// Blocks consists of:
    ///  - A 1 byte token
    ///      * A 4 bit integer $t_1$.
    ///      * A 4 bit integer $t_2$.
    ///  - A $n$ byte sequence of 0xFF bytes (if $t_1 \neq 15$, then $n = 0$).
    ///  - $x$ non-0xFF 8-bit integers, L (if $t_1 = 15$, $x = 1$, else $x = 0$).
    ///  - $t_1 + 15n + L$ bytes of uncompressed data (literals).
    ///  - 16-bits offset (little endian), $a$.
    ///  - A $m$ byte sequence of 0xFF bytes (if $t_2 \neq 15$, then $m = 0$).
    ///  - $y$ non-0xFF 8-bit integers, $c$ (if $t_2 = 15$, $y = 1$, else $y = 0$).
    ///
    /// First, the literals are copied directly and unprocessed to the output buffer, then (after
    /// the involved parameters are read) $t_2 + 15m + c$ bytes are copied from the output buffer
    /// at position $a + 4$ and appended to the output buffer. Note that this copy can be
    /// overlapping.
    #[inline]
    fn complete(&mut self) -> io::Result<()> {
        // Exhaust the decoder by reading and decompressing all blocks until the remaining buffer
        // is empty.
        while !self.input.is_empty() {
            // Read the token. The token is the first byte in a block. It is divided into two 4-bit
            // subtokens, the higher and the lower.
            self.token = self.take(1)?[0];

            // Now, we read the literals section.
            self.read_literal_section()?;

            // If the input stream is emptied, we break out of the loop. This is only the case
            // in the end of the stream, since the block is intact otherwise.
            if self.input.is_empty() { break; }

            // Now, we read the duplicates section.
            self.read_duplicate_section()?;
        }

        Ok(())
    }
}

/// Decompress all bytes of `input`.
pub fn decompress(input: &[u8]) -> io::Result<Vec<u8>> {
    // Allocate a vector to contain the decompressed stream.
    let mut vec = Vec::with_capacity(4096);

    // Decode into our vector.
    Decoder {
        input: input,
        output: &mut vec,
        token: 0,
    }.complete()?;

    Ok(vec)
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn aaaaaaaaaaa_lots_of_aaaaaaaaa() {
        assert_eq!(decompress(&[0x11, b'a', 1, 0]).unwrap(), b"aaaaaa");
    }

    #[test]
    fn multiple_repeated_blocks() {
        assert_eq!(decompress(&[0x11, b'a', 1, 0, 0x22, b'b', b'c', 2, 0]).unwrap(), b"aaaaaabcbcbcbc");
    }

    #[test]
    fn all_literal() {
        assert_eq!(decompress(&[0x30, b'a', b'4', b'9']).unwrap(), b"a49");
    }

    #[test]
    fn offset_oob() {
        decompress(&[0x10, b'a', 2, 0]).unwrap_err();
        decompress(&[0x40, b'a', 1, 0]).unwrap_err();
    }
}
<commit_msg>Fix overflow bug arising from misplaced parenthesizes.<commit_after>//! The decompression algorithm.

use std::io;
use byteorder::{LittleEndian, ByteOrder};

/// A LZ4 decoder.
///
/// This will decode in accordance to the LZ4 format. It represents a particular state of the
/// decompressor.
struct Decoder<'a> {
    /// The compressed input.
    input: &'a [u8],
    /// The decompressed output.
    output: &'a mut Vec<u8>,
    /// The current block's "token".
    ///
    /// This token contains to 4-bit "fields", a higher and a lower, representing the literals'
    /// length and the back reference's length, respectively. LSIC is used if either are their
    /// maximal values.
    token: u8,
}

impl<'a> Decoder<'a> {
    /// Internal (partial) function for `take`.
    #[inline]
    fn take_imp(input: &mut &'a [u8], n: usize) -> io::Result<&'a [u8]> {
        // Check if we have enough bytes left.
        if input.len() < n {
            // No extra bytes. This is clearly not expected, so we return an error.
            Err(io::Error::new(io::ErrorKind::InvalidData, "Expected another byte, found none."))
        } else {
            // Take the first n bytes.
            let res = Ok(&input[..n]);
            // Shift the stream to left, so that it is no longer the first byte.
            *input = &input[n..];

            // Return the former first byte.
            res
        }
    }

    /// Pop n bytes from the start of the input stream.
    fn take(&mut self, n: usize) -> io::Result<&[u8]> {
        Self::take_imp(&mut self.input, n)
    }

    /// Write a buffer to the output stream.
    ///
    /// The reason this doesn't take `&mut self` is that we need partial borrowing due to the rules
    /// of the borrow checker. For this reason, we instead take some number of segregated
    /// references so we can read and write them independently.
    fn output(output: &mut Vec<u8>, buf: &[u8]) {
        // We use simple memcpy to extend the vector.
        output.extend_from_slice(&buf[..buf.len()]);
    }

    /// Write an already decompressed match to the output stream.
    ///
    /// This is used for the essential part of the algorithm: deduplication. We start at some
    /// position `start` and then keep pushing the following element until we've added
    /// `match_length` elements.
    fn duplicate(&mut self, start: usize, match_length: usize) {
        // We cannot simply use memcpy or `extend_from_slice`, because these do not allow
        // self-referential copies: http://ticki.github.io/img/lz4_runs_encoding_diagram.svg
        for i in start..start + match_length {
            let b = self.output[i];
            self.output.push(b);
        }
    }

    /// Read an integer LSIC (linear small integer code) encoded.
    ///
    /// In LZ4, we encode small integers in a way that we can have an arbitrary number of bytes. In
    /// particular, we add the bytes repeatedly until we hit a non-0xFF byte. When we do, we add
    /// this byte to our sum and terminate the loop.
    ///
    /// # Example
    ///
    /// ```notest
    ///     255, 255, 255, 4, 2, 3, 4, 6, 7
    /// ```
    ///
    /// is encoded to _255 + 255 + 255 + 4 = 769_. The bytes after the first 4 is ignored, because
    /// 4 is the first non-0xFF byte.
    #[inline]
    fn read_integer(&mut self) -> io::Result<usize> {
        // We start at zero and count upwards.
        let mut n = 0;
        // If this byte takes value 255 (the maximum value it can take), another byte is read
        // and added to the sum. This repeats until a byte lower than 255 is read.
        while {
            // We add the next byte until we get a byte which we add to the counting variable.
            let extra = self.take(1)?[0];
            n += extra as usize;

            // We continue if we got 255.
            extra == 0xFF
        } {}

        Ok(n)
    }

    /// Read a little-endian 16-bit integer from the input stream.
    #[inline]
    fn read_u16(&mut self) -> io::Result<u16> {
        // We use byteorder to read an u16 in little endian.
        Ok(LittleEndian::read_u16(self.take(2)?))
    }

    /// Read the literals section of a block.
    ///
    /// The literals section encodes some bytes which are to be copied to the output without any
    /// modification.
    ///
    /// It consists of two parts:
    ///
    /// 1. An LSIC integer extension to the literals length as defined by the first part of the
    ///    token, if it takes the highest value (15).
    /// 2. The literals themself.
    fn read_literal_section(&mut self) -> io::Result<()> {
        // The higher token is the literals part of the token. It takes a value from 0 to 15.
        let mut literal = (self.token >> 4) as usize;
        // If the initial value is 15, it is indicated that another byte will be read and added to
        // it.
        if literal == 15 {
            // The literal length took the maximal value, indicating that there is more than 15
            // literal bytes. We read the extra integer.
            literal += self.read_integer()?;
        }

        // Now we know the literal length. The number will be used to indicate how long the
        // following literal copied to the output buffer is.

        // Read the literals segment and output them without processing.
        Self::output(&mut self.output, Self::take_imp(&mut self.input, literal)?);

        Ok(())
    }

    /// Read the duplicates section of the block.
    ///
    /// The duplicates section serves to reference an already decoded segment. This consists of two
    /// parts:
    ///
    /// 1. A 16-bit little-endian integer defining the "offset", i.e. how long back we need to go
    ///    in the decoded buffer and copy.
    /// 2. An LSIC integer extension to the duplicate length as defined by the first part of the
    ///    token, if it takes the highest value (15).
    fn read_duplicate_section(&mut self) -> io::Result<()> {
        // Now, we will obtain the offset which we will use to copy from the output. It is an
        // 16-bit integer.
        let offset = self.read_u16()?;

        // Obtain the initial match length. The match length is the length of the duplicate segment
        // which will later be copied from data previously decompressed into the output buffer. The
        // initial length is derived from the second part of the token (the lower nibble), we read
        // earlier. Since having a match length of less than 4 would mean negative compression
        // ratio, we start at 4.
        let mut match_length = (4 + (self.token & 0xF)) as usize;

        // The intial match length can maximally be 19. As with the literal length, this indicates
        // that there are more bytes to read.
        if match_length == 4 + 15 {
            // The match length took the maximal value, indicating that there is more bytes. We
            // read the extra integer.
            match_length += self.read_integer()?;
        }

        // We now copy from the already decompressed buffer. This allows us for storing duplicates
        // by simply referencing the other location.

        // Calculate the start of this duplicate segment. We use wrapping subtraction to avoid
        // overflow checks, which we will catch later.
        let start = self.output.len().wrapping_sub(offset as usize);

        // We'll do a bound check to avoid panicking.
        if start < self.output.len() {
            // Write the duplicate segment to the output buffer.
            self.duplicate(start, match_length);

            Ok(())
        } else {
            Err(io::Error::new(io::ErrorKind::InvalidData, "The offset to copy is not contained in \
                               the decompressed buffer."))
        }
    }

    /// Complete the decompression by reading all the blocks.
    ///
    /// # Decompressing a block
    ///
    /// Blocks consists of:
    ///  - A 1 byte token
    ///      * A 4 bit integer $t_1$.
    ///      * A 4 bit integer $t_2$.
    ///  - A $n$ byte sequence of 0xFF bytes (if $t_1 \neq 15$, then $n = 0$).
    ///  - $x$ non-0xFF 8-bit integers, L (if $t_1 = 15$, $x = 1$, else $x = 0$).
    ///  - $t_1 + 15n + L$ bytes of uncompressed data (literals).
    ///  - 16-bits offset (little endian), $a$.
    ///  - A $m$ byte sequence of 0xFF bytes (if $t_2 \neq 15$, then $m = 0$).
    ///  - $y$ non-0xFF 8-bit integers, $c$ (if $t_2 = 15$, $y = 1$, else $y = 0$).
    ///
    /// First, the literals are copied directly and unprocessed to the output buffer, then (after
    /// the involved parameters are read) $t_2 + 15m + c$ bytes are copied from the output buffer
    /// at position $a + 4$ and appended to the output buffer. Note that this copy can be
    /// overlapping.
    #[inline]
    fn complete(&mut self) -> io::Result<()> {
        // Exhaust the decoder by reading and decompressing all blocks until the remaining buffer
        // is empty.
        while !self.input.is_empty() {
            // Read the token. The token is the first byte in a block. It is divided into two 4-bit
            // subtokens, the higher and the lower.
            self.token = self.take(1)?[0];

            // Now, we read the literals section.
            self.read_literal_section()?;

            // If the input stream is emptied, we break out of the loop. This is only the case
            // in the end of the stream, since the block is intact otherwise.
            if self.input.is_empty() { break; }

            // Now, we read the duplicates section.
            self.read_duplicate_section()?;
        }

        Ok(())
    }
}

/// Decompress all bytes of `input`.
pub fn decompress(input: &[u8]) -> io::Result<Vec<u8>> {
    // Allocate a vector to contain the decompressed stream.
    let mut vec = Vec::with_capacity(4096);

    // Decode into our vector.
    Decoder {
        input: input,
        output: &mut vec,
        token: 0,
    }.complete()?;

    Ok(vec)
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn aaaaaaaaaaa_lots_of_aaaaaaaaa() {
        assert_eq!(decompress(&[0x11, b'a', 1, 0]).unwrap(), b"aaaaaa");
    }

    #[test]
    fn multiple_repeated_blocks() {
        assert_eq!(decompress(&[0x11, b'a', 1, 0, 0x22, b'b', b'c', 2, 0]).unwrap(), b"aaaaaabcbcbcbc");
    }

    #[test]
    fn all_literal() {
        assert_eq!(decompress(&[0x30, b'a', b'4', b'9']).unwrap(), b"a49");
    }

    #[test]
    fn offset_oob() {
        decompress(&[0x10, b'a', 2, 0]).unwrap_err();
        decompress(&[0x40, b'a', 1, 0]).unwrap_err();
    }
}
<|endoftext|>