<commit_before>// Connection Pool-esque construct used to queue up and asynchronously execute
// Postgres queries ane optionally evaluate callbacks based on the results.

use std::collections::VecDeque;
use std::thread;
use std::sync::{Arc, Mutex};

use postgres;
use futures::stream::{Stream, channel, Sender, Receiver};
use futures::Future;

use transport::postgres::get_client;

type QueryError = postgres::error::Error;
type SenderQueue = Arc<Mutex<VecDeque<Sender<(String, Sender<(), ()>), ()>>>>;
type QueryQueue = Arc<Mutex<VecDeque<String>>>;

pub struct QueryServer {
    conn_count: usize, // how many connections to open
    query_queue: QueryQueue, // internal query queue
    conn_queue: SenderQueue, // Database connection objects
}

// locks the QueryQueue and returns a queued query, if there are any.
fn try_get_new_query(query_queue: QueryQueue) -> Option<String> {
    let mut qq_inner = query_queue.lock().unwrap();
    // there is a queued query
    if !qq_inner.is_empty() {
        return Some(qq_inner.pop_front().unwrap())
    }else{
        // No new queries
        return None
    }
}

// executes the query and blocks the calling thread until it completes
#[allow(unused_must_use)]
fn execute_query(query: String, client: &postgres::Connection) {
    client.execute(query.as_str(), &[])
        .map_err(|err| println!("Error saving tick: {:?}", err) );
}

// Creates a query processor that awaits requests
fn init_query_processor(rx: Receiver<(String, Sender<(), ()>), ()>, query_queue: QueryQueue) {
    // get a connection to the postgres database
    let client = get_client().expect("Couldn't create postgres connection.");
    // Handler for new queries from main thread
    // This blocks the worker thread until a new message is received
    // .wait() consumes the stream immediately, so the main thread has to wait
    // for the worker to push a message saying it's done before sending more messages
    for tup in rx.wait() {
        let (query, done_tx) = tup.unwrap();
        execute_query(query, &client);
        // keep trying to get queued queries to exeucte until the queue is empty
        while let Some(new_query) = try_get_new_query(query_queue.clone()) {
            execute_query(new_query, &client);
        }
        // Let the main thread know it's safe to use the sender again
        // This essentially indicates that the worker thread is idle
        done_tx.send(Ok(()));
    }
}

impl QueryServer {
    pub fn new(conn_count: usize) -> QueryServer {
        let mut conn_queue = VecDeque::with_capacity(conn_count);
        let query_queue = Arc::new(Mutex::new(VecDeque::new()));
        for _ in 0..conn_count {
            // channel for getting the Sender back from the worker thread
            let (tx, rx) = channel::<(String, Sender<(), ()>), ()>();
            let qq_copy = query_queue.clone();
            thread::spawn(move || { init_query_processor(rx, qq_copy) });
            // store the sender which can be used to send queries
            // to the worker in the connection queue
            conn_queue.push_back(tx);
        }

        QueryServer {
            conn_count: conn_count,
            query_queue: query_queue,
            conn_queue: Arc::new(Mutex::new(conn_queue))
        }
    }

    // queues up a query to execute that doesn't return a result
    pub fn execute(&mut self, query: String) {
        // no connections available
        let temp_lock_res = self.conn_queue.lock().unwrap().is_empty();
        // Force the guard locking conn_queue to go out of scope
        // this prevents the lock from being held through the entire if/else
        let copy_res = temp_lock_res.clone();
        if copy_res {
            // push query to the query queue
            self.query_queue.lock().unwrap().push_back(query);
        }else{
            let tx = self.conn_queue.lock().unwrap().pop_front().unwrap();
            let cq_clone = self.conn_queue.clone();
            // channel for notifying main thread when query is done and worker is idle
            let (tx_tx, tx_rx) = channel::<(), ()>();
            tx.send(Ok((query, tx_tx))).and_then(|new_tx| {
                // Wait until the worker thread signals that it is idle
                tx_rx.take(1).into_future().and_then(move |_| {
                    // Put the Sender for the newly idle worker into the connection queue
                    cq_clone.lock().unwrap().push_back(new_tx);
                    Ok(())
                }).forget();
                Ok(())
            }).forget();
        }
    }
}
<commit_msg>Replaced second channel in QueryServer with a Oneshot<commit_after>// Connection Pool-esque construct used to queue up
// and asynchronously execute Postgres queries

use std::collections::VecDeque;
use std::thread;
use std::time::Duration;
use std::sync::{Arc, Mutex};

use postgres;
use futures::stream::{Stream, channel, Sender, Receiver};
use futures::{Future, oneshot, Complete};

use transport::postgres::get_client;

type QueryError = postgres::error::Error;
type SenderQueue = Arc<Mutex<VecDeque<Sender<(String, Complete<()>), ()>>>>;
type QueryQueue = Arc<Mutex<VecDeque<String>>>;

pub struct QueryServer {
    conn_count: usize, // how many connections to open
    query_queue: QueryQueue, // internal query queue
    conn_queue: SenderQueue, // senders for idle query threads
}

// locks the QueryQueue and returns a queued query, if there are any.
fn try_get_new_query(query_queue: QueryQueue) -> Option<String> {
    let mut qq_inner = query_queue.lock().unwrap();
    // there is a queued query
    if !qq_inner.is_empty() {
        return Some(qq_inner.pop_front().unwrap())
    }else{
        // No new queries
        return None
    }
}

// executes the query and blocks the calling thread until it completes
#[allow(unused_must_use)]
fn execute_query(query: String, client: &postgres::Connection) {
    client.execute(query.as_str(), &[])
        .map_err(|err| println!("Error saving tick: {:?}", err) );
    thread::sleep(Duration::new(5, 0));
}

// Creates a query processor that awaits requests
fn init_query_processor(rx: Receiver<(String, Complete<()>), ()>, query_queue: QueryQueue) {
    // get a connection to the postgres database
    let client = get_client().expect("Couldn't create postgres connection.");
    // Handler for new queries from main thread
    // This blocks the worker thread until a new message is received
    // .wait() consumes the stream immediately, so the main thread has to wait
    // for the worker to push a message saying it's done before sending more messages
    for tup in rx.wait() {
        let (query, done_tx) = tup.unwrap();
        execute_query(query, &client);
        // keep trying to get queued queries to exeucte until the queue is empty
        while let Some(new_query) = try_get_new_query(query_queue.clone()) {
            execute_query(new_query, &client);
        }
        // Let the main thread know it's safe to use the sender again
        // This essentially indicates that the worker thread is idle
        done_tx.complete(());
    }
}

impl QueryServer {
    pub fn new(conn_count: usize) -> QueryServer {
        let mut conn_queue = VecDeque::with_capacity(conn_count);
        let query_queue = Arc::new(Mutex::new(VecDeque::new()));
        for _ in 0..conn_count {
            // channel for getting the Sender back from the worker thread
            let (tx, rx) = channel::<(String, Complete<()>), ()>();
            let qq_copy = query_queue.clone();
            thread::spawn(move || { init_query_processor(rx, qq_copy) });
            // store the sender which can be used to send queries
            // to the worker in the connection queue
            conn_queue.push_back(tx);
        }

        QueryServer {
            conn_count: conn_count,
            query_queue: query_queue,
            conn_queue: Arc::new(Mutex::new(conn_queue))
        }
    }

