<commit_before>use std::borrow::Borrow;

use rustc_ast::ast;
use rustc_span::symbol::sym;
use rustc_span::Span;
use rustc_target::spec::abi::Abi;

use rustc_index::bit_set::BitSet;
use rustc_middle::mir::MirPass;
use rustc_middle::mir::{self, Body, Local, Location};
use rustc_middle::ty::{self, Ty, TyCtxt};

use crate::impls::{
    DefinitelyInitializedPlaces, MaybeInitializedPlaces, MaybeLiveLocals, MaybeMutBorrowedLocals,
    MaybeUninitializedPlaces,
};
use crate::move_paths::{HasMoveData, MoveData};
use crate::move_paths::{LookupResult, MovePathIndex};
use crate::MoveDataParamEnv;
use crate::{Analysis, JoinSemiLattice, Results, ResultsCursor};

pub struct SanityCheck;

impl<'tcx> MirPass<'tcx> for SanityCheck {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        use crate::has_rustc_mir_with;
        let def_id = body.source.def_id();
        if !tcx.has_attr(def_id, sym::rustc_mir) {
            debug!("skipping rustc_peek::SanityCheck on {}", tcx.def_path_str(def_id));
            return;
        } else {
            debug!("running rustc_peek::SanityCheck on {}", tcx.def_path_str(def_id));
        }

        let attributes = tcx.get_attrs(def_id);
        let param_env = tcx.param_env(def_id);
        let move_data = MoveData::gather_moves(body, tcx, param_env).unwrap();
        let mdpe = MoveDataParamEnv { move_data, param_env };
        let sess = &tcx.sess;

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_maybe_init).is_some() {
            let flow_inits = MaybeInitializedPlaces::new(tcx, body, &mdpe)
                .into_engine(tcx, body)
                .iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_inits);
        }

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_maybe_uninit).is_some() {
            let flow_uninits = MaybeUninitializedPlaces::new(tcx, body, &mdpe)
                .into_engine(tcx, body)
                .iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_uninits);
        }

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_definite_init).is_some() {
            let flow_def_inits = DefinitelyInitializedPlaces::new(tcx, body, &mdpe)
                .into_engine(tcx, body)
                .iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_def_inits);
        }

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_indirectly_mutable).is_some() {
            let flow_mut_borrowed = MaybeMutBorrowedLocals::mut_borrows_only(tcx, body, param_env)
                .into_engine(tcx, body)
                .iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_mut_borrowed);
        }

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_liveness).is_some() {
            let flow_liveness = MaybeLiveLocals.into_engine(tcx, body).iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_liveness);
        }

        if has_rustc_mir_with(sess, &attributes, sym::stop_after_dataflow).is_some() {
            tcx.sess.fatal("stop_after_dataflow ended compilation");
        }
    }
}

/// This function scans `mir` for all calls to the intrinsic
/// `rustc_peek` that have the expression form `rustc_peek(&expr)`.
///
/// For each such call, determines what the dataflow bit-state is for
/// the L-value corresponding to `expr`; if the bit-state is a 1, then
/// that call to `rustc_peek` is ignored by the sanity check. If the
/// bit-state is a 0, then this pass emits an error message saying
/// "rustc_peek: bit not set".
///
/// The intention is that one can write unit tests for dataflow by
/// putting code into a UI test and using `rustc_peek` to
/// make observations about the results of dataflow static analyses.
///
/// (If there are any calls to `rustc_peek` that do not match the
/// expression form above, then that emits an error as well, but those
/// errors are not intended to be used for unit tests.)
pub fn sanity_check_via_rustc_peek<'tcx, A>(
    tcx: TyCtxt<'tcx>,
    body: &Body<'tcx>,
    _attributes: &[ast::Attribute],
    results: &Results<'tcx, A>,
) where
    A: RustcPeekAt<'tcx>,
{
    let def_id = body.source.def_id();
    debug!("sanity_check_via_rustc_peek def_id: {:?}", def_id);

    let mut cursor = ResultsCursor::new(body, results);

    let peek_calls = body.basic_blocks().iter_enumerated().filter_map(|(bb, block_data)| {
        PeekCall::from_terminator(tcx, block_data.terminator()).map(|call| (bb, block_data, call))
    });

    for (bb, block_data, call) in peek_calls {
        // Look for a sequence like the following to indicate that we should be peeking at `_1`:
        //    _2 = &_1;
        //    rustc_peek(_2);
        //
        //    /* or */
        //
        //    _2 = _1;
        //    rustc_peek(_2);
        let (statement_index, peek_rval) = block_data
            .statements
            .iter()
            .enumerate()
            .find_map(|(i, stmt)| value_assigned_to_local(stmt, call.arg).map(|rval| (i, rval)))
            .expect(
                "call to rustc_peek should be preceded by \
                    assignment to temporary holding its argument",
            );

        match (call.kind, peek_rval) {
            (PeekCallKind::ByRef, mir::Rvalue::Ref(_, _, place))
            | (
                PeekCallKind::ByVal,
                mir::Rvalue::Use(mir::Operand::Move(place) | mir::Operand::Copy(place)),
            ) => {
                let loc = Location { block: bb, statement_index };
                cursor.seek_before_primary_effect(loc);
                let state = cursor.get();
                results.analysis.peek_at(tcx, *place, state, call);
            }

            _ => {
                let msg = "rustc_peek: argument expression \
                           must be either `place` or `&place`";
                tcx.sess.span_err(call.span, msg);
            }
        }
    }
}

/// If `stmt` is an assignment where the LHS is the given local (with no projections), returns the
/// RHS of the assignment.
fn value_assigned_to_local<'a, 'tcx>(
    stmt: &'a mir::Statement<'tcx>,
    local: Local,
) -> Option<&'a mir::Rvalue<'tcx>> {
    if let mir::StatementKind::Assign(box (place, rvalue)) = &stmt.kind {
        if let Some(l) = place.as_local() {
            if local == l {
                return Some(&*rvalue);
            }
        }
    }

    None
}

#[derive(Clone, Copy, Debug)]
enum PeekCallKind {
    ByVal,
    ByRef,
}

impl PeekCallKind {
    fn from_arg_ty(arg: Ty<'_>) -> Self {
        match arg.kind() {
            ty::Ref(_, _, _) => PeekCallKind::ByRef,
            _ => PeekCallKind::ByVal,
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct PeekCall {
    arg: Local,
    kind: PeekCallKind,
    span: Span,
}

impl PeekCall {
    fn from_terminator<'tcx>(
        tcx: TyCtxt<'tcx>,
        terminator: &mir::Terminator<'tcx>,
    ) -> Option<Self> {
        use mir::Operand;

        let span = terminator.source_info.span;
        if let mir::TerminatorKind::Call { func: Operand::Constant(func), args, .. } =
            &terminator.kind
        {
            if let ty::FnDef(def_id, substs) = *func.literal.ty().kind() {
                let sig = tcx.fn_sig(def_id);
                let name = tcx.item_name(def_id);
                if sig.abi() != Abi::RustIntrinsic || name != sym::rustc_peek {
                    return None;
                }

                assert_eq!(args.len(), 1);
                let kind = PeekCallKind::from_arg_ty(substs.type_at(0));
                let arg = match &args[0] {
                    Operand::Copy(place) | Operand::Move(place) => {
                        if let Some(local) = place.as_local() {
                            local
                        } else {
                            tcx.sess.diagnostic().span_err(
                                span,
                                "dataflow::sanity_check cannot feed a non-temp to rustc_peek.",
                            );
                            return None;
                        }
                    }
                    _ => {
                        tcx.sess.diagnostic().span_err(
                            span,
                            "dataflow::sanity_check cannot feed a non-temp to rustc_peek.",
                        );
                        return None;
                    }
                };

                return Some(PeekCall { arg, kind, span });
            }
        }

        None
    }
}

pub trait RustcPeekAt<'tcx>: Analysis<'tcx> {
    fn peek_at(
        &self,
        tcx: TyCtxt<'tcx>,
        place: mir::Place<'tcx>,
        flow_state: &Self::Domain,
        call: PeekCall,
    );
}

impl<'tcx, A, D> RustcPeekAt<'tcx> for A
where
    A: Analysis<'tcx, Domain = D> + HasMoveData<'tcx>,
    D: JoinSemiLattice + Clone + Borrow<BitSet<MovePathIndex>>,
{
    fn peek_at(
        &self,
        tcx: TyCtxt<'tcx>,
        place: mir::Place<'tcx>,
        flow_state: &Self::Domain,
        call: PeekCall,
    ) {
        match self.move_data().rev_lookup.find(place.as_ref()) {
            LookupResult::Exact(peek_mpi) => {
                let bit_state = flow_state.borrow().contains(peek_mpi);
                debug!("rustc_peek({:?} = &{:?}) bit_state: {}", call.arg, place, bit_state);
                if !bit_state {
                    tcx.sess.span_err(call.span, "rustc_peek: bit not set");
                }
            }

            LookupResult::Parent(..) => {
                tcx.sess.span_err(call.span, "rustc_peek: argument untracked");
            }
        }
    }
}

impl<'tcx> RustcPeekAt<'tcx> for MaybeMutBorrowedLocals<'_, 'tcx> {
    fn peek_at(
        &self,
        tcx: TyCtxt<'tcx>,
        place: mir::Place<'tcx>,
        flow_state: &BitSet<Local>,
        call: PeekCall,
    ) {
        warn!("peek_at: place={:?}", place);
        let local = if let Some(l) = place.as_local() {
            l
        } else {
            tcx.sess.span_err(call.span, "rustc_peek: argument was not a local");
            return;
        };

        if !flow_state.contains(local) {
            tcx.sess.span_err(call.span, "rustc_peek: bit not set");
        }
    }
}

impl<'tcx> RustcPeekAt<'tcx> for MaybeLiveLocals {
    fn peek_at(
        &self,
        tcx: TyCtxt<'tcx>,
        place: mir::Place<'tcx>,
        flow_state: &BitSet<Local>,
        call: PeekCall,
    ) {
        warn!("peek_at: place={:?}", place);
        let local = if let Some(l) = place.as_local() {
            l
        } else {
            tcx.sess.span_err(call.span, "rustc_peek: argument was not a local");
            return;
        };

        if !flow_state.contains(local) {
            tcx.sess.span_err(call.span, "rustc_peek: bit not set");
        }
    }
}
<commit_msg>demote `rustc_peek` traces look not user-facing<commit_after>use std::borrow::Borrow;

use rustc_ast::ast;
use rustc_span::symbol::sym;
use rustc_span::Span;
use rustc_target::spec::abi::Abi;

use rustc_index::bit_set::BitSet;
use rustc_middle::mir::MirPass;
use rustc_middle::mir::{self, Body, Local, Location};
use rustc_middle::ty::{self, Ty, TyCtxt};

use crate::impls::{
    DefinitelyInitializedPlaces, MaybeInitializedPlaces, MaybeLiveLocals, MaybeMutBorrowedLocals,
    MaybeUninitializedPlaces,
};
use crate::move_paths::{HasMoveData, MoveData};
use crate::move_paths::{LookupResult, MovePathIndex};
use crate::MoveDataParamEnv;
use crate::{Analysis, JoinSemiLattice, Results, ResultsCursor};

pub struct SanityCheck;

impl<'tcx> MirPass<'tcx> for SanityCheck {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
        use crate::has_rustc_mir_with;
        let def_id = body.source.def_id();
        if !tcx.has_attr(def_id, sym::rustc_mir) {
            debug!("skipping rustc_peek::SanityCheck on {}", tcx.def_path_str(def_id));
            return;
        } else {
            debug!("running rustc_peek::SanityCheck on {}", tcx.def_path_str(def_id));
        }

        let attributes = tcx.get_attrs(def_id);
        let param_env = tcx.param_env(def_id);
        let move_data = MoveData::gather_moves(body, tcx, param_env).unwrap();
        let mdpe = MoveDataParamEnv { move_data, param_env };
        let sess = &tcx.sess;

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_maybe_init).is_some() {
            let flow_inits = MaybeInitializedPlaces::new(tcx, body, &mdpe)
                .into_engine(tcx, body)
                .iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_inits);
        }

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_maybe_uninit).is_some() {
            let flow_uninits = MaybeUninitializedPlaces::new(tcx, body, &mdpe)
                .into_engine(tcx, body)
                .iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_uninits);
        }

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_definite_init).is_some() {
            let flow_def_inits = DefinitelyInitializedPlaces::new(tcx, body, &mdpe)
                .into_engine(tcx, body)
                .iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_def_inits);
        }

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_indirectly_mutable).is_some() {
            let flow_mut_borrowed = MaybeMutBorrowedLocals::mut_borrows_only(tcx, body, param_env)
                .into_engine(tcx, body)
                .iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_mut_borrowed);
        }

        if has_rustc_mir_with(sess, &attributes, sym::rustc_peek_liveness).is_some() {
            let flow_liveness = MaybeLiveLocals.into_engine(tcx, body).iterate_to_fixpoint();

            sanity_check_via_rustc_peek(tcx, body, &attributes, &flow_liveness);
        }

        if has_rustc_mir_with(sess, &attributes, sym::stop_after_dataflow).is_some() {
            tcx.sess.fatal("stop_after_dataflow ended compilation");
        }
    }
}

