<commit_before>use crate::{file, full, gen};
use anyhow::Result;
use proc_macro2::{Ident, Span, TokenStream};
use quote::{format_ident, quote};
use syn::Index;
use syn_codegen::{Data, Definitions, Features, Node, Type};

const FOLD_SRC: &str = "../src/gen/fold.rs";

fn simple_visit(item: &str, name: &TokenStream) -> TokenStream {
    let ident = gen::under_name(item);
    let method = format_ident!("fold_{}", ident);
    quote! {
        f.#method(#name)
    }
}

fn visit(
    ty: &Type,
    features: &Features,
    defs: &Definitions,
    name: &TokenStream,
) -> Option<TokenStream> {
    match ty {
        Type::Box(t) => {
            let res = visit(t, features, defs, &quote!(*#name))?;
            Some(quote! {
                Box::new(#res)
            })
        }
        Type::Vec(t) => {
            let operand = quote!(it);
            let val = visit(t, features, defs, &operand)?;
            Some(quote! {
                FoldHelper::lift(#name, |it| { #val })
            })
        }
        Type::Punctuated(p) => {
            let operand = quote!(it);
            let val = visit(&p.element, features, defs, &operand)?;
            Some(quote! {
                FoldHelper::lift(#name, |it| { #val })
            })
        }
        Type::Option(t) => {
            let it = quote!(it);
            let val = visit(t, features, defs, &it)?;
            Some(quote! {
                (#name).map(|it| { #val })
            })
        }
        Type::Tuple(t) => {
            let mut code = TokenStream::new();
            for (i, elem) in t.iter().enumerate() {
                let i = Index::from(i);
                let it = quote!((#name).#i);
                let val = visit(elem, features, defs, &it).unwrap_or(it);
                code.extend(val);
                code.extend(quote!(,));
            }
            Some(quote! {
                (#code)
            })
        }
        Type::Token(t) => {
            let repr = &defs.tokens[t];
            let is_keyword = repr.chars().next().unwrap().is_alphabetic();
            let spans = if is_keyword {
                quote!(span)
            } else {
                quote!(spans)
            };
            let ty = if repr == "await" {
                quote!(crate::token::Await)
            } else {
                syn::parse_str(&format!("Token![{}]", repr)).unwrap()
            };
            Some(quote! {
                #ty(tokens_helper(f, &#name.#spans))
            })
        }
        Type::Group(t) => {
            let ty = Ident::new(t, Span::call_site());
            Some(quote! {
                #ty(tokens_helper(f, &#name.span))
            })
        }
        Type::Syn(t) => {
            fn requires_full(features: &Features) -> bool {
                features.any.contains("full") && features.any.len() == 1
            }
            let mut res = simple_visit(t, name);
            let target = defs.types.iter().find(|ty| ty.ident == *t).unwrap();
            if requires_full(&target.features) && !requires_full(features) {
                res = quote!(full!(#res));
            }
            Some(res)
        }
        Type::Ext(t) if gen::TERMINAL_TYPES.contains(&&t[..]) => Some(simple_visit(t, name)),
        Type::Ext(_) | Type::Std(_) => None,
    }
}

fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Definitions) {
    let under_name = gen::under_name(&s.ident);
    let ty = Ident::new(&s.ident, Span::call_site());
    let fold_fn = format_ident!("fold_{}", under_name);

    let mut fold_impl = TokenStream::new();

    match &s.data {
        Data::Enum(variants) => {
            let mut fold_variants = TokenStream::new();

            for (variant, fields) in variants {
                let variant_ident = Ident::new(variant, Span::call_site());

                if fields.is_empty() {
                    fold_variants.extend(quote! {
                        #ty::#variant_ident => {
                            #ty::#variant_ident
                        }
                    });
                } else {
                    let mut bind_fold_fields = TokenStream::new();
                    let mut fold_fields = TokenStream::new();

                    for (idx, ty) in fields.iter().enumerate() {
                        let binding = format_ident!("_binding_{}", idx);

                        bind_fold_fields.extend(quote! {
                            #binding,
                        });

                        let owned_binding = quote!(#binding);

                        fold_fields.extend(
                            visit(ty, &s.features, defs, &owned_binding).unwrap_or(owned_binding),
                        );

                        fold_fields.extend(quote!(,));
                    }

                    fold_variants.extend(quote! {
                        #ty::#variant_ident(#bind_fold_fields) => {
                            #ty::#variant_ident(
                                #fold_fields
                            )
                        }
                    });
                }
            }

            let nonexhaustive = if s.exhaustive {
                None
            } else {
                Some(quote!(_ => unreachable!()))
            };

            fold_impl.extend(quote! {
                match node {
                    #fold_variants
                    #nonexhaustive
                }
            });
        }
        Data::Struct(fields) => {
            let mut fold_fields = TokenStream::new();

            for (field, ty) in fields {
                let id = Ident::new(&field, Span::call_site());
                let ref_toks = quote!(node.#id);

                if let Type::Syn(ty) = ty {
                    if ty == "Reserved" {
                        fold_fields.extend(quote! {
                            #id: #ref_toks,
                        });
                        continue;
                    }
                }

                let fold = visit(&ty, &s.features, defs, &ref_toks).unwrap_or(ref_toks);

                fold_fields.extend(quote! {
                    #id: #fold,
                });
            }

            if !fields.is_empty() {
                fold_impl.extend(quote! {
                    #ty {
                        #fold_fields
                    }
                })
            } else {
                if ty == "Ident" {
                    fold_impl.extend(quote! {
                        let mut node = node;
                        let span = f.fold_span(node.span());
                        node.set_span(span);
                    });
                }
                fold_impl.extend(quote! {
                    node
                });
            }
        }
        Data::Private => {
            if ty == "Ident" {
                fold_impl.extend(quote! {
                    let mut node = node;
                    let span = f.fold_span(node.span());
                    node.set_span(span);
                });
            }
            fold_impl.extend(quote! {
                node
            });
        }
    }

    let fold_span_only =
        s.data == Data::Private && !gen::TERMINAL_TYPES.contains(&s.ident.as_str());
    if fold_span_only {
        fold_impl = quote! {
            let span = f.fold_span(node.span());
            let mut node = node;
            node.set_span(span);
            node
        };
    }

    traits.extend(quote! {
        fn #fold_fn(&mut self, i: #ty) -> #ty {
            #fold_fn(self, i)
        }
    });

    impls.extend(quote! {
        pub fn #fold_fn<F>(f: &mut F, node: #ty) -> #ty
        where
            F: Fold + ?Sized,
        {
            #fold_impl
        }
    });
}

pub fn generate(defs: &Definitions) -> Result<()> {
    let (traits, impls) = gen::traverse(defs, node);
    let full_macro = full::get_macro();
    file::write(
        FOLD_SRC,
        quote! {
            // Unreachable code is generated sometimes without the full feature.
            #![allow(unreachable_code, unused_variables)]
            #![allow(clippy::match_wildcard_for_single_variants)]

            use crate::*;
            #[cfg(any(feature = "full", feature = "derive"))]
            use crate::token::{Brace, Bracket, Paren, Group};
            use proc_macro2::Span;
            #[cfg(any(feature = "full", feature = "derive"))]
            use crate::gen::helper::fold::*;

            #full_macro

            /// Syntax tree traversal to transform the nodes of an owned syntax tree.
            ///
            /// See the [module documentation] for details.
            ///
            /// [module documentation]: self
            ///
            /// *This trait is available only if Syn is built with the `"fold"` feature.*
            pub trait Fold {
                #traits
            }

            #impls
        },
    )?;
    Ok(())
}
<commit_msg>Eliminate unneeded closure braces from fold impl<commit_after>use crate::{file, full, gen};
use anyhow::Result;
use proc_macro2::{Ident, Span, TokenStream};
use quote::{format_ident, quote};
use syn::Index;
use syn_codegen::{Data, Definitions, Features, Node, Type};

const FOLD_SRC: &str = "../src/gen/fold.rs";

fn simple_visit(item: &str, name: &TokenStream) -> TokenStream {
    let ident = gen::under_name(item);
    let method = format_ident!("fold_{}", ident);
    quote! {
        f.#method(#name)
    }
}

fn visit(
    ty: &Type,
    features: &Features,
    defs: &Definitions,
    name: &TokenStream,
) -> Option<TokenStream> {
    match ty {
        Type::Box(t) => {
            let res = visit(t, features, defs, &quote!(*#name))?;
            Some(quote! {
                Box::new(#res)
            })
        }
        Type::Vec(t) => {
            let operand = quote!(it);
            let val = visit(t, features, defs, &operand)?;
            Some(quote! {
                FoldHelper::lift(#name, |it| #val)
            })
        }
        Type::Punctuated(p) => {
            let operand = quote!(it);
            let val = visit(&p.element, features, defs, &operand)?;
            Some(quote! {
                FoldHelper::lift(#name, |it| #val)
            })
        }
        Type::Option(t) => {
            let it = quote!(it);
            let val = visit(t, features, defs, &it)?;
            Some(quote! {
                (#name).map(|it| #val)
            })
        }
        Type::Tuple(t) => {
            let mut code = TokenStream::new();
            for (i, elem) in t.iter().enumerate() {
                let i = Index::from(i);
                let it = quote!((#name).#i);
                let val = visit(elem, features, defs, &it).unwrap_or(it);
                code.extend(val);
                code.extend(quote!(,));
            }
            Some(quote! {
                (#code)
            })
        }
        Type::Token(t) => {
            let repr = &defs.tokens[t];
            let is_keyword = repr.chars().next().unwrap().is_alphabetic();
            let spans = if is_keyword {
                quote!(span)
            } else {
                quote!(spans)
            };
            let ty = if repr == "await" {
                quote!(crate::token::Await)
            } else {
                syn::parse_str(&format!("Token![{}]", repr)).unwrap()
            };
            Some(quote! {
                #ty(tokens_helper(f, &#name.#spans))
            })
        }
        Type::Group(t) => {
            let ty = Ident::new(t, Span::call_site());
            Some(quote! {
                #ty(tokens_helper(f, &#name.span))
            })
        }
        Type::Syn(t) => {
            fn requires_full(features: &Features) -> bool {
                features.any.contains("full") && features.any.len() == 1
            }
            let mut res = simple_visit(t, name);
            let target = defs.types.iter().find(|ty| ty.ident == *t).unwrap();
            if requires_full(&target.features) && !requires_full(features) {
                res = quote!(full!(#res));
            }
            Some(res)
        }
        Type::Ext(t) if gen::TERMINAL_TYPES.contains(&&t[..]) => Some(simple_visit(t, name)),
        Type::Ext(_) | Type::Std(_) => None,
    }
}

fn node(traits: &mut TokenStream, impls: &mut TokenStream, s: &Node, defs: &Definitions) {
    let under_name = gen::under_name(&s.ident);
    let ty = Ident::new(&s.ident, Span::call_site());
    let fold_fn = format_ident!("fold_{}", under_name);

    let mut fold_impl = TokenStream::new();

    match &s.data {
        Data::Enum(variants) => {
            let mut fold_variants = TokenStream::new();

            for (variant, fields) in variants {
                let variant_ident = Ident::new(variant, Span::call_site());

                if fields.is_empty() {
                    fold_variants.extend(quote! {
                        #ty::#variant_ident => {
                            #ty::#variant_ident
                        }
                    });
                } else {
                    let mut bind_fold_fields = TokenStream::new();
                    let mut fold_fields = TokenStream::new();

                    for (idx, ty) in fields.iter().enumerate() {
                        let binding = format_ident!("_binding_{}", idx);

                        bind_fold_fields.extend(quote! {
                            #binding,
                        });

                        let owned_binding = quote!(#binding);

                        fold_fields.extend(
                            visit(ty, &s.features, defs, &owned_binding).unwrap_or(owned_binding),
                        );

                        fold_fields.extend(quote!(,));
                    }

                    fold_variants.extend(quote! {
                        #ty::#variant_ident(#bind_fold_fields) => {
                            #ty::#variant_ident(
                                #fold_fields
                            )
                        }
                    });
                }
            }

            let nonexhaustive = if s.exhaustive {
                None
            } else {
                Some(quote!(_ => unreachable!()))
            };

            fold_impl.extend(quote! {
                match node {
                    #fold_variants
                    #nonexhaustive
                }
            });
        }
        Data::Struct(fields) => {
            let mut fold_fields = TokenStream::new();

            for (field, ty) in fields {
                let id = Ident::new(&field, Span::call_site());
                let ref_toks = quote!(node.#id);

                if let Type::Syn(ty) = ty {
                    if ty == "Reserved" {
                        fold_fields.extend(quote! {
                            #id: #ref_toks,
                        });
                        continue;
                    }
                }

                let fold = visit(&ty, &s.features, defs, &ref_toks).unwrap_or(ref_toks);

                fold_fields.extend(quote! {
                    #id: #fold,
                });
            }

            if !fields.is_empty() {
                fold_impl.extend(quote! {
                    #ty {
                        #fold_fields
                    }
                })
            } else {
                if ty == "Ident" {
                    fold_impl.extend(quote! {
                        let mut node = node;
                        let span = f.fold_span(node.span());
                        node.set_span(span);
                    });
                }
                fold_impl.extend(quote! {
                    node
                });
            }
        }
        Data::Private => {
            if ty == "Ident" {
                fold_impl.extend(quote! {
                    let mut node = node;
                    let span = f.fold_span(node.span());
                    node.set_span(span);
                });
            }
            fold_impl.extend(quote! {
                node
            });
        }
    }

    let fold_span_only =
        s.data == Data::Private && !gen::TERMINAL_TYPES.contains(&s.ident.as_str());
    if fold_span_only {
        fold_impl = quote! {
            let span = f.fold_span(node.span());
            let mut node = node;
            node.set_span(span);
            node
        };
    }

    traits.extend(quote! {
        fn #fold_fn(&mut self, i: #ty) -> #ty {
            #fold_fn(self, i)
        }
    });

    impls.extend(quote! {
        pub fn #fold_fn<F>(f: &mut F, node: #ty) -> #ty
        where
            F: Fold + ?Sized,
        {
            #fold_impl
        }
    });
}

pub fn generate(defs: &Definitions) -> Result<()> {
    let (traits, impls) = gen::traverse(defs, node);
    let full_macro = full::get_macro();
    file::write(
        FOLD_SRC,
        quote! {
            // Unreachable code is generated sometimes without the full feature.
            #![allow(unreachable_code, unused_variables)]
            #![allow(clippy::match_wildcard_for_single_variants)]

            use crate::*;
            #[cfg(any(feature = "full", feature = "derive"))]
            use crate::token::{Brace, Bracket, Paren, Group};
            use proc_macro2::Span;
            #[cfg(any(feature = "full", feature = "derive"))]
            use crate::gen::helper::fold::*;

            #full_macro

            /// Syntax tree traversal to transform the nodes of an owned syntax tree.
            ///
            /// See the [module documentation] for details.
            ///
            /// [module documentation]: self
            ///
            /// *This trait is available only if Syn is built with the `"fold"` feature.*
            pub trait Fold {
                #traits
            }

            #impls
        },
    )?;
    Ok(())
}
<|endoftext|>


<commit_before>extern crate argparse;
extern crate libc;
extern crate lithos;
extern crate quire;
extern crate regex;
extern crate serde_json;
extern crate unshare;
#[macro_use] extern crate log;

use std::env;
use std::str::FromStr;
use std::process::exit;
use std::path::{Path, PathBuf};
use std::io::{stderr, Write};
use std::collections::BTreeMap;

use argparse::{ArgumentParser, Parse, List, StoreTrue, StoreOption, Print};
use libc::getpid;
use quire::{parse_config, Options};
use regex::Regex;
use serde_json::to_string;
use unshare::{Command, Namespace};

use lithos::setup::{clean_child, init_logging};
use lithos::master_config::{MasterConfig, create_master_dirs};
use lithos::sandbox_config::SandboxConfig;
use lithos::child_config::{ChildConfig, ChildKind};

fn run(master_cfg: &Path, sandbox_name: String,
    command_name: String, args: Vec<String>,
    log_stderr: bool, log_level: Option<log::LogLevel>)
    -> Result<(), String>
{
    let master: MasterConfig = try!(parse_config(&master_cfg,
        &MasterConfig::validator(), &Options::default())
        .map_err(|e| format!("Error reading master config: {}", e)));
    try!(create_master_dirs(&master));

    if !Regex::new(r"^[\w-]+$").unwrap().is_match(&sandbox_name) {
        return Err(format!("Wrong sandbox name: {}", sandbox_name));
    }
    if !Regex::new(r"^[\w-]+$").unwrap().is_match(&command_name) {
        return Err(format!("Wrong command name: {}", command_name));
    }

    let sandbox: SandboxConfig = try!(parse_config(
        &master_cfg.parent().unwrap()
         .join(&master.sandboxes_dir).join(sandbox_name.clone() + ".yaml"),
        &SandboxConfig::validator(), &Options::default())
        .map_err(|e| format!("Error reading sandbox config: {}", e)));

    let log_file;
    if let Some(ref fname) = sandbox.log_file {
        log_file = master.default_log_dir.join(fname);
    } else {
        log_file = master.default_log_dir.join(format!("{}.log", sandbox_name));
    }
    try!(init_logging(&master, &log_file,
        &format!("{}-{}", master.syslog_app_name, sandbox_name),
        log_stderr,
        log_level
            .or(sandbox.log_level
                .and_then(|x| FromStr::from_str(&x).ok()))
            .or_else(|| FromStr::from_str(&master.log_level).ok())
            .unwrap_or(log::LogLevel::Warn)));

    let cfg = master_cfg.parent().unwrap()
        .join(&master.processes_dir)
        .join(sandbox.config_file.as_ref().unwrap_or(
            &PathBuf::from(&(sandbox_name.clone() + ".yaml"))));
    debug!("Children config {:?}", cfg);
    let sandbox_children: BTreeMap<String, ChildConfig>;
    sandbox_children = try!(parse_config(&cfg,
            &ChildConfig::mapping_validator(), &Options::default())
        .map_err(|e| format!("Error reading children config: {}", e)));
    let child_cfg = try!(sandbox_children.get(&command_name)
        .ok_or(format!("Command {:?} not found", command_name)));

    if child_cfg.kind != ChildKind::Command {
        return Err(format!("The target container is: {:?}", child_cfg.kind));
    }

    let child_cfg = child_cfg.instantiate(0)
        .map_err(|e| format!("can't instantiate: {}", e))?;

    let name = format!("{}/cmd.{}.{}", sandbox_name,
        command_name, unsafe { getpid() });

    let mut cmd = Command::new(env::current_exe().unwrap()
                     .parent().unwrap().join("lithos_knot"));

    // Name is first here, so it's easily visible in ps
    cmd.arg("--name");
    cmd.arg(&name);

    cmd.arg("--master");
    cmd.arg(master_cfg);
    cmd.arg("--config");
    cmd.arg(to_string(&child_cfg).unwrap());
    cmd.env_clear();
    cmd.env("TERM", env::var("TERM").unwrap_or("dumb".to_string()));
    if let Ok(x) = env::var("RUST_LOG") {
        cmd.env("RUST_LOG", x);
    }
    if let Ok(x) = env::var("RUST_BACKTRACE") {
        cmd.env("RUST_BACKTRACE", x);
    }
    cmd.arg("--");
    cmd.args(&args);
    cmd.unshare(&[Namespace::Mount, Namespace::Uts,
                 Namespace::Ipc, Namespace::Pid]);
    if sandbox.bridged_network.is_some() {
        cmd.unshare(&[Namespace::Net]);
    }

    info!("Running {:?}", cmd);

    let res = match cmd.status() {
        Ok(x) if x.success() => {
            info!("Command {:?} {}", cmd, x);
            Ok(())
        }
        Ok(x) => Err(format!("Command {:?} {}", cmd, x)),
        Err(e) => Err(format!("Can't run {:?}: {}", cmd, e)),
    };

    clean_child(&name, &master, false);

    return res;
}

fn main() {
    let mut master_config = PathBuf::from("/etc/lithos/master.yaml");
    let mut command_name = "".to_string();
    let mut sandbox_name = "".to_string();
    let mut args = vec!();
    let mut log_stderr: bool = false;
    let mut log_level: Option<log::LogLevel> = None;
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Runs single ad-hoc command");
        ap.refer(&mut master_config)
          .add_option(&["--master"], Parse,
            "Name of the master configuration file \
             (default /etc/lithos/master.yaml)")
          .metavar("FILE");
        ap.refer(&mut log_stderr)
          .add_option(&["--log-stderr"], StoreTrue,
            "Print debugging info to stderr");
        ap.refer(&mut log_level)
          .add_option(&["--log-level"], StoreOption,
            "Set log level (default info for now)");
        ap.refer(&mut sandbox_name)
          .add_argument("sandbox", Parse,
            "Name of the sandbox to run command for")
          .required();
        ap.refer(&mut command_name)
          .add_argument("name", Parse,
            "Name of the command to run")
          .required();
        ap.refer(&mut args)
          .add_argument("argument", List,
            "Arguments for the command");
        ap.add_option(&["--version"],
            Print(env!("CARGO_PKG_VERSION").to_string()),
            "Show version");
        ap.stop_on_first_argument(true);
        match ap.parse_args() {
            Ok(()) => {}
            Err(x) => {
                exit(x);
            }
        }
    }
    match run(&master_config, sandbox_name, command_name, args,
              log_stderr, log_level)
    {
        Ok(()) => {
            exit(0);
        }
        Err(e) => {
            write!(&mut stderr(), "Fatal error: {}\n", e).ok();
            error!("Fatal error: {}", e);
            exit(1);
        }
    }
}
<commit_msg>Fix lithos_cmd for new networking support<commit_after>extern crate argparse;
extern crate libc;
extern crate lithos;
extern crate quire;
extern crate regex;
extern crate serde_json;
extern crate unshare;
#[macro_use] extern crate log;

use std::env;
use std::str::FromStr;
use std::process::exit;
use std::path::{Path, PathBuf};
use std::io::{stderr, Write};
use std::collections::BTreeMap;

use argparse::{ArgumentParser, Parse, List, StoreTrue, StoreOption, Print};
use libc::getpid;
use quire::{parse_config, Options};
use regex::Regex;
use serde_json::to_string;
use unshare::{Command, Namespace};

use lithos::setup::{clean_child, init_logging};
use lithos::master_config::{MasterConfig, create_master_dirs};
use lithos::sandbox_config::SandboxConfig;
use lithos::child_config::{ChildConfig, ChildKind};

fn run(master_cfg: &Path, sandbox_name: String,
    command_name: String, args: Vec<String>,
    log_stderr: bool, log_level: Option<log::LogLevel>)
    -> Result<(), String>
{
    let master: MasterConfig = try!(parse_config(&master_cfg,
        &MasterConfig::validator(), &Options::default())
        .map_err(|e| format!("Error reading master config: {}", e)));
    try!(create_master_dirs(&master));

    if !Regex::new(r"^[\w-]+$").unwrap().is_match(&sandbox_name) {
        return Err(format!("Wrong sandbox name: {}", sandbox_name));
    }
    if !Regex::new(r"^[\w-]+$").unwrap().is_match(&command_name) {
        return Err(format!("Wrong command name: {}", command_name));
    }

    let sandbox: SandboxConfig = try!(parse_config(
        &master_cfg.parent().unwrap()
         .join(&master.sandboxes_dir).join(sandbox_name.clone() + ".yaml"),
        &SandboxConfig::validator(), &Options::default())
        .map_err(|e| format!("Error reading sandbox config: {}", e)));

    let log_file;
    if let Some(ref fname) = sandbox.log_file {
        log_file = master.default_log_dir.join(fname);
    } else {
        log_file = master.default_log_dir.join(format!("{}.log", sandbox_name));
    }
    try!(init_logging(&master, &log_file,
        &format!("{}-{}", master.syslog_app_name, sandbox_name),
        log_stderr,
        log_level
            .or(sandbox.log_level
                .and_then(|x| FromStr::from_str(&x).ok()))
            .or_else(|| FromStr::from_str(&master.log_level).ok())
            .unwrap_or(log::LogLevel::Warn)));

    let cfg = master_cfg.parent().unwrap()
        .join(&master.processes_dir)
        .join(sandbox.config_file.as_ref().unwrap_or(
            &PathBuf::from(&(sandbox_name.clone() + ".yaml"))));
    debug!("Children config {:?}", cfg);
    let sandbox_children: BTreeMap<String, ChildConfig>;
    sandbox_children = try!(parse_config(&cfg,
            &ChildConfig::mapping_validator(), &Options::default())
        .map_err(|e| format!("Error reading children config: {}", e)));
    let child_cfg = try!(sandbox_children.get(&command_name)
        .ok_or(format!("Command {:?} not found", command_name)));

    if child_cfg.kind != ChildKind::Command {
        return Err(format!("The target container is: {:?}", child_cfg.kind));
    }

    let child_cfg = child_cfg.instantiate(0)
        .map_err(|e| format!("can't instantiate: {}", e))?;

    let name = format!("{}/cmd.{}.{}", sandbox_name,
        command_name, unsafe { getpid() });

    let mut cmd = Command::new(env::current_exe().unwrap()
                     .parent().unwrap().join("lithos_knot"));

    // Name is first here, so it's easily visible in ps
    cmd.arg("--name");
    cmd.arg(&name);

    cmd.arg("--master");
    cmd.arg(master_cfg);
    cmd.arg("--config");
    cmd.arg(to_string(&child_cfg).unwrap());
    cmd.env_clear();
    cmd.env("TERM", env::var("TERM").unwrap_or("dumb".to_string()));
    if let Ok(x) = env::var("RUST_LOG") {
        cmd.env("RUST_LOG", x);
    }
    if let Ok(x) = env::var("RUST_BACKTRACE") {
        cmd.env("RUST_BACKTRACE", x);
    }
    cmd.arg("--");
    cmd.args(&args);
    cmd.unshare(&[Namespace::Mount, Namespace::Uts,
                 Namespace::Ipc, Namespace::Pid]);

    info!("Running {:?}", cmd);

    let res = match cmd.status() {
        Ok(x) if x.success() => {
            info!("Command {:?} {}", cmd, x);
            Ok(())
        }
        Ok(x) => Err(format!("Command {:?} {}", cmd, x)),
        Err(e) => Err(format!("Can't run {:?}: {}", cmd, e)),
    };

    clean_child(&name, &master, false);

    return res;
}

fn main() {
    let mut master_config = PathBuf::from("/etc/lithos/master.yaml");
    let mut command_name = "".to_string();
    let mut sandbox_name = "".to_string();
    let mut args = vec!();
    let mut log_stderr: bool = false;
    let mut log_level: Option<log::LogLevel> = None;
    {
        let mut ap = ArgumentParser::new();
        ap.set_description("Runs single ad-hoc command");
        ap.refer(&mut master_config)
          .add_option(&["--master"], Parse,
            "Name of the master configuration file \
             (default /etc/lithos/master.yaml)")
          .metavar("FILE");
        ap.refer(&mut log_stderr)
          .add_option(&["--log-stderr"], StoreTrue,
            "Print debugging info to stderr");
        ap.refer(&mut log_level)
          .add_option(&["--log-level"], StoreOption,
            "Set log level (default info for now)");
        ap.refer(&mut sandbox_name)
          .add_argument("sandbox", Parse,
            "Name of the sandbox to run command for")
          .required();
        ap.refer(&mut command_name)
          .add_argument("name", Parse,
            "Name of the command to run")
          .required();
        ap.refer(&mut args)
          .add_argument("argument", List,
            "Arguments for the command");
        ap.add_option(&["--version"],
            Print(env!("CARGO_PKG_VERSION").to_string()),
            "Show version");
        ap.stop_on_first_argument(true);
        match ap.parse_args() {
            Ok(()) => {}
            Err(x) => {
                exit(x);
            }
        }
    }
    match run(&master_config, sandbox_name, command_name, args,
              log_stderr, log_level)
    {
        Ok(()) => {
            exit(0);
        }
        Err(e) => {
            write!(&mut stderr(), "Fatal error: {}\n", e).ok();
            error!("Fatal error: {}", e);
            exit(1);
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Implements parallel traversal over the DOM tree.
//!
//! This traversal is based on Rayon, and therefore its safety is largely
//! verified by the type system.
//!
//! The primary trickiness and fine print for the above relates to the
//! thread safety of the DOM nodes themselves. Accessing a DOM element
//! concurrently on multiple threads is actually mostly "safe", since all
//! the mutable state is protected by an AtomicRefCell, and so we'll
//! generally panic if something goes wrong. Still, we try to to enforce our
//! thread invariants at compile time whenever possible. As such, TNode and
//! TElement are not Send, so ordinary style system code cannot accidentally
//! share them with other threads. In the parallel traversal, we explicitly
//! invoke |unsafe { SendNode::new(n) }| to put nodes in containers that may
//! be sent to other threads. This occurs in only a handful of places and is
//! easy to grep for. At the time of this writing, there is no other unsafe
//! code in the parallel traversal.

#![deny(missing_docs)]

use arrayvec::ArrayVec;
use context::{StyleContext, ThreadLocalStyleContext};
use dom::{OpaqueNode, SendNode, TElement};
use itertools::Itertools;
use rayon;
use scoped_tls::ScopedTLS;
use smallvec::SmallVec;
use traversal::{DomTraversal, PerLevelTraversalData};

/// The minimum stack size for a thread in the styling pool, in kilobytes.
pub const STYLE_THREAD_STACK_SIZE_KB: usize = 128;

/// The stack margin. If we get this deep in the stack, we will skip recursive
/// optimizations to ensure that there is sufficient room for non-recursive work.
///
/// When measured with 128KB stacks and 40KB margin, we could support 53
/// levels of recursion before the limiter kicks in, on x86_64-Linux [1].
///
/// We currently use 45KB margins, because that seems to be the minimum amount
/// of head room required by an unoptimized debug build, as measured on [2]. We
/// could probably get away with a smaller margin on optimized release builds,
/// but that would be a pain, and the extra padding reduces stability risk.
///
/// [1] See Gecko bug 1376883 for more discussion on the measurements.
/// [2] layout/style/crashtests/1383981-3.html
///
pub const STACK_SAFETY_MARGIN_KB: usize = 45;

/// The maximum number of child nodes that we will process as a single unit.
///
/// Larger values will increase style sharing cache hits and general DOM
/// locality at the expense of decreased opportunities for parallelism.  There
/// are some measurements in
/// https://bugzilla.mozilla.org/show_bug.cgi?id=1385982#c11 and comments 12
/// and 13 that investigate some slightly different values for the work unit
/// size.  If the size is significantly increased, make sure to adjust the
/// condition for kicking off a new work unit in top_down_dom, because
/// otherwise we're likely to end up doing too much work serially.  For
/// example, the condition there could become some fraction of WORK_UNIT_MAX
/// instead of WORK_UNIT_MAX.
pub const WORK_UNIT_MAX: usize = 16;

/// A set of nodes, sized to the work unit. This gets copied when sent to other
/// threads, so we keep it compact.
type WorkUnit<N> = ArrayVec<[SendNode<N>; WORK_UNIT_MAX]>;

/// A callback to create our thread local context.  This needs to be
/// out of line so we don't allocate stack space for the entire struct
/// in the caller.
#[inline(never)]
fn create_thread_local_context<'scope, E, D>(
    traversal: &'scope D,
    slot: &mut Option<ThreadLocalStyleContext<E>>)
    where E: TElement + 'scope,
          D: DomTraversal<E>
{
    *slot = Some(ThreadLocalStyleContext::new(traversal.shared_context()));
}

/// A parallel top-down DOM traversal.
///
/// This algorithm traverses the DOM in a breadth-first, top-down manner. The
/// goals are:
/// * Never process a child before its parent (since child style depends on
///   parent style). If this were to happen, the styling algorithm would panic.
/// * Prioritize discovering nodes as quickly as possible to maximize
///   opportunities for parallelism.  But this needs to be weighed against
///   styling cousins on a single thread to improve sharing.
/// * Style all the children of a given node (i.e. all sibling nodes) on
///   a single thread (with an upper bound to handle nodes with an
///   abnormally large number of children). This is important because we use
///   a thread-local cache to share styles between siblings.
#[inline(always)]
#[allow(unsafe_code)]
fn top_down_dom<'a, 'scope, E, D>(nodes: &'a [SendNode<E::ConcreteNode>],
                                  root: OpaqueNode,
                                  mut traversal_data: PerLevelTraversalData,
                                  scope: &'a rayon::Scope<'scope>,
                                  pool: &'scope rayon::ThreadPool,
                                  traversal: &'scope D,
                                  tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>)
    where E: TElement + 'scope,
          D: DomTraversal<E>,
{
    debug_assert!(nodes.len() <= WORK_UNIT_MAX);

    // We set this below, when we have a borrow of the thread-local-context
    // available.
    let recursion_ok;

    // Collect all the children of the elements in our work unit. This will
    // contain the combined children of up to WORK_UNIT_MAX nodes, which may
    // be numerous. As such, we store it in a large SmallVec to minimize heap-
    // spilling, and never move it.
    let mut discovered_child_nodes = SmallVec::<[SendNode<E::ConcreteNode>; 128]>::new();
    {
        // Scope the borrow of the TLS so that the borrow is dropped before
        // a potential recursive call when we pass TailCall.
        let mut tlc = tls.ensure(
            |slot: &mut Option<ThreadLocalStyleContext<E>>| create_thread_local_context(traversal, slot));

        // Check that we're not in danger of running out of stack.
        recursion_ok = !tlc.stack_limit_checker.limit_exceeded();

        let mut context = StyleContext {
            shared: traversal.shared_context(),
            thread_local: &mut *tlc,
        };

        for n in nodes {
            // If the last node we processed produced children, we may want to
            // spawn them off into a work item. We do this at the beginning of
            // the loop (rather than at the end) so that we can traverse our
            // last bits of work directly on this thread without a spawn call.
            //
            // This has the important effect of removing the allocation and
            // context-switching overhead of the parallel traversal for perfectly
            // linear regions of the DOM, i.e.:
            //
            // <russian><doll><tag><nesting></nesting></tag></doll></russian>
            //
            // which are not at all uncommon.
            //
            // There's a tension here between spawning off a work item as soon
            // as discovered_child_nodes is nonempty and waiting until we have a
            // full work item to do so.  The former optimizes for speed of
            // discovery (we'll start discovering the kids of the things in
            // "nodes" ASAP).  The latter gives us better sharing (e.g. we can
            // share between cousins much better, because we don't hand them off
            // as separate work items, which are likely to end up on separate
            // threads) and gives us a chance to just handle everything on this
            // thread for small DOM subtrees, as in the linear example above.
            //
            // There are performance and "number of ComputedValues"
            // measurements for various testcases in
            // https://bugzilla.mozilla.org/show_bug.cgi?id=1385982#c10 and
            // following.
            //
            // The worst case behavior for waiting until we have a full work
            // item is a deep tree which has WORK_UNIT_MAX "linear" branches,
            // hence WORK_UNIT_MAX elements at each level.  Such a tree would
            // end up getting processed entirely sequentially, because we would
            // process each level one at a time as a single work unit, whether
            // via our end-of-loop tail call or not.  If we kicked off a
            // traversal as soon as we discovered kids, we would instead
            // process such a tree more or less with a thread-per-branch,
            // multiplexed across our actual threadpool.
            if discovered_child_nodes.len() >= WORK_UNIT_MAX {
                let mut traversal_data_copy = traversal_data.clone();
                traversal_data_copy.current_dom_depth += 1;
                traverse_nodes(discovered_child_nodes.drain(),
                               DispatchMode::NotTailCall,
                               recursion_ok,
                               root,
                               traversal_data_copy,
                               scope,
                               pool,
                               traversal,
                               tls);
            }

            let node = **n;
            let mut children_to_process = 0isize;
            traversal.process_preorder(&traversal_data, &mut context, node, |n| {
                children_to_process += 1;
                let send_n = unsafe { SendNode::new(n) };
                discovered_child_nodes.push(send_n);
            });

            traversal.handle_postorder_traversal(&mut context, root, node,
                                                 children_to_process);
        }
    }

    // Handle whatever elements we have queued up but not kicked off traversals
    // for yet.  If any exist, we can process them (or at least one work unit's
    // worth of them) directly on this thread by passing TailCall.
    if !discovered_child_nodes.is_empty() {
        traversal_data.current_dom_depth += 1;
        traverse_nodes(discovered_child_nodes.drain(),
                       DispatchMode::TailCall,
                       recursion_ok,
                       root,
                       traversal_data,
                       scope,
                       pool,
                       traversal,
                       tls);
    }
}

/// Controls whether traverse_nodes may make a recursive call to continue
/// doing work, or whether it should always dispatch work asynchronously.
#[derive(Clone, Copy, PartialEq)]
pub enum DispatchMode {
    /// This is the last operation by the caller.
    TailCall,
    /// This is not the last operation by the caller.
    NotTailCall,
}

impl DispatchMode {
    fn is_tail_call(&self) -> bool { matches!(*self, DispatchMode::TailCall) }
}

/// Enqueues |nodes| for processing, possibly on this thread if the tail call
/// conditions are met.
#[inline]
pub fn traverse_nodes<'a, 'scope, E, D, I>(
    nodes: I,
    mode: DispatchMode,
    recursion_ok: bool,
    root: OpaqueNode,
    traversal_data: PerLevelTraversalData,
    scope: &'a rayon::Scope<'scope>,
    pool: &'scope rayon::ThreadPool,
    traversal: &'scope D,
    tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>
)
where
    E: TElement + 'scope,
    D: DomTraversal<E>,
    I: ExactSizeIterator<Item = SendNode<E::ConcreteNode>>
{
    debug_assert_ne!(nodes.len(), 0);

    // This is a tail call from the perspective of the caller. However, we only
    // want to actually dispatch the job as a tail call if there's nothing left
    // in our local queue. Otherwise we need to return to it to maintain proper
    // breadth-first ordering. We also need to take care to avoid stack
    // overflow due to excessive tail recursion. The stack overflow avoidance
    // isn't observable to content -- we're still completely correct, just not
    // using tail recursion any more. See Gecko bugs 1368302 and 1376883.
    let may_dispatch_tail = mode.is_tail_call() &&
        recursion_ok &&
        !pool.current_thread_has_pending_tasks().unwrap();

    // In the common case, our children fit within a single work unit, in which
    // case we can pass the SmallVec directly and avoid extra allocation.
    if nodes.len() <= WORK_UNIT_MAX {
        let work: WorkUnit<E::ConcreteNode> = nodes.collect();
        if may_dispatch_tail {
            top_down_dom(&work, root,
                         traversal_data, scope, pool, traversal, tls);
        } else {
            scope.spawn(move |scope| {
                let work = work;
                top_down_dom(&work, root,
                             traversal_data, scope, pool, traversal, tls);
            });
        }
    } else {
        for chunk in nodes.chunks(WORK_UNIT_MAX).into_iter() {
            let nodes: WorkUnit<E::ConcreteNode> = chunk.collect();
            let traversal_data_copy = traversal_data.clone();
            scope.spawn(move |scope| {
                let n = nodes;
                top_down_dom(&*n, root,
                             traversal_data_copy, scope, pool, traversal, tls)
            });
        }
    }
}
<commit_msg>Auto merge of #18412 - bholley:bump_style_stacks, r=bholley<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Implements parallel traversal over the DOM tree.
//!
//! This traversal is based on Rayon, and therefore its safety is largely
//! verified by the type system.
//!
//! The primary trickiness and fine print for the above relates to the
//! thread safety of the DOM nodes themselves. Accessing a DOM element
//! concurrently on multiple threads is actually mostly "safe", since all
//! the mutable state is protected by an AtomicRefCell, and so we'll
//! generally panic if something goes wrong. Still, we try to to enforce our
//! thread invariants at compile time whenever possible. As such, TNode and
//! TElement are not Send, so ordinary style system code cannot accidentally
//! share them with other threads. In the parallel traversal, we explicitly
//! invoke |unsafe { SendNode::new(n) }| to put nodes in containers that may
//! be sent to other threads. This occurs in only a handful of places and is
//! easy to grep for. At the time of this writing, there is no other unsafe
//! code in the parallel traversal.

#![deny(missing_docs)]

use arrayvec::ArrayVec;
use context::{StyleContext, ThreadLocalStyleContext};
use dom::{OpaqueNode, SendNode, TElement};
use itertools::Itertools;
use rayon;
use scoped_tls::ScopedTLS;
use smallvec::SmallVec;
use traversal::{DomTraversal, PerLevelTraversalData};

/// The minimum stack size for a thread in the styling pool, in kilobytes.
pub const STYLE_THREAD_STACK_SIZE_KB: usize = 256;

/// The stack margin. If we get this deep in the stack, we will skip recursive
/// optimizations to ensure that there is sufficient room for non-recursive work.
///
/// We allocate large safety margins because certain OS calls can use very large
/// amounts of stack space [1]. Reserving a larger-than-necessary stack costs us
/// address space, but if we keep our safety margin big, we will generally avoid
/// committing those extra pages, and only use them in edge cases that would
/// otherwise cause crashes.
///
/// When measured with 128KB stacks and 40KB margin, we could support 53
/// levels of recursion before the limiter kicks in, on x86_64-Linux [2]. When
/// we doubled the stack size, we added it all to the safety margin, so we should
/// be able to get the same amount of recursion.
///
/// [1] https://bugzilla.mozilla.org/show_bug.cgi?id=1395708#c15
/// [2] See Gecko bug 1376883 for more discussion on the measurements.
///
pub const STACK_SAFETY_MARGIN_KB: usize = 168;

/// The maximum number of child nodes that we will process as a single unit.
///
/// Larger values will increase style sharing cache hits and general DOM
/// locality at the expense of decreased opportunities for parallelism.  There
/// are some measurements in
/// https://bugzilla.mozilla.org/show_bug.cgi?id=1385982#c11 and comments 12
/// and 13 that investigate some slightly different values for the work unit
/// size.  If the size is significantly increased, make sure to adjust the
/// condition for kicking off a new work unit in top_down_dom, because
/// otherwise we're likely to end up doing too much work serially.  For
/// example, the condition there could become some fraction of WORK_UNIT_MAX
/// instead of WORK_UNIT_MAX.
pub const WORK_UNIT_MAX: usize = 16;

/// A set of nodes, sized to the work unit. This gets copied when sent to other
/// threads, so we keep it compact.
type WorkUnit<N> = ArrayVec<[SendNode<N>; WORK_UNIT_MAX]>;

/// A callback to create our thread local context.  This needs to be
/// out of line so we don't allocate stack space for the entire struct
/// in the caller.
#[inline(never)]
fn create_thread_local_context<'scope, E, D>(
    traversal: &'scope D,
    slot: &mut Option<ThreadLocalStyleContext<E>>)
    where E: TElement + 'scope,
          D: DomTraversal<E>
{
    *slot = Some(ThreadLocalStyleContext::new(traversal.shared_context()));
}

/// A parallel top-down DOM traversal.
///
/// This algorithm traverses the DOM in a breadth-first, top-down manner. The
/// goals are:
/// * Never process a child before its parent (since child style depends on
///   parent style). If this were to happen, the styling algorithm would panic.
/// * Prioritize discovering nodes as quickly as possible to maximize
///   opportunities for parallelism.  But this needs to be weighed against
///   styling cousins on a single thread to improve sharing.
/// * Style all the children of a given node (i.e. all sibling nodes) on
///   a single thread (with an upper bound to handle nodes with an
///   abnormally large number of children). This is important because we use
///   a thread-local cache to share styles between siblings.
#[inline(always)]
#[allow(unsafe_code)]
fn top_down_dom<'a, 'scope, E, D>(nodes: &'a [SendNode<E::ConcreteNode>],
                                  root: OpaqueNode,
                                  mut traversal_data: PerLevelTraversalData,
                                  scope: &'a rayon::Scope<'scope>,
                                  pool: &'scope rayon::ThreadPool,
                                  traversal: &'scope D,
                                  tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>)
    where E: TElement + 'scope,
          D: DomTraversal<E>,
{
    debug_assert!(nodes.len() <= WORK_UNIT_MAX);

    // We set this below, when we have a borrow of the thread-local-context
    // available.
    let recursion_ok;

    // Collect all the children of the elements in our work unit. This will
    // contain the combined children of up to WORK_UNIT_MAX nodes, which may
    // be numerous. As such, we store it in a large SmallVec to minimize heap-
    // spilling, and never move it.
    let mut discovered_child_nodes = SmallVec::<[SendNode<E::ConcreteNode>; 128]>::new();
    {
        // Scope the borrow of the TLS so that the borrow is dropped before
        // a potential recursive call when we pass TailCall.
        let mut tlc = tls.ensure(
            |slot: &mut Option<ThreadLocalStyleContext<E>>| create_thread_local_context(traversal, slot));

        // Check that we're not in danger of running out of stack.
        recursion_ok = !tlc.stack_limit_checker.limit_exceeded();

        let mut context = StyleContext {
            shared: traversal.shared_context(),
            thread_local: &mut *tlc,
        };

        for n in nodes {
            // If the last node we processed produced children, we may want to
            // spawn them off into a work item. We do this at the beginning of
            // the loop (rather than at the end) so that we can traverse our
            // last bits of work directly on this thread without a spawn call.
            //
            // This has the important effect of removing the allocation and
            // context-switching overhead of the parallel traversal for perfectly
            // linear regions of the DOM, i.e.:
            //
            // <russian><doll><tag><nesting></nesting></tag></doll></russian>
            //
            // which are not at all uncommon.
            //
            // There's a tension here between spawning off a work item as soon
            // as discovered_child_nodes is nonempty and waiting until we have a
            // full work item to do so.  The former optimizes for speed of
            // discovery (we'll start discovering the kids of the things in
            // "nodes" ASAP).  The latter gives us better sharing (e.g. we can
            // share between cousins much better, because we don't hand them off
            // as separate work items, which are likely to end up on separate
            // threads) and gives us a chance to just handle everything on this
            // thread for small DOM subtrees, as in the linear example above.
            //
            // There are performance and "number of ComputedValues"
            // measurements for various testcases in
            // https://bugzilla.mozilla.org/show_bug.cgi?id=1385982#c10 and
            // following.
            //
            // The worst case behavior for waiting until we have a full work
            // item is a deep tree which has WORK_UNIT_MAX "linear" branches,
            // hence WORK_UNIT_MAX elements at each level.  Such a tree would
            // end up getting processed entirely sequentially, because we would
            // process each level one at a time as a single work unit, whether
            // via our end-of-loop tail call or not.  If we kicked off a
            // traversal as soon as we discovered kids, we would instead
            // process such a tree more or less with a thread-per-branch,
            // multiplexed across our actual threadpool.
            if discovered_child_nodes.len() >= WORK_UNIT_MAX {
                let mut traversal_data_copy = traversal_data.clone();
                traversal_data_copy.current_dom_depth += 1;
                traverse_nodes(discovered_child_nodes.drain(),
                               DispatchMode::NotTailCall,
                               recursion_ok,
                               root,
                               traversal_data_copy,
                               scope,
                               pool,
                               traversal,
                               tls);
            }

            let node = **n;
            let mut children_to_process = 0isize;
            traversal.process_preorder(&traversal_data, &mut context, node, |n| {
                children_to_process += 1;
                let send_n = unsafe { SendNode::new(n) };
                discovered_child_nodes.push(send_n);
            });

            traversal.handle_postorder_traversal(&mut context, root, node,
                                                 children_to_process);
        }
    }

    // Handle whatever elements we have queued up but not kicked off traversals
    // for yet.  If any exist, we can process them (or at least one work unit's
    // worth of them) directly on this thread by passing TailCall.
    if !discovered_child_nodes.is_empty() {
        traversal_data.current_dom_depth += 1;
        traverse_nodes(discovered_child_nodes.drain(),
                       DispatchMode::TailCall,
                       recursion_ok,
                       root,
                       traversal_data,
                       scope,
                       pool,
                       traversal,
                       tls);
    }
}

/// Controls whether traverse_nodes may make a recursive call to continue
/// doing work, or whether it should always dispatch work asynchronously.
#[derive(Clone, Copy, PartialEq)]
pub enum DispatchMode {
    /// This is the last operation by the caller.
    TailCall,
    /// This is not the last operation by the caller.
    NotTailCall,
}

impl DispatchMode {
    fn is_tail_call(&self) -> bool { matches!(*self, DispatchMode::TailCall) }
}

/// Enqueues |nodes| for processing, possibly on this thread if the tail call
/// conditions are met.
#[inline]
pub fn traverse_nodes<'a, 'scope, E, D, I>(
    nodes: I,
    mode: DispatchMode,
    recursion_ok: bool,
    root: OpaqueNode,
    traversal_data: PerLevelTraversalData,
    scope: &'a rayon::Scope<'scope>,
    pool: &'scope rayon::ThreadPool,
    traversal: &'scope D,
    tls: &'scope ScopedTLS<'scope, ThreadLocalStyleContext<E>>
)
where
    E: TElement + 'scope,
    D: DomTraversal<E>,
    I: ExactSizeIterator<Item = SendNode<E::ConcreteNode>>
{
    debug_assert_ne!(nodes.len(), 0);

    // This is a tail call from the perspective of the caller. However, we only
    // want to actually dispatch the job as a tail call if there's nothing left
    // in our local queue. Otherwise we need to return to it to maintain proper
    // breadth-first ordering. We also need to take care to avoid stack
    // overflow due to excessive tail recursion. The stack overflow avoidance
    // isn't observable to content -- we're still completely correct, just not
    // using tail recursion any more. See Gecko bugs 1368302 and 1376883.
    let may_dispatch_tail = mode.is_tail_call() &&
        recursion_ok &&
        !pool.current_thread_has_pending_tasks().unwrap();

    // In the common case, our children fit within a single work unit, in which
    // case we can pass the SmallVec directly and avoid extra allocation.
    if nodes.len() <= WORK_UNIT_MAX {
        let work: WorkUnit<E::ConcreteNode> = nodes.collect();
        if may_dispatch_tail {
            top_down_dom(&work, root,
                         traversal_data, scope, pool, traversal, tls);
        } else {
            scope.spawn(move |scope| {
                let work = work;
                top_down_dom(&work, root,
                             traversal_data, scope, pool, traversal, tls);
            });
        }
    } else {
        for chunk in nodes.chunks(WORK_UNIT_MAX).into_iter() {
            let nodes: WorkUnit<E::ConcreteNode> = chunk.collect();
            let traversal_data_copy = traversal_data.clone();
            scope.spawn(move |scope| {
                let n = nodes;
                top_down_dom(&*n, root,
                             traversal_data_copy, scope, pool, traversal, tls)
            });
        }
    }
}
<|endoftext|>


<commit_before>use bidir_map::BidirMap;

use hmac::{Hmac, Mac};
use sha2::{Sha256};

use serde_json;
use serde_json::Value;
use serde_json::value::Map;

use std::thread;
use std::time::Duration;

use error::*;
use helpers;
use types::Currency;
use types::Pair;
use types::Pair::*;

lazy_static! {
    static ref PAIRS_STRING: BidirMap<Pair, &'static str> = {
        let mut m = BidirMap::new();
        m.insert(BTC_USD, "btcusd");
        m.insert(BTC_EUR, "btceur");
        m.insert(EUR_USD, "eurusd");
        m.insert(XRP_USD, "xrpusd");
        m.insert(XRP_EUR, "xrpeur");
        m.insert(XRP_BTC, "xrpbtc");
        m
    };
}

/// Return the name associated to pair used by Bitstamp
/// If the Pair is not supported, None is returned.
pub fn get_pair_string(pair: &Pair) -> Option<&&str> {
    PAIRS_STRING.get_by_first(pair)
}

/// Return the Pair enum associated to the string used by Bitstamp
/// If the Pair is not supported, None is returned.
pub fn get_pair_enum(pair: &str) -> Option<&Pair> {
    PAIRS_STRING.get_by_second(&pair)
}

pub fn block_or_continue(last_request: i64) {
    let threshold = 1000; // 600 requests per 10 mins = 1 request per second
    let delay = helpers::get_unix_timestamp_ms() - last_request;
    if delay < threshold {
        let duration_ms = Duration::from_millis(delay as u64);
        thread::sleep(duration_ms);
    }
}

pub fn build_signature(nonce: &str,
                       customer_id: &str,
                       api_key: &str,
                       api_secret: &str)
                       -> Result<String> {
    const C: &'static [u8] = b"0123456789ABCDEF";

    let message = nonce.to_owned() + customer_id + api_key;

    let mut mac = Hmac::<Sha256>::new(api_secret.as_bytes());

    mac.input(message.as_bytes());
    let result = mac.result();

    let raw_signature = result.code();
    let mut signature = Vec::with_capacity(raw_signature.len() * 2);
    for &byte in raw_signature {
        signature.push(C[(byte >> 4) as usize]);
        signature.push(C[(byte & 0xf) as usize]);
    }
    // TODO: Handle correctly the from_utf8 errors with error_chain.
    Ok(String::from_utf8(signature)?)
}

pub fn build_url(method: &str, pair: &str) -> String {
    "https://www.bitstamp.net/api/v2/".to_string() + method + "/" + pair + "/"
}

pub fn deserialize_json(json_string: &str) -> Result<Map<String, Value>> {
    let data: Value = match serde_json::from_str(json_string) {
        Ok(data) => data,
        Err(_) => return Err(ErrorKind::BadParse.into()),
    };

    match data.as_object() {
        Some(value) => Ok(value.clone()),
        None => Err(ErrorKind::BadParse.into()),
    }
}

pub fn generate_nonce(fixed_nonce: Option<String>) -> String {
    match fixed_nonce {
        Some(v) => v,
        None => helpers::get_unix_timestamp_ms().to_string(),
    }
}

/// If error array is null, return the result (encoded in a json object)
/// else return the error string found in array
pub fn parse_result(response: &Map<String, Value>) -> Result<Map<String, Value>> {
    let error_msg = match response.get("error") {
        Some(error) => {
            error
                .as_str()
                .ok_or_else(|| ErrorKind::InvalidFieldFormat("error".to_string()))?
        }
        None => return Ok(response.clone()),
    };

    match error_msg.as_ref() {
        "Invalid command." => Err(ErrorKind::InvalidArguments.into()),
        "Invalid API key/secret pair." => Err(ErrorKind::BadCredentials.into()),
        "Total must be at least 0.0001." => Err(ErrorKind::InsufficientOrderSize.into()),
        other => Err(ErrorKind::ExchangeSpecificError(other.to_string()).into()),
    }
}

/// Return the currency enum associated with the
/// string used by Bitstamp. If no currency is found,
/// return None
/// # Examples
///
/// ```
/// use coinnect::bitstamp::utils::get_currency_enum;
/// use coinnect::types::Currency;
///
/// let currency = get_currency_enum("usd_balance");
/// assert_eq!(Some(Currency::USD), currency);
/// ```
pub fn get_currency_enum(currency: &str) -> Option<Currency> {
    match currency {
        "usd_balance" => Some(Currency::USD),
        "btc_balance" => Some(Currency::BTC),
        "eur_balance" => Some(Currency::EUR),
        "xrp_balance" => Some(Currency::XRP),
        _ => None,
    }
}

/// Return the currency string associated with the
/// enum used by Bitstamp. If no currency is found,
/// return None
/// # Examples
///
/// ```
/// use coinnect::bitstamp::utils::get_currency_string;
/// use coinnect::types::Currency;
///
/// let currency = get_currency_string(Currency::USD);
/// assert_eq!(currency, Some("USD".to_string()));
/// ```
pub fn get_currency_string(currency: Currency) -> Option<String> {
    match currency {
        Currency::USD => Some("USD".to_string()),
        Currency::BTC => Some("BTC".to_string()),
        Currency::EUR => Some("EUR".to_string()),
        Currency::XRP => Some("XRP".to_string()),
        _ => None,
    }
}
<commit_msg>Fix `block_or_continue` bug.<commit_after>use bidir_map::BidirMap;

use hmac::{Hmac, Mac};
use sha2::{Sha256};

use serde_json;
use serde_json::Value;
use serde_json::value::Map;

use std::thread;
use std::time::Duration;

use error::*;
use helpers;
use types::Currency;
use types::Pair;
use types::Pair::*;

lazy_static! {
    static ref PAIRS_STRING: BidirMap<Pair, &'static str> = {
        let mut m = BidirMap::new();
        m.insert(BTC_USD, "btcusd");
        m.insert(BTC_EUR, "btceur");
        m.insert(EUR_USD, "eurusd");
        m.insert(XRP_USD, "xrpusd");
        m.insert(XRP_EUR, "xrpeur");
        m.insert(XRP_BTC, "xrpbtc");
        m
    };
}

/// Return the name associated to pair used by Bitstamp
/// If the Pair is not supported, None is returned.
pub fn get_pair_string(pair: &Pair) -> Option<&&str> {
    PAIRS_STRING.get_by_first(pair)
}

/// Return the Pair enum associated to the string used by Bitstamp
/// If the Pair is not supported, None is returned.
pub fn get_pair_enum(pair: &str) -> Option<&Pair> {
    PAIRS_STRING.get_by_second(&pair)
}

pub fn block_or_continue(last_request: i64) {
    let threshold: u64 = 1000; // 600 requests per 10 mins = 1 request per second
    let offset: u64 = helpers::get_unix_timestamp_ms() as u64 - last_request as u64;
    if offset < threshold {
        let wait_ms = Duration::from_millis(threshold - offset);
        thread::sleep(wait_ms);
    }
}

pub fn build_signature(nonce: &str,
                       customer_id: &str,
                       api_key: &str,
                       api_secret: &str)
                       -> Result<String> {
    const C: &'static [u8] = b"0123456789ABCDEF";

    let message = nonce.to_owned() + customer_id + api_key;

    let mut mac = Hmac::<Sha256>::new(api_secret.as_bytes());

    mac.input(message.as_bytes());
    let result = mac.result();

    let raw_signature = result.code();
    let mut signature = Vec::with_capacity(raw_signature.len() * 2);
    for &byte in raw_signature {
        signature.push(C[(byte >> 4) as usize]);
        signature.push(C[(byte & 0xf) as usize]);
    }
    // TODO: Handle correctly the from_utf8 errors with error_chain.
    Ok(String::from_utf8(signature)?)
}

pub fn build_url(method: &str, pair: &str) -> String {
    "https://www.bitstamp.net/api/v2/".to_string() + method + "/" + pair + "/"
}

pub fn deserialize_json(json_string: &str) -> Result<Map<String, Value>> {
    let data: Value = match serde_json::from_str(json_string) {
        Ok(data) => data,
        Err(_) => return Err(ErrorKind::BadParse.into()),
    };

    match data.as_object() {
        Some(value) => Ok(value.clone()),
        None => Err(ErrorKind::BadParse.into()),
    }
}

pub fn generate_nonce(fixed_nonce: Option<String>) -> String {
    match fixed_nonce {
        Some(v) => v,
        None => helpers::get_unix_timestamp_ms().to_string(),
    }
}

/// If error array is null, return the result (encoded in a json object)
/// else return the error string found in array
pub fn parse_result(response: &Map<String, Value>) -> Result<Map<String, Value>> {
    let error_msg = match response.get("error") {
        Some(error) => {
            error
                .as_str()
                .ok_or_else(|| ErrorKind::InvalidFieldFormat("error".to_string()))?
        }
        None => return Ok(response.clone()),
    };

    match error_msg.as_ref() {
        "Invalid command." => Err(ErrorKind::InvalidArguments.into()),
        "Invalid API key/secret pair." => Err(ErrorKind::BadCredentials.into()),
        "Total must be at least 0.0001." => Err(ErrorKind::InsufficientOrderSize.into()),
        other => Err(ErrorKind::ExchangeSpecificError(other.to_string()).into()),
    }
}

/// Return the currency enum associated with the
/// string used by Bitstamp. If no currency is found,
/// return None
/// # Examples
///
/// ```
/// use coinnect::bitstamp::utils::get_currency_enum;
/// use coinnect::types::Currency;
///
/// let currency = get_currency_enum("usd_balance");
/// assert_eq!(Some(Currency::USD), currency);
/// ```
pub fn get_currency_enum(currency: &str) -> Option<Currency> {
    match currency {
        "usd_balance" => Some(Currency::USD),
        "btc_balance" => Some(Currency::BTC),
        "eur_balance" => Some(Currency::EUR),
        "xrp_balance" => Some(Currency::XRP),
        _ => None,
    }
}

/// Return the currency string associated with the
/// enum used by Bitstamp. If no currency is found,
/// return None
/// # Examples
///
/// ```
/// use coinnect::bitstamp::utils::get_currency_string;
/// use coinnect::types::Currency;
///
/// let currency = get_currency_string(Currency::USD);
/// assert_eq!(currency, Some("USD".to_string()));
/// ```
pub fn get_currency_string(currency: Currency) -> Option<String> {
    match currency {
        Currency::USD => Some("USD".to_string()),
        Currency::BTC => Some("BTC".to_string()),
        Currency::EUR => Some("EUR".to_string()),
        Currency::XRP => Some("XRP".to_string()),
        _ => None,
    }
}

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

    #[test]
    fn should_block_when_enabled() {
        let start = helpers::get_unix_timestamp_ms();
        block_or_continue(start);
        let end = helpers::get_unix_timestamp_ms();
        assert!(end - start >= 1000)
    }
}<|endoftext|>


<commit_before>//! TODO Fill in

use std::fmt::{self, Display, Formatter};
use std::default::Default;
use std::rc::Rc;
use rlua::{self, Table, Lua, UserData, ToLua, Value};
use super::object::{Object, Objectable, ObjectBuilder};
use super::property::Property;
use super::class::{self, Class, ClassBuilder};

#[derive(Clone, Debug)]
pub struct ScreenState {
    // TODO Fill in
    dummy: i32
}

pub struct Screen<'lua>(Table<'lua>);

impl_objectable!(Screen, ScreenState);

impl Display for ScreenState {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Screen: {:p}", self)
    }
}

impl <'lua> ToLua<'lua> for Screen<'lua> {
    fn to_lua(self, lua: &'lua Lua) -> rlua::Result<Value<'lua>> {
        self.0.to_lua(lua)
    }
}

impl UserData for ScreenState {}

impl Default for ScreenState {
    fn default() -> Self {
        ScreenState {
            dummy: 0
        }
    }
}

impl <'lua> Screen<'lua> {
    fn new(lua: &Lua) -> rlua::Result<Object> {
        let class = class::class_setup(lua, "screen")?;
        Ok(object_setup(lua, Screen::allocate(lua, class)?)?.build())
    }
}

pub fn init(lua: &Lua) -> rlua::Result<Class> {
    method_setup(lua, Class::builder(lua, Some(Rc::new(Screen::new)), None, None)?)?
        .save_class("screen")?
        .build()
}

fn method_setup<'lua>(lua: &'lua Lua, builder: ClassBuilder<'lua>) -> rlua::Result<ClassBuilder<'lua>> {
    // TODO Do properly
    use super::dummy;
    builder.method("connect_signal".into(), lua.create_function(dummy))?
           .method("__call".into(), lua.create_function(dummy))
}

fn object_setup<'lua>(lua: &'lua Lua, builder: ObjectBuilder<'lua>) -> rlua::Result<ObjectBuilder<'lua>> {
    // TODO FIXME ughhhhhh

    //// TODO Make sure I'm doing this right...
    //let screen_table = Screen::new(lua)?.to_lua(lua)?;
    //table.set("screen", screen_table)?;
    //meta_table.set("__tostring", lua.create_function(|_, val: Table| {
    //    Ok(format!("{}", val.get::<_, ScreenState>("__data")?))
    //}))?;
    //table.set("screen", screen_table)?;
    builder//.add_to_meta(table)?
           .property(Property::new("screen".into(),
                                   // TODO
                                   Some(lua.create_function(screen_new)),
                                   Some(lua.create_function(get_visible)),
                                   Some(lua.create_function(set_visible))))
}

fn screen_new<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<()> {
    unimplemented!()
}

fn get_visible<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<()> {
    unimplemented!()
}

fn set_visible<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<()> {
    unimplemented!()
}
<commit_msg>Removed bad object setup code in screen<commit_after>//! TODO Fill in

use std::fmt::{self, Display, Formatter};
use std::default::Default;
use std::rc::Rc;
use rlua::{self, Table, Lua, UserData, ToLua, Value};
use super::object::{Object, Objectable, ObjectBuilder};
use super::property::Property;
use super::class::{self, Class, ClassBuilder};

#[derive(Clone, Debug)]
pub struct ScreenState {
    // TODO Fill in
    dummy: i32
}

pub struct Screen<'lua>(Table<'lua>);

impl_objectable!(Screen, ScreenState);

impl Display for ScreenState {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "Screen: {:p}", self)
    }
}

impl <'lua> ToLua<'lua> for Screen<'lua> {
    fn to_lua(self, lua: &'lua Lua) -> rlua::Result<Value<'lua>> {
        self.0.to_lua(lua)
    }
}

impl UserData for ScreenState {}

impl Default for ScreenState {
    fn default() -> Self {
        ScreenState {
            dummy: 0
        }
    }
}

impl <'lua> Screen<'lua> {
    fn new(lua: &Lua) -> rlua::Result<Object> {
        let class = class::class_setup(lua, "screen")?;
        Ok(object_setup(lua, Screen::allocate(lua, class)?)?.build())
    }
}

pub fn init(lua: &Lua) -> rlua::Result<Class> {
    method_setup(lua, Class::builder(lua, Some(Rc::new(Screen::new)), None, None)?)?
        .save_class("screen")?
        .build()
}

fn method_setup<'lua>(lua: &'lua Lua, builder: ClassBuilder<'lua>) -> rlua::Result<ClassBuilder<'lua>> {
    // TODO Do properly
    use super::dummy;
    builder.method("connect_signal".into(), lua.create_function(dummy))?
           .method("__call".into(), lua.create_function(dummy))
}

fn object_setup<'lua>(lua: &'lua Lua, builder: ObjectBuilder<'lua>) -> rlua::Result<ObjectBuilder<'lua>> {
    builder
           .property(Property::new("screen".into(),
                                   // TODO Implement
                                   Some(lua.create_function(screen_new)),
                                   Some(lua.create_function(get_visible)),
                                   Some(lua.create_function(set_visible))))
}

fn screen_new<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<()> {
    unimplemented!()
}

fn get_visible<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<()> {
    unimplemented!()
}

fn set_visible<'lua>(lua: &'lua Lua, _: ()) -> rlua::Result<()> {
    unimplemented!()
}
<|endoftext|>


<commit_before>use super::invoke::{Invoke, Identity};
use super::code::{self, SignaturesBuilder, FunctionBuilder};
use super::import;
use elements;

/// Module builder
pub struct ModuleBuilder<F=Identity> {
    callback: F,
    module: ModuleScaffold,
}

/// Location of the internal module function
pub struct CodeLocation {
    /// Location (index in 'functions' section) of the signature
    pub signature: u32,
    /// Location (index in the 'code' section) of the body
    pub body: u32,
}

#[derive(Default)]
struct ModuleScaffold {
    pub functions: elements::FunctionsSection,
    pub types: elements::TypeSection,
    pub import: elements::ImportSection,
    pub code: elements::CodeSection,
    pub other: Vec<elements::Section>,
}

impl From<elements::Module> for ModuleScaffold {
    fn from(module: elements::Module) -> Self {
        let mut funcs: Option<elements::FunctionsSection> = None;
        let mut types: Option<elements::TypeSection> = None;
        let mut import: Option<elements::ImportSection> = None;
        let mut code: Option<elements::CodeSection> = None;

        let mut sections = module.into_sections();
        while let Some(section) = sections.pop() {
            match section {
                elements::Section::Type(sect) => { types = Some(sect); }
                elements::Section::Function(sect) => { funcs = Some(sect); }
                elements::Section::Import(sect) => { import = Some(sect); }
                elements::Section::Code(sect) => { code = Some(sect); }
                _ => {}
            }
        }

        ModuleScaffold {
            functions: funcs.unwrap_or_default(),
            types: types.unwrap_or_default(),
            import: import.unwrap_or_default(),
            code: code.unwrap_or_default(),
            other: sections,
        }
    }
}

impl From<ModuleScaffold> for elements::Module {
    fn from(module: ModuleScaffold) -> Self {
        let mut sections = Vec::new();

        let types = module.types;
        if types.types().len() > 0 {
            sections.push(elements::Section::Type(types));
        }
        let import = module.import;
        if import.entries().len() > 0 {
            sections.push(elements::Section::Import(import));
        }
        let functions = module.functions;
        if functions.entries().len() > 0 {
            sections.push(elements::Section::Function(functions));
        }
        let code = module.code;
        if code.bodies().len() > 0 {
            sections.push(elements::Section::Code(code));
        }
        sections.extend(module.other);
        elements::Module::new(sections)
    }
}

impl ModuleBuilder {
    /// New empty module builder
    pub fn new() -> Self {
        ModuleBuilder::with_callback(Identity)
    }
}

impl<F> ModuleBuilder<F> where F: Invoke<elements::Module> {
    /// New module builder with bound callback
    pub fn with_callback(callback: F) -> Self {
        ModuleBuilder {
            callback: callback,
            module: Default::default(),
        }
    }

    /// Builder from raw module
    pub fn with_module(mut self, module: elements::Module) -> Self {
        self.module = module.into();
        self
    }

    /// Fill module with sections from iterator
    pub fn with_sections<I>(mut self, sections: I) -> Self
        where I: IntoIterator<Item=elements::Section>
    {
        self.module.other.extend(sections);
        self
    }

    /// Add additional section
    pub fn with_section(mut self, section: elements::Section) -> Self {
        self.module.other.push(section);
        self
    }

    /// Binds to the type section, creates additional types when required
    pub fn with_signatures(mut self, bindings: code::SignatureBindings) -> Self {
        self.push_signatures(bindings);
        self
    }

    /// Push stand-alone function definition, creating sections, signature and code blocks
    /// in corresponding sections.
    /// `FunctionDefinition` can be build using `builder::function` builder
    pub fn push_function(&mut self, func: code::FunctionDefinition) -> CodeLocation {
        let signature = func.signature;
        let body = func.code;

        let type_ref = self.resolve_type_ref(signature);

        self.module.functions.entries_mut().push(elements::Func::new(type_ref));
        let signature_index = self.module.functions.entries_mut().len() as u32 - 1;
        self.module.code.bodies_mut().push(body);
        let body_index = self.module.code.bodies_mut().len() as u32 - 1;

        CodeLocation {
            signature: signature_index,
            body: body_index,
        }
    }

    fn resolve_type_ref(&mut self, signature: code::Signature) -> u32 {
        match signature {
            code::Signature::Inline(func_type) => {
                self.module.types.types_mut().push(elements::Type::Function(func_type));
                self.module.types.types().len() as u32 - 1
            }
            code::Signature::TypeReference(type_ref) => {
                type_ref
            }
        }
    }

    /// Push one function signature, returning it's calling index.
    /// Can create corresponding type in type section.
    pub fn push_signature(&mut self, signature: code::Signature) -> u32 {
        self.resolve_type_ref(signature)
    }

    /// Push signatures in the module, returning corresponding indices of pushed signatures
    pub fn push_signatures(&mut self, signatures: code::SignatureBindings) -> Vec<u32> {
        let mut result = Vec::new();

        // todo: maybe reuse existing types with the equal signatures
        let raw_functions: Vec<u32> = signatures.into_iter().map(|binding|
            match binding {
                code::Signature::Inline(func_type) => {
                    self.module.types.types_mut().push(elements::Type::Function(func_type));
                    self.module.types.types().len() as u32 - 1
                }
                code::Signature::TypeReference(type_ref) => {
                    type_ref
                }
            }
        ).collect();

        for function in raw_functions {
            self.module.functions.entries_mut().push(elements::Func::new(function));
            result.push(self.module.functions.entries_mut().len() as u32 - 1);
        }

        result
    }

    pub fn function(self) -> FunctionBuilder<Self> {
        FunctionBuilder::with_callback(self)
    }

    /// Define functions section
    pub fn functions(self) -> SignaturesBuilder<Self> {
        SignaturesBuilder::with_callback(self)
    }

    /// With inserted import entry
    pub fn with_import(mut self, entry: elements::ImportEntry) -> Self {
        self.module.import.entries_mut().push(entry);
        self
    }

    /// Import entry builder
    pub fn import(self) -> import::ImportBuilder<Self> {
        import::ImportBuilder::with_callback(self)
    }

    /// Build module (final step)
    pub fn build(self) -> F::Result {
        self.callback.invoke(self.module.into())
    }
}

impl<F> Invoke<elements::FunctionsSection> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, section: elements::FunctionsSection) -> Self {
		self.with_section(elements::Section::Function(section))
    }
}

impl<F> Invoke<code::SignatureBindings> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, bindings: code::SignatureBindings) -> Self {
        self.with_signatures(bindings)
    }
}

impl<F> Invoke<code::FunctionDefinition> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, def: code::FunctionDefinition) -> Self {
        let mut b = self;
        b.push_function(def);
        b
    }
}