    // queues up a query to execute that doesn't return a result
    pub fn execute(&mut self, query: String) {
        // no connections available
        let temp_lock_res = self.conn_queue.lock().unwrap().is_empty();
        // Force the guard locking conn_queue to go out of scope
        // this prevents the lock from being held through the entire if/else
        let copy_res = temp_lock_res.clone();
        if copy_res {
            // push query to the query queue
            self.query_queue.lock().unwrap().push_back(query);
        }else{
            let tx = self.conn_queue.lock().unwrap().pop_front().unwrap();
            let cq_clone = self.conn_queue.clone();
            // future for notifying main thread when query is done and worker is idle
            let (c, o) = oneshot::<()>();
            tx.send(Ok((query, c))).and_then(|new_tx| {
                // Wait until the worker thread signals that it is idle
                o.and_then(move |_| {
                    // Put the Sender for the newly idle worker into the connection queue
                    cq_clone.lock().unwrap().push_back(new_tx);
                    Ok(())
                }).forget();
                Ok(())
            }).forget();
        }
    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2017, 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0
use argh::FromArgs;
use slog_scope::info;

#[derive(FromArgs)]
/// Top-level command.
struct TopLevel {
    #[argh(subcommand)]
    entry_point: EntryPoints,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum EntryPoints {
    Cli(CliOptions),
    Server(ServerOptions),
}

#[derive(FromArgs)]
/// Cli subcommand
#[argh(subcommand, name = "cli")]
struct CliOptions {
    #[argh(subcommand)]
    commands: CliCommands,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum CliCommands {
    Probe(SubCommandProbe),
    ProbeCustom(SubCommandProbeCustom),
    Info(SubCommandInfo),
    Log(SubCommandLog),
    Abort(SubCommandAbort),
    Download(SubCommandDownload),
}

#[derive(FromArgs)]
/// Probe subcommand
#[argh(subcommand, name = "probe")]
struct SubCommandProbe {}

#[derive(FromArgs)]
/// Info subcommand
#[argh(subcommand, name = "info")]
struct SubCommandInfo {}

#[derive(FromArgs)]
/// Log subcommand
#[argh(subcommand, name = "log")]
struct SubCommandLog {}

#[derive(FromArgs)]
/// Abort subcommand
#[argh(subcommand, name = "abort")]
struct SubCommandAbort {}

#[derive(FromArgs)]
/// Download subcommand
#[argh(subcommand, name = "download")]
struct SubCommandDownload {}

#[derive(FromArgs)]
/// Probe custom subcommand
#[argh(subcommand, name = "probe-custom")]
struct SubCommandProbeCustom {
    /// custom address to try probe
    #[argh(positional)]
    server: String,
}

#[derive(FromArgs)]
/// Server subcommand
#[argh(subcommand, name = "server")]
struct ServerOptions {
    /// increase the verboseness level
    #[argh(option, short = 'v', from_str_fn(verbosity_level))]
    verbose: Option<slog::Level>,
}

fn verbosity_level(value: &str) -> Result<slog::Level, String> {
    use std::str::FromStr;
    slog::Level::from_str(value).map_err(|_| format!("Failed to parse verbosity level: {}", value))
}

async fn server_main(cmd: ServerOptions) -> updatehub::Result<()> {
    updatehub::logger::init(cmd.verbose.unwrap_or(slog::Level::Info));
    info!("Starting UpdateHub Agent {}", updatehub::version());

    let settings = updatehub::Settings::load()?;
    updatehub::run(settings).await?;

    Ok(())
}

async fn cli_main(cmd: CliCommands) -> updatehub::Result<()> {
    let client = sdk::Client::new("localhost:8080");

    let res = match cmd {
        CliCommands::Info(_) => client.info().await,
        CliCommands::Probe(_) => todo!("probe"),
        CliCommands::Log(_) => todo!("log"),
        CliCommands::Abort(_) => todo!("abort"),
        CliCommands::Download(_) => todo!("Download"),
        CliCommands::ProbeCustom(input) => todo!("Probe custom {:?}", input.server),
    };

    dbg!(res).unwrap();

    Ok(())
}

#[actix_rt::main]
async fn main() {
    let cmd: TopLevel = argh::from_env();

    let res = match cmd.entry_point {
        EntryPoints::Cli(cli) => cli_main(cli.commands).await,
        EntryPoints::Server(cmd) => server_main(cmd).await,
    };

    if let Err(e) = res {
        eprintln!("{}", e);
        std::process::exit(1);
    }
}
<commit_msg>updatehub: Finish commands on client<commit_after>// Copyright (C) 2017, 2018 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0
use argh::FromArgs;
use slog_scope::info;
use std::path::PathBuf;

#[derive(FromArgs)]
/// Top-level command.
struct TopLevel {
    #[argh(subcommand)]
    entry_point: EntryPoints,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum EntryPoints {
    Client(ClientOptions),
    Server(ServerOptions),
}

#[derive(FromArgs)]
/// Client subcommand
#[argh(subcommand, name = "client")]
struct ClientOptions {
    #[argh(subcommand)]
    commands: ClientCommands,
}

#[derive(FromArgs)]
#[argh(subcommand)]
enum ClientCommands {
    Info(Info),
    Log(Log),
    Probe(Probe),
    AbortDownload(AbortDownload),
    LocalInstall(LocalInstall),
    RemoteInstall(RemoteInstall),
}

#[derive(FromArgs)]
/// Fetches information about the current state of the agent
#[argh(subcommand, name = "info")]
struct Info {}

#[derive(FromArgs)]
/// Fetches the available log entries for the last update cycle
#[argh(subcommand, name = "log")]
struct Log {}

#[derive(FromArgs)]
/// Checks if the server has a new update for this device.
///
/// A custom server for the update cycle can be specified via the ´--server´
#[argh(subcommand, name = "probe")]
struct Probe {
    /// custom address to try probe
    #[argh(option)]
    server: Option<String>,
}

#[derive(FromArgs)]
/// Abort current running download
#[argh(subcommand, name = "abort-download")]
struct AbortDownload {}

#[derive(FromArgs)]
/// Request agent to install a local update package
#[argh(subcommand, name = "local-install")]
struct LocalInstall {
    /// path to the update package
    #[argh(positional)]
    file: PathBuf,
}

#[derive(FromArgs)]
/// Request agent to download and install a package from a direct URL
#[argh(subcommand, name = "remote-install")]
struct RemoteInstall {
    /// the URL to get the update package
    #[argh(positional)]
    url: String,
}

#[derive(FromArgs)]
/// Server subcommand
#[argh(subcommand, name = "server")]
struct ServerOptions {
    /// increase the verboseness level
    #[argh(option, short = 'v', from_str_fn(verbosity_level))]
    verbosity: Option<slog::Level>,
}

fn verbosity_level(value: &str) -> Result<slog::Level, String> {
    use std::str::FromStr;
    slog::Level::from_str(value).map_err(|_| format!("Failed to parse verbosity level: {}", value))
}

async fn server_main(cmd: ServerOptions) -> updatehub::Result<()> {
    updatehub::logger::init(cmd.verbosity.unwrap_or(slog::Level::Info));
    info!("Starting UpdateHub Agent {}", updatehub::version());

    let settings = updatehub::Settings::load()?;
    updatehub::run(settings).await?;

    Ok(())
}

async fn client_main(cmd: ClientCommands) -> updatehub::Result<()> {
    let client = sdk::Client::new("localhost:8080");

    match cmd {
        ClientCommands::Info(_) => println!("{:#?}", client.info().await),
        ClientCommands::Log(_) => println!("{:#?}", client.log().await),
        ClientCommands::Probe(Probe { server }) => println!("{:#?}", client.probe(server).await),
        ClientCommands::AbortDownload(_) => println!("{:#?}", client.abort_download().await),
        ClientCommands::LocalInstall(LocalInstall { file }) => {
            let file =
                if file.is_absolute() { file } else { std::env::current_dir().unwrap().join(file) };
            println!("{:#?}", client.local_install(file).await)
        }
        ClientCommands::RemoteInstall(RemoteInstall { url }) => {
            println!("{:#?}", client.remote_install(url).await)
        }
    }

    Ok(())
}

#[actix_rt::main]
async fn main() {
    let cmd: TopLevel = argh::from_env();

    let res = match cmd.entry_point {
        EntryPoints::Client(client) => client_main(client.commands).await,
        EntryPoints::Server(cmd) => server_main(cmd).await,
    };

    if let Err(e) = res {
        eprintln!("{}", e);
        std::process::exit(1);
    }
}
<|endoftext|>


<commit_before>extern crate num_cpus;
extern crate pkg_config;
extern crate cmake;

use std::path::Path;
use std::process::{Command, self};
use std::io::Write;
use std::env;

macro_rules! println_stderr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

fn run_command_or_fail(dir: &str, cmd: &str, args: &[&str]) {
    println_stderr!("Running command: \"{} {}\" in dir: {}", cmd, args.join(" "), dir);
    let ret = Command::new(cmd).current_dir(dir).args(args).status();
    match ret.map(|status| (status.success(), status.code())) {
        Ok((true, _)) => { return },
        Ok((false, Some(c))) => { panic!("Command failed with error code {}", c) },
        Ok((false, None)) => { panic!("Command got killed") },
        Err(e) => { panic!("Command failed with error: {}", e) },
    }
}

fn main() {
    let librdkafka_version = env!("CARGO_PKG_VERSION")
        .split('-')
        .next()
        .expect("Crate version is not valid");

    if env::var("CARGO_FEATURE_DYNAMIC_LINKING").is_ok() {
        println_stderr!("librdkafka will be linked dynamically");
        let pkg_probe = pkg_config::Config::new()
            .cargo_metadata(true)
            .atleast_version(librdkafka_version)
            .probe("rdkafka");

        match pkg_probe {
            Ok(library) => {
                println_stderr!("librdkafka found on the system:");
                println_stderr!("  Name: {:?}", library.libs);
                println_stderr!("  Path: {:?}", library.link_paths);
                println_stderr!("  Version: {}", library.version);
            }
            Err(_) => {
                println_stderr!("librdkafka {} cannot be found on the system", librdkafka_version);
                println_stderr!("Dynamic linking failed. Exiting.");
                process::exit(1);
            }
        }
    } else {
        println_stderr!("Building and linking librdkafka statically");
        build_librdkafka();
    }
}

fn build_librdkafka() {
    let mut configure_flags = Vec::new();

    if env::var("CARGO_FEATURE_SASL").is_ok() {
        configure_flags.push("--enable-sasl");
    } else {
        configure_flags.push("--disable-sasl");
    }

    if env::var("CARGO_FEATURE_SSL").is_ok() {
        configure_flags.push("--enable-ssl");
    } else {
        configure_flags.push("--disable-ssl");
    }

    if env::var("CARGO_FEATURE_EXTERNAL_LZ4").is_ok() {
        configure_flags.push("--enable-lz4");
    } else {
        configure_flags.push("--disable-lz4");
    }

    configure_flags.push("--enable-static");

    if !Path::new("librdkafka/LICENSE").exists() {
        println_stderr!("Setting up submodules");
        run_command_or_fail("../", "git", &["submodule", "update", "--init"]);
    }

    env::set_var("NUM_JOBS", num_cpus::get().to_string());
    let mut config = cmake::Config::new("librdkafka");
    config.define("RDKAFKA_BUILD_STATIC", "1")
          .build_target("rdkafka");
    if env::var("CARGO_FEATURE_SASL").is_ok() {
        config.define("WITH_SASL", "1");
    }
    let dst = config.build();
    println!("cargo:rustc-link-search=native={}/build/src", dst.display());
    println!("cargo:rustc-link-lib=static=rdkafka");
}
<commit_msg>Define WITH_SASL by default<commit_after>extern crate num_cpus;
extern crate pkg_config;
extern crate cmake;

use std::path::Path;
use std::process::{Command, self};
use std::io::Write;
use std::env;

macro_rules! println_stderr(
    ($($arg:tt)*) => { {
        let r = writeln!(&mut ::std::io::stderr(), $($arg)*);
        r.expect("failed printing to stderr");
    } }
);

fn run_command_or_fail(dir: &str, cmd: &str, args: &[&str]) {
    println_stderr!("Running command: \"{} {}\" in dir: {}", cmd, args.join(" "), dir);
    let ret = Command::new(cmd).current_dir(dir).args(args).status();
    match ret.map(|status| (status.success(), status.code())) {
        Ok((true, _)) => { return },
        Ok((false, Some(c))) => { panic!("Command failed with error code {}", c) },
        Ok((false, None)) => { panic!("Command got killed") },
        Err(e) => { panic!("Command failed with error: {}", e) },
    }
}

fn main() {
    let librdkafka_version = env!("CARGO_PKG_VERSION")
        .split('-')
        .next()
        .expect("Crate version is not valid");

    if env::var("CARGO_FEATURE_DYNAMIC_LINKING").is_ok() {
        println_stderr!("librdkafka will be linked dynamically");
        let pkg_probe = pkg_config::Config::new()
            .cargo_metadata(true)
            .atleast_version(librdkafka_version)
            .probe("rdkafka");

        match pkg_probe {
            Ok(library) => {
                println_stderr!("librdkafka found on the system:");
                println_stderr!("  Name: {:?}", library.libs);
                println_stderr!("  Path: {:?}", library.link_paths);
                println_stderr!("  Version: {}", library.version);
            }
            Err(_) => {
                println_stderr!("librdkafka {} cannot be found on the system", librdkafka_version);
                println_stderr!("Dynamic linking failed. Exiting.");
                process::exit(1);
            }
        }
    } else {
        println_stderr!("Building and linking librdkafka statically");
        build_librdkafka();
    }
}

fn build_librdkafka() {
    let mut configure_flags = Vec::new();

    if env::var("CARGO_FEATURE_SASL").is_ok() {
        configure_flags.push("--enable-sasl");
    } else {
        configure_flags.push("--disable-sasl");
    }

    if env::var("CARGO_FEATURE_SSL").is_ok() {
        configure_flags.push("--enable-ssl");
    } else {
        configure_flags.push("--disable-ssl");
    }

    if env::var("CARGO_FEATURE_EXTERNAL_LZ4").is_ok() {
        configure_flags.push("--enable-lz4");
    } else {
        configure_flags.push("--disable-lz4");
    }

    configure_flags.push("--enable-static");

    if !Path::new("librdkafka/LICENSE").exists() {
        println_stderr!("Setting up submodules");
        run_command_or_fail("../", "git", &["submodule", "update", "--init"]);
    }

    env::set_var("NUM_JOBS", num_cpus::get().to_string());
    let mut config = cmake::Config::new("librdkafka");
    config.define("RDKAFKA_BUILD_STATIC", "1")
          .build_target("rdkafka");
    if env::var("CARGO_FEATURE_SASL").is_ok() {
        config.define("WITH_SASL", "1");
    } else {
        config.define("WITH_SASL", "0");
    }
    let dst = config.build();
    println!("cargo:rustc-link-search=native={}/build/src", dst.display());
    println!("cargo:rustc-link-lib=static=rdkafka");
}
<|endoftext|>


<commit_before>use clap::{App, Arg};
use regex::Regex;

fn integer_value(v: String) -> Result<(), String> {
  v.parse::<u16>().map(|_| ()).map_err(|e| format!("'{}' is not a valid port value: {}", v, e) )
}

pub(crate) fn setup_app<'a, 'b>(program: String, version: &'b str) -> App<'a, 'b> {
  App::new(program)
    .version(version)
    .about("Standalone Pact verifier")
    .version_short("v")
    .arg(Arg::with_name("loglevel")
      .short("l")
      .long("loglevel")
      .takes_value(true)
      .use_delimiter(false)
      .possible_values(&["error", "warn", "info", "debug", "trace", "none"])
      .help("Log level (defaults to warn)"))
    .arg(Arg::with_name("file")
      .short("f")
      .long("file")
      .required_unless_one(&["dir", "url", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("Pact file to verify (can be repeated)"))
    .arg(Arg::with_name("dir")
      .short("d")
      .long("dir")
      .required_unless_one(&["file", "url", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("Directory of pact files to verify (can be repeated)"))
    .arg(Arg::with_name("url")
      .short("u")
      .long("url")
      .required_unless_one(&["file", "dir", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of pact file to verify (can be repeated)"))
    .arg(Arg::with_name("broker-url")
      .short("b")
      .long("broker-url")
      .env("PACT_BROKER_BASE_URL")
      .required_unless_one(&["file", "dir", "url"])
      .requires("provider-name")
      .takes_value(true)
      .use_delimiter(false)
      .empty_values(false)
      .help("URL of the pact broker to fetch pacts from to verify (requires the provider name parameter)"))
    .arg(Arg::with_name("hostname")
      .short("h")
      .long("hostname")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider hostname (defaults to localhost)"))
    .arg(Arg::with_name("port")
      .short("p")
      .long("port")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider port (defaults to protocol default 80/443)")
      .validator(integer_value))
    .arg(Arg::with_name("scheme")
      .long("scheme")
      .takes_value(true)
      .possible_values(&["http", "https"])
      .default_value("http")
      .help("Provider URI scheme (defaults to http)"))
    .arg(Arg::with_name("provider-name")
      .short("n")
      .long("provider-name")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider name (defaults to provider)"))
    .arg(Arg::with_name("state-change-url")
      .short("s")
      .long("state-change-url")
      .takes_value(true)
      .use_delimiter(false)
      .help("URL to post state change requests to"))
    .arg(Arg::with_name("state-change-as-query")
      .long("state-change-as-query")
      .help("State change request data will be sent as query parameters instead of in the request body"))
    .arg(Arg::with_name("state-change-teardown")
      .long("state-change-teardown")
      .help("State change teardown requests are to be made after each interaction"))
    .arg(Arg::with_name("filter-description")
      .long("filter-description")
      .env("PACT_DESCRIPTION")
      .takes_value(true)
      .use_delimiter(false)
      .validator(|val| Regex::new(&val)
        .map(|_| ())
        .map_err(|err| format!("'{}' is an invalid filter value: {}", val, err)))
      .help("Only validate interactions whose descriptions match this filter"))
    .arg(Arg::with_name("filter-state")
      .long("filter-state")
      .env("PACT_PROVIDER_STATE")
      .takes_value(true)
      .use_delimiter(false)
      .conflicts_with("filter-no-state")
      .validator(|val| Regex::new(&val)
        .map(|_| ())
        .map_err(|err| format!("'{}' is an invalid filter value: {}", val, err)))
      .help("Only validate interactions whose provider states match this filter"))
    .arg(Arg::with_name("filter-no-state")
      .long("filter-no-state")
      .env("PACT_PROVIDER_NO_STATE")
      .conflicts_with("filter-state")
      .help("Only validate interactions that have no defined provider state"))
    .arg(Arg::with_name("filter-consumer")
      .short("c")
      .long("filter-consumer")
      .takes_value(true)
      .multiple(true)
      .empty_values(false)
      .help("Consumer name to filter the pacts to be verified (can be repeated)"))
    .arg(Arg::with_name("user")
      .long("user")
      .env("PACT_BROKER_USERNAME")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("token")
      .help("Username to use when fetching pacts from URLS"))
    .arg(Arg::with_name("password")
      .long("password")
      .env("PACT_BROKER_PASSWORD")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("token")
      .help("Password to use when fetching pacts from URLS"))
    .arg(Arg::with_name("token")
      .short("t")
      .long("token")
      .env("PACT_BROKER_TOKEN")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("user")
      .help("Bearer token to use when fetching pacts from URLS"))
    .arg(Arg::with_name("publish")
      .long("publish")
      .requires("broker-url")
      .requires("provider-version")
      .help("Enables publishing of verification results back to the Pact Broker. Requires the broker-url and provider-version parameters."))
    .arg(Arg::with_name("provider-version")
      .long("provider-version")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .help("Provider version that is being verified. This is required when publishing results."))
    .arg(Arg::with_name("build-url")
      .long("build-url")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of the build to associate with the published verification results."))
    .arg(Arg::with_name("provider-tags")
      .long("provider-tags")
      .takes_value(true)
      .use_delimiter(true)
      .empty_values(false)
      .help("Provider tags to use when publishing results. Accepts comma-separated values."))
    .arg(Arg::with_name("base-path")
      .long("base-path")
      .takes_value(true)
      .use_delimiter(false)
      .empty_values(false)
      .help("Base path to add to all requests"))
    .arg(Arg::with_name("consumer-version-tags")
      .long("consumer-version-tags")
      .takes_value(true)
      .use_delimiter(true)
      .empty_values(false)
      .requires("broker-url")
      .conflicts_with("consumer-version-selectors")
      .help("Consumer tags to use when fetching pacts from the Broker. Accepts comma-separated values."))
    .arg(Arg::with_name("consumer-version-selectors")
      .long("consumer-version-selectors")
      .takes_value(true)
      .use_delimiter(false)
      .empty_values(false)
      .requires("broker-url")
      .conflicts_with("consumer-version-tags")
      .help("Consumer version selectors to use when fetching pacts from the Broker. Accepts a JSON string as per https://docs.pact.io/pact_broker/advanced_topics/consumer_version_selectors/"))
    .arg(Arg::with_name("disable-ssl-verification")
      .long("disable-ssl-verification")
      .takes_value(false)
      .help("Disables validation of SSL certificates"))
    .arg(Arg::with_name("enable-pending")
      .long("enable-pending")
      .requires("broker-url")
      .help("Enables Pending Pacts"))
    .arg(Arg::with_name("include-wip-pacts-since")
      .long("include-wip-pacts-since")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .requires("broker-url")
      .help("Allow pacts that don't match given consumer selectors (or tags) to  be verified, without causing the overall task to fail. For more information, see https://pact.io/wip"))
    .arg(Arg::with_name("request-timeout")
      .long("request-timeout")
      .takes_value(true)
      .empty_values(false)
      .validator(integer_value)
      .help("Sets the HTTP request timeout in milliseconds for requests to the target API and for state change requests."))
    }

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

  #[test]
  fn validates_integer_value() {
    expect!(integer_value("1234".to_string())).to(be_ok());
    expect!(integer_value("1234x".to_string())).to(be_err());
  }
}
<commit_msg>fix: add missing provider-branch to verifier CLI<commit_after>use clap::{App, Arg};
use regex::Regex;

fn integer_value(v: String) -> Result<(), String> {
  v.parse::<u16>().map(|_| ()).map_err(|e| format!("'{}' is not a valid port value: {}", v, e) )
}

pub(crate) fn setup_app<'a, 'b>(program: String, version: &'b str) -> App<'a, 'b> {
  App::new(program)
    .version(version)
    .about("Standalone Pact verifier")
    .version_short("v")
    .arg(Arg::with_name("loglevel")
      .short("l")
      .long("loglevel")
      .takes_value(true)
      .use_delimiter(false)
      .possible_values(&["error", "warn", "info", "debug", "trace", "none"])
      .help("Log level (defaults to warn)"))
    .arg(Arg::with_name("file")
      .short("f")
      .long("file")
      .required_unless_one(&["dir", "url", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("Pact file to verify (can be repeated)"))
    .arg(Arg::with_name("dir")
      .short("d")
      .long("dir")
      .required_unless_one(&["file", "url", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("Directory of pact files to verify (can be repeated)"))
    .arg(Arg::with_name("url")
      .short("u")
      .long("url")
      .required_unless_one(&["file", "dir", "broker-url"])
      .takes_value(true)
      .use_delimiter(false)
      .multiple(true)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of pact file to verify (can be repeated)"))
    .arg(Arg::with_name("broker-url")
      .short("b")
      .long("broker-url")
      .env("PACT_BROKER_BASE_URL")
      .required_unless_one(&["file", "dir", "url"])
      .requires("provider-name")
      .takes_value(true)
      .use_delimiter(false)
      .empty_values(false)
      .help("URL of the pact broker to fetch pacts from to verify (requires the provider name parameter)"))
    .arg(Arg::with_name("hostname")
      .short("h")
      .long("hostname")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider hostname (defaults to localhost)"))
    .arg(Arg::with_name("port")
      .short("p")
      .long("port")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider port (defaults to protocol default 80/443)")
      .validator(integer_value))
    .arg(Arg::with_name("scheme")
      .long("scheme")
      .takes_value(true)
      .possible_values(&["http", "https"])
      .default_value("http")
      .help("Provider URI scheme (defaults to http)"))
    .arg(Arg::with_name("provider-name")
      .short("n")
      .long("provider-name")
      .takes_value(true)
      .use_delimiter(false)
      .help("Provider name (defaults to provider)"))
    .arg(Arg::with_name("state-change-url")
      .short("s")
      .long("state-change-url")
      .takes_value(true)
      .use_delimiter(false)
      .help("URL to post state change requests to"))
    .arg(Arg::with_name("state-change-as-query")
      .long("state-change-as-query")
      .help("State change request data will be sent as query parameters instead of in the request body"))
    .arg(Arg::with_name("state-change-teardown")
      .long("state-change-teardown")
      .help("State change teardown requests are to be made after each interaction"))
    .arg(Arg::with_name("filter-description")
      .long("filter-description")
      .env("PACT_DESCRIPTION")
      .takes_value(true)
      .use_delimiter(false)
      .validator(|val| Regex::new(&val)
        .map(|_| ())
        .map_err(|err| format!("'{}' is an invalid filter value: {}", val, err)))
      .help("Only validate interactions whose descriptions match this filter"))
    .arg(Arg::with_name("filter-state")
      .long("filter-state")
      .env("PACT_PROVIDER_STATE")
      .takes_value(true)
      .use_delimiter(false)
      .conflicts_with("filter-no-state")
      .validator(|val| Regex::new(&val)
        .map(|_| ())
        .map_err(|err| format!("'{}' is an invalid filter value: {}", val, err)))
      .help("Only validate interactions whose provider states match this filter"))
    .arg(Arg::with_name("filter-no-state")
      .long("filter-no-state")
      .env("PACT_PROVIDER_NO_STATE")
      .conflicts_with("filter-state")
      .help("Only validate interactions that have no defined provider state"))
    .arg(Arg::with_name("filter-consumer")
      .short("c")
      .long("filter-consumer")
      .takes_value(true)
      .multiple(true)
      .empty_values(false)
      .help("Consumer name to filter the pacts to be verified (can be repeated)"))
    .arg(Arg::with_name("user")
      .long("user")
      .env("PACT_BROKER_USERNAME")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("token")
      .help("Username to use when fetching pacts from URLS"))
    .arg(Arg::with_name("password")
      .long("password")
      .env("PACT_BROKER_PASSWORD")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("token")
      .help("Password to use when fetching pacts from URLS"))
    .arg(Arg::with_name("token")
      .short("t")
      .long("token")
      .env("PACT_BROKER_TOKEN")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .conflicts_with("user")
      .help("Bearer token to use when fetching pacts from URLS"))
    .arg(Arg::with_name("publish")
      .long("publish")
      .requires("broker-url")
      .requires("provider-version")
      .help("Enables publishing of verification results back to the Pact Broker. Requires the broker-url and provider-version parameters."))
    .arg(Arg::with_name("provider-version")
      .long("provider-version")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .help("Provider version that is being verified. This is required when publishing results."))
    .arg(Arg::with_name("build-url")
      .long("build-url")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .help("URL of the build to associate with the published verification results."))
    .arg(Arg::with_name("provider-tags")
      .long("provider-tags")
      .takes_value(true)
      .use_delimiter(true)
      .empty_values(false)
      .help("Provider tags to use when publishing results. Accepts comma-separated values."))
    .arg(Arg::with_name("provider-branch")
      .long("provider-branch")
      .takes_value(true)
      .empty_values(false)
      .help("Provider branch to use when publishing results"))
    .arg(Arg::with_name("base-path")
      .long("base-path")
      .takes_value(true)
      .use_delimiter(false)
      .empty_values(false)
      .help("Base path to add to all requests"))
    .arg(Arg::with_name("consumer-version-tags")
      .long("consumer-version-tags")
      .takes_value(true)
      .use_delimiter(true)
      .empty_values(false)
      .requires("broker-url")
      .conflicts_with("consumer-version-selectors")
      .help("Consumer tags to use when fetching pacts from the Broker. Accepts comma-separated values."))
    .arg(Arg::with_name("consumer-version-selectors")
      .long("consumer-version-selectors")
      .takes_value(true)
      .use_delimiter(false)
      .empty_values(false)
      .requires("broker-url")
      .conflicts_with("consumer-version-tags")
      .help("Consumer version selectors to use when fetching pacts from the Broker. Accepts a JSON string as per https://docs.pact.io/pact_broker/advanced_topics/consumer_version_selectors/"))
    .arg(Arg::with_name("disable-ssl-verification")
      .long("disable-ssl-verification")
      .takes_value(false)
      .help("Disables validation of SSL certificates"))
    .arg(Arg::with_name("enable-pending")
      .long("enable-pending")
      .requires("broker-url")
      .help("Enables Pending Pacts"))
    .arg(Arg::with_name("include-wip-pacts-since")
      .long("include-wip-pacts-since")
      .takes_value(true)
      .use_delimiter(false)
      .number_of_values(1)
      .empty_values(false)
      .requires("broker-url")
      .help("Allow pacts that don't match given consumer selectors (or tags) to  be verified, without causing the overall task to fail. For more information, see https://pact.io/wip"))
    .arg(Arg::with_name("request-timeout")
      .long("request-timeout")
      .takes_value(true)
      .empty_values(false)
      .validator(integer_value)
      .help("Sets the HTTP request timeout in milliseconds for requests to the target API and for state change requests."))
    }

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

  #[test]
  fn validates_integer_value() {
    expect!(integer_value("1234".to_string())).to(be_ok());
    expect!(integer_value("1234x".to_string())).to(be_err());
  }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::rc::Rc;

use serde::de;

use errors::*;
use header::*;

use super::Deserializer;
use super::read::DeRead;
use super::reentrant::ReentrantGuard;

// This is ugly because #[serde(borrow)] can't be used with collections directly
// at the moment. See
// https://github.com/serde-rs/serde/issues/914#issuecomment-298801226 for more.
#[derive(Clone, Debug, Deserialize)]
pub struct Key<'a>(#[serde(borrow)] Cow<'a, [u8]>);

/// A BSER template is logically an array of objects, all with the same or
/// similar keys.
///
/// A template is serialized as
/// `<BSER_TEMPLATE><BSER_ARRAY of keys><number of objects><object 1 values><object 2 values>...`
///
/// and gets deserialized as
///
/// ```text
/// [
///    {key1: obj1_value1, key2: obj1_value2, ...},
///    {key2: obj2_value1, key2: obj2_value2, ...},
///    ...
/// ]
/// ```
///
/// The special value BSER_SKIP is used if a particular object doesn't have a
/// key.
pub struct Template<'a, 'de, R: 'a> {
    de: &'a mut Deserializer<R>,
    keys: Rc<Vec<Key<'de>>>,
    remaining: usize,
}

impl<'a, 'de, R: 'a> Template<'a, 'de, R> {
    /// Create a new `Template`.
    ///
    /// `_guard` makes sure the caller is accounting for the recursion limit.
    pub fn new(
        de: &'a mut Deserializer<R>,
        keys: Vec<Key<'de>>,
        nitems: usize,
        _guard: &ReentrantGuard,
    ) -> Self {
        Template {
            de: de,
            keys: Rc::new(keys),
            remaining: nitems,
        }
    }
}

impl<'a, 'de, R> de::SeqAccess<'de> for Template<'a, 'de, R>
where
    R: 'a + DeRead<'de>,
{
    type Error = Error;

    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>>
    where
        T: de::DeserializeSeed<'de>,
    {
        if self.remaining == 0 {
            Ok(None)
        } else {
            self.remaining -= 1;
            let obj_de = ObjectDeserializer {
                de: &mut *self.de,
                keys: self.keys.clone(),
            };
            let value = seed.deserialize(obj_de)?;
            Ok(Some(value))
        }
    }
}

struct ObjectDeserializer<'a, 'de, R: 'a> {
    de: &'a mut Deserializer<R>,
    keys: Rc<Vec<Key<'de>>>,
}

impl<'a, 'de, R> de::Deserializer<'de> for ObjectDeserializer<'a, 'de, R>
where
    R: 'a + DeRead<'de>,
{
    type Error = Error;

    #[inline]
    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        visitor.visit_map(TemplateObject::new(&mut *self.de, self.keys))
    }

    #[inline]
    fn deserialize_newtype_struct<V>(self, _name: &'static str, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        // This is e.g. E(T). Ignore the E.
        visitor.visit_newtype_struct(self)
    }

    fn deserialize_enum<V>(
        self,
        _name: &'static str,
        _variants: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        visitor.visit_map(TemplateObject::new(&mut *self.de, self.keys))
    }

    forward_to_deserialize_any! {
        bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string bytes
        byte_buf unit unit_struct seq tuple tuple_struct map struct identifier
        ignored_any option
    }
}

struct TemplateObject<'a, 'de, R: 'a> {
    de: &'a mut Deserializer<R>,
    keys: Rc<Vec<Key<'de>>>,
    cur: usize,
}

impl<'a, 'de, R> TemplateObject<'a, 'de, R>
where
    R: 'a + DeRead<'de>,
{
    fn new(de: &'a mut Deserializer<R>, keys: Rc<Vec<Key<'de>>>) -> Self {
        TemplateObject {
            de: de,
            keys: keys,
            cur: 0,
        }
    }
}

impl<'a, 'de, R> de::MapAccess<'de> for TemplateObject<'a, 'de, R>
where
    R: 'a + DeRead<'de>,
{
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>>
    where
        K: de::DeserializeSeed<'de>,
    {
        if self.cur == self.keys.len() {
            Ok(None)
        } else {
            let cur = self.cur;
            self.cur += 1;
            let obj_de = KeyDeserializer {
                key: &self.keys[cur],
            };
            let value = seed.deserialize(obj_de)?;
            Ok(Some(value))
        }
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value>
    where
        V: de::DeserializeSeed<'de>,
    {
        seed.deserialize(ValueDeserializer { de: &mut *self.de })
    }
}

struct KeyDeserializer<'a, 'de: 'a> {
    key: &'a Key<'de>,
}

impl<'a, 'de: 'a> de::Deserializer<'de> for KeyDeserializer<'a, 'de> {
    type Error = Error;

    #[inline]
    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        match self.key.0 {
            Cow::Borrowed(s) => visitor.visit_borrowed_bytes(s),
            Cow::Owned(ref s) => visitor.visit_bytes(s),
        }
    }

    forward_to_deserialize_any! {
        bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string bytes
        byte_buf unit unit_struct seq tuple tuple_struct map struct identifier
        ignored_any option enum newtype_struct
    }
}

struct ValueDeserializer<'a, R: 'a> {
    de: &'a mut Deserializer<R>,
}

impl<'a, 'de, R> de::Deserializer<'de> for ValueDeserializer<'a, R>
where
    R: 'a + DeRead<'de>,
{
    type Error = Error;

    #[inline]
    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        // The only new thing here compared to the main Deserializer is that it
        // is possible to skip this value with BSER_SKIP.
        match self.de.bunser.peek()? {
            BSER_SKIP => {
                self.de.bunser.discard();
                visitor.visit_none()
            }
            _ => self.de.deserialize_any(visitor),
        }
    }

    /// Parse a BSER_SKIP or a null as a None, and anything else as a
    /// `Some(...)`.
    #[inline]
    fn deserialize_option<V>(self, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        match self.de.bunser.peek()? {
            BSER_SKIP | BSER_NULL => {
                self.de.bunser.discard();
                visitor.visit_none()
            }
            _ => visitor.visit_some(self),
        }
    }

    #[inline]
    fn deserialize_newtype_struct<V>(self, _name: &str, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        // This is e.g. E(T). Ignore the E.
        visitor.visit_newtype_struct(self)
    }

    // TODO: do we also need to do enum here?

    forward_to_deserialize_any! {
        bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string bytes
        byte_buf unit unit_struct seq tuple tuple_struct map struct identifier
        ignored_any enum
    }
}
<commit_msg>always deserialize template keys as strings<commit_after>use std::borrow::Cow;
use std::rc::Rc;

use serde::de;

use errors::*;
use header::*;

use super::Deserializer;
use super::read::DeRead;
use super::reentrant::ReentrantGuard;

// This is ugly because #[serde(borrow)] can't be used with collections directly
// at the moment. See
// https://github.com/serde-rs/serde/issues/914#issuecomment-298801226 for more.
// Note that keys are always ASCII, so treating them as str is fine, even if they're
// serialized as BSER_BYTESTRING instances. (These keys are used as struct field
// identifiers, which are Unicode strings so can't be directly matched up with
// bytestrings.)
#[derive(Clone, Debug, Deserialize)]
pub struct Key<'a>(#[serde(borrow)] Cow<'a, str>);

/// A BSER template is logically an array of objects, all with the same or
/// similar keys.
///
/// A template is serialized as
/// `<BSER_TEMPLATE><BSER_ARRAY of keys><number of objects><object 1 values><object 2 values>...`
///
/// and gets deserialized as
///
/// ```text
/// [
///    {key1: obj1_value1, key2: obj1_value2, ...},
///    {key2: obj2_value1, key2: obj2_value2, ...},
///    ...
/// ]
/// ```
///
/// The special value BSER_SKIP is used if a particular object doesn't have a
/// key.
pub struct Template<'a, 'de, R: 'a> {
    de: &'a mut Deserializer<R>,
    keys: Rc<Vec<Key<'de>>>,
    remaining: usize,
}

impl<'a, 'de, R: 'a> Template<'a, 'de, R> {
    /// Create a new `Template`.
    ///
    /// `_guard` makes sure the caller is accounting for the recursion limit.
    pub fn new(
        de: &'a mut Deserializer<R>,
        keys: Vec<Key<'de>>,
        nitems: usize,
        _guard: &ReentrantGuard,
    ) -> Self {
        Template {
            de: de,
            keys: Rc::new(keys),
            remaining: nitems,
        }
    }
}

impl<'a, 'de, R> de::SeqAccess<'de> for Template<'a, 'de, R>
where
    R: 'a + DeRead<'de>,
{
    type Error = Error;

    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>>
    where
        T: de::DeserializeSeed<'de>,
    {
        if self.remaining == 0 {
            Ok(None)
        } else {
            self.remaining -= 1;
            let obj_de = ObjectDeserializer {
                de: &mut *self.de,
                keys: self.keys.clone(),
            };
            let value = seed.deserialize(obj_de)?;
            Ok(Some(value))
        }
    }
}

struct ObjectDeserializer<'a, 'de, R: 'a> {
    de: &'a mut Deserializer<R>,
    keys: Rc<Vec<Key<'de>>>,
}

impl<'a, 'de, R> de::Deserializer<'de> for ObjectDeserializer<'a, 'de, R>
where
    R: 'a + DeRead<'de>,
{
    type Error = Error;

    #[inline]
    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        visitor.visit_map(TemplateObject::new(&mut *self.de, self.keys))
    }

    #[inline]
    fn deserialize_newtype_struct<V>(self, _name: &'static str, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        // This is e.g. E(T). Ignore the E.
        visitor.visit_newtype_struct(self)
    }

    fn deserialize_enum<V>(
        self,
        _name: &'static str,
        _variants: &'static [&'static str],
        visitor: V,
    ) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        visitor.visit_map(TemplateObject::new(&mut *self.de, self.keys))
    }

    forward_to_deserialize_any! {
        bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string bytes
        byte_buf unit unit_struct seq tuple tuple_struct map struct identifier
        ignored_any option
    }
}

struct TemplateObject<'a, 'de, R: 'a> {
    de: &'a mut Deserializer<R>,
    keys: Rc<Vec<Key<'de>>>,
    cur: usize,
}

impl<'a, 'de, R> TemplateObject<'a, 'de, R>
where
    R: 'a + DeRead<'de>,
{
    fn new(de: &'a mut Deserializer<R>, keys: Rc<Vec<Key<'de>>>) -> Self {
        TemplateObject {
            de: de,
            keys: keys,
            cur: 0,
        }
    }
}

impl<'a, 'de, R> de::MapAccess<'de> for TemplateObject<'a, 'de, R>
where
    R: 'a + DeRead<'de>,
{
    type Error = Error;

    fn next_key_seed<K>(&mut self, seed: K) -> Result<Option<K::Value>>
    where
        K: de::DeserializeSeed<'de>,
    {
        if self.cur == self.keys.len() {
            Ok(None)
        } else {
            let cur = self.cur;
            self.cur += 1;
            let obj_de = KeyDeserializer {
                key: &self.keys[cur],
            };
            let value = seed.deserialize(obj_de)?;
            Ok(Some(value))
        }
    }

    fn next_value_seed<V>(&mut self, seed: V) -> Result<V::Value>
    where
        V: de::DeserializeSeed<'de>,
    {
        seed.deserialize(ValueDeserializer { de: &mut *self.de })
    }
}

struct KeyDeserializer<'a, 'de: 'a> {
    key: &'a Key<'de>,
}

impl<'a, 'de: 'a> de::Deserializer<'de> for KeyDeserializer<'a, 'de> {
    type Error = Error;

    #[inline]
    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        match self.key.0 {
            Cow::Borrowed(s) => visitor.visit_borrowed_str(s),
            Cow::Owned(ref s) => visitor.visit_str(s),
        }
    }

    forward_to_deserialize_any! {
        bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string bytes
        byte_buf unit unit_struct seq tuple tuple_struct map struct identifier
        ignored_any option enum newtype_struct
    }
}

struct ValueDeserializer<'a, R: 'a> {
    de: &'a mut Deserializer<R>,
}

impl<'a, 'de, R> de::Deserializer<'de> for ValueDeserializer<'a, R>
where
    R: 'a + DeRead<'de>,
{
    type Error = Error;

    #[inline]
    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        // The only new thing here compared to the main Deserializer is that it
        // is possible to skip this value with BSER_SKIP.
        match self.de.bunser.peek()? {
            BSER_SKIP => {
                self.de.bunser.discard();
                visitor.visit_none()
            }
            _ => self.de.deserialize_any(visitor),
        }
    }

    /// Parse a BSER_SKIP or a null as a None, and anything else as a
    /// `Some(...)`.
    #[inline]
    fn deserialize_option<V>(self, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        match self.de.bunser.peek()? {
            BSER_SKIP | BSER_NULL => {
                self.de.bunser.discard();
                visitor.visit_none()
            }
            _ => visitor.visit_some(self),
        }
    }

    #[inline]
    fn deserialize_newtype_struct<V>(self, _name: &str, visitor: V) -> Result<V::Value>
    where
        V: de::Visitor<'de>,
    {
        // This is e.g. E(T). Ignore the E.
        visitor.visit_newtype_struct(self)
    }

    // TODO: do we also need to do enum here?

    forward_to_deserialize_any! {
        bool i8 i16 i32 i64 u8 u16 u32 u64 f32 f64 char str string bytes
        byte_buf unit unit_struct seq tuple tuple_struct map struct identifier
        ignored_any enum
    }
}
<|endoftext|>


<commit_before>use super::allocators::{Allocator, DefaultHeap};
use super::compact::Compact;
use super::compact_vec::CompactVec;

/// A simple linear-search key-value dictionary,
/// implemented using two `CompactVec`'s, one for keys, one for values.
///
/// The API loosely follows that of `std::collections::HashMap`.
/// Spilling behaviour using `Allocator` is equivalent to `CompactVec`.
pub struct CompactDict<K: Copy, V: Compact + Clone, A: Allocator = DefaultHeap> {
    keys: CompactVec<K, A>,
    values: CompactVec<V, A>,
}

impl<K: Eq + Copy, V: Compact + Clone, A: Allocator> CompactDict<K, V, A> {
    /// Create new, empty dictionary
    pub fn new() -> Self {
        CompactDict {
            keys: CompactVec::new(),
            values: CompactVec::new(),
        }
    }

