<commit_before>/// Custom-derive providing macros for the field-by-field Traits
///
/// See the tests/derive.rs file for examples of use.

extern crate field_by_field;
extern crate proc_macro;
#[macro_use]
extern crate quote;
extern crate syn;

use proc_macro::TokenStream;

#[proc_macro_derive(FieldByField)]
pub fn derive(input: TokenStream) -> TokenStream {
    expand_derive(input, build_trait_field_by_field)
}

fn expand_derive(input: TokenStream, f: fn(syn::MacroInput) -> quote::Tokens) -> TokenStream {
    let item = syn::parse_macro_input(&input.to_string()).unwrap();
    f(item).to_string().parse().unwrap()
}

fn build_trait_field_by_field(input: syn::MacroInput) -> quote::Tokens {
    let name = &input.ident;

    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

    match input.body {
        syn::Body::Struct(ref data) => {
            let fn_fields_not_equal = build_fn_fields_not_equal(&data.fields());
            let fn_assert_equal_field_by_field = build_fn_assert_equal_field_by_field(&name);

            quote! {
                impl #impl_generics ::field_by_field::EqualFieldByField
                for #name #ty_generics #where_clause {
                    #fn_fields_not_equal

                    #fn_assert_equal_field_by_field
                }
            }
        }
        syn::Body::Enum(ref data) => {
            let fn_assert = build_fn_assert_equal_field_by_field(&name);
            let is_multivariant = data.len() > 1;
            let variants = data.iter()
                .map(|var| {
                    let var_name = &var.ident;
                    match var.data {
                        syn::VariantData::Unit =>
                            build_match_unit_variant(&name, &var_name, is_multivariant),
                        syn::VariantData::Tuple(ref fields) =>
                            build_match_tuple_variant(&name, &var_name, fields, is_multivariant),
                        syn::VariantData::Struct(ref fields) =>
                            build_match_struct_variant(&name, &var_name, fields, is_multivariant),
                    }
                }
            );

            quote! {
                impl #impl_generics ::field_by_field::EqualFieldByField
                for #name #ty_generics #where_clause {

                    fn fields_not_equal(&self, other: &Self)
                    -> Vec<::field_by_field::UnequalField> {
                        // This is never modified in the single-variant enum case.
                        #[allow(unused_mut)]
                        let mut list: Vec<::field_by_field::UnequalField> = Vec::new();

                        match (self, other) {
                            #(#variants)*
                        }

                        list
                    }

                    #fn_assert
                }
            }
        }
    }
}

