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

extern crate bytes;
extern crate httpbis;
extern crate futures;
extern crate native_tls;
extern crate tokio_core;
extern crate tokio_tls;
#[macro_use]
extern crate log;
extern crate env_logger;

use bytes::Bytes;

use futures::Future;

mod test_misc;

use httpbis::solicit::header::*;
use httpbis::error::Error;
use httpbis::error::ErrorCode;
use httpbis::*;
use httpbis::for_test::*;

use test_misc::*;

#[test]
fn stream_count() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    let mut server_tester = server.accept();
    server_tester.recv_preface();
    server_tester.settings_xchg();

    let state: ConnectionStateSnapshot = client.dump_state().wait().expect("state");
    assert_eq!(0, state.streams.len());

    let req = client.start_post("/foobar", "localhost", Bytes::from(&b"xxyy"[..])).collect();

    let headers = server_tester.recv_frame_headers_check(1, false);
    assert_eq!("POST", headers.get(":method"));
    assert_eq!("/foobar", headers.get(":path"));

    let data = server_tester.recv_frame_data_check(1, false);
    assert_eq!(b"xxyy", &data[..]);

    server_tester.recv_frame_data_check_empty_end(1);

    let mut resp_headers = Headers::new();
    resp_headers.add(":status", "200");
    server_tester.send_headers(1, resp_headers, false);

    server_tester.send_data(1, b"aabb", true);

    let message = req.wait().expect("r");
    assert_eq!((b"aabb"[..]).to_owned(), message.body);

    let state: ConnectionStateSnapshot = client.dump_state().wait().expect("state");
    assert_eq!(0, state.streams.len(), "{:?}", state);
}

#[test]
fn rst_is_error() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    let mut server_tester = server.accept();
    server_tester.recv_preface();
    server_tester.settings_xchg();

    let req = client.start_get("/fgfg", "localhost").collect();

    let get = server_tester.recv_message(1);
    assert_eq!("GET", get.headers.method());

    server_tester.send_headers(1, Headers::ok_200(), false);
    server_tester.send_rst(1, ErrorCode::InadequateSecurity);

    match req.wait() {
        Ok(..) => panic!("expected error"),
        Err(Error::CodeError(ErrorCode::InadequateSecurity)) => {},
        Err(e) => panic!("wrong error: {:?}", e),
    }

    let state: ConnectionStateSnapshot = client.dump_state().wait().expect("state");
    assert_eq!(0, state.streams.len(), "{:?}", state);
}

#[test]
fn client_call_dropped() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    let mut server_tester = server.accept();
    server_tester.recv_preface();
    server_tester.settings_xchg();

    {
        let req = client.start_get("/fgfg", "localhost");

        server_tester.recv_message(1);

        drop(req);

        server_tester.send_headers(1, Headers::ok_200(), true);
    }

    {
        let req = client.start_get("/fgfg", "localhost").collect();
        server_tester.recv_message(3);
        server_tester.send_headers(3, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());
    }

    let state: ConnectionStateSnapshot = client.dump_state().wait().expect("state");
    assert_eq!(0, state.streams.len(), "{:?}", state);
}

#[test]
fn reconnect_on_disconnect() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    let mut server_tester = server.accept();
    server_tester.recv_preface();
    server_tester.settings_xchg();

    {
        let req = client.start_get("/111", "localhost").collect();
        server_tester.recv_message(1);
        server_tester.send_headers(1, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());
    }

    // drop server connection
    drop(server_tester);

    // waiting for client connection to die
    while let Ok(_) = client.dump_state().wait() {
        thread::sleep(Duration::from_millis(1));
    }

    {
        let req = client.start_get("/222", "localhost").collect();

        let mut server_tester = server.accept();
        server_tester.recv_preface();
        server_tester.settings_xchg_but_ack();

        server_tester.recv_message(1);
        server_tester.send_headers(1, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());
    }
}

#[test]
fn reconnect_on_goaway() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    {
        let mut server_tester = server.accept();
        server_tester.recv_preface();
        server_tester.settings_xchg();

        let req = client.start_get("/111", "localhost").collect();
        server_tester.recv_message(1);
        server_tester.send_headers(1, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());

        server_tester.send_goaway(1);

        // TODO
        server_tester.recv_eof();
    }

    {
        let req = client.start_get("/111", "localhost").collect();

        let mut server_tester = server.accept();
        server_tester.recv_preface();
        server_tester.settings_xchg_but_ack();

        server_tester.recv_message(1);
        server_tester.send_headers(1, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());
    }
}
<commit_msg>Outdated todo<commit_after>use std::str;
use std::thread;
use std::time::Duration;

extern crate bytes;
extern crate httpbis;
extern crate futures;
extern crate native_tls;
extern crate tokio_core;
extern crate tokio_tls;
#[macro_use]
extern crate log;
extern crate env_logger;

use bytes::Bytes;

use futures::Future;

mod test_misc;

use httpbis::solicit::header::*;
use httpbis::error::Error;
use httpbis::error::ErrorCode;
use httpbis::*;
use httpbis::for_test::*;

use test_misc::*;

#[test]
fn stream_count() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    let mut server_tester = server.accept();
    server_tester.recv_preface();
    server_tester.settings_xchg();

    let state: ConnectionStateSnapshot = client.dump_state().wait().expect("state");
    assert_eq!(0, state.streams.len());

    let req = client.start_post("/foobar", "localhost", Bytes::from(&b"xxyy"[..])).collect();

    let headers = server_tester.recv_frame_headers_check(1, false);
    assert_eq!("POST", headers.get(":method"));
    assert_eq!("/foobar", headers.get(":path"));

    let data = server_tester.recv_frame_data_check(1, false);
    assert_eq!(b"xxyy", &data[..]);

    server_tester.recv_frame_data_check_empty_end(1);

    let mut resp_headers = Headers::new();
    resp_headers.add(":status", "200");
    server_tester.send_headers(1, resp_headers, false);

    server_tester.send_data(1, b"aabb", true);

    let message = req.wait().expect("r");
    assert_eq!((b"aabb"[..]).to_owned(), message.body);

    let state: ConnectionStateSnapshot = client.dump_state().wait().expect("state");
    assert_eq!(0, state.streams.len(), "{:?}", state);
}

#[test]
fn rst_is_error() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    let mut server_tester = server.accept();
    server_tester.recv_preface();
    server_tester.settings_xchg();

    let req = client.start_get("/fgfg", "localhost").collect();

    let get = server_tester.recv_message(1);
    assert_eq!("GET", get.headers.method());

    server_tester.send_headers(1, Headers::ok_200(), false);
    server_tester.send_rst(1, ErrorCode::InadequateSecurity);

    match req.wait() {
        Ok(..) => panic!("expected error"),
        Err(Error::CodeError(ErrorCode::InadequateSecurity)) => {},
        Err(e) => panic!("wrong error: {:?}", e),
    }

    let state: ConnectionStateSnapshot = client.dump_state().wait().expect("state");
    assert_eq!(0, state.streams.len(), "{:?}", state);
}

#[test]
fn client_call_dropped() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    let mut server_tester = server.accept();
    server_tester.recv_preface();
    server_tester.settings_xchg();

    {
        let req = client.start_get("/fgfg", "localhost");

        server_tester.recv_message(1);

        drop(req);

        server_tester.send_headers(1, Headers::ok_200(), true);
    }

    {
        let req = client.start_get("/fgfg", "localhost").collect();
        server_tester.recv_message(3);
        server_tester.send_headers(3, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());
    }

    let state: ConnectionStateSnapshot = client.dump_state().wait().expect("state");
    assert_eq!(0, state.streams.len(), "{:?}", state);
}

#[test]
fn reconnect_on_disconnect() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    let mut server_tester = server.accept();
    server_tester.recv_preface();
    server_tester.settings_xchg();

    {
        let req = client.start_get("/111", "localhost").collect();
        server_tester.recv_message(1);
        server_tester.send_headers(1, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());
    }

    // drop server connection
    drop(server_tester);

    // waiting for client connection to die
    while let Ok(_) = client.dump_state().wait() {
        thread::sleep(Duration::from_millis(1));
    }

    {
        let req = client.start_get("/222", "localhost").collect();

        let mut server_tester = server.accept();
        server_tester.recv_preface();
        server_tester.settings_xchg_but_ack();

        server_tester.recv_message(1);
        server_tester.send_headers(1, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());
    }
}

