<commit_before>use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{quote, quote_spanned, ToTokens};

#[derive(Clone, Copy, PartialEq)]
enum RuntimeFlavor {
    CurrentThread,
    Threaded,
}

impl RuntimeFlavor {
    fn from_str(s: &str) -> Result<RuntimeFlavor, String> {
        match s {
            "current_thread" => Ok(RuntimeFlavor::CurrentThread),
            "multi_thread" => Ok(RuntimeFlavor::Threaded),
            "single_thread" => Err("The single threaded runtime flavor is called `current_thread`.".to_string()),
            "basic_scheduler" => Err("The `basic_scheduler` runtime flavor has been renamed to `current_thread`.".to_string()),
            "threaded_scheduler" => Err("The `threaded_scheduler` runtime flavor has been renamed to `multi_thread`.".to_string()),
            _ => Err(format!("No such runtime flavor `{}`. The runtime flavors are `current_thread` and `multi_thread`.", s)),
        }
    }
}

struct FinalConfig {
    flavor: RuntimeFlavor,
    worker_threads: Option<usize>,
    start_paused: Option<bool>,
}

struct Configuration {
    rt_multi_thread_available: bool,
    default_flavor: RuntimeFlavor,
    flavor: Option<RuntimeFlavor>,
    worker_threads: Option<(usize, Span)>,
    start_paused: Option<(bool, Span)>,
    is_test: bool,
}

impl Configuration {
    fn new(is_test: bool, rt_multi_thread: bool) -> Self {
        Configuration {
            rt_multi_thread_available: rt_multi_thread,
            default_flavor: match is_test {
                true => RuntimeFlavor::CurrentThread,
                false => RuntimeFlavor::Threaded,
            },
            flavor: None,
            worker_threads: None,
            start_paused: None,
            is_test,
        }
    }

    fn set_flavor(&mut self, runtime: syn::Lit, span: Span) -> Result<(), syn::Error> {
        if self.flavor.is_some() {
            return Err(syn::Error::new(span, "`flavor` set multiple times."));
        }

        let runtime_str = parse_string(runtime, span, "flavor")?;
        let runtime =
            RuntimeFlavor::from_str(&runtime_str).map_err(|err| syn::Error::new(span, err))?;
        self.flavor = Some(runtime);
        Ok(())
    }

    fn set_worker_threads(
        &mut self,
        worker_threads: syn::Lit,
        span: Span,
    ) -> Result<(), syn::Error> {
        if self.worker_threads.is_some() {
            return Err(syn::Error::new(
                span,
                "`worker_threads` set multiple times.",
            ));
        }

        let worker_threads = parse_int(worker_threads, span, "worker_threads")?;
        if worker_threads == 0 {
            return Err(syn::Error::new(span, "`worker_threads` may not be 0."));
        }
        self.worker_threads = Some((worker_threads, span));
        Ok(())
    }

    fn set_start_paused(&mut self, start_paused: syn::Lit, span: Span) -> Result<(), syn::Error> {
        if self.start_paused.is_some() {
            return Err(syn::Error::new(span, "`start_paused` set multiple times."));
        }

        let start_paused = parse_bool(start_paused, span, "start_paused")?;
        self.start_paused = Some((start_paused, span));
        Ok(())
    }

    fn macro_name(&self) -> &'static str {
        if self.is_test {
            "tokio::test"
        } else {
            "tokio::main"
        }
    }

    fn build(&self) -> Result<FinalConfig, syn::Error> {
        let flavor = self.flavor.unwrap_or(self.default_flavor);
        use RuntimeFlavor::*;

        let worker_threads = match (flavor, self.worker_threads) {
            (CurrentThread, Some((_, worker_threads_span))) => {
                let msg = format!(
                    "The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[{}(flavor = \"multi_thread\")]`",
                    self.macro_name(),
                );
                return Err(syn::Error::new(worker_threads_span, msg));
            }
            (CurrentThread, None) => None,
            (Threaded, worker_threads) if self.rt_multi_thread_available => {
                worker_threads.map(|(val, _span)| val)
            }
            (Threaded, _) => {
                let msg = if self.flavor.is_none() {
                    "The default runtime flavor is `multi_thread`, but the `rt-multi-thread` feature is disabled."
                } else {
                    "The runtime flavor `multi_thread` requires the `rt-multi-thread` feature."
                };
                return Err(syn::Error::new(Span::call_site(), msg));
            }
        };

        let start_paused = match (flavor, self.start_paused) {
            (Threaded, Some((_, start_paused_span))) => {
                let msg = format!(
                    "The `start_paused` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
                    self.macro_name(),
                );
                return Err(syn::Error::new(start_paused_span, msg));
            }
            (CurrentThread, Some((start_paused, _))) => Some(start_paused),
            (_, None) => None,
        };

        Ok(FinalConfig {
            flavor,
            worker_threads,
            start_paused,
        })
    }
}

fn parse_int(int: syn::Lit, span: Span, field: &str) -> Result<usize, syn::Error> {
    match int {
        syn::Lit::Int(lit) => match lit.base10_parse::<usize>() {
            Ok(value) => Ok(value),
            Err(e) => Err(syn::Error::new(
                span,
                format!("Failed to parse value of `{}` as integer: {}", field, e),
            )),
        },
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as integer.", field),
        )),
    }
}

fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::Error> {
    match int {
        syn::Lit::Str(s) => Ok(s.value()),
        syn::Lit::Verbatim(s) => Ok(s.to_string()),
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as string.", field),
        )),
    }
}

fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Error> {
    match bool {
        syn::Lit::Bool(b) => Ok(b.value),
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as bool.", field),
        )),
    }
}