impl<F> Invoke<elements::ImportEntry> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, entry: elements::ImportEntry) -> Self::Result {
        self.with_import(entry)
    }
}

/// Start new module builder
pub fn module() -> ModuleBuilder {
    ModuleBuilder::new()
}

/// Start builder to extend existing module
pub fn from_module(module: elements::Module) -> ModuleBuilder {
    ModuleBuilder::new().with_module(module)
}

#[cfg(test)]
mod tests {

    use super::module;

    #[test]
    fn smoky() {
        let module = module().build();
        assert_eq!(module.sections().len(), 0);
    }

    #[test]
    fn functions() {
        let module = module()
            .function()
                .signature().param().i32().build()
                .body().build()
                .build()
            .build();

        assert_eq!(module.type_section().expect("type section to exist").types().len(), 1);
        assert_eq!(module.functions_section().expect("function section to exist").entries().len(), 1);
        assert_eq!(module.code_section().expect("code section to exist").bodies().len(), 1);
    }

}
<commit_msg>fix also for signature collection<commit_after>use super::invoke::{Invoke, Identity};
use super::code::{self, SignaturesBuilder, FunctionBuilder};
use super::import;
use elements;

/// Module builder
pub struct ModuleBuilder<F=Identity> {
    callback: F,
    module: ModuleScaffold,
}