/// Build a function that compares all the items in a simple struct
///
/// This emits just a long list of `if self.name != other.name {
/// vec.push(UnequalField); }` tokens.
fn build_fn_fields_not_equal(fields: &[syn::Field]) -> quote::Tokens {
    let find_unequal_fields = fields.iter().map(|f| {
        let f_name = &f.ident;
        let f_str = f_name.as_ref().map(|v| v.to_string())
            .expect(&format!("Couldn't convert field to str: {:?}", f));
        quote! {
            if self.#f_name != other.#f_name {
                list.push(::field_by_field::UnequalField {
                    field_name: #f_str.to_string(),
                    actually: Box::new(self.#f_name.clone()),
                    expected: Box::new(other.#f_name.clone()),
                });
            }
        }
    });

    quote! {
        fn fields_not_equal(&self, other: &Self)
                            -> Vec<::field_by_field::UnequalField>
        {
            let mut list: Vec<::field_by_field::UnequalField> = Vec::new();

            #(#find_unequal_fields)*

            list
        }
    }
}

/// Build a match statement that compares the self variant to the other variant
///
/// Since this is for unit variants this doesn't check the actual value if it
/// isn't an exact match.
fn build_match_unit_variant(name: &syn::Ident, var_name: &syn::Ident, is_multivariant: bool)
-> quote::Tokens {
    let left_str = format!("{}::{}", &name, &var_name);

    if is_multivariant {
        quote! {
            ( &#name::#var_name, &#name::#var_name ) => {}
            ( &#name::#var_name, ref expected ) => {
                list.push(::field_by_field::UnequalField {
                    field_name: #left_str.to_string(),
                    actually: Box::new(#left_str.to_string()),
                    expected: Box::new(format!("{:?}", &expected)),
                });
            }
        }
    } else {
        // TODO: This causes a "variable does not need to be mutable" warning
        quote! {
            ( &#name::#var_name, &#name::#var_name ) => {}
        }
    }
}

/// Build a match statement that compares a tuple enum against others
fn build_match_tuple_variant(name: &syn::Ident,
                             var_name: &syn::Ident,
                             var_fields: &[syn::Field],
                             is_multivariant: bool)
-> quote::Tokens {
    let actually_fields: Vec<_> = var_fields.iter()
        .enumerate()
        .map(|(i, f)| {
            assert!(f.ident.is_none(),
                    format!("Fields in tuples should be unnamed, not {:?}", f.ident));
            syn::Ident::from(format!("_{}", i))
        })
        .collect();
    let actually_field_refs: Vec<_> = actually_fields.iter()
        .map(|name| quote! { ref #name }).collect();
    let expected_fields: Vec<_> = var_fields.iter()
        .enumerate()
        .map(|(i, f)| {
            assert!(f.ident.is_none(),
                    format!("Fields in tuples should be unnamed, not {:?}", f.ident));
            syn::Ident::from(format!("_e{}", i))
        })
        .collect();
    let expected_field_refs: Vec<_> = expected_fields.iter()
        .map(|name| quote! { ref #name }).collect();

    let name_str = name.to_string();
    let var_name_str = var_name.to_string();

    let comparisons: Vec<_> = actually_fields.iter()
        .zip(&expected_fields)
        .enumerate()
        .map(|(i, (actually, expected))| {
             let field_name = format!("{}::{}.{}", name_str, var_name_str, i);
             quote! {
                 if #actually != #expected {
                     list.push(::field_by_field::UnequalField {
                         field_name: #field_name.into(),
                         actually: Box::new(#actually.clone()),
                         expected: Box::new(#expected.clone()),
                     });
                 }
             }
        })
        .collect();

    let field_match = quote! {
        ( &#name::#var_name(#(#actually_field_refs),*),
          &#name::#var_name(#(#expected_field_refs),*) ) => {
            #(#comparisons)*
        }
    };

    if is_multivariant {
        let var_field_name = format!("{}::{}", name_str, var_name_str);
        // TODO: improve error messages if the variants differ
        quote! {
            #field_match
            (ref actually @ &#name::#var_name(..), ref expected) => {
                list.push(::field_by_field::UnequalField {
                    field_name: #var_field_name.into(),
                    actually: Box::new((*actually).clone()),
                    expected: Box::new((*expected).clone()),
                });
            }
        }
    } else {
        field_match
    }
}

/// Build a match arm that compares structs variants against themselves or other
fn build_match_struct_variant(name: &syn::Ident,
                              var_name: &syn::Ident,
                              fields: &[syn::Field],
                              is_multivariant: bool)
-> quote::Tokens {
    let field_names = fields.iter().cloned()
        .map(|field| {
             let ident = field.ident;
             ident
                .unwrap_or_else(|| panic!("Unable to get name for field in struct-like enum"))
        })
        .collect::<Vec<_>>();
    let expected_names = field_names.iter()
        .map(|name| format!("expected_{}", name).into())
        .collect::<Vec<syn::Ident>>();
    let expected_name_bindings = field_names.iter().zip(&expected_names)
        .map(|(name, expected_name)| {
            quote! { #name: ref #expected_name }
        })
        .collect::<Vec<_>>();
    let comparisons = field_names.iter()
        .zip(&expected_names)
        .map(|(name, other_name)| {
            let vname = name.as_ref();
            quote! {
                if #name != #other_name {
                    list.push(::field_by_field::UnequalField {
                        field_name: #vname.into(),
                        actually: Box::new(#name.clone()),
                        expected: Box::new(#other_name.clone()),
                    })
                }
            }
        })
        .collect::<Vec<_>>();
    let by_field_match = quote! {
        (&#name::#var_name { #(ref #field_names),* },
         &#name::#var_name { #(#expected_name_bindings,)* }) => {
            #(#comparisons)*
        }
    };

    if is_multivariant {
        let vname = format!("{}::{}", name, var_name);
        quote! {
            #by_field_match
            (ref actually @ &#name::#var_name { .. }, expected) => {
                list.push(::field_by_field::UnequalField {
                    field_name: #vname.into(),
                    actually: Box::new((*actually).clone()),
                    expected: Box::new((*expected).clone()),
                })
            }
        }
    } else {
        by_field_match
    }
}

/// Build a function that panics if the result of fiels_not_equal is non-empty
///
/// The Function will panic with a nice error message
fn build_fn_assert_equal_field_by_field(name: &syn::Ident) -> quote::Tokens {
    let overall_errmsgfmt = quote! {
        let ac_exp = format!("    actually: {:?}\n    \
                              expected: {:?}\n", self, other);
        errmsg.push_str(&ac_exp);
    };

    let field_errmsgfmt = quote! {
        errmsg.push_str(&format!("        {}: {:?} != {:?}\n",
                                 field_err.field_name,
                                 field_err.actually, field_err.expected));

    };

    quote! {
        fn assert_equal_field_by_field(&self, other: &#name) {
            let errs = self.fields_not_equal(other);

            if errs.len() > 0 {
                let mut errmsg = String::from("\n    Items are not equal:\n");
                for field_err in errs {
                    #field_errmsgfmt
                }

                #overall_errmsgfmt

                panic!("{}", errmsg);
            }
        }
    }
}
<commit_msg>Dedupe some code between enum and struct cases<commit_after>/// Custom-derive providing macros for the field-by-field Traits
///
/// See the tests/derive.rs file for examples of use.

extern crate field_by_field;
extern crate proc_macro;
#[macro_use]
extern crate quote;
extern crate syn;

use proc_macro::TokenStream;

#[proc_macro_derive(FieldByField)]
pub fn derive(input: TokenStream) -> TokenStream {
    expand_derive(input, build_trait_field_by_field)
}

fn expand_derive(input: TokenStream, f: fn(syn::MacroInput) -> quote::Tokens) -> TokenStream {
    let item = syn::parse_macro_input(&input.to_string()).unwrap();
    f(item).to_string().parse().unwrap()
}

fn build_trait_field_by_field(input: syn::MacroInput) -> quote::Tokens {
    let name = &input.ident;

    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

    let fn_assert_equal_field_by_field = build_fn_assert_equal_field_by_field(&name);

    let fn_fields_not_equal = match input.body {
        syn::Body::Struct(ref data) => build_fn_fields_not_equal(&data.fields()),
        syn::Body::Enum(ref data) => build_fn_fields_not_equal_enum(&name, &data)
    };
    quote! {
        impl #impl_generics ::field_by_field::EqualFieldByField
        for #name #ty_generics #where_clause {

            #fn_fields_not_equal

            #fn_assert_equal_field_by_field
        }
    }
}

/// Build a function that compares all the items in a simple struct
///
/// This emits just a long list of `if self.name != other.name {
/// vec.push(UnequalField); }` tokens.
fn build_fn_fields_not_equal(fields: &[syn::Field]) -> quote::Tokens {
    let find_unequal_fields = fields.iter()
        .map(|f| {
            let f_name = &f.ident;
            let f_str = f_name.as_ref().map(|v| v.to_string())
                .expect(&format!("Couldn't convert field to str: {:?}", f));
            quote! {
                if self.#f_name != other.#f_name {
                    list.push(::field_by_field::UnequalField {
                        field_name: #f_str.to_string(),
                        actually: Box::new(self.#f_name.clone()),
                        expected: Box::new(other.#f_name.clone()),
                    });
                }
            }
        });

    quote! {
        fn fields_not_equal(&self, other: &Self)
        -> Vec<::field_by_field::UnequalField> {
            let mut list: Vec<::field_by_field::UnequalField> = Vec::new();

            #(#find_unequal_fields)*

            list
        }
    }
}

/// Build a function that compares enum variants and their nested fields
///
/// Same as build_fn_fields_not_equal, but a little more complex because it's
/// iterating over variants and then checking their fields, instead of just
/// checking fields
fn build_fn_fields_not_equal_enum(ty_name: &syn::Ident, data: &[syn::Variant]) -> quote::Tokens {
    let is_multivariant = data.len() > 1;
    let variants = data.iter()
        .map(|var| {
            let variant_name = &var.ident;
            match var.data {
                syn::VariantData::Unit =>
                    build_match_unit_variant(&ty_name, &variant_name, is_multivariant),
                syn::VariantData::Tuple(ref fields) =>
                    build_match_tuple_variant(&ty_name, &variant_name, fields, is_multivariant),
                syn::VariantData::Struct(ref fields) =>
                    build_match_struct_variant(&ty_name, &variant_name, fields, is_multivariant),
            }
        });
    quote! {
        fn fields_not_equal(&self, other: &Self)
        -> Vec<::field_by_field::UnequalField> {
            // This is never modified in the single-variant enum case.
            #![allow(unused_mut)]
            let mut list: Vec<::field_by_field::UnequalField> = Vec::new();

            match (self, other) {
                #(#variants)*
            }

            list
        }
    }
}

/// Build a match statement that compares the self variant to the other variant
///
/// Since this is for unit variants this doesn't check the actual value if it
/// isn't an exact match.
fn build_match_unit_variant(name: &syn::Ident, var_name: &syn::Ident, is_multivariant: bool)
-> quote::Tokens {
    let left_str = format!("{}::{}", &name, &var_name);

    if is_multivariant {
        quote! {
            ( &#name::#var_name, &#name::#var_name ) => {}
            ( &#name::#var_name, ref expected ) => {
                list.push(::field_by_field::UnequalField {
                    field_name: #left_str.to_string(),
                    actually: Box::new(#left_str.to_string()),
                    expected: Box::new(format!("{:?}", &expected)),
                });
            }
        }
    } else {
        // TODO: This causes a "variable does not need to be mutable" warning
        quote! {
            ( &#name::#var_name, &#name::#var_name ) => {}
        }
    }
}

/// Build a match statement that compares a tuple enum against others
fn build_match_tuple_variant(name: &syn::Ident,
                             var_name: &syn::Ident,
                             var_fields: &[syn::Field],
                             is_multivariant: bool)
-> quote::Tokens {
    let actually_fields: Vec<_> = var_fields.iter()
        .enumerate()
        .map(|(i, f)| {
            assert!(f.ident.is_none(),
                    format!("Fields in tuples should be unnamed, not {:?}", f.ident));
            syn::Ident::from(format!("_{}", i))
        })
        .collect();
    let actually_field_refs: Vec<_> = actually_fields.iter()
        .map(|name| quote! { ref #name }).collect();
    let expected_fields: Vec<_> = var_fields.iter()
        .enumerate()
        .map(|(i, f)| {
            assert!(f.ident.is_none(),
                    format!("Fields in tuples should be unnamed, not {:?}", f.ident));
            syn::Ident::from(format!("_e{}", i))
        })
        .collect();
    let expected_field_refs: Vec<_> = expected_fields.iter()
        .map(|name| quote! { ref #name }).collect();

    let name_str = name.to_string();
    let var_name_str = var_name.to_string();

    let comparisons: Vec<_> = actually_fields.iter()
        .zip(&expected_fields)
        .enumerate()
        .map(|(i, (actually, expected))| {
             let field_name = format!("{}::{}.{}", name_str, var_name_str, i);
             quote! {
                 if #actually != #expected {
                     list.push(::field_by_field::UnequalField {
                         field_name: #field_name.into(),
                         actually: Box::new(#actually.clone()),
                         expected: Box::new(#expected.clone()),
                     });
                 }
             }
        })
        .collect();

    let field_match = quote! {
        ( &#name::#var_name(#(#actually_field_refs),*),
          &#name::#var_name(#(#expected_field_refs),*) ) => {
            #(#comparisons)*
        }
    };

    if is_multivariant {
        let var_field_name = format!("{}::{}", name_str, var_name_str);
        // TODO: improve error messages if the variants differ
        quote! {
            #field_match
            (ref actually @ &#name::#var_name(..), ref expected) => {
                list.push(::field_by_field::UnequalField {
                    field_name: #var_field_name.into(),
                    actually: Box::new((*actually).clone()),
                    expected: Box::new((*expected).clone()),
                });
            }
        }
    } else {
        field_match
    }
}

/// Build a match arm that compares structs variants against themselves or other
fn build_match_struct_variant(name: &syn::Ident,
                              var_name: &syn::Ident,
                              fields: &[syn::Field],
                              is_multivariant: bool)
-> quote::Tokens {
    let field_names = fields.iter().cloned()
        .map(|field| {
             let ident = field.ident;
             ident
                .unwrap_or_else(|| panic!("Unable to get name for field in struct-like enum"))
        })
        .collect::<Vec<_>>();
    let expected_names = field_names.iter()
        .map(|name| format!("expected_{}", name).into())
        .collect::<Vec<syn::Ident>>();
    let expected_name_bindings = field_names.iter().zip(&expected_names)
        .map(|(name, expected_name)| {
            quote! { #name: ref #expected_name }
        })
        .collect::<Vec<_>>();
    let comparisons = field_names.iter()
        .zip(&expected_names)
        .map(|(name, other_name)| {
            let vname = name.as_ref();
            quote! {
                if #name != #other_name {
                    list.push(::field_by_field::UnequalField {
                        field_name: #vname.into(),
                        actually: Box::new(#name.clone()),
                        expected: Box::new(#other_name.clone()),
                    })
                }
            }
        })
        .collect::<Vec<_>>();
    let by_field_match = quote! {
        (&#name::#var_name { #(ref #field_names),* },
         &#name::#var_name { #(#expected_name_bindings,)* }) => {
            #(#comparisons)*
        }
    };

    if is_multivariant {
        let vname = format!("{}::{}", name, var_name);
        quote! {
            #by_field_match
            (ref actually @ &#name::#var_name { .. }, expected) => {
                list.push(::field_by_field::UnequalField {
                    field_name: #vname.into(),
                    actually: Box::new((*actually).clone()),
                    expected: Box::new((*expected).clone()),
                })
            }
        }
    } else {
        by_field_match
    }
}

/// Build a function that panics if the result of fiels_not_equal is non-empty
///
/// The Function will panic with a nice error message
fn build_fn_assert_equal_field_by_field(name: &syn::Ident) -> quote::Tokens {
    let overall_errmsgfmt = quote! {
        let ac_exp = format!("    actually: {:?}\n    \
                              expected: {:?}\n", self, other);
        errmsg.push_str(&ac_exp);
    };

    let field_errmsgfmt = quote! {
        errmsg.push_str(&format!("        {}: {:?} != {:?}\n",
                                 field_err.field_name,
                                 field_err.actually, field_err.expected));

    };

    quote! {
        fn assert_equal_field_by_field(&self, other: &#name) {
            let errs = self.fields_not_equal(other);

            if errs.len() > 0 {
                let mut errmsg = String::from("\n    Items are not equal:\n");
                for field_err in errs {
                    #field_errmsgfmt
                }

                #overall_errmsgfmt

                panic!("{}", errmsg);
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Starts drivers for found PCI devices,
//! and then reponds to device queries

#![no_std]
#![feature(allocator_api)]
#![deny(unused_must_use)]

#[macro_use]
extern crate alloc;

#[macro_use]
extern crate libd7;

use alloc::borrow::ToOwned;
use alloc::string::String;
use alloc::vec::Vec;

use hashbrown::HashMap;
use serde::Deserialize;

use libd7::{ipc, process::Process, syscall};

#[derive(Debug, Deserialize)]
struct ConfigDevice {
    shortname: Option<String>,
    name: String,
    driver: Option<ConfigDeviceDriver>,
}

#[derive(Debug, Deserialize)]
struct ConfigDeviceDriver {
    from_initrd: bool,
    executable: String,
}

#[no_mangle]
fn main() -> ! {
    syscall::debug_print("PCI driver starting");

    libd7::service::register("driver_pci", false);

    let s: Vec<u8> = ipc::request("initrd/read", "pci_devices.json".to_owned()).unwrap();
    let config_devices: HashMap<String, ConfigDevice> = serde_json::from_slice(&s).unwrap();

    let devices = unsafe { d7pci::list_devices() };

    for device in &devices {
        let vendor_and_id = format!("{:x}:{:x}", device.vendor, device.id);

        if let Some(device_config) = config_devices.get(&vendor_and_id) {
            println!(
                "PCI device: {}/{:?} {} ({})",
                vendor_and_id,
                device.location,
                device_config.name,
                device_config
                    .driver
                    .as_ref()
                    .map(|d| d.executable.as_ref())
                    .unwrap_or("no driver")
            );
            if let Some(driver) = &device_config.driver {
                assert!(
                    driver.from_initrd,
                    "Non-initrd executables are not supported yet"
                );
                Process::spawn(&driver.executable).unwrap();
            }
        } else {
            println!("Ignoring unknown PCI device {}", vendor_and_id);
        }
    }

    // (DriverName|"vendor:id") -> d7pci::Device
    let server: ipc::Server<String, Option<d7pci::Device>> =
        ipc::Server::exact("pci/device").unwrap();

    loop {
        server
            .handle(|name| {
                for device in &devices {
                    let vendor_and_id = format!("{:x}:{:x}", device.vendor, device.id);
                    if name == vendor_and_id {
                        return Ok(Some(device.clone()));
                    }

                    if let Some(device_config) = config_devices.get(&vendor_and_id) {
                        if name == device_config.name
                            || Some(&name) == device_config.shortname.as_ref()
                        {
                            return Ok(Some(device.clone()));
                        }
                    }
                }

                Ok(None)
            })
            .unwrap();
    }
}
<commit_msg>pci: Only register service after subscripting<commit_after>//! Starts drivers for found PCI devices,
//! and then reponds to device queries

#![no_std]
#![feature(allocator_api)]
#![deny(unused_must_use)]

#[macro_use]
extern crate alloc;

#[macro_use]
extern crate libd7;

use alloc::borrow::ToOwned;
use alloc::string::String;
use alloc::vec::Vec;

use hashbrown::HashMap;
use serde::Deserialize;

use libd7::{ipc, process::Process, syscall};

#[derive(Debug, Deserialize)]
struct ConfigDevice {
    shortname: Option<String>,
    name: String,
    driver: Option<ConfigDeviceDriver>,
}

#[derive(Debug, Deserialize)]
struct ConfigDeviceDriver {
    from_initrd: bool,
    executable: String,
}

#[no_mangle]
fn main() -> ! {
    syscall::debug_print("PCI driver starting");

    // (DriverName|"vendor:id") -> d7pci::Device
    let server: ipc::Server<String, Option<d7pci::Device>> =
        ipc::Server::exact("pci/device").unwrap();

    libd7::service::register("driver_pci", false);

    let s: Vec<u8> = ipc::request("initrd/read", "pci_devices.json".to_owned()).unwrap();
    let config_devices: HashMap<String, ConfigDevice> = serde_json::from_slice(&s).unwrap();

    let devices = unsafe { d7pci::list_devices() };

    for device in &devices {
        let vendor_and_id = format!("{:x}:{:x}", device.vendor, device.id);

        if let Some(device_config) = config_devices.get(&vendor_and_id) {
            println!(
                "PCI device: {}/{:?} {} ({})",
                vendor_and_id,
                device.location,
                device_config.name,
                device_config
                    .driver
                    .as_ref()
                    .map(|d| d.executable.as_ref())
                    .unwrap_or("no driver")
            );
            if let Some(driver) = &device_config.driver {
                assert!(
                    driver.from_initrd,
                    "Non-initrd executables are not supported yet"
                );
                Process::spawn(&driver.executable).unwrap();
            }
        } else {
            println!("Ignoring unknown PCI device {}", vendor_and_id);
        }
    }

    loop {
        server
            .handle(|name| {
                for device in &devices {
                    let vendor_and_id = format!("{:x}:{:x}", device.vendor, device.id);
                    if name == vendor_and_id {
                        return Ok(Some(device.clone()));
                    }

                    if let Some(device_config) = config_devices.get(&vendor_and_id) {
                        if name == device_config.name
                            || Some(&name) == device_config.shortname.as_ref()
                        {
                            return Ok(Some(device.clone()));
                        }
                    }
                }

                Ok(None)
            })
            .unwrap();
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate quick_protobuf;
extern crate test;

#[macro_use]
extern crate lazy_static;

use test::{Bencher, black_box};

use quick_protobuf::{Reader, Writer};

const LEN: i32 = 10_000;

lazy_static! {
    static ref BUFFER: Vec<u8> = {
        let mut buf = Vec::new();
        {
            let mut writer = Writer::new(&mut buf);
            for i in 0..LEN {
                writer.write_int32(i).unwrap();
            }
        }
        buf
    };
}

#[bench]
fn read_varint(b: &mut Bencher) {
    b.iter(|| {
        let buf_len = BUFFER.len();
        let mut buf: &[u8] = &*BUFFER;
        let mut reader = Reader::from_reader(&mut buf, buf_len);
        for _ in 0..LEN {
            let _ = black_box(reader.read_varint().unwrap());
        }
        assert!(reader.is_eof());
    })
}

#[bench]
fn read_varint_and_is_eof(b: &mut Bencher) {
    b.iter(|| {
        let buf_len = BUFFER.len();
        let mut buf: &[u8] = &*BUFFER;
        let mut reader = Reader::from_reader(&mut buf, buf_len);
        for _ in 0..LEN {
            assert!(!reader.is_eof());
            let _ = black_box(reader.read_varint().unwrap());
        }
        assert!(reader.is_eof());
    })
}
<commit_msg>fix varint bench<commit_after>#![feature(test)]

extern crate quick_protobuf;
extern crate test;

#[macro_use]
extern crate lazy_static;

use test::{Bencher, black_box};

use quick_protobuf::{Reader, Writer};

const LEN: i32 = 10_000;

lazy_static! {
    static ref BUFFER: Vec<u8> = {
        let mut buf = Vec::new();
        {
            let mut writer = Writer::new(&mut buf);
            for i in 0..LEN {
                writer.write_int32(i).unwrap();
            }
        }
        buf
    };
}

#[bench]
fn read_varint(b: &mut Bencher) {
    b.iter(|| {
        let mut reader = Reader::from_bytes(&BUFFER);
        for _ in 0..LEN {
            let _ = black_box(reader.read_varint(&BUFFER).unwrap());
        }
        assert!(reader.is_eof());
    })
}

#[bench]
fn read_varint_and_is_eof(b: &mut Bencher) {
    b.iter(|| {
        let mut reader = Reader::from_bytes(&BUFFER);
        for _ in 0..LEN {
            assert!(!reader.is_eof());
            let _ = black_box(reader.read_varint(&BUFFER).unwrap());
        }
    })
}
<|endoftext|>


<commit_before>// Copyright 2016 coroutine-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

#![feature(test)]

extern crate context;
extern crate test;

use test::Bencher;

use context::{Context, Transfer};
use context::stack::FixedSizeStack;
use std::mem;

#[bench]
fn resume_reference_perf(b: &mut Bencher) {
    #[inline(never)]
    extern "C" fn yielder(t: Transfer) -> Transfer {
        test::black_box(t)
    }

    b.iter(|| yielder(unsafe { Transfer::empty(0) }));
}

#[bench]
fn resume(b: &mut Bencher) {
    extern "C" fn yielder(mut t: Transfer) -> ! {
        loop {
            t = t.context.resume(1);
        }
    }

    let stack = FixedSizeStack::default();
    let mut t = Transfer::new(Context::new(&stack, yielder), 0);

    b.iter(|| unsafe { t = mem::replace(&mut t, mem::uninitialized()).context.resume(0) });
}
<commit_msg>add bench resume_ontop<commit_after>// Copyright 2016 coroutine-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

#![feature(test)]

extern crate context;
extern crate test;

use test::Bencher;

use context::{Context, Transfer};
use context::stack::FixedSizeStack;
use std::mem;

#[bench]
fn resume_reference_perf(b: &mut Bencher) {
    #[inline(never)]
    extern "C" fn yielder(t: Transfer) -> Transfer {
        test::black_box(t)
    }

    b.iter(|| yielder(unsafe { Transfer::empty(0) }));
}

#[bench]
fn resume(b: &mut Bencher) {
    extern "C" fn yielder(mut t: Transfer) -> ! {
        loop {
            t = t.context.resume(1);
        }
    }

    let stack = FixedSizeStack::default();
    let mut t = Transfer::new(Context::new(&stack, yielder), 0);

    b.iter(|| unsafe { t = mem::replace(&mut t, mem::uninitialized()).context.resume(0) });
}

#[bench]
fn resume_ontop(b: &mut Bencher) {
    extern "C" fn yielder(mut t: Transfer) -> ! {
        loop {
            t = t.context.resume(1);
        }
    }

    extern "C" fn ontop_function(mut t: Transfer) -> Transfer {
        t
    }

    let stack = FixedSizeStack::default();
    let mut t = Transfer::new(Context::new(&stack, yielder), 0);
    b.iter(|| unsafe { t = mem::replace(&mut t, mem::uninitialized()).context.resume_ontop(0, ontop_function) });
}
<|endoftext|>


<commit_before>//! Transactions

use std::cell::Cell;
use std::fmt;

use rows::Rows;
use text_rows::TextRows;
use stmt::Statement;
use types::ToSql;
use {bad_response, Connection, Result};

/// An enumeration of transaction isolation levels.
///
/// See the [Postgres documentation](http://www.postgresql.org/docs/9.4/static/transaction-iso.html)
/// for full details on the semantics of each level.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IsolationLevel {
    /// The "read uncommitted" level.
    ///
    /// In current versions of Postgres, this behaves identically to
    /// `ReadCommitted`.
    ReadUncommitted,
    /// The "read committed" level.
    ///
    /// This is the default isolation level in Postgres.
    ReadCommitted,
    /// The "repeatable read" level.
    RepeatableRead,
    /// The "serializable" level.
    Serializable,
}

impl IsolationLevel {
    pub(crate) fn new(raw: &str) -> Result<IsolationLevel> {
        if raw.eq_ignore_ascii_case("READ UNCOMMITTED") {
            Ok(IsolationLevel::ReadUncommitted)
        } else if raw.eq_ignore_ascii_case("READ COMMITTED") {
            Ok(IsolationLevel::ReadCommitted)
        } else if raw.eq_ignore_ascii_case("REPEATABLE READ") {
            Ok(IsolationLevel::RepeatableRead)
        } else if raw.eq_ignore_ascii_case("SERIALIZABLE") {
            Ok(IsolationLevel::Serializable)
        } else {
            Err(bad_response().into())
        }
    }

    fn to_sql(&self) -> &'static str {
        match *self {
            IsolationLevel::ReadUncommitted => "READ UNCOMMITTED",
            IsolationLevel::ReadCommitted => "READ COMMITTED",
            IsolationLevel::RepeatableRead => "REPEATABLE READ",
            IsolationLevel::Serializable => "SERIALIZABLE",
        }
    }
}

/// Configuration of a transaction.
#[derive(Debug)]
pub struct Config {
    isolation_level: Option<IsolationLevel>,
    read_only: Option<bool>,
    deferrable: Option<bool>,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            isolation_level: None,
            read_only: None,
            deferrable: None,
        }
    }
}

impl Config {
    pub(crate) fn build_command(&self, s: &mut String) {
        let mut first = true;

        if let Some(isolation_level) = self.isolation_level {
            s.push_str(" ISOLATION LEVEL ");
            s.push_str(isolation_level.to_sql());
            first = false;
        }

        if let Some(read_only) = self.read_only {
            if !first {
                s.push(',');
            }
            if read_only {
                s.push_str(" READ ONLY");
            } else {
                s.push_str(" READ WRITE");
            }
            first = false;
        }

        if let Some(deferrable) = self.deferrable {
            if !first {
                s.push(',');
            }
            if deferrable {
                s.push_str(" DEFERRABLE");
            } else {
                s.push_str(" NOT DEFERRABLE");
            }
        }
    }

    /// Creates a new `Config` with no configuration overrides.
    pub fn new() -> Config {
        Config::default()
    }

    /// Sets the isolation level of the configuration.
    pub fn isolation_level(&mut self, isolation_level: IsolationLevel) -> &mut Config {
        self.isolation_level = Some(isolation_level);
        self
    }

    /// Sets the read-only property of a transaction.
    ///
    /// If enabled, a transaction will be unable to modify any persistent
    /// database state.
    pub fn read_only(&mut self, read_only: bool) -> &mut Config {
        self.read_only = Some(read_only);
        self
    }

    /// Sets the deferrable property of a transaction.
    ///
    /// If enabled in a read only, serializable transaction, the transaction may
    /// block when created, after which it will run without the normal overhead
    /// of a serializable transaction and will not be forced to roll back due
    /// to serialization failures.
    pub fn deferrable(&mut self, deferrable: bool) -> &mut Config {
        self.deferrable = Some(deferrable);
        self
    }
}

/// A transaction on a database connection.
///
/// The transaction will roll back by default.
pub struct Transaction<'conn> {
    conn: &'conn Connection,
    depth: u32,
    savepoint_name: Option<String>,
    commit: Cell<bool>,
    finished: bool,
}

impl<'a> fmt::Debug for Transaction<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Transaction")
            .field("commit", &self.commit.get())
            .field("depth", &self.depth)
            .finish()
    }
}

impl<'conn> Drop for Transaction<'conn> {
    fn drop(&mut self) {
        if !self.finished {
            let _ = self.finish_inner();
        }
    }
}

impl<'conn> Transaction<'conn> {
    pub(crate) fn new(conn: &'conn Connection, depth: u32) -> Transaction<'conn> {
        Transaction {
            conn: conn,
            depth: depth,
            savepoint_name: None,
            commit: Cell::new(false),
            finished: false,
        }
    }

    pub(crate) fn conn(&self) -> &'conn Connection {
        self.conn
    }

    pub(crate) fn depth(&self) -> u32 {
        self.depth
    }

    fn finish_inner(&mut self) -> Result<()> {
        let mut conn = self.conn.0.borrow_mut();
        debug_assert!(self.depth == conn.trans_depth);
        conn.trans_depth -= 1;
        match (self.commit.get(), &self.savepoint_name) {
            (false, &Some(ref sp)) => conn.quick_query(&format!("ROLLBACK TO {}", sp))?,
            (false, &None) => conn.quick_query("ROLLBACK")?,
            (true, &Some(ref sp)) => conn.quick_query(&format!("RELEASE {}", sp))?,
            (true, &None) => conn.quick_query("COMMIT")?,
        };

        Ok(())
    }

    /// Like `Connection::prepare`.
    pub fn prepare(&self, query: &str) -> Result<Statement<'conn>> {
        self.conn.prepare(query)
    }

    /// Like `Connection::prepare_cached`.
    ///
    /// # Note
    ///
    /// The statement will be cached for the duration of the
    /// connection, not just the duration of this transaction.
    pub fn prepare_cached(&self, query: &str) -> Result<Statement<'conn>> {
        self.conn.prepare_cached(query)
    }

    /// Like `Connection::execute`.
    pub fn execute(&self, query: &str, params: &[&ToSql]) -> Result<u64> {
        self.conn.execute(query, params)
    }

    /// Like `Connection::query`.
    pub fn query<'a>(&'a self, query: &str, params: &[&ToSql]) -> Result<Rows> {
        self.conn.query(query, params)
    }

    /// Like `Connection::batch_execute`.
    #[deprecated(since="0.15.3", note="please use `simple_query` instead")]
    pub fn batch_execute(&self, query: &str) -> Result<()> {
        self.simple_query(query)
            .map(|_| ())
    }

    /// Like `Connection::simple_query`.
    pub fn simple_query(&self, query: &str) -> Result<Vec<TextRows>> {
        self.conn.simple_query(query)
    }

    /// Like `Connection::transaction`, but creates a nested transaction via
    /// a savepoint.
    ///
    /// # Panics
    ///
    /// Panics if there is an active nested transaction.
    pub fn transaction<'a>(&'a self) -> Result<Transaction<'a>> {
        self.savepoint("sp")
    }

    /// Like `Connection::transaction`, but creates a nested transaction via
    /// a savepoint with the specified name.
    ///
    /// # Panics
    ///
    /// Panics if there is an active nested transaction.
    pub fn savepoint<'a>(&'a self, name: &str) -> Result<Transaction<'a>> {
        let mut conn = self.conn.0.borrow_mut();
        check_desync!(conn);
        assert!(
            conn.trans_depth == self.depth,
            "`savepoint` may only be called on the active transaction"
        );
        conn.quick_query(&format!("SAVEPOINT {}", name))?;
        conn.trans_depth += 1;
        Ok(Transaction {
            conn: self.conn,
            depth: self.depth + 1,
            savepoint_name: Some(name.to_owned()),
            commit: Cell::new(false),
            finished: false,
        })
    }

    /// Returns a reference to the `Transaction`'s `Connection`.
    pub fn connection(&self) -> &'conn Connection {
        self.conn
    }

    /// Like `Connection::is_active`.
    pub fn is_active(&self) -> bool {
        self.conn.0.borrow().trans_depth == self.depth
    }

    /// Alters the configuration of the active transaction.
    pub fn set_config(&self, config: &Config) -> Result<()> {
        let mut command = "SET TRANSACTION".to_owned();
        config.build_command(&mut command);
        self.simple_query(&command)
            .map(|_| ())
    }

    /// Determines if the transaction is currently set to commit or roll back.
    pub fn will_commit(&self) -> bool {
        self.commit.get()
    }

    /// Sets the transaction to commit at its completion.
    pub fn set_commit(&self) {
        self.commit.set(true);
    }

    /// Sets the transaction to roll back at its completion.
    pub fn set_rollback(&self) {
        self.commit.set(false);
    }

    /// A convenience method which consumes and commits a transaction.
    pub fn commit(self) -> Result<()> {
        self.set_commit();
        self.finish()
    }

    /// Consumes the transaction, commiting or rolling it back as appropriate.
    ///
    /// Functionally equivalent to the `Drop` implementation of `Transaction`
    /// except that it returns any error to the caller.
    pub fn finish(mut self) -> Result<()> {
        self.finished = true;
        self.finish_inner()
    }
}
<commit_msg>Fixes #371 by not reusing the same savepoint name.<commit_after>//! Transactions

use std::cell::Cell;
use std::fmt;

use rows::Rows;
use text_rows::TextRows;
use stmt::Statement;
use types::ToSql;
use {bad_response, Connection, Result};

/// An enumeration of transaction isolation levels.
///
/// See the [Postgres documentation](http://www.postgresql.org/docs/9.4/static/transaction-iso.html)
/// for full details on the semantics of each level.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum IsolationLevel {
    /// The "read uncommitted" level.
    ///
    /// In current versions of Postgres, this behaves identically to
    /// `ReadCommitted`.
    ReadUncommitted,
    /// The "read committed" level.
    ///
    /// This is the default isolation level in Postgres.
    ReadCommitted,
    /// The "repeatable read" level.
    RepeatableRead,
    /// The "serializable" level.
    Serializable,
}

impl IsolationLevel {
    pub(crate) fn new(raw: &str) -> Result<IsolationLevel> {
        if raw.eq_ignore_ascii_case("READ UNCOMMITTED") {
            Ok(IsolationLevel::ReadUncommitted)
        } else if raw.eq_ignore_ascii_case("READ COMMITTED") {
            Ok(IsolationLevel::ReadCommitted)
        } else if raw.eq_ignore_ascii_case("REPEATABLE READ") {
            Ok(IsolationLevel::RepeatableRead)
        } else if raw.eq_ignore_ascii_case("SERIALIZABLE") {
            Ok(IsolationLevel::Serializable)
        } else {
            Err(bad_response().into())
        }
    }

    fn to_sql(&self) -> &'static str {
        match *self {
            IsolationLevel::ReadUncommitted => "READ UNCOMMITTED",
            IsolationLevel::ReadCommitted => "READ COMMITTED",
            IsolationLevel::RepeatableRead => "REPEATABLE READ",
            IsolationLevel::Serializable => "SERIALIZABLE",
        }
    }
}

/// Configuration of a transaction.
#[derive(Debug)]
pub struct Config {
    isolation_level: Option<IsolationLevel>,
    read_only: Option<bool>,
    deferrable: Option<bool>,
}

impl Default for Config {
    fn default() -> Config {
        Config {
            isolation_level: None,
            read_only: None,
            deferrable: None,
        }
    }
}

impl Config {
    pub(crate) fn build_command(&self, s: &mut String) {
        let mut first = true;

        if let Some(isolation_level) = self.isolation_level {
            s.push_str(" ISOLATION LEVEL ");
            s.push_str(isolation_level.to_sql());
            first = false;
        }

        if let Some(read_only) = self.read_only {
            if !first {
                s.push(',');
            }
            if read_only {
                s.push_str(" READ ONLY");
            } else {
                s.push_str(" READ WRITE");
            }
            first = false;
        }

        if let Some(deferrable) = self.deferrable {
            if !first {
                s.push(',');
            }
            if deferrable {
                s.push_str(" DEFERRABLE");
            } else {
                s.push_str(" NOT DEFERRABLE");
            }
        }
    }

    /// Creates a new `Config` with no configuration overrides.
    pub fn new() -> Config {
        Config::default()
    }

    /// Sets the isolation level of the configuration.
    pub fn isolation_level(&mut self, isolation_level: IsolationLevel) -> &mut Config {
        self.isolation_level = Some(isolation_level);
        self
    }

    /// Sets the read-only property of a transaction.
    ///
    /// If enabled, a transaction will be unable to modify any persistent
    /// database state.
    pub fn read_only(&mut self, read_only: bool) -> &mut Config {
        self.read_only = Some(read_only);
        self
    }

    /// Sets the deferrable property of a transaction.
    ///
    /// If enabled in a read only, serializable transaction, the transaction may
    /// block when created, after which it will run without the normal overhead
    /// of a serializable transaction and will not be forced to roll back due
    /// to serialization failures.
    pub fn deferrable(&mut self, deferrable: bool) -> &mut Config {
        self.deferrable = Some(deferrable);
        self
    }
}

/// A transaction on a database connection.
///
/// The transaction will roll back by default.
pub struct Transaction<'conn> {
    conn: &'conn Connection,
    depth: u32,
    savepoint_name: Option<String>,
    commit: Cell<bool>,
    finished: bool,
}

impl<'a> fmt::Debug for Transaction<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Transaction")
            .field("commit", &self.commit.get())
            .field("depth", &self.depth)
            .finish()
    }
}

impl<'conn> Drop for Transaction<'conn> {
    fn drop(&mut self) {
        if !self.finished {
            let _ = self.finish_inner();
        }
    }
}

impl<'conn> Transaction<'conn> {
    pub(crate) fn new(conn: &'conn Connection, depth: u32) -> Transaction<'conn> {
        Transaction {
            conn: conn,
            depth: depth,
            savepoint_name: None,
            commit: Cell::new(false),
            finished: false,
        }
    }

    pub(crate) fn conn(&self) -> &'conn Connection {
        self.conn
    }

    pub(crate) fn depth(&self) -> u32 {
        self.depth
    }

    fn finish_inner(&mut self) -> Result<()> {
        let mut conn = self.conn.0.borrow_mut();
        debug_assert!(self.depth == conn.trans_depth);
        conn.trans_depth -= 1;
        match (self.commit.get(), &self.savepoint_name) {
            (false, &Some(ref sp)) => conn.quick_query(&format!("ROLLBACK TO {}", sp))?,
            (false, &None) => conn.quick_query("ROLLBACK")?,
            (true, &Some(ref sp)) => conn.quick_query(&format!("RELEASE {}", sp))?,
            (true, &None) => conn.quick_query("COMMIT")?,
        };

        Ok(())
    }

    /// Like `Connection::prepare`.
    pub fn prepare(&self, query: &str) -> Result<Statement<'conn>> {
        self.conn.prepare(query)
    }

    /// Like `Connection::prepare_cached`.
    ///
    /// # Note
    ///
    /// The statement will be cached for the duration of the
    /// connection, not just the duration of this transaction.
    pub fn prepare_cached(&self, query: &str) -> Result<Statement<'conn>> {
        self.conn.prepare_cached(query)
    }

    /// Like `Connection::execute`.
    pub fn execute(&self, query: &str, params: &[&ToSql]) -> Result<u64> {
        self.conn.execute(query, params)
    }

    /// Like `Connection::query`.
    pub fn query<'a>(&'a self, query: &str, params: &[&ToSql]) -> Result<Rows> {
        self.conn.query(query, params)
    }

    /// Like `Connection::batch_execute`.
    #[deprecated(since="0.15.3", note="please use `simple_query` instead")]
    pub fn batch_execute(&self, query: &str) -> Result<()> {
        self.simple_query(query)
            .map(|_| ())
    }

    /// Like `Connection::simple_query`.
    pub fn simple_query(&self, query: &str) -> Result<Vec<TextRows>> {
        self.conn.simple_query(query)
    }

    /// Like `Connection::transaction`, but creates a nested transaction via
    /// a savepoint.
    ///
    /// # Panics
    ///
    /// Panics if there is an active nested transaction.
    pub fn transaction<'a>(&'a self) -> Result<Transaction<'a>> {
        self.savepoint(format!("sp_{}", self.depth()))
    }

    /// Like `Connection::transaction`, but creates a nested transaction via
    /// a savepoint with the specified name.
    ///
    /// # Panics
    ///
    /// Panics if there is an active nested transaction.
    #[inline]
    pub fn savepoint<'a, I>(&'a self, name: I) -> Result<Transaction<'a>>
        where I: Into<String>
    {
        self._savepoint(name.into())
    }

    fn _savepoint<'a>(&'a self, name: String) -> Result<Transaction<'a>> {
        let mut conn = self.conn.0.borrow_mut();
        check_desync!(conn);
        assert!(
            conn.trans_depth == self.depth,
            "`savepoint` may only be called on the active transaction"
        );
        conn.quick_query(&format!("SAVEPOINT {}", name))?;
        conn.trans_depth += 1;
        Ok(Transaction {
            conn: self.conn,
            depth: self.depth + 1,
            savepoint_name: Some(name),
            commit: Cell::new(false),
            finished: false,
        })
    }

    /// Returns a reference to the `Transaction`'s `Connection`.
    pub fn connection(&self) -> &'conn Connection {
        self.conn
    }

    /// Like `Connection::is_active`.
    pub fn is_active(&self) -> bool {
        self.conn.0.borrow().trans_depth == self.depth
    }

    /// Alters the configuration of the active transaction.
    pub fn set_config(&self, config: &Config) -> Result<()> {
        let mut command = "SET TRANSACTION".to_owned();
        config.build_command(&mut command);
        self.simple_query(&command)
            .map(|_| ())
    }

    /// Determines if the transaction is currently set to commit or roll back.
    pub fn will_commit(&self) -> bool {
        self.commit.get()
    }

    /// Sets the transaction to commit at its completion.
    pub fn set_commit(&self) {
        self.commit.set(true);
    }

    /// Sets the transaction to roll back at its completion.
    pub fn set_rollback(&self) {
        self.commit.set(false);
    }

    /// A convenience method which consumes and commits a transaction.
    pub fn commit(self) -> Result<()> {
        self.set_commit();
        self.finish()
    }

    /// Consumes the transaction, commiting or rolling it back as appropriate.
    ///
    /// Functionally equivalent to the `Drop` implementation of `Transaction`
    /// except that it returns any error to the caller.
    pub fn finish(mut self) -> Result<()> {
        self.finished = true;
        self.finish_inner()
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate serde_derive;
extern crate docopt;
extern crate chan_signal;
extern crate rand;
extern crate sled;

use std::thread;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;

use chan_signal::Signal;
use docopt::Docopt;
use rand::{Rng, thread_rng};

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
";

#[derive(Deserialize)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
}