/// This function scans `mir` for all calls to the intrinsic
/// `rustc_peek` that have the expression form `rustc_peek(&expr)`.
///
/// For each such call, determines what the dataflow bit-state is for
/// the L-value corresponding to `expr`; if the bit-state is a 1, then
/// that call to `rustc_peek` is ignored by the sanity check. If the
/// bit-state is a 0, then this pass emits an error message saying
/// "rustc_peek: bit not set".
///
/// The intention is that one can write unit tests for dataflow by
/// putting code into a UI test and using `rustc_peek` to
/// make observations about the results of dataflow static analyses.
///
/// (If there are any calls to `rustc_peek` that do not match the
/// expression form above, then that emits an error as well, but those
/// errors are not intended to be used for unit tests.)
pub fn sanity_check_via_rustc_peek<'tcx, A>(
    tcx: TyCtxt<'tcx>,
    body: &Body<'tcx>,
    _attributes: &[ast::Attribute],
    results: &Results<'tcx, A>,
) where
    A: RustcPeekAt<'tcx>,
{
    let def_id = body.source.def_id();
    debug!("sanity_check_via_rustc_peek def_id: {:?}", def_id);

    let mut cursor = ResultsCursor::new(body, results);

    let peek_calls = body.basic_blocks().iter_enumerated().filter_map(|(bb, block_data)| {
        PeekCall::from_terminator(tcx, block_data.terminator()).map(|call| (bb, block_data, call))
    });

    for (bb, block_data, call) in peek_calls {
        // Look for a sequence like the following to indicate that we should be peeking at `_1`:
        //    _2 = &_1;
        //    rustc_peek(_2);
        //
        //    /* or */
        //
        //    _2 = _1;
        //    rustc_peek(_2);
        let (statement_index, peek_rval) = block_data
            .statements
            .iter()
            .enumerate()
            .find_map(|(i, stmt)| value_assigned_to_local(stmt, call.arg).map(|rval| (i, rval)))
            .expect(
                "call to rustc_peek should be preceded by \
                    assignment to temporary holding its argument",
            );

        match (call.kind, peek_rval) {
            (PeekCallKind::ByRef, mir::Rvalue::Ref(_, _, place))
            | (
                PeekCallKind::ByVal,
                mir::Rvalue::Use(mir::Operand::Move(place) | mir::Operand::Copy(place)),
            ) => {
                let loc = Location { block: bb, statement_index };
                cursor.seek_before_primary_effect(loc);
                let state = cursor.get();
                results.analysis.peek_at(tcx, *place, state, call);
            }

            _ => {
                let msg = "rustc_peek: argument expression \
                           must be either `place` or `&place`";
                tcx.sess.span_err(call.span, msg);
            }
        }
    }
}

/// If `stmt` is an assignment where the LHS is the given local (with no projections), returns the
/// RHS of the assignment.
fn value_assigned_to_local<'a, 'tcx>(
    stmt: &'a mir::Statement<'tcx>,
    local: Local,
) -> Option<&'a mir::Rvalue<'tcx>> {
    if let mir::StatementKind::Assign(box (place, rvalue)) = &stmt.kind {
        if let Some(l) = place.as_local() {
            if local == l {
                return Some(&*rvalue);
            }
        }
    }

    None
}

#[derive(Clone, Copy, Debug)]
enum PeekCallKind {
    ByVal,
    ByRef,
}