    /// Amount of entries in the dictionary
    pub fn len(&self) -> usize {
        self.keys.len()
    }

    /// Is the dictionary empty?
    pub fn is_empty(&self) -> bool {
        self.keys.is_empty()
    }

    /// Look up the value for key `query`, if it exists
    pub fn get(&self, query: K) -> Option<&V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                return Some(&self.values[i]);
            }
        }
        None
    }

    /// Look up the value for keu `query` mutably, if it exists
    pub fn get_mut(&mut self, query: K) -> Option<&mut V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                return Some(&mut self.values[i]);
            }
        }
        None
    }

    /// Lookup up the value for key `query`, if it exists, but also swap the entry
    /// to the beginning of the key/value vectors, so a repeated lookup for that item will be faster
    pub fn get_mru(&mut self, query: K) -> Option<&V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                self.keys.swap(0, i);
                self.values.swap(0, i);
                return Some(&self.values[0]);
            }
        }
        None
    }

    /// Lookup up the value for key `query`, if it exists, but also swap the entry
    /// one index towards the beginning of the key/value vectors, so frequently repeated lookups
    /// for that item will be faster
    pub fn get_mfu(&mut self, query: K) -> Option<&V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                if i > 0 {
                    self.keys.swap(i - 1, i);
                    self.values.swap(i - 1, i);
                    return Some(&self.values[i - 1]);
                } else {
                    return Some(&self.values[0]);
                }
            }
        }
        None
    }

    /// Does the dictionary contain a value for `query`?
    pub fn contains_key(&self, query: K) -> bool {
        self.keys.contains(&query)
    }

    /// Insert new value at key `query` and return the previous value at that key, if any existed
    pub fn insert(&mut self, query: K, new_value: V) -> Option<V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                let old_val = self.values[i].clone();
                self.values[i] = new_value;
                return Some(old_val);
            }
        }
        self.keys.push(query);
        self.values.push(new_value);
        None
    }

    /// Remove value at key `query` and return it, if it existed
    pub fn remove(&mut self, query: K) -> Option<V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                let old_val = self.values[i].clone();
                self.keys.remove(i);
                self.values.remove(i);
                return Some(old_val);
            }
        }
        None
    }

    /// Iterator over all keys in the dictionary
    pub fn keys(&self) -> ::std::slice::Iter<K> {
        self.keys.iter()
    }

    /// Iterator over all values in the dictionary
    pub fn values(&self) -> ::std::slice::Iter<V> {
        self.values.iter()
    }

    /// Iterator over mutable references to all values in the dictionary
    pub fn values_mut(&mut self) -> ::std::slice::IterMut<V> {
        self.values.iter_mut()
    }

    /// Iterator over all key-value pairs in the dictionary
    pub fn pairs<'a>(&'a self) -> impl Iterator<Item = (&'a K, &'a V)> + Clone + 'a {
        self.keys().zip(self.values())
    }
}