#[test]
fn reconnect_on_goaway() {
    env_logger::init().ok();

    let server = HttpServerTester::new();

    let client: Client =
        Client::new("::1", server.port(), false, Default::default()).expect("connect");

    {
        let mut server_tester = server.accept();
        server_tester.recv_preface();
        server_tester.settings_xchg();

        let req = client.start_get("/111", "localhost").collect();
        server_tester.recv_message(1);
        server_tester.send_headers(1, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());

        server_tester.send_goaway(1);

        server_tester.recv_eof();
    }

    {
        let req = client.start_get("/111", "localhost").collect();

        let mut server_tester = server.accept();
        server_tester.recv_preface();
        server_tester.settings_xchg_but_ack();

        server_tester.recv_message(1);
        server_tester.send_headers(1, Headers::ok_200(), true);
        let resp = req.wait().expect("OK");
        assert_eq!(200, resp.headers.status());
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Berkus Decker <berkus+github@metta.systems>
//
// 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 example shows simple packet_fanout processing under linux.
extern crate pnet;
extern crate pnet_datalink;

use pnet::datalink::{self, Config, FanoutOption, FanoutType, NetworkInterface};
use std::env;
use std::io::{self, Write};
use std::process;
use std::thread;

#[cfg(not(target_os = "linux"))]
fn main() {
    writeln!(io::stderr(), "fanout is only supported on Linux").unwrap();
    process::exit(1);
}

#[cfg(target_os = "linux")]
fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = match env::args().nth(1) {
        Some(n) => n,
        None => {
            writeln!(io::stderr(), "USAGE: fanout <NETWORK INTERFACE> [hash|*round-robin*|cpu|rollover|rnd|qm|cbpf|ebpf] [group-id:123]").unwrap();
            process::exit(1);
        }
    };
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces
        .into_iter()
        .filter(interface_names_match)
        .next()
        .unwrap();

    let fanout_type = match env::args().nth(2) {
        Some(n) => match n.to_lowercase().as_str() {
            "hash" => FanoutType::HASH,
            "round-robin" => FanoutType::LB,
            "cpu" => FanoutType::CPU,
            "rollover" => FanoutType::ROLLOVER,
            "rnd" => FanoutType::RND,
            "qm" => FanoutType::QM,
            "cbpf" => FanoutType::CBPF,
            "ebpf" => FanoutType::EBPF,
            _ => panic!("Unsupported fanout type, use one of hash, round-robin, cpu, rollover, rnd, qm, cbpf or ebpf")
        },
        None => FanoutType::LB,
    };

    let group_id = match env::args().nth(3) {
        Some(n) => n.parse::<u16>().unwrap(),
        None => 123,
    };

    let mut config: Config = Default::default();
    config.linux_fanout = Some(FanoutOption {
        group_id: group_id,
        fanout_type: fanout_type,
        defrag: true,
        rollover: false,
    });

    let mut threads = vec![];
    for x in 0..3 {
        let itf = interface.clone();
        let thread = thread::Builder::new()
            .name(format!("thread{}", x))
            .spawn(move || {
                // Create a channel to receive on
                let (_, mut rx) = match datalink::channel(&itf, config) {
                    Ok(Ethernet(tx, rx)) => (tx, rx),
                    Ok(_) => panic!("packetdump: unhandled channel type: {}"),
                    Err(e) => panic!("packetdump: unable to create channel: {}", e),
                };

                let handle = thread::current();

                loop {
                    match rx.next() {
                        Ok(_packet) => {
                            writeln!(
                                io::stdout(),
                                "Received packet on thread {:?}",
                                handle.name()
                            ).unwrap();
                        }
                        Err(e) => panic!("packetdump: unable to receive packet: {}", e),
                    }
                }
            })
            .unwrap();
        threads.push(thread);
    }

    for t in threads {
        t.join().unwrap();
    }
}
<commit_msg>Add description for fanout example<commit_after>// Copyright (c) 2018 Berkus Decker <berkus+github@metta.systems>
//
// 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 example shows simple packet_fanout processing under linux.
/// PACKET_FANOUT in linux allows to offload packet processing to multiple threads.
/// See [man 7 packet](http://man7.org/linux/man-pages/man7/packet.7.html) for more details.
extern crate pnet;
extern crate pnet_datalink;

use pnet::datalink::{self, Config, FanoutOption, FanoutType, NetworkInterface};
use std::env;
use std::io::{self, Write};
use std::process;
use std::thread;

#[cfg(not(target_os = "linux"))]
fn main() {
    writeln!(io::stderr(), "fanout is only supported on Linux").unwrap();
    process::exit(1);
}

#[cfg(target_os = "linux")]
fn main() {
    use pnet::datalink::Channel::Ethernet;

    let iface_name = match env::args().nth(1) {
        Some(n) => n,
        None => {
            writeln!(io::stderr(), "USAGE: fanout <NETWORK INTERFACE> [hash|*round-robin*|cpu|rollover|rnd|qm|cbpf|ebpf] [group-id:123]").unwrap();
            process::exit(1);
        }
    };
    let interface_names_match = |iface: &NetworkInterface| iface.name == iface_name;

    // Find the network interface with the provided name
    let interfaces = datalink::interfaces();
    let interface = interfaces
        .into_iter()
        .filter(interface_names_match)
        .next()
        .unwrap();

    let fanout_type = match env::args().nth(2) {
        Some(n) => match n.to_lowercase().as_str() {
            "hash" => FanoutType::HASH,
            "round-robin" => FanoutType::LB,
            "cpu" => FanoutType::CPU,
            "rollover" => FanoutType::ROLLOVER,
            "rnd" => FanoutType::RND,
            "qm" => FanoutType::QM,
            "cbpf" => FanoutType::CBPF,
            "ebpf" => FanoutType::EBPF,
            _ => panic!("Unsupported fanout type, use one of hash, round-robin, cpu, rollover, rnd, qm, cbpf or ebpf")
        },
        None => FanoutType::LB,
    };

    let group_id = match env::args().nth(3) {
        Some(n) => n.parse::<u16>().unwrap(),
        None => 123,
    };

    let mut config: Config = Default::default();
    config.linux_fanout = Some(FanoutOption {
        group_id: group_id,
        fanout_type: fanout_type,
        defrag: true,
        rollover: false,
    });

    let mut threads = vec![];
    for x in 0..3 {
        let itf = interface.clone();
        let thread = thread::Builder::new()
            .name(format!("thread{}", x))
            .spawn(move || {
                // Create a channel to receive on
                let (_, mut rx) = match datalink::channel(&itf, config) {
                    Ok(Ethernet(tx, rx)) => (tx, rx),
                    Ok(_) => panic!("packetdump: unhandled channel type: {}"),
                    Err(e) => panic!("packetdump: unable to create channel: {}", e),
                };

                let handle = thread::current();

                loop {
                    match rx.next() {
                        Ok(_packet) => {
                            writeln!(
                                io::stdout(),
                                "Received packet on thread {:?}",
                                handle.name()
                            ).unwrap();
                        }
                        Err(e) => panic!("packetdump: unable to receive packet: {}", e),
                    }
                }
            })
            .unwrap();
        threads.push(thread);
    }

    for t in threads {
        t.join().unwrap();
    }
}
<|endoftext|>


<commit_before>#![feature(unboxed_closures)]

extern crate synthrs;

use synthrs::synthesizer::{ make_samples, quantize_samples, peak_normalize };
use synthrs::wave::{ SineWave, SquareWave, SawtoothWave, TriangleWave, TangentWave, Bell, KarplusStrong, Noise };
use synthrs::writer::{ write_pcm, write_wav };

fn main() {
    // This creates a sine wave for 1.0s at 44100Hz
    // 0. SineWave create a 440Hz sine function
    // 1. make_samples creates a Vec of samples from the SineWave function
    //    from 0.0 to 1.0 seconds at a 44100Hz sample rate
    // 2. quantize_samples::<i16> quantizes the floating-point samples as a signed 16-bit int
    // 3. write_pcm writes the samples to a PCM file
    write_pcm("out/sin.pcm",
        &quantize_samples::<i16>(&make_samples(1.0, 44100, SineWave(440.0)))
    ).expect("failed");

    write_wav("out/sin.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, SineWave(440.0)))
    ).expect("failed");

    write_wav("out/square.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, SquareWave(440.0)))
    ).expect("failed");

    write_wav("out/sawtooth.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, SawtoothWave(440.0)))
    ).expect("failed");

    write_wav("out/triangle.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, TriangleWave(440.0)))
    ).expect("failed");

    write_wav("out/tangent.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, TangentWave(440.0)))
    ).expect("failed");

    write_wav("out/noise.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, Noise))
    ).expect("failed");

    // Custom function for tone generation, t is in seconds
    write_wav("out/wolftone.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(1.0, 44100, |t: f64| -> f64 {
                (SquareWave(1000.0)(t) + SquareWave(1020.0)(t)) / 2.0
            })
        )
    ).expect("failed");

    write_wav("out/rising.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(1.0, 44100, |t: f64| -> f64 {
                let (min_f, max_f) = (1000.0, 8000.0);
                let max_t = 1.0; // Duration of clip in seconds
                let range = max_f - min_f;
                let f = max_f - (max_t - t) * range;
                SineWave(f)(t)
            })
        )
    ).expect("failed");

    write_wav("out/bell.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(10.0, 44100, |t: f64| -> f64 {
                Bell(200.0, 0.003, 0.5)(t)
            })
        )
    ).expect("failed");

    write_wav("out/karplusstrong.wav", 44100,
        &quantize_samples::<i16>(
            &peak_normalize(
                &make_samples(5.0, 44100, |t: f64| -> f64 {
                    KarplusStrong(SawtoothWave(440.0), 0.01, 1.0, 0.9, 44100.0)(t)
                })
            )
        )
    ).expect("failed");

    write_wav("out/racecar.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(15.0, 44100, |t: f64| -> f64 {
                let mut out = 0.0;
                if t < 14.0 { out += SawtoothWave(40.63 * (t / 2.0))(t); } // Engine
                if t < 1.0 { out += SawtoothWave(30.0)(t) * 10.0; } // Engine start
                if t < 14.0 && t.tan() > 0.0 { out += SquareWave(t.fract() * 130.8125)(t); } // Gear
                if t < 14.0 && t.tan() < 0.0 { out += SquareWave(t.sin() * 261.625)(t); } // Gear
                if t < 14.0 && t.tan() > 0.866 { out += SquareWave(t.fract() * 523.25)(t); } // Gear
                if t > 14.0 { out += TangentWave(100.0)(t) } // Tree
                (out / 4.0).min(1.0)
            })
        )
    ).expect("failed");

    write_wav("out/shepard.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(30.0, 44100, |t: f64| -> f64 {
                let length = 10.0;
                let t_mod = t % length;
                let progress = t_mod / length;
                let tone_a2 = SineWave(110.0 / (1.0 + progress));
                let tone_a3 = SineWave(220.0 / (1.0 + progress));
                let tone_a4 = SineWave(440.0 / (1.0 + progress));
                let tone_a5 = SineWave(880.0 / (1.0 + progress));
                (tone_a2(t_mod) * (1.0 - progress)
                + tone_a5(t_mod) * progress
                + tone_a3(t_mod) + tone_a4(t_mod)) / 4.0
            })
        )
    ).expect("failed");
}
<commit_msg>Add code blocks to comment for write_pcm example<commit_after>#![feature(unboxed_closures)]

extern crate synthrs;

use synthrs::synthesizer::{ make_samples, quantize_samples, peak_normalize };
use synthrs::wave::{ SineWave, SquareWave, SawtoothWave, TriangleWave, TangentWave, Bell, KarplusStrong, Noise };
use synthrs::writer::{ write_pcm, write_wav };