fn report(shutdown: Arc<AtomicBool>, total: Arc<AtomicUsize>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = total.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn run(
    tree: Arc<sled::Tree>,
    shutdown: Arc<AtomicBool>,
    total: Arc<AtomicUsize>,
) {
    let mut rng = thread_rng();
    let mut byte = || vec![rng.gen::<u8>()];
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        total.fetch_add(1, Ordering::Release);
        let choice = rng.gen_range(0, 5);

        match choice {
            0 => {
                tree.set(byte(), byte());
            }
            1 => {
                tree.get(&*byte());
            }
            2 => {
                tree.del(&*byte());
            }
            3 => {
                if let Err(_) = tree.cas(byte(), Some(byte()), Some(byte())) {};
            }
            4 => {
                tree.scan(&*byte())
                    .take(rng.gen_range(0, 15))
                    .collect::<Vec<_>>();
            }
            _ => panic!("impossible choice"),
        }

    }
}

fn main() {
    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.argv(std::env::args().into_iter()).deserialize())
        .unwrap_or_else(|e| e.exit());

    let total = Arc::new(AtomicUsize::new(0));
    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(2)
        .io_buf_size(8_000_000)
        .blink_fanout(32)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000)
        .flush_every_ms(Some(100))
        .snapshot_after_ops(1000000)
        .build();

    let tree = Arc::new(sled::Tree::start(config));

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();
        let total = total.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown, total))
        } else {
            thread::spawn(move || run(tree, shutdown, total))
        };

        threads.push(t);
    }

    if args.flag_burn_in {
        println!("waiting on signal");
        signal.recv();
        println!("got shutdown signal, cleaning up...");
    } else {
        thread::sleep(std::time::Duration::from_secs(args.flag_duration));
    }

    shutdown.store(true, Ordering::SeqCst);

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = total.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}
<commit_msg>account for new result return types<commit_after>#[macro_use]
extern crate serde_derive;
extern crate docopt;
extern crate chan_signal;
extern crate rand;
extern crate sled;