/// Location of the internal module function
pub struct CodeLocation {
    /// Location (index in 'functions' section) of the signature
    pub signature: u32,
    /// Location (index in the 'code' section) of the body
    pub body: u32,
}

#[derive(Default)]
struct ModuleScaffold {
    pub functions: elements::FunctionsSection,
    pub types: elements::TypeSection,
    pub import: elements::ImportSection,
    pub code: elements::CodeSection,
    pub other: Vec<elements::Section>,
}

impl From<elements::Module> for ModuleScaffold {
    fn from(module: elements::Module) -> Self {
        let mut funcs: Option<elements::FunctionsSection> = None;
        let mut types: Option<elements::TypeSection> = None;
        let mut import: Option<elements::ImportSection> = None;
        let mut code: Option<elements::CodeSection> = None;

        let mut sections = module.into_sections();
        while let Some(section) = sections.pop() {
            match section {
                elements::Section::Type(sect) => { types = Some(sect); }
                elements::Section::Function(sect) => { funcs = Some(sect); }
                elements::Section::Import(sect) => { import = Some(sect); }
                elements::Section::Code(sect) => { code = Some(sect); }
                _ => {}
            }
        }

        ModuleScaffold {
            functions: funcs.unwrap_or_default(),
            types: types.unwrap_or_default(),
            import: import.unwrap_or_default(),
            code: code.unwrap_or_default(),
            other: sections,
        }
    }
}