fn parse_knobs(
    mut input: syn::ItemFn,
    args: syn::AttributeArgs,
    is_test: bool,
    rt_multi_thread: bool,
) -> Result<TokenStream, syn::Error> {
    if input.sig.asyncness.take().is_none() {
        let msg = "the `async` keyword is missing from the function declaration";
        return Err(syn::Error::new_spanned(input.sig.fn_token, msg));
    }

    let mut config = Configuration::new(is_test, rt_multi_thread);
    let macro_name = config.macro_name();

    for arg in args {
        match arg {
            syn::NestedMeta::Meta(syn::Meta::NameValue(namevalue)) => {
                let ident = namevalue
                    .path
                    .get_ident()
                    .ok_or_else(|| {
                        syn::Error::new_spanned(&namevalue, "Must have specified ident")
                    })?
                    .to_string()
                    .to_lowercase();
                match ident.as_str() {
                    "worker_threads" => {
                        config.set_worker_threads(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "flavor" => {
                        config.set_flavor(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "start_paused" => {
                        config.set_start_paused(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "core_threads" => {
                        let msg = "Attribute `core_threads` is renamed to `worker_threads`";
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                    name => {
                        let msg = format!(
                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`",
                            name,
                        );
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                }
            }
            syn::NestedMeta::Meta(syn::Meta::Path(path)) => {
                let name = path
                    .get_ident()
                    .ok_or_else(|| syn::Error::new_spanned(&path, "Must have specified ident"))?
                    .to_string()
                    .to_lowercase();
                let msg = match name.as_str() {
                    "threaded_scheduler" | "multi_thread" => {
                        format!(
                            "Set the runtime flavor with #[{}(flavor = \"multi_thread\")].",
                            macro_name
                        )
                    }
                    "basic_scheduler" | "current_thread" | "single_threaded" => {
                        format!(
                            "Set the runtime flavor with #[{}(flavor = \"current_thread\")].",
                            macro_name
                        )
                    }
                    "flavor" | "worker_threads" | "start_paused" => {
                        format!("The `{}` attribute requires an argument.", name)
                    }
                    name => {
                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`", name)
                    }
                };
                return Err(syn::Error::new_spanned(path, msg));
            }
            other => {
                return Err(syn::Error::new_spanned(
                    other,
                    "Unknown attribute inside the macro",
                ));
            }
        }
    }

    let config = config.build()?;

    // If type mismatch occurs, the current rustc points to the last statement.
    let (last_stmt_start_span, last_stmt_end_span) = {
        let mut last_stmt = input
            .block
            .stmts
            .last()
            .map(ToTokens::into_token_stream)
            .unwrap_or_default()
            .into_iter();
        // `Span` on stable Rust has a limitation that only points to the first
        // token, not the whole tokens. We can work around this limitation by
        // using the first/last span of the tokens like
        // `syn::Error::new_spanned` does.
        let start = last_stmt.next().map_or_else(Span::call_site, |t| t.span());
        let end = last_stmt.last().map_or(start, |t| t.span());
        (start, end)
    };

    let mut rt = match config.flavor {
        RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
            tokio::runtime::Builder::new_current_thread()
        },
        RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
            tokio::runtime::Builder::new_multi_thread()
        },
    };
    if let Some(v) = config.worker_threads {
        rt = quote! { #rt.worker_threads(#v) };
    }
    if let Some(v) = config.start_paused {
        rt = quote! { #rt.start_paused(#v) };
    }

    let header = if is_test {
        quote! {
            #[::core::prelude::v1::test]
        }
    } else {
        quote! {}
    };

    let body = &input.block;
    let brace_token = input.block.brace_token;
    let (tail_return, tail_semicolon) = match body.stmts.last() {
        Some(syn::Stmt::Semi(expr, _)) => (
            match expr {
                syn::Expr::Return(_) => quote! { return },
                _ => quote! {},
            },
            quote! {
                ;
            },
        ),
        _ => (quote! {}, quote! {}),
    };
    input.block = syn::parse2(quote_spanned! {last_stmt_end_span=>
        {
            let body = async #body;
            #[allow(clippy::expect_used)]
            #tail_return #rt
                .enable_all()
                .build()
                .expect("Failed building the Runtime")
                .block_on(body)#tail_semicolon
        }
    })
    .expect("Parsing failure");
    input.block.brace_token = brace_token;

    let result = quote! {
        #header
        #input
    };

    Ok(result.into())
}

#[cfg(not(test))] // Work around for rust-lang/rust#62127
pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
    let input = syn::parse_macro_input!(item as syn::ItemFn);
    let args = syn::parse_macro_input!(args as syn::AttributeArgs);

    if input.sig.ident == "main" && !input.sig.inputs.is_empty() {
        let msg = "the main function cannot accept arguments";
        return syn::Error::new_spanned(&input.sig.ident, msg)
            .to_compile_error()
            .into();
    }

    parse_knobs(input, args, false, rt_multi_thread).unwrap_or_else(|e| e.to_compile_error().into())
}

pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
    let input = syn::parse_macro_input!(item as syn::ItemFn);
    let args = syn::parse_macro_input!(args as syn::AttributeArgs);

    for attr in &input.attrs {
        if attr.path.is_ident("test") {
            let msg = "second test attribute is supplied";
            return syn::Error::new_spanned(&attr, msg)
                .to_compile_error()
                .into();
        }
    }

    parse_knobs(input, args, true, rt_multi_thread).unwrap_or_else(|e| e.to_compile_error().into())
}
<commit_msg>macros: make tokio-macros attributes more IDE friendly (#4162)<commit_after>use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{quote, quote_spanned, ToTokens};
use syn::parse::Parser;

// syn::AttributeArgs does not implement syn::Parse
type AttributeArgs = syn::punctuated::Punctuated<syn::NestedMeta, syn::Token![,]>;

#[derive(Clone, Copy, PartialEq)]
enum RuntimeFlavor {
    CurrentThread,
    Threaded,
}

impl RuntimeFlavor {
    fn from_str(s: &str) -> Result<RuntimeFlavor, String> {
        match s {
            "current_thread" => Ok(RuntimeFlavor::CurrentThread),
            "multi_thread" => Ok(RuntimeFlavor::Threaded),
            "single_thread" => Err("The single threaded runtime flavor is called `current_thread`.".to_string()),
            "basic_scheduler" => Err("The `basic_scheduler` runtime flavor has been renamed to `current_thread`.".to_string()),
            "threaded_scheduler" => Err("The `threaded_scheduler` runtime flavor has been renamed to `multi_thread`.".to_string()),
            _ => Err(format!("No such runtime flavor `{}`. The runtime flavors are `current_thread` and `multi_thread`.", s)),
        }
    }
}

struct FinalConfig {
    flavor: RuntimeFlavor,
    worker_threads: Option<usize>,
    start_paused: Option<bool>,
}

/// Config used in case of the attribute not being able to build a valid config
const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
    flavor: RuntimeFlavor::CurrentThread,
    worker_threads: None,
    start_paused: None,
};

struct Configuration {
    rt_multi_thread_available: bool,
    default_flavor: RuntimeFlavor,
    flavor: Option<RuntimeFlavor>,
    worker_threads: Option<(usize, Span)>,
    start_paused: Option<(bool, Span)>,
    is_test: bool,
}

impl Configuration {
    fn new(is_test: bool, rt_multi_thread: bool) -> Self {
        Configuration {
            rt_multi_thread_available: rt_multi_thread,
            default_flavor: match is_test {
                true => RuntimeFlavor::CurrentThread,
                false => RuntimeFlavor::Threaded,
            },
            flavor: None,
            worker_threads: None,
            start_paused: None,
            is_test,
        }
    }

    fn set_flavor(&mut self, runtime: syn::Lit, span: Span) -> Result<(), syn::Error> {
        if self.flavor.is_some() {
            return Err(syn::Error::new(span, "`flavor` set multiple times."));
        }

        let runtime_str = parse_string(runtime, span, "flavor")?;
        let runtime =
            RuntimeFlavor::from_str(&runtime_str).map_err(|err| syn::Error::new(span, err))?;
        self.flavor = Some(runtime);
        Ok(())
    }

    fn set_worker_threads(
        &mut self,
        worker_threads: syn::Lit,
        span: Span,
    ) -> Result<(), syn::Error> {
        if self.worker_threads.is_some() {
            return Err(syn::Error::new(
                span,
                "`worker_threads` set multiple times.",
            ));
        }

        let worker_threads = parse_int(worker_threads, span, "worker_threads")?;
        if worker_threads == 0 {
            return Err(syn::Error::new(span, "`worker_threads` may not be 0."));
        }
        self.worker_threads = Some((worker_threads, span));
        Ok(())
    }

    fn set_start_paused(&mut self, start_paused: syn::Lit, span: Span) -> Result<(), syn::Error> {
        if self.start_paused.is_some() {
            return Err(syn::Error::new(span, "`start_paused` set multiple times."));
        }

        let start_paused = parse_bool(start_paused, span, "start_paused")?;
        self.start_paused = Some((start_paused, span));
        Ok(())
    }

    fn macro_name(&self) -> &'static str {
        if self.is_test {
            "tokio::test"
        } else {
            "tokio::main"
        }
    }

    fn build(&self) -> Result<FinalConfig, syn::Error> {
        let flavor = self.flavor.unwrap_or(self.default_flavor);
        use RuntimeFlavor::*;

        let worker_threads = match (flavor, self.worker_threads) {
            (CurrentThread, Some((_, worker_threads_span))) => {
                let msg = format!(
                    "The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[{}(flavor = \"multi_thread\")]`",
                    self.macro_name(),
                );
                return Err(syn::Error::new(worker_threads_span, msg));
            }
            (CurrentThread, None) => None,
            (Threaded, worker_threads) if self.rt_multi_thread_available => {
                worker_threads.map(|(val, _span)| val)
            }
            (Threaded, _) => {
                let msg = if self.flavor.is_none() {
                    "The default runtime flavor is `multi_thread`, but the `rt-multi-thread` feature is disabled."
                } else {
                    "The runtime flavor `multi_thread` requires the `rt-multi-thread` feature."
                };
                return Err(syn::Error::new(Span::call_site(), msg));
            }
        };

        let start_paused = match (flavor, self.start_paused) {
            (Threaded, Some((_, start_paused_span))) => {
                let msg = format!(
                    "The `start_paused` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
                    self.macro_name(),
                );
                return Err(syn::Error::new(start_paused_span, msg));
            }
            (CurrentThread, Some((start_paused, _))) => Some(start_paused),
            (_, None) => None,
        };

        Ok(FinalConfig {
            flavor,
            worker_threads,
            start_paused,
        })
    }
}

fn parse_int(int: syn::Lit, span: Span, field: &str) -> Result<usize, syn::Error> {
    match int {
        syn::Lit::Int(lit) => match lit.base10_parse::<usize>() {
            Ok(value) => Ok(value),
            Err(e) => Err(syn::Error::new(
                span,
                format!("Failed to parse value of `{}` as integer: {}", field, e),
            )),
        },
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as integer.", field),
        )),
    }
}

fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::Error> {
    match int {
        syn::Lit::Str(s) => Ok(s.value()),
        syn::Lit::Verbatim(s) => Ok(s.to_string()),
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as string.", field),
        )),
    }
}

fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Error> {
    match bool {
        syn::Lit::Bool(b) => Ok(b.value),
        _ => Err(syn::Error::new(
            span,
            format!("Failed to parse value of `{}` as bool.", field),
        )),
    }
}

fn build_config(
    input: syn::ItemFn,
    args: AttributeArgs,
    is_test: bool,
    rt_multi_thread: bool,
) -> Result<FinalConfig, syn::Error> {
    if input.sig.asyncness.is_none() {
        let msg = "the `async` keyword is missing from the function declaration";
        return Err(syn::Error::new_spanned(input.sig.fn_token, msg));
    }

    let mut config = Configuration::new(is_test, rt_multi_thread);
    let macro_name = config.macro_name();

    for arg in args {
        match arg {
            syn::NestedMeta::Meta(syn::Meta::NameValue(namevalue)) => {
                let ident = namevalue
                    .path
                    .get_ident()
                    .ok_or_else(|| {
                        syn::Error::new_spanned(&namevalue, "Must have specified ident")
                    })?
                    .to_string()
                    .to_lowercase();
                match ident.as_str() {
                    "worker_threads" => {
                        config.set_worker_threads(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "flavor" => {
                        config.set_flavor(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "start_paused" => {
                        config.set_start_paused(
                            namevalue.lit.clone(),
                            syn::spanned::Spanned::span(&namevalue.lit),
                        )?;
                    }
                    "core_threads" => {
                        let msg = "Attribute `core_threads` is renamed to `worker_threads`";
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                    name => {
                        let msg = format!(
                            "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`",
                            name,
                        );
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                }
            }
            syn::NestedMeta::Meta(syn::Meta::Path(path)) => {
                let name = path
                    .get_ident()
                    .ok_or_else(|| syn::Error::new_spanned(&path, "Must have specified ident"))?
                    .to_string()
                    .to_lowercase();
                let msg = match name.as_str() {
                    "threaded_scheduler" | "multi_thread" => {
                        format!(
                            "Set the runtime flavor with #[{}(flavor = \"multi_thread\")].",
                            macro_name
                        )
                    }
                    "basic_scheduler" | "current_thread" | "single_threaded" => {
                        format!(
                            "Set the runtime flavor with #[{}(flavor = \"current_thread\")].",
                            macro_name
                        )
                    }
                    "flavor" | "worker_threads" | "start_paused" => {
                        format!("The `{}` attribute requires an argument.", name)
                    }
                    name => {
                        format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`", name)
                    }
                };
                return Err(syn::Error::new_spanned(path, msg));
            }
            other => {
                return Err(syn::Error::new_spanned(
                    other,
                    "Unknown attribute inside the macro",
                ));
            }
        }
    }

    config.build()
}

fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> TokenStream {
    input.sig.asyncness = None;

    // If type mismatch occurs, the current rustc points to the last statement.
    let (last_stmt_start_span, last_stmt_end_span) = {
        let mut last_stmt = input
            .block
            .stmts
            .last()
            .map(ToTokens::into_token_stream)
            .unwrap_or_default()
            .into_iter();
        // `Span` on stable Rust has a limitation that only points to the first
        // token, not the whole tokens. We can work around this limitation by
        // using the first/last span of the tokens like
        // `syn::Error::new_spanned` does.
        let start = last_stmt.next().map_or_else(Span::call_site, |t| t.span());
        let end = last_stmt.last().map_or(start, |t| t.span());
        (start, end)
    };

    let mut rt = match config.flavor {
        RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
            tokio::runtime::Builder::new_current_thread()
        },
        RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
            tokio::runtime::Builder::new_multi_thread()
        },
    };
    if let Some(v) = config.worker_threads {
        rt = quote! { #rt.worker_threads(#v) };
    }
    if let Some(v) = config.start_paused {
        rt = quote! { #rt.start_paused(#v) };
    }

    let header = if is_test {
        quote! {
            #[::core::prelude::v1::test]
        }
    } else {
        quote! {}
    };

    let body = &input.block;
    let brace_token = input.block.brace_token;
    let (tail_return, tail_semicolon) = match body.stmts.last() {
        Some(syn::Stmt::Semi(expr, _)) => (
            match expr {
                syn::Expr::Return(_) => quote! { return },
                _ => quote! {},
            },
            quote! {
                ;
            },
        ),
        _ => (quote! {}, quote! {}),
    };
    input.block = syn::parse2(quote_spanned! {last_stmt_end_span=>
        {
            let body = async #body;
            #[allow(clippy::expect_used)]
            #tail_return #rt
                .enable_all()
                .build()
                .expect("Failed building the Runtime")
                .block_on(body)#tail_semicolon
        }
    })
    .expect("Parsing failure");
    input.block.brace_token = brace_token;

    let result = quote! {
        #header
        #input
    };

    result.into()
}

fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream {
    tokens.extend(TokenStream::from(error.into_compile_error()));
    tokens
}

#[cfg(not(test))] // Work around for rust-lang/rust#62127
pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
    // If any of the steps for this macro fail, we still want to expand to an item that is as close
    // to the expected output as possible. This helps out IDEs such that completions and other
    // related features keep working.
    let input: syn::ItemFn = match syn::parse(item.clone()) {
        Ok(it) => it,
        Err(e) => return token_stream_with_error(item, e),
    };

    let config = if input.sig.ident == "main" && !input.sig.inputs.is_empty() {
        let msg = "the main function cannot accept arguments";
        Err(syn::Error::new_spanned(&input.sig.ident, msg))
    } else {
        AttributeArgs::parse_terminated
            .parse(args)
            .and_then(|args| build_config(input.clone(), args, false, rt_multi_thread))
    };

    match config {
        Ok(config) => parse_knobs(input, false, config),
        Err(e) => token_stream_with_error(parse_knobs(input, false, DEFAULT_ERROR_CONFIG), e),
    }
}

pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
    // If any of the steps for this macro fail, we still want to expand to an item that is as close
    // to the expected output as possible. This helps out IDEs such that completions and other
    // related features keep working.
    let input: syn::ItemFn = match syn::parse(item.clone()) {
        Ok(it) => it,
        Err(e) => return token_stream_with_error(item, e),
    };
    let config = if let Some(attr) = input.attrs.iter().find(|attr| attr.path.is_ident("test")) {
        let msg = "second test attribute is supplied";
        Err(syn::Error::new_spanned(&attr, msg))
    } else {
        AttributeArgs::parse_terminated
            .parse(args)
            .and_then(|args| build_config(input.clone(), args, true, rt_multi_thread))
    };

    match config {
        Ok(config) => parse_knobs(input, true, config),
        Err(e) => token_stream_with_error(parse_knobs(input, true, DEFAULT_ERROR_CONFIG), e),
    }
}
<|endoftext|>


<commit_before>// Anima Engine. The quirky game engine
// Copyright (C) 2016  Dragoș Tiselice
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use math::Vector;

/// A `struct` useful for creating square and cubic Bézier curves.
pub struct Bezier {
    v1: Vector,
    v2: Vector,
    v3: Vector,
    v4: Option<Vector>
}

impl Bezier {
    /// Creates a square Bézier from `v1` to `v3` curving towards `v2`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// let b = Bezier::new_sqr(
    ///     Vector::new(0.0, 0.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0),
    ///     Vector::new(2.0, 0.0, 0.0)
    /// );
    /// ```
    pub fn new_sqr(v1: Vector, v2: Vector, v3: Vector) -> Bezier {
        Bezier {
            v1: v1,
            v2: v2,
            v3: v3,
            v4: None
        }
    }

    /// Creates a cubic Bézier from `v1` to `v4` curving towards `v2` and `v3`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// let b = Bezier::new_cub(
    ///     Vector::new(0.0, 0.0, 0.0),
    ///     Vector::new(0.0, 1.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0),
    ///     Vector::new(1.0, 0.0, 0.0)
    /// );
    /// ```
    pub fn new_cub(v1: Vector, v2: Vector, v3: Vector, v4: Vector) -> Bezier {
        Bezier {
            v1: v1,
            v2: v2,
            v3: v3,
            v4: Some(v4)
        }
    }

    /// Computes the vector on a Bézier curve correspoding to a ratio (between `0.0` and `1.0`).
    ///
    /// # Examples
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// let b = Bezier::new_sqr(
    ///     Vector::new(1.0, 0.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0),
    ///     Vector::new(0.0, 1.0, 0.0)
    /// );
    ///
    /// assert_eq!(b.interpolate(0.5), Vector::new(0.75, 0.75, 0.0));
    /// ```
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// let b = Bezier::new_cub(
    ///     Vector::new(0.0, 0.0, 0.0),
    ///     Vector::new(0.0, 1.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0),
    ///     Vector::new(1.0, 0.0, 0.0)
    /// );
    ///
    /// assert_eq!(b.interpolate(0.5), Vector::new(0.5, 0.75, 0.0));
    /// ```
    pub fn interpolate(&self, ratio: f32) -> Vector {
        match self.v4 {
            Some(v4) => {
                self.v1 * (1.0 - ratio).powi(3) +
                self.v2 * 3.0 * (1.0 - ratio).powi(2) * ratio +
                self.v3 * 3.0 * (1.0 - ratio) * ratio.powi(2) +
                v4 * ratio.powi(3)
            },
            None => {
                self.v1 * (1.0 - ratio).powi(2) +
                self.v2 * 2.0 * (1.0 - ratio) * ratio +
                self.v3 * ratio.powi(2)
            }
        }
    }
}
<commit_msg>Added length approximation to Bézier curves.<commit_after>// Anima Engine. The quirky game engine
// Copyright (C) 2016  Dragoș Tiselice
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use math::Vector;

/// A `struct` useful for creating square and cubic Bézier curves.
pub struct Bezier {
    v1: Vector,
    v2: Vector,
    v3: Vector,
    v4: Option<Vector>
}

impl Bezier {
    /// Creates a square Bézier from `v1` to `v3` curving towards `v2`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// let b = Bezier::new_sqr(
    ///     Vector::new(0.0, 0.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0),
    ///     Vector::new(2.0, 0.0, 0.0)
    /// );
    /// ```
    pub fn new_sqr(v1: Vector, v2: Vector, v3: Vector) -> Bezier {
        Bezier {
            v1: v1,
            v2: v2,
            v3: v3,
            v4: None
        }
    }

    /// Creates a cubic Bézier from `v1` to `v4` curving towards `v2` and `v3`.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// let b = Bezier::new_cub(
    ///     Vector::new(0.0, 0.0, 0.0),
    ///     Vector::new(0.0, 1.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0),
    ///     Vector::new(1.0, 0.0, 0.0)
    /// );
    /// ```
    pub fn new_cub(v1: Vector, v2: Vector, v3: Vector, v4: Vector) -> Bezier {
        Bezier {
            v1: v1,
            v2: v2,
            v3: v3,
            v4: Some(v4)
        }
    }

    /// Computes the vector on a Bézier curve correspoding to a ratio (between `0.0` and `1.0`).
    ///
    /// # Examples
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// let b = Bezier::new_sqr(
    ///     Vector::new(1.0, 0.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0),
    ///     Vector::new(0.0, 1.0, 0.0)
    /// );
    ///
    /// assert_eq!(b.interpolate(0.5), Vector::new(0.75, 0.75, 0.0));
    /// ```
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// let b = Bezier::new_cub(
    ///     Vector::new(0.0, 0.0, 0.0),
    ///     Vector::new(0.0, 1.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0),
    ///     Vector::new(1.0, 0.0, 0.0)
    /// );
    ///
    /// assert_eq!(b.interpolate(0.5), Vector::new(0.5, 0.75, 0.0));
    /// ```
    pub fn interpolate(&self, ratio: f32) -> Vector {
        match self.v4 {
            Some(v4) => {
                self.v1 * (1.0 - ratio).powi(3) +
                self.v2 * 3.0 * (1.0 - ratio).powi(2) * ratio +
                self.v3 * 3.0 * (1.0 - ratio) * ratio.powi(2) +
                v4 * ratio.powi(3)
            },
            None => {
                self.v1 * (1.0 - ratio).powi(2) +
                self.v2 * 2.0 * (1.0 - ratio) * ratio +
                self.v3 * ratio.powi(2)
            }
        }
    }

    /// Computes the approximated length of a Bézier curve by summing the distances between `steps`
    /// uniformly distrubuted, consecutive points.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anima::math::Bezier;
    /// # use anima::math::Vector;
    /// # use std::f32::consts;
    /// // approximation of radius 1.0 circle arc
    /// let b = Bezier::new_cub(
    ///     Vector::new(0.0, 0.0, 0.0),
    ///     Vector::new(0.0, 0.55228, 0.0),
    ///     Vector::new(0.44772, 1.0, 0.0),
    ///     Vector::new(1.0, 1.0, 0.0)
    /// );
    ///
    /// const EPSILON: f32 = 0.001;
    ///
    /// assert!((b.length(20) - consts::PI / 2.0).abs() < EPSILON);
    /// ```
    pub fn length(&self, steps: i32) -> f32 {
        let (length, _) = (1..steps + 1).fold((0.0, self.v1), |(l, v), i| {
            let n = self.interpolate((i as f32) / (steps as f32));

            (l + v.dist(n), n)
        });

        length
    }
}
<|endoftext|>


<commit_before>use index::btree::external::ExtNode;
use index::btree::internal::InNode;
use index::btree::node::is_node_locked;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_non_empty;
use index::btree::node::write_targeted;
use index::btree::node::write_unchecked;
use index::btree::node::EmptyNode;
use index::btree::node::Node;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use index::btree::{external, BPlusTree};
use index::lsmtree::tree::LEVEL_PAGE_DIFF_MULTIPLIER;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::collections::BTreeSet;
use std::fmt::Debug;
use std::mem;
use std::sync::atomic::Ordering::Relaxed;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

enum PruningSearch {
    DeepestInnode,
    Innode(NodeCellRef),
}

struct NodeRemoval<'a> {
    empty_pages: Vec<&'a EntryKey>,
    split: Vec<(NodeCellRef, EntryKey)>,
}

impl<'a> NodeRemoval<'a> {
    fn new() -> Self {
        Self {
            empty_pages: vec![],
            split: vec![],
        }
    }
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let mut collected = vec![write_node(node)];
            while collected.len() < LEVEL_PAGE_DIFF_MULTIPLIER {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn merge_innode_remnant<'a, KS, PS>(
    current_node: &mut NodeWriteGuard<KS, PS>,
    prev_key: &'a EntryKey,
    removal: &mut NodeRemoval<'a>,
) where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let curr_right_ref = current_node.right_ref_mut_no_empty().unwrap().clone();
    let curr_innode = current_node.innode_mut();
    let curr_right_bound = &curr_innode.right_bound;
    let curr_last_child = mem::replace(&mut curr_innode.ptrs.as_slice()[0], NodeCellRef::default());
    debug_assert_eq!(curr_innode.len, 0);
    removal.empty_pages.push(prev_key);
    if curr_last_child.is_default() {
        return;
    }
    let mut next_node = write_targeted(write_node::<KS, PS>(&curr_right_ref), curr_right_bound);
    let pos = next_node.search(curr_right_bound);
    let next_innode = next_node.innode_mut();
    next_innode.debug_check_integrity();
    if next_innode.len == KS::slice_len() {
        removal.split.push(next_innode.split_insert(
            curr_right_bound.clone(),
            curr_last_child,
            pos,
            false,
        ))
    } else {
        next_innode.insert_in_place(curr_right_bound.clone(), curr_last_child, pos, false);
    }
}

fn apply_removal<'a, KS, PS>(
    cursor_guard: &mut NodeWriteGuard<KS, PS>,
    poses: &mut BTreeSet<usize>,
    node_removal: &mut NodeRemoval<'a>,
    prev_key: &Option<&'a EntryKey>,
    remove_children_right_nodes: bool,
) where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    if poses.is_empty() {
        return;
    }
    debug!(
        "Applying removal, guard key {:?}, poses {:?}",
        cursor_guard.first_key(),
        poses
    );
    {
        let innode = cursor_guard.innode_mut();
        let mut new_keys = KS::init();
        let mut new_ptrs = PS::init();
        {
            if remove_children_right_nodes {
                // A node will be reclaimed if and only if it have zero references
                // If there is any sequential empty nodes, they have to be unlinked each other
                // The upper level will do the rest to unlink it from the tree
                innode.ptrs.as_slice()[..innode.len + 1]
                    .iter()
                    .enumerate()
                    .filter(|(i, _)| poses.contains(i))
                    .for_each(|(_, r)| {
                        write_node::<KS, PS>(r).right_ref_mut_no_empty();
                    });
            }
            let ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..innode.len + 1]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .map(|(_, p)| p)
                .collect();

            let keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                .iter_mut()
                .enumerate()
                .filter(|(i, _)| !poses.contains(i))
                .take(if ptrs.len() == 0 { 0 } else { ptrs.len() - 1 })
                .map(|(_, k)| k)
                .collect();
            innode.len = keys.len();
            debug!("Prune filtered page have keys {:?}", &keys);
            for (i, key) in keys.into_iter().enumerate() {
                debug_assert!(key > &mut smallvec!(0));
                mem::swap(&mut new_keys.as_slice()[i], key);
            }
            for (i, ptr) in ptrs.into_iter().enumerate() {
                debug_assert!(!ptr.is_default());
                mem::swap(&mut new_ptrs.as_slice()[i], ptr);
            }
        }
        innode.keys = new_keys;
        innode.ptrs = new_ptrs;
        innode.debug_check_integrity();

        debug!(
            "Pruned page have keys {:?}",
            &innode.keys.as_slice_immute()[..innode.len]
        );
    }

    if cursor_guard.is_empty() {
        if let &Some(k) = prev_key {
            debug_assert!(!cursor_guard.is_ext());
            merge_innode_remnant(cursor_guard, k, node_removal);
        }
        debug!("Pruned page is empty: {:?}", prev_key);
        cursor_guard.make_empty_node();
    }
    poses.clear();
}

fn prune_selected<'a, KS, PS>(
    node: &NodeCellRef,
    mut removal: Box<NodeRemoval<'a>>,
    level: usize,
) -> Box<NodeRemoval<'a>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, removal.empty_pages.first().unwrap());
    let pruning = match first_search {
        MutSearchResult::Internal(sub_node) => {
            if read_unchecked::<KS, PS>(&NodeData::<KS, PS>::get_non_empty_node(&sub_node)).is_ext()
            {
                PruningSearch::DeepestInnode
            } else {
                PruningSearch::Innode(sub_node)
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut deepest = false;
    match pruning {
        PruningSearch::DeepestInnode => {
            debug!(
                "Removing in deepest nodes keys {:?}, level {}",
                &removal.empty_pages, level
            );
            deepest = true;
        }
        PruningSearch::Innode(sub_node) => {
            removal = prune_selected::<KS, PS>(&sub_node, removal, level + 1);
        }
    }
    // empty page references that will dealt with by upper level
    let mut upper_removal = NodeRemoval::new();
    if !removal.empty_pages.is_empty() {
        debug!(
            "Pruning page containing keys {:?}, level {}",
            &removal.empty_pages, level
        );
        // start delete
        let mut cursor_guard = write_non_empty(write_node::<KS, PS>(node));
        let mut guard_removing_poses = BTreeSet::new();
        let mut prev_key = None;
        debug!(
            "Prune deepest node starting key: {:?}, level {}",
            cursor_guard.first_key(),
            level
        );
        // iterate over all removal empty nodes
        // in the beginning cursor_guard contains the first key
        // if the empty node go out of the right bound of current cursor guard,
        // it will `apply_removal` and switch to the next right page
        for (_i, &key_to_del) in removal.empty_pages.iter().enumerate() {
            if key_to_del >= cursor_guard.right_bound() {
                debug!(
                    "Applying removal for overflow current page ({}/{}) key: {:?} >= bound: {:?}. guard keys: {:?}, level {}",
                    guard_removing_poses.len(),
                    cursor_guard.len() + 1,
                    key_to_del,
                    cursor_guard.right_bound(),
                    &cursor_guard.keys()[..cursor_guard.len()],
                    level
                );
                apply_removal(
                    &mut cursor_guard,
                    &mut guard_removing_poses,
                    &mut upper_removal,
                    &prev_key,
                    !deepest,
                );
                cursor_guard = write_targeted(cursor_guard, key_to_del);
                debug_assert!(!cursor_guard.is_empty_node());
                debug!(
                    "Applied removal for overflow current page ({}), level {}",
                    cursor_guard.len(),
                    level
                );
            }
            let pos = cursor_guard.search(key_to_del);
            debug!(
                "Key to delete have position {}, key: {:?}, level {}",
                pos, key_to_del, level
            );
            guard_removing_poses.insert(pos);
            prev_key = Some(key_to_del)
        }
        if !guard_removing_poses.is_empty() {
            debug!(
                "Applying removal for last keys {:?}, level {}",
                &guard_removing_poses, level
            );
            apply_removal(
                &mut cursor_guard,
                &mut guard_removing_poses,
                &mut upper_removal,
                &prev_key,
                !deepest,
            );
        }
    }
    if !removal.split.is_empty() {
        let mut cursor_guard = write_node::<KS, PS>(node);
        for (node, pivot) in &removal.split {
            cursor_guard = write_targeted(cursor_guard, &pivot);
            let pos = cursor_guard.search(&pivot);
            let innode = cursor_guard.innode_mut();
            let pivot = pivot.clone();
            let node = node.clone();
            if innode.len >= KS::slice_len() {
                upper_removal
                    .split
                    .push(innode.split_insert(pivot, node, pos, true));
            } else {
                innode.insert_in_place(pivot, node, pos, true)
            }
        }
    }
    debug!(
        "Have empty nodes {:?}, level {:?}",
        &upper_removal.empty_pages, level
    );
    box upper_removal
}

pub fn level_merge<KS, PS>(src_tree: &BPlusTree<KS, PS>, dest_tree: &LevelTree) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut left_most_leaf_guards = select::<KS, PS>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    let mut num_keys_removed = 0;
    let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();

    debug!("Merge selected {} pages", left_most_leaf_guards.len());
    debug_assert_eq!(left_most_id, src_tree.head_page_id, "{}", left_most_leaf_guards.first().unwrap().type_name());

    // merge to dest_tree
    {
        let mut deleted_keys = src_tree.deleted.write();
        let mut merged_deleted_keys = vec![];
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .filter(|&g| !g.is_empty_node())
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .filter(|&k| {
                if deleted_keys.contains(k) {
                    merged_deleted_keys.push(k.clone());
                    false
                } else {
                    true
                }
            })
            .cloned()
            .collect_vec();
        num_keys_removed = keys.len();
        debug!("Merge selected keys are {:?}", &keys);
        dest_tree.merge_with_keys(box keys);
        for rk in &merged_deleted_keys {
            deleted_keys.remove(rk);
        }
    }

    // cleanup upper level references
    {
        let page_keys = left_most_leaf_guards
            .iter()
            .filter(|g| !g.is_empty())
            .map(|g| g.first_key())
            .collect_vec();
        let mut removal = NodeRemoval::new();
        removal.empty_pages = page_keys;
        prune_selected::<KS, PS>(&src_tree.get_root(), box removal, 0);
    }

    // adjust leaf left, right references
    {
        let right_right_most = left_most_leaf_guards
            .last()
            .unwrap()
            .right_ref()
            .unwrap()
            .clone();

        let left_left_most = left_most_leaf_guards.first().unwrap().left_ref().unwrap().clone();

        debug_assert!(read_unchecked::<KS, PS>(&left_left_most).is_none());

        for mut g in &mut left_most_leaf_guards {
            external::make_deleted(&g.ext_id());
            **g = NodeData::Empty(box EmptyNode {
                left: Some(left_left_most.clone()),
                right: right_right_most.clone(),
            });
        }

        let mut new_first_node = write_node::<KS, PS>(&right_right_most);
        let mut new_first_node_ext = new_first_node.extnode_mut();
        new_first_node_ext.id = left_most_id;
        new_first_node_ext.prev = left_left_most;

        debug_assert_eq!(new_first_node_ext.id, src_tree.head_page_id);

        ExtNode::<KS, PS>::make_changed(&right_right_most, src_tree);
    }

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<commit_msg>prune by scanning<commit_after>use index::btree::external::ExtNode;
use index::btree::internal::InNode;
use index::btree::node::is_node_locked;
use index::btree::node::read_node;
use index::btree::node::read_unchecked;
use index::btree::node::write_node;
use index::btree::node::write_non_empty;
use index::btree::node::write_targeted;
use index::btree::node::write_unchecked;
use index::btree::node::EmptyNode;
use index::btree::node::Node;
use index::btree::node::NodeData;
use index::btree::node::NodeWriteGuard;
use index::btree::search::mut_search;
use index::btree::search::MutSearchResult;
use index::btree::LevelTree;
use index::btree::NodeCellRef;
use index::btree::{external, BPlusTree};
use index::lsmtree::tree::LEVEL_PAGE_DIFF_MULTIPLIER;
use index::EntryKey;
use index::Slice;
use itertools::Itertools;
use smallvec::SmallVec;
use std::collections::{BTreeSet, HashSet};
use std::fmt::Debug;
use std::mem;
use std::sync::atomic::Ordering::Relaxed;

enum Selection<KS, PS>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    Selected(Vec<NodeWriteGuard<KS, PS>>),
    Innode(NodeCellRef),
}

fn select<KS, PS>(node: &NodeCellRef) -> Vec<NodeWriteGuard<KS, PS>>
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let search = mut_search::<KS, PS>(node, &smallvec!());
    match search {
        MutSearchResult::External => {
            let mut collected = vec![write_node(node)];
            while collected.len() < LEVEL_PAGE_DIFF_MULTIPLIER {
                let right = write_node(
                    collected
                        .last_mut()
                        .unwrap()
                        .right_ref_mut_no_empty()
                        .unwrap(),
                );
                if right.is_none() {
                    break;
                } else {
                    collected.push(right);
                }
            }
            return collected;
        }
        MutSearchResult::Internal(node) => select::<KS, PS>(&node),
    }
}

fn prune_selected<'a, KS, PS>(node: &NodeCellRef, bound: &EntryKey, level: usize)
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let first_search = mut_search::<KS, PS>(node, &smallvec!());
    match first_search {
        MutSearchResult::Internal(sub_node_ref) => {
            let sub_node = read_unchecked::<KS, PS>(&sub_node_ref);
            if !sub_node.is_empty_node() && !sub_node.is_ext() {
                prune_selected::<KS, PS>(&sub_node_ref, bound, level + 1);
            }
        }
        MutSearchResult::External => unreachable!(),
    };
    let mut node_ref = node.clone();
    loop {
        let mut page = write_node::<KS, PS>(&node);
        if page.is_empty_node() {
            node_ref = page.right_ref().unwrap().clone();
        } else {
            let mut ptr_indices: HashSet<_> = page.innode()
                .ptrs
                .as_slice_immute()
                .iter()
                .enumerate()
                .filter_map(|(i, sub_level)| {
                    if read_unchecked::<KS, PS>(sub_level).is_empty_node() {
                        Some(i)
                    } else {
                        None
                    }
                })
                .collect();
            if ptr_indices.len() == page.len() + 1 {
                // all sub nodes are empty
                // will set current node empty either
                let next_page_ref = page.innode().right.clone();
                *page = NodeData::Empty(box EmptyNode {
                    left: None,
                    right: next_page_ref,
                });
            } else {
                let mut innode = page.innode_mut();
                let mut new_keys = KS::init();
                let mut new_ptrs = PS::init();
                {
                    let ptrs: Vec<&mut _> = innode.ptrs.as_slice()[..innode.len + 1]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&i))
                        .map(|(_, p)| p)
                        .collect();

                    let keys: Vec<&mut _> = innode.keys.as_slice()[..innode.len]
                        .iter_mut()
                        .enumerate()
                        .filter(|(i, _)| !ptr_indices.contains(&(i + 1)))
                        .map(|(_, k)| k)
                        .collect();
                    debug!("Prune filtered page have keys {:?}", &keys);
                    innode.len = keys.len();
                    for (i, key) in keys.into_iter().enumerate() {
                        debug_assert!(key > &mut smallvec!(0));
                        mem::swap(&mut new_keys.as_slice()[i], key);
                    }
                    for (i, ptr) in ptrs.into_iter().enumerate() {
                        debug_assert!(!ptr.is_default());
                        mem::swap(&mut new_ptrs.as_slice()[i], ptr);
                    }
                }
                innode.keys = new_keys;
                innode.ptrs = new_ptrs;
            }
        }
    }
}

pub fn level_merge<KS, PS>(src_tree: &BPlusTree<KS, PS>, dest_tree: &LevelTree) -> usize
where
    KS: Slice<EntryKey> + Debug + 'static,
    PS: Slice<NodeCellRef> + 'static,
{
    let mut left_most_leaf_guards = select::<KS, PS>(&src_tree.get_root());
    let merge_page_len = left_most_leaf_guards.len();
    let mut num_keys_removed = 0;
    let left_most_id = left_most_leaf_guards.first().unwrap().ext_id();
    let prune_bound = left_most_leaf_guards.last().unwrap().right_bound().clone();

    debug!("Merge selected {} pages", left_most_leaf_guards.len());
    debug_assert_eq!(
        left_most_id,
        src_tree.head_page_id,
        "{}",
        left_most_leaf_guards.first().unwrap().type_name()
    );

    // merge to dest_tree
    {
        let mut deleted_keys = src_tree.deleted.write();
        let mut merged_deleted_keys = vec![];
        let keys: Vec<EntryKey> = left_most_leaf_guards
            .iter()
            .filter(|&g| !g.is_empty_node())
            .map(|g| &g.keys()[..g.len()])
            .flatten()
            .filter(|&k| {
                if deleted_keys.contains(k) {
                    merged_deleted_keys.push(k.clone());
                    false
                } else {
                    true
                }
            })
            .cloned()
            .collect_vec();
        num_keys_removed = keys.len();
        debug!("Merge selected keys are {:?}", &keys);
        dest_tree.merge_with_keys(box keys);
        for rk in &merged_deleted_keys {
            deleted_keys.remove(rk);
        }
    }

    // adjust leaf left, right references
    {
        let right_right_most = left_most_leaf_guards
            .last()
            .unwrap()
            .right_ref()
            .unwrap()
            .clone();

        let left_left_most = left_most_leaf_guards
            .first()
            .unwrap()
            .left_ref()
            .unwrap()
            .clone();

        debug_assert!(read_unchecked::<KS, PS>(&left_left_most).is_none());

        for mut g in &mut left_most_leaf_guards {
            external::make_deleted(&g.ext_id());
            **g = NodeData::Empty(box EmptyNode {
                left: Some(left_left_most.clone()),
                right: right_right_most.clone(),
            });
        }

        let mut new_first_node = write_node::<KS, PS>(&right_right_most);
        let mut new_first_node_ext = new_first_node.extnode_mut();
        new_first_node_ext.id = left_most_id;
        new_first_node_ext.prev = left_left_most;

        debug_assert_eq!(new_first_node_ext.id, src_tree.head_page_id);

        ExtNode::<KS, PS>::make_changed(&right_right_most, src_tree);
    }

    // cleanup upper level references
    {}

    src_tree.len.fetch_sub(num_keys_removed, Relaxed);

    merge_page_len
}
<|endoftext|>


<commit_before>//! Defines types found in various posix_types.h files.
//!
//! If unsure, check include/uapi/asm-generic/posix_types.h **first**, then /arch/*/include/uapi/asm/posix_types.h
//!
//! For the best result, use http://lxr.free-electrons.com
//!
//! Due for a complete makeover for 0.1

#[cfg(any(target_arch="sparc", target_arch="parisc"))]
///A type that defines microseconds. Since you're using either SPARC or (pa)RISC, it is a c_int
pub type suseconds_t = ::c_int;

#[cfg(not(any(target_arch="sparc", target_arch="parisc")))]
///A type that defines microseconds. You're using a "regular" processor, so it is a c_long
pub type suseconds_t = ::c_long;

///A type that defines a clock ID. As of 4.1, this is always c_int
pub type clockid_t = ::c_int;

///A type that defines a Process or Thread Group ID
pub type pid_t = ::c_int;

///A type used for new uids
pub type __kernel_uid32_t = ::c_uint;

///A type used for new gids
pub type __kernel_gid32_t = ::c_uint;

///A type that defines a user ID
pub type uid_t = __kernel_uid32_t;

///A type that defines a group ID
pub type gid_t = __kernel_uid32_t;

#[cfg(not(any(target_arch = "cris",
              target_arch = "m68k",
              target_arch = "sh",
              target_arch = "arm",
              target_arch = "sparc",
              target_arch = "frv",
              target_arch = "mn10300",
              target_arch = "x86",
              target_arch = "x86_64",
              target_arch = "m32r")))]