use std::thread;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;

use chan_signal::Signal;
use docopt::Docopt;
use rand::{Rng, thread_rng};

const USAGE: &'static str = "
Usage: stress [--threads=<#>] [--burn-in] [--duration=<s>]

Options:
    --threads=<#>      Number of threads [default: 4].
    --burn-in          Don't halt until we receive a signal.
    --duration=<s>     Seconds to run for [default: 10].
";

#[derive(Deserialize)]
struct Args {
    flag_threads: usize,
    flag_burn_in: bool,
    flag_duration: u64,
}

fn report(shutdown: Arc<AtomicBool>, total: Arc<AtomicUsize>) {
    let mut last = 0;
    while !shutdown.load(Ordering::Relaxed) {
        thread::sleep(std::time::Duration::from_secs(1));
        let total = total.load(Ordering::Acquire);

        println!("did {} ops", total - last);

        last = total;
    }
}

fn run(
    tree: Arc<sled::Tree>,
    shutdown: Arc<AtomicBool>,
    total: Arc<AtomicUsize>,
) {
    let mut rng = thread_rng();
    let mut byte = || vec![rng.gen::<u8>()];
    let mut rng = thread_rng();

    while !shutdown.load(Ordering::Relaxed) {
        total.fetch_add(1, Ordering::Release);
        let choice = rng.gen_range(0, 5);

        match choice {
            0 => {
                tree.set(byte(), byte()).unwrap();
            }
            1 => {
                tree.get(&*byte()).unwrap();
            }
            2 => {
                tree.del(&*byte()).unwrap();
            }
            3 => {
                match tree.cas(byte(), Some(byte()), Some(byte())) {
                    Ok(_) |
                    Err(sled::Error::CasFailed(_)) => {}
                    other => panic!("operational error: {:?}", other),
                }
            }
            4 => {
                tree.scan(&*byte())
                    .unwrap()
                    .take(rng.gen_range(0, 15))
                    .collect::<Vec<_>>();
            }
            _ => panic!("impossible choice"),
        }
    }
}

fn main() {
    let signal = chan_signal::notify(&[Signal::INT, Signal::TERM]);

    let args: Args = Docopt::new(USAGE)
        .and_then(|d| d.argv(std::env::args().into_iter()).deserialize())
        .unwrap_or_else(|e| e.exit());

    let total = Arc::new(AtomicUsize::new(0));
    let shutdown = Arc::new(AtomicBool::new(false));

    let config = sled::ConfigBuilder::new()
        .io_bufs(2)
        .io_buf_size(8_000_000)
        .blink_fanout(32)
        .page_consolidation_threshold(10)
        .cache_bits(6)
        .cache_capacity(1_000_000)
        .flush_every_ms(Some(100))
        .snapshot_after_ops(1000000)
        .build();

    let tree = Arc::new(sled::Tree::start(config));

    let mut threads = vec![];

    let now = std::time::Instant::now();

    let n_threads = args.flag_threads;

    for i in 0..n_threads + 1 {
        let tree = tree.clone();
        let shutdown = shutdown.clone();
        let total = total.clone();

        let t = if i == 0 {
            thread::spawn(move || report(shutdown, total))
        } else {
            thread::spawn(move || run(tree, shutdown, total))
        };

        threads.push(t);
    }

    if args.flag_burn_in {
        println!("waiting on signal");
        signal.recv();
        println!("got shutdown signal, cleaning up...");
    } else {
        thread::sleep(std::time::Duration::from_secs(args.flag_duration));
    }

    shutdown.store(true, Ordering::SeqCst);

    for t in threads.into_iter() {
        t.join().unwrap();
    }

    let ops = total.load(Ordering::SeqCst);
    let time = now.elapsed().as_secs() as usize;

    println!(
        "did {} total ops in {} seconds. {} ops/s",
        ops,
        time,
        (ops * 1_000) / (time * 1_000)
    );
}
<|endoftext|>


<commit_before>use crate::i18n::LocalText;
use crate::*;

use colored::Colorize;
use regex::Regex;
use std::io::prelude::*;
use std::{ffi::OsString, io};
use subprocess::{Exec, ExitStatus, Redirection};

// FTL: help-subcommand-make
/// Build specified target(s)
pub fn run(target: Vec<String>) -> Result<()> {
    setup::is_setup()?;
    show_header("make-header");
    let mut makeflags: Vec<OsString> = Vec::new();
    let cpus = num_cpus::get();
    makeflags.push(OsString::from(format!("--jobs={}", cpus)));
    let mut makefiles: Vec<OsString> = Vec::new();
    makefiles.push(OsString::from("-f"));
    makefiles.push(OsString::from(format!(
        "{}{}",
        CONFIGURE_DATADIR, "rules/casile.mk"
    )));
    let rules = status::get_rules()?;
    for rule in rules {
        makefiles.push(OsString::from("-f"));
        let p = rule.into_os_string();
        makefiles.push(p);
    }
    makefiles.push(OsString::from("-f"));
    makefiles.push(OsString::from(format!(
        "{}{}",
        CONFIGURE_DATADIR, "rules/rules.mk"
    )));
    let mut process = Exec::cmd("make")
        .args(&makeflags)
        .args(&makefiles)
        .args(&target);
    let mut targets: Vec<_> = target.into_iter().collect();
    if status::is_gha()? {
        targets.push(String::from("_gha"));
    }
    let gitname = status::get_gitname()?;
    let git_version = status::get_git_version();
    process = process
        .env("CASILE_CLI", "true")
        .env("CASILEDIR", CONFIGURE_DATADIR)
        .env("CONTAINERIZED", status::is_container().to_string())
        .env("LANGUAGE", lang_to_language(CONF.get_string("language")?))
        .env("PROJECT", &gitname)
        .env("PROJECTDIR", CONF.get_string("path")?)
        .env("PROJECTVERSION", git_version);
    if CONF.get_bool("debug")? {
        process = process.env("DEBUG", "true");
    };
    if CONF.get_bool("quiet")? {
        process = process.env("QUIET", "true");
    };
    if CONF.get_bool("verbose")? || targets.contains(&"debug".into()) {
        process = process.env("VERBOSE", "true");
    };
    let repo = get_repo()?;
    let workdir = repo.workdir().unwrap();
    process = process.cwd(workdir);
    let process = process.stderr(Redirection::Merge).stdout(Redirection::Pipe);
    let mut popen = process.popen()?;
    let buf = io::BufReader::new(popen.stdout.as_mut().unwrap());
    let mut backlog: Vec<String> = Vec::new();
    let seps = Regex::new(r"").unwrap();
    let mut ret: u32 = 0;
    for line in buf.lines() {
        let text: &str = &line.unwrap();
        let fields: Vec<&str> = seps.splitn(text, 4).collect();
        match fields[0] {
            "CASILE" => match fields[1] {
                "PRE" => report_start(fields[2]),
                "STDOUT" => {
                    if status::is_gha()? {
                        println!("{}", fields[3]);
                    } else if CONF.get_bool("verbose")? {
                        report_line(fields[3]);
                    } else {
                        backlog.push(String::from(fields[3]));
                    }
                }
                "STDERR" => {
                    if CONF.get_bool("verbose")? {
                        report_line(fields[3]);
                    } else {
                        backlog.push(String::from(fields[3]));
                    }
                }
                "POST" => match fields[2] {
                    "0" => {
                        report_end(fields[3]);
                    }
                    val => {
                        report_fail(fields[3]);
                        ret = val.parse().unwrap_or(1);
                    }
                },
                _ => {
                    let errmsg = LocalText::new("make-error-unknown-code").fmt();
                    panic!(errmsg)
                }
            },
            _ => backlog.push(String::from(fields[0])),
        }
    }
    let status = popen.wait();
    match status {
        Ok(ExitStatus::Exited(int)) => {
            let code = int + ret;
            match code {
                0 => {
                    if targets.contains(&"debug".into()) {
                        dump_backlog(&backlog)
                    };
                    Ok(())
                }
                1 => {
                    dump_backlog(&backlog);
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        LocalText::new("make-error-unfinished").fmt(),
                    )))
                }
                2 => {
                    dump_backlog(&backlog);
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        LocalText::new("make-error-build").fmt(),
                    )))
                }
                3 => {
                    if !CONF.get_bool("verbose")? {
                        dump_backlog(&backlog);
                    }
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        LocalText::new("make-error-target").fmt(),
                    )))
                }
                _ => {
                    dump_backlog(&backlog);
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        LocalText::new("make-error-unknown").fmt(),
                    )))
                }
            }
        }
        _ => Err(Box::new(io::Error::new(
            io::ErrorKind::InvalidInput,
            LocalText::new("make-error").fmt(),
        ))),
    }
}

