<commit_before>// vim: tw=80
use super::*;

use quote::ToTokens;

use crate::mock_function::MockFunction;

pub(crate) struct MockItemStruct {
    generics: Generics,
    /// Inherent methods of the mock struct
    methods: Vec<MockFunction>,
    name: Ident,
    //traits: Vec<MockTraitImpl>
    vis: Visibility,
}

impl From<MockableStruct> for MockItemStruct {
    fn from(mockable: MockableStruct) -> MockItemStruct {
        let mock_ident = gen_mod_ident(&mockable.name, None);
        let methods = Vec::new();
        //let methods = mockable.methods.into_iter()
            //.map(|meth| MockFunction::from((&mock_ident, 1, meth)))
            //.collect::<Vec<_>>();
        MockItemStruct {
            generics: mockable.generics,
            methods,
            name: mockable.name,
            //traits: Vec::new()
            vis: mockable.vis
        }
    }
}

impl ToTokens for MockItemStruct {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let struct_name = &self.name;
        let (ig, tg, wc) = self.generics.split_for_impl(); //TODO
        let vis = &self.vis;
        let mut default_body = TokenStream::new();  // TODO
        quote!(
            //#[allow(non_snake_case)]
            //#[doc(hidden)]
            //pub mod mod_ident {
                //use super::*;
                ////mod_body
            //}
            #[allow(non_camel_case_types)]
            #[allow(non_snake_case)]
            #[allow(missing_docs)]
            //attr_ts
            #vis struct #struct_name
            //ig wc
            {
                //body
            }
            impl #ig ::std::default::Default for #struct_name #tg #wc {
                fn default() -> Self {
                    Self {
                        #default_body
                    }
                }
            }
            impl #ig #struct_name #tg #wc {}
        ).to_tokens(tokens);
    }
}
<commit_msg>WIP mocking structs<commit_after>// vim: tw=80
use super::*;

use quote::ToTokens;

use crate::mock_function::MockFunction;

pub(crate) struct MockItemStruct {
    generics: Generics,
    /// Inherent methods of the mock struct
    methods: Vec<MockFunction>,
    name: Ident,
    //traits: Vec<MockTraitImpl>
    vis: Visibility,
}

impl From<MockableStruct> for MockItemStruct {
    fn from(mockable: MockableStruct) -> MockItemStruct {
        let mock_ident = gen_mod_ident(&mockable.name, None);
        let vis = Visibility::Public(VisPublic{
            pub_token: Token![pub](Span::call_site())
        });
        let methods = mockable.methods.into_iter()
            .map(|meth|
                MockFunction::from((&mock_ident, 1, meth.sig, vis.clone()))
            ).collect::<Vec<_>>();
        MockItemStruct {
            generics: mockable.generics,
            methods,
            name: mockable.name,
            //traits: Vec::new()
            vis: mockable.vis
        }
    }
}

impl ToTokens for MockItemStruct {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let struct_name = &self.name;
        let (ig, tg, wc) = self.generics.split_for_impl(); //TODO
        let methods = &self.methods;
        let vis = &self.vis;
        let mut default_body = TokenStream::new();  // TODO
        quote!(
            #[allow(non_snake_case)]
            #[doc(hidden)]
            pub mod mod_ident {
                use super::*;
                #(#methods)*
                //mod_body
            }
            #[allow(non_camel_case_types)]
            #[allow(non_snake_case)]
            #[allow(missing_docs)]
            //attr_ts
            #vis struct #struct_name
            //ig wc
            {
                //body
            }
            impl #ig ::std::default::Default for #struct_name #tg #wc {
                fn default() -> Self {
                    Self {
                        #default_body
                    }
                }
            }
            impl #ig #struct_name #tg #wc {}
        ).to_tokens(tokens);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use codespan_reporting::diagnostic::{Diagnostic, Label};
use itertools::Itertools;
#[allow(unused_imports)]
use log::warn;
use std::collections::BTreeSet;

use builder::module_builder::ModuleBuilder;
use move_lang::{
    compiled_unit::{self, CompiledUnit},
    errors::Errors,
    expansion::ast::{ModuleDefinition, Program},
    move_continue_up_to, move_parse,
    parser::ast::ModuleIdent,
    shared::{unique_map::UniqueMap, Address, CompilationEnv, Flags},
    Pass as MovePass, PassResult as MovePassResult,
};
use vm::{
    access::ModuleAccess,
    file_format::{CompiledModule, FunctionDefinitionIndex, StructDefinitionIndex},
};

use crate::{
    ast::{ModuleName, Spec},
    builder::model_builder::ModelBuilder,
    model::{FunId, FunctionData, GlobalEnv, Loc, ModuleData, ModuleEnv, ModuleId, StructId},
};

pub mod ast;
mod builder;
pub mod code_writer;
pub mod exp_generator;
pub mod exp_rewriter;
pub mod model;
pub mod pragmas;
pub mod spec_translator;
pub mod symbol;
pub mod ty;

// =================================================================================================
// Entry Point