fn main() {
    // This creates a sine wave for 1.0s at 44100Hz
    // 0. `SineWave` create a 440Hz sine function
    // 1. `make_samples` creates a `Vec` of samples from the `SineWave` function
    //    from 0.0 to 1.0 seconds at a 44100Hz sample rate
    // 2. `quantize_samples::<i16>` quantizes the floating-point samples as a signed 16-bit int
    // 3. `write_pcm` writes the samples to a PCM file
    write_pcm("out/sin.pcm",
        &quantize_samples::<i16>(&make_samples(1.0, 44100, SineWave(440.0)))
    ).expect("failed");

    write_wav("out/sin.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, SineWave(440.0)))
    ).expect("failed");

    write_wav("out/square.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, SquareWave(440.0)))
    ).expect("failed");

    write_wav("out/sawtooth.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, SawtoothWave(440.0)))
    ).expect("failed");

    write_wav("out/triangle.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, TriangleWave(440.0)))
    ).expect("failed");

    write_wav("out/tangent.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, TangentWave(440.0)))
    ).expect("failed");

    write_wav("out/noise.wav", 44100,
        &quantize_samples::<i16>(&make_samples(1.0, 44100, Noise))
    ).expect("failed");

    // Custom function for tone generation, t is in seconds
    write_wav("out/wolftone.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(1.0, 44100, |t: f64| -> f64 {
                (SquareWave(1000.0)(t) + SquareWave(1020.0)(t)) / 2.0
            })
        )
    ).expect("failed");

    write_wav("out/rising.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(1.0, 44100, |t: f64| -> f64 {
                let (min_f, max_f) = (1000.0, 8000.0);
                let max_t = 1.0; // Duration of clip in seconds
                let range = max_f - min_f;
                let f = max_f - (max_t - t) * range;
                SineWave(f)(t)
            })
        )
    ).expect("failed");

    write_wav("out/bell.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(10.0, 44100, |t: f64| -> f64 {
                Bell(200.0, 0.003, 0.5)(t)
            })
        )
    ).expect("failed");

    write_wav("out/karplusstrong.wav", 44100,
        &quantize_samples::<i16>(
            &peak_normalize(
                &make_samples(5.0, 44100, |t: f64| -> f64 {
                    KarplusStrong(SawtoothWave(440.0), 0.01, 1.0, 0.9, 44100.0)(t)
                })
            )
        )
    ).expect("failed");

    write_wav("out/racecar.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(15.0, 44100, |t: f64| -> f64 {
                let mut out = 0.0;
                if t < 14.0 { out += SawtoothWave(40.63 * (t / 2.0))(t); } // Engine
                if t < 1.0 { out += SawtoothWave(30.0)(t) * 10.0; } // Engine start
                if t < 14.0 && t.tan() > 0.0 { out += SquareWave(t.fract() * 130.8125)(t); } // Gear
                if t < 14.0 && t.tan() < 0.0 { out += SquareWave(t.sin() * 261.625)(t); } // Gear
                if t < 14.0 && t.tan() > 0.866 { out += SquareWave(t.fract() * 523.25)(t); } // Gear
                if t > 14.0 { out += TangentWave(100.0)(t) } // Tree
                (out / 4.0).min(1.0)
            })
        )
    ).expect("failed");

    write_wav("out/shepard.wav", 44100,
        &quantize_samples::<i16>(
            &make_samples(30.0, 44100, |t: f64| -> f64 {
                let length = 10.0;
                let t_mod = t % length;
                let progress = t_mod / length;
                let tone_a2 = SineWave(110.0 / (1.0 + progress));
                let tone_a3 = SineWave(220.0 / (1.0 + progress));
                let tone_a4 = SineWave(440.0 / (1.0 + progress));
                let tone_a5 = SineWave(880.0 / (1.0 + progress));
                (tone_a2(t_mod) * (1.0 - progress)
                + tone_a5(t_mod) * progress
                + tone_a3(t_mod) + tone_a4(t_mod)) / 4.0
            })
        )
    ).expect("failed");
}
<|endoftext|>


<commit_before>extern crate polydraw;

use polydraw::os::win;
use polydraw::os::wgl;
use polydraw::os::gl;

fn main() {
   let width: usize = 1280;
   let height: usize = 720;

   let window = win::Window::create(width, height, "Win Min", "WinMinClass");
   window.show_normal();

   let dc = window.dc();
   println!("DC: {:?}", dc);

   match wgl::init_pixel_format(dc) {
      Ok(_) => {},
      Err(e) => {
         panic!(e.description);
      }
   };

   let context = match wgl::Context::create(dc) {
      Ok(context) => context,
      Err(e) => {
         panic!(e.description);
      }
   };

   println!("RC: {:?}", context.rc);

   match wgl::Context::current() {
      Ok(_) => {},
      Err(e) => {
         panic!(e.description);
      }
   };

   gl::load(wgl::Loader::new());

   loop {
      let message = match win::Message::get() {
         Some(message) => message,
         None => break
      };

      message.translate();
      message.dispatch();

      gl::clear_color(0.0, 0.7, 1.0, 1.0);
      gl::clear();
      gl::flush();

      wgl::swap_buffers(dc);
   }
}
<commit_msg>Not retreiving again the current WGL context<commit_after>extern crate polydraw;

use polydraw::os::win;
use polydraw::os::wgl;
use polydraw::os::gl;

fn main() {
   let width: usize = 1280;
   let height: usize = 720;

   let window = win::Window::create(width, height, "Win Min", "WinMinClass");
   window.show_normal();

   let dc = window.dc();
   println!("DC: {:?}", dc);

   match wgl::init_pixel_format(dc) {
      Ok(_) => {},
      Err(e) => {
         panic!(e.description);
      }
   };

   let context = match wgl::Context::create(dc) {
      Ok(context) => context,
      Err(e) => {
         panic!(e.description);
      }
   };

   println!("RC: {:?}", context.rc);

   gl::load(wgl::Loader::new());

   loop {
      let message = match win::Message::get() {
         Some(message) => message,
         None => break
      };

      message.translate();
      message.dispatch();

      gl::clear_color(0.0, 0.7, 1.0, 1.0);
      gl::clear();
      gl::flush();

      wgl::swap_buffers(dc);
   }
}
<|endoftext|>


<commit_before>extern crate rusty_machine as rm;
extern crate rand;

use rm::linalg::{Matrix, BaseMatrix};
use rand::{Rng, ThreadRng};
use std::collections::HashMap;

fn generate_arm(content_id: &usize, arms: &HashMap<usize, Arm>, feature_dim: &usize) -> Arm {
    let arm: Arm;
    if !arms.contains_key(content_id) {
        arm = Arm::new(*feature_dim, *content_id, 0.0001, 0.0);
    }else{
        arm = arms.get(content_id).unwrap().clone();
    };
    arm
}

#[derive(Debug, Clone)]
struct Arm{
    content_id : usize,
    alpha : f64,
    norm_mean : Matrix<f64>,
    cov_matrix: Matrix<f64>,
    win_rate : f64,
    win : f64,
    lose : f64,
}

impl Arm {
    fn new(feature_dim: usize, content_id: usize, alpha: f64, win_rate: f64) -> Arm {
        Arm {
            content_id: content_id,
            norm_mean: Matrix::zeros(1, feature_dim),
            cov_matrix: Matrix::ones(1, feature_dim),
            alpha: alpha,
            win_rate: win_rate,
            win: 0.0,
            lose: 0.0,
        }
    }

    fn update(&mut self, features: Matrix<f64>, is_click: bool) {
        // error of diag
        self.cov_matrix += Matrix::new(self.cov_matrix.rows(),
                                        self.cov_matrix.cols(),
                                        (&features.transpose()*&features).diag());
        let cost_of_click = is_click as i8 as f64;

        let feat_mul_cost_vec = features.iter().map(|x| cost_of_click*x).collect::<Vec<f64>>();
        let feat_mul_cost = Matrix::new(features.rows(), features.cols(), feat_mul_cost_vec);
        self.norm_mean += feat_mul_cost;

        if is_click{
            self.win += 1.;
        } else{
            self.lose += 1.;
        }
        self.win_rate = &self.win/(&self.win + &self.lose);
    }

    fn predict(&mut self, features: &Matrix<f64>) -> f64 {
        let one_div_cov_vec = self.cov_matrix.iter().map(|x| 1.0/x).collect::<Vec<f64>>();
        let one_div_cov = Matrix::new(self.cov_matrix.rows(),
                                      self.cov_matrix.cols(), one_div_cov_vec);
        // Since the covariance matrix preserves only the diagonal components,
        // it suffices to take the inverse matrix
        let theta = &one_div_cov.elemul(&self.norm_mean);
        // Again, the inverse matrix of the covariance matrix
        // is computed by taking reciprocal
        let mut tmp: f64 = ((features.elemul(&one_div_cov))*features.transpose()).data()[0];
        tmp = &tmp.sqrt()*&self.alpha;
        (theta*features.transpose()).data()[0] + &tmp
    }

    fn print_result(&mut self) {
        println!("content_id:{}, total_num:{}, win_rate:{}",
                &self.content_id, &self.win+&self.lose, &self.win_rate);
    }
}

struct Viewer{
    gender: String,
    rng: ThreadRng,
}

impl Viewer {
    fn new(gender: String) -> Viewer{
        Viewer{
            gender: gender,
            rng: rand::thread_rng(),
        }
    }

    fn view(&mut self, content_id: &usize) -> bool{
        if &self.gender=="man" {
            // Men are easy to click on ads with id 5 or less
            if *content_id < 6 {
                return Some(self.rng.next_f32()).and_then(|n| if n>0.3 {Some(true)} else {Some(false)}).unwrap();
            } else{
                return Some(self.rng.next_f32()).and_then(|n| if n>0.7 {Some(true)} else {Some(false)}).unwrap();
            }
        } else {
            // Women are easy to click on ads with id 6 or higher
            if *content_id > 5{
                return Some(self.rng.next_f32()).and_then(|n| if n>0.3 {Some(true)} else {Some(false)}).unwrap();
            } else {
                return Some(self.rng.next_f32()).and_then(|n| if n>0.7 {Some(true)} else {Some(false)}).unwrap();
            }
        }
    }
}

struct Rulet{
    rng: ThreadRng,
}

impl Rulet{
    fn new() -> Rulet {
        Rulet{
            rng: rand::thread_rng(),
        }
    }

    fn generate_features(&mut self, viewer: &Viewer) -> Matrix<f64> {
        let features = Some(&viewer.gender)
            .and_then(|gender|
                if gender=="man" {Some(Matrix::new(1,2,vec![1.,0.]))}
                else {Some(Matrix::new(1,2,vec![0.,1.]))}
            ).unwrap();
        features
    }

    fn generate_content(&mut self) -> usize{
        self.rng.gen_range(0, 10)
    }

    fn generate_gender(&mut self) -> String {
        if self.rng.next_f32() > 0.5{
            return "man".to_string();
        } else{
            return "women".to_string();
        }
    }
}

fn main() {
    /*Context is for men and women only
    Men are easy to click on ads with id 5 or less
    Women are easy to click on ads with id 6 or higher
    */
    let feature_dim = 2;
    let num_of_views = 10000;
    let mut rulet = Rulet::new();

    let mut content_id: usize;
    let mut features: Matrix<f64>;
    let mut is_clicked: bool;
    let mut arms: HashMap<usize, Arm> = HashMap::new();

    for _ in 0..num_of_views {
        let mut viewer = Viewer::new(rulet.generate_gender());

        content_id = rulet.generate_content();
        features = rulet.generate_features(&viewer);
        is_clicked = viewer.view(&content_id);

        let mut arm = generate_arm(&content_id, &arms, &feature_dim);
        arm.update(features, is_clicked);
        arms.remove(&content_id);
        arms.insert(content_id, arm.clone());
    }

    let man_mat: Matrix<f64> = Matrix::new(1,2, vec![1.,0.]);
    let woman_mat: Matrix<f64> = Matrix::new(1,2, vec![0.,1.]);

    println!("print result======");
    for (_, arm) in arms.iter() {
        let mut arm = arm.clone();  // impl needs mutable self
        &arm.print_result();
        println!("Click rate when men browse: {}", &arm.predict(&man_mat) );
        println!("Click rate when women browse: {}", &arm.predict(&woman_mat) );
    }

}
<commit_msg>remove comment<commit_after>extern crate rusty_machine as rm;
extern crate rand;