fn dump_backlog(backlog: &[String]) {
    let start = LocalText::new("make-backlog-start").fmt();
    eprintln!("{} {}", "┖┄".cyan(), start);
    for line in backlog.iter() {
        eprintln!("{}", line);
    }
    let end = LocalText::new("make-backlog-end").fmt();
    eprintln!("{} {}", "┎┄".cyan(), end);
}

fn report_line(line: &str) {
    eprintln!("{} {}", "┠╎".cyan(), line.dimmed());
}

fn report_start(target: &str) {
    let text = LocalText::new("make-report-start")
        .arg("target", target.white().bold())
        .fmt();
    eprintln!("{} {}", "┠┄".cyan(), text.yellow());
}

fn report_end(target: &str) {
    let text = LocalText::new("make-report-end")
        .arg("target", target.white().bold())
        .fmt();
    eprintln!("{} {}", "┠┄".cyan(), text.green());
}

fn report_fail(target: &str) {
    let text = LocalText::new("make-report-fail")
        .arg("target", target.white().bold())
        .fmt();
    eprintln!("{} {}", "┠┄".cyan(), text.red());
}
<commit_msg>feat(actions): Always invoke install-dist when run as action<commit_after>use crate::i18n::LocalText;
use crate::*;

use colored::Colorize;
use regex::Regex;
use std::io::prelude::*;
use std::{ffi::OsString, io};
use subprocess::{Exec, ExitStatus, Redirection};