impl From<ModuleScaffold> for elements::Module {
    fn from(module: ModuleScaffold) -> Self {
        let mut sections = Vec::new();

        let types = module.types;
        if types.types().len() > 0 {
            sections.push(elements::Section::Type(types));
        }
        let import = module.import;
        if import.entries().len() > 0 {
            sections.push(elements::Section::Import(import));
        }
        let functions = module.functions;
        if functions.entries().len() > 0 {
            sections.push(elements::Section::Function(functions));
        }
        let code = module.code;
        if code.bodies().len() > 0 {
            sections.push(elements::Section::Code(code));
        }
        sections.extend(module.other);
        elements::Module::new(sections)
    }
}

impl ModuleBuilder {
    /// New empty module builder
    pub fn new() -> Self {
        ModuleBuilder::with_callback(Identity)
    }
}

impl<F> ModuleBuilder<F> where F: Invoke<elements::Module> {
    /// New module builder with bound callback
    pub fn with_callback(callback: F) -> Self {
        ModuleBuilder {
            callback: callback,
            module: Default::default(),
        }
    }

    /// Builder from raw module
    pub fn with_module(mut self, module: elements::Module) -> Self {
        self.module = module.into();
        self
    }

    /// Fill module with sections from iterator
    pub fn with_sections<I>(mut self, sections: I) -> Self
        where I: IntoIterator<Item=elements::Section>
    {
        self.module.other.extend(sections);
        self
    }