use rm::linalg::{Matrix, BaseMatrix};
use rand::{Rng, ThreadRng};
use std::collections::HashMap;

fn generate_arm(content_id: &usize, arms: &HashMap<usize, Arm>, feature_dim: &usize) -> Arm {
    let arm: Arm;
    if !arms.contains_key(content_id) {
        arm = Arm::new(*feature_dim, *content_id, 0.0001, 0.0);
    }else{
        arm = arms.get(content_id).unwrap().clone();
    };
    arm
}

#[derive(Debug, Clone)]
struct Arm{
    content_id : usize,
    alpha : f64,
    norm_mean : Matrix<f64>,
    cov_matrix: Matrix<f64>,
    win_rate : f64,
    win : f64,
    lose : f64,
}

impl Arm {
    fn new(feature_dim: usize, content_id: usize, alpha: f64, win_rate: f64) -> Arm {
        Arm {
            content_id: content_id,
            norm_mean: Matrix::zeros(1, feature_dim),
            cov_matrix: Matrix::ones(1, feature_dim),
            alpha: alpha,
            win_rate: win_rate,
            win: 0.0,
            lose: 0.0,
        }
    }

    fn update(&mut self, features: Matrix<f64>, is_click: bool) {
        // error of diag
        self.cov_matrix += Matrix::new(self.cov_matrix.rows(),
                                        self.cov_matrix.cols(),
                                        (&features.transpose()*&features).diag());
        let cost_of_click = is_click as i8 as f64;

        let feat_mul_cost_vec = features.iter().map(|x| cost_of_click*x).collect::<Vec<f64>>();
        let feat_mul_cost = Matrix::new(features.rows(), features.cols(), feat_mul_cost_vec);
        self.norm_mean += feat_mul_cost;

        if is_click{
            self.win += 1.;
        } else{
            self.lose += 1.;
        }
        self.win_rate = &self.win/(&self.win + &self.lose);
    }

    fn predict(&mut self, features: &Matrix<f64>) -> f64 {
        let one_div_cov_vec = self.cov_matrix.iter().map(|x| 1.0/x).collect::<Vec<f64>>();
        let one_div_cov = Matrix::new(self.cov_matrix.rows(),
                                      self.cov_matrix.cols(), one_div_cov_vec);
        // Since the covariance matrix preserves only the diagonal components,
        // it suffices to take the inverse matrix
        let theta = &one_div_cov.elemul(&self.norm_mean);
        // Again, the inverse matrix of the covariance matrix
        // is computed by taking reciprocal
        let mut tmp: f64 = ((features.elemul(&one_div_cov))*features.transpose()).data()[0];
        tmp = &tmp.sqrt()*&self.alpha;
        (theta*features.transpose()).data()[0] + &tmp
    }

    fn print_result(&mut self) {
        println!("content_id:{}, total_num:{}, win_rate:{}",
                &self.content_id, &self.win+&self.lose, &self.win_rate);
    }
}

struct Viewer{
    gender: String,
    rng: ThreadRng,
}

impl Viewer {
    fn new(gender: String) -> Viewer{
        Viewer{
            gender: gender,
            rng: rand::thread_rng(),
        }
    }

    fn view(&mut self, content_id: &usize) -> bool{
        if &self.gender=="man" {
            // Men are easy to click on ads with id 5 or less
            if *content_id < 6 {
                return Some(self.rng.next_f32()).and_then(|n| if n>0.3 {Some(true)} else {Some(false)}).unwrap();
            } else{
                return Some(self.rng.next_f32()).and_then(|n| if n>0.7 {Some(true)} else {Some(false)}).unwrap();
            }
        } else {
            // Women are easy to click on ads with id 6 or higher
            if *content_id > 5{
                return Some(self.rng.next_f32()).and_then(|n| if n>0.3 {Some(true)} else {Some(false)}).unwrap();
            } else {
                return Some(self.rng.next_f32()).and_then(|n| if n>0.7 {Some(true)} else {Some(false)}).unwrap();
            }
        }
    }
}

struct Rulet{
    rng: ThreadRng,
}

impl Rulet{
    fn new() -> Rulet {
        Rulet{
            rng: rand::thread_rng(),
        }
    }

    fn generate_features(&mut self, viewer: &Viewer) -> Matrix<f64> {
        let features = Some(&viewer.gender)
            .and_then(|gender|
                if gender=="man" {Some(Matrix::new(1,2,vec![1.,0.]))}
                else {Some(Matrix::new(1,2,vec![0.,1.]))}
            ).unwrap();
        features
    }

    fn generate_content(&mut self) -> usize{
        self.rng.gen_range(0, 10)
    }

    fn generate_gender(&mut self) -> String {
        if self.rng.next_f32() > 0.5{
            return "man".to_string();
        } else{
            return "women".to_string();
        }
    }
}

fn main() {
    /*Context is for men and women only
    Men are easy to click on ads with id 5 or less
    Women are easy to click on ads with id 6 or higher
    */
    let feature_dim = 2;
    let num_of_views = 10000;
    let mut rulet = Rulet::new();

    let mut content_id: usize;
    let mut features: Matrix<f64>;
    let mut is_clicked: bool;
    let mut arms: HashMap<usize, Arm> = HashMap::new();

    for _ in 0..num_of_views {
        let mut viewer = Viewer::new(rulet.generate_gender());

        content_id = rulet.generate_content();
        features = rulet.generate_features(&viewer);
        is_clicked = viewer.view(&content_id);

        let mut arm = generate_arm(&content_id, &arms, &feature_dim);
        arm.update(features, is_clicked);
        arms.remove(&content_id);
        arms.insert(content_id, arm.clone());
    }

    let man_mat: Matrix<f64> = Matrix::new(1,2, vec![1.,0.]);
    let woman_mat: Matrix<f64> = Matrix::new(1,2, vec![0.,1.]);

    println!("print result======");
    for (_, arm) in arms.iter() {
        let mut arm = arm.clone();
        &arm.print_result();
        println!("Click rate when men browse: {}", &arm.predict(&man_mat) );
        println!("Click rate when women browse: {}", &arm.predict(&woman_mat) );
    }

}
<|endoftext|>


<commit_before>use winapi::um::winuser::{WS_VISIBLE, WS_DISABLED};
use crate::win32::{
    base_helper::check_hwnd,
    window_helper as wh,
    resources_helper as rh
};
use super::{ControlBase, ControlHandle};
use crate::{Bitmap, Icon, NwgError, RawEventHandler, unbind_raw_event_handler};
use std::cell::RefCell;

const NOT_BOUND: &'static str = "ImageFrame is not yet bound to a winapi object";
const BAD_HANDLE: &'static str = "INTERNAL ERROR: ImageFrame handle is not HWND!";

bitflags! {
    pub struct ImageFrameFlags: u32 {
        const VISIBLE = WS_VISIBLE;
        const DISABLED = WS_DISABLED;
    }
}

/**
An image frame is a control that displays a `Bitmap` or a `Icon` image resource. It can also triggers mouse clicks.

ImageFrame is not behind any features.

**Builder parameters:**
  * `parent`:           **Required.** The image frame parent container.
  * `size`:             The image frame size.
  * `position`:         The image frame position.
  * `flags`:            A combination of the ImageFrameFlags values.
  * `background_color`: The background color of the image frame. Used if the image is smaller than the control
  * `bitmap`:           A bitmap to display. If this value is set, icon is ignored.
  * `icon`:             An icon to display

**Control events:**
  * `OnImageFrameClick`: When the image frame is clicked once by the user
  * `OnImageFrameDoubleClick`: When the image frame is clicked twice rapidly by the user
  * `MousePress(_)`: Generic mouse press events on the button
  * `OnMouseMove`: Generic mouse mouse event
  * `OnMouseWheel`: Generic mouse wheel event

```rust
use native_windows_gui as nwg;
fn build_frame(button: &mut nwg::ImageFrame, window: &nwg::Window, ico: &nwg::Icon) {
    nwg::ImageFrame::builder()
        .parent(window)
        .build(button);
}
```
*/
#[derive(Default)]
pub struct ImageFrame {
    pub handle: ControlHandle,
    handler0: RefCell<Option<RawEventHandler>>,
}

impl ImageFrame {

    pub fn builder<'a>() -> ImageFrameBuilder<'a> {
        ImageFrameBuilder {
            size: (100, 100),
            position: (0, 0),
            flags: None,
            bitmap: None,
            icon: None,
            parent: None,
            background_color: None
        }
    }