// FTL: help-subcommand-make
/// Build specified target(s)
pub fn run(target: Vec<String>) -> Result<()> {
    setup::is_setup()?;
    show_header("make-header");
    let mut makeflags: Vec<OsString> = Vec::new();
    let cpus = num_cpus::get();
    makeflags.push(OsString::from(format!("--jobs={}", cpus)));
    let mut makefiles: Vec<OsString> = Vec::new();
    makefiles.push(OsString::from("-f"));
    makefiles.push(OsString::from(format!(
        "{}{}",
        CONFIGURE_DATADIR, "rules/casile.mk"
    )));
    let rules = status::get_rules()?;
    for rule in rules {
        makefiles.push(OsString::from("-f"));
        let p = rule.into_os_string();
        makefiles.push(p);
    }
    makefiles.push(OsString::from("-f"));
    makefiles.push(OsString::from(format!(
        "{}{}",
        CONFIGURE_DATADIR, "rules/rules.mk"
    )));
    let mut process = Exec::cmd("make")
        .args(&makeflags)
        .args(&makefiles)
        .args(&target);
    let mut targets: Vec<_> = target.into_iter().collect();
    if status::is_gha()? {
        targets.push(String::from("_gha"));
        if targets.len() == 1 {
            targets.push(String::from("default"));
        }
        targets.push(String::from("install-dist"));
    }
    let gitname = status::get_gitname()?;
    let git_version = status::get_git_version();
    process = process
        .env("CASILE_CLI", "true")
        .env("CASILEDIR", CONFIGURE_DATADIR)
        .env("CONTAINERIZED", status::is_container().to_string())
        .env("LANGUAGE", lang_to_language(CONF.get_string("language")?))
        .env("PROJECT", &gitname)
        .env("PROJECTDIR", CONF.get_string("path")?)
        .env("PROJECTVERSION", git_version);
    if CONF.get_bool("debug")? {
        process = process.env("DEBUG", "true");
    };
    if CONF.get_bool("quiet")? {
        process = process.env("QUIET", "true");
    };
    if CONF.get_bool("verbose")? || targets.contains(&"debug".into()) {
        process = process.env("VERBOSE", "true");
    };
    let repo = get_repo()?;
    let workdir = repo.workdir().unwrap();
    process = process.cwd(workdir);
    let process = process.stderr(Redirection::Merge).stdout(Redirection::Pipe);
    let mut popen = process.popen()?;
    let buf = io::BufReader::new(popen.stdout.as_mut().unwrap());
    let mut backlog: Vec<String> = Vec::new();
    let seps = Regex::new(r"").unwrap();
    let mut ret: u32 = 0;
    for line in buf.lines() {
        let text: &str = &line.unwrap();
        let fields: Vec<&str> = seps.splitn(text, 4).collect();
        match fields[0] {
            "CASILE" => match fields[1] {
                "PRE" => report_start(fields[2]),
                "STDOUT" => {
                    if status::is_gha()? {
                        println!("{}", fields[3]);
                    } else if CONF.get_bool("verbose")? {
                        report_line(fields[3]);
                    } else {
                        backlog.push(String::from(fields[3]));
                    }
                }
                "STDERR" => {
                    if CONF.get_bool("verbose")? {
                        report_line(fields[3]);
                    } else {
                        backlog.push(String::from(fields[3]));
                    }
                }
                "POST" => match fields[2] {
                    "0" => {
                        report_end(fields[3]);
                    }
                    val => {
                        report_fail(fields[3]);
                        ret = val.parse().unwrap_or(1);
                    }
                },
                _ => {
                    let errmsg = LocalText::new("make-error-unknown-code").fmt();
                    panic!(errmsg)
                }
            },
            _ => backlog.push(String::from(fields[0])),
        }
    }
    let status = popen.wait();
    match status {
        Ok(ExitStatus::Exited(int)) => {
            let code = int + ret;
            match code {
                0 => {
                    if targets.contains(&"debug".into()) {
                        dump_backlog(&backlog)
                    };
                    Ok(())
                }
                1 => {
                    dump_backlog(&backlog);
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        LocalText::new("make-error-unfinished").fmt(),
                    )))
                }
                2 => {
                    dump_backlog(&backlog);
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        LocalText::new("make-error-build").fmt(),
                    )))
                }
                3 => {
                    if !CONF.get_bool("verbose")? {
                        dump_backlog(&backlog);
                    }
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        LocalText::new("make-error-target").fmt(),
                    )))
                }
                _ => {
                    dump_backlog(&backlog);
                    Err(Box::new(io::Error::new(
                        io::ErrorKind::InvalidInput,
                        LocalText::new("make-error-unknown").fmt(),
                    )))
                }
            }
        }
        _ => Err(Box::new(io::Error::new(
            io::ErrorKind::InvalidInput,
            LocalText::new("make-error").fmt(),
        ))),
    }
}