pub type __kernel_uid_t = ::c_uint;

#[cfg(any(target_arch = "cris",
         target_arch = "m68k",
         target_arch = "sh",
         target_arch = "arm",
         target_arch = "sparc",
         target_arch = "frv",
         target_arch = "mn10300",
         target_arch = "x86",
         target_arch = "x86_64",
         target_arch = "m32r"))]
pub type __kernel_uid_t = ::c_ushort;

#[cfg(not(any(target_arch = "cris",
              target_arch = "m68k",
              target_arch = "sh",
              target_arch = "arm",
              target_arch = "sparc",
              target_arch = "frv",
              target_arch = "mn10300",
              target_arch = "x86",
              target_arch = "x86_64",
              target_arch = "m32r")))]
pub type __kernel_gid_t = ::c_uint;

#[cfg(any(target_arch = "cris",
         target_arch = "m68k",
         target_arch = "sh",
         target_arch = "arm",
         target_arch = "sparc",
         target_arch = "frv",
         target_arch = "mn10300",
         target_arch = "x86",
         target_arch = "x86_64",
         target_arch = "m32r"))]
pub type __kernel_gid_t = ::c_ushort;

#[cfg(not(any(target_arch = "cris",
              target_arch = "parisc",
              target_arch = "microblaze",
              target_arch = "m68k",
              target_arch = "sh",
              target_arch = "arm",
              target_arch = "avr32",
              target_arch = "sparc",
              target_arch = "frv",
              target_arch = "blackfin",
              target_arch = "mn10300",
              target_arch = "x86",
              target_arch = "x86_64",
              target_arch = "m32r")))]