    /// Sets the bitmap image of the image frame. Replace the current bitmap or icon.
    /// Set `image` to `None` to remove the image
    pub fn set_bitmap<'a>(&self, image: Option<&'a Bitmap>) {
        use winapi::um::winuser::{STM_SETIMAGE, IMAGE_BITMAP};
        use winapi::shared::minwindef::{WPARAM, LPARAM};

        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);

        let image_handle = image.map(|i| i.handle as LPARAM).unwrap_or(0);
        wh::send_message(handle, STM_SETIMAGE, IMAGE_BITMAP as WPARAM, image_handle);
    }

    /// Sets the bitmap image of the image frame. Replace the current bitmap or icon.
    /// Set `image` to `None` to remove the image
    pub fn set_icon<'a>(&self, image: Option<&'a Icon>) {
        use winapi::um::winuser::{STM_SETIMAGE, IMAGE_ICON};
        use winapi::shared::minwindef::{WPARAM, LPARAM};

        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);

        let image_handle = image.map(|i| i.handle as LPARAM).unwrap_or(0);
        wh::send_message(handle, STM_SETIMAGE, IMAGE_ICON as WPARAM, image_handle);
    }

    /// Returns the current image in the image frame.
    /// If the image frame has a bitmap, the value will be returned in `bitmap`
    /// If the image frame has a icon, the value will be returned in `icon`
    pub fn image<'a>(&self, bitmap: &mut Option<Bitmap>, icon: &mut Option<Icon>) {
        use winapi::um::winuser::{STM_GETIMAGE, IMAGE_BITMAP, IMAGE_ICON};
        use winapi::shared::minwindef::WPARAM;
        use winapi::shared::windef::HBITMAP;
        use winapi::um::winnt::HANDLE;

        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        let bitmap_handle = wh::send_message(handle, STM_GETIMAGE, IMAGE_BITMAP as WPARAM, 0);
        let icon_handle = wh::send_message(handle, STM_GETIMAGE, IMAGE_ICON as WPARAM, 0);

        *bitmap = None;
        *icon = None;

        if bitmap_handle != 0 && rh::is_bitmap(bitmap_handle as HBITMAP) {
            *bitmap = Some(Bitmap { handle: bitmap_handle as HANDLE, owned: false });
        } else if icon_handle != 0 {
            *icon = Some(Icon { handle: icon_handle as HANDLE, owned: false });
        }
    }

    /// Return true if the control user can interact with the control, return false otherwise
    pub fn enabled(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_enabled(handle) }
    }

    /// Enable or disable the control
    pub fn set_enabled(&self, v: bool) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_enabled(handle, v) }
    }

    /// Return true if the control is visible to the user. Will return true even if the
    /// control is outside of the parent client view (ex: at the position (10000, 10000))
    pub fn visible(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_visibility(handle) }
    }

    /// Show or hide the control to the user
    pub fn set_visible(&self, v: bool) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_visibility(handle, v) }
    }

    /// Return the size of the image frame in the parent window
    pub fn size(&self) -> (u32, u32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_size(handle) }
    }

    /// Set the size of the image frame in the parent window
    pub fn set_size(&self, x: u32, y: u32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_size(handle, x, y, false) }
    }

    /// Return the position of the image frame in the parent window
    pub fn position(&self) -> (i32, i32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_position(handle) }
    }

    /// Set the position of the image frame in the parent window
    pub fn set_position(&self, x: i32, y: i32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_position(handle, x, y) }
    }

    /// Winapi class name used during control creation
    pub fn class_name(&self) -> &'static str {
        "STATIC"
    }

    /// Winapi base flags used during window creation
    pub fn flags(&self) -> u32 {
        WS_VISIBLE
    }

    /// Winapi flags required by the control
    pub fn forced_flags(&self) -> u32 {
        use winapi::um::winuser::{SS_NOTIFY, WS_CHILD, SS_CENTERIMAGE};

        WS_CHILD | SS_NOTIFY | SS_CENTERIMAGE
    }

    /// Change the label background color to transparent.
    /// Change the checkbox background color.
    fn hook_background_color(&self, c: [u8; 3]) {
        use crate::bind_raw_event_handler_inner;
        use winapi::um::winuser::{WM_CTLCOLORSTATIC};
        use winapi::shared::{basetsd::UINT_PTR, windef::{HWND}, minwindef::LRESULT};
        use winapi::um::wingdi::{CreateSolidBrush, RGB};

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        let parent_handle = ControlHandle::Hwnd(wh::get_window_parent(handle));
        let brush = unsafe { CreateSolidBrush(RGB(c[0], c[1], c[2])) };

        let handler = bind_raw_event_handler_inner(&parent_handle, handle as UINT_PTR, move |_hwnd, msg, _w, l| {
            match msg {
                WM_CTLCOLORSTATIC => {
                    let child = l as HWND;
                    if child == handle {
                        return Some(brush as LRESULT);
                    }
                },
                _ => {}
            }

            None
        });

        *self.handler0.borrow_mut() = Some(handler.unwrap());
    }

}

impl Drop for ImageFrame {
    fn drop(&mut self) {
        let handler = self.handler0.borrow();
        if let Some(h) = handler.as_ref() {
            drop(unbind_raw_event_handler(h));
        }

        self.handle.destroy();
    }
}

pub struct ImageFrameBuilder<'a> {
    size: (i32, i32),
    position: (i32, i32),
    flags: Option<ImageFrameFlags>,
    bitmap: Option<&'a Bitmap>,
    icon: Option<&'a Icon>,
    parent: Option<ControlHandle>,
    background_color: Option<[u8; 3]>,
}

impl<'a> ImageFrameBuilder<'a> {

    pub fn flags(mut self, flags: ImageFrameFlags) -> ImageFrameBuilder<'a> {
        self.flags = Some(flags);
        self
    }

    pub fn size(mut self, size: (i32, i32)) -> ImageFrameBuilder<'a> {
        self.size = size;
        self
    }

    pub fn position(mut self, pos: (i32, i32)) -> ImageFrameBuilder<'a> {
        self.position = pos;
        self
    }

    pub fn bitmap(mut self, bit: Option<&'a Bitmap>) -> ImageFrameBuilder<'a> {
        self.bitmap = bit;
        self
    }

    pub fn icon(mut self, ico: Option<&'a Icon>) -> ImageFrameBuilder<'a> {
        self.icon = ico;
        self
    }

    pub fn parent<C: Into<ControlHandle>>(mut self, p: C) -> ImageFrameBuilder<'a> {
        self.parent = Some(p.into());
        self
    }

    pub fn background_color(mut self, color: Option<[u8;3]>) -> ImageFrameBuilder<'a> {
        self.background_color = color;
        self
    }

    pub fn build(self, out: &mut ImageFrame) -> Result<(), NwgError> {
        use winapi::um::winuser::{SS_BITMAP, SS_ICON};

        let mut flags = self.flags.map(|f| f.bits()).unwrap_or(out.flags());
        if self.icon.is_some() {
            flags |= SS_ICON;
        } else {
            flags |= SS_BITMAP;
        }

        let parent = match self.parent {
            Some(p) => Ok(p),
            None => Err(NwgError::no_parent("ImageFrame"))
        }?;

        out.handle = ControlBase::build_hwnd()
            .class_name(out.class_name())
            .forced_flags(out.forced_flags())
            .flags(flags)
            .size(self.size)
            .position(self.position)
            .parent(Some(parent))
            .build()?;

        if self.bitmap.is_some() {
            out.set_bitmap(self.bitmap);
        } else if self.icon.is_some() {
            out.set_icon(self.icon);
        }

        if self.background_color.is_some() {
            out.hook_background_color(self.background_color.unwrap());
        }

        Ok(())
    }

}

impl PartialEq for ImageFrame {
    fn eq(&self, other: &Self) -> bool {
        self.handle == other.handle
    }
}
<commit_msg>fixed memory leaks in image controls<commit_after>use winapi::um::winuser::{WS_VISIBLE, WS_DISABLED};
use winapi::um::wingdi::DeleteObject;
use crate::win32::{
    base_helper::check_hwnd,
    window_helper as wh,
    resources_helper as rh
};
use super::{ControlBase, ControlHandle};
use crate::{Bitmap, Icon, NwgError, RawEventHandler, unbind_raw_event_handler};
use std::cell::RefCell;

const NOT_BOUND: &'static str = "ImageFrame is not yet bound to a winapi object";
const BAD_HANDLE: &'static str = "INTERNAL ERROR: ImageFrame handle is not HWND!";

bitflags! {
    pub struct ImageFrameFlags: u32 {
        const VISIBLE = WS_VISIBLE;
        const DISABLED = WS_DISABLED;
    }
}

/**
An image frame is a control that displays a `Bitmap` or a `Icon` image resource. It can also triggers mouse clicks.

ImageFrame is not behind any features.

**Builder parameters:**
  * `parent`:           **Required.** The image frame parent container.
  * `size`:             The image frame size.
  * `position`:         The image frame position.
  * `flags`:            A combination of the ImageFrameFlags values.
  * `background_color`: The background color of the image frame. Used if the image is smaller than the control
  * `bitmap`:           A bitmap to display. If this value is set, icon is ignored.
  * `icon`:             An icon to display

**Control events:**
  * `OnImageFrameClick`: When the image frame is clicked once by the user
  * `OnImageFrameDoubleClick`: When the image frame is clicked twice rapidly by the user
  * `MousePress(_)`: Generic mouse press events on the button
  * `OnMouseMove`: Generic mouse mouse event
  * `OnMouseWheel`: Generic mouse wheel event

```rust
use native_windows_gui as nwg;
fn build_frame(button: &mut nwg::ImageFrame, window: &nwg::Window, ico: &nwg::Icon) {
    nwg::ImageFrame::builder()
        .parent(window)
        .build(button);
}
```
*/
#[derive(Default)]
pub struct ImageFrame {
    pub handle: ControlHandle,
    handler0: RefCell<Option<RawEventHandler>>,
}

impl ImageFrame {

    pub fn builder<'a>() -> ImageFrameBuilder<'a> {
        ImageFrameBuilder {
            size: (100, 100),
            position: (0, 0),
            flags: None,
            bitmap: None,
            icon: None,
            parent: None,
            background_color: None
        }
    }

    /// Sets the bitmap image of the image frame. Replace the current bitmap or icon.
    /// Set `image` to `None` to remove the image
    pub fn set_bitmap<'a>(&self, image: Option<&'a Bitmap>) {
        use winapi::um::winuser::{STM_SETIMAGE, IMAGE_BITMAP};
        use winapi::shared::minwindef::{WPARAM, LPARAM};

        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);

        let image_handle = image.map(|i| i.handle as LPARAM).unwrap_or(0);
        let prev_img = wh::send_message(handle, STM_SETIMAGE, IMAGE_BITMAP as WPARAM, image_handle);
        if prev_img != 0 {
            unsafe { DeleteObject(prev_img as _); }
        }
    }

    /// Sets the bitmap image of the image frame. Replace the current bitmap or icon.
    /// Set `image` to `None` to remove the image
    pub fn set_icon<'a>(&self, image: Option<&'a Icon>) {
        use winapi::um::winuser::{STM_SETIMAGE, IMAGE_ICON};
        use winapi::shared::minwindef::{WPARAM, LPARAM};

        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);

        let image_handle = image.map(|i| i.handle as LPARAM).unwrap_or(0);
        let prev_img = wh::send_message(handle, STM_SETIMAGE, IMAGE_ICON as WPARAM, image_handle);
        if prev_img != 0 {
            unsafe { DeleteObject(prev_img as _); }
        }
    }

    /// Returns the current image in the image frame.
    /// If the image frame has a bitmap, the value will be returned in `bitmap`
    /// If the image frame has a icon, the value will be returned in `icon`
    pub fn image<'a>(&self, bitmap: &mut Option<Bitmap>, icon: &mut Option<Icon>) {
        use winapi::um::winuser::{STM_GETIMAGE, IMAGE_BITMAP, IMAGE_ICON};
        use winapi::shared::minwindef::WPARAM;
        use winapi::shared::windef::HBITMAP;
        use winapi::um::winnt::HANDLE;

        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        let bitmap_handle = wh::send_message(handle, STM_GETIMAGE, IMAGE_BITMAP as WPARAM, 0);
        let icon_handle = wh::send_message(handle, STM_GETIMAGE, IMAGE_ICON as WPARAM, 0);

        *bitmap = None;
        *icon = None;

        if bitmap_handle != 0 && rh::is_bitmap(bitmap_handle as HBITMAP) {
            *bitmap = Some(Bitmap { handle: bitmap_handle as HANDLE, owned: false });
        } else if icon_handle != 0 {
            *icon = Some(Icon { handle: icon_handle as HANDLE, owned: false });
        }
    }

    /// Return true if the control user can interact with the control, return false otherwise
    pub fn enabled(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_enabled(handle) }
    }

    /// Enable or disable the control
    pub fn set_enabled(&self, v: bool) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_enabled(handle, v) }
    }

    /// Return true if the control is visible to the user. Will return true even if the
    /// control is outside of the parent client view (ex: at the position (10000, 10000))
    pub fn visible(&self) -> bool {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_visibility(handle) }
    }

    /// Show or hide the control to the user
    pub fn set_visible(&self, v: bool) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_visibility(handle, v) }
    }

    /// Return the size of the image frame in the parent window
    pub fn size(&self) -> (u32, u32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_size(handle) }
    }

    /// Set the size of the image frame in the parent window
    pub fn set_size(&self, x: u32, y: u32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_size(handle, x, y, false) }
    }

    /// Return the position of the image frame in the parent window
    pub fn position(&self) -> (i32, i32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::get_window_position(handle) }
    }

    /// Set the position of the image frame in the parent window
    pub fn set_position(&self, x: i32, y: i32) {
        let handle = check_hwnd(&self.handle, NOT_BOUND, BAD_HANDLE);
        unsafe { wh::set_window_position(handle, x, y) }
    }

    /// Winapi class name used during control creation
    pub fn class_name(&self) -> &'static str {
        "STATIC"
    }

    /// Winapi base flags used during window creation
    pub fn flags(&self) -> u32 {
        WS_VISIBLE
    }

    /// Winapi flags required by the control
    pub fn forced_flags(&self) -> u32 {
        use winapi::um::winuser::{SS_NOTIFY, WS_CHILD, SS_CENTERIMAGE};

        WS_CHILD | SS_NOTIFY | SS_CENTERIMAGE
    }

    /// Change the label background color to transparent.
    /// Change the checkbox background color.
    fn hook_background_color(&self, c: [u8; 3]) {
        use crate::bind_raw_event_handler_inner;
        use winapi::um::winuser::{WM_CTLCOLORSTATIC};
        use winapi::shared::{basetsd::UINT_PTR, windef::{HWND}, minwindef::LRESULT};
        use winapi::um::wingdi::{CreateSolidBrush, RGB};

        if self.handle.blank() { panic!(NOT_BOUND); }
        let handle = self.handle.hwnd().expect(BAD_HANDLE);

        let parent_handle = ControlHandle::Hwnd(wh::get_window_parent(handle));
        let brush = unsafe { CreateSolidBrush(RGB(c[0], c[1], c[2])) };

        let handler = bind_raw_event_handler_inner(&parent_handle, handle as UINT_PTR, move |_hwnd, msg, _w, l| {
            match msg {
                WM_CTLCOLORSTATIC => {
                    let child = l as HWND;
                    if child == handle {
                        return Some(brush as LRESULT);
                    }
                },
                _ => {}
            }

            None
        });

        *self.handler0.borrow_mut() = Some(handler.unwrap());
    }

}