impl<K: Eq + Copy, I: Compact, A1: Allocator, A2: Allocator> CompactDict<K, CompactVec<I, A1>, A2> {
    /// Push a value onto the `CompactVec` at the key `query`
    pub fn push_at(&mut self, query: K, item: I) {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                self.values[i].push(item);
                return;
            }
        }
        self.keys.push(query);
        let mut vec = CompactVec::new();
        vec.push(item);
        self.values.push(vec);
    }

    /// Iterator over the `CompactVec` at the key `query`
    pub fn get_iter<'a>(&'a self, query: K) -> impl Iterator<Item = &'a I> + 'a {
        self.get(query).into_iter().flat_map(|vec_in_option| {
            vec_in_option.iter()
        })
    }

    /// Remove the `CompactVec` at the key `query` and iterate over its elements (if it existed)
    pub fn remove_iter<'a>(&'a mut self, query: K) -> impl Iterator<Item = I> + 'a {
        self.remove(query).into_iter().flat_map(|vec_in_option| {
            vec_in_option.into_iter()
        })
    }
}

impl<K: Copy, V: Compact + Clone, A: Allocator> Compact for CompactDict<K, V, A> {
    fn is_still_compact(&self) -> bool {
        self.keys.is_still_compact() && self.values.is_still_compact()
    }