/// Build the move model. This collects transitive dependencies for move sources
/// from the provided directory list.
pub fn run_model_builder(
    move_sources: &[String],
    deps_dir: &[String],
) -> anyhow::Result<GlobalEnv> {
    let mut env = GlobalEnv::new();

    // Step 1: parse the program to get comments and a separation of targets and dependencies.
    let (files, pprog_and_comments_res) = move_parse(
        move_sources,
        deps_dir,
        None,
        /* sources_shadow_deps */ true,
    )?;
    let (comment_map, parsed_prog) = match pprog_and_comments_res {
        Err(errors) => {
            // Add source files so that the env knows how to translate locations of parse errors
            for fname in files.keys().sorted() {
                let fsrc = &files[fname];
                env.add_source(fname, fsrc, /* is_dep */ false);
            }
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(res) => res,
    };
    // Add source files for targets and dependencies
    let dep_sources: BTreeSet<_> = parsed_prog
        .lib_definitions
        .iter()
        .map(|def| def.file())
        .collect();
    for fname in files.keys().sorted() {
        let fsrc = &files[fname];
        env.add_source(fname, fsrc, dep_sources.contains(fname));
    }
    // Add any documentation comments found by the Move compiler to the env.
    for (fname, documentation) in comment_map {
        let file_id = env.get_file_id(fname).expect("file name defined");
        env.add_documentation(file_id, documentation);
    }

    // Step 2: compile all sources in targets and dependencies, as we need bytecode for all of them.
    let all_sources: Vec<_> = files
        .into_iter()
        .map(|(fname, _)| fname.to_owned())
        .collect();
    // Run the compiler up to expansion and clone a copy of the expansion program ast
    let (_, pprog_and_comments_res) = move_parse(&all_sources, &[], None, false)?;
    let (_, parsed_prog) = match pprog_and_comments_res {
        Err(errors) => {
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(res) => res,
    };
    let mut compilation_env = CompilationEnv::new(Flags::empty());
    let (expansion_ast, expansion_result) = match move_continue_up_to(
        &mut compilation_env,
        None,
        MovePassResult::Parser(parsed_prog),
        MovePass::Expansion,
    ) {
        Err(errors) => {
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(MovePassResult::Expansion(eprog)) => (eprog.clone(), MovePassResult::Expansion(eprog)),
        Ok(_) => unreachable!(),
    };
    // Run the compiler fully to the compiled units
    let units = match move_continue_up_to(
        &mut compilation_env,
        None,
        expansion_result,
        MovePass::Compilation,
    ) {
        Err(errors) => {
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(MovePassResult::Compilation(units)) => units,
        Ok(_) => unreachable!(),
    };
    // Check for bytecode verifier errors (there should not be any)
    let (verified_units, errors) = compiled_unit::verify_units(units);
    if !errors.is_empty() {
        add_move_lang_errors(&mut env, errors);
        return Ok(env);
    }

    // Now that it is known that the program has no errors, run the spec checker on verified units
    // plus expanded AST. This will populate the environment including any errors.
    run_spec_checker(&mut env, verified_units, expansion_ast);
    Ok(env)
}

/// Build a `GlobalEnv` from a collection of `CompiledModule`'s. The `modules` list must be
/// topologically sorted by the dependency relation (i.e., a child node in the dependency graph
/// should appear earlier in the vector than its parents).
pub fn run_bytecode_model_builder(modules: Vec<CompiledModule>) -> anyhow::Result<GlobalEnv> {
    let mut env = GlobalEnv::new();
    for (i, m) in modules.into_iter().enumerate() {
        let id = m.self_id();
        let addr = ModuleEnv::addr_to_big_uint(id.address());
        let module_name = ModuleName::new(addr, env.symbol_pool().make(id.name().as_str()));
        let module_id = ModuleId::new(i);
        let mut module_data = ModuleData::stub(module_name.clone(), module_id, m.clone());

        // add functions
        for (def_idx, def) in m.function_defs().iter().enumerate() {
            let name = m.identifier_at(m.function_handle_at(def.function).name);
            let symbol = env.symbol_pool().make(name.as_str());
            let data = FunctionData::stub(
                symbol,
                FunctionDefinitionIndex(def_idx as u16),
                def.function,
            );
            module_data.function_data.insert(FunId::new(symbol), data);
        }

        // add structs
        for (def_idx, def) in m.struct_defs().iter().enumerate() {
            let name = m.identifier_at(m.struct_handle_at(def.struct_handle).name);
            let symbol = env.symbol_pool().make(name.as_str());
            let data = env.create_struct_data(
                &m,
                StructDefinitionIndex(def_idx as u16),
                symbol,
                Loc::default(),
                Spec::default(),
            );
            module_data.struct_data.insert(StructId::new(symbol), data);
        }

        env.module_data.push(module_data);
    }
    Ok(env)
}

fn add_move_lang_errors(env: &mut GlobalEnv, errors: Errors) {
    let mk_label = |env: &mut GlobalEnv, err: (move_ir_types::location::Loc, String)| {
        let loc = env.to_loc(&err.0);
        Label::new(loc.file_id(), loc.span(), err.1)
    };
    for mut error in errors {
        let primary = error.remove(0);
        let diag = Diagnostic::new_error("", mk_label(env, primary))
            .with_secondary_labels(error.into_iter().map(|e| mk_label(env, e)));
        env.add_diag(diag);
    }
}

#[allow(deprecated)]
fn run_spec_checker(env: &mut GlobalEnv, units: Vec<CompiledUnit>, mut eprog: Program) {
    let mut builder = ModelBuilder::new(env);
    // Merge the compiled units with the expanded program, preserving the order of the compiled
    // units which is topological w.r.t. use relation.
    let modules = units
        .into_iter()
        .flat_map(|unit| {
            Some(match unit {
                CompiledUnit::Module {
                    ident,
                    module,
                    source_map,
                    function_infos,
                } => {
                    let expanded_module = match eprog.modules.remove(&ident) {
                        Some(m) => m,
                        None => {
                            warn!(
                                "[internal] cannot associate bytecode module `{}` with AST",
                                ident
                            );
                            return None;
                        }
                    };
                    (ident, expanded_module, module, source_map, function_infos)
                }
                CompiledUnit::Script {
                    loc: _loc,
                    key,
                    script,
                    source_map,
                    function_info,
                } => {
                    let move_lang::expansion::ast::Script {
                        attributes,
                        loc,
                        dependency_summary,
                        function_name,
                        constants,
                        function,
                        specs,
                    } = match eprog.scripts.remove(&key) {
                        Some(s) => s,
                        None => {
                            warn!(
                                "[internal] cannot associate bytecode script `{}` with AST",
                                key
                            );
                            return None;
                        }
                    };
                    // Convert the script into a module.
                    let ident = ModuleIdent {
                        locs: (loc, loc),
                        value: (Address::default(), function_name.0.value.clone()),
                    };
                    let mut function_infos = UniqueMap::new();
                    function_infos
                        .add(function_name.clone(), function_info)
                        .unwrap();
                    // Construct a pseudo module definition.
                    let mut functions = UniqueMap::new();
                    functions.add(function_name, function).unwrap();
                    let expanded_module = ModuleDefinition {
                        attributes,
                        loc,
                        dependency_order: usize::MAX,
                        dependency_summary,
                        is_source_module: true,
                        friends: UniqueMap::new(),
                        structs: UniqueMap::new(),
                        constants,
                        functions,
                        specs,
                    };
                    let module = script.into_module().1;
                    (ident, expanded_module, module, source_map, function_infos)
                }
            })
        })
        .enumerate();
    for (module_count, (module_id, expanded_module, compiled_module, source_map, function_infos)) in
        modules
    {
        let loc = builder.to_loc(&expanded_module.loc);
        let module_name = ModuleName::from_str(
            &module_id.value.0.to_string(),
            builder.env.symbol_pool().make(&module_id.value.1),
        );
        let module_id = ModuleId::new(module_count);
        let mut module_translator = ModuleBuilder::new(&mut builder, module_id, module_name);
        module_translator.translate(
            loc,
            expanded_module,
            compiled_module,
            source_map,
            function_infos,
        );
    }
    // After all specs have been processed, warn about any unused schemas.
    builder.warn_unused_schemas();
}

// =================================================================================================
// Crate Helpers

/// Helper to project the 1st element from a vector of pairs.
pub(crate) fn project_1st<T: Clone, R>(v: &[(T, R)]) -> Vec<T> {
    v.iter().map(|(x, _)| x.clone()).collect()
}

/// Helper to project the 2nd element from a vector of pairs.
pub(crate) fn project_2nd<T, R: Clone>(v: &[(T, R)]) -> Vec<R> {
    v.iter().map(|(_, x)| x.clone()).collect()
}
<commit_msg>[move-model] compile only related modules in the last phase of module building<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use codespan_reporting::diagnostic::{Diagnostic, Label};
use itertools::Itertools;
#[allow(unused_imports)]
use log::warn;
use std::collections::BTreeSet;

use builder::module_builder::ModuleBuilder;
use move_lang::{
    compiled_unit::{self, CompiledUnit},
    errors::Errors,
    expansion::ast::{ModuleDefinition, Program},
    move_continue_up_to, move_parse,
    parser::ast::ModuleIdent,
    shared::{unique_map::UniqueMap, Address, CompilationEnv, Flags},
    Pass as MovePass, PassResult as MovePassResult,
};
use vm::{
    access::ModuleAccess,
    file_format::{CompiledModule, FunctionDefinitionIndex, StructDefinitionIndex},
};

use crate::{
    ast::{ModuleName, Spec},
    builder::model_builder::ModelBuilder,
    model::{FunId, FunctionData, GlobalEnv, Loc, ModuleData, ModuleEnv, ModuleId, StructId},
};

pub mod ast;
mod builder;
pub mod code_writer;
pub mod exp_generator;
pub mod exp_rewriter;
pub mod model;
pub mod pragmas;
pub mod spec_translator;
pub mod symbol;
pub mod ty;

// =================================================================================================
// Entry Point

/// Build the move model. This collects transitive dependencies for move sources
/// from the provided directory list.
pub fn run_model_builder(
    move_sources: &[String],
    deps_dir: &[String],
) -> anyhow::Result<GlobalEnv> {
    let mut env = GlobalEnv::new();

    // Step 1: parse the program to get comments and a separation of targets and dependencies.
    let (files, pprog_and_comments_res) = move_parse(
        move_sources,
        deps_dir,
        None,
        /* sources_shadow_deps */ true,
    )?;
    let (comment_map, parsed_prog) = match pprog_and_comments_res {
        Err(errors) => {
            // Add source files so that the env knows how to translate locations of parse errors
            for fname in files.keys().sorted() {
                let fsrc = &files[fname];
                env.add_source(fname, fsrc, /* is_dep */ false);
            }
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(res) => res,
    };
    // Add source files for targets and dependencies
    let dep_sources: BTreeSet<_> = parsed_prog
        .lib_definitions
        .iter()
        .map(|def| def.file())
        .collect();
    for fname in files.keys().sorted() {
        let fsrc = &files[fname];
        env.add_source(fname, fsrc, dep_sources.contains(fname));
    }
    // Add any documentation comments found by the Move compiler to the env.
    for (fname, documentation) in comment_map {
        let file_id = env.get_file_id(fname).expect("file name defined");
        env.add_documentation(file_id, documentation);
    }

    // Step 2: parse all sources in targets and dependencies, prepare for selective compilation.
    let all_sources: Vec<_> = files
        .into_iter()
        .map(|(fname, _)| fname.to_owned())
        .collect();
    // Run the compiler up to expansion
    let (_, pprog_and_comments_res) = move_parse(&all_sources, &[], None, false)?;
    let (_, parsed_prog) = match pprog_and_comments_res {
        Err(errors) => {
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(res) => res,
    };
    let mut compilation_env = CompilationEnv::new(Flags::empty());
    let expansion_ast = match move_continue_up_to(
        &mut compilation_env,
        None,
        MovePassResult::Parser(parsed_prog),
        MovePass::Expansion,
    ) {
        Err(errors) => {
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(MovePassResult::Expansion(eprog)) => eprog,
        Ok(_) => unreachable!(),
    };
    // Extract the module/script closure
    let mut visited_modules = BTreeSet::new();
    let mut selective_files = BTreeSet::new();
    for (mident, mdef) in expansion_ast.modules.key_cloned_iter() {
        let src_file = mdef.loc.file();
        if !dep_sources.contains(src_file) {
            selective_files.insert(src_file.to_owned());
            collect_related_modules_recursive(mident, &expansion_ast.modules, &mut visited_modules);
        }
    }
    for sdef in expansion_ast.scripts.values() {
        let src_file = sdef.loc.file();
        if !dep_sources.contains(src_file) {
            selective_files.insert(src_file.to_owned());
            for mident in &sdef.dependency_summary {
                collect_related_modules_recursive(
                    mident.clone(),
                    &expansion_ast.modules,
                    &mut visited_modules,
                );
            }
        }
    }
    for mident in visited_modules {
        selective_files.insert(
            expansion_ast
                .modules
                .get(&mident)
                .unwrap()
                .loc
                .file()
                .to_owned(),
        );
    }

    // Step 3: selective compilation.
    let selective_sources: Vec<_> = selective_files.into_iter().collect();

    // Run the compiler up to expansion and clone a copy of the expansion program ast
    let (_, pprog_and_comments_res) = move_parse(&selective_sources, &[], None, false)?;
    let (_, parsed_prog) = match pprog_and_comments_res {
        Err(errors) => {
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(res) => res,
    };
    let mut compilation_env = CompilationEnv::new(Flags::empty());
    let (expansion_ast, expansion_result) = match move_continue_up_to(
        &mut compilation_env,
        None,
        MovePassResult::Parser(parsed_prog),
        MovePass::Expansion,
    ) {
        Err(errors) => {
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(MovePassResult::Expansion(eprog)) => (eprog.clone(), MovePassResult::Expansion(eprog)),
        Ok(_) => unreachable!(),
    };
    // Run the compiler fully to the compiled units
    let units = match move_continue_up_to(
        &mut compilation_env,
        None,
        expansion_result,
        MovePass::Compilation,
    ) {
        Err(errors) => {
            add_move_lang_errors(&mut env, errors);
            return Ok(env);
        }
        Ok(MovePassResult::Compilation(units)) => units,
        Ok(_) => unreachable!(),
    };
    // Check for bytecode verifier errors (there should not be any)
    let (verified_units, errors) = compiled_unit::verify_units(units);
    if !errors.is_empty() {
        add_move_lang_errors(&mut env, errors);
        return Ok(env);
    }

    // Now that it is known that the program has no errors, run the spec checker on verified units
    // plus expanded AST. This will populate the environment including any errors.
    run_spec_checker(&mut env, verified_units, expansion_ast);
    Ok(env)
}

fn collect_related_modules_recursive(
    mident: ModuleIdent,
    modules: &UniqueMap<ModuleIdent, ModuleDefinition>,
    visited: &mut BTreeSet<ModuleIdent>,
) {
    if visited.contains(&mident) {
        return;
    }
    let mdef = modules.get(&mident).unwrap();
    let deps: BTreeSet<_> = mdef.dependency_summary.iter().cloned().collect();
    visited.insert(mident);
    for next_mident in deps {
        collect_related_modules_recursive(next_mident, modules, visited);
    }
}

/// Build a `GlobalEnv` from a collection of `CompiledModule`'s. The `modules` list must be
/// topologically sorted by the dependency relation (i.e., a child node in the dependency graph
/// should appear earlier in the vector than its parents).
pub fn run_bytecode_model_builder(modules: Vec<CompiledModule>) -> anyhow::Result<GlobalEnv> {
    let mut env = GlobalEnv::new();
    for (i, m) in modules.into_iter().enumerate() {
        let id = m.self_id();
        let addr = ModuleEnv::addr_to_big_uint(id.address());
        let module_name = ModuleName::new(addr, env.symbol_pool().make(id.name().as_str()));
        let module_id = ModuleId::new(i);
        let mut module_data = ModuleData::stub(module_name.clone(), module_id, m.clone());

        // add functions
        for (def_idx, def) in m.function_defs().iter().enumerate() {
            let name = m.identifier_at(m.function_handle_at(def.function).name);
            let symbol = env.symbol_pool().make(name.as_str());
            let data = FunctionData::stub(
                symbol,
                FunctionDefinitionIndex(def_idx as u16),
                def.function,
            );
            module_data.function_data.insert(FunId::new(symbol), data);
        }

        // add structs
        for (def_idx, def) in m.struct_defs().iter().enumerate() {
            let name = m.identifier_at(m.struct_handle_at(def.struct_handle).name);
            let symbol = env.symbol_pool().make(name.as_str());
            let data = env.create_struct_data(
                &m,
                StructDefinitionIndex(def_idx as u16),
                symbol,
                Loc::default(),
                Spec::default(),
            );
            module_data.struct_data.insert(StructId::new(symbol), data);
        }

        env.module_data.push(module_data);
    }
    Ok(env)
}

fn add_move_lang_errors(env: &mut GlobalEnv, errors: Errors) {
    let mk_label = |env: &mut GlobalEnv, err: (move_ir_types::location::Loc, String)| {
        let loc = env.to_loc(&err.0);
        Label::new(loc.file_id(), loc.span(), err.1)
    };
    for mut error in errors {
        let primary = error.remove(0);
        let diag = Diagnostic::new_error("", mk_label(env, primary))
            .with_secondary_labels(error.into_iter().map(|e| mk_label(env, e)));
        env.add_diag(diag);
    }
}

#[allow(deprecated)]
fn run_spec_checker(env: &mut GlobalEnv, units: Vec<CompiledUnit>, mut eprog: Program) {
    let mut builder = ModelBuilder::new(env);
    // Merge the compiled units with the expanded program, preserving the order of the compiled
    // units which is topological w.r.t. use relation.
    let modules = units
        .into_iter()
        .flat_map(|unit| {
            Some(match unit {
                CompiledUnit::Module {
                    ident,
                    module,
                    source_map,
                    function_infos,
                } => {
                    let expanded_module = match eprog.modules.remove(&ident) {
                        Some(m) => m,
                        None => {
                            warn!(
                                "[internal] cannot associate bytecode module `{}` with AST",
                                ident
                            );
                            return None;
                        }
                    };
                    (ident, expanded_module, module, source_map, function_infos)
                }
                CompiledUnit::Script {
                    loc: _loc,
                    key,
                    script,
                    source_map,
                    function_info,
                } => {
                    let move_lang::expansion::ast::Script {
                        attributes,
                        loc,
                        dependency_summary,
                        function_name,
                        constants,
                        function,
                        specs,
                    } = match eprog.scripts.remove(&key) {
                        Some(s) => s,
                        None => {
                            warn!(
                                "[internal] cannot associate bytecode script `{}` with AST",
                                key
                            );
                            return None;
                        }
                    };
                    // Convert the script into a module.
                    let ident = ModuleIdent {
                        locs: (loc, loc),
                        value: (Address::default(), function_name.0.value.clone()),
                    };
                    let mut function_infos = UniqueMap::new();
                    function_infos
                        .add(function_name.clone(), function_info)
                        .unwrap();
                    // Construct a pseudo module definition.
                    let mut functions = UniqueMap::new();
                    functions.add(function_name, function).unwrap();
                    let expanded_module = ModuleDefinition {
                        attributes,
                        loc,
                        dependency_order: usize::MAX,
                        dependency_summary,
                        is_source_module: true,
                        friends: UniqueMap::new(),
                        structs: UniqueMap::new(),
                        constants,
                        functions,
                        specs,
                    };
                    let module = script.into_module().1;
                    (ident, expanded_module, module, source_map, function_infos)
                }
            })
        })
        .enumerate();
    for (module_count, (module_id, expanded_module, compiled_module, source_map, function_infos)) in
        modules
    {
        let loc = builder.to_loc(&expanded_module.loc);
        let module_name = ModuleName::from_str(
            &module_id.value.0.to_string(),
            builder.env.symbol_pool().make(&module_id.value.1),
        );
        let module_id = ModuleId::new(module_count);
        let mut module_translator = ModuleBuilder::new(&mut builder, module_id, module_name);
        module_translator.translate(
            loc,
            expanded_module,
            compiled_module,
            source_map,
            function_infos,
        );
    }
    // After all specs have been processed, warn about any unused schemas.
    builder.warn_unused_schemas();
}

// =================================================================================================
// Crate Helpers

/// Helper to project the 1st element from a vector of pairs.
pub(crate) fn project_1st<T: Clone, R>(v: &[(T, R)]) -> Vec<T> {
    v.iter().map(|(x, _)| x.clone()).collect()
}

/// Helper to project the 2nd element from a vector of pairs.
pub(crate) fn project_2nd<T, R: Clone>(v: &[(T, R)]) -> Vec<R> {
    v.iter().map(|(_, x)| x.clone()).collect()
}
<|endoftext|>


<commit_before>use cgmath::EuclideanSpace;
use gilrs::{Axis, Button, Gilrs};
use std::f64::consts::PI;
use winit::{
    event,
    event_loop::{ControlFlow, EventLoop},
};

fn compute_projection_matrix(width: f32, height: f32) -> cgmath::Matrix4<f32> {
    let aspect = width / height;
    let f = 1.0 / (45.0f32.to_radians() / aspect).tan();
    let near = 0.1;

    #[cfg_attr(rustfmt, rustfmt_skip)]
    cgmath::Matrix4::new(
        f/aspect,  0.0,  0.0,   0.0,
        0.0,       f,    0.0,   0.0,
        0.0,       0.0,  0.0,  -1.0,
        0.0,       0.0,  near,  0.0)
}

fn make_swapchain(
    device: &wgpu::Device,
    surface: &wgpu::Surface,
    width: u32,
    height: u32,
) -> wgpu::SwapChain {
    device.create_swap_chain(
        &surface,
        &wgpu::SwapChainDescriptor {
            usage: wgpu::TextureUsage::OUTPUT_ATTACHMENT,
            format: wgpu::TextureFormat::Bgra8UnormSrgb,
            width,
            height,
            present_mode: wgpu::PresentMode::Fifo, // disable vsync by switching to Mailbox,
        },
    )
}
fn make_depth_buffer(device: &wgpu::Device, width: u32, height: u32) -> wgpu::TextureView {
    device
        .create_texture(&wgpu::TextureDescriptor {
            size: wgpu::Extent3d { width, height, depth: 1 },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Depth32Float,
            usage: wgpu::TextureUsage::OUTPUT_ATTACHMENT,
            label: None,
        })
        .create_default_view()
}

fn main() {
    env_logger::init();

    let mapfile = terra::MapFileBuilder::build().unwrap();

    let event_loop = EventLoop::new();
    let window = winit::window::Window::new(&event_loop).unwrap();
    for monitor in window.available_monitors() {
        if monitor.video_modes().any(|mode| mode.size().width == 1920) {
            window.set_fullscreen(Some(winit::window::Fullscreen::Borderless(monitor)));
            break;
        }
    }
    let instance = wgpu::Instance::new(wgpu::BackendBit::VULKAN | wgpu::BackendBit::DX12);
    let surface = unsafe { instance.create_surface(&window) };
    let adapter =
        futures::executor::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::Default,
            compatible_surface: Some(&surface),
        }))
        .expect("Unable to create compatible wgpu adapter");
    let (device, mut queue) = futures::executor::block_on(adapter.request_device(
        &wgpu::DeviceDescriptor {
            features: wgpu::Features::TEXTURE_COMPRESSION_BC,
            limits: wgpu::Limits::default(),
            shader_validation: true,
        },
        None,
    ))
    .expect("Unable to create compatible wgpu device");

    let mut size = window.inner_size();
    let mut swap_chain = make_swapchain(&device, &surface, size.width, size.height);
    let mut depth_buffer = make_depth_buffer(&device, size.width, size.height);
    let mut proj = compute_projection_matrix(size.width as f32, size.height as f32);

    let mut gilrs = Gilrs::new().unwrap();
    let mut current_gamepad = None;
    for (_id, gamepad) in gilrs.gamepads() {
        current_gamepad = Some(gamepad.id());
    }

    let planet_radius = 6371000.0;
    let mut angle = 0.0f64;
    let mut lat = 41.3851f64.to_radians();
    let mut long = 2.1734f64.to_radians();
    let mut altitude = 200000.0f64;

    let mut terrain = terra::Terrain::new(&device, &mut queue, mapfile).unwrap();

    event_loop.run(move |event, _, control_flow| {
        *control_flow = if cfg!(feature = "metal-auto-capture") {
            ControlFlow::Exit
        } else {
            ControlFlow::Poll
        };
        match event {
            event::Event::WindowEvent { event, .. } => match event {
                event::WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit;
                }
                event::WindowEvent::KeyboardInput {
                    input:
                        event::KeyboardInput {
                            virtual_keycode: Some(keycode),
                            state: event::ElementState::Pressed,
                            ..
                        },
                    ..
                } => match keycode {
                    event::VirtualKeyCode::Escape => *control_flow = ControlFlow::Exit,
                    event::VirtualKeyCode::Space => altitude += 0.1 * altitude,
                    event::VirtualKeyCode::Semicolon => altitude -= 0.1 * altitude,
                    event::VirtualKeyCode::Left => {
                        lat += -angle.sin() * -(0.0000001 * altitude).min(0.01);
                        long += angle.cos() * -(0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Right => {
                        lat += -angle.sin() * (0.0000001 * altitude).min(0.01);
                        long += angle.cos() * (0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Up => {
                        lat += angle.cos() * (0.0000001 * altitude).min(0.01);
                        long += -angle.sin() * (0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Down => {
                        lat += angle.cos() * -(0.0000001 * altitude).min(0.01);
                        long += -angle.sin() * -(0.0000001 * altitude).min(0.01);
                    }
                    _ => {}
                },
                event::WindowEvent::Resized(new_size) => {
                    size = new_size;
                    swap_chain = make_swapchain(&device, &surface, size.width, size.height);
                    depth_buffer = make_depth_buffer(&device, size.width, size.height);
                    proj = compute_projection_matrix(size.width as f32, size.height as f32);
                }
                _ => {}
            },
            event::Event::MainEventsCleared => {
                let frame = swap_chain.get_current_frame();
                let frame = match frame {
                    Ok(ref f) => &f.output.view,
                    Err(_) => return,
                };

                while let Some(gilrs::Event { id, event: _event, time: _ }) = gilrs.next_event() {
                    current_gamepad = Some(id);
                }
                if let Some(gamepad) = current_gamepad.map(|id| gilrs.gamepad(id)) {
                    lat += angle.cos() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        - angle.sin() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    long += -angle.sin() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        + angle.cos() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    angle -= gamepad.value(Axis::RightZ) as f64 * 0.01;

                    if gamepad.is_pressed(Button::DPadUp) {
                        altitude += 0.01 * altitude;
                    }
                    if gamepad.is_pressed(Button::DPadDown) {
                        altitude -= 0.01 * altitude;
                    }
                }

                lat = lat.max(-PI).min(PI);
                if long < -PI {
                    long += PI * 2.0;
                }
                if long > PI {
                    long -= PI * 2.0;
                }

                let r = altitude + planet_radius;
                let eye = cgmath::Point3::new(
                    r * lat.cos() * long.cos(),
                    r * lat.cos() * long.sin(),
                    r * lat.sin(),
                );

                let latc = lat + angle.cos() * 0.1;
                let longc = long - angle.sin() * 0.1;

                let center = cgmath::Point3::new(
                    planet_radius * latc.cos() * longc.cos() - eye.x,
                    planet_radius * latc.cos() * longc.sin() - eye.y,
                    planet_radius * latc.sin() - eye.z,
                );
                let up = cgmath::Vector3::new(eye.x as f32, eye.y as f32, eye.z as f32);

                let view = cgmath::Matrix4::look_at(
                    cgmath::Point3::origin(),
                    cgmath::Point3::new(center.x as f32, center.y as f32, center.z as f32),
                    up,
                );

                let view_proj = proj * view;
                let view_proj = mint::ColumnMatrix4 {
                    x: view_proj.x.into(),
                    y: view_proj.y.into(),
                    z: view_proj.z.into(),
                    w: view_proj.w.into(),
                };

                terrain.render(
                    &device,
                    &mut queue,
                    &frame,
                    &depth_buffer,
                    (size.width, size.height),
                    view_proj,
                    eye.into(),
                );
            }
            _ => (),
        }
    });
}
<commit_msg>Tweak camera controls<commit_after>use cgmath::EuclideanSpace;
use gilrs::{Axis, Button, Gilrs};
use std::f64::consts::PI;
use winit::{
    event,
    event_loop::{ControlFlow, EventLoop},
};

fn compute_projection_matrix(width: f32, height: f32) -> cgmath::Matrix4<f32> {
    let aspect = width / height;
    let f = 1.0 / (45.0f32.to_radians() / aspect).tan();
    let near = 0.1;

    #[cfg_attr(rustfmt, rustfmt_skip)]
    cgmath::Matrix4::new(
        f/aspect,  0.0,  0.0,   0.0,
        0.0,       f,    0.0,   0.0,
        0.0,       0.0,  0.0,  -1.0,
        0.0,       0.0,  near,  0.0)
}

fn make_swapchain(
    device: &wgpu::Device,
    surface: &wgpu::Surface,
    width: u32,
    height: u32,
) -> wgpu::SwapChain {
    device.create_swap_chain(
        &surface,
        &wgpu::SwapChainDescriptor {
            usage: wgpu::TextureUsage::OUTPUT_ATTACHMENT,
            format: wgpu::TextureFormat::Bgra8UnormSrgb,
            width,
            height,
            present_mode: wgpu::PresentMode::Fifo, // disable vsync by switching to Mailbox,
        },
    )
}
fn make_depth_buffer(device: &wgpu::Device, width: u32, height: u32) -> wgpu::TextureView {
    device
        .create_texture(&wgpu::TextureDescriptor {
            size: wgpu::Extent3d { width, height, depth: 1 },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Depth32Float,
            usage: wgpu::TextureUsage::OUTPUT_ATTACHMENT,
            label: None,
        })
        .create_default_view()
}

fn main() {
    env_logger::init();

    let mapfile = terra::MapFileBuilder::build().unwrap();

    let event_loop = EventLoop::new();
    let window = winit::window::Window::new(&event_loop).unwrap();
    for monitor in window.available_monitors() {
        if monitor.video_modes().any(|mode| mode.size().width == 1920) {
            window.set_fullscreen(Some(winit::window::Fullscreen::Borderless(monitor)));
            break;
        }
    }
    let instance = wgpu::Instance::new(wgpu::BackendBit::VULKAN | wgpu::BackendBit::DX12);
    let surface = unsafe { instance.create_surface(&window) };
    let adapter =
        futures::executor::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
            power_preference: wgpu::PowerPreference::Default,
            compatible_surface: Some(&surface),
        }))
        .expect("Unable to create compatible wgpu adapter");
    let (device, mut queue) = futures::executor::block_on(adapter.request_device(
        &wgpu::DeviceDescriptor {
            features: wgpu::Features::TEXTURE_COMPRESSION_BC,
            limits: wgpu::Limits::default(),
            shader_validation: true,
        },
        None,
    ))
    .expect("Unable to create compatible wgpu device");

    let mut size = window.inner_size();
    let mut swap_chain = make_swapchain(&device, &surface, size.width, size.height);
    let mut depth_buffer = make_depth_buffer(&device, size.width, size.height);
    let mut proj = compute_projection_matrix(size.width as f32, size.height as f32);

    let mut gilrs = Gilrs::new().unwrap();
    let mut current_gamepad = None;
    for (_id, gamepad) in gilrs.gamepads() {
        current_gamepad = Some(gamepad.id());
    }

    let planet_radius = 6371000.0;
    let mut angle = 0.0f64;
    let mut lat = 41.3851f64.to_radians();
    let mut long = 2.1734f64.to_radians();
    let mut altitude = 200000.0f64;

    let mut terrain = terra::Terrain::new(&device, &mut queue, mapfile).unwrap();

    event_loop.run(move |event, _, control_flow| {
        *control_flow = if cfg!(feature = "metal-auto-capture") {
            ControlFlow::Exit
        } else {
            ControlFlow::Poll
        };
        match event {
            event::Event::WindowEvent { event, .. } => match event {
                event::WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit;
                }
                event::WindowEvent::KeyboardInput {
                    input:
                        event::KeyboardInput {
                            virtual_keycode: Some(keycode),
                            state: event::ElementState::Pressed,
                            ..
                        },
                    ..
                } => match keycode {
                    event::VirtualKeyCode::Escape => *control_flow = ControlFlow::Exit,
                    event::VirtualKeyCode::Space => altitude += 0.1 * altitude,
                    event::VirtualKeyCode::Semicolon => altitude -= 0.1 * altitude,
                    event::VirtualKeyCode::Left => {
                        lat += -angle.sin() * -(0.0000001 * altitude).min(0.01);
                        long += angle.cos() * -(0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Right => {
                        lat += -angle.sin() * (0.0000001 * altitude).min(0.01);
                        long += angle.cos() * (0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Up => {
                        lat += angle.cos() * (0.0000001 * altitude).min(0.01);
                        long += -angle.sin() * (0.0000001 * altitude).min(0.01);
                    }
                    event::VirtualKeyCode::Down => {
                        lat += angle.cos() * -(0.0000001 * altitude).min(0.01);
                        long += -angle.sin() * -(0.0000001 * altitude).min(0.01);
                    }
                    _ => {}
                },
                event::WindowEvent::Resized(new_size) => {
                    size = new_size;
                    swap_chain = make_swapchain(&device, &surface, size.width, size.height);
                    depth_buffer = make_depth_buffer(&device, size.width, size.height);
                    proj = compute_projection_matrix(size.width as f32, size.height as f32);
                }
                _ => {}
            },
            event::Event::MainEventsCleared => {
                let frame = swap_chain.get_current_frame();
                let frame = match frame {
                    Ok(ref f) => &f.output.view,
                    Err(_) => return,
                };

                while let Some(gilrs::Event { id, event: _event, time: _ }) = gilrs.next_event() {
                    current_gamepad = Some(id);
                }
                if let Some(gamepad) = current_gamepad.map(|id| gilrs.gamepad(id)) {
                    lat += angle.cos() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        - angle.sin() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    long += -angle.sin() * gamepad.value(Axis::LeftStickY) as f64 * 0.01
                        + angle.cos() * gamepad.value(Axis::LeftStickX) as f64 * 0.01;

                    angle -= gamepad.value(Axis::RightZ) as f64 * 0.01;

                    if gamepad.is_pressed(Button::DPadUp) {
                        altitude += 0.01 * altitude;
                    }
                    if gamepad.is_pressed(Button::DPadDown) {
                        altitude -= 0.01 * altitude;
                    }
                }

                lat = lat.max(-PI).min(PI);
                if long < -PI {
                    long += PI * 2.0;
                }
                if long > PI {
                    long -= PI * 2.0;
                }

                let r = altitude + planet_radius;
                let eye = cgmath::Point3::new(
                    r * lat.cos() * long.cos(),
                    r * lat.cos() * long.sin(),
                    r * lat.sin(),
                );

                let dt = (planet_radius / (planet_radius + altitude)).acos() * 0.3;
                let latc = lat + angle.cos() *  dt;
                let longc = long - angle.sin() * dt;

                let center = cgmath::Point3::new(
                    planet_radius * latc.cos() * longc.cos() - eye.x,
                    planet_radius * latc.cos() * longc.sin() - eye.y,
                    planet_radius * latc.sin() - eye.z,
                );
                let up = cgmath::Vector3::new(eye.x as f32, eye.y as f32, eye.z as f32);

                let view = cgmath::Matrix4::look_at(
                    cgmath::Point3::origin(),
                    cgmath::Point3::new(center.x as f32, center.y as f32, center.z as f32),
                    up,
                );

                let view_proj = proj * view;
                let view_proj = mint::ColumnMatrix4 {
                    x: view_proj.x.into(),
                    y: view_proj.y.into(),
                    z: view_proj.z.into(),
                    w: view_proj.w.into(),
                };

                terrain.render(
                    &device,
                    &mut queue,
                    &frame,
                    &depth_buffer,
                    (size.width, size.height),
                    view_proj,
                    eye.into(),
                );
            }
            _ => (),
        }
    });
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::TcpStream;
use std::time::Duration;
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Sender};
use std::thread::{self, JoinHandle};
use util::codec::rpc;
use util::make_std_tcp_conn;
use protobuf::{self, MessageStatic};
use super::Result;

pub trait TRpcClient: Sync + Send {
    fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic;
    fn post<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<()>;
}

#[derive(Debug)]
struct RpcClientCore {
    addrs: Vec<String>,
    // Try to connect pd with round-robin.
    next_index: usize,
    stream: Option<TcpStream>,
}

fn send_msg<M: protobuf::Message + ?Sized>(stream: &mut TcpStream,
                                           msg_id: u64,
                                           message: &M)
                                           -> Result<(u64, Vec<u8>)> {
    try!(stream.set_write_timeout(Some(Duration::from_secs(SOCKET_WRITE_TIMEOUT))));
    try!(rpc::encode_msg(stream, msg_id, message));

    try!(stream.set_read_timeout(Some(Duration::from_secs(SOCKET_READ_TIMEOUT))));
    let (id, data) = try!(rpc::decode_data(stream));
    Ok((id, data))
}

impl RpcClientCore {
    fn new(dsn: &str) -> RpcClientCore {
        let addrs: Vec<String> = dsn.split(',').map(From::from).collect();
        RpcClientCore {
            addrs: addrs,
            next_index: 0,
            stream: None,
        }
    }

    fn try_connect(&mut self) -> Result<()> {
        let index = self.next_index;
        self.next_index = (self.next_index + 1) % self.addrs.len();

        let addr = self.addrs.get(index).unwrap();
        let stream = try!(make_std_tcp_conn(&**addr));
        self.stream = Some(stream);
        Ok(())
    }

    fn post_msg<M: protobuf::Message + ?Sized>(&mut self,
                                               msg_id: u64,
                                               message: &M)
                                               -> Result<Vec<u8>> {
        // If we post failed, we should retry other addresses.
        let n = self.addrs.len();
        for _ in 0..n {
            // If no stream, try connect first.
            if self.stream.is_none() {
                if let Err(e) = self.try_connect() {
                    warn!("connect pd failed {:?}", e);
                    continue;
                }
            }

            let mut stream = self.stream.take().unwrap();
            // We may send message to a not leader pd, retry.

            let (id, data) = match send_msg(&mut stream, msg_id, message) {
                Err(e) => {
                    warn!("send message to pd failed {:?}", e);
                    continue;
                }
                Ok((id, data)) => (id, data),
            };

            if id != msg_id {
                return Err(box_err!("pd response msg_id not match, want {}, got {}", msg_id, id));
            }

            self.stream = Some(stream);
            return Ok(data);
        }

        Err(box_err!("post message to pd failed"))
    }
}

#[derive(Debug)]
pub struct RpcClient {
    core: Mutex<RpcClientCore>,
}

const SOCKET_READ_TIMEOUT: u64 = 3;
const SOCKET_WRITE_TIMEOUT: u64 = 3;

impl RpcClient {
    pub fn new(dsn: &str) -> Result<RpcClient> {
        Ok(RpcClient { core: Mutex::new(RpcClientCore::new(dsn)) })
    }
}

impl TRpcClient for RpcClient {
    fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic
    {
        let mut resp = P::new();
        let data = try!(self.core.lock().unwrap().post_msg(msg_id, message));
        try!(rpc::decode_body(&data, &mut resp));
        Ok(resp)
    }

    fn post<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<()> {
        let _ = try!(self.core.lock().unwrap().post_msg(msg_id, message));
        Ok(())
    }
}

#[derive(Debug)]
enum Msg {
    Rpc(u64, Box<protobuf::Message + Send>),
    Quit,
}

pub struct Client<C: TRpcClient + 'static> {
    msg_id: AtomicUsize,
    rpc_client: Arc<C>,
    tx: Mutex<Sender<Msg>>,
    handle: Option<JoinHandle<()>>,
}

impl<C: TRpcClient + 'static> Client<C> {
    pub fn new(rpc_client: C) -> Result<Client<C>> {
        let shared = Arc::new(rpc_client);
        let (tx, rx) = mpsc::channel();
        let handle = {
            let rpc_client = shared.clone();
            try!(thread::Builder::new()
                     .name("pd-client background worker".to_owned())
                     .spawn(move || {
                         info!("pd-client background worker starting");
                         while let Msg::Rpc(msg_id, msg) = rx.recv().unwrap() {
                             if let Err(e) = rpc_client.post(msg_id, msg.as_ref()) {
                                 error!("pd post error: {}", e);
                             } else {
                                 debug!("pd post done: {}", msg_id);
                             }
                         }
                         info!("pd-client background worker closing");
                     }))
        };
        Ok(Client {
            msg_id: AtomicUsize::new(0),
            rpc_client: shared.clone(),
            tx: Mutex::new(tx),
            handle: Some(handle),
        })
    }

    fn close(&mut self) {
        if let Err(e) = self.clone_tx().send(Msg::Quit) {
            error!("failed to send quit message to pd-client background worker: {}",
                   e);
        }
        let handle = self.handle.take().unwrap();
        if handle.join().is_err() {
            error!("failed to wait pd-client background worker quit");
        }
        info!("pd-client closed.");
    }

    pub fn send<M, P>(&self, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic
    {
        let id = self.alloc_msg_id();
        self.rpc_client.send(id, message)
    }

    pub fn post<M: protobuf::Message + Send>(&self, message: M) -> Result<()> {
        let id = self.alloc_msg_id();
        if let Err(e) = self.clone_tx().send(Msg::Rpc(id, box message)) {
            return Err(box_err!(format!("SendError: {:?}", e)));
        }
        Ok(())
    }

    fn alloc_msg_id(&self) -> u64 {
        self.msg_id.fetch_add(1, Ordering::Relaxed) as u64
    }

    fn clone_tx(&self) -> Sender<Msg> {
        let tx = self.tx.lock().unwrap();
        tx.clone()
    }
}

impl<C: TRpcClient + 'static> Drop for Client<C> {
    fn drop(&mut self) {
        self.close()
    }
}

#[cfg(test)]
mod tests {
    use std::net::TcpListener;
    use std::thread;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering};

    use super::super::{Result, PdClient};
    use super::*;
    use protobuf;
    use bytes::ByteBuf;
    use util::codec::rpc;
    use kvproto::{metapb, pdpb};
    use rand;
    use util::make_std_tcp_conn;

    struct MockRpcClient;

    impl MockRpcClient {
        fn reply(&self, req: &pdpb::Request) -> pdpb::Response {
            assert!(req.has_header());
            let header = req.get_header();
            let (uuid, cluster_id) = (header.get_uuid(), header.get_cluster_id());

            let mut resp = pdpb::Response::new();
            let mut header = pdpb::ResponseHeader::new();
            header.set_uuid(uuid.to_vec());
            header.set_cluster_id(cluster_id);

            match req.get_cmd_type() {
                pdpb::CommandType::AllocId => {
                    let mut alloc = pdpb::AllocIdResponse::new();
                    alloc.set_id(42);
                    resp.set_alloc_id(alloc);
                }
                pdpb::CommandType::PutMeta => {
                    let meta = req.get_put_meta();
                    if meta.get_meta_type() == pdpb::MetaType::StoreType {
                        let store = meta.get_store();
                        assert_eq!(store.get_id(), 233);
                        assert_eq!(store.get_address(), "localhost");
                    }
                }
                _ => {
                    let mut err = pdpb::Error::new();
                    err.set_message(String::from("I've got a bad feeling about this"));
                    header.set_error(err);
                }
            }

            resp.set_header(header);
            resp.set_cmd_type(req.get_cmd_type());
            resp
        }
    }

    impl TRpcClient for MockRpcClient {
        fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

            where M: protobuf::Message,
                  P: protobuf::Message + Default
        {
            let mut buf = ByteBuf::mut_with_capacity(128);
            rpc::encode_msg(&mut buf, msg_id, message).unwrap();
            let mut req = pdpb::Request::new();
            rpc::decode_msg(&mut buf.flip(), &mut req).unwrap();

            let resp = self.reply(&req);
            let mut buf = ByteBuf::mut_with_capacity(128);
            rpc::encode_msg(&mut buf, msg_id, &resp).unwrap();
            let mut msg = P::default();
            rpc::decode_msg(&mut buf.flip(), &mut msg).unwrap();
            Ok(msg)
        }

        fn post<M: protobuf::Message + ?Sized>(&self, _: u64, _: &M) -> Result<()> {
            Ok(())
        }
    }

    #[test]
    fn test_pd_client() {
        let mut client = Client::new(MockRpcClient).unwrap();
        assert_eq!(client.alloc_id().unwrap(), 42u64);
        assert!(client.is_cluster_bootstrapped(1).is_err());

        let mut store = metapb::Store::new();
        store.set_id(233);
        store.set_address("localhost".to_owned());
        assert!(client.put_store(1, store).is_ok());
    }

    fn start_pd_server(index: usize,
                       leader_index: Arc<AtomicUsize>,
                       quit: Arc<AtomicBool>)
                       -> (thread::JoinHandle<()>, String) {
        let l = TcpListener::bind("127.0.0.1:0").unwrap();

        let addr = format!("{}", l.local_addr().unwrap());

        let h = thread::spawn(move || {
            loop {
                let (mut stream, _) = l.accept().unwrap();

                if quit.load(Ordering::SeqCst) {
                    return;
                }

                stream.set_nodelay(true).unwrap();

                let leader = leader_index.load(Ordering::SeqCst);
                if leader != index {
                    continue;
                }

                let (id, data) = rpc::decode_data(&mut stream).unwrap();
                rpc::encode_data(&mut stream, id, &data).unwrap();
            }
        });

        (h, addr)
    }

    #[test]
    fn test_rpc_client() {
        let leader = Arc::new(AtomicUsize::new(0));
        let quit = Arc::new(AtomicBool::new(false));
        let mut handlers = vec![];
        let mut addrs = vec![];

        let count = 3;
        for i in 0..count {
            let (h, addr) = start_pd_server(i, leader.clone(), quit.clone());
            handlers.push(h);
            addrs.push(addr);
        }

        let dsn = addrs.join(",");

        let msg = pdpb::Request::new();
        let client = RpcClient::new(&dsn).unwrap();

        for i in 0..10 {
            if let Err(_) = client.post(i as u64, &msg) {
                // If failed, next retry must be OK.
                client.post(i as u64, &msg).unwrap();
            }

            // select a leader randomly
            leader.store(rand::random::<usize>() % 3, Ordering::SeqCst);
        }

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

        // connect the server so that we can close the thread.
        for addr in addrs {
            make_std_tcp_conn(&*addr).unwrap();
        }

        for h in handlers {
            h.join().unwrap();
        }
    }
}
<commit_msg>pd: use max retry for sending msg. (#381)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::TcpStream;
use std::time::Duration;
use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Sender};
use std::thread::{self, JoinHandle};
use util::codec::rpc;
use util::make_std_tcp_conn;
use protobuf::{self, MessageStatic};
use super::Result;

const MAX_PD_SEND_RETRY_COUNT: usize = 20;

pub trait TRpcClient: Sync + Send {
    fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic;
    fn post<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<()>;
}

#[derive(Debug)]
struct RpcClientCore {
    addrs: Vec<String>,
    // Try to connect pd with round-robin.
    next_index: usize,
    stream: Option<TcpStream>,
}

fn send_msg<M: protobuf::Message + ?Sized>(stream: &mut TcpStream,
                                           msg_id: u64,
                                           message: &M)
                                           -> Result<(u64, Vec<u8>)> {
    try!(stream.set_write_timeout(Some(Duration::from_secs(SOCKET_WRITE_TIMEOUT))));
    try!(rpc::encode_msg(stream, msg_id, message));

    try!(stream.set_read_timeout(Some(Duration::from_secs(SOCKET_READ_TIMEOUT))));
    let (id, data) = try!(rpc::decode_data(stream));
    Ok((id, data))
}

impl RpcClientCore {
    fn new(dsn: &str) -> RpcClientCore {
        let addrs: Vec<String> = dsn.split(',').map(From::from).collect();
        RpcClientCore {
            addrs: addrs,
            next_index: 0,
            stream: None,
        }
    }

    fn try_connect(&mut self) -> Result<()> {
        let index = self.next_index;
        self.next_index = (self.next_index + 1) % self.addrs.len();

        let addr = self.addrs.get(index).unwrap();
        let stream = try!(make_std_tcp_conn(&**addr));
        self.stream = Some(stream);
        Ok(())
    }

    fn post_msg<M: protobuf::Message + ?Sized>(&mut self,
                                               msg_id: u64,
                                               message: &M)
                                               -> Result<Vec<u8>> {
        // If we post failed, we should retry.
        for _ in 0..MAX_PD_SEND_RETRY_COUNT {
            // If no stream, try connect first.
            if self.stream.is_none() {
                if let Err(e) = self.try_connect() {
                    warn!("connect pd failed {:?}", e);
                    continue;
                }
            }

            let mut stream = self.stream.take().unwrap();
            // We may send message to a not leader pd, retry.

            let (id, data) = match send_msg(&mut stream, msg_id, message) {
                Err(e) => {
                    warn!("send message to pd failed {:?}", e);
                    continue;
                }
                Ok((id, data)) => (id, data),
            };

            if id != msg_id {
                return Err(box_err!("pd response msg_id not match, want {}, got {}", msg_id, id));
            }

            self.stream = Some(stream);
            return Ok(data);
        }

        Err(box_err!("post message to pd failed"))
    }
}

#[derive(Debug)]
pub struct RpcClient {
    core: Mutex<RpcClientCore>,
}

const SOCKET_READ_TIMEOUT: u64 = 3;
const SOCKET_WRITE_TIMEOUT: u64 = 3;

impl RpcClient {
    pub fn new(dsn: &str) -> Result<RpcClient> {
        Ok(RpcClient { core: Mutex::new(RpcClientCore::new(dsn)) })
    }
}

impl TRpcClient for RpcClient {
    fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic
    {
        let mut resp = P::new();
        let data = try!(self.core.lock().unwrap().post_msg(msg_id, message));
        try!(rpc::decode_body(&data, &mut resp));
        Ok(resp)
    }

    fn post<M: protobuf::Message + ?Sized>(&self, msg_id: u64, message: &M) -> Result<()> {
        let _ = try!(self.core.lock().unwrap().post_msg(msg_id, message));
        Ok(())
    }
}

#[derive(Debug)]
enum Msg {
    Rpc(u64, Box<protobuf::Message + Send>),
    Quit,
}

pub struct Client<C: TRpcClient + 'static> {
    msg_id: AtomicUsize,
    rpc_client: Arc<C>,
    tx: Mutex<Sender<Msg>>,
    handle: Option<JoinHandle<()>>,
}

impl<C: TRpcClient + 'static> Client<C> {
    pub fn new(rpc_client: C) -> Result<Client<C>> {
        let shared = Arc::new(rpc_client);
        let (tx, rx) = mpsc::channel();
        let handle = {
            let rpc_client = shared.clone();
            try!(thread::Builder::new()
                     .name("pd-client background worker".to_owned())
                     .spawn(move || {
                         info!("pd-client background worker starting");
                         while let Msg::Rpc(msg_id, msg) = rx.recv().unwrap() {
                             if let Err(e) = rpc_client.post(msg_id, msg.as_ref()) {
                                 error!("pd post error: {}", e);
                             } else {
                                 debug!("pd post done: {}", msg_id);
                             }
                         }
                         info!("pd-client background worker closing");
                     }))
        };
        Ok(Client {
            msg_id: AtomicUsize::new(0),
            rpc_client: shared.clone(),
            tx: Mutex::new(tx),
            handle: Some(handle),
        })
    }

    fn close(&mut self) {
        if let Err(e) = self.clone_tx().send(Msg::Quit) {
            error!("failed to send quit message to pd-client background worker: {}",
                   e);
        }
        let handle = self.handle.take().unwrap();
        if handle.join().is_err() {
            error!("failed to wait pd-client background worker quit");
        }
        info!("pd-client closed.");
    }

    pub fn send<M, P>(&self, message: &M) -> Result

        where M: protobuf::Message,
              P: protobuf::Message + MessageStatic
    {
        let id = self.alloc_msg_id();
        self.rpc_client.send(id, message)
    }

    pub fn post<M: protobuf::Message + Send>(&self, message: M) -> Result<()> {
        let id = self.alloc_msg_id();
        if let Err(e) = self.clone_tx().send(Msg::Rpc(id, box message)) {
            return Err(box_err!(format!("SendError: {:?}", e)));
        }
        Ok(())
    }

    fn alloc_msg_id(&self) -> u64 {
        self.msg_id.fetch_add(1, Ordering::Relaxed) as u64
    }

    fn clone_tx(&self) -> Sender<Msg> {
        let tx = self.tx.lock().unwrap();
        tx.clone()
    }
}

impl<C: TRpcClient + 'static> Drop for Client<C> {
    fn drop(&mut self) {
        self.close()
    }
}

#[cfg(test)]
mod tests {
    use std::net::TcpListener;
    use std::thread;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, AtomicBool, Ordering};

    use super::super::{Result, PdClient};
    use super::*;
    use protobuf;
    use bytes::ByteBuf;
    use util::codec::rpc;
    use kvproto::{metapb, pdpb};
    use rand;
    use util::make_std_tcp_conn;

    struct MockRpcClient;

    impl MockRpcClient {
        fn reply(&self, req: &pdpb::Request) -> pdpb::Response {
            assert!(req.has_header());
            let header = req.get_header();
            let (uuid, cluster_id) = (header.get_uuid(), header.get_cluster_id());

            let mut resp = pdpb::Response::new();
            let mut header = pdpb::ResponseHeader::new();
            header.set_uuid(uuid.to_vec());
            header.set_cluster_id(cluster_id);

            match req.get_cmd_type() {
                pdpb::CommandType::AllocId => {
                    let mut alloc = pdpb::AllocIdResponse::new();
                    alloc.set_id(42);
                    resp.set_alloc_id(alloc);
                }
                pdpb::CommandType::PutMeta => {
                    let meta = req.get_put_meta();
                    if meta.get_meta_type() == pdpb::MetaType::StoreType {
                        let store = meta.get_store();
                        assert_eq!(store.get_id(), 233);
                        assert_eq!(store.get_address(), "localhost");
                    }
                }
                _ => {
                    let mut err = pdpb::Error::new();
                    err.set_message(String::from("I've got a bad feeling about this"));
                    header.set_error(err);
                }
            }

            resp.set_header(header);
            resp.set_cmd_type(req.get_cmd_type());
            resp
        }
    }

    impl TRpcClient for MockRpcClient {
        fn send<M, P>(&self, msg_id: u64, message: &M) -> Result

            where M: protobuf::Message,
                  P: protobuf::Message + Default
        {
            let mut buf = ByteBuf::mut_with_capacity(128);
            rpc::encode_msg(&mut buf, msg_id, message).unwrap();
            let mut req = pdpb::Request::new();
            rpc::decode_msg(&mut buf.flip(), &mut req).unwrap();

            let resp = self.reply(&req);
            let mut buf = ByteBuf::mut_with_capacity(128);
            rpc::encode_msg(&mut buf, msg_id, &resp).unwrap();
            let mut msg = P::default();
            rpc::decode_msg(&mut buf.flip(), &mut msg).unwrap();
            Ok(msg)
        }

        fn post<M: protobuf::Message + ?Sized>(&self, _: u64, _: &M) -> Result<()> {
            Ok(())
        }
    }

    #[test]
    fn test_pd_client() {
        let mut client = Client::new(MockRpcClient).unwrap();
        assert_eq!(client.alloc_id().unwrap(), 42u64);
        assert!(client.is_cluster_bootstrapped(1).is_err());

        let mut store = metapb::Store::new();
        store.set_id(233);
        store.set_address("localhost".to_owned());
        assert!(client.put_store(1, store).is_ok());
    }

    fn start_pd_server(index: usize,
                       leader_index: Arc<AtomicUsize>,
                       quit: Arc<AtomicBool>)
                       -> (thread::JoinHandle<()>, String) {
        let l = TcpListener::bind("127.0.0.1:0").unwrap();

        let addr = format!("{}", l.local_addr().unwrap());

        let h = thread::spawn(move || {
            loop {
                let (mut stream, _) = l.accept().unwrap();

                if quit.load(Ordering::SeqCst) {
                    return;
                }

                stream.set_nodelay(true).unwrap();

                let leader = leader_index.load(Ordering::SeqCst);
                if leader != index {
                    continue;
                }

                let (id, data) = rpc::decode_data(&mut stream).unwrap();
                rpc::encode_data(&mut stream, id, &data).unwrap();
            }
        });

        (h, addr)
    }

    #[test]
    fn test_rpc_client() {
        let leader = Arc::new(AtomicUsize::new(0));
        let quit = Arc::new(AtomicBool::new(false));
        let mut handlers = vec![];
        let mut addrs = vec![];

        let count = 3;
        for i in 0..count {
            let (h, addr) = start_pd_server(i, leader.clone(), quit.clone());
            handlers.push(h);
            addrs.push(addr);
        }

        let dsn = addrs.join(",");

        let msg = pdpb::Request::new();
        let client = RpcClient::new(&dsn).unwrap();

        for i in 0..10 {
            client.post(i as u64, &msg).unwrap();

            // select a leader randomly
            leader.store(rand::random::<usize>() % 3, Ordering::SeqCst);
        }

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

        // connect the server so that we can close the thread.
        for addr in addrs {
            make_std_tcp_conn(&*addr).unwrap();
        }

        for h in handlers {
            h.join().unwrap();
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2013-2014 The gl-rs developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![experimental]

use registry::{Registry, Ns};
use syntax::ast;
use syntax::ext::base::ExtCtxt;
use syntax::ext::quote::rt::ExtParseUtils;
use syntax::ptr::P;

pub struct StructGenerator;

impl super::Generator for StructGenerator {
    fn write(&self, ecx: &ExtCtxt, registry: &Registry, ns: Ns) -> Vec<P<ast::Item>> {
        let mut result = Vec::new();
        result.push(write_header(ecx));
        result.push(write_type_aliases(ecx, &ns));
        result.extend(write_enums(ecx, registry).into_iter());
        result.extend(write_fnptr_struct_def(ecx).into_iter());
        result.push(write_failing_fns(ecx, registry));
        result.push(write_struct(ecx, registry, &ns));
        result.push(write_impl(ecx, registry, &ns));
        result
    }
}

/// Creates a `__gl_imports` module which contains all the external symbols that we need for the
///  bindings.
fn write_header(ecx: &ExtCtxt) -> P<ast::Item> {
    (quote_item!(ecx,
        mod __gl_imports {
            extern crate gl_common;
            extern crate libc;
            pub use std::mem;
        }
    )).unwrap()
}

/// Creates a `types` module which contains all the type aliases.
///
/// See also `generators::gen_type_aliases`.
fn write_type_aliases(ecx: &ExtCtxt, ns: &Ns) -> P<ast::Item> {
    let aliases = super::gen_type_aliases(ecx, ns);

    (quote_item!(ecx,
        #[stable]
        pub mod types {
            #![allow(non_camel_case_types)]
            #![allow(non_snake_case)]
            #![allow(dead_code)]

            $aliases
        }
    )).unwrap()
}

/// Writes all the `<enum>` elements at the root of the bindings.
fn write_enums(ecx: &ExtCtxt, registry: &Registry) -> Vec<P<ast::Item>> {
    registry.enum_iter().map(|e| {
        super::gen_enum_item(ecx, e, "types::")
    }).collect()
}

/// Creates a `FnPtr` structure which contains the store for a single binding.
fn write_fnptr_struct_def(ecx: &ExtCtxt) -> Vec<P<ast::Item>> {
    let mut result = Vec::new();

    result.push((quote_item!(ecx,
        #[allow(dead_code)]
        pub struct FnPtr {
            /// The function pointer that will be used when calling the function.
            f: *const __gl_imports::libc::c_void,
            /// True if the pointer points to a real function, false if points to a `panic!` fn.
            is_loaded: bool,
        }
    )).unwrap());

    result.push((quote_item!(ecx,
        impl FnPtr {
            /// Creates a `FnPtr` from a load attempt.
            fn new(ptr: *const __gl_imports::libc::c_void,
                failing_fn: *const __gl_imports::libc::c_void) -> FnPtr {
                if ptr.is_null() {
                    FnPtr { f: failing_fn, is_loaded: false }
                } else {
                    FnPtr { f: ptr, is_loaded: true }
                }
            }

            /// Returns `true` if the function has been successfully loaded.
            ///
            /// If it returns `false`, calling the corresponding function will fail.
            #[inline]
            #[allow(dead_code)]
            pub fn is_loaded(&self) -> bool {
                self.is_loaded
            }
        }
    )).unwrap());

    result
}

/// Creates a `failing` module which contains one function per GL command.
///
/// These functions are the mocks that are called if the real function could not be loaded.
fn write_failing_fns(ecx: &ExtCtxt, registry: &Registry) -> P<ast::Item> {
    let fns = registry.cmd_iter().map(|c| {
        use syntax::ext::quote::rt::ToSource;

        ecx.parse_item(format!(r#"
            #[allow(unused_variables)]
            #[allow(non_snake_case)]
            #[allow(dead_code)]
            pub extern "system" fn {name}({params}) -> {return_suffix} {{
                panic!("`{name}` was not loaded")
            }}
            "#,
            name = c.proto.ident,
            params = super::gen_parameters(ecx, c).into_iter().map(|p| p.to_source()).collect::<Vec<String>>().connect(", "),
            return_suffix = super::gen_return_type(ecx, c).to_source()
        ))
    }).collect::<Vec<P<ast::Item>>>();

    (quote_item!(ecx,
        mod failing {
            use super::types;
            use super::__gl_imports;

            $fns
        }
    )).unwrap()
}

/// Creates a structure which stores all the `FnPtr` of the bindings.
///
/// The name of the struct corresponds to the namespace.
fn write_struct(ecx: &ExtCtxt, registry: &Registry, ns: &Ns) -> P<ast::Item> {
    ecx.parse_item(format!("
        #[allow(non_camel_case_types)]
        #[allow(non_snake_case)]
        #[allow(dead_code)]
        #[stable]
        pub struct {ns} {{
            {ptrs}
        }}",

        ns = ns.fmt_struct_name(),
        ptrs = registry.cmd_iter().map(|c| {
            let fallbacks = match registry.aliases.get(&c.proto.ident) {
                Some(v) => v.clone(),
                None => Vec::new(),
            };
            format!(
                "{doc} pub {name}: FnPtr,",
                doc = if fallbacks.len() == 0 { "".to_string() } else { format!("/** Fallbacks: {} */", fallbacks.connect(", ")) },
                name = c.proto.ident,
            )
        }).collect::<Vec<String>>().connect("\n")
    ))
}

/// Creates the `impl` of the structure created by `write_struct`.
fn write_impl(ecx: &ExtCtxt, registry: &Registry, ns: &Ns) -> P<ast::Item> {
    ecx.parse_item(format!("
        impl {ns} {{
            /// Load each OpenGL symbol using a custom load function. This allows for the
            /// use of functions like `glfwGetProcAddress` or `SDL_GL_GetProcAddress`.
            ///
            /// ~~~ignore
            /// let gl = Gl::load_with(|s| glfw.get_proc_address(s));
            /// ~~~
            #[unstable]
            #[allow(dead_code)]
            #[allow(unused_variables)]
            pub fn load_with(loadfn: |symbol: &str| -> *const __gl_imports::libc::c_void) -> {ns} {{
                let metaloadfn: |&str, &[&str]| -> *const __gl_imports::libc::c_void = |symbol, symbols| {{
                    let mut ptr = loadfn(symbol);
                    if ptr.is_null() {{
                        for &sym in symbols.iter() {{
                            ptr = loadfn(sym);
                            if !ptr.is_null() {{ break; }}
                        }}
                    }}
                    ptr
                }};
                {ns} {{
                    {loadings}
                }}
            }}

            /// Load each OpenGL symbol using a custom load function.
            ///
            /// ~~~ignore
            /// let gl = Gl::load(&glfw);
            /// ~~~
            #[unstable]
            #[allow(dead_code)]
            #[allow(unused_variables)]
            pub fn load<T: __gl_imports::gl_common::GlFunctionsSource>(loader: &T) -> {ns} {{
                {ns}::load_with(|name| loader.get_proc_addr(name))
            }}

            {modules}
        }}",

        ns = ns.fmt_struct_name(),

        loadings = registry.cmd_iter().map(|c| {
            let fallbacks = registry.aliases.get(&c.proto.ident);
            format!(
                "{name}: FnPtr::new(metaloadfn(\"{symbol}\", {fb}), failing::{name} as *const __gl_imports::libc::c_void),",
                name = c.proto.ident,
                symbol = super::gen_symbol_name(ns, c.proto.ident.as_slice()),
                fb = match fallbacks {
                    Some(fallbacks) => format!("&[{}]", fallbacks.iter().map(|name| format!("\"{}\"", name)).collect::<Vec<_>>().connect(", ")),
                    None => format!("&[]")
                },
            )
        }).collect::<Vec<String>>().connect("\n"),

        modules = registry.cmd_iter().map(|c| {
            use syntax::ext::quote::rt::ToSource;

            format!(
                "#[allow(non_snake_case)] #[allow(unused_variables)] #[allow(dead_code)]
                #[inline] #[unstable] pub unsafe fn {name}(&self, {typed_params}) -> {return_suffix} {{ \
                    __gl_imports::mem::transmute::<_, extern \"system\" fn({typed_params}) -> {return_suffix}>\
                        (self.{name}.f)({idents}) \
                }}",
                name = c.proto.ident,
                typed_params = super::gen_parameters(ecx, c).into_iter().map(|p| p.to_source()).collect::<Vec<String>>().connect(", "),
                return_suffix = super::gen_return_type(ecx, c).to_source(),
                idents = super::gen_parameters(ecx, c).into_iter().map(|p| p.pat.to_source()).collect::<Vec<String>>().connect(", "),
            )
        }).collect::<Vec<String>>().connect("\n")
    ))
}
<commit_msg>Add appropriate prefix for aliases too.<commit_after>// Copyright 2013-2014 The gl-rs developers. For a full listing of the authors,
// refer to the AUTHORS file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![experimental]

use registry::{Registry, Ns};
use syntax::ast;
use syntax::ext::base::ExtCtxt;
use syntax::ext::quote::rt::ExtParseUtils;
use syntax::ptr::P;

pub struct StructGenerator;

impl super::Generator for StructGenerator {
    fn write(&self, ecx: &ExtCtxt, registry: &Registry, ns: Ns) -> Vec<P<ast::Item>> {
        let mut result = Vec::new();
        result.push(write_header(ecx));
        result.push(write_type_aliases(ecx, &ns));
        result.extend(write_enums(ecx, registry).into_iter());
        result.extend(write_fnptr_struct_def(ecx).into_iter());
        result.push(write_failing_fns(ecx, registry));
        result.push(write_struct(ecx, registry, &ns));
        result.push(write_impl(ecx, registry, &ns));
        result
    }
}

/// Creates a `__gl_imports` module which contains all the external symbols that we need for the
///  bindings.
fn write_header(ecx: &ExtCtxt) -> P<ast::Item> {
    (quote_item!(ecx,
        mod __gl_imports {
            extern crate gl_common;
            extern crate libc;
            pub use std::mem;
        }
    )).unwrap()
}

/// Creates a `types` module which contains all the type aliases.
///
/// See also `generators::gen_type_aliases`.
fn write_type_aliases(ecx: &ExtCtxt, ns: &Ns) -> P<ast::Item> {
    let aliases = super::gen_type_aliases(ecx, ns);

    (quote_item!(ecx,
        #[stable]
        pub mod types {
            #![allow(non_camel_case_types)]
            #![allow(non_snake_case)]
            #![allow(dead_code)]

            $aliases
        }
    )).unwrap()
}

/// Writes all the `<enum>` elements at the root of the bindings.
fn write_enums(ecx: &ExtCtxt, registry: &Registry) -> Vec<P<ast::Item>> {
    registry.enum_iter().map(|e| {
        super::gen_enum_item(ecx, e, "types::")
    }).collect()
}

/// Creates a `FnPtr` structure which contains the store for a single binding.
fn write_fnptr_struct_def(ecx: &ExtCtxt) -> Vec<P<ast::Item>> {
    let mut result = Vec::new();

    result.push((quote_item!(ecx,
        #[allow(dead_code)]
        pub struct FnPtr {
            /// The function pointer that will be used when calling the function.
            f: *const __gl_imports::libc::c_void,
            /// True if the pointer points to a real function, false if points to a `panic!` fn.
            is_loaded: bool,
        }
    )).unwrap());

    result.push((quote_item!(ecx,
        impl FnPtr {
            /// Creates a `FnPtr` from a load attempt.
            fn new(ptr: *const __gl_imports::libc::c_void,
                failing_fn: *const __gl_imports::libc::c_void) -> FnPtr {
                if ptr.is_null() {
                    FnPtr { f: failing_fn, is_loaded: false }
                } else {
                    FnPtr { f: ptr, is_loaded: true }
                }
            }

            /// Returns `true` if the function has been successfully loaded.
            ///
            /// If it returns `false`, calling the corresponding function will fail.
            #[inline]
            #[allow(dead_code)]
            pub fn is_loaded(&self) -> bool {
                self.is_loaded
            }
        }
    )).unwrap());

    result
}

/// Creates a `failing` module which contains one function per GL command.
///
/// These functions are the mocks that are called if the real function could not be loaded.
fn write_failing_fns(ecx: &ExtCtxt, registry: &Registry) -> P<ast::Item> {
    let fns = registry.cmd_iter().map(|c| {
        use syntax::ext::quote::rt::ToSource;

        ecx.parse_item(format!(r#"
            #[allow(unused_variables)]
            #[allow(non_snake_case)]
            #[allow(dead_code)]
            pub extern "system" fn {name}({params}) -> {return_suffix} {{
                panic!("`{name}` was not loaded")
            }}
            "#,
            name = c.proto.ident,
            params = super::gen_parameters(ecx, c).into_iter().map(|p| p.to_source()).collect::<Vec<String>>().connect(", "),
            return_suffix = super::gen_return_type(ecx, c).to_source()
        ))
    }).collect::<Vec<P<ast::Item>>>();

    (quote_item!(ecx,
        mod failing {
            use super::types;
            use super::__gl_imports;

            $fns
        }
    )).unwrap()
}

/// Creates a structure which stores all the `FnPtr` of the bindings.
///
/// The name of the struct corresponds to the namespace.
fn write_struct(ecx: &ExtCtxt, registry: &Registry, ns: &Ns) -> P<ast::Item> {
    ecx.parse_item(format!("
        #[allow(non_camel_case_types)]
        #[allow(non_snake_case)]
        #[allow(dead_code)]
        #[stable]
        pub struct {ns} {{
            {ptrs}
        }}",

        ns = ns.fmt_struct_name(),
        ptrs = registry.cmd_iter().map(|c| {
            let fallbacks = match registry.aliases.get(&c.proto.ident) {
                Some(v) => v.clone(),
                None => Vec::new(),
            };
            format!(
                "{doc} pub {name}: FnPtr,",
                doc = if fallbacks.len() == 0 { "".to_string() } else { format!("/** Fallbacks: {} */", fallbacks.connect(", ")) },
                name = c.proto.ident,
            )
        }).collect::<Vec<String>>().connect("\n")
    ))
}

/// Creates the `impl` of the structure created by `write_struct`.
fn write_impl(ecx: &ExtCtxt, registry: &Registry, ns: &Ns) -> P<ast::Item> {
    ecx.parse_item(format!("
        impl {ns} {{
            /// Load each OpenGL symbol using a custom load function. This allows for the
            /// use of functions like `glfwGetProcAddress` or `SDL_GL_GetProcAddress`.
            ///
            /// ~~~ignore
            /// let gl = Gl::load_with(|s| glfw.get_proc_address(s));
            /// ~~~
            #[unstable]
            #[allow(dead_code)]
            #[allow(unused_variables)]
            pub fn load_with(loadfn: |symbol: &str| -> *const __gl_imports::libc::c_void) -> {ns} {{
                let metaloadfn: |&str, &[&str]| -> *const __gl_imports::libc::c_void = |symbol, symbols| {{
                    let mut ptr = loadfn(symbol);
                    if ptr.is_null() {{
                        for &sym in symbols.iter() {{
                            ptr = loadfn(sym);
                            if !ptr.is_null() {{ break; }}
                        }}
                    }}
                    ptr
                }};
                {ns} {{
                    {loadings}
                }}
            }}

            /// Load each OpenGL symbol using a custom load function.
            ///
            /// ~~~ignore
            /// let gl = Gl::load(&glfw);
            /// ~~~
            #[unstable]
            #[allow(dead_code)]
            #[allow(unused_variables)]
            pub fn load<T: __gl_imports::gl_common::GlFunctionsSource>(loader: &T) -> {ns} {{
                {ns}::load_with(|name| loader.get_proc_addr(name))
            }}

            {modules}
        }}",

        ns = ns.fmt_struct_name(),

        loadings = registry.cmd_iter().map(|c| {
            let fallbacks = registry.aliases.get(&c.proto.ident);
            format!(
                "{name}: FnPtr::new(metaloadfn(\"{symbol}\", {fb}), failing::{name} as *const __gl_imports::libc::c_void),",
                name = c.proto.ident,
                symbol = super::gen_symbol_name(ns, c.proto.ident.as_slice()),
                fb = match fallbacks {
                    Some(fallbacks) => format!("&[{}]", fallbacks.iter().map(|name| format!("\"{}\"", super::gen_symbol_name(ns, name.as_slice()))).collect::<Vec<_>>().connect(", ")),
                    None => format!("&[]")
                },
            )
        }).collect::<Vec<String>>().connect("\n"),

        modules = registry.cmd_iter().map(|c| {
            use syntax::ext::quote::rt::ToSource;

            format!(
                "#[allow(non_snake_case)] #[allow(unused_variables)] #[allow(dead_code)]
                #[inline] #[unstable] pub unsafe fn {name}(&self, {typed_params}) -> {return_suffix} {{ \
                    __gl_imports::mem::transmute::<_, extern \"system\" fn({typed_params}) -> {return_suffix}>\
                        (self.{name}.f)({idents}) \
                }}",
                name = c.proto.ident,
                typed_params = super::gen_parameters(ecx, c).into_iter().map(|p| p.to_source()).collect::<Vec<String>>().connect(", "),
                return_suffix = super::gen_return_type(ecx, c).to_source(),
                idents = super::gen_parameters(ecx, c).into_iter().map(|p| p.pat.to_source()).collect::<Vec<String>>().connect(", "),
            )
        }).collect::<Vec<String>>().connect("\n")
    ))
}
<|endoftext|>