pub type __kernel_mode_t = ::c_uint;

#[cfg(any(target_arch = "cris",
         target_arch = "parisc",
         target_arch = "microblaze",
         target_arch = "m68k",
         target_arch = "sh",
         target_arch = "arm",
         target_arch = "avr32",
         target_arch = "sparc",
         target_arch = "frv",
         target_arch = "blackfin",
         target_arch = "mn10300",
         target_arch = "x86",
         target_arch = "x86_64",
         target_arch = "m32r"))]
pub type __kernel_mode_t = ::c_ushort;

pub type __kernel_key_t = ::c_int;

pub type key_t = __kernel_key_t;

pub type mode_t = __kernel_mode_t;

#[cfg(target_arch="x86")]
///A kernel type for long, set to c_longlong snice you're on x86
pub type __kernel_long_t = ::c_longlong;

#[cfg(not(target_arch="x86"))]
///A kernel type for long, set to c_long since you're not on x86
pub type __kernel_long_t = ::c_long;

#[cfg(target_arch="x86")]
///A kernel type for ulong, set to c_ulonglong snice you're on x86
pub type __kernel_ulong_t = ::c_ulonglong;

#[cfg(not(target_arch="x86"))]
///A kernel type for ulong, set to c_ulong since you're not on x86
pub type __kernel_ulong_t = ::c_ulong;<commit_msg>Don't re-export mode_t yet<commit_after>//! Defines types found in various posix_types.h files.
//!
//! If unsure, check include/uapi/asm-generic/posix_types.h **first**, then /arch/*/include/uapi/asm/posix_types.h
//!
//! For the best result, use http://lxr.free-electrons.com
//!
//! Due for a complete makeover for 0.1