impl PeekCallKind {
    fn from_arg_ty(arg: Ty<'_>) -> Self {
        match arg.kind() {
            ty::Ref(_, _, _) => PeekCallKind::ByRef,
            _ => PeekCallKind::ByVal,
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct PeekCall {
    arg: Local,
    kind: PeekCallKind,
    span: Span,
}

impl PeekCall {
    fn from_terminator<'tcx>(
        tcx: TyCtxt<'tcx>,
        terminator: &mir::Terminator<'tcx>,
    ) -> Option<Self> {
        use mir::Operand;

        let span = terminator.source_info.span;
        if let mir::TerminatorKind::Call { func: Operand::Constant(func), args, .. } =
            &terminator.kind
        {
            if let ty::FnDef(def_id, substs) = *func.literal.ty().kind() {
                let sig = tcx.fn_sig(def_id);
                let name = tcx.item_name(def_id);
                if sig.abi() != Abi::RustIntrinsic || name != sym::rustc_peek {
                    return None;
                }

                assert_eq!(args.len(), 1);
                let kind = PeekCallKind::from_arg_ty(substs.type_at(0));
                let arg = match &args[0] {
                    Operand::Copy(place) | Operand::Move(place) => {
                        if let Some(local) = place.as_local() {
                            local
                        } else {
                            tcx.sess.diagnostic().span_err(
                                span,
                                "dataflow::sanity_check cannot feed a non-temp to rustc_peek.",
                            );
                            return None;
                        }
                    }
                    _ => {
                        tcx.sess.diagnostic().span_err(
                            span,
                            "dataflow::sanity_check cannot feed a non-temp to rustc_peek.",
                        );
                        return None;
                    }
                };

                return Some(PeekCall { arg, kind, span });
            }
        }

        None
    }
}

pub trait RustcPeekAt<'tcx>: Analysis<'tcx> {
    fn peek_at(
        &self,
        tcx: TyCtxt<'tcx>,
        place: mir::Place<'tcx>,
        flow_state: &Self::Domain,
        call: PeekCall,
    );
}

impl<'tcx, A, D> RustcPeekAt<'tcx> for A
where
    A: Analysis<'tcx, Domain = D> + HasMoveData<'tcx>,
    D: JoinSemiLattice + Clone + Borrow<BitSet<MovePathIndex>>,
{
    fn peek_at(
        &self,
        tcx: TyCtxt<'tcx>,
        place: mir::Place<'tcx>,
        flow_state: &Self::Domain,
        call: PeekCall,
    ) {
        match self.move_data().rev_lookup.find(place.as_ref()) {
            LookupResult::Exact(peek_mpi) => {
                let bit_state = flow_state.borrow().contains(peek_mpi);
                debug!("rustc_peek({:?} = &{:?}) bit_state: {}", call.arg, place, bit_state);
                if !bit_state {
                    tcx.sess.span_err(call.span, "rustc_peek: bit not set");
                }
            }

            LookupResult::Parent(..) => {
                tcx.sess.span_err(call.span, "rustc_peek: argument untracked");
            }
        }
    }
}

impl<'tcx> RustcPeekAt<'tcx> for MaybeMutBorrowedLocals<'_, 'tcx> {
    fn peek_at(
        &self,
        tcx: TyCtxt<'tcx>,
        place: mir::Place<'tcx>,
        flow_state: &BitSet<Local>,
        call: PeekCall,
    ) {
        info!("peek_at: place={:?}", place);
        let local = if let Some(l) = place.as_local() {
            l
        } else {
            tcx.sess.span_err(call.span, "rustc_peek: argument was not a local");
            return;
        };

        if !flow_state.contains(local) {
            tcx.sess.span_err(call.span, "rustc_peek: bit not set");
        }
    }
}

impl<'tcx> RustcPeekAt<'tcx> for MaybeLiveLocals {
    fn peek_at(
        &self,
        tcx: TyCtxt<'tcx>,
        place: mir::Place<'tcx>,
        flow_state: &BitSet<Local>,
        call: PeekCall,
    ) {
        info!("peek_at: place={:?}", place);
        let local = if let Some(l) = place.as_local() {
            l
        } else {
            tcx.sess.span_err(call.span, "rustc_peek: argument was not a local");
            return;
        };

        if !flow_state.contains(local) {
            tcx.sess.span_err(call.span, "rustc_peek: bit not set");
        }
    }
}
<|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/. */

use font_template::{FontTemplate, FontTemplateDescriptor};
use ipc_channel::ipc;
use ipc_channel::router::ROUTER;
use net_traits::{AsyncResponseTarget, PendingAsyncLoad, ResourceTask, ResponseAction};
use platform::font_context::FontContextHandle;
use platform::font_list::for_each_available_family;
use platform::font_list::for_each_variation;
use platform::font_list::last_resort_font_families;
use platform::font_list::system_default_family;
use platform::font_template::FontTemplateData;
use std::borrow::ToOwned;
use std::collections::HashMap;
use std::mem;
use std::sync::mpsc::{Sender, Receiver, channel};
use std::sync::{Arc, Mutex};
use string_cache::Atom;
use style::font_face::Source;
use url::Url;
use util::str::LowercaseString;
use util::task::spawn_named;

/// A list of font templates that make up a given font family.
struct FontFamily {
    templates: Vec<FontTemplate>,
}

impl FontFamily {
    fn new() -> FontFamily {
        FontFamily {
            templates: vec!(),
        }
    }

    /// Find a font in this family that matches a given descriptor.
    fn find_font_for_style<'a>(&'a mut self, desc: &FontTemplateDescriptor, fctx: &FontContextHandle)
                               -> Option<Arc<FontTemplateData>> {
        // TODO(Issue #189): optimize lookup for
        // regular/bold/italic/bolditalic with fixed offsets and a
        // static decision table for fallback between these values.

        // TODO(Issue #190): if not in the fast path above, do
        // expensive matching of weights, etc.
        for template in &mut self.templates {
            let maybe_template = template.data_for_descriptor(fctx, desc);
            if maybe_template.is_some() {
                return maybe_template;
            }
        }

        // If a request is made for a font family that exists,
        // pick the first valid font in the family if we failed
        // to find an exact match for the descriptor.
        for template in &mut self.templates {
            let maybe_template = template.get();
            if maybe_template.is_some() {
                return maybe_template;
            }
        }

        None
    }

    fn add_template(&mut self, identifier: Atom, maybe_data: Option<Vec<u8>>) {
        for template in &self.templates {
            if *template.identifier() == identifier {
                return;
            }
        }

        let template = FontTemplate::new(identifier, maybe_data);
        self.templates.push(template);
    }
}

/// Commands that the FontContext sends to the font cache task.
pub enum Command {
    GetFontTemplate(String, FontTemplateDescriptor, Sender<Reply>),
    GetLastResortFontTemplate(FontTemplateDescriptor, Sender<Reply>),
    AddWebFont(Atom, Source, Sender<()>),
    AddDownloadedWebFont(LowercaseString, Url, Vec<u8>, Sender<()>),
    Exit(Sender<()>),
}

/// Reply messages sent from the font cache task to the FontContext caller.
pub enum Reply {
    GetFontTemplateReply(Option<Arc<FontTemplateData>>),
}

/// The font cache task itself. It maintains a list of reference counted
/// font templates that are currently in use.
struct FontCache {
    port: Receiver<Command>,
    channel_to_self: Sender<Command>,
    generic_fonts: HashMap<LowercaseString, LowercaseString>,
    local_families: HashMap<LowercaseString, FontFamily>,
    web_families: HashMap<LowercaseString, FontFamily>,
    font_context: FontContextHandle,
    resource_task: ResourceTask,
}

fn add_generic_font(generic_fonts: &mut HashMap<LowercaseString, LowercaseString>,
                    generic_name: &str, mapped_name: &str) {
    let opt_system_default = system_default_family(generic_name);
    let family_name = match opt_system_default {
        Some(system_default) => LowercaseString::new(&system_default),
        None => LowercaseString::new(mapped_name),
    };
    generic_fonts.insert(LowercaseString::new(generic_name), family_name);
}

impl FontCache {
    fn run(&mut self) {
        loop {
            let msg = self.port.recv().unwrap();

            match msg {
                Command::GetFontTemplate(family, descriptor, result) => {
                    let family = LowercaseString::new(&family);
                    let maybe_font_template = self.find_font_template(&family, &descriptor);
                    result.send(Reply::GetFontTemplateReply(maybe_font_template)).unwrap();
                }
                Command::GetLastResortFontTemplate(descriptor, result) => {
                    let font_template = self.last_resort_font_template(&descriptor);
                    result.send(Reply::GetFontTemplateReply(Some(font_template))).unwrap();
                }
                Command::AddWebFont(family_name, src, result) => {
                    let family_name = LowercaseString::new(&family_name);
                    if !self.web_families.contains_key(&family_name) {
                        let family = FontFamily::new();
                        self.web_families.insert(family_name.clone(), family);
                    }

                    match src {
                        Source::Url(ref url_source) => {
                            let url = &url_source.url;
                            let load = PendingAsyncLoad::new(self.resource_task.clone(),
                                                             url.clone(),
                                                             None);
                            let (data_sender, data_receiver) = ipc::channel().unwrap();
                            let data_target = AsyncResponseTarget {
                                sender: data_sender,
                            };
                            load.load_async(data_target);
                            let channel_to_self = self.channel_to_self.clone();
                            let url = (*url).clone();
                            let bytes = Mutex::new(Vec::new());
                            ROUTER.add_route(data_receiver.to_opaque(), box move |message| {
                                let response: ResponseAction = message.to().unwrap();
                                match response {
                                    ResponseAction::HeadersAvailable(_) |
                                    ResponseAction::ResponseComplete(Err(_)) => {}
                                    ResponseAction::DataAvailable(new_bytes) => {
                                        bytes.lock().unwrap().extend(new_bytes.into_iter())
                                    }
                                    ResponseAction::ResponseComplete(Ok(_)) => {
                                        let mut bytes = bytes.lock().unwrap();
                                        let bytes = mem::replace(&mut *bytes, Vec::new());
                                        let command =
                                            Command::AddDownloadedWebFont(family_name.clone(),
                                                                          url.clone(),
                                                                          bytes,
                                                                          result.clone());
                                        channel_to_self.send(command).unwrap();
                                    }
                                }
                            });
                        }
                        Source::Local(ref local_family_name) => {
                            let family = &mut self.web_families.get_mut(&family_name).unwrap();
                            for_each_variation(&local_family_name, |path| {
                                family.add_template(Atom::from_slice(&path), None);
                            });
                            result.send(()).unwrap();
                        }
                    }
                }
                Command::AddDownloadedWebFont(family_name, url, bytes, result) => {
                    let family = &mut self.web_families.get_mut(&family_name).unwrap();
                    family.add_template(Atom::from_slice(&url.to_string()), Some(bytes));
                    drop(result.send(()));
                }
                Command::Exit(result) => {
                    result.send(()).unwrap();
                    break;
                }
            }
        }
    }

    fn refresh_local_families(&mut self) {
        self.local_families.clear();
        for_each_available_family(|family_name| {
            let family_name = LowercaseString::new(&family_name);
            if !self.local_families.contains_key(&family_name) {
                let family = FontFamily::new();
                self.local_families.insert(family_name, family);
            }
        });
    }

    fn transform_family(&self, family: &LowercaseString) -> LowercaseString {
        match self.generic_fonts.get(family) {
            None => family.clone(),
            Some(mapped_family) => (*mapped_family).clone()
        }
    }

    fn find_font_in_local_family<'a>(&'a mut self, family_name: &LowercaseString, desc: &FontTemplateDescriptor)
                                -> Option<Arc<FontTemplateData>> {
        // TODO(Issue #188): look up localized font family names if canonical name not found
        // look up canonical name
        if self.local_families.contains_key(family_name) {
            debug!("FontList: Found font family with name={}", &**family_name);
            let s = self.local_families.get_mut(family_name).unwrap();

            if s.templates.is_empty() {
                for_each_variation(family_name, |path| {
                    s.add_template(Atom::from_slice(&path), None);
                });
            }

            // TODO(Issue #192: handle generic font families, like 'serif' and 'sans-serif'.
            // if such family exists, try to match style to a font
            let result = s.find_font_for_style(desc, &self.font_context);
            if result.is_some() {
                return result;
            }

            None
        } else {
            debug!("FontList: Couldn't find font family with name={}", &**family_name);
            None
        }
    }

    fn find_font_in_web_family<'a>(&'a mut self, family_name: &LowercaseString, desc: &FontTemplateDescriptor)
                                -> Option<Arc<FontTemplateData>> {
        if self.web_families.contains_key(family_name) {
            let family = self.web_families.get_mut(family_name).unwrap();
            let maybe_font = family.find_font_for_style(desc, &self.font_context);
            maybe_font
        } else {
            None
        }
    }

    fn find_font_template(&mut self, family: &LowercaseString, desc: &FontTemplateDescriptor)
                            -> Option<Arc<FontTemplateData>> {
        let transformed_family_name = self.transform_family(family);
        let mut maybe_template = self.find_font_in_web_family(&transformed_family_name, desc);
        if maybe_template.is_none() {
            maybe_template = self.find_font_in_local_family(&transformed_family_name, desc);
        }
        maybe_template
    }

    fn last_resort_font_template(&mut self, desc: &FontTemplateDescriptor)
                                        -> Arc<FontTemplateData> {
        let last_resort = last_resort_font_families();

        for family in &last_resort {
            let family = LowercaseString::new(family);
            let maybe_font_in_family = self.find_font_in_local_family(&family, desc);
            if maybe_font_in_family.is_some() {
                return maybe_font_in_family.unwrap();
            }
        }

        panic!("Unable to find any fonts that match (do you have fallback fonts installed?)");
    }
}

/// The public interface to the font cache task, used exclusively by
/// the per-thread/task FontContext structures.
#[derive(Clone)]
pub struct FontCacheTask {
    chan: Sender<Command>,
}