    /// Add additional section
    pub fn with_section(mut self, section: elements::Section) -> Self {
        self.module.other.push(section);
        self
    }

    /// Binds to the type section, creates additional types when required
    pub fn with_signatures(mut self, bindings: code::SignatureBindings) -> Self {
        self.push_signatures(bindings);
        self
    }

    /// Push stand-alone function definition, creating sections, signature and code blocks
    /// in corresponding sections.
    /// `FunctionDefinition` can be build using `builder::function` builder
    pub fn push_function(&mut self, func: code::FunctionDefinition) -> CodeLocation {
        let signature = func.signature;
        let body = func.code;

        let type_ref = self.resolve_type_ref(signature);

        self.module.functions.entries_mut().push(elements::Func::new(type_ref));
        let signature_index = self.module.functions.entries_mut().len() as u32 - 1;
        self.module.code.bodies_mut().push(body);
        let body_index = self.module.code.bodies_mut().len() as u32 - 1;

        CodeLocation {
            signature: signature_index,
            body: body_index,
        }
    }

    fn resolve_type_ref(&mut self, signature: code::Signature) -> u32 {
        match signature {
            code::Signature::Inline(func_type) => {
                // todo: maybe search for existing type
                self.module.types.types_mut().push(elements::Type::Function(func_type));
                self.module.types.types().len() as u32 - 1
            }
            code::Signature::TypeReference(type_ref) => {
                type_ref
            }
        }
    }