#[cfg(any(target_arch="sparc", target_arch="parisc"))]
///A type that defines microseconds. Since you're using either SPARC or (pa)RISC, it is a c_int
pub type suseconds_t = ::c_int;

#[cfg(not(any(target_arch="sparc", target_arch="parisc")))]
///A type that defines microseconds. You're using a "regular" processor, so it is a c_long
pub type suseconds_t = ::c_long;

///A type that defines a clock ID. As of 4.1, this is always c_int
pub type clockid_t = ::c_int;

///A type that defines a Process or Thread Group ID
pub type pid_t = ::c_int;

///A type used for new uids
pub type __kernel_uid32_t = ::c_uint;

///A type used for new gids
pub type __kernel_gid32_t = ::c_uint;

///A type that defines a user ID
pub type uid_t = __kernel_uid32_t;

///A type that defines a group ID
pub type gid_t = __kernel_uid32_t;

#[cfg(not(any(target_arch = "cris",
              target_arch = "m68k",
              target_arch = "sh",
              target_arch = "arm",
              target_arch = "sparc",
              target_arch = "frv",
              target_arch = "mn10300",
              target_arch = "x86",
              target_arch = "x86_64",
              target_arch = "m32r")))]
pub type __kernel_uid_t = ::c_uint;

#[cfg(any(target_arch = "cris",
         target_arch = "m68k",
         target_arch = "sh",
         target_arch = "arm",
         target_arch = "sparc",
         target_arch = "frv",
         target_arch = "mn10300",
         target_arch = "x86",
         target_arch = "x86_64",
         target_arch = "m32r"))]