impl FontCacheTask {
    pub fn new(resource_task: ResourceTask) -> FontCacheTask {
        let (chan, port) = channel();

        let channel_to_self = chan.clone();
        spawn_named("FontCacheTask".to_owned(), move || {
            // TODO: Allow users to specify these.
            let mut generic_fonts = HashMap::with_capacity(5);
            add_generic_font(&mut generic_fonts, "serif", "Times New Roman");
            add_generic_font(&mut generic_fonts, "sans-serif", "Arial");
            add_generic_font(&mut generic_fonts, "cursive", "Apple Chancery");
            add_generic_font(&mut generic_fonts, "fantasy", "Papyrus");
            add_generic_font(&mut generic_fonts, "monospace", "Menlo");

            let mut cache = FontCache {
                port: port,
                channel_to_self: channel_to_self,
                generic_fonts: generic_fonts,
                local_families: HashMap::new(),
                web_families: HashMap::new(),
                font_context: FontContextHandle::new(),
                resource_task: resource_task,
            };

            cache.refresh_local_families();
            cache.run();
        });

        FontCacheTask {
            chan: chan,
        }
    }

    pub fn find_font_template(&self, family: String, desc: FontTemplateDescriptor)
                                                -> Option<Arc<FontTemplateData>> {

        let (response_chan, response_port) = channel();
        self.chan.send(Command::GetFontTemplate(family, desc, response_chan)).unwrap();

        let reply = response_port.recv().unwrap();

        match reply {
            Reply::GetFontTemplateReply(data) => {
                data
            }
        }
    }

    pub fn last_resort_font_template(&self, desc: FontTemplateDescriptor)
                                                -> Arc<FontTemplateData> {

        let (response_chan, response_port) = channel();
        self.chan.send(Command::GetLastResortFontTemplate(desc, response_chan)).unwrap();

        let reply = response_port.recv().unwrap();

        match reply {
            Reply::GetFontTemplateReply(data) => {
                data.unwrap()
            }
        }
    }

    pub fn add_web_font(&self, family: Atom, src: Source, sender: Sender<()>) {
        self.chan.send(Command::AddWebFont(family, src, sender)).unwrap();
    }

    pub fn exit(&self) {
        let (response_chan, response_port) = channel();
        self.chan.send(Command::Exit(response_chan)).unwrap();
        response_port.recv().unwrap();
    }
}

<commit_msg>Clean up use of Option in `font_cache_task`.<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/. */

use font_template::{FontTemplate, FontTemplateDescriptor};
use ipc_channel::ipc;
use ipc_channel::router::ROUTER;
use net_traits::{AsyncResponseTarget, PendingAsyncLoad, ResourceTask, ResponseAction};
use platform::font_context::FontContextHandle;
use platform::font_list::for_each_available_family;
use platform::font_list::for_each_variation;
use platform::font_list::last_resort_font_families;
use platform::font_list::system_default_family;
use platform::font_template::FontTemplateData;
use std::borrow::ToOwned;
use std::collections::HashMap;
use std::mem;
use std::sync::mpsc::{Sender, Receiver, channel};
use std::sync::{Arc, Mutex};
use string_cache::Atom;
use style::font_face::Source;
use url::Url;
use util::str::LowercaseString;
use util::task::spawn_named;

/// A list of font templates that make up a given font family.
struct FontFamily {
    templates: Vec<FontTemplate>,
}

impl FontFamily {
    fn new() -> FontFamily {
        FontFamily {
            templates: vec!(),
        }
    }

    /// Find a font in this family that matches a given descriptor.
    fn find_font_for_style<'a>(&'a mut self, desc: &FontTemplateDescriptor, fctx: &FontContextHandle)
                               -> Option<Arc<FontTemplateData>> {
        // TODO(Issue #189): optimize lookup for
        // regular/bold/italic/bolditalic with fixed offsets and a
        // static decision table for fallback between these values.

        // TODO(Issue #190): if not in the fast path above, do
        // expensive matching of weights, etc.
        for template in &mut self.templates {
            let maybe_template = template.data_for_descriptor(fctx, desc);
            if maybe_template.is_some() {
                return maybe_template;
            }
        }

        // If a request is made for a font family that exists,
        // pick the first valid font in the family if we failed
        // to find an exact match for the descriptor.
        for template in &mut self.templates {
            let maybe_template = template.get();
            if maybe_template.is_some() {
                return maybe_template;
            }
        }

        None
    }

    fn add_template(&mut self, identifier: Atom, maybe_data: Option<Vec<u8>>) {
        for template in &self.templates {
            if *template.identifier() == identifier {
                return;
            }
        }

        let template = FontTemplate::new(identifier, maybe_data);
        self.templates.push(template);
    }
}

/// Commands that the FontContext sends to the font cache task.
pub enum Command {
    GetFontTemplate(String, FontTemplateDescriptor, Sender<Reply>),
    GetLastResortFontTemplate(FontTemplateDescriptor, Sender<Reply>),
    AddWebFont(Atom, Source, Sender<()>),
    AddDownloadedWebFont(LowercaseString, Url, Vec<u8>, Sender<()>),
    Exit(Sender<()>),
}

/// Reply messages sent from the font cache task to the FontContext caller.
pub enum Reply {
    GetFontTemplateReply(Option<Arc<FontTemplateData>>),
}

/// The font cache task itself. It maintains a list of reference counted
/// font templates that are currently in use.
struct FontCache {
    port: Receiver<Command>,
    channel_to_self: Sender<Command>,
    generic_fonts: HashMap<LowercaseString, LowercaseString>,
    local_families: HashMap<LowercaseString, FontFamily>,
    web_families: HashMap<LowercaseString, FontFamily>,
    font_context: FontContextHandle,
    resource_task: ResourceTask,
}

fn add_generic_font(generic_fonts: &mut HashMap<LowercaseString, LowercaseString>,
                    generic_name: &str, mapped_name: &str) {
    let opt_system_default = system_default_family(generic_name);
    let family_name = match opt_system_default {
        Some(system_default) => LowercaseString::new(&system_default),
        None => LowercaseString::new(mapped_name),
    };
    generic_fonts.insert(LowercaseString::new(generic_name), family_name);
}

impl FontCache {
    fn run(&mut self) {
        loop {
            let msg = self.port.recv().unwrap();

            match msg {
                Command::GetFontTemplate(family, descriptor, result) => {
                    let family = LowercaseString::new(&family);
                    let maybe_font_template = self.find_font_template(&family, &descriptor);
                    result.send(Reply::GetFontTemplateReply(maybe_font_template)).unwrap();
                }
                Command::GetLastResortFontTemplate(descriptor, result) => {
                    let font_template = self.last_resort_font_template(&descriptor);
                    result.send(Reply::GetFontTemplateReply(Some(font_template))).unwrap();
                }
                Command::AddWebFont(family_name, src, result) => {
                    let family_name = LowercaseString::new(&family_name);
                    if !self.web_families.contains_key(&family_name) {
                        let family = FontFamily::new();
                        self.web_families.insert(family_name.clone(), family);
                    }

                    match src {
                        Source::Url(ref url_source) => {
                            let url = &url_source.url;
                            let load = PendingAsyncLoad::new(self.resource_task.clone(),
                                                             url.clone(),
                                                             None);
                            let (data_sender, data_receiver) = ipc::channel().unwrap();
                            let data_target = AsyncResponseTarget {
                                sender: data_sender,
                            };
                            load.load_async(data_target);
                            let channel_to_self = self.channel_to_self.clone();
                            let url = (*url).clone();
                            let bytes = Mutex::new(Vec::new());
                            ROUTER.add_route(data_receiver.to_opaque(), box move |message| {
                                let response: ResponseAction = message.to().unwrap();
                                match response {
                                    ResponseAction::HeadersAvailable(_) |
                                    ResponseAction::ResponseComplete(Err(_)) => {}
                                    ResponseAction::DataAvailable(new_bytes) => {
                                        bytes.lock().unwrap().extend(new_bytes.into_iter())
                                    }
                                    ResponseAction::ResponseComplete(Ok(_)) => {
                                        let mut bytes = bytes.lock().unwrap();
                                        let bytes = mem::replace(&mut *bytes, Vec::new());
                                        let command =
                                            Command::AddDownloadedWebFont(family_name.clone(),
                                                                          url.clone(),
                                                                          bytes,
                                                                          result.clone());
                                        channel_to_self.send(command).unwrap();
                                    }
                                }
                            });
                        }
                        Source::Local(ref local_family_name) => {
                            let family = &mut self.web_families.get_mut(&family_name).unwrap();
                            for_each_variation(&local_family_name, |path| {
                                family.add_template(Atom::from_slice(&path), None);
                            });
                            result.send(()).unwrap();
                        }
                    }
                }
                Command::AddDownloadedWebFont(family_name, url, bytes, result) => {
                    let family = &mut self.web_families.get_mut(&family_name).unwrap();
                    family.add_template(Atom::from_slice(&url.to_string()), Some(bytes));
                    drop(result.send(()));
                }
                Command::Exit(result) => {
                    result.send(()).unwrap();
                    break;
                }
            }
        }
    }

    fn refresh_local_families(&mut self) {
        self.local_families.clear();
        for_each_available_family(|family_name| {
            let family_name = LowercaseString::new(&family_name);
            if !self.local_families.contains_key(&family_name) {
                let family = FontFamily::new();
                self.local_families.insert(family_name, family);
            }
        });
    }

    fn transform_family(&self, family: &LowercaseString) -> LowercaseString {
        match self.generic_fonts.get(family) {
            None => family.clone(),
            Some(mapped_family) => (*mapped_family).clone()
        }
    }

    fn find_font_in_local_family<'a>(&'a mut self, family_name: &LowercaseString, desc: &FontTemplateDescriptor)
                                -> Option<Arc<FontTemplateData>> {
        // TODO(Issue #188): look up localized font family names if canonical name not found
        // look up canonical name
        if self.local_families.contains_key(family_name) {
            debug!("FontList: Found font family with name={}", &**family_name);
            let s = self.local_families.get_mut(family_name).unwrap();

            if s.templates.is_empty() {
                for_each_variation(family_name, |path| {
                    s.add_template(Atom::from_slice(&path), None);
                });
            }

            // TODO(Issue #192: handle generic font families, like 'serif' and 'sans-serif'.
            // if such family exists, try to match style to a font

            s.find_font_for_style(desc, &self.font_context)
        } else {
            debug!("FontList: Couldn't find font family with name={}", &**family_name);
            None
        }
    }

    fn find_font_in_web_family<'a>(&'a mut self, family_name: &LowercaseString, desc: &FontTemplateDescriptor)
                                -> Option<Arc<FontTemplateData>> {
        if self.web_families.contains_key(family_name) {
            let family = self.web_families.get_mut(family_name).unwrap();
            let maybe_font = family.find_font_for_style(desc, &self.font_context);
            maybe_font
        } else {
            None
        }
    }

    fn find_font_template(&mut self, family: &LowercaseString, desc: &FontTemplateDescriptor)
                            -> Option<Arc<FontTemplateData>> {
        let transformed_family_name = self.transform_family(family);
        self.find_font_in_web_family(&transformed_family_name, desc)
            .or_else(|| self.find_font_in_local_family(&transformed_family_name, desc))
    }

    fn last_resort_font_template(&mut self, desc: &FontTemplateDescriptor)
                                        -> Arc<FontTemplateData> {
        let last_resort = last_resort_font_families();

        for family in &last_resort {
            let family = LowercaseString::new(family);
            let maybe_font_in_family = self.find_font_in_local_family(&family, desc);
            if let Some(family) = maybe_font_in_family {
                return family;
            }
        }

        panic!("Unable to find any fonts that match (do you have fallback fonts installed?)");
    }
}