    fn dynamic_size_bytes(&self) -> usize {
        self.keys.dynamic_size_bytes() + self.values.dynamic_size_bytes()
    }

    unsafe fn compact(source: *mut Self, dest: *mut Self, new_dynamic_part: *mut u8) {
        let values_offset = (*source).keys.dynamic_size_bytes() as isize;
        Compact::compact(&mut (*source).keys, &mut (*dest).keys, new_dynamic_part);
        Compact::compact(
            &mut (*source).values,
            &mut (*dest).values,
            new_dynamic_part.offset(values_offset),
        );
    }

    unsafe fn decompact(source: *const Self) -> CompactDict<K, V, A> {
        CompactDict {
            keys: Compact::decompact(&(*source).keys),
            values: Compact::decompact(&(*source).values),
        }
    }
}

impl<K: Copy, V: Compact + Clone, A: Allocator> Clone for CompactDict<K, V, A> {
    fn clone(&self) -> Self {
        CompactDict {
            keys: self.keys.clone(),
            values: self.values.clone(),
        }
    }
}

impl<K: Copy + Eq, V: Compact + Clone, A: Allocator> Default for CompactDict<K, V, A> {
    fn default() -> Self {
        CompactDict::new()
    }
}

impl<K: Copy + Eq, V: Compact + Clone, A: Allocator> ::std::iter::FromIterator<(K, V)>
    for CompactDict<K, V, A> {
    /// Construct a compact dictionary from an interator over key-value pairs
    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
        let mut dict = Self::new();
        for (key, value) in iter {
            dict.insert(key, value);
        }
        dict
    }
}