fn dump_backlog(backlog: &[String]) {
    let start = LocalText::new("make-backlog-start").fmt();
    eprintln!("{} {}", "┖┄".cyan(), start);
    for line in backlog.iter() {
        eprintln!("{}", line);
    }
    let end = LocalText::new("make-backlog-end").fmt();
    eprintln!("{} {}", "┎┄".cyan(), end);
}

fn report_line(line: &str) {
    eprintln!("{} {}", "┠╎".cyan(), line.dimmed());
}

fn report_start(target: &str) {
    let text = LocalText::new("make-report-start")
        .arg("target", target.white().bold())
        .fmt();
    eprintln!("{} {}", "┠┄".cyan(), text.yellow());
}

fn report_end(target: &str) {
    let text = LocalText::new("make-report-end")
        .arg("target", target.white().bold())
        .fmt();
    eprintln!("{} {}", "┠┄".cyan(), text.green());
}

fn report_fail(target: &str) {
    let text = LocalText::new("make-report-fail")
        .arg("target", target.white().bold())
        .fmt();
    eprintln!("{} {}", "┠┄".cyan(), text.red());
}
<|endoftext|>


<commit_before>use futures::{Future, BoxFuture};
use linear_map::LinearMap;
use protobuf;

use mercury::MercuryError;
use protocol;
use session::Session;
use util::{SpotifyId, FileId, StrChunksExt};

pub use protocol::metadata::AudioFile_Format as FileFormat;

fn countrylist_contains(list: &str, country: &str) -> bool {
    list.chunks(2).any(|cc| cc == country)
}

fn parse_restrictions<'s, I>(restrictions: I, country: &str, catalogue: &str) -> bool
    where I: IntoIterator<Item = &'s protocol::metadata::Restriction>
{
    restrictions.into_iter()
                .filter(|r| r.get_catalogue_str().contains(&catalogue.to_owned()))
                .all(|r| {
                    !countrylist_contains(r.get_countries_forbidden(), country) &&
                    (!r.has_countries_allowed() ||
                     countrylist_contains(r.get_countries_allowed(), country))
                })
}

pub trait MetadataTrait : Send + 'static {
    type Message: protobuf::MessageStatic;

    fn base_url() -> &'static str;
    fn parse(msg: &Self::Message, session: &Session) -> Self;
}

#[derive(Debug, Clone)]
pub struct Track {
    pub id: SpotifyId,
    pub name: String,
    pub album: SpotifyId,
    pub artists: Vec<SpotifyId>,
    pub files: LinearMap<FileFormat, FileId>,
    pub alternatives: Vec<SpotifyId>,
    pub available: bool,
}