/// The public interface to the font cache task, used exclusively by
/// the per-thread/task FontContext structures.
#[derive(Clone)]
pub struct FontCacheTask {
    chan: Sender<Command>,
}

impl FontCacheTask {
    pub fn new(resource_task: ResourceTask) -> FontCacheTask {
        let (chan, port) = channel();

        let channel_to_self = chan.clone();
        spawn_named("FontCacheTask".to_owned(), move || {
            // TODO: Allow users to specify these.
            let mut generic_fonts = HashMap::with_capacity(5);
            add_generic_font(&mut generic_fonts, "serif", "Times New Roman");
            add_generic_font(&mut generic_fonts, "sans-serif", "Arial");
            add_generic_font(&mut generic_fonts, "cursive", "Apple Chancery");
            add_generic_font(&mut generic_fonts, "fantasy", "Papyrus");
            add_generic_font(&mut generic_fonts, "monospace", "Menlo");

            let mut cache = FontCache {
                port: port,
                channel_to_self: channel_to_self,
                generic_fonts: generic_fonts,
                local_families: HashMap::new(),
                web_families: HashMap::new(),
                font_context: FontContextHandle::new(),
                resource_task: resource_task,
            };

            cache.refresh_local_families();
            cache.run();
        });

        FontCacheTask {
            chan: chan,
        }
    }

    pub fn find_font_template(&self, family: String, desc: FontTemplateDescriptor)
                                                -> Option<Arc<FontTemplateData>> {

        let (response_chan, response_port) = channel();
        self.chan.send(Command::GetFontTemplate(family, desc, response_chan)).unwrap();

        let reply = response_port.recv().unwrap();

        match reply {
            Reply::GetFontTemplateReply(data) => {
                data
            }
        }
    }

    pub fn last_resort_font_template(&self, desc: FontTemplateDescriptor)
                                                -> Arc<FontTemplateData> {

        let (response_chan, response_port) = channel();
        self.chan.send(Command::GetLastResortFontTemplate(desc, response_chan)).unwrap();

        let reply = response_port.recv().unwrap();

        match reply {
            Reply::GetFontTemplateReply(data) => {
                data.unwrap()
            }
        }
    }

    pub fn add_web_font(&self, family: Atom, src: Source, sender: Sender<()>) {
        self.chan.send(Command::AddWebFont(family, src, sender)).unwrap();
    }

    pub fn exit(&self) {
        let (response_chan, response_port) = channel();
        self.chan.send(Command::Exit(response_chan)).unwrap();
        response_port.recv().unwrap();
    }
}

<|endoftext|>


<commit_before>//!
//! A static environment is used to store a fixed set of tools.
//!

use std::sync::Arc;
use std::collections::HashMap;
use std::result::Result;
use serde_json::*;

use super::super::tool::*;
use super::super::environment::*;
use super::toolset::*;
use super::basic_toolset::*;

///
/// A static environment just contains a fixed set of tools
///
pub struct StaticEnvironment {
    /// The tools in this environment
    tools: HashMap<String, Arc<Box<Tool>>>
}

///
/// Wrapper for a tool from a static environment
///
struct StaticEnvironmentTool {
    /// Reference to the tool within the environment
    tool: Arc<Box<Tool>>
}

impl Tool for StaticEnvironmentTool {
    fn invoke_json(&self, input: Value, environment: &Environment) -> Result<Value, Value> {
        (**self.tool).invoke_json(input, environment)
    }
}

impl Environment for StaticEnvironment {
    ///
    /// Retrieves a tool using a JSON interface by name
    ///
    fn get_json_tool(&self, name: &str) -> Result<Box<Tool>, RetrieveToolError> {
        let tool = self.tools.get(&String::from(name));

        match tool {
            Some(tool)  => Ok(Box::new(StaticEnvironmentTool { tool: tool.clone() })),
            None        => Err(RetrieveToolError::not_found())
        }
    }
}

impl StaticEnvironment {
    ///
    /// Creates a new static environment from a toolset
    ///
    pub fn from_toolset<T: ToolSet>(set: T, environment: &Environment) -> StaticEnvironment {
        let tools           = set.create_tools(environment);
        let mut tool_hash   = HashMap::new();

        for tool_and_name in tools {
            let (name, tool) = tool_and_name;

            tool_hash.insert(name, Arc::new(tool));
        }

        StaticEnvironment { tools: tool_hash }
    }

    ///
    /// Creates a new static environment containing a single tool
    ///
    pub fn from_tool<T: Tool+'static>(tool_name: &str, tool: T, environment: &Environment) -> StaticEnvironment {
        // Create a toolset containing this tool
        let single_toolset = BasicToolSet::from(vec![ (String::from(tool_name), tool) ]);

        // Pass through to from_toolset
        StaticEnvironment::from_toolset(single_toolset, environment)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use super::super::functional_tool::*;
    use super::super::empty_environment::*;
    use super::super::basic_toolset::*;

    #[test]
    fn can_get_tool_by_name() {
        let toolset = BasicToolSet::from(vec![
            ("add-1", make_pure_tool(|x: i32| { x+1 })),
            ("add-2", make_pure_tool(|x: i32| { x+2 }))
        ]);
        let environment = StaticEnvironment::from_toolset(toolset, &EmptyEnvironment::new());

        assert!(environment.get_json_tool("add-1").is_ok());
        assert!(environment.get_json_tool("add-2").is_ok());
    }

    #[test]
    fn error_for_missing_tool() {
        let toolset = BasicToolSet::from(vec![
            ("add-1", make_pure_tool(|x: i32| { x+1 })),
            ("add-2", make_pure_tool(|x: i32| { x+2 }))
        ]);
        let environment = StaticEnvironment::from_toolset(toolset, &EmptyEnvironment::new());

        assert!(environment.get_json_tool("add-3").is_err());
    }

    #[test]
    fn tools_are_right() {
        let toolset = BasicToolSet::from(vec![
            ("add-1", make_pure_tool(|x: i32| { x+1 })),
            ("add-2", make_pure_tool(|x: i32| { x+2 }))
        ]);
        let environment = StaticEnvironment::from_toolset(toolset, &EmptyEnvironment::new());

        let add1 = environment.get_json_tool("add-1").unwrap();
        let add2 = environment.get_json_tool("add-2").unwrap();

        assert!(add1.invoke_json(json![ 2 ], &environment) == Ok(json![ 3 ]));
        assert!(add2.invoke_json(json![ 2 ], &environment) == Ok(json![ 4 ]));
    }
}
<commit_msg>Fix warning<commit_after>//!
//! A static environment is used to store a fixed set of tools.
//!

use std::sync::Arc;
use std::collections::HashMap;
use std::result::Result;
use serde_json::*;

use super::super::tool::*;
use super::super::environment::*;
use super::toolset::*;
use super::basic_toolset::*;

///
/// A static environment just contains a fixed set of tools
///
pub struct StaticEnvironment {
    /// The tools in this environment
    tools: HashMap<String, Arc<Box<Tool>>>
}

///
/// Wrapper for a tool from a static environment
///
struct StaticEnvironmentTool {
    /// Reference to the tool within the environment
    tool: Arc<Box<Tool>>
}

impl Tool for StaticEnvironmentTool {
    fn invoke_json(&self, input: Value, environment: &Environment) -> Result<Value, Value> {
        (**self.tool).invoke_json(input, environment)
    }
}

impl Environment for StaticEnvironment {
    ///
    /// Retrieves a tool using a JSON interface by name
    ///
    fn get_json_tool(&self, name: &str) -> Result<Box<Tool>, RetrieveToolError> {
        let tool = self.tools.get(&String::from(name));

        match tool {
            Some(tool)  => Ok(Box::new(StaticEnvironmentTool { tool: tool.clone() })),
            None        => Err(RetrieveToolError::not_found())
        }
    }
}

impl StaticEnvironment {
    ///
    /// Creates a new static environment from a toolset
    ///
    pub fn from_toolset<T: ToolSet>(set: T, environment: &Environment) -> StaticEnvironment {
        let tools           = set.create_tools(environment);
        let mut tool_hash   = HashMap::new();

        for tool_and_name in tools {
            let (name, tool) = tool_and_name;

            tool_hash.insert(name, Arc::new(tool));
        }

        StaticEnvironment { tools: tool_hash }
    }

    ///
    /// Creates a new static environment containing a single tool
    ///
    pub fn from_tool<T: Tool+'static>(tool_name: &str, tool: T, environment: &Environment) -> StaticEnvironment {
        // Create a toolset containing this tool
        let single_toolset = BasicToolSet::from(vec![ (String::from(tool_name), tool) ]);

        // Pass through to from_toolset
        StaticEnvironment::from_toolset(single_toolset, environment)
    }
}

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

    #[test]
    fn can_get_tool_by_name() {
        let toolset = BasicToolSet::from(vec![
            ("add-1", make_pure_tool(|x: i32| { x+1 })),
            ("add-2", make_pure_tool(|x: i32| { x+2 }))
        ]);
        let environment = StaticEnvironment::from_toolset(toolset, &EmptyEnvironment::new());

        assert!(environment.get_json_tool("add-1").is_ok());
        assert!(environment.get_json_tool("add-2").is_ok());
    }

    #[test]
    fn error_for_missing_tool() {
        let toolset = BasicToolSet::from(vec![
            ("add-1", make_pure_tool(|x: i32| { x+1 })),
            ("add-2", make_pure_tool(|x: i32| { x+2 }))
        ]);
        let environment = StaticEnvironment::from_toolset(toolset, &EmptyEnvironment::new());

        assert!(environment.get_json_tool("add-3").is_err());
    }

    #[test]
    fn tools_are_right() {
        let toolset = BasicToolSet::from(vec![
            ("add-1", make_pure_tool(|x: i32| { x+1 })),
            ("add-2", make_pure_tool(|x: i32| { x+2 }))
        ]);
        let environment = StaticEnvironment::from_toolset(toolset, &EmptyEnvironment::new());

        let add1 = environment.get_json_tool("add-1").unwrap();
        let add2 = environment.get_json_tool("add-2").unwrap();

        assert!(add1.invoke_json(json![ 2 ], &environment) == Ok(json![ 3 ]));
        assert!(add2.invoke_json(json![ 2 ], &environment) == Ok(json![ 4 ]));
    }
}
<|endoftext|>