impl<K: Copy + Eq, V: Compact + Clone, A: Allocator> ::std::iter::Extend<(K, V)>
    for CompactDict<K, V, A> {
    /// Extend a compact dictionary from an iterator over key-value pairs
    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
        for (key, value) in iter {
            self.insert(key, value);
        }
    }
}

#[test]
fn elem(n: usize) -> usize {
    (n * n) as usize
}

#[test]
fn very_basic() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    map.insert(0, 54);
    assert!(*map.get(0).unwrap() == 54);
    map.insert(1, 48);
    assert!(*map.get(1).unwrap() == 48);
}

#[test]
fn very_basic2() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    map.insert(0, 54);
    map.insert(1, 48);
    assert!(*map.get(0).unwrap() == 54);
    assert!(*map.get(1).unwrap() == 48);
}

#[test]
fn basic() {
    let n: usize = 1000;
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    assert!(map.is_empty() == true);
    for i in 0..n {
        let e = elem(i);
        map.insert(i, e);
    }
    assert!(map.is_empty() == false);
    for i in 0..n {
        let test = map.get(i).unwrap();
        let exp = elem(i);
        assert!(*test == exp, " failed exp {:?}  was {:?}", exp, test);
    }
    assert!(map.len() == n as usize);
    assert!(*map.get_mru(n - 1).unwrap() == elem(n - 1));
    assert!(*map.get_mfu(n - 100).unwrap() == elem(n - 100));
    assert!(map.contains_key(n - 300) == true);
    assert!(map.contains_key(n + 1) == false);
    assert!(map.remove(500) == Some(elem(500)));
    assert!(map.get_mru(500).is_none());
}

#[test]
fn iter() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    let n = 10;
    assert!(map.is_empty() == true);
    for i in 0..n {
        map.insert(i, i * i);
    }
    for k in map.keys() {
        println!(" k {:?}", k);
    }
    for h in 0..n {
        let mut keys = map.keys();
        assert!(keys.find(|&i| *i == n - 1 - h).is_some());
    }
    for h in 0..n {
        let mut values = map.values();
        assert!(values.find(|i| **i == elem(h)).is_some());
    }
}

#[test]
fn values_mut() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    assert!(map.is_empty() == true);
    for n in 0..100 {
        map.insert(n, n * n);
    }
    {
        let mut values_mut = map.values_mut();
        for i in &mut values_mut {
            *i = *i + 1;
        }
    }
    for i in 0..100 {
        assert!(*map.get(i).unwrap() == i * i + 1);
    }
}

#[test]
fn pairs() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    assert!(map.is_empty() == true);
    for n in 0..100 {
        map.insert(n, n * n);
    }
    for (key, value) in map.pairs() {
        assert!(elem(*key) == *value);
    }
}

#[test]
fn push_at() {
    let mut map: CompactDict<usize, CompactVec<usize>> = CompactDict::new();
    assert!(map.is_empty() == true);
    for n in 0..100 {
        map.push_at(n, elem(n));
        map.push_at(n, elem(n) + 1);
    }

    for n in 0..100 {
        let mut iter = map.get_iter(n);
        assert!(iter.find(|i| **i == elem(n)).is_some());
        assert!(iter.find(|i| **i == elem(n) + 1).is_some());
    }
}

#[test]
fn remove_iter() {
    let mut map: CompactDict<usize, CompactVec<usize>> = CompactDict::new();
    assert!(map.is_empty() == true);
    for n in 0..100 {
        map.push_at(n, elem(n));
        map.push_at(n, elem(n) + 1);
    }
    let mut iter = map.remove_iter(50);
    assert!(iter.find(|i| *i == elem(50)).is_some());
    assert!(iter.find(|i| *i == elem(50) + 1).is_some());
}
<commit_msg>elem is used only in test, so cfg(test)<commit_after>use super::allocators::{Allocator, DefaultHeap};
use super::compact::Compact;
use super::compact_vec::CompactVec;

/// A simple linear-search key-value dictionary,
/// implemented using two `CompactVec`'s, one for keys, one for values.
///
/// The API loosely follows that of `std::collections::HashMap`.
/// Spilling behaviour using `Allocator` is equivalent to `CompactVec`.
pub struct CompactDict<K: Copy, V: Compact + Clone, A: Allocator = DefaultHeap> {
    keys: CompactVec<K, A>,
    values: CompactVec<V, A>,
}

impl<K: Eq + Copy, V: Compact + Clone, A: Allocator> CompactDict<K, V, A> {
    /// Create new, empty dictionary
    pub fn new() -> Self {
        CompactDict {
            keys: CompactVec::new(),
            values: CompactVec::new(),
        }
    }

    /// Amount of entries in the dictionary
    pub fn len(&self) -> usize {
        self.keys.len()
    }

    /// Is the dictionary empty?
    pub fn is_empty(&self) -> bool {
        self.keys.is_empty()
    }

    /// Look up the value for key `query`, if it exists
    pub fn get(&self, query: K) -> Option<&V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                return Some(&self.values[i]);
            }
        }
        None
    }

    /// Look up the value for keu `query` mutably, if it exists
    pub fn get_mut(&mut self, query: K) -> Option<&mut V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                return Some(&mut self.values[i]);
            }
        }
        None
    }

    /// Lookup up the value for key `query`, if it exists, but also swap the entry
    /// to the beginning of the key/value vectors, so a repeated lookup for that item will be faster
    pub fn get_mru(&mut self, query: K) -> Option<&V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                self.keys.swap(0, i);
                self.values.swap(0, i);
                return Some(&self.values[0]);
            }
        }
        None
    }

    /// Lookup up the value for key `query`, if it exists, but also swap the entry
    /// one index towards the beginning of the key/value vectors, so frequently repeated lookups
    /// for that item will be faster
    pub fn get_mfu(&mut self, query: K) -> Option<&V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                if i > 0 {
                    self.keys.swap(i - 1, i);
                    self.values.swap(i - 1, i);
                    return Some(&self.values[i - 1]);
                } else {
                    return Some(&self.values[0]);
                }
            }
        }
        None
    }

    /// Does the dictionary contain a value for `query`?
    pub fn contains_key(&self, query: K) -> bool {
        self.keys.contains(&query)
    }

    /// Insert new value at key `query` and return the previous value at that key, if any existed
    pub fn insert(&mut self, query: K, new_value: V) -> Option<V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                let old_val = self.values[i].clone();
                self.values[i] = new_value;
                return Some(old_val);
            }
        }
        self.keys.push(query);
        self.values.push(new_value);
        None
    }

    /// Remove value at key `query` and return it, if it existed
    pub fn remove(&mut self, query: K) -> Option<V> {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                let old_val = self.values[i].clone();
                self.keys.remove(i);
                self.values.remove(i);
                return Some(old_val);
            }
        }
        None
    }

    /// Iterator over all keys in the dictionary
    pub fn keys(&self) -> ::std::slice::Iter<K> {
        self.keys.iter()
    }

    /// Iterator over all values in the dictionary
    pub fn values(&self) -> ::std::slice::Iter<V> {
        self.values.iter()
    }

    /// Iterator over mutable references to all values in the dictionary
    pub fn values_mut(&mut self) -> ::std::slice::IterMut<V> {
        self.values.iter_mut()
    }

    /// Iterator over all key-value pairs in the dictionary
    pub fn pairs<'a>(&'a self) -> impl Iterator<Item = (&'a K, &'a V)> + Clone + 'a {
        self.keys().zip(self.values())
    }
}

impl<K: Eq + Copy, I: Compact, A1: Allocator, A2: Allocator> CompactDict<K, CompactVec<I, A1>, A2> {
    /// Push a value onto the `CompactVec` at the key `query`
    pub fn push_at(&mut self, query: K, item: I) {
        for i in 0..self.keys.len() {
            if self.keys[i] == query {
                self.values[i].push(item);
                return;
            }
        }
        self.keys.push(query);
        let mut vec = CompactVec::new();
        vec.push(item);
        self.values.push(vec);
    }

    /// Iterator over the `CompactVec` at the key `query`
    pub fn get_iter<'a>(&'a self, query: K) -> impl Iterator<Item = &'a I> + 'a {
        self.get(query).into_iter().flat_map(|vec_in_option| {
            vec_in_option.iter()
        })
    }

    /// Remove the `CompactVec` at the key `query` and iterate over its elements (if it existed)
    pub fn remove_iter<'a>(&'a mut self, query: K) -> impl Iterator<Item = I> + 'a {
        self.remove(query).into_iter().flat_map(|vec_in_option| {
            vec_in_option.into_iter()
        })
    }
}

impl<K: Copy, V: Compact + Clone, A: Allocator> Compact for CompactDict<K, V, A> {
    fn is_still_compact(&self) -> bool {
        self.keys.is_still_compact() && self.values.is_still_compact()
    }

    fn dynamic_size_bytes(&self) -> usize {
        self.keys.dynamic_size_bytes() + self.values.dynamic_size_bytes()
    }

    unsafe fn compact(source: *mut Self, dest: *mut Self, new_dynamic_part: *mut u8) {
        let values_offset = (*source).keys.dynamic_size_bytes() as isize;
        Compact::compact(&mut (*source).keys, &mut (*dest).keys, new_dynamic_part);
        Compact::compact(
            &mut (*source).values,
            &mut (*dest).values,
            new_dynamic_part.offset(values_offset),
        );
    }

    unsafe fn decompact(source: *const Self) -> CompactDict<K, V, A> {
        CompactDict {
            keys: Compact::decompact(&(*source).keys),
            values: Compact::decompact(&(*source).values),
        }
    }
}

impl<K: Copy, V: Compact + Clone, A: Allocator> Clone for CompactDict<K, V, A> {
    fn clone(&self) -> Self {
        CompactDict {
            keys: self.keys.clone(),
            values: self.values.clone(),
        }
    }
}

impl<K: Copy + Eq, V: Compact + Clone, A: Allocator> Default for CompactDict<K, V, A> {
    fn default() -> Self {
        CompactDict::new()
    }
}

impl<K: Copy + Eq, V: Compact + Clone, A: Allocator> ::std::iter::FromIterator<(K, V)>
    for CompactDict<K, V, A> {
    /// Construct a compact dictionary from an interator over key-value pairs
    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
        let mut dict = Self::new();
        for (key, value) in iter {
            dict.insert(key, value);
        }
        dict
    }
}