impl Drop for ImageFrame {
    fn drop(&mut self) {
        let handler = self.handler0.borrow();
        if let Some(h) = handler.as_ref() {
            drop(unbind_raw_event_handler(h));
        }

        self.handle.destroy();
    }
}

pub struct ImageFrameBuilder<'a> {
    size: (i32, i32),
    position: (i32, i32),
    flags: Option<ImageFrameFlags>,
    bitmap: Option<&'a Bitmap>,
    icon: Option<&'a Icon>,
    parent: Option<ControlHandle>,
    background_color: Option<[u8; 3]>,
}

impl<'a> ImageFrameBuilder<'a> {

    pub fn flags(mut self, flags: ImageFrameFlags) -> ImageFrameBuilder<'a> {
        self.flags = Some(flags);
        self
    }

    pub fn size(mut self, size: (i32, i32)) -> ImageFrameBuilder<'a> {
        self.size = size;
        self
    }

    pub fn position(mut self, pos: (i32, i32)) -> ImageFrameBuilder<'a> {
        self.position = pos;
        self
    }

    pub fn bitmap(mut self, bit: Option<&'a Bitmap>) -> ImageFrameBuilder<'a> {
        self.bitmap = bit;
        self
    }

    pub fn icon(mut self, ico: Option<&'a Icon>) -> ImageFrameBuilder<'a> {
        self.icon = ico;
        self
    }

    pub fn parent<C: Into<ControlHandle>>(mut self, p: C) -> ImageFrameBuilder<'a> {
        self.parent = Some(p.into());
        self
    }

    pub fn background_color(mut self, color: Option<[u8;3]>) -> ImageFrameBuilder<'a> {
        self.background_color = color;
        self
    }

    pub fn build(self, out: &mut ImageFrame) -> Result<(), NwgError> {
        use winapi::um::winuser::{SS_BITMAP, SS_ICON};

        let mut flags = self.flags.map(|f| f.bits()).unwrap_or(out.flags());
        if self.icon.is_some() {
            flags |= SS_ICON;
        } else {
            flags |= SS_BITMAP;
        }

        let parent = match self.parent {
            Some(p) => Ok(p),
            None => Err(NwgError::no_parent("ImageFrame"))
        }?;

        out.handle = ControlBase::build_hwnd()
            .class_name(out.class_name())
            .forced_flags(out.forced_flags())
            .flags(flags)
            .size(self.size)
            .position(self.position)
            .parent(Some(parent))
            .build()?;

        if self.bitmap.is_some() {
            out.set_bitmap(self.bitmap);
        } else if self.icon.is_some() {
            out.set_icon(self.icon);
        }

        if self.background_color.is_some() {
            out.hook_background_color(self.background_color.unwrap());
        }

        Ok(())
    }

}

impl PartialEq for ImageFrame {
    fn eq(&self, other: &Self) -> bool {
        self.handle == other.handle
    }
}
<|endoftext|>


<commit_before>// Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/).

#![deny(warnings)]

use cargo_metadata;

use std::cmp::Ordering;
use std::collections::{BTreeMap, BTreeSet};
use std::env;
use std::fs;
use std::hash::{Hash, Hasher};
use std::io::{self, Write};
use std::iter::FromIterator;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::str;

use structopt::StructOpt;

#[derive(StructOpt, Debug)]
#[structopt(
    bin_name = "cargo fmt",
    author = "",
    about = "This utility formats all bin and lib files of \
             the current crate using rustfmt."
)]
pub struct Opts {
    /// No output printed to stdout
    #[structopt(short = "q", long = "quiet")]
    quiet: bool,

    /// Use verbose output
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,

    /// Print rustfmt version and exit
    #[structopt(long = "version")]
    version: bool,

    /// Specify package to format (only usable in workspaces)
    #[structopt(short = "p", long = "package", value_name = "package")]
    packages: Vec<String>,

    /// Options passed to rustfmt
    // 'raw = true' to make `--` explicit.
    #[structopt(name = "rustfmt_options", raw(raw = "true"))]
    rustfmt_options: Vec<String>,

    /// Format all packages (only usable in workspaces)
    #[structopt(long = "all")]
    format_all: bool,
}

fn main() {
    let exit_status = execute();
    std::io::stdout().flush().unwrap();
    std::process::exit(exit_status);
}

const SUCCESS: i32 = 0;
const FAILURE: i32 = 1;

fn execute() -> i32 {
    // Drop extra `fmt` argument provided by `cargo`.
    let mut found_fmt = false;
    let args = env::args().filter(|x| {
        if found_fmt {
            true
        } else {
            found_fmt = x == "fmt";
            x != "fmt"
        }
    });

    let opts = Opts::from_iter(args);

    let verbosity = match (opts.verbose, opts.quiet) {
        (false, false) => Verbosity::Normal,
        (false, true) => Verbosity::Quiet,
        (true, false) => Verbosity::Verbose,
        (true, true) => {
            print_usage_to_stderr("quiet mode and verbose mode are not compatible");
            return FAILURE;
        }
    };

    if opts.version {
        return handle_command_status(get_version());
    }

    let strategy = CargoFmtStrategy::from_opts(&opts);

    handle_command_status(format_crate(verbosity, &strategy, opts.rustfmt_options))
}

fn print_usage_to_stderr(reason: &str) {
    eprintln!("{}", reason);
    let app = Opts::clap();
    app.write_help(&mut io::stderr())
        .expect("failed to write to stderr");
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Verbosity {
    Verbose,
    Normal,
    Quiet,
}

fn handle_command_status(status: Result<i32, io::Error>) -> i32 {
    match status {
        Err(e) => {
            print_usage_to_stderr(&e.to_string());
            FAILURE
        }
        Ok(status) => status,
    }
}

fn get_version() -> Result<i32, io::Error> {
    let mut command = Command::new("rustfmt")
        .stdout(std::process::Stdio::inherit())
        .args(&[String::from("--version")])
        .spawn()
        .map_err(|e| match e.kind() {
            io::ErrorKind::NotFound => io::Error::new(
                io::ErrorKind::Other,
                "Could not run rustfmt, please make sure it is in your PATH.",
            ),
            _ => e,
        })?;
    let result = command.wait()?;
    if result.success() {
        Ok(SUCCESS)
    } else {
        Ok(result.code().unwrap_or(SUCCESS))
    }
}

fn format_crate(
    verbosity: Verbosity,
    strategy: &CargoFmtStrategy,
    rustfmt_args: Vec<String>,
) -> Result<i32, io::Error> {
    let targets = if rustfmt_args
        .iter()
        .any(|s| ["--print-config", "-h", "--help", "-V", "--version"].contains(&s.as_str()))
    {
        BTreeSet::new()
    } else {
        get_targets(strategy)?
    };

    // Currently only bin and lib files get formatted.
    run_rustfmt(&targets, &rustfmt_args, verbosity)
}

/// Target uses a `path` field for equality and hashing.
#[derive(Debug)]
pub struct Target {
    /// A path to the main source file of the target.
    path: PathBuf,
    /// A kind of target (e.g., lib, bin, example, ...).
    kind: String,
    /// Rust edition for this target.
    edition: String,
}

impl Target {
    pub fn from_target(target: &cargo_metadata::Target) -> Self {
        let path = PathBuf::from(&target.src_path);
        let canonicalized = fs::canonicalize(&path).unwrap_or(path);

        Target {
            path: canonicalized,
            kind: target.kind[0].clone(),
            edition: target.edition.clone(),
        }
    }
}

impl PartialEq for Target {
    fn eq(&self, other: &Target) -> bool {
        self.path == other.path
    }
}

impl PartialOrd for Target {
    fn partial_cmp(&self, other: &Target) -> Option<Ordering> {
        Some(self.path.cmp(&other.path))
    }
}

impl Ord for Target {
    fn cmp(&self, other: &Target) -> Ordering {
        self.path.cmp(&other.path)
    }
}

impl Eq for Target {}

impl Hash for Target {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.path.hash(state);
    }
}

#[derive(Debug, PartialEq, Eq)]
pub enum CargoFmtStrategy {
    /// Format every packages and dependencies.
    All,
    /// Format packages that are specified by the command line argument.
    Some(Vec<String>),
    /// Format the root packages only.
    Root,
}