<commit_before>use syn::{
    Attribute,
    Meta,
    NestedMeta,
    Lit,
};
use regex::Regex;

// Get the nested items of a a #[graphql(...)] attribute.
pub fn get_graphl_attr(attrs: &Vec<Attribute>) -> Option<Vec<NestedMeta>> {
    for attr in attrs {
        match attr.interpret_meta() {
            Some(Meta::List(ref list)) if list.ident == "graphql" => {
                return Some(list.nested.iter().map(|x| x.clone()).collect());
            },
            _ => {}
        }
    }
    None
}

pub fn keyed_item_value(item: &NestedMeta, name: &str, must_be_string: bool) -> Option<String> {
    match item {
        &NestedMeta::Meta(Meta::NameValue(ref nameval)) if nameval.ident == name => {
            match &nameval.lit {
                &Lit::Str(ref strlit) => {
                    Some(strlit.value())
                },
                _ => if must_be_string {
                    panic!(format!(
                        "Invalid format for attribute \"{:?}\": expected a string",
                        item
                    ));
                } else {
                    None
                },
            }
        },
        _ => None,
    }
}

// Note: duplicated from juniper crate!
#[doc(hidden)]
pub fn to_camel_case(s: &str) -> String {
    let mut dest = String::new();

    for (i, part) in s.split('_').enumerate() {
        if i > 0 && part.len() == 1 {
            dest.push_str(&part.to_uppercase());
        } else if i > 0 && part.len() > 1 {
            let first = part.chars()
                .next()
                .unwrap()
                .to_uppercase()
                .collect::<String>();
            let second = &part[1..];

            dest.push_str(&first);
            dest.push_str(second);
        } else if i == 0 {
            dest.push_str(part);
        }
    }

    dest
}

#[test]
fn test_to_camel_case() {
    assert_eq!(&to_camel_case("test")[..], "test");
    assert_eq!(&to_camel_case("_test")[..], "Test");
    assert_eq!(&to_camel_case("first_second")[..], "firstSecond");
    assert_eq!(&to_camel_case("first_")[..], "first");
    assert_eq!(&to_camel_case("a_b_c")[..], "aBC");
    assert_eq!(&to_camel_case("a_bc")[..], "aBc");
    assert_eq!(&to_camel_case("a_b")[..], "aB");
    assert_eq!(&to_camel_case("a")[..], "a");
    assert_eq!(&to_camel_case("")[..], "");
}

pub(crate) fn to_upper_snake_case(s: &str) -> String {
    let mut last_lower = false;
    let mut upper = String::new();
    for c in s.chars() {
        if c == '_' {
            last_lower = false;
        } else if c.is_lowercase() {
            last_lower = true;
        } else if c.is_uppercase() {
            if last_lower {
                upper.push('_');
            }
            last_lower = false;
        }

        for u in c.to_uppercase() {
            upper.push(u);
        }
    }
    upper
}

#[test]
fn test_to_upper_snake_case() {
    assert_eq!(to_upper_snake_case("abc"), "ABC");
    assert_eq!(to_upper_snake_case("a_bc"), "A_BC");
    assert_eq!(to_upper_snake_case("ABC"), "ABC");
    assert_eq!(to_upper_snake_case("A_BC"), "A_BC");
    assert_eq!(to_upper_snake_case("SomeInput"), "SOME_INPUT");
    assert_eq!(to_upper_snake_case("someInput"), "SOME_INPUT");
    assert_eq!(to_upper_snake_case("someINpuT"), "SOME_INPU_T");
    assert_eq!(to_upper_snake_case("some_INpuT"), "SOME_INPU_T");
}

#[doc(hidden)]
pub fn is_valid_name(field_name: &str) -> bool {
    lazy_static!{
        static ref CAMELCASE: Regex = Regex::new("^[_A-Za-z][_0-9A-Za-z]*$").unwrap();
    }
    CAMELCASE.is_match(field_name)
}

#[test]
fn test_is_valid_name(){
    assert_eq!(is_valid_name("yesItIs"), true);
    assert_eq!(is_valid_name("NoitIsnt"), true);
    assert_eq!(is_valid_name("iso6301"), true);
    assert_eq!(is_valid_name("thisIsATest"), true);
    assert_eq!(is_valid_name("i6Op"), true);
    assert_eq!(is_valid_name("i!"), false);
    assert_eq!(is_valid_name(""), false);
    assert_eq!(is_valid_name("aTest"), true);
    assert_eq!(is_valid_name("__Atest90"), true);
}
<commit_msg>Rename variable to be more descriptive.<commit_after>use syn::{
    Attribute,
    Meta,
    NestedMeta,
    Lit,
};
use regex::Regex;

// Get the nested items of a a #[graphql(...)] attribute.
pub fn get_graphl_attr(attrs: &Vec<Attribute>) -> Option<Vec<NestedMeta>> {
    for attr in attrs {
        match attr.interpret_meta() {
            Some(Meta::List(ref list)) if list.ident == "graphql" => {
                return Some(list.nested.iter().map(|x| x.clone()).collect());
            },
            _ => {}
        }
    }
    None
}

pub fn keyed_item_value(item: &NestedMeta, name: &str, must_be_string: bool) -> Option<String> {
    match item {
        &NestedMeta::Meta(Meta::NameValue(ref nameval)) if nameval.ident == name => {
            match &nameval.lit {
                &Lit::Str(ref strlit) => {
                    Some(strlit.value())
                },
                _ => if must_be_string {
                    panic!(format!(
                        "Invalid format for attribute \"{:?}\": expected a string",
                        item
                    ));
                } else {
                    None
                },
            }
        },
        _ => None,
    }
}

// Note: duplicated from juniper crate!
#[doc(hidden)]
pub fn to_camel_case(s: &str) -> String {
    let mut dest = String::new();

    for (i, part) in s.split('_').enumerate() {
        if i > 0 && part.len() == 1 {
            dest.push_str(&part.to_uppercase());
        } else if i > 0 && part.len() > 1 {
            let first = part.chars()
                .next()
                .unwrap()
                .to_uppercase()
                .collect::<String>();
            let second = &part[1..];

            dest.push_str(&first);
            dest.push_str(second);
        } else if i == 0 {
            dest.push_str(part);
        }
    }

    dest
}

#[test]
fn test_to_camel_case() {
    assert_eq!(&to_camel_case("test")[..], "test");
    assert_eq!(&to_camel_case("_test")[..], "Test");
    assert_eq!(&to_camel_case("first_second")[..], "firstSecond");
    assert_eq!(&to_camel_case("first_")[..], "first");
    assert_eq!(&to_camel_case("a_b_c")[..], "aBC");
    assert_eq!(&to_camel_case("a_bc")[..], "aBc");
    assert_eq!(&to_camel_case("a_b")[..], "aB");
    assert_eq!(&to_camel_case("a")[..], "a");
    assert_eq!(&to_camel_case("")[..], "");
}

pub(crate) fn to_upper_snake_case(s: &str) -> String {
    let mut last_lower = false;
    let mut upper = String::new();
    for c in s.chars() {
        if c == '_' {
            last_lower = false;
        } else if c.is_lowercase() {
            last_lower = true;
        } else if c.is_uppercase() {
            if last_lower {
                upper.push('_');
            }
            last_lower = false;
        }

        for u in c.to_uppercase() {
            upper.push(u);
        }
    }
    upper
}

#[test]
fn test_to_upper_snake_case() {
    assert_eq!(to_upper_snake_case("abc"), "ABC");
    assert_eq!(to_upper_snake_case("a_bc"), "A_BC");
    assert_eq!(to_upper_snake_case("ABC"), "ABC");
    assert_eq!(to_upper_snake_case("A_BC"), "A_BC");
    assert_eq!(to_upper_snake_case("SomeInput"), "SOME_INPUT");
    assert_eq!(to_upper_snake_case("someInput"), "SOME_INPUT");
    assert_eq!(to_upper_snake_case("someINpuT"), "SOME_INPU_T");
    assert_eq!(to_upper_snake_case("some_INpuT"), "SOME_INPU_T");
}

#[doc(hidden)]
pub fn is_valid_name(field_name: &str) -> bool {
    lazy_static!{
        static ref GRAPHQL_NAME_SPEC: Regex = Regex::new("^[_A-Za-z][_0-9A-Za-z]*$").unwrap();
    }
    GRAPHQL_NAME_SPEC.is_match(field_name)
}

#[test]
fn test_is_valid_name(){
    assert_eq!(is_valid_name("yesItIs"), true);
    assert_eq!(is_valid_name("NoitIsnt"), true);
    assert_eq!(is_valid_name("iso6301"), true);
    assert_eq!(is_valid_name("thisIsATest"), true);
    assert_eq!(is_valid_name("i6Op"), true);
    assert_eq!(is_valid_name("i!"), false);
    assert_eq!(is_valid_name(""), false);
    assert_eq!(is_valid_name("aTest"), true);
    assert_eq!(is_valid_name("__Atest90"), true);
}
<|endoftext|>


<commit_before>use amq_protocol::frame::{GenError, Offset, gen_frame, parse_frame};
use log::{error, trace};
use mio::{Evented, Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use parking_lot::Mutex;

use std::{
  io::{self, Read, Write},
  sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
  },
  thread::{self, Builder as ThreadBuilder, JoinHandle},
  time::{Duration, Instant},
};

use crate::{
  buffer::Buffer,
  connection::Connection,
  connection_status::ConnectionState,
  error::{Error, ErrorKind},
};

const SOCKET:   Token = Token(1);
const DATA:     Token = Token(2);
const CONTINUE: Token = Token(3);

const FRAMES_STORAGE: usize = 32;