impl<K: Copy + Eq, V: Compact + Clone, A: Allocator> ::std::iter::Extend<(K, V)>
    for CompactDict<K, V, A> {
    /// Extend a compact dictionary from an iterator over key-value pairs
    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
        for (key, value) in iter {
            self.insert(key, value);
        }
    }
}

#[cfg(test)]
fn elem(n: usize) -> usize {
    (n * n) as usize
}

#[test]
fn very_basic() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    map.insert(0, 54);
    assert!(*map.get(0).unwrap() == 54);
    map.insert(1, 48);
    assert!(*map.get(1).unwrap() == 48);
}

#[test]
fn very_basic2() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    map.insert(0, 54);
    map.insert(1, 48);
    assert!(*map.get(0).unwrap() == 54);
    assert!(*map.get(1).unwrap() == 48);
}

#[test]
fn basic() {
    let n: usize = 1000;
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    assert!(map.is_empty() == true);
    for i in 0..n {
        let e = elem(i);
        map.insert(i, e);
    }
    assert!(map.is_empty() == false);
    for i in 0..n {
        let test = map.get(i).unwrap();
        let exp = elem(i);
        assert!(*test == exp, " failed exp {:?}  was {:?}", exp, test);
    }
    assert!(map.len() == n as usize);
    assert!(*map.get_mru(n - 1).unwrap() == elem(n - 1));
    assert!(*map.get_mfu(n - 100).unwrap() == elem(n - 100));
    assert!(map.contains_key(n - 300) == true);
    assert!(map.contains_key(n + 1) == false);
    assert!(map.remove(500) == Some(elem(500)));
    assert!(map.get_mru(500).is_none());
}

#[test]
fn iter() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    let n = 10;
    assert!(map.is_empty() == true);
    for i in 0..n {
        map.insert(i, i * i);
    }
    for k in map.keys() {
        println!(" k {:?}", k);
    }
    for h in 0..n {
        let mut keys = map.keys();
        assert!(keys.find(|&i| *i == n - 1 - h).is_some());
    }
    for h in 0..n {
        let mut values = map.values();
        assert!(values.find(|i| **i == elem(h)).is_some());
    }
}

#[test]
fn values_mut() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    assert!(map.is_empty() == true);
    for n in 0..100 {
        map.insert(n, n * n);
    }
    {
        let mut values_mut = map.values_mut();
        for i in &mut values_mut {
            *i = *i + 1;
        }
    }
    for i in 0..100 {
        assert!(*map.get(i).unwrap() == i * i + 1);
    }
}

#[test]
fn pairs() {
    let mut map: CompactDict<usize, usize> = CompactDict::new();
    assert!(map.is_empty() == true);
    for n in 0..100 {
        map.insert(n, n * n);
    }
    for (key, value) in map.pairs() {
        assert!(elem(*key) == *value);
    }
}

#[test]
fn push_at() {
    let mut map: CompactDict<usize, CompactVec<usize>> = CompactDict::new();
    assert!(map.is_empty() == true);
    for n in 0..100 {
        map.push_at(n, elem(n));
        map.push_at(n, elem(n) + 1);
    }

    for n in 0..100 {
        let mut iter = map.get_iter(n);
        assert!(iter.find(|i| **i == elem(n)).is_some());
        assert!(iter.find(|i| **i == elem(n) + 1).is_some());
    }
}

#[test]
fn remove_iter() {
    let mut map: CompactDict<usize, CompactVec<usize>> = CompactDict::new();
    assert!(map.is_empty() == true);
    for n in 0..100 {
        map.push_at(n, elem(n));
        map.push_at(n, elem(n) + 1);
    }
    let mut iter = map.remove_iter(50);
    assert!(iter.find(|i| *i == elem(50)).is_some());
    assert!(iter.find(|i| *i == elem(50) + 1).is_some());
}
<|endoftext|>


<commit_before>//! A small GPU compute shader demonstration.
//!
//! Here we use a compute shader to calculate the amplitude of `OSCILLATOR_COUNT` number of
//! oscillators. The oscillator amplitudes are then laid out across the screen using rectangles
//! with a gray value equal to the amplitude. Real-time interaction is demonstrated by providing
//! access to time, frequency (mouse `x`) and the number of oscillators via uniform data.

use nannou::prelude::*;
use nannou::wgpu::BufferInitDescriptor;
use std::sync::{Arc, Mutex};

struct Model {
    compute: Compute,
    oscillators: Arc<Mutex<Vec<f32>>>,
    threadpool: futures::executor::ThreadPool,
}

struct Compute {
    oscillator_buffer: wgpu::Buffer,
    oscillator_buffer_size: wgpu::BufferAddress,
    uniform_buffer: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
    pipeline: wgpu::ComputePipeline,
}

#[repr(C)]
#[derive(Copy, Clone)]
pub struct Uniforms {
    time: f32,
    freq: f32,
    oscillator_count: u32,
}

const OSCILLATOR_COUNT: u32 = 128;

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

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

    // Create the compute shader module.
    let cs_mod = wgpu::shader_from_spirv_bytes(device, include_bytes!("shaders/comp.spv"));

    // Create the buffer that will store the result of our compute operation.
    let oscillator_buffer_size =
        (OSCILLATOR_COUNT as usize * std::mem::size_of::<f32>()) as wgpu::BufferAddress;
    let oscillator_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("oscillators"),
        size: oscillator_buffer_size,
        usage: wgpu::BufferUsage::STORAGE
            | wgpu::BufferUsage::COPY_DST
            | wgpu::BufferUsage::COPY_SRC,
        mapped_at_creation: true,
    });

    // Create the buffer that will store time.
    let uniforms = create_uniforms(app.time, app.mouse.x, window.rect());
    let uniforms_bytes = uniforms_as_bytes(&uniforms);
    let usage = wgpu::BufferUsage::UNIFORM | wgpu::BufferUsage::COPY_DST;
    let uniform_buffer = device.create_buffer_init(&BufferInitDescriptor {
        label: None,
        contents: uniforms_bytes,
        usage,
    });

    // Create the bind group and pipeline.
    let bind_group_layout = create_bind_group_layout(device);
    let bind_group = create_bind_group(
        device,
        &bind_group_layout,
        &oscillator_buffer,
        oscillator_buffer_size,
        &uniform_buffer,
    );
    let pipeline_layout = create_pipeline_layout(device, &bind_group_layout);
    let pipeline = create_compute_pipeline(device, &pipeline_layout, &cs_mod);

    let compute = Compute {
        oscillator_buffer,
        oscillator_buffer_size,
        uniform_buffer,
        bind_group,
        pipeline,
    };

    // The vector that we will write oscillator values to.
    let oscillators = Arc::new(Mutex::new(vec![0.0; OSCILLATOR_COUNT as usize]));

    // Create a thread pool capable of running our GPU buffer read futures.
    let threadpool = futures::executor::ThreadPool::new().unwrap();

    Model {
        compute,
        oscillators,
        threadpool,
    }
}

fn update(app: &App, model: &mut Model, _update: Update) {
    let window = app.main_window();
    let device = window.swap_chain_device();
    let win_rect = window.rect();
    let compute = &mut model.compute;

    // The buffer into which we'll read some data.
    let read_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("read_oscillators"),
        size: compute.oscillator_buffer_size,
        usage: wgpu::BufferUsage::MAP_READ
            | wgpu::BufferUsage::COPY_DST
            | wgpu::BufferUsage::COPY_SRC,
        mapped_at_creation: true,
    });

    // An update for the uniform buffer with the current time.
    let uniforms = create_uniforms(app.time, app.mouse.x, win_rect);
    let uniforms_size = std::mem::size_of::<Uniforms>() as wgpu::BufferAddress;
    let uniforms_bytes = uniforms_as_bytes(&uniforms);
    let usage = wgpu::BufferUsage::COPY_SRC;
    let new_uniform_buffer = device.create_buffer_init(&BufferInitDescriptor {
        label: None,
        contents: uniforms_bytes,
        usage,
    });

    // The encoder we'll use to encode the compute pass.
    let desc = wgpu::CommandEncoderDescriptor {
        label: Some("oscillator_compute"),
    };
    let mut encoder = device.create_command_encoder(&desc);
    encoder.copy_buffer_to_buffer(
        &new_uniform_buffer,
        0,
        &compute.uniform_buffer,
        0,
        uniforms_size,
    );
    {
        let mut cpass = encoder.begin_compute_pass();
        cpass.set_pipeline(&compute.pipeline);
        cpass.set_bind_group(0, &compute.bind_group, &[]);
        cpass.dispatch(OSCILLATOR_COUNT as u32, 1, 1);
    }
    encoder.copy_buffer_to_buffer(
        &compute.oscillator_buffer,
        0,
        &read_buffer,
        0,
        compute.oscillator_buffer_size,
    );

    // Submit the compute pass to the device's queue.
    window
        .swap_chain_queue()
        .submit(std::iter::once(encoder.finish()));

    // Spawn a future that reads the result of the compute pass.
    let oscillators = model.oscillators.clone();
    let future = async move {
        let slice = read_buffer.slice(..);
        if let Ok(_) = slice.map_async(wgpu::MapMode::Read).await {
            if let Ok(mut oscillators) = oscillators.lock() {
                let bytes = &slice.get_mapped_range()[..];
                // "Cast" the slice of bytes to a slice of floats as required.
                let floats = {
                    let len = bytes.len() / std::mem::size_of::<f32>();
                    let ptr = bytes.as_ptr() as *const f32;
                    unsafe { std::slice::from_raw_parts(ptr, len) }
                };
                oscillators.copy_from_slice(floats);
            }
        }
    };
    model.threadpool.spawn_ok(future);

    // Check for resource cleanups and mapping callbacks.
    //
    // Note that this line is not necessary in our case, as the device we are using already gets
    // polled when nannou submits the command buffer for drawing and presentation after `view`
    // completes. If we were to use a standalone device to create our buffer and perform our
    // compute (rather than the device requested during window creation), calling `poll` regularly
    // would be a must.
    //
    // device.poll(false);
}

fn view(app: &App, model: &Model, frame: Frame) {
    frame.clear(BLACK);
    let draw = app.draw();
    let window = app.window(frame.window_id()).unwrap();
    let rect = window.rect();

    if let Ok(oscillators) = model.oscillators.lock() {
        let w = rect.w() / OSCILLATOR_COUNT as f32;
        let h = rect.h();
        let half_w = w * 0.5;
        for (i, &osc) in oscillators.iter().enumerate() {
            let x = half_w + map_range(i as u32, 0, OSCILLATOR_COUNT, rect.left(), rect.right());
            draw.rect().w_h(w, h).x(x).color(gray(osc));
        }
    }

    draw.to_frame(app, &frame).unwrap();
}

fn create_uniforms(time: f32, mouse_x: f32, win_rect: geom::Rect) -> Uniforms {
    let freq = map_range(
        mouse_x,
        win_rect.left(),
        win_rect.right(),
        0.0,
        win_rect.w(),
    );
    let oscillator_count = OSCILLATOR_COUNT;
    Uniforms {
        time,
        freq,
        oscillator_count,
    }
}