pub type __kernel_uid_t = ::c_ushort;

#[cfg(not(any(target_arch = "cris",
              target_arch = "m68k",
              target_arch = "sh",
              target_arch = "arm",
              target_arch = "sparc",
              target_arch = "frv",
              target_arch = "mn10300",
              target_arch = "x86",
              target_arch = "x86_64",
              target_arch = "m32r")))]
pub type __kernel_gid_t = ::c_uint;

#[cfg(any(target_arch = "cris",
         target_arch = "m68k",
         target_arch = "sh",
         target_arch = "arm",
         target_arch = "sparc",
         target_arch = "frv",
         target_arch = "mn10300",
         target_arch = "x86",
         target_arch = "x86_64",
         target_arch = "m32r"))]
pub type __kernel_gid_t = ::c_ushort;

#[cfg(not(any(target_arch = "cris",
              target_arch = "parisc",
              target_arch = "microblaze",
              target_arch = "m68k",
              target_arch = "sh",
              target_arch = "arm",
              target_arch = "avr32",
              target_arch = "sparc",
              target_arch = "frv",
              target_arch = "blackfin",
              target_arch = "mn10300",
              target_arch = "x86",
              target_arch = "x86_64",
              target_arch = "m32r")))]
pub type __kernel_mode_t = ::c_uint;

#[cfg(any(target_arch = "cris",
         target_arch = "parisc",
         target_arch = "microblaze",
         target_arch = "m68k",
         target_arch = "sh",
         target_arch = "arm",
         target_arch = "avr32",
         target_arch = "sparc",
         target_arch = "frv",
         target_arch = "blackfin",
         target_arch = "mn10300",
         target_arch = "x86",
         target_arch = "x86_64",
         target_arch = "m32r"))]
pub type __kernel_mode_t = ::c_ushort;

pub type __kernel_key_t = ::c_int;

pub type key_t = __kernel_key_t;

//pub type mode_t = __kernel_mode_t; TODO: Re-export this in Rust

#[cfg(target_arch="x86")]
///A kernel type for long, set to c_longlong snice you're on x86
pub type __kernel_long_t = ::c_longlong;

#[cfg(not(target_arch="x86"))]
///A kernel type for long, set to c_long since you're not on x86
pub type __kernel_long_t = ::c_long;

#[cfg(target_arch="x86")]
///A kernel type for ulong, set to c_ulonglong snice you're on x86
pub type __kernel_ulong_t = ::c_ulonglong;

#[cfg(not(target_arch="x86"))]
///A kernel type for ulong, set to c_ulong since you're not on x86
pub type __kernel_ulong_t = ::c_ulong;<|endoftext|>


<commit_before>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use futures::future;
use futures::future::FutureResult;
use httparse;
use hyper;
use hyper::{Method, StatusCode, Headers};
use hyper::header::{IfNoneMatch, ETag, EntityTag};
use hyper::server::{Request, Response};
use farmhash;

use super::header::ProxyHeader;
use super::tunnel::ProxyTunnelBuilder;
use header::janitor::HeaderJanitor;
use header::request_shard::HeaderRequestBloomRequestShard;
use header::status::{HeaderBloomStatus, HeaderBloomStatusValue};
use cache::read::CacheRead;
use cache::write::CacheWrite;
use cache::route::CacheRoute;
use LINE_FEED;