#[derive(Clone, Debug)]
pub(crate) struct IoLoopHandle {
  handle: Arc<Mutex<Option<JoinHandle<Result<(), Error>>>>>,
}

impl Default for IoLoopHandle {
  fn default() -> Self {
    Self { handle: Arc::new(Mutex::new(None)) }
  }
}

impl IoLoopHandle {
  pub(crate) fn register(&self, handle: JoinHandle<Result<(), Error>>) {
    *self.handle.lock() = Some(handle);
  }

  pub(crate) fn wait(&self) -> Result<(), Error> {
    if let Some(handle) = self.handle.lock().take() {
      handle.join().expect("io loop")?
    }
    Ok(())
  }
}

#[derive(Debug, PartialEq)]
enum Status {
  Initial,
  Setup,
  Stop,
}

pub(crate) struct IoLoop<T> {
  connection:     Connection,
  socket:         T,
  status:         Status,
  poll:           Poll,
  registration:   Registration,
  set_readiness:  SetReadiness,
  hb_handle:      Option<JoinHandle<()>>,
  frame_size:     usize,
  receive_buffer: Buffer,
  send_buffer:    Buffer,
  can_write:      bool,
  can_read:       bool,
  has_data:       bool,
  send_heartbeat: Arc<AtomicBool>,
}

impl<T: Evented + Read + Write + Send + 'static> IoLoop<T> {
  pub(crate) fn new(connection: Connection, socket: T) -> Result<Self, Error> {
    let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
    let (registration, set_readiness) = Registration::new2();
    let inner = Self {
      connection,
      socket,
      status:         Status::Initial,
      poll:           Poll::new().map_err(ErrorKind::IOError)?,
      registration,
      set_readiness,
      hb_handle:      None,
      frame_size,
      receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
      send_buffer:    Buffer::with_capacity(FRAMES_STORAGE * frame_size),
      can_write:      false,
      can_read:       false,
      has_data:       false,
      send_heartbeat: Arc::new(AtomicBool::new(false)),
    };
    inner.poll.register(&inner.socket, SOCKET, Ready::readable() | Ready::writable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    inner.poll.register(&inner.connection, DATA, Ready::readable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    inner.poll.register(&inner.registration, CONTINUE, Ready::readable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    Ok(inner)
  }

  fn start_heartbeat(&mut self, interval: Duration) -> Result<(), Error> {
    let connection    = self.connection.clone();
    let send_hartbeat = self.send_heartbeat.clone();
    let hb_handle = ThreadBuilder::new().name("heartbeat".to_owned()).spawn(move || {
      while connection.status().connected() {
        let start         = Instant::now();
        let mut remaining = interval;

        loop {
          thread::park_timeout(remaining);
          let elapsed = start.elapsed();
          if elapsed >= remaining {
            break
          }
          remaining -= interval - elapsed;
        }

        send_hartbeat.store(true, Ordering::Relaxed);
      }
    }).map_err(ErrorKind::IOError)?;
    self.hb_handle = Some(hb_handle);
    Ok(())
  }

  fn heartbeat(&mut self) -> Result<(), Error> {
    self.connection.send_heartbeat()?;
    self.send_heartbeat.store(false, Ordering::Relaxed);
    Ok(())
  }

  fn ensure_setup(&mut self) -> Result<(), Error> {
    if self.status != Status::Setup && self.connection.status().connected() {
      let frame_max = self.connection.configuration().frame_max() as usize;
      self.frame_size = std::cmp::max(self.frame_size, frame_max);
      self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
      self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
      let heartbeat = self.connection.configuration().heartbeat();
      if heartbeat != 0 {
        trace!("io_loop: start heartbeat");
        self.start_heartbeat(Duration::from_secs(heartbeat as u64))?;
        trace!("io_loop: heartbeat started");
      }
      self.status = Status::Setup;
    }
    Ok(())
  }

  fn wants_to_write(&self) -> bool {
    self.can_write && self.has_data
  }

  fn wants_to_read(&self) -> bool {
    self.can_read
  }

  fn should_continue(&self) -> bool {
    let connection_status = self.connection.status();
    (self.status == Status::Initial || connection_status.connected() || connection_status.closing()) && self.status != Status::Stop && !connection_status.errored()
  }

  pub(crate) fn run(mut self) -> Result<(), Error> {
    self.connection.clone().set_io_loop(ThreadBuilder::new().name("io_loop".to_owned()).spawn(move || {
      let mut events = Events::with_capacity(1024);
      while self.should_continue() {
        self.do_run(&mut events)?;
      }
      if let Some(hb_handle) = self.hb_handle.take() {
        hb_handle.thread().unpark();
        hb_handle.join().expect("heartbeat loop failed");
      }
      Ok(())
    }).map_err(ErrorKind::IOError)?);
    Ok(())
  }

  fn do_run(&mut self, events: &mut Events) -> Result<(), Error> {
    // First, update our internal state
    trace!("io_loop run");
    self.ensure_setup()?;
    trace!("io_loop poll");
    self.poll.poll(events, None).map_err(ErrorKind::IOError)?;
    trace!("io_loop poll done");
    for event in events.iter() {
      match event.token() {
        SOCKET    => {
          if event.readiness().is_readable() {
            self.can_read = true;
          }
          if event.readiness().is_writable() {
            self.can_write = true;
          }
        },
        DATA      => self.has_data = true,
        _         => {},
      }
    }

    // Then, actually run an iteration of the event loop
    trace!("io_loop do_run; can_read={}, can_write={}, has_data={}", self.can_read, self.can_write, self.has_data);
    loop {
      if self.send_heartbeat.load(Ordering::Relaxed) {
        self.heartbeat()?;
      }
      if self.wants_to_write() && !self.connection.status().blocked() {
        if let Err(e) = self.write_to_stream() {
          match e.kind() {
            ErrorKind::IOError(e) if e.kind() == io::ErrorKind::WouldBlock => self.can_write = false,
            _ => {
              error!("error writing: {:?}", e);
              if let ConnectionState::SentProtocolHeader(wait_handle, ..) = self.connection.status().state() {
                wait_handle.error(ErrorKind::ConnectionRefused.into());
                self.status = Status::Stop;
              }
              self.connection.set_error()?;
              return Err(e);
            }
          }
        }
        self.send_buffer.shift_unless_available(self.frame_size);
      }
      if self.connection.status().closed() {
        self.status = Status::Stop;
      }
      if self.should_continue() && self.wants_to_read() {
        if let Err(e) = self.read_from_stream() {
          match e.kind() {
            ErrorKind::IOError(e) if e.kind() == io::ErrorKind::WouldBlock => self.can_read = false,
            _ => {
              error!("error reading: {:?}", e);
              self.connection.set_error()?;
              return Err(e);
            }
          }
        }
        self.receive_buffer.shift_unless_available(self.frame_size);
      }
      if self.can_parse() {
        self.parse()?;
      }
      if !self.wants_to_read() || !self.wants_to_write() || self.status == Status::Stop || self.connection.status().errored() {
        if self.status != Status::Stop && (self.wants_to_read() || self.can_parse() || self.has_data) {
          trace!("io_loop send continue; can_read={}, can_write={}, has_data={}", self.can_read, self.can_write, self.has_data);
          self.set_readiness.set_readiness(Ready::readable()).map_err(ErrorKind::IOError)?;
        }
        break;
      }
    }
    trace!("io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}", self.can_read, self.can_write, self.has_data, self.status);
    Ok(())
  }

  fn can_parse(&self) -> bool {
    self.receive_buffer.available_data() > 0
  }

  fn write_to_stream(&mut self) -> Result<(), Error> {
    self.serialize()?;

    self.socket.write(&self.send_buffer.data()).map(|sz| {
      trace!("wrote {} bytes", sz);
      self.send_buffer.consume(sz);
    }).map_err(|e| ErrorKind::IOError(e).into())
  }

  fn read_from_stream(&mut self) -> Result<(), Error> {
    match self.connection.status().state() {
      ConnectionState::Closed => Ok(()),
      ConnectionState::Error  => Err(ErrorKind::InvalidConnectionState(ConnectionState::Error).into()),
      _                       => self.socket.read(&mut self.receive_buffer.space()).map(|sz| {
        trace!("read {} bytes", sz);
        self.receive_buffer.fill(sz);
      }).map_err(|e| ErrorKind::IOError(e).into()),
    }
  }

  fn serialize(&mut self) -> Result<(), Error> {
    if let Some((send_id, next_msg)) = self.connection.next_frame() {
      trace!("will write to buffer: {:?}", next_msg);
      match gen_frame(self.send_buffer.space(), &next_msg).map(|tup| tup.0) {
        Ok(sz) => {
          self.send_buffer.fill(sz);
          self.connection.mark_sent(send_id);
          Ok(())
        },
        Err(e) => {
          match e {
            GenError::BufferTooSmall(_) => {
              // Requeue msg
              self.connection.requeue_frame(send_id, next_msg)?;
              self.send_buffer.shift();
              Ok(())
            },
            GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
              error!("error generating frame: {:?}", e);
              self.connection.set_error()?;
              Err(ErrorKind::SerialisationError(e).into())
            }
          }
        }
      }
    } else {
      self.has_data = false;
      Ok(())
    }
  }

  fn parse(&mut self) -> Result<(), Error> {
    match parse_frame(self.receive_buffer.data()) {
      Ok((i, f)) => {
        let consumed = self.receive_buffer.data().offset(i);
        self.receive_buffer.consume(consumed);

        if let Err(e) = self.connection.handle_frame(f) {
          self.connection.set_error()?;
          Err(e)
        } else {
          Ok(())
        }
      },
      Err(e) => {
        if e.is_incomplete() {
          self.receive_buffer.shift();
          Ok(())
        } else {
          error!("parse error: {:?}", e);
          self.connection.set_error()?;
          Err(ErrorKind::ParsingError(format!("{:?}", e)).into())
        }
      }
    }
  }
}
<commit_msg>log when we send heartbeat<commit_after>use amq_protocol::frame::{GenError, Offset, gen_frame, parse_frame};
use log::{error, trace};
use mio::{Evented, Events, Poll, PollOpt, Ready, Registration, SetReadiness, Token};
use parking_lot::Mutex;

use std::{
  io::{self, Read, Write},
  sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
  },
  thread::{self, Builder as ThreadBuilder, JoinHandle},
  time::{Duration, Instant},
};

