<commit_before>use ide_db::defs::{Definition, NameRefClass};
use syntax::{ast, AstNode, SyntaxKind, T};

use crate::{
    assist_context::{AssistContext, Assists},
    AssistId, AssistKind,
};

// Assist: add_turbo_fish
//
// Adds `::<_>` to a call of a generic method or function.
//
// ```
// fn make<T>() -> T { todo!() }
// fn main() {
//     let x = make$0();
// }
// ```
// ->
// ```
// fn make<T>() -> T { todo!() }
// fn main() {
//     let x = make::<${0:_}>();
// }
// ```
pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let ident = ctx.find_token_syntax_at_offset(SyntaxKind::IDENT).or_else(|| {
        let arg_list = ctx.find_node_at_offset::<ast::ArgList>()?;
        if arg_list.args().next().is_some() {
            return None;
        }
        cov_mark::hit!(add_turbo_fish_after_call);
        cov_mark::hit!(add_type_ascription_after_call);
        arg_list.l_paren_token()?.prev_token().filter(|it| it.kind() == SyntaxKind::IDENT)
    })?;
    let next_token = ident.next_token()?;
    if next_token.kind() == T![::] {
        cov_mark::hit!(add_turbo_fish_one_fish_is_enough);
        return None;
    }
    let name_ref = ast::NameRef::cast(ident.parent()?)?;
    let def = match NameRefClass::classify(&ctx.sema, &name_ref)? {
        NameRefClass::Definition(def) => def,
        NameRefClass::FieldShorthand { .. } => return None,
    };
    let fun = match def {
        Definition::Function(it) => it,
        _ => return None,
    };
    let generics = hir::GenericDef::Function(fun).params(ctx.sema.db);
    if generics.is_empty() {
        cov_mark::hit!(add_turbo_fish_non_generic);
        return None;
    }

    if let Some(let_stmt) = ctx.find_node_at_offset::<ast::LetStmt>() {
        if let_stmt.colon_token().is_none() {
            let type_pos = let_stmt.pat()?.syntax().last_token()?.text_range().end();
            let semi_pos = let_stmt.syntax().last_token()?.text_range().end();

            acc.add(
                AssistId("add_type_ascription", AssistKind::RefactorRewrite),
                "Add `: _` before assignment operator",
                ident.text_range(),
                |builder| {
                    if let_stmt.semicolon_token().is_none() {
                        builder.insert(semi_pos, ";");
                    }
                    match ctx.config.snippet_cap {
                        Some(cap) => builder.insert_snippet(cap, type_pos, ": ${0:_}"),
                        None => builder.insert(type_pos, ": _"),
                    }
                },
            )?
        } else {
            cov_mark::hit!(add_type_ascription_already_typed);
        }
    }

    let number_of_arguments = generics
        .iter()
        .filter(|param| match param {
            hir::GenericParam::TypeParam(_) => true,
            _ => false,
        })
        .count();
    let fish_head = std::iter::repeat("_").take(number_of_arguments).collect::<Vec<_>>().join(",");

    acc.add(
        AssistId("add_turbo_fish", AssistKind::RefactorRewrite),
        "Add `::<>`",
        ident.text_range(),
        |builder| match ctx.config.snippet_cap {
            Some(cap) => {
                let snip = format!("::<${{0:{}}}>", fish_head);
                builder.insert_snippet(cap, ident.text_range().end(), snip)
            }
            None => {
                let snip = format!("::<{}>", fish_head);
                builder.insert(ident.text_range().end(), snip);
            }
        },
    )
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_by_label, check_assist_not_applicable};

    use super::*;

    #[test]
    fn add_turbo_fish_function() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    make$0();
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    make::<${0:_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_function_multiple_generic_types() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<T, A>() -> T {}
fn main() {
    make$0();
}
"#,
            r#"
fn make<T, A>() -> T {}
fn main() {
    make::<${0:_,_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_function_many_generic_types() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<T, A, B, C, D, E, F>() -> T {}
fn main() {
    make$0();
}
"#,
            r#"
fn make<T, A, B, C, D, E, F>() -> T {}
fn main() {
    make::<${0:_,_,_,_,_,_,_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_after_call() {
        cov_mark::check!(add_turbo_fish_after_call);
        check_assist(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    make()$0;
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    make::<${0:_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_method() {
        check_assist(
            add_turbo_fish,
            r#"
struct S;
impl S {
    fn make<T>(&self) -> T {}
}
fn main() {
    S.make$0();
}
"#,
            r#"
struct S;
impl S {
    fn make<T>(&self) -> T {}
}
fn main() {
    S.make::<${0:_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_one_fish_is_enough() {
        cov_mark::check!(add_turbo_fish_one_fish_is_enough);
        check_assist_not_applicable(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    make$0::<()>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_non_generic() {
        cov_mark::check!(add_turbo_fish_non_generic);
        check_assist_not_applicable(
            add_turbo_fish,
            r#"
fn make() -> () {}
fn main() {
    make$0();
}
"#,
        );
    }

    #[test]
    fn add_type_ascription_function() {
        check_assist_by_label(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    let x = make$0();
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: ${0:_} = make();
}
"#,
            "Add `: _` before assignment operator",
        );
    }

    #[test]
    fn add_type_ascription_after_call() {
        cov_mark::check!(add_type_ascription_after_call);
        check_assist_by_label(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    let x = make()$0;
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: ${0:_} = make();
}
"#,
            "Add `: _` before assignment operator",
        );
    }

    #[test]
    fn add_type_ascription_method() {
        check_assist_by_label(
            add_turbo_fish,
            r#"
struct S;
impl S {
    fn make<T>(&self) -> T {}
}
fn main() {
    let x = S.make$0();
}
"#,
            r#"
struct S;
impl S {
    fn make<T>(&self) -> T {}
}
fn main() {
    let x: ${0:_} = S.make();
}
"#,
            "Add `: _` before assignment operator",
        );
    }

    #[test]
    fn add_type_ascription_already_typed() {
        cov_mark::check!(add_type_ascription_already_typed);
        check_assist(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: () = make$0();
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: () = make::<${0:_}>();
}
"#,
        );
    }

    #[test]
    fn add_type_ascription_append_semicolon() {
        check_assist_by_label(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    let x = make$0()
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: ${0:_} = make();
}
"#,
            "Add `: _` before assignment operator",
        );
    }

    #[test]
    fn add_turbo_fish_function_lifetime_parameter() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<'a, T, A>(t: T, a: A) {}
fn main() {
    make$0(5, 2);
}
"#,
            r#"
fn make<'a, T, A>(t: T, a: A) {}
fn main() {
    make::<${0:_,_}>(5, 2);
}
"#,
        );
    }
}
<commit_msg>support const params<commit_after>use ide_db::defs::{Definition, NameRefClass};
use syntax::{ast, AstNode, SyntaxKind, T};

use crate::{
    assist_context::{AssistContext, Assists},
    AssistId, AssistKind,
};

// Assist: add_turbo_fish
//
// Adds `::<_>` to a call of a generic method or function.
//
// ```
// fn make<T>() -> T { todo!() }
// fn main() {
//     let x = make$0();
// }
// ```
// ->
// ```
// fn make<T>() -> T { todo!() }
// fn main() {
//     let x = make::<${0:_}>();
// }
// ```
pub(crate) fn add_turbo_fish(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let ident = ctx.find_token_syntax_at_offset(SyntaxKind::IDENT).or_else(|| {
        let arg_list = ctx.find_node_at_offset::<ast::ArgList>()?;
        if arg_list.args().next().is_some() {
            return None;
        }
        cov_mark::hit!(add_turbo_fish_after_call);
        cov_mark::hit!(add_type_ascription_after_call);
        arg_list.l_paren_token()?.prev_token().filter(|it| it.kind() == SyntaxKind::IDENT)
    })?;
    let next_token = ident.next_token()?;
    if next_token.kind() == T![::] {
        cov_mark::hit!(add_turbo_fish_one_fish_is_enough);
        return None;
    }
    let name_ref = ast::NameRef::cast(ident.parent()?)?;
    let def = match NameRefClass::classify(&ctx.sema, &name_ref)? {
        NameRefClass::Definition(def) => def,
        NameRefClass::FieldShorthand { .. } => return None,
    };
    let fun = match def {
        Definition::Function(it) => it,
        _ => return None,
    };
    let generics = hir::GenericDef::Function(fun).params(ctx.sema.db);
    if generics.is_empty() {
        cov_mark::hit!(add_turbo_fish_non_generic);
        return None;
    }

    if let Some(let_stmt) = ctx.find_node_at_offset::<ast::LetStmt>() {
        if let_stmt.colon_token().is_none() {
            let type_pos = let_stmt.pat()?.syntax().last_token()?.text_range().end();
            let semi_pos = let_stmt.syntax().last_token()?.text_range().end();

            acc.add(
                AssistId("add_type_ascription", AssistKind::RefactorRewrite),
                "Add `: _` before assignment operator",
                ident.text_range(),
                |builder| {
                    if let_stmt.semicolon_token().is_none() {
                        builder.insert(semi_pos, ";");
                    }
                    match ctx.config.snippet_cap {
                        Some(cap) => builder.insert_snippet(cap, type_pos, ": ${0:_}"),
                        None => builder.insert(type_pos, ": _"),
                    }
                },
            )?
        } else {
            cov_mark::hit!(add_type_ascription_already_typed);
        }
    }

    let number_of_arguments = generics
        .iter()
        .filter(|param| match param {
            hir::GenericParam::TypeParam(_) | hir::GenericParam::ConstParam(_) => true,
            _ => false,
        })
        .count();
    let fish_head = std::iter::repeat("_").take(number_of_arguments).collect::<Vec<_>>().join(",");

    acc.add(
        AssistId("add_turbo_fish", AssistKind::RefactorRewrite),
        "Add `::<>`",
        ident.text_range(),
        |builder| match ctx.config.snippet_cap {
            Some(cap) => {
                let snip = format!("::<${{0:{}}}>", fish_head);
                builder.insert_snippet(cap, ident.text_range().end(), snip)
            }
            None => {
                let snip = format!("::<{}>", fish_head);
                builder.insert(ident.text_range().end(), snip);
            }
        },
    )
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_by_label, check_assist_not_applicable};

    use super::*;

    #[test]
    fn add_turbo_fish_function() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    make$0();
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    make::<${0:_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_function_multiple_generic_types() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<T, A>() -> T {}
fn main() {
    make$0();
}
"#,
            r#"
fn make<T, A>() -> T {}
fn main() {
    make::<${0:_,_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_function_many_generic_types() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<T, A, B, C, D, E, F>() -> T {}
fn main() {
    make$0();
}
"#,
            r#"
fn make<T, A, B, C, D, E, F>() -> T {}
fn main() {
    make::<${0:_,_,_,_,_,_,_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_after_call() {
        cov_mark::check!(add_turbo_fish_after_call);
        check_assist(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    make()$0;
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    make::<${0:_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_method() {
        check_assist(
            add_turbo_fish,
            r#"
struct S;
impl S {
    fn make<T>(&self) -> T {}
}
fn main() {
    S.make$0();
}
"#,
            r#"
struct S;
impl S {
    fn make<T>(&self) -> T {}
}
fn main() {
    S.make::<${0:_}>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_one_fish_is_enough() {
        cov_mark::check!(add_turbo_fish_one_fish_is_enough);
        check_assist_not_applicable(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    make$0::<()>();
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_non_generic() {
        cov_mark::check!(add_turbo_fish_non_generic);
        check_assist_not_applicable(
            add_turbo_fish,
            r#"
fn make() -> () {}
fn main() {
    make$0();
}
"#,
        );
    }

    #[test]
    fn add_type_ascription_function() {
        check_assist_by_label(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    let x = make$0();
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: ${0:_} = make();
}
"#,
            "Add `: _` before assignment operator",
        );
    }

    #[test]
    fn add_type_ascription_after_call() {
        cov_mark::check!(add_type_ascription_after_call);
        check_assist_by_label(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    let x = make()$0;
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: ${0:_} = make();
}
"#,
            "Add `: _` before assignment operator",
        );
    }

    #[test]
    fn add_type_ascription_method() {
        check_assist_by_label(
            add_turbo_fish,
            r#"
struct S;
impl S {
    fn make<T>(&self) -> T {}
}
fn main() {
    let x = S.make$0();
}
"#,
            r#"
struct S;
impl S {
    fn make<T>(&self) -> T {}
}
fn main() {
    let x: ${0:_} = S.make();
}
"#,
            "Add `: _` before assignment operator",
        );
    }

    #[test]
    fn add_type_ascription_already_typed() {
        cov_mark::check!(add_type_ascription_already_typed);
        check_assist(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: () = make$0();
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: () = make::<${0:_}>();
}
"#,
        );
    }

    #[test]
    fn add_type_ascription_append_semicolon() {
        check_assist_by_label(
            add_turbo_fish,
            r#"
fn make<T>() -> T {}
fn main() {
    let x = make$0()
}
"#,
            r#"
fn make<T>() -> T {}
fn main() {
    let x: ${0:_} = make();
}
"#,
            "Add `: _` before assignment operator",
        );
    }

    #[test]
    fn add_turbo_fish_function_lifetime_parameter() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<'a, T, A>(t: T, a: A) {}
fn main() {
    make$0(5, 2);
}
"#,
            r#"
fn make<'a, T, A>(t: T, a: A) {}
fn main() {
    make::<${0:_,_}>(5, 2);
}
"#,
        );
    }

    #[test]
    fn add_turbo_fish_function_const_parameter() {
        check_assist(
            add_turbo_fish,
            r#"
fn make<T, const N: usize>(t: T) {}
fn main() {
    make$0(3);
}
"#,
            r#"
fn make<T, const N: usize>(t: T) {}
fn main() {
    make::<${0:_,_}>(3);
}
"#,
        );
    }
}
<|endoftext|>


<commit_before>//! The main scheduler logic.
//!
//! The scheduler is implemented as a singleton in order to make it easy for code anywhere in the
//! project to make use of async functionality. The actual scheduler instance is not publicly
//! accessible, instead we use various standalone functions like `start()` and `wait_for()` to
//! safely manage access to the scheduler.
//!
//! # Scheduling Work
//!
//! Use `scheduler::start()` to run some work asynchronously, getting an `Async<T>` representing the
//! result of the work. Use `Async::await()` to suspend the current fiber until the work completes
//! and get the result. By default, dropping an `Async<T>` will suspend the current fiber until
//! the work finishes, but you can use `Async::forget()` to ignore the result without blocking.
//!
//! # Sharing Data Between Work
//!
//! Unlike with `std::thread::spawn()`, it's possible for work started with `scheduler::start()`
//! to borrow data from the caller:
//!
//! ```
//!
//! ```

use fiber::{self, Fiber, FiberId};
use cell_extras::AtomicInitCell;
use std::boxed::FnBox;
use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt::{self, Debug, Formatter};
use std::marker::PhantomData;
use std::mem;
use std::sync::{Condvar, Mutex, Once, ONCE_INIT};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Receiver};

const DEFAULT_STACK_SIZE: usize = 64 * 1024;

static CONDVAR: AtomicInitCell<Condvar> = AtomicInitCell::new();
static INSTANCE: AtomicInitCell<Mutex<Scheduler>> = AtomicInitCell::new();
static INSTANCE_INIT: Once = ONCE_INIT;
static WORK_COUNTER: AtomicUsize = AtomicUsize::new(1);

/// Represents the result of a computation that may finish at some point in the future.
///
/// Use `scheduler::start()` to run some work asynchronously, getting an `Async<T>` representing the
/// result of the work. Use `Async::await()` to suspend the current fiber until the work completes
/// and get the result. By default, dropping an `Async<T>` will suspend the current fiber until
/// the work finishes, but you can use `Async::forget()` to ignore the result without blocking.
///
/// # Sharing Data Across Work
///
/// It's possible to share
#[derive(Debug)]
pub struct Async<'a, T> {
    work: WorkId,
    receiver: Receiver<T>,
    _phantom: PhantomData<&'a FnMut()>,
}

impl<'a, T> Async<'a, T> {
    /// Suspend the current fiber until the async operation finishes.
    pub fn await(self) -> T {
        let result = {
            let Async { work, ref receiver, .. } = self;
            work.await();
            receiver.try_recv().expect("Failed to receive result of async computation")
        };

        // Don't run the destuctor since we've already waited on the result.
        mem::forget(self);

        result
    }
}

impl<T> Async<'static, T> {
    pub fn work_id(&self) -> WorkId {
        self.work
    }

    pub fn forget(self) {
        mem::forget(self);
    }
}

impl<'a, T> Drop for Async<'a, T> {
    fn drop(&mut self) {
        self.work.await();
    }
}

/// A shareable reference to a work unit, counterpart to `Async<T>`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct WorkId(usize);

impl WorkId {
    /// Suspends the current fiber until this work unit has completed.
    ///
    /// If the work unit has already finished then `await()` will return immediately.
    pub fn await(self) {
        if Scheduler::with(|scheduler| scheduler.add_dependency(self)) {
            suspend();
        }
    }
}

/// Initializes a newly-spawned worker thread.
///
/// Prepares the worker thread by initializing it for Fiber usage.
// TODO: This should probably only be public within the crate. Only the engine should be using this,
// and only at startup, we probably don't want user code to be spawning threads anyway.
pub fn init_thread() {
    // Make sure the scheduler is initialized before first use.
    Scheduler::with(|_| {});

    // Manually convert current thread into a fiber because reasons.
    fiber::init();
}

// TODO: This should probably only be public within the crate. Only the engine should be using this,
// and only at startup, we probably don't want user code to be spawning threads anyway.
pub fn run_wait_fiber() {
    // Make sure the scheduler is initialized before first use.
    Scheduler::with(|_| {});

    // Setup this thread for running fibers and create an initial fiber for it. This will become
    // the wait fiber for this thread.
    fiber::init();

    fiber_routine();
}

pub fn start<'a, F, T>(func: F) -> Async<'a, T>
    where
    F: FnOnce() -> T,
    F: 'a + Send,
    T: 'a + Send,
{
    // Normally we can't box a closure with a non 'static lifetime because it could outlive its
    // borrowed data. In this case the lifetime parameter on the returned `Async` ensures that
    // the closure can't outlive the borrowed data, so we use this evil magic to convince the
    // compiler to allow us to box the closure.
    unsafe fn erase_lifetime<'a, F>(func: F) -> Box<FnBox()>
        where
        F: FnOnce(),
        F: 'a + Send,
    {
        let boxed_proc = Box::new(func);
        let proc_ptr = Box::into_raw(boxed_proc) as *mut FnBox();
        Box::from_raw(::std::mem::transmute(proc_ptr))
    }

    // Create the channel that'll be used to send the result of the operation to the `Async` object.
    let (sender, receiver) = mpsc::sync_channel(1);

    let work_id = WorkId(WORK_COUNTER.fetch_add(1, Ordering::Relaxed));

    let work_proc = unsafe {
        erase_lifetime(move || {
            let result = func();
            sender.try_send(result).expect("Failed to send async result");
        })
    };

    Scheduler::with(move |scheduler| scheduler.schedule_work(Work {
        func: work_proc,
        id: work_id,
    }));

    Async {
        work: work_id,
        receiver: receiver,
        _phantom: PhantomData,
    }
}

/// Suspends the current fiber and makes the wait fiber active.
///
/// Generally you shouldn't need to call this directly, but if you have one piece of code that
/// runs synchronously for a long time you can use `suspend()` to yield time to other work.
pub fn suspend() {
    let next_fiber = Scheduler::with(|scheduler| scheduler.next_fiber());

    let suspended = unsafe { next_fiber.resume() };

    Scheduler::with(move |scheduler| scheduler.handle_suspended(suspended));
}

fn fiber_routine() -> ! {
    loop {
        match Scheduler::with(|scheduler| scheduler.next()) {
            Some(NextWork::Work(Work { func, id })) => {
                Scheduler::with(|scheduler| scheduler.start_work(id));
                func();
                Scheduler::with(|scheduler| scheduler.finish_work(id));
            },
            Some(NextWork::Fiber(fiber)) => {
                let suspended = unsafe { fiber.resume() };
                Scheduler::with(move |scheduler| scheduler.handle_suspended(suspended));
            },
            None => {
                // If there's no new work and no fibers ready to run then we want to block the
                // thread until some becomes available.
                let mutex = INSTANCE.borrow();
                let condvar = CONDVAR.borrow();

                let _ = condvar
                    .wait(mutex.lock().expect("Scheduler mutex was poisoned"))
                    .expect("Scheduler mutex was poisoned");
            },
        }
    }
}

struct Work {
    func: Box<FnBox()>,
    id: WorkId,
}

impl Debug for Work {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), fmt::Error> {
        write!(formatter, "Work {{ id: {:?} }}", self.id)
    }
}

enum NextWork {
    Work(Work),
    Fiber(Fiber),
}

struct Scheduler {
    /// Work units that are currently in progress.
    running_work: HashSet<WorkId>,

    work_map: HashMap<FiberId, WorkId>,

    /// New units of work that haven't been started on a fiber yet.
    ///
    /// These are ready to be made active at any time.
    new_work: VecDeque<Work>,

    /// Fibers that have no pending dependencies.
    ///
    /// These are ready to be made active at any time.
    ready_fibers: VecDeque<Fiber>,

    /// A map specifying which pending fibers depend on which others.
    ///
    /// Once all of a fiber's dependencies complete it should be moved to `new_work`.
    dependencies: HashMap<FiberId, (Option<Fiber>, HashSet<WorkId>)>,

    // TODO: Should we distinguise between "finished" and "ready" fibers? My intuition is that we'd
    // want to give fibers that actively have work CPU time before we resume fibers that would be
    // pulling new work, but maybe not? If we threw them all into one queue I guess the worst case
    // scenario would be there's no work left, a bunch of empty fibers, and only a few fibers with
    // active work. In which case we might have to cycle through a bunch of fibers before we can
    // start doing actual work.
    finished: VecDeque<Fiber>,

    finished_work: HashSet<WorkId>,
}

unsafe impl Send for Scheduler {}

impl Scheduler {
    /// Provides safe access to the scheduler instance.
    ///
    /// # Fiber Switches
    ///
    /// Note that it is an error to call `Fiber::make_active()` within `func`. Doing so will cause
    /// the `Mutex` guard on the instance to never unlock, making the scheduler instance
    /// inaccessible. All standalone functions that access the scheduler and wish to switch fibers
    /// should use `Scheduler::next()` to return the next fiber from `with()` and then call
    /// `make_active()` *after* `with()` has returned.
    fn with<F, T>(func: F) -> T
        where F: FnOnce(&mut Scheduler) -> T
    {
        INSTANCE_INIT.call_once(|| {
            let scheduler = Scheduler {
                running_work: HashSet::new(),
                work_map: HashMap::new(),
                new_work: VecDeque::new(),
                ready_fibers: VecDeque::new(),
                dependencies: HashMap::new(),
                finished: VecDeque::new(),
                finished_work: HashSet::new(),
            };

            INSTANCE.init(Mutex::new(scheduler));
            CONDVAR.init(Condvar::new());
        });

        let instance = INSTANCE.borrow();

        let mut guard = instance.lock().expect("Scheduler mutex was poisoned");
        func(&mut *guard)
    }

    /// Add a new unit of work to the pending queue.
    fn schedule_work(&mut self, work: Work) {
        self.new_work.push_back(work);
        CONDVAR.borrow().notify_one();
    }

    /// Adds `dependency` as a dependency of the currently running fiber.
    ///
    /// Returns `true` if work is still in progress and was added as a dependency, false otherwise.
    fn add_dependency(&mut self, dependency: WorkId) -> bool {
        let pending = fiber::current().unwrap();

        if !self.finished_work.contains(&dependency) {
            debug_assert!(
                !self.dependencies.contains_key(&pending),
                "Marking a fiber as pending but it is already pending: {:?}",
                pending,
            );

            // Add `pending` to set of pending fibers and list `dependencies` as dependencies.
            let &mut (_, ref mut dependencies_set) =
            self.dependencies
            .entry(pending)
            .or_insert((None, HashSet::new()));

            // Add `fibers` to the list of ready fibers.
            dependencies_set.insert(dependency);

            true
        } else {
            false
        }
    }

    fn start_work(&mut self, new_work: WorkId) {
        let current = fiber::current().unwrap();
        self.running_work.insert(new_work);
        self.work_map.insert(current, new_work);
    }

    /// Removes the specified unit of work from the scheduler, updating any dependent work.
    fn finish_work(&mut self, finished_work: WorkId) {
        // Iterate over all suspended work units, removing `finished_work` as a dependency where
        // necessary. If any of the work units no longer have dependencies then
        let mut ready = Vec::new();
        for (&pending_fiber, &mut (_, ref mut dependencies)) in &mut self.dependencies {
            dependencies.remove(&finished_work);
            if dependencies.len() == 0 {
                ready.push(pending_fiber);
            }
        }

        for ready_work in ready {
            let (maybe_fiber, _) = self.dependencies.remove(&ready_work).unwrap();
            if let Some(ready_fiber) = maybe_fiber {
                self.ready_fibers.push_back(ready_fiber);
                CONDVAR.borrow().notify_one();
            }
        }

        let fiber = fiber::current().unwrap();
        assert!(self.running_work.remove(&finished_work), "{:?} wasn't in running work set when it finished", finished_work);
        assert!(self.work_map.remove(&fiber).is_some(), "{:?} didn't have {:?} associated in the work map", fiber, finished_work);

        self.finished_work.insert(finished_work);
    }

    /// Performs the necessary bookkeeping when a fiber becomes active.
    fn handle_suspended(&mut self, suspended: Fiber) {
        // If the suspended fiber has dependencies then update the dependencies map with the
        // actual fiber, that way when its dependencies complete it can be resumed. Otherwise, the
        // fiber is done and ready to take on more work. This means that we need to make sure that
        // we always call `add_dependencies()` before suspending a fiber, otherwise a fiber could
        // be marked as done before it's ready.
        if let Some(&mut (ref mut none_fiber, _)) = self.dependencies.get_mut(&suspended.id()) {
            debug_assert!(none_fiber.is_none(), "Dependencies map already had a fiber assicated with fiber ID");
            mem::replace(none_fiber, Some(suspended));
        } else if self.work_map.contains_key(&suspended.id()) {
            self.ready_fibers.push_back(suspended);
        } else {
            self.finished.push_back(suspended);
        }
    }

    /// Gets the next ready fiber, or creates a new one if necessary.
    fn next_fiber(&mut self) -> Fiber {
        self.ready_fibers.pop_front()
            .or_else(|| self.finished.pop_front())
            .unwrap_or_else(|| {
                fn fiber_proc(prev: Fiber) -> ! {
                    // The current fiber has been resumed. Let the scheduler know that the previous fiber is no
                    // longer active.
                    Scheduler::with(|scheduler| scheduler.handle_suspended(prev));

                    fiber_routine();
                }

                Fiber::new(DEFAULT_STACK_SIZE, fiber_proc)
            })
    }

    /// Gets the next available work for a thread, either a new unit of work or a ready fiber.
    ///
    /// Prioritizes new work over pending fibers, and will only return ready fibers that already
    /// have work. To get *any* next fiber, including ones without active work or a new one if no
    /// existing fibers are available, use `next_fiber()`.
    fn next(&mut self) -> Option<NextWork> {
        if let Some(work) = self.new_work.pop_front() {
            Some(NextWork::Work(work))
        } else {
            self.ready_fibers.pop_front().map(|fiber| NextWork::Fiber(fiber))
        }
    }
}
<commit_msg>Don't track finished work IDs<commit_after>//! The main scheduler logic.
//!
//! The scheduler is implemented as a singleton in order to make it easy for code anywhere in the
//! project to make use of async functionality. The actual scheduler instance is not publicly
//! accessible, instead we use various standalone functions like `start()` and `wait_for()` to
//! safely manage access to the scheduler.
//!
//! # Scheduling Work
//!
//! Use `scheduler::start()` to run some work asynchronously, getting an `Async<T>` representing the
//! result of the work. Use `Async::await()` to suspend the current fiber until the work completes
//! and get the result. By default, dropping an `Async<T>` will suspend the current fiber until
//! the work finishes, but you can use `Async::forget()` to ignore the result without blocking.
//!
//! # Sharing Data Between Work
//!
//! Unlike with `std::thread::spawn()`, it's possible for work started with `scheduler::start()`
//! to borrow data from the caller:
//!
//! ```
//!
//! ```

use fiber::{self, Fiber, FiberId};
use cell_extras::AtomicInitCell;
use std::boxed::FnBox;
use std::collections::{HashMap, HashSet, VecDeque};
use std::fmt::{self, Debug, Formatter};
use std::marker::PhantomData;
use std::mem;
use std::sync::{Condvar, Mutex, Once, ONCE_INIT};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::mpsc::{self, Receiver};

const DEFAULT_STACK_SIZE: usize = 64 * 1024;

static CONDVAR: AtomicInitCell<Condvar> = AtomicInitCell::new();
static INSTANCE: AtomicInitCell<Mutex<Scheduler>> = AtomicInitCell::new();
static INSTANCE_INIT: Once = ONCE_INIT;
static WORK_COUNTER: AtomicUsize = AtomicUsize::new(1);

/// Represents the result of a computation that may finish at some point in the future.
///
/// Use `scheduler::start()` to run some work asynchronously, getting an `Async<T>` representing the
/// result of the work. Use `Async::await()` to suspend the current fiber until the work completes
/// and get the result. By default, dropping an `Async<T>` will suspend the current fiber until
/// the work finishes, but you can use `Async::forget()` to ignore the result without blocking.
///
/// # Sharing Data Across Work
///
/// It's possible to share
#[derive(Debug)]
pub struct Async<'a, T> {
    work: WorkId,
    receiver: Receiver<T>,
    _phantom: PhantomData<&'a FnMut()>,
}

impl<'a, T> Async<'a, T> {
    /// Suspend the current fiber until the async operation finishes.
    pub fn await(self) -> T {
        let result = {
            let Async { work, ref receiver, .. } = self;
            work.await();
            receiver.try_recv().expect("Failed to receive result of async computation")
        };

        // Don't run the destuctor since we've already waited on the result.
        mem::forget(self);

        result
    }
}

impl<T> Async<'static, T> {
    pub fn work_id(&self) -> WorkId {
        self.work
    }

    pub fn forget(self) {
        mem::forget(self);
    }
}

impl<'a, T> Drop for Async<'a, T> {
    fn drop(&mut self) {
        self.work.await();
    }
}

/// A shareable reference to a work unit, counterpart to `Async<T>`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct WorkId(usize);

impl WorkId {
    /// Suspends the current fiber until this work unit has completed.
    ///
    /// If the work unit has already finished then `await()` will return immediately.
    pub fn await(self) {
        if Scheduler::with(|scheduler| scheduler.add_dependency(self)) {
            suspend();
        }
    }
}

/// Initializes a newly-spawned worker thread.
///
/// Prepares the worker thread by initializing it for Fiber usage.
// TODO: This should probably only be public within the crate. Only the engine should be using this,
// and only at startup, we probably don't want user code to be spawning threads anyway.
pub fn init_thread() {
    // Make sure the scheduler is initialized before first use.
    Scheduler::with(|_| {});

    // Manually convert current thread into a fiber because reasons.
    fiber::init();
}

// TODO: This should probably only be public within the crate. Only the engine should be using this,
// and only at startup, we probably don't want user code to be spawning threads anyway.
pub fn run_wait_fiber() {
    // Make sure the scheduler is initialized before first use.
    Scheduler::with(|_| {});

    // Setup this thread for running fibers and create an initial fiber for it. This will become
    // the wait fiber for this thread.
    fiber::init();

    fiber_routine();
}

pub fn start<'a, F, T>(func: F) -> Async<'a, T>
    where
    F: FnOnce() -> T,
    F: 'a + Send,
    T: 'a + Send,
{
    // Normally we can't box a closure with a non 'static lifetime because it could outlive its
    // borrowed data. In this case the lifetime parameter on the returned `Async` ensures that
    // the closure can't outlive the borrowed data, so we use this evil magic to convince the
    // compiler to allow us to box the closure.
    unsafe fn erase_lifetime<'a, F>(func: F) -> Box<FnBox()>
        where
        F: FnOnce(),
        F: 'a + Send,
    {
        let boxed_proc = Box::new(func);
        let proc_ptr = Box::into_raw(boxed_proc) as *mut FnBox();
        Box::from_raw(::std::mem::transmute(proc_ptr))
    }

    // Create the channel that'll be used to send the result of the operation to the `Async` object.
    let (sender, receiver) = mpsc::sync_channel(1);

    let work_id = WorkId(WORK_COUNTER.fetch_add(1, Ordering::Relaxed));

    let work_proc = unsafe {
        erase_lifetime(move || {
            let result = func();
            sender.try_send(result).expect("Failed to send async result");
        })
    };

    Scheduler::with(move |scheduler| scheduler.schedule_work(Work {
        func: work_proc,
        id: work_id,
    }));

    Async {
        work: work_id,
        receiver: receiver,
        _phantom: PhantomData,
    }
}

/// Suspends the current fiber and makes the wait fiber active.
///
/// Generally you shouldn't need to call this directly, but if you have one piece of code that
/// runs synchronously for a long time you can use `suspend()` to yield time to other work.
pub fn suspend() {
    let next_fiber = Scheduler::with(|scheduler| scheduler.next_fiber());

    let suspended = unsafe { next_fiber.resume() };

    Scheduler::with(move |scheduler| scheduler.handle_suspended(suspended));
}

fn fiber_routine() -> ! {
    loop {
        match Scheduler::with(|scheduler| scheduler.next()) {
            Some(NextWork::Work(Work { func, id })) => {
                Scheduler::with(|scheduler| scheduler.start_work(id));
                func();
                Scheduler::with(|scheduler| scheduler.finish_work(id));
            },
            Some(NextWork::Fiber(fiber)) => {
                let suspended = unsafe { fiber.resume() };
                Scheduler::with(move |scheduler| scheduler.handle_suspended(suspended));
            },
            None => {
                // If there's no new work and no fibers ready to run then we want to block the
                // thread until some becomes available.
                let mutex = INSTANCE.borrow();
                let condvar = CONDVAR.borrow();

                let _ = condvar
                    .wait(mutex.lock().expect("Scheduler mutex was poisoned"))
                    .expect("Scheduler mutex was poisoned");
            },
        }
    }
}

struct Work {
    func: Box<FnBox()>,
    id: WorkId,
}

impl Debug for Work {
    fn fmt(&self, formatter: &mut Formatter) -> Result<(), fmt::Error> {
        write!(formatter, "Work {{ id: {:?} }}", self.id)
    }
}

enum NextWork {
    Work(Work),
    Fiber(Fiber),
}

struct Scheduler {
    /// Work units that are currently pending or in progress.
    current_work: HashSet<WorkId>,

    work_map: HashMap<FiberId, WorkId>,

    /// New units of work that haven't been started on a fiber yet.
    ///
    /// These are ready to be made active at any time.
    new_work: VecDeque<Work>,

    /// Fibers that have no pending dependencies.
    ///
    /// These are ready to be made active at any time.
    ready_fibers: VecDeque<Fiber>,

    /// A map specifying which pending fibers depend on which others.
    ///
    /// Once all of a fiber's dependencies complete it should be moved to `new_work`.
    dependencies: HashMap<FiberId, (Option<Fiber>, HashSet<WorkId>)>,

    // TODO: Should we distinguise between "finished" and "ready" fibers? My intuition is that we'd
    // want to give fibers that actively have work CPU time before we resume fibers that would be
    // pulling new work, but maybe not? If we threw them all into one queue I guess the worst case
    // scenario would be there's no work left, a bunch of empty fibers, and only a few fibers with
    // active work. In which case we might have to cycle through a bunch of fibers before we can
    // start doing actual work.
    finished: VecDeque<Fiber>,
}

unsafe impl Send for Scheduler {}

impl Scheduler {
    /// Provides safe access to the scheduler instance.
    ///
    /// # Fiber Switches
    ///
    /// Note that it is an error to call `Fiber::make_active()` within `func`. Doing so will cause
    /// the `Mutex` guard on the instance to never unlock, making the scheduler instance
    /// inaccessible. All standalone functions that access the scheduler and wish to switch fibers
    /// should use `Scheduler::next()` to return the next fiber from `with()` and then call
    /// `make_active()` *after* `with()` has returned.
    fn with<F, T>(func: F) -> T
        where F: FnOnce(&mut Scheduler) -> T
    {
        INSTANCE_INIT.call_once(|| {
            let scheduler = Scheduler {
                current_work: HashSet::new(),
                work_map: HashMap::new(),
                new_work: VecDeque::new(),
                ready_fibers: VecDeque::new(),
                dependencies: HashMap::new(),
                finished: VecDeque::new(),
            };

            INSTANCE.init(Mutex::new(scheduler));
            CONDVAR.init(Condvar::new());
        });

        let instance = INSTANCE.borrow();

        let mut guard = instance.lock().expect("Scheduler mutex was poisoned");
        func(&mut *guard)
    }

    /// Add a new unit of work to the pending queue.
    fn schedule_work(&mut self, work: Work) {
        assert!(self.current_work.insert(work.id), "Work's ID was already present in current work set");
        self.new_work.push_back(work);
        CONDVAR.borrow().notify_one();
    }

    /// Adds `dependency` as a dependency of the currently running fiber.
    ///
    /// Returns `true` if work is still in progress and was added as a dependency, false otherwise.
    fn add_dependency(&mut self, dependency: WorkId) -> bool {
        let pending = fiber::current().unwrap();

        if self.current_work.contains(&dependency) {
            debug_assert!(
                !self.dependencies.contains_key(&pending),
                "Marking a fiber as pending but it is already pending: {:?}",
                pending,
            );

            // Add `pending` to set of pending fibers and list `dependencies` as dependencies.
            let &mut (_, ref mut dependencies_set) =
            self.dependencies
            .entry(pending)
            .or_insert((None, HashSet::new()));

            // Add `fibers` to the list of ready fibers.
            dependencies_set.insert(dependency);

            true
        } else {
            false
        }
    }

    fn start_work(&mut self, new_work: WorkId) {
        debug_assert!(self.current_work.contains(&new_work), "Work ID was not in current work set");

        let current = fiber::current().unwrap();
        self.work_map.insert(current, new_work);
    }

    /// Removes the specified unit of work from the scheduler, updating any dependent work.
    fn finish_work(&mut self, finished_work: WorkId) {
        // Iterate over all suspended work units, removing `finished_work` as a dependency where
        // necessary. If any of the work units no longer have dependencies then
        let mut ready = Vec::new();
        for (&pending_fiber, &mut (_, ref mut dependencies)) in &mut self.dependencies {
            dependencies.remove(&finished_work);
            if dependencies.len() == 0 {
                ready.push(pending_fiber);
            }
        }

        for ready_work in ready {
            let (maybe_fiber, _) = self.dependencies.remove(&ready_work).unwrap();
            if let Some(ready_fiber) = maybe_fiber {
                self.ready_fibers.push_back(ready_fiber);
                CONDVAR.borrow().notify_one();
            }
        }

        let fiber = fiber::current().unwrap();
        assert!(self.current_work.remove(&finished_work), "{:?} wasn't in current work set when it finished", finished_work);
        assert!(self.work_map.remove(&fiber).is_some(), "{:?} didn't have {:?} associated in the work map", fiber, finished_work);
    }

    /// Performs the necessary bookkeeping when a fiber becomes active.
    fn handle_suspended(&mut self, suspended: Fiber) {
        // If the suspended fiber has dependencies then update the dependencies map with the
        // actual fiber, that way when its dependencies complete it can be resumed. Otherwise, the
        // fiber is done and ready to take on more work. This means that we need to make sure that
        // we always call `add_dependencies()` before suspending a fiber, otherwise a fiber could
        // be marked as done before it's ready.
        if let Some(&mut (ref mut none_fiber, _)) = self.dependencies.get_mut(&suspended.id()) {
            debug_assert!(none_fiber.is_none(), "Dependencies map already had a fiber assicated with fiber ID");
            mem::replace(none_fiber, Some(suspended));
        } else if self.work_map.contains_key(&suspended.id()) {
            self.ready_fibers.push_back(suspended);
        } else {
            self.finished.push_back(suspended);
        }
    }

    /// Gets the next ready fiber, or creates a new one if necessary.
    fn next_fiber(&mut self) -> Fiber {
        self.ready_fibers.pop_front()
            .or_else(|| self.finished.pop_front())
            .unwrap_or_else(|| {
                fn fiber_proc(prev: Fiber) -> ! {
                    // The current fiber has been resumed. Let the scheduler know that the previous fiber is no
                    // longer active.
                    Scheduler::with(|scheduler| scheduler.handle_suspended(prev));

                    fiber_routine();
                }

                Fiber::new(DEFAULT_STACK_SIZE, fiber_proc)
            })
    }

    /// Gets the next available work for a thread, either a new unit of work or a ready fiber.
    ///
    /// Prioritizes new work over pending fibers, and will only return ready fibers that already
    /// have work. To get *any* next fiber, including ones without active work or a new one if no
    /// existing fibers are available, use `next_fiber()`.
    fn next(&mut self) -> Option<NextWork> {
        if let Some(work) = self.new_work.pop_front() {
            Some(NextWork::Work(work))
        } else {
            self.ready_fibers.pop_front().map(|fiber| NextWork::Fiber(fiber))
        }
    }
}
<|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/. */

//! A middleware for authenticating requests based on
//! [JWT](https://jwt.io/introduction/).
//!
//! # The Auth Middleware
//!
//! Here _auth_ stands for _authentication_. A
//! [POST `/login`](https://github.com/fxbox/users/blob/master/doc/API.md#post-login)
//! request will authenticate a user. If so, a session JWT is returned
//! in the body of the response. This token must be sent with any further request to
//! keep track of the session.

use super::users_db::{ReadFilter, User, UsersDb};
use super::errors::*;

use crypto::sha2::Sha256;
use iron::{AroundMiddleware, Handler, headers, status};
use iron::method::Method;
use iron::prelude::*;
use jwt::{self, Error, Header, Token};

/// Structure representing [JWT claims section](https://jwt.io/introduction/).
///
/// Claims made by the authentication protocol includes `id` and `name` with
/// database unique id and username respectively.
#[derive(Default, RustcDecodable, RustcEncodable)]
pub struct SessionClaims{
    pub id: i32,
    pub name: String
}

/// Factory to create a session token `String` for a user in the database.
pub struct SessionToken;

impl SessionToken {
    pub fn for_user(user: &User) -> Result<String, Error> {
        let jwt_header: jwt::Header = Default::default();
        let claims = SessionClaims {
            id: user.id.unwrap(),
            name: user.name.to_owned()
        };
        let token = jwt::Token::new(jwt_header, claims);
        token.signed(
            user.secret.to_owned().as_bytes(),
            Sha256::new()
        )
    }
}

/// Represents an authorized endpoint.
///
/// When initializing [AuthMiddleware](./struct.AuthMiddleware.html) you need to
/// pass some routes to be authenticated, these are instances of `AuthEndpoint`.
///
/// `AuthEndpoints` take a vector of methods and a string representing the path
/// of the endpoint to be authenticated. This path can contain wildcard
/// parts. For example:
///
/// AuthEndpoint(vec![Method::Get, Method::Post], "/a/path/:foo/bar/:baz")
///
/// would match with a GET or POST request to /a/path/whatever/bar/whatever
#[derive(Debug)]
pub struct AuthEndpoint(pub Vec<Method>, pub String);

impl PartialEq for AuthEndpoint {
    fn eq(&self, other: &AuthEndpoint) -> bool {
        let AuthEndpoint(ref self_method, ref self_path) = *self;
        let AuthEndpoint(ref other_method, ref other_path) = *other;

        let self_path: Vec<&str> = if self_path.starts_with('/') {
            self_path[1..].split('/').collect()
        } else {
            self_path[0..].split('/').collect()
        };
        let other_path: Vec<&str> = if other_path.starts_with('/') {
            other_path[1..].split('/').collect()
        } else {
            other_path[0..].split('/').collect()
        };

        if self_path.len() != other_path.len() {
            return false;
        }

        let mut contains_method = false;
        for (_, method) in self_method.iter().enumerate() {
            if other_method.contains(method) {
                contains_method = true;
            }
        }

        if !contains_method {
            return false;
        }

        for (i, path) in self_path.iter().enumerate() {
            if &other_path[i] != path && !other_path[i].starts_with(':') {
                return false;
            }
        }
        true
    }
}

struct AuthHandler<H: Handler> {
    handler: H,
    auth_endpoints: Vec<AuthEndpoint>
}

impl<H: Handler> Handler for AuthHandler<H> {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        {
            let endpoint = AuthEndpoint(vec![req.method.clone()],
                                        req.url.path.join("/"));
            // If this is not an authenticated endpoint, just proceed with the
            // original request.
            if !self.auth_endpoints.contains(&endpoint) {
                return self.handler.handle(req);
            }
        }
        // Otherwise, we need to verify the authorization header.
        match req.headers.get::<headers::Authorization<headers::Bearer>>() {
            Some(&headers::Authorization(headers::Bearer { ref token })) => {
                let token = match Token::<Header, SessionClaims>::parse(token) {
                    Ok(token) => token,
                    Err(_) => return EndpointError::with(status::Unauthorized, 401)
                };

                // To verify the token we need to get the secret associated to
                // user id contained in the token claim.
                let db = UsersDb::new();
                match db.read(ReadFilter::Id(token.claims.id)) {
                    Ok(users) => {
                        if users.len() != 1 {
                            return EndpointError::with(status::Unauthorized, 401)
                        }
                        if !token.verify(users[0].secret.to_owned().as_bytes(),
                                         Sha256::new()) {
                            return EndpointError::with(status::Unauthorized, 401)
                        }
                    },
                    Err(_) => return EndpointError::with(status::Unauthorized, 401)
                }
            },
            _ => return EndpointError::with(status::Unauthorized, 401)
        };
        self.handler.handle(req)
    }
}

/// Handle JWT authentication on specified endpoints.
///
/// # Examples
///
/// Before passing it as middleware, you should specify the authenticated
/// endpoints:
///
/// ```
/// extern crate iron;
/// extern crate foxbox_users;
///
/// fn main() {
///     use foxbox_users::users_router::UsersRouter;
///     use foxbox_users::auth_middleware::AuthMiddleware;
///     use iron::prelude::{Chain, Iron};
///
///     let router = UsersRouter::init();
///     let mut chain = Chain::new(router);
///     chain.around(AuthMiddleware{
///         auth_endpoints: vec![]
///     });
/// # if false {
///     Iron::new(chain).http("localhost:3000").unwrap();
/// # }
/// }
/// ```
pub struct AuthMiddleware {
    /// `Vec<AuthEndpoint>` containing the set of endpoints to be authenticated.
    pub auth_endpoints: Vec<AuthEndpoint>
}

impl AroundMiddleware for AuthMiddleware {
    fn around(self, handler: Box<Handler>) -> Box<Handler> {
        Box::new(AuthHandler {
            handler: handler,
            auth_endpoints: self.auth_endpoints
        }) as Box<Handler>
    }
}

#[cfg(test)]
describe! auth_middleware_tests {
    before_each {
        use iron::headers::Headers;
        use iron::prelude::*;
        use iron::method::Method;
        use iron::status::Status;
        use iron_test::request;
        use router::Router;

        fn not_implemented(_: &mut Request) -> IronResult<Response> {
            Ok(Response::with(Status::NotImplemented))
        }

        let mut router = Router::new();
        router.get("/authenticated", not_implemented);
        router.get("/authenticated/:foo/bar/:baz", not_implemented);
        router.delete("/authenticated", not_implemented);
        router.get("/not_authenticated", not_implemented);

        let mut chain = Chain::new(router);
        chain.around(AuthMiddleware {
            auth_endpoints: vec![
                AuthEndpoint(vec![Method::Get, Method::Delete],
                             "/authenticated".to_owned()),
                AuthEndpoint(vec![Method::Get],
                             "/authenticated/:foo/bar/:baz".to_owned())
            ]
        });
    }

    it "should allow request to not authenticated endpoint" {
        match request::get("http://localhost:3000/not_authenticated",
                           Headers::new(), &chain) {
            Ok(res) => {
                assert_eq!(res.status.unwrap(), Status::NotImplemented)
            },
            Err(_) => assert!(false)
        }
    }

    it "should reject request to authenticated endpoint" {
        match request::get("http://localhost:3000/authenticated",
                           Headers::new(), &chain) {
            Ok(_) => assert!(false),
            Err(err) => {
                assert_eq!(err.response.status.unwrap(), Status::Unauthorized)
            }
        };
        match request::get("http://localhost:3000/authenticated/afoo/bar/abaz",
                           Headers::new(), &chain) {
            Ok(_) => assert!(false),
            Err(err) => {
                assert_eq!(err.response.status.unwrap(), Status::Unauthorized)
            }
        };
        match request::delete("http://localhost:3000/authenticated",
                              Headers::new(), &chain) {
            Ok(_) => assert!(false),
            Err(err) => {
                assert_eq!(err.response.status.unwrap(), Status::Unauthorized)
            }
        }
    }

    it "should allow request to authenticated endpoint" {
        use super::super::users_db::{UserBuilder, UsersDb};

        use iron::headers::{Authorization, Bearer};
        use crypto::sha2::Sha256;
        use jwt;

        let db = UsersDb::new();
        db.clear().ok();
        let user = UserBuilder::new()
            .id(1).name(String::from("username"))
            .password(String::from("password"))
            .email(String::from("username@example.com"))
            .secret(String::from("secret"))
            .finalize().unwrap();
        db.create(&user).ok();
        let jwt_header: jwt::Header = Default::default();
        let claims = SessionClaims {
            id: user.id.unwrap(),
            name: user.name.to_owned()
        };
        let token = jwt::Token::new(jwt_header, claims);
        let signed = token.signed(
            user.secret.to_owned().as_bytes(),
            Sha256::new()
        ).ok().unwrap();
        let mut headers = Headers::new();
        headers.set(Authorization(Bearer { token: signed.to_owned() }));
        match request::get("http://localhost:3000/authenticated",
                           headers.clone(), &chain) {
            Ok(res) => {
                assert_eq!(res.status.unwrap(), Status::NotImplemented)
            },
            Err(_) => assert!(false)
        };
        match request::get("http://localhost:3000/authenticated/afoo/bar/abaz",
                           headers.clone(), &chain) {
            Ok(res) => {
                assert_eq!(res.status.unwrap(), Status::NotImplemented)
            },
            Err(_) => assert!(false)
        };
        match request::delete("http://localhost:3000/authenticated",
                              headers, &chain) {
            Ok(res) => {
                assert_eq!(res.status.unwrap(), Status::NotImplemented)
            },
            Err(_) => assert!(false)
        }
    }
}
<commit_msg>Expose a AuthenticationMiddleware.verify method. r=francisco<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/. */

//! A middleware for authenticating requests based on
//! [JWT](https://jwt.io/introduction/).
//!
//! # The Auth Middleware
//!
//! Here _auth_ stands for _authentication_. A
//! [POST `/login`](https://github.com/fxbox/users/blob/master/doc/API.md#post-login)
//! request will authenticate a user. If so, a session JWT is returned
//! in the body of the response. This token must be sent with any further request to
//! keep track of the session.

use super::users_db::{ReadFilter, User, UsersDb};
use super::errors::*;

use crypto::sha2::Sha256;
use iron::{AroundMiddleware, Handler, headers, status};
use iron::method::Method;
use iron::prelude::*;
use jwt::{self, Error, Header, Token};

/// Structure representing [JWT claims section](https://jwt.io/introduction/).
///
/// Claims made by the authentication protocol includes `id` and `name` with
/// database unique id and username respectively.
#[derive(Default, RustcDecodable, RustcEncodable)]
pub struct SessionClaims{
    pub id: i32,
    pub name: String
}

/// Factory to create a session token `String` for a user in the database.
pub struct SessionToken;

impl SessionToken {
    pub fn for_user(user: &User) -> Result<String, Error> {
        let jwt_header: jwt::Header = Default::default();
        let claims = SessionClaims {
            id: user.id.unwrap(),
            name: user.name.to_owned()
        };
        let token = jwt::Token::new(jwt_header, claims);
        token.signed(
            user.secret.to_owned().as_bytes(),
            Sha256::new()
        )
    }
}

/// Represents an authorized endpoint.
///
/// When initializing [AuthMiddleware](./struct.AuthMiddleware.html) you need to
/// pass some routes to be authenticated, these are instances of `AuthEndpoint`.
///
/// `AuthEndpoints` take a vector of methods and a string representing the path
/// of the endpoint to be authenticated. This path can contain wildcard
/// parts. For example:
///
/// AuthEndpoint(vec![Method::Get, Method::Post], "/a/path/:foo/bar/:baz")
///
/// would match with a GET or POST request to /a/path/whatever/bar/whatever
#[derive(Debug)]
pub struct AuthEndpoint(pub Vec<Method>, pub String);

impl PartialEq for AuthEndpoint {
    fn eq(&self, other: &AuthEndpoint) -> bool {
        let AuthEndpoint(ref self_method, ref self_path) = *self;
        let AuthEndpoint(ref other_method, ref other_path) = *other;

        let self_path: Vec<&str> = if self_path.starts_with('/') {
            self_path[1..].split('/').collect()
        } else {
            self_path[0..].split('/').collect()
        };
        let other_path: Vec<&str> = if other_path.starts_with('/') {
            other_path[1..].split('/').collect()
        } else {
            other_path[0..].split('/').collect()
        };

        if self_path.len() != other_path.len() {
            return false;
        }

        let mut contains_method = false;
        for (_, method) in self_method.iter().enumerate() {
            if other_method.contains(method) {
                contains_method = true;
            }
        }

        if !contains_method {
            return false;
        }

        for (i, path) in self_path.iter().enumerate() {
            if &other_path[i] != path && !other_path[i].starts_with(':') {
                return false;
            }
        }
        true
    }
}

struct AuthHandler<H: Handler> {
    handler: H,
    auth_endpoints: Vec<AuthEndpoint>
}

impl<H: Handler> Handler for AuthHandler<H> {
    fn handle(&self, req: &mut Request) -> IronResult<Response> {
        {
            let endpoint = AuthEndpoint(vec![req.method.clone()],
                                        req.url.path.join("/"));
            // If this is not an authenticated endpoint, just proceed with the
            // original request.
            if !self.auth_endpoints.contains(&endpoint) {
                return self.handler.handle(req);
            }
        }
        // Otherwise, we need to verify the authorization header.
        match req.headers.get::<headers::Authorization<headers::Bearer>>() {
            Some(&headers::Authorization(headers::Bearer { ref token })) => {
                if let Err(_) = AuthMiddleware::verify(token) {
                    return EndpointError::with(status::Unauthorized, 401)
                }
            },
            _ => return EndpointError::with(status::Unauthorized, 401)
        };
        self.handler.handle(req)
    }
}

/// Handle JWT authentication on specified endpoints.
///
/// # Examples
///
/// Before passing it as middleware, you should specify the authenticated
/// endpoints:
///
/// ```
/// extern crate iron;
/// extern crate foxbox_users;
///
/// fn main() {
///     use foxbox_users::users_router::UsersRouter;
///     use foxbox_users::auth_middleware::AuthMiddleware;
///     use iron::prelude::{Chain, Iron};
///
///     let router = UsersRouter::init();
///     let mut chain = Chain::new(router);
///     chain.around(AuthMiddleware{
///         auth_endpoints: vec![]
///     });
/// # if false {
///     Iron::new(chain).http("localhost:3000").unwrap();
/// # }
/// }
/// ```
pub struct AuthMiddleware {
    /// `Vec<AuthEndpoint>` containing the set of endpoints to be authenticated.
    pub auth_endpoints: Vec<AuthEndpoint>
}

impl AroundMiddleware for AuthMiddleware {
    fn around(self, handler: Box<Handler>) -> Box<Handler> {
        Box::new(AuthHandler {
            handler: handler,
            auth_endpoints: self.auth_endpoints
        }) as Box<Handler>
    }
}

impl AuthMiddleware {
    pub fn verify(token: &str) -> Result<(), ()> {
        if token.is_empty() {
            return Err(());
        }

        let token = match Token::<Header, SessionClaims>::parse(token) {
            Ok(token) => token,
            Err(_) => return Err(())
        };

        // To verify the token we need to get the secret associated to
        // user id contained in the token claim.
        let db = UsersDb::new();
        match db.read(ReadFilter::Id(token.claims.id)) {
            Ok(users) => {
                if users.len() != 1 {
                    return Err(());
                }
                if !token.verify(users[0].secret.to_owned().as_bytes(),
                                 Sha256::new()) {
                    return Err(());
                }
            },
            Err(_) => {
                return Err(());
            }
        };

        Ok(())
    }
}

#[cfg(test)]
describe! auth_middleware_tests {
    before_each {
        use iron::headers::Headers;
        use iron::prelude::*;
        use iron::method::Method;
        use iron::status::Status;
        use iron_test::request;
        use router::Router;

        fn not_implemented(_: &mut Request) -> IronResult<Response> {
            Ok(Response::with(Status::NotImplemented))
        }

        let mut router = Router::new();
        router.get("/authenticated", not_implemented);
        router.get("/authenticated/:foo/bar/:baz", not_implemented);
        router.delete("/authenticated", not_implemented);
        router.get("/not_authenticated", not_implemented);

        let mut chain = Chain::new(router);
        chain.around(AuthMiddleware {
            auth_endpoints: vec![
                AuthEndpoint(vec![Method::Get, Method::Delete],
                             "/authenticated".to_owned()),
                AuthEndpoint(vec![Method::Get],
                             "/authenticated/:foo/bar/:baz".to_owned())
            ]
        });
    }

    it "should allow request to not authenticated endpoint" {
        match request::get("http://localhost:3000/not_authenticated",
                           Headers::new(), &chain) {
            Ok(res) => {
                assert_eq!(res.status.unwrap(), Status::NotImplemented)
            },
            Err(_) => assert!(false)
        }
    }

    it "should reject request to authenticated endpoint" {
        match request::get("http://localhost:3000/authenticated",
                           Headers::new(), &chain) {
            Ok(_) => assert!(false),
            Err(err) => {
                assert_eq!(err.response.status.unwrap(), Status::Unauthorized)
            }
        };
        match request::get("http://localhost:3000/authenticated/afoo/bar/abaz",
                           Headers::new(), &chain) {
            Ok(_) => assert!(false),
            Err(err) => {
                assert_eq!(err.response.status.unwrap(), Status::Unauthorized)
            }
        };
        match request::delete("http://localhost:3000/authenticated",
                              Headers::new(), &chain) {
            Ok(_) => assert!(false),
            Err(err) => {
                assert_eq!(err.response.status.unwrap(), Status::Unauthorized)
            }
        }
    }

    it "should allow request to authenticated endpoint" {
        use super::super::users_db::{UserBuilder, UsersDb};

        use iron::headers::{Authorization, Bearer};
        use crypto::sha2::Sha256;
        use jwt;

        let db = UsersDb::new();
        db.clear().ok();
        let user = UserBuilder::new()
            .id(1).name(String::from("username"))
            .password(String::from("password"))
            .email(String::from("username@example.com"))
            .secret(String::from("secret"))
            .finalize().unwrap();
        db.create(&user).ok();
        let jwt_header: jwt::Header = Default::default();
        let claims = SessionClaims {
            id: user.id.unwrap(),
            name: user.name.to_owned()
        };
        let token = jwt::Token::new(jwt_header, claims);
        let signed = token.signed(
            user.secret.to_owned().as_bytes(),
            Sha256::new()
        ).ok().unwrap();
        let mut headers = Headers::new();
        headers.set(Authorization(Bearer { token: signed.to_owned() }));
        match request::get("http://localhost:3000/authenticated",
                           headers.clone(), &chain) {
            Ok(res) => {
                assert_eq!(res.status.unwrap(), Status::NotImplemented)
            },
            Err(_) => assert!(false)
        };
        match request::get("http://localhost:3000/authenticated/afoo/bar/abaz",
                           headers.clone(), &chain) {
            Ok(res) => {
                assert_eq!(res.status.unwrap(), Status::NotImplemented)
            },
            Err(_) => assert!(false)
        };
        match request::delete("http://localhost:3000/authenticated",
                              headers, &chain) {
            Ok(res) => {
                assert_eq!(res.status.unwrap(), Status::NotImplemented)
            },
            Err(_) => assert!(false)
        }
    }
}
<|endoftext|>


<commit_before>// Cooperative optimization solver
use std::collections::BinaryHeap;
use std::io::Write;
extern crate rand;

#[macro_use(max)]
extern crate gelpia_utils;
extern crate ga;
extern crate gr;

use ga::{ea, Individual};

use gelpia_utils::{Quple, INF, NINF, Flt, Parameters};

use gr::{GI, width_box, split_box, midpoint_box, eps_tol};

use std::sync::{Barrier, RwLock, Arc, RwLockWriteGuard};

use std::sync::atomic::{AtomicBool, Ordering};

use std::thread;

use std::time::Duration;

extern crate function;
use function::FuncObj;

extern crate args;
use args::{process_args};

extern crate time;

/// Returns the guaranteed upperbound for the algorithm
/// from the queue.
fn get_upper_bound(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
                   f_best_high: f64) -> f64{
    let mut max = f_best_high;
    for qi in q.iter() {
        max = max!{max, qi.fdata.upper()};
    }
    max
}

fn log_max(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
           f_best_low: f64,
           f_best_high: f64) {
    let max = get_upper_bound(q, f_best_high);
    let _ = writeln!(&mut std::io::stderr(),
                     "lb: {}, possible ub: {}, guaranteed ub: {}",
                     f_best_low,
                     f_best_high,
                     max);
}

// Returns the upper bound, the domain where this bound occurs and a status
// flag indicating whether the answer is complete for the problem.
fn ibba(x_0: Vec<GI>, e_x: Flt, e_f: Flt,
        f_bestag: Arc<RwLock<Flt>>,
        f_best_shared: Arc<RwLock<Flt>>,
        x_bestbb: Arc<RwLock<Vec<GI>>>,
        b1: Arc<Barrier>, b2: Arc<Barrier>,
        q: Arc<RwLock<BinaryHeap<Quple>>>,
        sync: Arc<AtomicBool>, stop: Arc<AtomicBool>,
        f: FuncObj,
        logging: bool)
        -> (Flt, Vec<GI>, bool) {
    let mut f_best_high = NINF;
    let mut f_best_low  = NINF;
    let mut best_x = x_0.clone();

    let mut iters: u32 = 0;
    q.write().unwrap().push(Quple{p: INF, pf: 0, data: x_0.clone(),
                                  fdata: f.call(&x_0)});

    while q.read().unwrap().len() != 0 && !stop.load(Ordering::Acquire) {
        if sync.load(Ordering::Acquire) {
            // Ugly: Update the update thread's view of the best branch bound.
            *f_best_shared.write().unwrap() = f_best_low;
            b1.wait();
            b2.wait();
        }
        // Take q as writable during an iteration
        let mut q = q.write().unwrap();
        let fbl_orig = f_best_low;
        f_best_low = max!(f_best_low, *f_bestag.read().unwrap());

        if iters % 2048 == 0 {
            let guaranteed_bound = get_upper_bound(&q, f_best_high);
            if (guaranteed_bound - f_best_high).abs() < e_f {
                f_best_high = guaranteed_bound;
                break;
            }

        }

        if fbl_orig != f_best_low && logging {
            log_max(&q, f_best_low, f_best_high);
        }

        let (ref x, fx) =
            match q.pop() {
                Some(y) => (y.data, y.fdata),
                None    => unreachable!()
            };

        if fx.upper() < f_best_low ||
            width_box(x, e_x) ||
            eps_tol(fx, e_f) {
                if f_best_high < fx.upper() {
                    f_best_high = fx.upper();
                    best_x = x.clone();
                    if logging {
                        log_max(&q, f_best_low, f_best_high);
                    }
                }
                continue;
            }
        else {
            let x_s = split_box(&x);
            for sx in x_s {
                let mid = midpoint_box(&sx);
                let est_m = f.call(&mid);
                let fsx = f.call(&sx);
                let est_max = if est_m.lower() > fsx.lower() {est_m} else {fsx};
                if f_best_low < est_max.lower()  {
                    f_best_low = est_max.lower();
                    *x_bestbb.write().unwrap() = sx.clone();
                }
                iters += 1;
                q.push(Quple{p: est_max.upper(),
                             pf: iters,
                             data: sx,
                             fdata: fsx});
            }
        }
    }
    if !stop.load(Ordering::Acquire) {
        // Exiting normally
        // Tell GA thread to stop
        stop.store(true, Ordering::Release);
        (f_best_high, best_x, true)
    }
    else {
        // Exiting by some other condition
        (f_best_high, best_x, false)
    }
}

fn distance(p: &Vec<GI>, x: &Vec<GI>) -> Flt {
    let mut result = 0.0;
    for i in 0..x.len() {
        let dx = max!(x[i].lower() - p[i].lower(),
                      0.0, &p[i].lower() - x[i].upper());
        result += dx*dx;
    }
    result.sqrt()
}

fn update(q: Arc<RwLock<BinaryHeap<Quple>>>, population: Arc<RwLock<Vec<Individual>>>,
          f_best_shared: Arc<RwLock<Flt>>,
          stop: Arc<AtomicBool>, sync: Arc<AtomicBool>,
          b1: Arc<Barrier>, b2: Arc<Barrier>,
          f: FuncObj,
          upd_interval: u32,
          timeout: u32) {
    let start = time::get_time();

    'out: while !stop.load(Ordering::Acquire) {
        // Timer code...
        let last_update = time::get_time();
        while (time::get_time() - last_update).num_seconds() <= upd_interval as i64 {
            thread::sleep(Duration::new(1, 0));
            if timeout > 0 &&
                (time::get_time() - start).num_seconds() >= timeout as i64 {
                    let _ = writeln!(&mut std::io::stderr(), "Stopping early...");
                    stop.store(true, Ordering::Release);
                    break 'out;
                }
            if stop.load(Ordering::Acquire) { // Check if we've already stopped
                break 'out;
            }
        }
        // Signal EA and IBBA threads.
        sync.store(true, Ordering::Release);

        // Wait for EA and IBBA threads to stop.
        b1.wait();                                                               // Dead lock here?
        // Do update bizness.
        let mut q_u = q.write().unwrap();
        let mut q = Vec::new();
        for qi in q_u.iter() {
            q.push(qi.clone());
        }
        let mut pop = population.write().unwrap();
        let fbest = f_best_shared.read().unwrap();

        for i in 0..pop.len() {
            let mut d_min = INF;
            let mut j = 0;
            let mut x: Vec<GI>;
            let px: Flt = NINF;
            let mut x_c: Vec<GI> = q[0].data.clone();
            while j < q.len() && d_min != 0.0 {
                x = q[j].data.clone();
                if f.call(&x).upper() < *fbest {
                    q.remove(j);
                    continue;
                }
                let d = distance(&pop[i].solution, &x);

                if d == 0.0 {d_min = 0.0; continue;}
                else {
                    if d < d_min {d_min = d; x_c = x.clone();}
                }
                j += 1;
            }
            if d_min == 0.0 {
                let fp = f.call(&pop[i].solution).lower();
                if px < fp {
                    q[j] = Quple{p: fp, pf: q[j].pf,
                                 data: q[j].data.clone(),
                                 fdata: q[j].fdata }.clone();
                }
            }
            else {
                // Individual is outside the search region.
                // Project it back into the nearest current search space.
                project(&mut pop[i], &x_c, f.clone());
            }
        }
        // Restore the q for the ibba thread.
        (*q_u) = BinaryHeap::from(q);
        // Clear sync flag.
        sync.store(false, Ordering::SeqCst);

        // Resume EA and IBBA threads.
        b2.wait();                                                               // Dead lock here?
    }
}

/* Projects the box x into the box x_c */
fn project(p: &mut Individual, x_c: &Vec<GI>, f: FuncObj) {
    for i in 0..x_c.len() {
        if p.solution[i].lower() < x_c[i].lower()
            || p.solution[i].lower() > x_c[i].upper() {
                if x_c[i].upper() < p.solution[i].lower()
                {p.solution[i] = GI::new_d(x_c[i].upper(),
                                           x_c[i].upper());}
                else {p.solution[i] = GI::new_d(x_c[i].lower(),
                                                x_c[i].lower());}
            }
    }
    p.fitness = f.call(&p.solution).lower();
}

fn main() {
    let args = process_args();

    let ref x_0 = args.domain;
    let ref fo = args.function;
    let x_err = args.x_error;
    let y_err = args.y_error;

    // Early out if there are no input variables...
    if x_0.len() == 0 {
        let result = fo.call(&x_0);
        println!("[{}, {{}}]", result.upper());
        return
    }

    let q_inner: BinaryHeap<Quple> = BinaryHeap::new();
    let q = Arc::new(RwLock::new(q_inner));

    let population_inner: Vec<Individual> = Vec::new();
    let population = Arc::new(RwLock::new(population_inner));

    let b1 = Arc::new(Barrier::new(3));
    let b2 = Arc::new(Barrier::new(3));

    let sync = Arc::new(AtomicBool::new(false));
    let stop = Arc::new(AtomicBool::new(false));

    let f_bestag: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));
    let f_best_shared: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));

    let x_e = x_0.clone();
    let x_i = x_0.clone();

    let x_bestbb = Arc::new(RwLock::new(x_0.clone()));

    let ibba_thread =
    {
        let q = q.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let f_bestag = f_bestag.clone();
        let f_best_shared = f_best_shared.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let fo_c = fo.clone();
        let logging = args.logging;
        thread::Builder::new().name("IBBA".to_string()).spawn(move || {
            ibba(x_i, x_err, y_err,
                 f_bestag, f_best_shared,
                 x_bestbb,
                 b1, b2, q, sync, stop, fo_c, logging)
        })};

    let _ea_thread =
    {
        let population = population.clone();
        let f_bestag = f_bestag.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        thread::Builder::new().name("EA".to_string()).spawn(move || {
            ea(x_e, Parameters{population: 2000, //1000,
                               selection: 8, //4,
                               elitism: 5, //2,
                               mutation: 0.4_f64,//0.3_f64,
                               crossover: 0.3_f64 // 0.5_f64
            },
               population,
               f_bestag,
               x_bestbb,
               b1, b2,
               stop, sync, fo_c)
        })};

    // pending finding out how to kill threads
    //let update_thread =
    {
        let q = q.clone();
        let population = population.clone();
        let f_best_shared = f_best_shared.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        let to = args.timeout.clone();
        let ui = args.update_interval.clone();
        let _ = thread::Builder::new().name("Update".to_string()).spawn(move || {
            update(q, population, f_best_shared,
                   stop, sync, b1, b2, fo_c, ui, to)
        });};

    let result = ibba_thread.unwrap().join();
    /*let ea_result = ea_thread.unwrap().join();
    if !ea_result.is_ok() {
        unreachable!();
    }*/

    // Join EA and Update here pending stop signaling.
    if result.is_ok() {
        let (mut max, mut interval, completed) = result.unwrap();
        if !completed {
            // Go through all remaining intervals from IBBA to find the true
            // max
            let mut lq = q.write().unwrap();
            while lq.len() != 0 {
                let v = lq.pop().unwrap();
                let ref l_data = &v.data;
                let ub = fo.call(l_data).upper();
                if ub > max {
                    max = ub;
                    interval =v.data.clone();
                }
            }
        }
        println!("[{}, {{", max);
        for i in 0..args.names.len() {
            println!("'{}' : {},", args.names[i], interval[i].to_string());
        }
        println!("}}]");

    }
    else {println!("error")}
}
<commit_msg>Max iters complet<commit_after>// Cooperative optimization solver
use std::collections::BinaryHeap;
use std::io::Write;
extern crate rand;

#[macro_use(max)]
extern crate gelpia_utils;
extern crate ga;
extern crate gr;

use ga::{ea, Individual};

use gelpia_utils::{Quple, INF, NINF, Flt, Parameters};

use gr::{GI, width_box, split_box, midpoint_box, eps_tol};

use std::sync::{Barrier, RwLock, Arc, RwLockWriteGuard};

use std::sync::atomic::{AtomicBool, Ordering};

use std::thread;

use std::time::Duration;

extern crate function;
use function::FuncObj;

extern crate args;
use args::{process_args};

extern crate time;

/// Returns the guaranteed upperbound for the algorithm
/// from the queue.
fn get_upper_bound(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
                   f_best_high: f64) -> f64{
    let mut max = f_best_high;
    for qi in q.iter() {
        max = max!{max, qi.fdata.upper()};
    }
    max
}

fn log_max(q: &RwLockWriteGuard<BinaryHeap<Quple>>,
           f_best_low: f64,
           f_best_high: f64) {
    let max = get_upper_bound(q, f_best_high);
    let _ = writeln!(&mut std::io::stderr(),
                     "lb: {}, possible ub: {}, guaranteed ub: {}",
                     f_best_low,
                     f_best_high,
                     max);
}

// Returns the upper bound, the domain where this bound occurs and a status
// flag indicating whether the answer is complete for the problem.
fn ibba(x_0: Vec<GI>, e_x: Flt, e_f: Flt,
        f_bestag: Arc<RwLock<Flt>>,
        f_best_shared: Arc<RwLock<Flt>>,
        x_bestbb: Arc<RwLock<Vec<GI>>>,
        b1: Arc<Barrier>, b2: Arc<Barrier>,
        q: Arc<RwLock<BinaryHeap<Quple>>>,
        sync: Arc<AtomicBool>, stop: Arc<AtomicBool>,
        f: FuncObj,
        logging: bool, max_iters: u32)
        -> (Flt, Vec<GI>, bool) {
    let mut f_best_high = NINF;
    let mut f_best_low  = NINF;
    let mut best_x = x_0.clone();

    let mut iters: u32 = 0;
    q.write().unwrap().push(Quple{p: INF, pf: 0, data: x_0.clone(),
                                  fdata: f.call(&x_0)});

    while q.read().unwrap().len() != 0 && !stop.load(Ordering::Acquire)
          && iters < max_iters {
        if sync.load(Ordering::Acquire) {
            // Ugly: Update the update thread's view of the best branch bound.
            *f_best_shared.write().unwrap() = f_best_low;
            b1.wait();
            b2.wait();
        }
        // Take q as writable during an iteration
        let mut q = q.write().unwrap();
        let fbl_orig = f_best_low;
        f_best_low = max!(f_best_low, *f_bestag.read().unwrap());

        if iters % 2048 == 0 {
            let guaranteed_bound = get_upper_bound(&q, f_best_high);
            if (guaranteed_bound - f_best_high).abs() < e_f {
                f_best_high = guaranteed_bound;
                break;
            }

        }

        if fbl_orig != f_best_low && logging {
            log_max(&q, f_best_low, f_best_high);
        }

        let (ref x, fx) =
            match q.pop() {
                Some(y) => (y.data, y.fdata),
                None    => unreachable!()
            };

        if fx.upper() < f_best_low ||
            width_box(x, e_x) ||
            eps_tol(fx, e_f) {
                if f_best_high < fx.upper() {
                    f_best_high = fx.upper();
                    best_x = x.clone();
                    if logging {
                        log_max(&q, f_best_low, f_best_high);
                    }
                }
                continue;
            }
        else {
            let x_s = split_box(&x);
            for sx in x_s {
                let mid = midpoint_box(&sx);
                let est_m = f.call(&mid);
                let fsx = f.call(&sx);
                let est_max = if est_m.lower() > fsx.lower() {est_m} else {fsx};
                if f_best_low < est_max.lower()  {
                    f_best_low = est_max.lower();
                    *x_bestbb.write().unwrap() = sx.clone();
                }
                iters += 1;
                q.push(Quple{p: est_max.upper(),
                             pf: iters,
                             data: sx,
                             fdata: fsx});
            }
        }
    }
    if !stop.load(Ordering::Acquire) {
        // Exiting normally
        // Tell GA thread to stop
        stop.store(true, Ordering::Release);
        (f_best_high, best_x, true)
    }
    else {
        // Exiting by some other condition
        (f_best_high, best_x, false)
    }
}

fn distance(p: &Vec<GI>, x: &Vec<GI>) -> Flt {
    let mut result = 0.0;
    for i in 0..x.len() {
        let dx = max!(x[i].lower() - p[i].lower(),
                      0.0, &p[i].lower() - x[i].upper());
        result += dx*dx;
    }
    result.sqrt()
}

fn update(q: Arc<RwLock<BinaryHeap<Quple>>>, population: Arc<RwLock<Vec<Individual>>>,
          f_best_shared: Arc<RwLock<Flt>>,
          stop: Arc<AtomicBool>, sync: Arc<AtomicBool>,
          b1: Arc<Barrier>, b2: Arc<Barrier>,
          f: FuncObj,
          upd_interval: u32,
          timeout: u32) {
    let start = time::get_time();

    'out: while !stop.load(Ordering::Acquire) {
        // Timer code...
        let last_update = time::get_time();
        while (time::get_time() - last_update).num_seconds() <= upd_interval as i64 {
            thread::sleep(Duration::new(1, 0));
            if timeout > 0 &&
                (time::get_time() - start).num_seconds() >= timeout as i64 {
                    let _ = writeln!(&mut std::io::stderr(), "Stopping early...");
                    stop.store(true, Ordering::Release);
                    break 'out;
                }
            if stop.load(Ordering::Acquire) { // Check if we've already stopped
                break 'out;
            }
        }
        // Signal EA and IBBA threads.
        sync.store(true, Ordering::Release);

        // Wait for EA and IBBA threads to stop.
        b1.wait();                                                               // Dead lock here?
        // Do update bizness.
        let mut q_u = q.write().unwrap();
        let mut q = Vec::new();
        for qi in q_u.iter() {
            q.push(qi.clone());
        }
        let mut pop = population.write().unwrap();
        let fbest = f_best_shared.read().unwrap();

        for i in 0..pop.len() {
            let mut d_min = INF;
            let mut j = 0;
            let mut x: Vec<GI>;
            let px: Flt = NINF;
            let mut x_c: Vec<GI> = q[0].data.clone();
            while j < q.len() && d_min != 0.0 {
                x = q[j].data.clone();
                if f.call(&x).upper() < *fbest {
                    q.remove(j);
                    continue;
                }
                let d = distance(&pop[i].solution, &x);

                if d == 0.0 {d_min = 0.0; continue;}
                else {
                    if d < d_min {d_min = d; x_c = x.clone();}
                }
                j += 1;
            }
            if d_min == 0.0 {
                let fp = f.call(&pop[i].solution).lower();
                if px < fp {
                    q[j] = Quple{p: fp, pf: q[j].pf,
                                 data: q[j].data.clone(),
                                 fdata: q[j].fdata }.clone();
                }
            }
            else {
                // Individual is outside the search region.
                // Project it back into the nearest current search space.
                project(&mut pop[i], &x_c, f.clone());
            }
        }
        // Restore the q for the ibba thread.
        (*q_u) = BinaryHeap::from(q);
        // Clear sync flag.
        sync.store(false, Ordering::SeqCst);

        // Resume EA and IBBA threads.
        b2.wait();                                                               // Dead lock here?
    }
}

/* Projects the box x into the box x_c */
fn project(p: &mut Individual, x_c: &Vec<GI>, f: FuncObj) {
    for i in 0..x_c.len() {
        if p.solution[i].lower() < x_c[i].lower()
            || p.solution[i].lower() > x_c[i].upper() {
                if x_c[i].upper() < p.solution[i].lower()
                {p.solution[i] = GI::new_d(x_c[i].upper(),
                                           x_c[i].upper());}
                else {p.solution[i] = GI::new_d(x_c[i].lower(),
                                                x_c[i].lower());}
            }
    }
    p.fitness = f.call(&p.solution).lower();
}

fn main() {
    let args = process_args();

    let ref x_0 = args.domain;
    let ref fo = args.function;
    let x_err = args.x_error;
    let y_err = args.y_error;

    // Early out if there are no input variables...
    if x_0.len() == 0 {
        let result = fo.call(&x_0);
        println!("[{}, {{}}]", result.upper());
        return
    }

    let q_inner: BinaryHeap<Quple> = BinaryHeap::new();
    let q = Arc::new(RwLock::new(q_inner));

    let population_inner: Vec<Individual> = Vec::new();
    let population = Arc::new(RwLock::new(population_inner));

    let b1 = Arc::new(Barrier::new(3));
    let b2 = Arc::new(Barrier::new(3));

    let sync = Arc::new(AtomicBool::new(false));
    let stop = Arc::new(AtomicBool::new(false));

    let f_bestag: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));
    let f_best_shared: Arc<RwLock<Flt>> = Arc::new(RwLock::new(NINF));

    let x_e = x_0.clone();
    let x_i = x_0.clone();

    let x_bestbb = Arc::new(RwLock::new(x_0.clone()));

    let ibba_thread =
    {
        let q = q.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let f_bestag = f_bestag.clone();
        let f_best_shared = f_best_shared.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let fo_c = fo.clone();
        let logging = args.logging;
        let iters= args.iters;
        thread::Builder::new().name("IBBA".to_string()).spawn(move || {
            ibba(x_i, x_err, y_err,
                 f_bestag, f_best_shared,
                 x_bestbb,
                 b1, b2, q, sync, stop, fo_c, logging, iters)
        })};

    let _ea_thread =
    {
        let population = population.clone();
        let f_bestag = f_bestag.clone();
        let x_bestbb = x_bestbb.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        thread::Builder::new().name("EA".to_string()).spawn(move || {
            ea(x_e, Parameters{population: 2000, //1000,
                               selection: 8, //4,
                               elitism: 5, //2,
                               mutation: 0.4_f64,//0.3_f64,
                               crossover: 0.3_f64 // 0.5_f64
            },
               population,
               f_bestag,
               x_bestbb,
               b1, b2,
               stop, sync, fo_c)
        })};

    // pending finding out how to kill threads
    //let update_thread =
    {
        let q = q.clone();
        let population = population.clone();
        let f_best_shared = f_best_shared.clone();
        let sync = sync.clone();
        let stop = stop.clone();
        let b1 = b1.clone();
        let b2 = b2.clone();
        let fo_c = fo.clone();
        let to = args.timeout.clone();
        let ui = args.update_interval.clone();
        let _ = thread::Builder::new().name("Update".to_string()).spawn(move || {
            update(q, population, f_best_shared,
                   stop, sync, b1, b2, fo_c, ui, to)
        });};

    let result = ibba_thread.unwrap().join();
    /*let ea_result = ea_thread.unwrap().join();
    if !ea_result.is_ok() {
        unreachable!();
    }*/

    // Join EA and Update here pending stop signaling.
    if result.is_ok() {
        let (mut max, mut interval, completed) = result.unwrap();
        if !completed {
            // Go through all remaining intervals from IBBA to find the true
            // max
            let mut lq = q.write().unwrap();
            while lq.len() != 0 {
                let v = lq.pop().unwrap();
                let ref l_data = &v.data;
                let ub = fo.call(l_data).upper();
                if ub > max {
                    max = ub;
                    interval =v.data.clone();
                }
            }
        }
        println!("[{}, {{", max);
        for i in 0..args.names.len() {
            println!("'{}' : {},", args.names[i], interval[i].to_string());
        }
        println!("}}]");

    }
    else {println!("error")}
}
<|endoftext|>