pub struct ProxyServeBuilder;
pub struct ProxyServe;

const CACHED_PARSE_MAX_HEADERS: usize = 100;

pub type ProxyServeFuture = FutureResult<Response, hyper::Error>;

impl ProxyServeBuilder {
    pub fn new() -> ProxyServe {
        ProxyServe {}
    }
}

impl ProxyServe {
    pub fn ensure(&self) -> Result<(), ()> {
        Ok(())
    }

    pub fn handle(&self, req: Request) -> ProxyServeFuture {
        info!("handled request: {} on {}", req.method(), req.path());

        if req.headers().has::<HeaderRequestBloomRequestShard>() == true {
            match *req.method() {
                Method::Options | Method::Head | Method::Get | Method::Post | Method::Patch |
                Method::Put | Method::Delete => self.accept(req),
                _ => self.reject(req, StatusCode::MethodNotAllowed),
            }
        } else {
            self.reject(req, StatusCode::NotExtended)
        }
    }

    fn accept(&self, req: Request) -> ProxyServeFuture {
        self.tunnel(req)
    }

    fn reject(&self, req: Request, status: StatusCode) -> ProxyServeFuture {
        let mut headers = Headers::new();

        headers.set::<HeaderBloomStatus>(HeaderBloomStatus(HeaderBloomStatusValue::Reject));

        self.respond(&req.method(), status, headers, format!("{}", status))
    }

    fn tunnel(&self, req: Request) -> ProxyServeFuture {
        let (method, uri, version, headers, body) = req.deconstruct();
        let (auth, shard) = ProxyHeader::parse_from_request(&headers);

        let ns = CacheRoute::gen_ns(shard, auth, version, &method, uri.path(), uri.query());

        info!("tunneling for ns = {}", ns);

        match CacheRead::acquire(ns.as_ref()) {
            Ok(cached_value) => self.dispatch_cached(&method, &headers, &cached_value),
            Err(_) => {
                match ProxyTunnelBuilder::new().run(&method, &uri, &headers, body, shard) {
                    Ok(tunnel_res) => {
                        let ref status = tunnel_res.status();

                        // TODO: maybe avoid cloning?
                        let headers = tunnel_res.headers().clone();

                        let result = CacheWrite::save(
                            ns.as_ref(),
                            &method,
                            &version,
                            status,
                            &headers,
                            tunnel_res.body(),
                        );

                        match result.body {
                            Ok(body_string) => {
                                self.dispatch_fetched(
                                    &method,
                                    status,
                                    headers,
                                    HeaderBloomStatusValue::Miss,
                                    body_string,
                                    result.value,
                                )
                            }
                            Err(body_string_values) => {
                                match body_string_values {
                                    Some(body_string) => {
                                        self.dispatch_fetched(
                                            &method,
                                            status,
                                            headers,
                                            HeaderBloomStatusValue::Direct,
                                            body_string,
                                            result.value,
                                        )
                                    }
                                    _ => self.dispatch_failure(&method),
                                }
                            }
                        }
                    }
                    _ => self.dispatch_failure(&method),
                }
            }
        }
    }

    fn dispatch_cached(
        &self,
        method: &Method,
        headers: &Headers,
        res_string: &str,
    ) -> ProxyServeFuture {
        // Process ETag for cached content
        let (res_hash, res_etag) = self.body_fingerprint(res_string);

        let isnt_modified = match headers.get::<IfNoneMatch>() {
            Some(req_if_none_match) => {
                (*req_if_none_match == IfNoneMatch::Items(vec![EntityTag::new(false, res_hash)]))
            }
            _ => false,
        };

        // Response not modified for client?
        if isnt_modified == true {
            // Process non-modified + cached headers
            let mut headers = Headers::new();

            headers.set::<ETag>(res_etag);
            headers.set::<HeaderBloomStatus>(HeaderBloomStatus(HeaderBloomStatusValue::Hit));

            // Serve non-modified response
            return self.respond(method, StatusCode::NotModified, headers, String::from(""));
        }

        // Response modified
        let mut headers = [httparse::EMPTY_HEADER; CACHED_PARSE_MAX_HEADERS];
        let mut res = httparse::Response::new(&mut headers);

        // Split headers from body
        let mut res_body_string = String::new();
        let mut is_last_line_empty = false;

        for res_line in res_string.lines() {
            if res_body_string.is_empty() == false || is_last_line_empty == true {
                // Write to body
                res_body_string.push_str(res_line.as_ref());
                res_body_string.push_str(LINE_FEED);
            }

            is_last_line_empty = res_line.is_empty();
        }

        match res.parse(res_string.as_bytes()) {
            Ok(_) => {
                // Process cached status
                let code = res.code.unwrap_or(500u16);
                let status = StatusCode::try_from(code).unwrap_or(StatusCode::Unregistered(code));

                // Process cached headers
                let mut headers = Headers::new();

                for header in res.headers {
                    headers.set_raw(
                        String::from_utf8(Vec::from(header.name)).unwrap(),
                        String::from_utf8(Vec::from(header.value)).unwrap(),
                    );
                }

                headers.set::<ETag>(res_etag);
                headers.set::<HeaderBloomStatus>(HeaderBloomStatus(HeaderBloomStatusValue::Hit));

                // Serve cached response
                self.respond(method, status, headers, res_body_string)
            }
            Err(err) => {
                error!("failed parsing cached response: {}", err);

                self.dispatch_failure(method)
            }
        }
    }

    fn dispatch_fetched(
        &self,
        method: &Method,
        status: &StatusCode,
        mut headers: Headers,
        bloom_status: HeaderBloomStatusValue,
        body_string: String,
        result_string: Option<String>,
    ) -> ProxyServeFuture {
        // Map headers to clean-up
        let mut headers_remove: Vec<String> = Vec::new();

        for header_view in headers.iter() {
            // Do not forward contextual and internal headers \
            //   (ie. 'Bloom-Response-*' headers)
            if HeaderJanitor::is_contextual(&header_view) == true ||
                HeaderJanitor::is_internal(&header_view) == true
            {
                headers_remove.push(String::from(header_view.name()));
            }
        }

        // Proceed headers clean-up
        for header_remove in headers_remove.iter() {
            headers.remove_raw(header_remove.as_ref());
        }

        // Process ETag for content?
        if result_string.is_some() == true {
            let (_, res_etag) = self.body_fingerprint(&result_string.unwrap());

            headers.set::<ETag>(res_etag);
        }

        headers.set(HeaderBloomStatus(bloom_status));

        self.respond(method, *status, headers, body_string)
    }

    fn dispatch_failure(&self, method: &Method) -> ProxyServeFuture {
        let status = StatusCode::BadGateway;

        let mut headers = Headers::new();

        headers.set::<HeaderBloomStatus>(HeaderBloomStatus(HeaderBloomStatusValue::Offline));

        self.respond(method, status, headers, format!("{}", status))
    }

    fn body_fingerprint(&self, body_string: &str) -> (String, ETag) {
        let body_hash = format!("{:x}", farmhash::fingerprint64(body_string.as_bytes()));
        let body_etag = ETag(EntityTag::new(false, body_hash.clone()));

        (body_hash, body_etag)
    }

    fn respond(
        &self,
        method: &Method,
        status: StatusCode,
        headers: Headers,
        mut body_string: String,
    ) -> ProxyServeFuture {
        let res = match method {
            &Method::Get | &Method::Post | &Method::Patch | &Method::Put => {
                // Ensure body string ends w/ a new line in any case, this \
                //   fixes an 'infinite loop' issue w/ Hyper
                if body_string.ends_with(LINE_FEED) == false {
                    body_string.push_str(LINE_FEED);
                }

                Response::new()
                    .with_status(status)
                    .with_headers(headers)
                    .with_body(body_string)
            }
            _ => Response::new().with_status(status).with_headers(headers),
        };

        future::ok(res)
    }
}
<commit_msg>Fix<commit_after>// Bloom
//
// HTTP REST API caching middleware
// Copyright: 2017, Valerian Saliou <valerian@valeriansaliou.name>
// License: Mozilla Public License v2.0 (MPL v2.0)

use futures::future;
use futures::future::FutureResult;
use httparse;
use hyper;
use hyper::{Method, StatusCode, Headers};
use hyper::header::{IfNoneMatch, ETag, EntityTag};
use hyper::server::{Request, Response};
use farmhash;

use super::header::ProxyHeader;
use super::tunnel::ProxyTunnelBuilder;
use header::janitor::HeaderJanitor;
use header::request_shard::HeaderRequestBloomRequestShard;
use header::status::{HeaderBloomStatus, HeaderBloomStatusValue};
use cache::read::CacheRead;
use cache::write::CacheWrite;
use cache::route::CacheRoute;
use LINE_FEED;

pub struct ProxyServeBuilder;
pub struct ProxyServe;

const CACHED_PARSE_MAX_HEADERS: usize = 100;

pub type ProxyServeFuture = FutureResult<Response, hyper::Error>;

impl ProxyServeBuilder {
    pub fn new() -> ProxyServe {
        ProxyServe {}
    }
}

impl ProxyServe {
    pub fn ensure(&self) -> Result<(), ()> {
        Ok(())
    }

    pub fn handle(&self, req: Request) -> ProxyServeFuture {
        info!("handled request: {} on {}", req.method(), req.path());

        if req.headers().has::<HeaderRequestBloomRequestShard>() == true {
            match *req.method() {
                Method::Options | Method::Head | Method::Get | Method::Post | Method::Patch |
                Method::Put | Method::Delete => self.accept(req),
                _ => self.reject(req, StatusCode::MethodNotAllowed),
            }
        } else {
            self.reject(req, StatusCode::NotExtended)
        }
    }