impl CargoFmtStrategy {
    pub fn from_opts(opts: &Opts) -> CargoFmtStrategy {
        match (opts.format_all, opts.packages.is_empty()) {
            (false, true) => CargoFmtStrategy::Root,
            (true, _) => CargoFmtStrategy::All,
            (false, false) => CargoFmtStrategy::Some(opts.packages.clone()),
        }
    }
}

/// Based on the specified `CargoFmtStrategy`, returns a set of main source files.
fn get_targets(strategy: &CargoFmtStrategy) -> Result<BTreeSet<Target>, io::Error> {
    let mut targets = BTreeSet::new();

    match *strategy {
        CargoFmtStrategy::Root => get_targets_root_only(&mut targets)?,
        CargoFmtStrategy::All => get_targets_recursive(None, &mut targets, &mut BTreeSet::new())?,
        CargoFmtStrategy::Some(ref hitlist) => get_targets_with_hitlist(hitlist, &mut targets)?,
    }

    if targets.is_empty() {
        Err(io::Error::new(
            io::ErrorKind::Other,
            "Failed to find targets".to_owned(),
        ))
    } else {
        Ok(targets)
    }
}

fn get_targets_root_only(targets: &mut BTreeSet<Target>) -> Result<(), io::Error> {
    let metadata = get_cargo_metadata(None)?;
    let current_dir = env::current_dir()?.canonicalize()?;
    let current_dir_manifest = current_dir.join("Cargo.toml");
    let workspace_root_path = PathBuf::from(&metadata.workspace_root).canonicalize()?;
    let in_workspace_root = workspace_root_path == current_dir;

    for package in metadata.packages {
        if in_workspace_root || PathBuf::from(&package.manifest_path) == current_dir_manifest {
            for target in package.targets {
                targets.insert(Target::from_target(&target));
            }
        }
    }

    Ok(())
}

fn get_targets_recursive(
    manifest_path: Option<&Path>,
    mut targets: &mut BTreeSet<Target>,
    visited: &mut BTreeSet<String>,
) -> Result<(), io::Error> {
    let metadata = get_cargo_metadata(manifest_path)?;

    for package in metadata.packages {
        add_targets(&package.targets, &mut targets);

        // Look for local dependencies.
        for dependency in package.dependencies {
            if dependency.source.is_some() || visited.contains(&dependency.name) {
                continue;
            }

            let mut manifest_path = PathBuf::from(&package.manifest_path);

            manifest_path.pop();
            manifest_path.push(&dependency.name);
            manifest_path.push("Cargo.toml");

            if manifest_path.exists() {
                visited.insert(dependency.name);
                get_targets_recursive(Some(&manifest_path), &mut targets, visited)?;
            }
        }
    }

    Ok(())
}

fn get_targets_with_hitlist(
    hitlist: &[String],
    targets: &mut BTreeSet<Target>,
) -> Result<(), io::Error> {
    let metadata = get_cargo_metadata(None)?;

    let mut workspace_hitlist: BTreeSet<&String> = BTreeSet::from_iter(hitlist);

    for package in metadata.packages {
        if workspace_hitlist.remove(&package.name) {
            for target in package.targets {
                targets.insert(Target::from_target(&target));
            }
        }
    }

    if workspace_hitlist.is_empty() {
        Ok(())
    } else {
        let package = workspace_hitlist.iter().next().unwrap();
        Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("package `{}` is not a member of the workspace", package),
        ))
    }
}

fn add_targets(target_paths: &[cargo_metadata::Target], targets: &mut BTreeSet<Target>) {
    for target in target_paths {
        targets.insert(Target::from_target(target));
    }
}

fn run_rustfmt(
    targets: &BTreeSet<Target>,
    fmt_args: &[String],
    verbosity: Verbosity,
) -> Result<i32, io::Error> {
    let by_edition = targets
        .iter()
        .inspect(|t| {
            if verbosity == Verbosity::Verbose {
                println!("[{} ({})] {:?}", t.kind, t.edition, t.path)
            }
        })
        .fold(BTreeMap::new(), |mut h, t| {
            h.entry(&t.edition).or_insert_with(Vec::new).push(&t.path);
            h
        });

    let mut status = vec![];
    for (edition, files) in by_edition {
        let stdout = if verbosity == Verbosity::Quiet {
            std::process::Stdio::null()
        } else {
            std::process::Stdio::inherit()
        };

        if verbosity == Verbosity::Verbose {
            print!("rustfmt");
            print!(" --edition {}", edition);
            fmt_args.iter().for_each(|f| print!(" {}", f));
            files.iter().for_each(|f| print!(" {}", f.display()));
            println!();
        }

        let mut command = Command::new("rustfmt")
            .stdout(stdout)
            .args(files)
            .args(&["--edition", edition])
            .args(fmt_args)
            .spawn()
            .map_err(|e| match e.kind() {
                io::ErrorKind::NotFound => io::Error::new(
                    io::ErrorKind::Other,
                    "Could not run rustfmt, please make sure it is in your PATH.",
                ),
                _ => e,
            })?;

        status.push(command.wait()?);
    }

    Ok(status
        .iter()
        .filter_map(|s| if s.success() { None } else { s.code() })
        .next()
        .unwrap_or(SUCCESS))
}

fn get_cargo_metadata(manifest_path: Option<&Path>) -> Result<cargo_metadata::Metadata, io::Error> {
    let mut cmd = cargo_metadata::MetadataCommand::new();
    cmd.no_deps();
    if let Some(manifest_path) = manifest_path {
        cmd.manifest_path(manifest_path);
    }
    match cmd.exec() {
        Ok(metadata) => Ok(metadata),
        Err(error) => Err(io::Error::new(io::ErrorKind::Other, error.to_string())),
    }
}

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

    #[test]
    fn default_options() {
        let empty: Vec<String> = vec![];
        let o = Opts::from_iter(&empty);
        assert_eq!(false, o.quiet);
        assert_eq!(false, o.verbose);
        assert_eq!(false, o.version);
        assert_eq!(empty, o.packages);
        assert_eq!(empty, o.rustfmt_options);
        assert_eq!(false, o.format_all);
    }

    #[test]
    fn good_options() {
        let o = Opts::from_iter(&[
            "test",
            "-q",
            "-p",
            "p1",
            "-p",
            "p2",
            "--",
            "--edition",
            "2018",
        ]);
        assert_eq!(true, o.quiet);
        assert_eq!(false, o.verbose);
        assert_eq!(false, o.version);
        assert_eq!(vec!["p1", "p2"], o.packages);
        assert_eq!(vec!["--edition", "2018"], o.rustfmt_options);
        assert_eq!(false, o.format_all);
    }

    #[test]
    fn unexpected_option() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "unexpected"])
                .is_err()
        );
    }

    #[test]
    fn unexpected_flag() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "--flag"])
                .is_err()
        );
    }

    #[test]
    fn mandatory_separator() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "--check"])
                .is_err()
        );
        assert!(
            !Opts::clap()
                .get_matches_from_safe(&["test", "--", "--check"])
                .is_err()
        );
    }

    #[test]
    fn multiple_packages_one_by_one() {
        let o = Opts::from_iter(&[
            "test",
            "-p",
            "package1",
            "--package",
            "package2",
            "-p",
            "package3",
        ]);
        assert_eq!(3, o.packages.len());
    }

    #[test]
    fn multiple_packages_grouped() {
        let o = Opts::from_iter(&[
            "test",
            "--package",
            "package1",
            "package2",
            "-p",
            "package3",
            "package4",
        ]);
        assert_eq!(4, o.packages.len());
    }

    #[test]
    fn empty_packages_1() {
        assert!(Opts::clap().get_matches_from_safe(&["test", "-p"]).is_err());
    }

    #[test]
    fn empty_packages_2() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "-p", "--", "--check"])
                .is_err()
        );
    }

    #[test]
    fn empty_packages_3() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "-p", "--verbose"])
                .is_err()
        );
    }

    #[test]
    fn empty_packages_4() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "-p", "--check"])
                .is_err()
        );
    }
}
<commit_msg>Add newline after help text (#3650)<commit_after>// Inspired by Paul Woolcock's cargo-fmt (https://github.com/pwoolcoc/cargo-fmt/).

#![deny(warnings)]

use cargo_metadata;

use std::cmp::Ordering;
use std::collections::{BTreeMap, BTreeSet};
use std::env;
use std::fs;
use std::hash::{Hash, Hasher};
use std::io::{self, Write};
use std::iter::FromIterator;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::str;

use structopt::StructOpt;

#[derive(StructOpt, Debug)]
#[structopt(
    bin_name = "cargo fmt",
    author = "",
    about = "This utility formats all bin and lib files of \
             the current crate using rustfmt."
)]
pub struct Opts {
    /// No output printed to stdout
    #[structopt(short = "q", long = "quiet")]
    quiet: bool,

    /// Use verbose output
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,

    /// Print rustfmt version and exit
    #[structopt(long = "version")]
    version: bool,

    /// Specify package to format (only usable in workspaces)
    #[structopt(short = "p", long = "package", value_name = "package")]
    packages: Vec<String>,

    /// Options passed to rustfmt
    // 'raw = true' to make `--` explicit.
    #[structopt(name = "rustfmt_options", raw(raw = "true"))]
    rustfmt_options: Vec<String>,

    /// Format all packages (only usable in workspaces)
    #[structopt(long = "all")]
    format_all: bool,
}

fn main() {
    let exit_status = execute();
    std::io::stdout().flush().unwrap();
    std::process::exit(exit_status);
}

const SUCCESS: i32 = 0;
const FAILURE: i32 = 1;

fn execute() -> i32 {
    // Drop extra `fmt` argument provided by `cargo`.
    let mut found_fmt = false;
    let args = env::args().filter(|x| {
        if found_fmt {
            true
        } else {
            found_fmt = x == "fmt";
            x != "fmt"
        }
    });

    let opts = Opts::from_iter(args);

    let verbosity = match (opts.verbose, opts.quiet) {
        (false, false) => Verbosity::Normal,
        (false, true) => Verbosity::Quiet,
        (true, false) => Verbosity::Verbose,
        (true, true) => {
            print_usage_to_stderr("quiet mode and verbose mode are not compatible");
            return FAILURE;
        }
    };

    if opts.version {
        return handle_command_status(get_version());
    }

    let strategy = CargoFmtStrategy::from_opts(&opts);

    handle_command_status(format_crate(verbosity, &strategy, opts.rustfmt_options))
}