#[derive(Debug, Clone)]
pub struct Album {
    pub id: SpotifyId,
    pub name: String,
    pub artists: Vec<SpotifyId>,
    pub tracks: Vec<SpotifyId>,
    pub covers: Vec<FileId>,
}

#[derive(Debug, Clone)]
pub struct Artist {
    pub id: SpotifyId,
    pub name: String,
    pub top_tracks: Vec<SpotifyId>,
}

impl MetadataTrait for Track {
    type Message = protocol::metadata::Track;

    fn base_url() -> &'static str {
        "hm://metadata/3/track"
    }

    fn parse(msg: &Self::Message, session: &Session) -> Self {
        let country = session.country();

        let artists = msg.get_artist()
                         .iter()
                         .filter(|artist| artist.has_gid())
                         .map(|artist| SpotifyId::from_raw(artist.get_gid()))
                         .collect::<Vec<_>>();

        let files = msg.get_file()
                       .iter()
                       .filter(|file| file.has_file_id())
                       .map(|file| {
                           let mut dst = [0u8; 20];
                           dst.clone_from_slice(file.get_file_id());
                           (file.get_format(), FileId(dst))
                       })
                       .collect();

        Track {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            album: SpotifyId::from_raw(msg.get_album().get_gid()),
            artists: artists,
            files: files,
            alternatives: msg.get_alternative()
                             .iter()
                             .map(|alt| SpotifyId::from_raw(alt.get_gid()))
                             .collect(),
            available: parse_restrictions(msg.get_restriction(),
                                          &country,
                                          "premium"),
        }
    }
}

impl MetadataTrait for Album {
    type Message = protocol::metadata::Album;

    fn base_url() -> &'static str {
        "hm://metadata/3/album"
    }

    fn parse(msg: &Self::Message, _: &Session) -> Self {
        let artists = msg.get_artist()
                         .iter()
                         .filter(|artist| artist.has_gid())
                         .map(|artist| SpotifyId::from_raw(artist.get_gid()))
                         .collect::<Vec<_>>();

        let tracks = msg.get_disc()
                        .iter()
                        .flat_map(|disc| disc.get_track())
                        .filter(|track| track.has_gid())
                        .map(|track| SpotifyId::from_raw(track.get_gid()))
                        .collect::<Vec<_>>();

        let covers = msg.get_cover_group()
                        .get_image()
                        .iter()
                        .filter(|image| image.has_file_id())
                        .map(|image| {
                            let mut dst = [0u8; 20];
                            dst.clone_from_slice(image.get_file_id());
                            FileId(dst)
                        })
                        .collect::<Vec<_>>();

        Album {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            artists: artists,
            tracks: tracks,
            covers: covers,
        }
    }
}

impl MetadataTrait for Artist {
    type Message = protocol::metadata::Artist;

    fn base_url() -> &'static str {
        "hm://metadata/3/artist"
    }

    fn parse(msg: &Self::Message, session: &Session) -> Self {
        let country = session.country();

        let top_tracks = msg.get_top_track()
                            .iter()
                            .find(|tt| !tt.has_country() || countrylist_contains(tt.get_country(), &country))
                            .unwrap()
                            .get_track()
                            .iter()
                            .filter(|track| track.has_gid())
                            .map(|track| SpotifyId::from_raw(track.get_gid()))
                            .collect::<Vec<_>>();

        Artist {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            top_tracks: top_tracks
        }
    }
}

component! {
    MetadataManager : MetadataManagerInner { }
}

impl MetadataManager {
    pub fn get<T: MetadataTrait>(&self, id: SpotifyId) -> BoxFuture<T, MercuryError> {
        let session = self.session();

        let uri = format!("{}/{}", T::base_url(), id.to_base16());
        let request = session.mercury().get(uri);

        request.and_then(move |response| {
            let data = response.payload.first().expect("Empty payload");
            let msg: T::Message = protobuf::parse_from_bytes(data).unwrap();

            Ok(T::parse(&msg, &session))
        }).boxed()
    }
}
<commit_msg>Fix #173 (a track cannot be played with librespot)<commit_after>use futures::{Future, BoxFuture};
use linear_map::LinearMap;
use protobuf;

use mercury::MercuryError;
use protocol;
use session::Session;
use util::{SpotifyId, FileId, StrChunksExt};

pub use protocol::metadata::AudioFile_Format as FileFormat;

fn countrylist_contains(list: &str, country: &str) -> bool {
    list.chunks(2).any(|cc| cc == country)
}

fn parse_restrictions<'s, I>(restrictions: I, country: &str, catalogue: &str) -> bool
    where I: IntoIterator<Item = &'s protocol::metadata::Restriction>
{
    let mut forbidden = "".to_string();
    let mut allowed = "".to_string();
    let rs = restrictions.into_iter().filter(|r|
        r.get_catalogue_str().contains(&catalogue.to_owned())
    );

    for r in rs {
        forbidden.push_str(r.get_countries_forbidden());
        allowed.push_str(r.get_countries_allowed());
    }

    (forbidden == "" || !countrylist_contains(forbidden.as_str(), country)) &&
    (allowed == "" || countrylist_contains(allowed.as_str(), country))
}

pub trait MetadataTrait : Send + 'static {
    type Message: protobuf::MessageStatic;

    fn base_url() -> &'static str;
    fn parse(msg: &Self::Message, session: &Session) -> Self;
}

#[derive(Debug, Clone)]
pub struct Track {
    pub id: SpotifyId,
    pub name: String,
    pub album: SpotifyId,
    pub artists: Vec<SpotifyId>,
    pub files: LinearMap<FileFormat, FileId>,
    pub alternatives: Vec<SpotifyId>,
    pub available: bool,
}

#[derive(Debug, Clone)]
pub struct Album {
    pub id: SpotifyId,
    pub name: String,
    pub artists: Vec<SpotifyId>,
    pub tracks: Vec<SpotifyId>,
    pub covers: Vec<FileId>,
}

#[derive(Debug, Clone)]
pub struct Artist {
    pub id: SpotifyId,
    pub name: String,
    pub top_tracks: Vec<SpotifyId>,
}

impl MetadataTrait for Track {
    type Message = protocol::metadata::Track;

    fn base_url() -> &'static str {
        "hm://metadata/3/track"
    }

    fn parse(msg: &Self::Message, session: &Session) -> Self {
        let country = session.country();

        let artists = msg.get_artist()
                         .iter()
                         .filter(|artist| artist.has_gid())
                         .map(|artist| SpotifyId::from_raw(artist.get_gid()))
                         .collect::<Vec<_>>();

        let files = msg.get_file()
                       .iter()
                       .filter(|file| file.has_file_id())
                       .map(|file| {
                           let mut dst = [0u8; 20];
                           dst.clone_from_slice(file.get_file_id());
                           (file.get_format(), FileId(dst))
                       })
                       .collect();

        Track {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            album: SpotifyId::from_raw(msg.get_album().get_gid()),
            artists: artists,
            files: files,
            alternatives: msg.get_alternative()
                             .iter()
                             .map(|alt| SpotifyId::from_raw(alt.get_gid()))
                             .collect(),
            available: parse_restrictions(msg.get_restriction(),
                                          &country,
                                          "premium"),
        }
    }
}

impl MetadataTrait for Album {
    type Message = protocol::metadata::Album;

    fn base_url() -> &'static str {
        "hm://metadata/3/album"
    }

    fn parse(msg: &Self::Message, _: &Session) -> Self {
        let artists = msg.get_artist()
                         .iter()
                         .filter(|artist| artist.has_gid())
                         .map(|artist| SpotifyId::from_raw(artist.get_gid()))
                         .collect::<Vec<_>>();

        let tracks = msg.get_disc()
                        .iter()
                        .flat_map(|disc| disc.get_track())
                        .filter(|track| track.has_gid())
                        .map(|track| SpotifyId::from_raw(track.get_gid()))
                        .collect::<Vec<_>>();

        let covers = msg.get_cover_group()
                        .get_image()
                        .iter()
                        .filter(|image| image.has_file_id())
                        .map(|image| {
                            let mut dst = [0u8; 20];
                            dst.clone_from_slice(image.get_file_id());
                            FileId(dst)
                        })
                        .collect::<Vec<_>>();

        Album {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            artists: artists,
            tracks: tracks,
            covers: covers,
        }
    }
}

impl MetadataTrait for Artist {
    type Message = protocol::metadata::Artist;

    fn base_url() -> &'static str {
        "hm://metadata/3/artist"
    }

    fn parse(msg: &Self::Message, session: &Session) -> Self {
        let country = session.country();

        let top_tracks = msg.get_top_track()
                            .iter()
                            .find(|tt| !tt.has_country() || countrylist_contains(tt.get_country(), &country))
                            .unwrap()
                            .get_track()
                            .iter()
                            .filter(|track| track.has_gid())
                            .map(|track| SpotifyId::from_raw(track.get_gid()))
                            .collect::<Vec<_>>();

        Artist {
            id: SpotifyId::from_raw(msg.get_gid()),
            name: msg.get_name().to_owned(),
            top_tracks: top_tracks
        }
    }
}

component! {
    MetadataManager : MetadataManagerInner { }
}

impl MetadataManager {
    pub fn get<T: MetadataTrait>(&self, id: SpotifyId) -> BoxFuture<T, MercuryError> {
        let session = self.session();

        let uri = format!("{}/{}", T::base_url(), id.to_base16());
        let request = session.mercury().get(uri);

        request.and_then(move |response| {
            let data = response.payload.first().expect("Empty payload");
            let msg: T::Message = protobuf::parse_from_bytes(data).unwrap();

            Ok(T::parse(&msg, &session))
        }).boxed()
    }
}
<|endoftext|>