    /// Push one function signature, returning it's calling index.
    /// Can create corresponding type in type section.
    pub fn push_signature(&mut self, signature: code::Signature) -> u32 {
        self.resolve_type_ref(signature)
    }

    /// Push signatures in the module, returning corresponding indices of pushed signatures
    pub fn push_signatures(&mut self, signatures: code::SignatureBindings) -> Vec<u32> {
        signatures.into_iter().map(|binding|
            self.resolve_type_ref(binding)
        ).collect()
    }

    pub fn function(self) -> FunctionBuilder<Self> {
        FunctionBuilder::with_callback(self)
    }

    /// Define functions section
    pub fn functions(self) -> SignaturesBuilder<Self> {
        SignaturesBuilder::with_callback(self)
    }

    /// With inserted import entry
    pub fn with_import(mut self, entry: elements::ImportEntry) -> Self {
        self.module.import.entries_mut().push(entry);
        self
    }

    /// Import entry builder
    pub fn import(self) -> import::ImportBuilder<Self> {
        import::ImportBuilder::with_callback(self)
    }

    /// Build module (final step)
    pub fn build(self) -> F::Result {
        self.callback.invoke(self.module.into())
    }
}

impl<F> Invoke<elements::FunctionsSection> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, section: elements::FunctionsSection) -> Self {
		self.with_section(elements::Section::Function(section))
    }
}