fn create_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    let storage_dynamic = false;
    let storage_readonly = false;
    let uniform_dynamic = false;
    wgpu::BindGroupLayoutBuilder::new()
        .storage_buffer(
            wgpu::ShaderStage::COMPUTE,
            storage_dynamic,
            storage_readonly,
        )
        .uniform_buffer(wgpu::ShaderStage::COMPUTE, uniform_dynamic)
        .build(device)
}

fn create_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    oscillator_buffer: &wgpu::Buffer,
    oscillator_buffer_size: wgpu::BufferAddress,
    uniform_buffer: &wgpu::Buffer,
) -> wgpu::BindGroup {
    wgpu::BindGroupBuilder::new()
        .buffer_bytes(oscillator_buffer, 0..oscillator_buffer_size)
        .buffer::<Uniforms>(uniform_buffer, 0..1)
        .build(device, layout)
}

fn create_pipeline_layout(
    device: &wgpu::Device,
    bind_group_layout: &wgpu::BindGroupLayout,
) -> wgpu::PipelineLayout {
    device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: Some("nannou"),
        bind_group_layouts: &[&bind_group_layout],
        push_constant_ranges: &[],
    })
}

fn create_compute_pipeline(
    device: &wgpu::Device,
    layout: &wgpu::PipelineLayout,
    cs_mod: &wgpu::ShaderModule,
) -> wgpu::ComputePipeline {
    let compute_stage = wgpu::ProgrammableStageDescriptor {
        module: &cs_mod,
        entry_point: "main",
    };
    let desc = wgpu::ComputePipelineDescriptor {
        label: Some("nannou"),
        layout: Some(layout),
        compute_stage,
    };
    device.create_compute_pipeline(&desc)
}

// See `nannou::wgpu::bytes` docs for why these are necessary.

fn uniforms_as_bytes(uniforms: &Uniforms) -> &[u8] {
    unsafe { wgpu::bytes::from(uniforms) }
}
<commit_msg>Fixes a `panic` in wgpu_compute_shader.rs example<commit_after>//! A small GPU compute shader demonstration.
//!
//! Here we use a compute shader to calculate the amplitude of `OSCILLATOR_COUNT` number of
//! oscillators. The oscillator amplitudes are then laid out across the screen using rectangles
//! with a gray value equal to the amplitude. Real-time interaction is demonstrated by providing
//! access to time, frequency (mouse `x`) and the number of oscillators via uniform data.

use nannou::prelude::*;
use nannou::wgpu::BufferInitDescriptor;
use std::sync::{Arc, Mutex};

struct Model {
    compute: Compute,
    oscillators: Arc<Mutex<Vec<f32>>>,
    threadpool: futures::executor::ThreadPool,
}

struct Compute {
    oscillator_buffer: wgpu::Buffer,
    oscillator_buffer_size: wgpu::BufferAddress,
    uniform_buffer: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
    pipeline: wgpu::ComputePipeline,
}

#[repr(C)]
#[derive(Copy, Clone)]
pub struct Uniforms {
    time: f32,
    freq: f32,
    oscillator_count: u32,
}

const OSCILLATOR_COUNT: u32 = 128;

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

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

    // Create the compute shader module.
    let cs_mod = wgpu::shader_from_spirv_bytes(device, include_bytes!("shaders/comp.spv"));

    // Create the buffer that will store the result of our compute operation.
    let oscillator_buffer_size =
        (OSCILLATOR_COUNT as usize * std::mem::size_of::<f32>()) as wgpu::BufferAddress;
    let oscillator_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("oscillators"),
        size: oscillator_buffer_size,
        usage: wgpu::BufferUsage::STORAGE
            | wgpu::BufferUsage::COPY_DST
            | wgpu::BufferUsage::COPY_SRC,
        mapped_at_creation: false,
    });

    // Create the buffer that will store time.
    let uniforms = create_uniforms(app.time, app.mouse.x, window.rect());
    let uniforms_bytes = uniforms_as_bytes(&uniforms);
    let usage = wgpu::BufferUsage::UNIFORM | wgpu::BufferUsage::COPY_DST;
    let uniform_buffer = device.create_buffer_init(&BufferInitDescriptor {
        label: None,
        contents: uniforms_bytes,
        usage,
    });

    // Create the bind group and pipeline.
    let bind_group_layout = create_bind_group_layout(device);
    let bind_group = create_bind_group(
        device,
        &bind_group_layout,
        &oscillator_buffer,
        oscillator_buffer_size,
        &uniform_buffer,
    );
    let pipeline_layout = create_pipeline_layout(device, &bind_group_layout);
    let pipeline = create_compute_pipeline(device, &pipeline_layout, &cs_mod);

    let compute = Compute {
        oscillator_buffer,
        oscillator_buffer_size,
        uniform_buffer,
        bind_group,
        pipeline,
    };

    // The vector that we will write oscillator values to.
    let oscillators = Arc::new(Mutex::new(vec![0.0; OSCILLATOR_COUNT as usize]));

    // Create a thread pool capable of running our GPU buffer read futures.
    let threadpool = futures::executor::ThreadPool::new().unwrap();

    Model {
        compute,
        oscillators,
        threadpool,
    }
}

fn update(app: &App, model: &mut Model, _update: Update) {
    let window = app.main_window();
    let device = window.swap_chain_device();
    let win_rect = window.rect();
    let compute = &mut model.compute;

    // The buffer into which we'll read some data.
    let read_buffer = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("read_oscillators"),
        size: compute.oscillator_buffer_size,
        usage: wgpu::BufferUsage::MAP_READ | wgpu::BufferUsage::COPY_DST,
        mapped_at_creation: false,
    });

    // An update for the uniform buffer with the current time.
    let uniforms = create_uniforms(app.time, app.mouse.x, win_rect);
    let uniforms_size = std::mem::size_of::<Uniforms>() as wgpu::BufferAddress;
    let uniforms_bytes = uniforms_as_bytes(&uniforms);
    let usage = wgpu::BufferUsage::COPY_SRC;
    let new_uniform_buffer = device.create_buffer_init(&BufferInitDescriptor {
        label: None,
        contents: uniforms_bytes,
        usage,
    });

    // The encoder we'll use to encode the compute pass.
    let desc = wgpu::CommandEncoderDescriptor {
        label: Some("oscillator_compute"),
    };
    let mut encoder = device.create_command_encoder(&desc);
    encoder.copy_buffer_to_buffer(
        &new_uniform_buffer,
        0,
        &compute.uniform_buffer,
        0,
        uniforms_size,
    );
    {
        let mut cpass = encoder.begin_compute_pass();
        cpass.set_pipeline(&compute.pipeline);
        cpass.set_bind_group(0, &compute.bind_group, &[]);
        cpass.dispatch(OSCILLATOR_COUNT as u32, 1, 1);
    }
    encoder.copy_buffer_to_buffer(
        &compute.oscillator_buffer,
        0,
        &read_buffer,
        0,
        compute.oscillator_buffer_size,
    );

    // Submit the compute pass to the device's queue.
    window
        .swap_chain_queue()
        .submit(std::iter::once(encoder.finish()));

    // Spawn a future that reads the result of the compute pass.
    let oscillators = model.oscillators.clone();
    let future = async move {
        let slice = read_buffer.slice(..);
        if let Ok(_) = slice.map_async(wgpu::MapMode::Read).await {
            if let Ok(mut oscillators) = oscillators.lock() {
                let bytes = &slice.get_mapped_range()[..];
                // "Cast" the slice of bytes to a slice of floats as required.
                let floats = {
                    let len = bytes.len() / std::mem::size_of::<f32>();
                    let ptr = bytes.as_ptr() as *const f32;
                    unsafe { std::slice::from_raw_parts(ptr, len) }
                };
                oscillators.copy_from_slice(floats);
            }
        }
    };
    model.threadpool.spawn_ok(future);

    // Check for resource cleanups and mapping callbacks.
    //
    // Note that this line is not necessary in our case, as the device we are using already gets
    // polled when nannou submits the command buffer for drawing and presentation after `view`
    // completes. If we were to use a standalone device to create our buffer and perform our
    // compute (rather than the device requested during window creation), calling `poll` regularly
    // would be a must.
    //
    // device.poll(false);
}

fn view(app: &App, model: &Model, frame: Frame) {
    frame.clear(BLACK);
    let draw = app.draw();
    let window = app.window(frame.window_id()).unwrap();
    let rect = window.rect();

    if let Ok(oscillators) = model.oscillators.lock() {
        let w = rect.w() / OSCILLATOR_COUNT as f32;
        let h = rect.h();
        let half_w = w * 0.5;
        for (i, &osc) in oscillators.iter().enumerate() {
            let x = half_w + map_range(i as u32, 0, OSCILLATOR_COUNT, rect.left(), rect.right());
            draw.rect().w_h(w, h).x(x).color(gray(osc));
        }
    }

    draw.to_frame(app, &frame).unwrap();
}

fn create_uniforms(time: f32, mouse_x: f32, win_rect: geom::Rect) -> Uniforms {
    let freq = map_range(
        mouse_x,
        win_rect.left(),
        win_rect.right(),
        0.0,
        win_rect.w(),
    );
    let oscillator_count = OSCILLATOR_COUNT;
    Uniforms {
        time,
        freq,
        oscillator_count,
    }
}

fn create_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    let storage_dynamic = false;
    let storage_readonly = false;
    let uniform_dynamic = false;
    wgpu::BindGroupLayoutBuilder::new()
        .storage_buffer(
            wgpu::ShaderStage::COMPUTE,
            storage_dynamic,
            storage_readonly,
        )
        .uniform_buffer(wgpu::ShaderStage::COMPUTE, uniform_dynamic)
        .build(device)
}

fn create_bind_group(
    device: &wgpu::Device,
    layout: &wgpu::BindGroupLayout,
    oscillator_buffer: &wgpu::Buffer,
    oscillator_buffer_size: wgpu::BufferAddress,
    uniform_buffer: &wgpu::Buffer,
) -> wgpu::BindGroup {
    wgpu::BindGroupBuilder::new()
        .buffer_bytes(oscillator_buffer, 0..oscillator_buffer_size)
        .buffer::<Uniforms>(uniform_buffer, 0..1)
        .build(device, layout)
}

fn create_pipeline_layout(
    device: &wgpu::Device,
    bind_group_layout: &wgpu::BindGroupLayout,
) -> wgpu::PipelineLayout {
    device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: Some("nannou"),
        bind_group_layouts: &[&bind_group_layout],
        push_constant_ranges: &[],
    })
}

fn create_compute_pipeline(
    device: &wgpu::Device,
    layout: &wgpu::PipelineLayout,
    cs_mod: &wgpu::ShaderModule,
) -> wgpu::ComputePipeline {
    let compute_stage = wgpu::ProgrammableStageDescriptor {
        module: &cs_mod,
        entry_point: "main",
    };
    let desc = wgpu::ComputePipelineDescriptor {
        label: Some("nannou"),
        layout: Some(layout),
        compute_stage,
    };
    device.create_compute_pipeline(&desc)
}

// See `nannou::wgpu::bytes` docs for why these are necessary.

fn uniforms_as_bytes(uniforms: &Uniforms) -> &[u8] {
    unsafe { wgpu::bytes::from(uniforms) }
}
<|endoftext|>