fn print_usage_to_stderr(reason: &str) {
    eprintln!("{}", reason);
    let app = Opts::clap();
    app.after_help("")
        .write_help(&mut io::stderr())
        .expect("failed to write to stderr");
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Verbosity {
    Verbose,
    Normal,
    Quiet,
}

fn handle_command_status(status: Result<i32, io::Error>) -> i32 {
    match status {
        Err(e) => {
            print_usage_to_stderr(&e.to_string());
            FAILURE
        }
        Ok(status) => status,
    }
}

fn get_version() -> Result<i32, io::Error> {
    let mut command = Command::new("rustfmt")
        .stdout(std::process::Stdio::inherit())
        .args(&[String::from("--version")])
        .spawn()
        .map_err(|e| match e.kind() {
            io::ErrorKind::NotFound => io::Error::new(
                io::ErrorKind::Other,
                "Could not run rustfmt, please make sure it is in your PATH.",
            ),
            _ => e,
        })?;
    let result = command.wait()?;
    if result.success() {
        Ok(SUCCESS)
    } else {
        Ok(result.code().unwrap_or(SUCCESS))
    }
}

fn format_crate(
    verbosity: Verbosity,
    strategy: &CargoFmtStrategy,
    rustfmt_args: Vec<String>,
) -> Result<i32, io::Error> {
    let targets = if rustfmt_args
        .iter()
        .any(|s| ["--print-config", "-h", "--help", "-V", "--version"].contains(&s.as_str()))
    {
        BTreeSet::new()
    } else {
        get_targets(strategy)?
    };

    // Currently only bin and lib files get formatted.
    run_rustfmt(&targets, &rustfmt_args, verbosity)
}

/// Target uses a `path` field for equality and hashing.
#[derive(Debug)]
pub struct Target {
    /// A path to the main source file of the target.
    path: PathBuf,
    /// A kind of target (e.g., lib, bin, example, ...).
    kind: String,
    /// Rust edition for this target.
    edition: String,
}

impl Target {
    pub fn from_target(target: &cargo_metadata::Target) -> Self {
        let path = PathBuf::from(&target.src_path);
        let canonicalized = fs::canonicalize(&path).unwrap_or(path);

        Target {
            path: canonicalized,
            kind: target.kind[0].clone(),
            edition: target.edition.clone(),
        }
    }
}

impl PartialEq for Target {
    fn eq(&self, other: &Target) -> bool {
        self.path == other.path
    }
}

impl PartialOrd for Target {
    fn partial_cmp(&self, other: &Target) -> Option<Ordering> {
        Some(self.path.cmp(&other.path))
    }
}

impl Ord for Target {
    fn cmp(&self, other: &Target) -> Ordering {
        self.path.cmp(&other.path)
    }
}

impl Eq for Target {}

impl Hash for Target {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.path.hash(state);
    }
}

#[derive(Debug, PartialEq, Eq)]
pub enum CargoFmtStrategy {
    /// Format every packages and dependencies.
    All,
    /// Format packages that are specified by the command line argument.
    Some(Vec<String>),
    /// Format the root packages only.
    Root,
}

impl CargoFmtStrategy {
    pub fn from_opts(opts: &Opts) -> CargoFmtStrategy {
        match (opts.format_all, opts.packages.is_empty()) {
            (false, true) => CargoFmtStrategy::Root,
            (true, _) => CargoFmtStrategy::All,
            (false, false) => CargoFmtStrategy::Some(opts.packages.clone()),
        }
    }
}

/// Based on the specified `CargoFmtStrategy`, returns a set of main source files.
fn get_targets(strategy: &CargoFmtStrategy) -> Result<BTreeSet<Target>, io::Error> {
    let mut targets = BTreeSet::new();

    match *strategy {
        CargoFmtStrategy::Root => get_targets_root_only(&mut targets)?,
        CargoFmtStrategy::All => get_targets_recursive(None, &mut targets, &mut BTreeSet::new())?,
        CargoFmtStrategy::Some(ref hitlist) => get_targets_with_hitlist(hitlist, &mut targets)?,
    }

    if targets.is_empty() {
        Err(io::Error::new(
            io::ErrorKind::Other,
            "Failed to find targets".to_owned(),
        ))
    } else {
        Ok(targets)
    }
}

fn get_targets_root_only(targets: &mut BTreeSet<Target>) -> Result<(), io::Error> {
    let metadata = get_cargo_metadata(None)?;
    let current_dir = env::current_dir()?.canonicalize()?;
    let current_dir_manifest = current_dir.join("Cargo.toml");
    let workspace_root_path = PathBuf::from(&metadata.workspace_root).canonicalize()?;
    let in_workspace_root = workspace_root_path == current_dir;

    for package in metadata.packages {
        if in_workspace_root || PathBuf::from(&package.manifest_path) == current_dir_manifest {
            for target in package.targets {
                targets.insert(Target::from_target(&target));
            }
        }
    }

    Ok(())
}

fn get_targets_recursive(
    manifest_path: Option<&Path>,
    mut targets: &mut BTreeSet<Target>,
    visited: &mut BTreeSet<String>,
) -> Result<(), io::Error> {
    let metadata = get_cargo_metadata(manifest_path)?;

    for package in metadata.packages {
        add_targets(&package.targets, &mut targets);

        // Look for local dependencies.
        for dependency in package.dependencies {
            if dependency.source.is_some() || visited.contains(&dependency.name) {
                continue;
            }

            let mut manifest_path = PathBuf::from(&package.manifest_path);

            manifest_path.pop();
            manifest_path.push(&dependency.name);
            manifest_path.push("Cargo.toml");

            if manifest_path.exists() {
                visited.insert(dependency.name);
                get_targets_recursive(Some(&manifest_path), &mut targets, visited)?;
            }
        }
    }

    Ok(())
}

fn get_targets_with_hitlist(
    hitlist: &[String],
    targets: &mut BTreeSet<Target>,
) -> Result<(), io::Error> {
    let metadata = get_cargo_metadata(None)?;

    let mut workspace_hitlist: BTreeSet<&String> = BTreeSet::from_iter(hitlist);

    for package in metadata.packages {
        if workspace_hitlist.remove(&package.name) {
            for target in package.targets {
                targets.insert(Target::from_target(&target));
            }
        }
    }

    if workspace_hitlist.is_empty() {
        Ok(())
    } else {
        let package = workspace_hitlist.iter().next().unwrap();
        Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!("package `{}` is not a member of the workspace", package),
        ))
    }
}

fn add_targets(target_paths: &[cargo_metadata::Target], targets: &mut BTreeSet<Target>) {
    for target in target_paths {
        targets.insert(Target::from_target(target));
    }
}

fn run_rustfmt(
    targets: &BTreeSet<Target>,
    fmt_args: &[String],
    verbosity: Verbosity,
) -> Result<i32, io::Error> {
    let by_edition = targets
        .iter()
        .inspect(|t| {
            if verbosity == Verbosity::Verbose {
                println!("[{} ({})] {:?}", t.kind, t.edition, t.path)
            }
        })
        .fold(BTreeMap::new(), |mut h, t| {
            h.entry(&t.edition).or_insert_with(Vec::new).push(&t.path);
            h
        });

    let mut status = vec![];
    for (edition, files) in by_edition {
        let stdout = if verbosity == Verbosity::Quiet {
            std::process::Stdio::null()
        } else {
            std::process::Stdio::inherit()
        };

        if verbosity == Verbosity::Verbose {
            print!("rustfmt");
            print!(" --edition {}", edition);
            fmt_args.iter().for_each(|f| print!(" {}", f));
            files.iter().for_each(|f| print!(" {}", f.display()));
            println!();
        }

        let mut command = Command::new("rustfmt")
            .stdout(stdout)
            .args(files)
            .args(&["--edition", edition])
            .args(fmt_args)
            .spawn()
            .map_err(|e| match e.kind() {
                io::ErrorKind::NotFound => io::Error::new(
                    io::ErrorKind::Other,
                    "Could not run rustfmt, please make sure it is in your PATH.",
                ),
                _ => e,
            })?;

        status.push(command.wait()?);
    }

    Ok(status
        .iter()
        .filter_map(|s| if s.success() { None } else { s.code() })
        .next()
        .unwrap_or(SUCCESS))
}

fn get_cargo_metadata(manifest_path: Option<&Path>) -> Result<cargo_metadata::Metadata, io::Error> {
    let mut cmd = cargo_metadata::MetadataCommand::new();
    cmd.no_deps();
    if let Some(manifest_path) = manifest_path {
        cmd.manifest_path(manifest_path);
    }
    match cmd.exec() {
        Ok(metadata) => Ok(metadata),
        Err(error) => Err(io::Error::new(io::ErrorKind::Other, error.to_string())),
    }
}

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

    #[test]
    fn default_options() {
        let empty: Vec<String> = vec![];
        let o = Opts::from_iter(&empty);
        assert_eq!(false, o.quiet);
        assert_eq!(false, o.verbose);
        assert_eq!(false, o.version);
        assert_eq!(empty, o.packages);
        assert_eq!(empty, o.rustfmt_options);
        assert_eq!(false, o.format_all);
    }

    #[test]
    fn good_options() {
        let o = Opts::from_iter(&[
            "test",
            "-q",
            "-p",
            "p1",
            "-p",
            "p2",
            "--",
            "--edition",
            "2018",
        ]);
        assert_eq!(true, o.quiet);
        assert_eq!(false, o.verbose);
        assert_eq!(false, o.version);
        assert_eq!(vec!["p1", "p2"], o.packages);
        assert_eq!(vec!["--edition", "2018"], o.rustfmt_options);
        assert_eq!(false, o.format_all);
    }

    #[test]
    fn unexpected_option() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "unexpected"])
                .is_err()
        );
    }

    #[test]
    fn unexpected_flag() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "--flag"])
                .is_err()
        );
    }

    #[test]
    fn mandatory_separator() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "--check"])
                .is_err()
        );
        assert!(
            !Opts::clap()
                .get_matches_from_safe(&["test", "--", "--check"])
                .is_err()
        );
    }

    #[test]
    fn multiple_packages_one_by_one() {
        let o = Opts::from_iter(&[
            "test",
            "-p",
            "package1",
            "--package",
            "package2",
            "-p",
            "package3",
        ]);
        assert_eq!(3, o.packages.len());
    }

    #[test]
    fn multiple_packages_grouped() {
        let o = Opts::from_iter(&[
            "test",
            "--package",
            "package1",
            "package2",
            "-p",
            "package3",
            "package4",
        ]);
        assert_eq!(4, o.packages.len());
    }

    #[test]
    fn empty_packages_1() {
        assert!(Opts::clap().get_matches_from_safe(&["test", "-p"]).is_err());
    }

    #[test]
    fn empty_packages_2() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "-p", "--", "--check"])
                .is_err()
        );
    }

    #[test]
    fn empty_packages_3() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "-p", "--verbose"])
                .is_err()
        );
    }

    #[test]
    fn empty_packages_4() {
        assert!(
            Opts::clap()
                .get_matches_from_safe(&["test", "-p", "--check"])
                .is_err()
        );
    }
}
<|endoftext|>