impl<F> Invoke<code::SignatureBindings> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, bindings: code::SignatureBindings) -> Self {
        self.with_signatures(bindings)
    }
}

impl<F> Invoke<code::FunctionDefinition> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, def: code::FunctionDefinition) -> Self {
        let mut b = self;
        b.push_function(def);
        b
    }
}

impl<F> Invoke<elements::ImportEntry> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, entry: elements::ImportEntry) -> Self::Result {
        self.with_import(entry)
    }
}

/// Start new module builder
pub fn module() -> ModuleBuilder {
    ModuleBuilder::new()
}

/// Start builder to extend existing module
pub fn from_module(module: elements::Module) -> ModuleBuilder {
    ModuleBuilder::new().with_module(module)
}

#[cfg(test)]
mod tests {

    use super::module;

    #[test]
    fn smoky() {
        let module = module().build();
        assert_eq!(module.sections().len(), 0);
    }

    #[test]
    fn functions() {
        let module = module()
            .function()
                .signature().param().i32().build()
                .body().build()
                .build()
            .build();

        assert_eq!(module.type_section().expect("type section to exist").types().len(), 1);
        assert_eq!(module.functions_section().expect("function section to exist").entries().len(), 1);
        assert_eq!(module.code_section().expect("code section to exist").bodies().len(), 1);
    }

}
<|endoftext|>