use crate::{
  buffer::Buffer,
  connection::Connection,
  connection_status::ConnectionState,
  error::{Error, ErrorKind},
};

const SOCKET:   Token = Token(1);
const DATA:     Token = Token(2);
const CONTINUE: Token = Token(3);

const FRAMES_STORAGE: usize = 32;

#[derive(Clone, Debug)]
pub(crate) struct IoLoopHandle {
  handle: Arc<Mutex<Option<JoinHandle<Result<(), Error>>>>>,
}

impl Default for IoLoopHandle {
  fn default() -> Self {
    Self { handle: Arc::new(Mutex::new(None)) }
  }
}

impl IoLoopHandle {
  pub(crate) fn register(&self, handle: JoinHandle<Result<(), Error>>) {
    *self.handle.lock() = Some(handle);
  }

  pub(crate) fn wait(&self) -> Result<(), Error> {
    if let Some(handle) = self.handle.lock().take() {
      handle.join().expect("io loop")?
    }
    Ok(())
  }
}

#[derive(Debug, PartialEq)]
enum Status {
  Initial,
  Setup,
  Stop,
}

pub(crate) struct IoLoop<T> {
  connection:     Connection,
  socket:         T,
  status:         Status,
  poll:           Poll,
  registration:   Registration,
  set_readiness:  SetReadiness,
  hb_handle:      Option<JoinHandle<()>>,
  frame_size:     usize,
  receive_buffer: Buffer,
  send_buffer:    Buffer,
  can_write:      bool,
  can_read:       bool,
  has_data:       bool,
  send_heartbeat: Arc<AtomicBool>,
}

impl<T: Evented + Read + Write + Send + 'static> IoLoop<T> {
  pub(crate) fn new(connection: Connection, socket: T) -> Result<Self, Error> {
    let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
    let (registration, set_readiness) = Registration::new2();
    let inner = Self {
      connection,
      socket,
      status:         Status::Initial,
      poll:           Poll::new().map_err(ErrorKind::IOError)?,
      registration,
      set_readiness,
      hb_handle:      None,
      frame_size,
      receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
      send_buffer:    Buffer::with_capacity(FRAMES_STORAGE * frame_size),
      can_write:      false,
      can_read:       false,
      has_data:       false,
      send_heartbeat: Arc::new(AtomicBool::new(false)),
    };
    inner.poll.register(&inner.socket, SOCKET, Ready::readable() | Ready::writable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    inner.poll.register(&inner.connection, DATA, Ready::readable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    inner.poll.register(&inner.registration, CONTINUE, Ready::readable(), PollOpt::edge()).map_err(ErrorKind::IOError)?;
    Ok(inner)
  }

  fn start_heartbeat(&mut self, interval: Duration) -> Result<(), Error> {
    let connection    = self.connection.clone();
    let send_hartbeat = self.send_heartbeat.clone();
    let hb_handle = ThreadBuilder::new().name("heartbeat".to_owned()).spawn(move || {
      while connection.status().connected() {
        let start         = Instant::now();
        let mut remaining = interval;

        loop {
          thread::park_timeout(remaining);
          let elapsed = start.elapsed();
          if elapsed >= remaining {
            break
          }
          remaining -= interval - elapsed;
        }

        send_hartbeat.store(true, Ordering::Relaxed);
      }
    }).map_err(ErrorKind::IOError)?;
    self.hb_handle = Some(hb_handle);
    Ok(())
  }

  fn heartbeat(&mut self) -> Result<(), Error> {
    trace!("send heartbeat");
    self.connection.send_heartbeat()?;
    self.send_heartbeat.store(false, Ordering::Relaxed);
    Ok(())
  }

  fn ensure_setup(&mut self) -> Result<(), Error> {
    if self.status != Status::Setup && self.connection.status().connected() {
      let frame_max = self.connection.configuration().frame_max() as usize;
      self.frame_size = std::cmp::max(self.frame_size, frame_max);
      self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
      self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
      let heartbeat = self.connection.configuration().heartbeat();
      if heartbeat != 0 {
        trace!("io_loop: start heartbeat");
        self.start_heartbeat(Duration::from_secs(heartbeat as u64))?;
        trace!("io_loop: heartbeat started");
      }
      self.status = Status::Setup;
    }
    Ok(())
  }

  fn wants_to_write(&self) -> bool {
    self.can_write && self.has_data
  }

  fn wants_to_read(&self) -> bool {
    self.can_read
  }

  fn should_continue(&self) -> bool {
    let connection_status = self.connection.status();
    (self.status == Status::Initial || connection_status.connected() || connection_status.closing()) && self.status != Status::Stop && !connection_status.errored()
  }

  pub(crate) fn run(mut self) -> Result<(), Error> {
    self.connection.clone().set_io_loop(ThreadBuilder::new().name("io_loop".to_owned()).spawn(move || {
      let mut events = Events::with_capacity(1024);
      while self.should_continue() {
        self.do_run(&mut events)?;
      }
      if let Some(hb_handle) = self.hb_handle.take() {
        hb_handle.thread().unpark();
        hb_handle.join().expect("heartbeat loop failed");
      }
      Ok(())
    }).map_err(ErrorKind::IOError)?);
    Ok(())
  }

  fn do_run(&mut self, events: &mut Events) -> Result<(), Error> {
    // First, update our internal state
    trace!("io_loop run");
    self.ensure_setup()?;
    trace!("io_loop poll");
    self.poll.poll(events, None).map_err(ErrorKind::IOError)?;
    trace!("io_loop poll done");
    for event in events.iter() {
      match event.token() {
        SOCKET    => {
          if event.readiness().is_readable() {
            self.can_read = true;
          }
          if event.readiness().is_writable() {
            self.can_write = true;
          }
        },
        DATA      => self.has_data = true,
        _         => {},
      }
    }

    // Then, actually run an iteration of the event loop
    trace!("io_loop do_run; can_read={}, can_write={}, has_data={}", self.can_read, self.can_write, self.has_data);
    loop {
      if self.send_heartbeat.load(Ordering::Relaxed) {
        self.heartbeat()?;
      }
      if self.wants_to_write() && !self.connection.status().blocked() {
        if let Err(e) = self.write_to_stream() {
          match e.kind() {
            ErrorKind::IOError(e) if e.kind() == io::ErrorKind::WouldBlock => self.can_write = false,
            _ => {
              error!("error writing: {:?}", e);
              if let ConnectionState::SentProtocolHeader(wait_handle, ..) = self.connection.status().state() {
                wait_handle.error(ErrorKind::ConnectionRefused.into());
                self.status = Status::Stop;
              }
              self.connection.set_error()?;
              return Err(e);
            }
          }
        }
        self.send_buffer.shift_unless_available(self.frame_size);
      }
      if self.connection.status().closed() {
        self.status = Status::Stop;
      }
      if self.should_continue() && self.wants_to_read() {
        if let Err(e) = self.read_from_stream() {
          match e.kind() {
            ErrorKind::IOError(e) if e.kind() == io::ErrorKind::WouldBlock => self.can_read = false,
            _ => {
              error!("error reading: {:?}", e);
              self.connection.set_error()?;
              return Err(e);
            }
          }
        }
        self.receive_buffer.shift_unless_available(self.frame_size);
      }
      if self.can_parse() {
        self.parse()?;
      }
      if !self.wants_to_read() || !self.wants_to_write() || self.status == Status::Stop || self.connection.status().errored() {
        if self.status != Status::Stop && (self.wants_to_read() || self.can_parse() || self.has_data) {
          trace!("io_loop send continue; can_read={}, can_write={}, has_data={}", self.can_read, self.can_write, self.has_data);
          self.set_readiness.set_readiness(Ready::readable()).map_err(ErrorKind::IOError)?;
        }
        break;
      }
    }
    trace!("io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}", self.can_read, self.can_write, self.has_data, self.status);
    Ok(())
  }

  fn can_parse(&self) -> bool {
    self.receive_buffer.available_data() > 0
  }

  fn write_to_stream(&mut self) -> Result<(), Error> {
    self.serialize()?;

    self.socket.write(&self.send_buffer.data()).map(|sz| {
      trace!("wrote {} bytes", sz);
      self.send_buffer.consume(sz);
    }).map_err(|e| ErrorKind::IOError(e).into())
  }

  fn read_from_stream(&mut self) -> Result<(), Error> {
    match self.connection.status().state() {
      ConnectionState::Closed => Ok(()),
      ConnectionState::Error  => Err(ErrorKind::InvalidConnectionState(ConnectionState::Error).into()),
      _                       => self.socket.read(&mut self.receive_buffer.space()).map(|sz| {
        trace!("read {} bytes", sz);
        self.receive_buffer.fill(sz);
      }).map_err(|e| ErrorKind::IOError(e).into()),
    }
  }

  fn serialize(&mut self) -> Result<(), Error> {
    if let Some((send_id, next_msg)) = self.connection.next_frame() {
      trace!("will write to buffer: {:?}", next_msg);
      match gen_frame(self.send_buffer.space(), &next_msg).map(|tup| tup.0) {
        Ok(sz) => {
          self.send_buffer.fill(sz);
          self.connection.mark_sent(send_id);
          Ok(())
        },
        Err(e) => {
          match e {
            GenError::BufferTooSmall(_) => {
              // Requeue msg
              self.connection.requeue_frame(send_id, next_msg)?;
              self.send_buffer.shift();
              Ok(())
            },
            GenError::InvalidOffset | GenError::CustomError(_) | GenError::NotYetImplemented => {
              error!("error generating frame: {:?}", e);
              self.connection.set_error()?;
              Err(ErrorKind::SerialisationError(e).into())
            }
          }
        }
      }
    } else {
      self.has_data = false;
      Ok(())
    }
  }

  fn parse(&mut self) -> Result<(), Error> {
    match parse_frame(self.receive_buffer.data()) {
      Ok((i, f)) => {
        let consumed = self.receive_buffer.data().offset(i);
        self.receive_buffer.consume(consumed);

        if let Err(e) = self.connection.handle_frame(f) {
          self.connection.set_error()?;
          Err(e)
        } else {
          Ok(())
        }
      },
      Err(e) => {
        if e.is_incomplete() {
          self.receive_buffer.shift();
          Ok(())
        } else {
          error!("parse error: {:?}", e);
          self.connection.set_error()?;
          Err(ErrorKind::ParsingError(format!("{:?}", e)).into())
        }
      }
    }
  }
}
<|endoftext|>