    fn accept(&self, req: Request) -> ProxyServeFuture {
        self.tunnel(req)
    }

    fn reject(&self, req: Request, status: StatusCode) -> ProxyServeFuture {
        let mut headers = Headers::new();

        headers.set::<HeaderBloomStatus>(HeaderBloomStatus(HeaderBloomStatusValue::Reject));

        self.respond(&req.method(), status, headers, format!("{}", status))
    }

    fn tunnel(&self, req: Request) -> ProxyServeFuture {
        let (method, uri, version, headers, body) = req.deconstruct();
        let (auth, shard) = ProxyHeader::parse_from_request(&headers);

        let ns = CacheRoute::gen_ns(shard, auth, version, &method, uri.path(), uri.query());

        info!("tunneling for ns = {}", ns);

        match CacheRead::acquire(ns.as_ref()) {
            Ok(cached_value) => self.dispatch_cached(&method, &headers, &cached_value),
            Err(_) => {
                match ProxyTunnelBuilder::new().run(&method, &uri, &headers, body, shard) {
                    Ok(tunnel_res) => {
                        let ref status = tunnel_res.status();
                        let headers = tunnel_res.headers().clone();

                        let result = CacheWrite::save(
                            ns.as_ref(),
                            &method,
                            &version,
                            status,
                            &headers,
                            tunnel_res.body(),
                        );

                        match result.body {
                            Ok(body_string) => {
                                self.dispatch_fetched(
                                    &method,
                                    status,
                                    headers,
                                    HeaderBloomStatusValue::Miss,
                                    body_string,
                                    result.value,
                                )
                            }
                            Err(body_string_values) => {
                                match body_string_values {
                                    Some(body_string) => {
                                        self.dispatch_fetched(
                                            &method,
                                            status,
                                            headers,
                                            HeaderBloomStatusValue::Direct,
                                            body_string,
                                            result.value,
                                        )
                                    }
                                    _ => self.dispatch_failure(&method),
                                }
                            }
                        }
                    }
                    _ => self.dispatch_failure(&method),
                }
            }
        }
    }

    fn dispatch_cached(
        &self,
        method: &Method,
        headers: &Headers,
        res_string: &str,
    ) -> ProxyServeFuture {
        // Process ETag for cached content
        let (res_hash, res_etag) = self.body_fingerprint(res_string);

        let isnt_modified = match headers.get::<IfNoneMatch>() {
            Some(req_if_none_match) => {
                (*req_if_none_match == IfNoneMatch::Items(vec![EntityTag::new(false, res_hash)]))
            }
            _ => false,
        };

        // Response not modified for client?
        if isnt_modified == true {
            // Process non-modified + cached headers
            let mut headers = Headers::new();

            headers.set::<ETag>(res_etag);
            headers.set::<HeaderBloomStatus>(HeaderBloomStatus(HeaderBloomStatusValue::Hit));

            // Serve non-modified response
            return self.respond(method, StatusCode::NotModified, headers, String::from(""));
        }

        // Response modified
        let mut headers = [httparse::EMPTY_HEADER; CACHED_PARSE_MAX_HEADERS];
        let mut res = httparse::Response::new(&mut headers);

        // Split headers from body
        let mut res_body_string = String::new();
        let mut is_last_line_empty = false;

        for res_line in res_string.lines() {
            if res_body_string.is_empty() == false || is_last_line_empty == true {
                // Write to body
                res_body_string.push_str(res_line.as_ref());
                res_body_string.push_str(LINE_FEED);
            }

            is_last_line_empty = res_line.is_empty();
        }

        match res.parse(res_string.as_bytes()) {
            Ok(_) => {
                // Process cached status
                let code = res.code.unwrap_or(500u16);
                let status = StatusCode::try_from(code).unwrap_or(StatusCode::Unregistered(code));

                // Process cached headers
                let mut headers = Headers::new();

                for header in res.headers {
                    headers.set_raw(
                        String::from_utf8(Vec::from(header.name)).unwrap(),
                        String::from_utf8(Vec::from(header.value)).unwrap(),
                    );
                }

                headers.set::<ETag>(res_etag);
                headers.set::<HeaderBloomStatus>(HeaderBloomStatus(HeaderBloomStatusValue::Hit));

                // Serve cached response
                self.respond(method, status, headers, res_body_string)
            }
            Err(err) => {
                error!("failed parsing cached response: {}", err);

                self.dispatch_failure(method)
            }
        }
    }

    fn dispatch_fetched(
        &self,
        method: &Method,
        status: &StatusCode,
        mut headers: Headers,
        bloom_status: HeaderBloomStatusValue,
        body_string: String,
        result_string: Option<String>,
    ) -> ProxyServeFuture {
        // Map headers to clean-up
        let mut headers_remove: Vec<String> = Vec::new();

        for header_view in headers.iter() {
            // Do not forward contextual and internal headers \
            //   (ie. 'Bloom-Response-*' headers)
            if HeaderJanitor::is_contextual(&header_view) == true ||
                HeaderJanitor::is_internal(&header_view) == true
            {
                headers_remove.push(String::from(header_view.name()));
            }
        }

        // Proceed headers clean-up
        for header_remove in headers_remove.iter() {
            headers.remove_raw(header_remove.as_ref());
        }

        // Process ETag for content?
        if result_string.is_some() == true {
            let (_, res_etag) = self.body_fingerprint(&result_string.unwrap());

            headers.set::<ETag>(res_etag);
        }

        headers.set(HeaderBloomStatus(bloom_status));

        self.respond(method, *status, headers, body_string)
    }

    fn dispatch_failure(&self, method: &Method) -> ProxyServeFuture {
        let status = StatusCode::BadGateway;

        let mut headers = Headers::new();

        headers.set::<HeaderBloomStatus>(HeaderBloomStatus(HeaderBloomStatusValue::Offline));

        self.respond(method, status, headers, format!("{}", status))
    }

    fn body_fingerprint(&self, body_string: &str) -> (String, ETag) {
        let body_hash = format!("{:x}", farmhash::fingerprint64(body_string.as_bytes()));
        let body_etag = ETag(EntityTag::new(false, body_hash.clone()));

        (body_hash, body_etag)
    }

    fn respond(
        &self,
        method: &Method,
        status: StatusCode,
        headers: Headers,
        mut body_string: String,
    ) -> ProxyServeFuture {
        let res = match method {
            &Method::Get | &Method::Post | &Method::Patch | &Method::Put => {
                // Ensure body string ends w/ a new line in any case, this \
                //   fixes an 'infinite loop' issue w/ Hyper
                if body_string.ends_with(LINE_FEED) == false {
                    body_string.push_str(LINE_FEED);
                }

                Response::new()
                    .with_status(status)
                    .with_headers(headers)
                    .with_body(body_string)
            }
            _ => Response::new().with_status(status).with_headers(headers),
        };

        future::ok(res)
    }
}
<|endoftext|>


<commit_before>use nalgebra::na::{AlgebraicVecExt, Rotate, Transform};
use bounding_volume::AABB;
use geom::Box;
use ray::{Ray, RayCast, RayCastWithTransform};

impl<N: Primitive + Orderable + Algebraic,
     V: AlgebraicVecExt<N> + Clone>
RayCast<N, V> for Box<N, V> {
    #[inline]
    fn toi_with_ray(&self, ray: &Ray<V>) -> Option<N> {
        AABB::new(-self.half_extents(), self.half_extents()).toi_with_ray(ray)
    }

    #[inline]
    fn toi_and_normal_with_ray(&self, ray: &Ray<V>) -> Option<(N, V)> {
        AABB::new(-self.half_extents(), self.half_extents()).toi_and_normal_with_ray(ray)
    }
}

impl<N: Primitive + Orderable + Algebraic,
     V: AlgebraicVecExt<N> + Clone,
     M: Rotate<V> + Transform<V>>
RayCastWithTransform<N, V, M> for Box<N, V> { }
<commit_msg>Use the gjk-based ray cast when a box geometry is not 0.<commit_after>use nalgebra::na::{AlgebraicVecExt, Rotate, Transform, Identity, Cast};
use narrow::algorithm::johnson_simplex::JohnsonSimplex;
use bounding_volume::AABB;
use geom::{HasMargin, Box};
use ray::{Ray, RayCast, RayCastWithTransform};
use ray::ray_implicit::gjk_toi_and_normal_with_ray;

impl<N: Primitive + Orderable + Algebraic + Float + Cast<f32>,
     V: AlgebraicVecExt<N> + Clone>
RayCast<N, V> for Box<N, V> {
    #[inline]
    fn toi_with_ray(&self, ray: &Ray<V>) -> Option<N> {
        if self.margin().is_zero() {
            gjk_toi_and_normal_with_ray(
                &Identity::new(),
                self,
                &mut JohnsonSimplex::<N, V>::new_w_tls(),
                ray).map(|(n, _)| n)
        }
        else {
            AABB::new(-self.half_extents(), self.half_extents()).toi_with_ray(ray)
        }
    }

    #[inline]
    fn toi_and_normal_with_ray(&self, ray: &Ray<V>) -> Option<(N, V)> {
        if self.margin().is_zero() {
            AABB::new(-self.half_extents(), self.half_extents()).toi_and_normal_with_ray(ray)
        }
        else {
            gjk_toi_and_normal_with_ray(
                &Identity::new(),
                self,
                &mut JohnsonSimplex::<N, V>::new_w_tls(),
                ray)
        }
    }
}

impl<N: Primitive + Orderable + Algebraic + Float + Cast<f32>,
     V: AlgebraicVecExt<N> + Clone,
     M: Rotate<V> + Transform<V>>
RayCastWithTransform<N, V, M> for Box<N, V> { }
<|endoftext|>