<commit_before>use crate::hir::def::Namespace;
use crate::hir::{self, Body, FunctionRetTy, Expr, ExprKind, HirId, Local, Pat};
use crate::hir::intravisit::{self, Visitor, NestedVisitorMap};
use crate::infer::InferCtxt;
use crate::infer::type_variable::TypeVariableOriginKind;
use crate::ty::{self, Ty, Infer, TyVar};
use crate::ty::print::Print;
use syntax::source_map::DesugaringKind;
use syntax_pos::Span;
use errors::{Applicability, DiagnosticBuilder};

struct FindLocalByTypeVisitor<'a, 'tcx> {
    infcx: &'a InferCtxt<'a, 'tcx>,
    target_ty: Ty<'tcx>,
    hir_map: &'a hir::map::Map<'tcx>,
    found_local_pattern: Option<&'tcx Pat>,
    found_arg_pattern: Option<&'tcx Pat>,
    found_ty: Option<Ty<'tcx>>,
    found_closure: Option<&'tcx ExprKind>,
}

impl<'a, 'tcx> FindLocalByTypeVisitor<'a, 'tcx> {
    fn new(
        infcx: &'a InferCtxt<'a, 'tcx>,
        target_ty: Ty<'tcx>,
        hir_map: &'a hir::map::Map<'tcx>,
    ) -> FindLocalByTypeVisitor<'a, 'tcx> {
        FindLocalByTypeVisitor {
            infcx,
            target_ty,
            hir_map,
            found_local_pattern: None,
            found_arg_pattern: None,
            found_ty: None,
            found_closure: None,
        }
    }

    fn node_matches_type(&mut self, hir_id: HirId) -> Option<Ty<'tcx>> {
        let ty_opt = self.infcx.in_progress_tables.and_then(|tables| {
            tables.borrow().node_type_opt(hir_id)
        });
        match ty_opt {
            Some(ty) => {
                let ty = self.infcx.resolve_vars_if_possible(&ty);
                if ty.walk().any(|inner_ty| {
                    inner_ty == self.target_ty || match (&inner_ty.sty, &self.target_ty.sty) {
                        (&Infer(TyVar(a_vid)), &Infer(TyVar(b_vid))) => {
                            self.infcx
                                .type_variables
                                .borrow_mut()
                                .sub_unified(a_vid, b_vid)
                        }
                        _ => false,
                    }
                }) {
                    Some(ty)
                } else {
                    None
                }
            }
            None => None,
        }
    }
}

impl<'a, 'tcx> Visitor<'tcx> for FindLocalByTypeVisitor<'a, 'tcx> {
    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
        NestedVisitorMap::OnlyBodies(&self.hir_map)
    }

    fn visit_local(&mut self, local: &'tcx Local) {
        if let (None, Some(ty)) = (self.found_local_pattern, self.node_matches_type(local.hir_id)) {
            self.found_local_pattern = Some(&*local.pat);
            self.found_ty = Some(ty);
        }
        intravisit::walk_local(self, local);
    }

    fn visit_body(&mut self, body: &'tcx Body) {
        for argument in &body.arguments {
            if let (None, Some(ty)) = (
                self.found_arg_pattern,
                self.node_matches_type(argument.hir_id),
            ) {
                self.found_arg_pattern = Some(&*argument.pat);
                self.found_ty = Some(ty);
            }
        }
        intravisit::walk_body(self, body);
    }

    fn visit_expr(&mut self, expr: &'tcx Expr) {
        if let (ExprKind::Closure(_, _fn_decl, _id, _sp, _), Some(_)) = (
            &expr.node,
            self.node_matches_type(expr.hir_id),
        ) {
            self.found_closure = Some(&expr.node);
        }
        intravisit::walk_expr(self, expr);
    }
}

impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
    pub fn extract_type_name(
        &self,
        ty: Ty<'tcx>,
        highlight: Option<ty::print::RegionHighlightMode>,
    ) -> String {
        if let ty::Infer(ty::TyVar(ty_vid)) = ty.sty {
            let ty_vars = self.type_variables.borrow();
            if let TypeVariableOriginKind::TypeParameterDefinition(name) =
                ty_vars.var_origin(ty_vid).kind {
                return name.to_string();
            }
        }

        let mut s = String::new();
        let mut printer = ty::print::FmtPrinter::new(self.tcx, &mut s, Namespace::TypeNS);
        if let Some(highlight) = highlight {
            printer.region_highlight_mode = highlight;
        }
        let _ = ty.print(printer);
        s
    }

    pub fn need_type_info_err(
        &self,
        body_id: Option<hir::BodyId>,
        span: Span,
        ty: Ty<'tcx>,
    ) -> DiagnosticBuilder<'tcx> {
        let ty = self.resolve_vars_if_possible(&ty);
        let name = self.extract_type_name(&ty, None);

        let mut local_visitor = FindLocalByTypeVisitor::new(&self, ty, &self.tcx.hir());
        let ty_to_string = |ty: Ty<'tcx>| -> String {
            let mut s = String::new();
            let mut printer = ty::print::FmtPrinter::new(self.tcx, &mut s, Namespace::TypeNS);
            let ty_vars = self.type_variables.borrow();
            let getter = move |ty_vid| {
                if let TypeVariableOriginKind::TypeParameterDefinition(name) =
                    ty_vars.var_origin(ty_vid).kind {
                    return Some(name.to_string());
                }
                None
            };
            printer.name_resolver = Some(Box::new(&getter));
            let _ = ty.print(printer);
            s
        };

        if let Some(body_id) = body_id {
            let expr = self.tcx.hir().expect_expr(body_id.hir_id);
            local_visitor.visit_expr(expr);
        }
        let err_span = if let Some(pattern) = local_visitor.found_arg_pattern {
            pattern.span
        } else {
            span
        };

        let ty_msg = match local_visitor.found_ty {
            Some(ty::TyS { sty: ty::Closure(def_id, substs), .. }) => {
                let fn_sig = substs.closure_sig(*def_id, self.tcx);
                let args = fn_sig.inputs()
                    .skip_binder()
                    .iter()
                    .next()
                    .map(|args| args.tuple_fields()
                        .map(|arg| arg.to_string())
                        .collect::<Vec<_>>().join(", "))
                    .unwrap_or_default();
                let ret = fn_sig.output().skip_binder().to_string();
                format!(" for the closure `fn({}) -> {}`", args, ret)
            }
            Some(ty) if &ty.to_string() != "_" &&
                // FIXME: Remove this check after `impl_trait_in_bindings` is stabilized. #63527
                (!ty.is_impl_trait() || self.tcx.features().impl_trait_in_bindings) =>
            {
                let ty = ty_to_string(ty);
                format!(" for `{}`", ty)
            }
            _ => String::new(),
        };

        // When `name` corresponds to a type argument, show the path of the full type we're
        // trying to infer. In the following example, `ty_msg` contains
        // " in `std::result::Result<i32, E>`":
        // ```
        // error[E0282]: type annotations needed for `std::result::Result<i32, E>`
        //  --> file.rs:L:CC
        //   |
        // L |     let b = Ok(4);
        //   |         -   ^^ cannot infer type for `E` in `std::result::Result<i32, E>`
        //   |         |
        //   |         consider giving `b` the explicit type `std::result::Result<i32, E>`, where
        //   |         the type parameter `E` is specified
        // ```
        let mut err = struct_span_err!(
            self.tcx.sess,
            err_span,
            E0282,
            "type annotations needed{}",
            ty_msg,
        );

        let suffix = match local_visitor.found_ty {
            Some(ty::TyS { sty: ty::Closure(def_id, substs), .. }) => {
                let fn_sig = substs.closure_sig(*def_id, self.tcx);
                let ret = fn_sig.output().skip_binder().to_string();

                if let Some(ExprKind::Closure(_, decl, body_id, ..)) = local_visitor.found_closure {
                    let (arrow, post) = match decl.output {
                        FunctionRetTy::DefaultReturn(_) => ("-> ", " "),
                        _ => ("", ""),
                    };
                    if let Some(body) = self.tcx.hir().krate().bodies.get(body_id) {
                        let suggestion = match body.value.node {
                            ExprKind::Block(..) => {
                                vec![(decl.output.span(), format!("{}{}{}", arrow, ret, post))]
                            }
                            _ => {
                                vec![
                                    (decl.output.span(), format!("{}{}{}{{ ", arrow, ret, post)),
                                    (body.value.span.shrink_to_hi(), " }".to_string()),
                                ]
                            }
                        };
                        err.multipart_suggestion(
                            "give this closure an explicit return type without `_` placeholders",
                            suggestion,
                            Applicability::HasPlaceholders,
                        );
                        err.span_label(span, InferCtxt::missing_type_msg(&name));
                        return err;
                    }
                }

                // This shouldn't be reachable, but just in case we leave a reasonable fallback.
                let args = fn_sig.inputs()
                    .skip_binder()
                    .iter()
                    .next()
                    .map(|args| args.tuple_fields()
                        .map(|arg| arg.to_string())
                        .collect::<Vec<_>>().join(", "))
                    .unwrap_or_default();
                // This suggestion is incomplete, as the user will get further type inference
                // errors due to the `_` placeholders and the introduction of `Box`, but it does
                // nudge them in the right direction.
                format!("a boxed closure type like `Box<dyn Fn({}) -> {}>`", args, ret)
            }
            Some(ty) if &ty.to_string() != "_" &&
                name == "_" &&
                // FIXME: Remove this check after `impl_trait_in_bindings` is stabilized. #63527
                (!ty.is_impl_trait() || self.tcx.features().impl_trait_in_bindings) =>
            {
                let ty = ty_to_string(ty);
                format!("the explicit type `{}`, with the type parameters specified", ty)
            }
            Some(ty) if &ty.to_string() != "_" &&
                ty.to_string() != name &&
                // FIXME: Remove this check after `impl_trait_in_bindings` is stabilized. #63527
                (!ty.is_impl_trait() || self.tcx.features().impl_trait_in_bindings) =>
            {
                let ty = ty_to_string(ty);
                format!(
                    "the explicit type `{}`, where the type parameter `{}` is specified",
                    ty,
                    name,
                )
            }
            _ => "a type".to_string(),
        };

        if let Some(pattern) = local_visitor.found_arg_pattern {
            // We don't want to show the default label for closures.
            //
            // So, before clearing, the output would look something like this:
            // ```
            // let x = |_| {  };
            //          -  ^^^^ cannot infer type for `[_; 0]`
            //          |
            //          consider giving this closure parameter a type
            // ```
            //
            // After clearing, it looks something like this:
            // ```
            // let x = |_| {  };
            //          ^ consider giving this closure parameter the type `[_; 0]`
            //            with the type parameter `_` specified
            // ```
            err.span_label(
                pattern.span,
                format!("consider giving this closure parameter {}", suffix),
            );
        } else if let Some(pattern) = local_visitor.found_local_pattern {
            if let Some(simple_ident) = pattern.simple_ident() {
                match pattern.span.desugaring_kind() {
                    None => {
                        err.span_label(
                            pattern.span,
                            format!("consider giving `{}` {}", simple_ident, suffix),
                        );
                    }
                    Some(DesugaringKind::ForLoop) => {
                        err.span_label(
                            pattern.span,
                            "the element type for this iterator is not specified".to_string(),
                        );
                    }
                    _ => {}
                }
            } else {
                err.span_label(pattern.span, format!("consider giving this pattern {}", suffix));
            }
        }
        if !err.span.span_labels().iter().any(|span_label| {
                span_label.label.is_some() && span_label.span == span
            }) && local_visitor.found_arg_pattern.is_none()
        { // Avoid multiple labels pointing at `span`.
            err.span_label(span, InferCtxt::missing_type_msg(&name));
        }

        err
    }

    pub fn need_type_info_err_in_generator(
        &self,
        kind: hir::GeneratorKind,
        span: Span,
        ty: Ty<'tcx>,
    ) -> DiagnosticBuilder<'tcx> {
        let ty = self.resolve_vars_if_possible(&ty);
        let name = self.extract_type_name(&ty, None);
        let mut err = struct_span_err!(
            self.tcx.sess, span, E0698, "type inside {} must be known in this context", kind,
        );
        err.span_label(span, InferCtxt::missing_type_msg(&name));
        err
    }

    fn missing_type_msg(type_name: &str) -> String {
        if type_name == "_" {
            "cannot infer type".to_owned()
        } else {
            format!("cannot infer type for `{}`", type_name)
        }
    }
}
<commit_msg>review comments<commit_after>use crate::hir::def::Namespace;
use crate::hir::{self, Body, FunctionRetTy, Expr, ExprKind, HirId, Local, Pat};
use crate::hir::intravisit::{self, Visitor, NestedVisitorMap};
use crate::infer::InferCtxt;
use crate::infer::type_variable::TypeVariableOriginKind;
use crate::ty::{self, Ty, Infer, TyVar};
use crate::ty::print::Print;
use syntax::source_map::DesugaringKind;
use syntax_pos::Span;
use errors::{Applicability, DiagnosticBuilder};

struct FindLocalByTypeVisitor<'a, 'tcx> {
    infcx: &'a InferCtxt<'a, 'tcx>,
    target_ty: Ty<'tcx>,
    hir_map: &'a hir::map::Map<'tcx>,
    found_local_pattern: Option<&'tcx Pat>,
    found_arg_pattern: Option<&'tcx Pat>,
    found_ty: Option<Ty<'tcx>>,
    found_closure: Option<&'tcx ExprKind>,
}

impl<'a, 'tcx> FindLocalByTypeVisitor<'a, 'tcx> {
    fn new(
        infcx: &'a InferCtxt<'a, 'tcx>,
        target_ty: Ty<'tcx>,
        hir_map: &'a hir::map::Map<'tcx>,
    ) -> Self {
        Self {
            infcx,
            target_ty,
            hir_map,
            found_local_pattern: None,
            found_arg_pattern: None,
            found_ty: None,
            found_closure: None,
        }
    }

    fn node_matches_type(&mut self, hir_id: HirId) -> Option<Ty<'tcx>> {
        let ty_opt = self.infcx.in_progress_tables.and_then(|tables| {
            tables.borrow().node_type_opt(hir_id)
        });
        match ty_opt {
            Some(ty) => {
                let ty = self.infcx.resolve_vars_if_possible(&ty);
                if ty.walk().any(|inner_ty| {
                    inner_ty == self.target_ty || match (&inner_ty.sty, &self.target_ty.sty) {
                        (&Infer(TyVar(a_vid)), &Infer(TyVar(b_vid))) => {
                            self.infcx
                                .type_variables
                                .borrow_mut()
                                .sub_unified(a_vid, b_vid)
                        }
                        _ => false,
                    }
                }) {
                    Some(ty)
                } else {
                    None
                }
            }
            None => None,
        }
    }
}

impl<'a, 'tcx> Visitor<'tcx> for FindLocalByTypeVisitor<'a, 'tcx> {
    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
        NestedVisitorMap::OnlyBodies(&self.hir_map)
    }

    fn visit_local(&mut self, local: &'tcx Local) {
        if let (None, Some(ty)) = (self.found_local_pattern, self.node_matches_type(local.hir_id)) {
            self.found_local_pattern = Some(&*local.pat);
            self.found_ty = Some(ty);
        }
        intravisit::walk_local(self, local);
    }

    fn visit_body(&mut self, body: &'tcx Body) {
        for argument in &body.arguments {
            if let (None, Some(ty)) = (
                self.found_arg_pattern,
                self.node_matches_type(argument.hir_id),
            ) {
                self.found_arg_pattern = Some(&*argument.pat);
                self.found_ty = Some(ty);
            }
        }
        intravisit::walk_body(self, body);
    }

    fn visit_expr(&mut self, expr: &'tcx Expr) {
        if let (ExprKind::Closure(_, _fn_decl, _id, _sp, _), Some(_)) = (
            &expr.node,
            self.node_matches_type(expr.hir_id),
        ) {
            self.found_closure = Some(&expr.node);
        }
        intravisit::walk_expr(self, expr);
    }
}

/// Suggest giving an appropriate return type to a closure expression.
fn closure_return_type_suggestion(
    span: Span,
    err: &mut DiagnosticBuilder<'_>,
    output: &FunctionRetTy,
    body: &Body,
    name: &str,
    ret: &str,
) {
    let (arrow, post) = match output {
        FunctionRetTy::DefaultReturn(_) => ("-> ", " "),
        _ => ("", ""),
    };
    let suggestion = match body.value.node {
        ExprKind::Block(..) => {
            vec![(output.span(), format!("{}{}{}", arrow, ret, post))]
        }
        _ => {
            vec![
                (output.span(), format!("{}{}{}{{ ", arrow, ret, post)),
                (body.value.span.shrink_to_hi(), " }".to_string()),
            ]
        }
    };
    err.multipart_suggestion(
        "give this closure an explicit return type without `_` placeholders",
        suggestion,
        Applicability::HasPlaceholders,
    );
    err.span_label(span, InferCtxt::missing_type_msg(&name));
}

/// Given a closure signature, return a `String` containing a list of all its argument types.
fn closure_args(fn_sig: &ty::PolyFnSig<'_>) -> String {
    fn_sig.inputs()
        .skip_binder()
        .iter()
        .next()
        .map(|args| args.tuple_fields()
            .map(|arg| arg.to_string())
            .collect::<Vec<_>>().join(", "))
        .unwrap_or_default()
}

impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
    pub fn extract_type_name(
        &self,
        ty: Ty<'tcx>,
        highlight: Option<ty::print::RegionHighlightMode>,
    ) -> String {
        if let ty::Infer(ty::TyVar(ty_vid)) = ty.sty {
            let ty_vars = self.type_variables.borrow();
            if let TypeVariableOriginKind::TypeParameterDefinition(name) =
                ty_vars.var_origin(ty_vid).kind {
                return name.to_string();
            }
        }

        let mut s = String::new();
        let mut printer = ty::print::FmtPrinter::new(self.tcx, &mut s, Namespace::TypeNS);
        if let Some(highlight) = highlight {
            printer.region_highlight_mode = highlight;
        }
        let _ = ty.print(printer);
        s
    }

    pub fn need_type_info_err(
        &self,
        body_id: Option<hir::BodyId>,
        span: Span,
        ty: Ty<'tcx>,
    ) -> DiagnosticBuilder<'tcx> {
        let ty = self.resolve_vars_if_possible(&ty);
        let name = self.extract_type_name(&ty, None);

        let mut local_visitor = FindLocalByTypeVisitor::new(&self, ty, &self.tcx.hir());
        let ty_to_string = |ty: Ty<'tcx>| -> String {
            let mut s = String::new();
            let mut printer = ty::print::FmtPrinter::new(self.tcx, &mut s, Namespace::TypeNS);
            let ty_vars = self.type_variables.borrow();
            let getter = move |ty_vid| {
                if let TypeVariableOriginKind::TypeParameterDefinition(name) =
                    ty_vars.var_origin(ty_vid).kind {
                    return Some(name.to_string());
                }
                None
            };
            printer.name_resolver = Some(Box::new(&getter));
            let _ = ty.print(printer);
            s
        };

        if let Some(body_id) = body_id {
            let expr = self.tcx.hir().expect_expr(body_id.hir_id);
            local_visitor.visit_expr(expr);
        }
        let err_span = if let Some(pattern) = local_visitor.found_arg_pattern {
            pattern.span
        } else {
            span
        };

        let is_named_and_not_impl_trait = |ty: Ty<'_>| {
            &ty.to_string() != "_" &&
                // FIXME: Remove this check after `impl_trait_in_bindings` is stabilized. #63527
                (!ty.is_impl_trait() || self.tcx.features().impl_trait_in_bindings)
        };

        let ty_msg = match local_visitor.found_ty {
            Some(ty::TyS { sty: ty::Closure(def_id, substs), .. }) => {
                let fn_sig = substs.closure_sig(*def_id, self.tcx);
                let args = closure_args(&fn_sig);
                let ret = fn_sig.output().skip_binder().to_string();
                format!(" for the closure `fn({}) -> {}`", args, ret)
            }
            Some(ty) if is_named_and_not_impl_trait(ty) => {
                let ty = ty_to_string(ty);
                format!(" for `{}`", ty)
            }
            _ => String::new(),
        };

        // When `name` corresponds to a type argument, show the path of the full type we're
        // trying to infer. In the following example, `ty_msg` contains
        // " in `std::result::Result<i32, E>`":
        // ```
        // error[E0282]: type annotations needed for `std::result::Result<i32, E>`
        //  --> file.rs:L:CC
        //   |
        // L |     let b = Ok(4);
        //   |         -   ^^ cannot infer type for `E` in `std::result::Result<i32, E>`
        //   |         |
        //   |         consider giving `b` the explicit type `std::result::Result<i32, E>`, where
        //   |         the type parameter `E` is specified
        // ```
        let mut err = struct_span_err!(
            self.tcx.sess,
            err_span,
            E0282,
            "type annotations needed{}",
            ty_msg,
        );

        let suffix = match local_visitor.found_ty {
            Some(ty::TyS { sty: ty::Closure(def_id, substs), .. }) => {
                let fn_sig = substs.closure_sig(*def_id, self.tcx);
                let ret = fn_sig.output().skip_binder().to_string();

                if let Some(ExprKind::Closure(_, decl, body_id, ..)) = local_visitor.found_closure {
                    if let Some(body) = self.tcx.hir().krate().bodies.get(body_id) {
                        closure_return_type_suggestion(
                            span,
                            &mut err,
                            &decl.output,
                            &body,
                            &name,
                            &ret,
                        );
                        // We don't want to give the other suggestions when the problem is the
                        // closure return type.
                        return err;
                    }
                }

                // This shouldn't be reachable, but just in case we leave a reasonable fallback.
                let args = closure_args(&fn_sig);
                // This suggestion is incomplete, as the user will get further type inference
                // errors due to the `_` placeholders and the introduction of `Box`, but it does
                // nudge them in the right direction.
                format!("a boxed closure type like `Box<dyn Fn({}) -> {}>`", args, ret)
            }
            Some(ty) if is_named_and_not_impl_trait(ty) && name == "_" => {
                let ty = ty_to_string(ty);
                format!("the explicit type `{}`, with the type parameters specified", ty)
            }
            Some(ty) if is_named_and_not_impl_trait(ty) && ty.to_string() != name => {
                let ty = ty_to_string(ty);
                format!(
                    "the explicit type `{}`, where the type parameter `{}` is specified",
                    ty,
                    name,
                )
            }
            _ => "a type".to_string(),
        };

        if let Some(pattern) = local_visitor.found_arg_pattern {
            // We don't want to show the default label for closures.
            //
            // So, before clearing, the output would look something like this:
            // ```
            // let x = |_| {  };
            //          -  ^^^^ cannot infer type for `[_; 0]`
            //          |
            //          consider giving this closure parameter a type
            // ```
            //
            // After clearing, it looks something like this:
            // ```
            // let x = |_| {  };
            //          ^ consider giving this closure parameter the type `[_; 0]`
            //            with the type parameter `_` specified
            // ```
            err.span_label(
                pattern.span,
                format!("consider giving this closure parameter {}", suffix),
            );
        } else if let Some(pattern) = local_visitor.found_local_pattern {
            let msg = if let Some(simple_ident) = pattern.simple_ident() {
                match pattern.span.desugaring_kind() {
                    None => {
                        format!("consider giving `{}` {}", simple_ident, suffix)
                    }
                    Some(DesugaringKind::ForLoop) => {
                        "the element type for this iterator is not specified".to_string()
                    }
                    _ => format!("this needs {}", suffix),
                }
            } else {
                format!("consider giving this pattern {}", suffix)
            };
            err.span_label(pattern.span, msg);
        }
        if !err.span.span_labels().iter().any(|span_label| {
                span_label.label.is_some() && span_label.span == span
            }) && local_visitor.found_arg_pattern.is_none()
        { // Avoid multiple labels pointing at `span`.
            err.span_label(span, InferCtxt::missing_type_msg(&name));
        }

        err
    }

    pub fn need_type_info_err_in_generator(
        &self,
        kind: hir::GeneratorKind,
        span: Span,
        ty: Ty<'tcx>,
    ) -> DiagnosticBuilder<'tcx> {
        let ty = self.resolve_vars_if_possible(&ty);
        let name = self.extract_type_name(&ty, None);
        let mut err = struct_span_err!(
            self.tcx.sess, span, E0698, "type inside {} must be known in this context", kind,
        );
        err.span_label(span, InferCtxt::missing_type_msg(&name));
        err
    }

    fn missing_type_msg(type_name: &str) -> String {
        if type_name == "_" {
            "cannot infer type".to_owned()
        } else {
            format!("cannot infer type for `{}`", type_name)
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    sync::{
        atomic::{self, AtomicU16, AtomicU32, AtomicU64},
        Arc, Mutex,
    },
    time::Instant,
};

use anyhow::{Context as _, Error};
use gba::{Button, ButtonSet};
use sdl2::{
    audio::{AudioCallback, AudioSpec, AudioSpecDesired},
    event::Event,
    keyboard::Keycode,
    pixels::{Color, PixelFormatEnum},
};

const LOCK_FPS: bool = false;

fn main() {
    if let Err(err) = run() {
        eprintln!("{:#}", err);
        std::process::exit(1);
    }
    std::process::exit(0)
}

fn run() -> anyhow::Result<()> {
    pretty_env_logger::formatted_builder()
        .parse_filters(&std::env::var("RUST_LOG").unwrap_or_else(|_| String::from("info")))
        .try_init()
        .context("failed to initiaize logger")?;

    let config = pyrite::config::from_toml_path("pyrite.toml")
        .map_err(|err| {
            let err = anyhow::Error::from(err);
            log::error!("{:#}", err);
        })
        .unwrap_or_default();

    let frame_delay_target = std::time::Duration::from_secs_f64(1.0 / 60.0);

    let sdl_context = sdl2::init()
        .map_err(Error::msg)
        .context("failed to initialize SDL")?;
    let video_subsystem = sdl_context
        .video()
        .map_err(Error::msg)
        .context("failed to initialize SDL video")?;
    let audio_subsystem = sdl_context
        .audio()
        .map_err(Error::msg)
        .context("failed to initialize SDL audio")?;

    let desired_spec = AudioSpecDesired {
        freq: Some(44100),
        channels: Some(1),
        samples: None,
    };
    let device = audio_subsystem
        .open_playback(None, &desired_spec, |spec| GbaAudio { spec })
        .map_err(Error::msg)
        .context("failed to open audio playback")?;
    device.resume();

    let window = video_subsystem
        .window("Pyrite", 480, 320)
        .position_centered()
        .resizable()
        .build()
        .context("failed to create SDL window")?;
    let mut event_pump = sdl_context
        .event_pump()
        .map_err(Error::msg)
        .context("failed to initialize SDL event pump")?;
    let mut canvas_builder = window.into_canvas().accelerated();

    if config.graphics.vsync.unwrap_or(true) {
        canvas_builder = canvas_builder.present_vsync();
    }

    let mut canvas = canvas_builder
        .build()
        .context("failed to initialize SDL canvas")?;
    let texture_creator = canvas.texture_creator();

    let mut gba_frame_texture = texture_creator
        .create_texture_streaming(PixelFormatEnum::BGR555, 240, 160)
        .context("failed to create GBA frame texture")?;
    gba_frame_texture
        .with_lock(None, |buffer: &mut [u8], _: usize| buffer.fill(0))
        .map_err(Error::msg)
        .context("failed to update GBA frame texture")?;
    let screen_buffer = Arc::new(Mutex::new([0u16; gba::SCREEN_PIXEL_COUNT]));

    let gba = pyrite::GbaHandle::new();

    let rom = get_rom_from_args().context("error occurred retrieving ROM from args")?;
    let bios = config.gba.bios_path.and_then(|path| {
        match std::fs::read(&path)
            .with_context(|| format!("failed to read path {}", path.display()))
        {
            Ok(data) => Some(data),
            Err(err) => {
                log::error!("error occurred while loading BIOS: {:?}", err);
                None
            }
        }
    });
    let boot_from_bios = config.gba.boot_from_bios.unwrap_or(true);

    gba.after_frame_wait(move |gba, _| {
        gba.set_gamepak(rom);
        gba.set_bios(bios);
        gba.reset(boot_from_bios);
    });

    let mut buttons = ButtonSet::default();
    let buttons_u16 = Arc::new(AtomicU16::new(buttons.into()));
    let buttons_u16_gba = buttons_u16.clone();

    let frame_duration_accumulate = Arc::new(AtomicU64::new(0));
    let frame_duration_count = Arc::new(AtomicU32::new(0));

    let frame_duration_accumulate_gba = frame_duration_accumulate.clone();
    let frame_duration_count_gba = frame_duration_count.clone();
    let screen = screen_buffer.clone();
    gba.on_frame(move |gba, state| {
        let buttons = buttons_u16_gba.load(atomic::Ordering::Acquire);
        gba.set_buttons(ButtonSet::from(buttons));

        let mut screen = screen.lock().expect("failed to lock screen buffer");
        screen.copy_from_slice(gba.video().screen());

        let frame_duration_us: u64 = state.frame_duration().as_micros().try_into().unwrap();
        frame_duration_accumulate_gba.fetch_add(frame_duration_us, atomic::Ordering::Release);
        frame_duration_count_gba.fetch_add(1, atomic::Ordering::Release);
    });
    gba.set_paused(false);

    let mut fps_counter = FPSCounter::default();

    log::trace!("starting SDL loop...");
    canvas.set_draw_color(Color::RGB(255, 0, 255));
    'main_loop: loop {
        let frame_start_time = Instant::now();

        for event in event_pump.poll_iter() {
            match event {
                Event::Quit { .. }
                | Event::KeyUp {
                    keycode: Some(Keycode::Escape),
                    ..
                } => {
                    log::trace!("exiting SDL loop...");
                    break 'main_loop;
                }

                Event::KeyDown {
                    keycode: Some(keycode),
                    ..
                } => {
                    if let Some(button) = keycode_to_button(keycode) {
                        buttons.set_pressed(button, true)
                    }
                }

                Event::KeyUp {
                    keycode: Some(keycode),
                    ..
                } => {
                    if let Some(button) = keycode_to_button(keycode) {
                        buttons.set_pressed(button, false)
                    }
                }

                _ => (),
            }
        }

        buttons_u16.store(buttons.into(), atomic::Ordering::Release);

        gba_frame_texture
            .with_lock(None, |buffer: &mut [u8], pitch: usize| {
                let screen = screen_buffer.lock().expect("failed to lock screen buffer");
                for y in 0..160 {
                    for x in 0..240 {
                        let offset_src = y * 240 + x;
                        let offset_dst = y * pitch + x * 2;

                        buffer[offset_dst] = screen[offset_src] as u8;
                        buffer[offset_dst + 1] = (screen[offset_src] >> 8) as u8;
                    }
                }
            })
            .map_err(Error::msg)
            .context("failed to update GBA frame texture")?;

        canvas.clear();
        canvas
            .copy(&gba_frame_texture, None, None)
            .map_err(Error::msg)
            .context("failed to copy GBA frame texture to canvas")?;
        canvas.present();

        let frame_end_time = Instant::now();
        let frame_duration = frame_end_time.duration_since(frame_start_time);
        if frame_duration < frame_delay_target && LOCK_FPS {
            let grace = std::time::Duration::from_millis(1);
            std::thread::sleep((frame_delay_target - frame_duration).saturating_sub(grace));
        }

        if let Some(fps) = fps_counter.count(frame_end_time) {
            const GBA_TARGET_FRAME_DURATION_US: f64 = 1000000.0 / 60.0;

            let frame_duration_accumulate_value =
                frame_duration_accumulate.load(atomic::Ordering::Acquire);
            let frame_duration_count_value = frame_duration_count.load(atomic::Ordering::Acquire);
            let average_frame_duration_us =
                frame_duration_accumulate_value as f64 / frame_duration_count_value as f64;
            frame_duration_accumulate
                .fetch_sub(frame_duration_accumulate_value, atomic::Ordering::Release);
            frame_duration_count.fetch_sub(frame_duration_count_value, atomic::Ordering::Release);

            let gba_average_frame_duration = average_frame_duration_us / 1000.0;
            let gba_performance_percentage =
                (GBA_TARGET_FRAME_DURATION_US / average_frame_duration_us) * 100.0;

            let title = format!(
                "Pyrite ({:.1} FPS) (GBA: {:.1} ms | {:.1} %)",
                fps, gba_average_frame_duration, gba_performance_percentage
            );
            canvas
                .window_mut()
                .set_title(&title)
                .context("failed to set window title")?;
        }
    }

    device.pause();
    gba.set_paused(true);
    gba.after_frame_wait(|gba, _| {
        log::debug!(
            "end address: 0x{:08X} (THUMB: {})",
            gba.cpu().next_exec_pc(),
            gba.cpu().registers.getf_t()
        );
    });

    log::info!("exiting...");

    Ok(())
}

pub struct GbaAudio {
    spec: AudioSpec,
}

impl AudioCallback for GbaAudio {
    type Channel = f32;

    fn callback(&mut self, out: &mut [Self::Channel]) {
        static mut PHASE: f32 = 0.0;
        static VOLUME: f32 = 0.001;

        let phase_inc = 440.0 / self.spec.freq as f32;
        out.iter_mut().for_each(|sample| {
            *sample = if unsafe { PHASE } <= 0.5 {
                VOLUME
            } else {
                -VOLUME
            };
            unsafe { PHASE = (PHASE + phase_inc) % 1.0 };
        });
    }
}

fn get_rom_from_args() -> anyhow::Result<Vec<u8>> {
    let rom_path = std::env::args()
        .nth(1)
        .map(std::path::PathBuf::from)
        .ok_or_else(|| anyhow::anyhow!("expected ROM path as first argument"))?;
    std::fs::read(&rom_path)
        .with_context(|| format!("failed to read ROM path `{}`", rom_path.display()))
}

#[derive(Default)]
struct FPSCounter {
    start_time: Option<std::time::Instant>,
    frames: u32,
}

impl FPSCounter {
    pub fn count(&mut self, now: std::time::Instant) -> Option<f64> {
        self.frames += 1;
        if self.start_time.is_none() {
            self.start_time = Some(now);
            return None;
        }

        let elapsed = now.duration_since(self.start_time.unwrap());
        if elapsed < std::time::Duration::from_secs(1) {
            return None;
        }
        self.start_time = Some(now);

        let seconds = elapsed.as_secs_f64();
        let fps = self.frames as f64 / seconds;
        self.frames = 0;

        Some(fps)
    }
}

fn keycode_to_button(keycode: sdl2::keyboard::Keycode) -> Option<Button> {
    match keycode {
        sdl2::keyboard::Keycode::Z => Some(Button::A),
        sdl2::keyboard::Keycode::X => Some(Button::B),
        sdl2::keyboard::Keycode::Left => Some(Button::Left),
        sdl2::keyboard::Keycode::Right => Some(Button::Right),
        sdl2::keyboard::Keycode::Up => Some(Button::Up),
        sdl2::keyboard::Keycode::Down => Some(Button::Down),
        sdl2::keyboard::Keycode::A => Some(Button::L),
        sdl2::keyboard::Keycode::S => Some(Button::R),
        sdl2::keyboard::Keycode::Return => Some(Button::Start),
        sdl2::keyboard::Keycode::Backspace => Some(Button::Select),
        _ => None,
    }
}
<commit_msg>Remove audio test for now.<commit_after>use std::{
    sync::{
        atomic::{self, AtomicU16, AtomicU32, AtomicU64},
        Arc, Mutex,
    },
    time::Instant,
};

use anyhow::{Context as _, Error};
use gba::{Button, ButtonSet};
use sdl2::{
    audio::{AudioCallback, AudioSpec, AudioSpecDesired},
    event::Event,
    keyboard::Keycode,
    pixels::{Color, PixelFormatEnum},
};

const LOCK_FPS: bool = false;

fn main() {
    if let Err(err) = run() {
        eprintln!("{:#}", err);
        std::process::exit(1);
    }
    std::process::exit(0)
}

fn run() -> anyhow::Result<()> {
    pretty_env_logger::formatted_builder()
        .parse_filters(&std::env::var("RUST_LOG").unwrap_or_else(|_| String::from("info")))
        .try_init()
        .context("failed to initiaize logger")?;

    let config = pyrite::config::from_toml_path("pyrite.toml")
        .map_err(|err| {
            let err = anyhow::Error::from(err);
            log::error!("{:#}", err);
        })
        .unwrap_or_default();

    let frame_delay_target = std::time::Duration::from_secs_f64(1.0 / 60.0);

    let sdl_context = sdl2::init()
        .map_err(Error::msg)
        .context("failed to initialize SDL")?;
    let video_subsystem = sdl_context
        .video()
        .map_err(Error::msg)
        .context("failed to initialize SDL video")?;
    let audio_subsystem = sdl_context
        .audio()
        .map_err(Error::msg)
        .context("failed to initialize SDL audio")?;

    let desired_spec = AudioSpecDesired {
        freq: Some(44100),
        channels: Some(1),
        samples: None,
    };
    let device = audio_subsystem
        .open_playback(None, &desired_spec, |spec| GbaAudio { spec })
        .map_err(Error::msg)
        .context("failed to open audio playback")?;
    device.resume();

    let window = video_subsystem
        .window("Pyrite", 480, 320)
        .position_centered()
        .resizable()
        .build()
        .context("failed to create SDL window")?;
    let mut event_pump = sdl_context
        .event_pump()
        .map_err(Error::msg)
        .context("failed to initialize SDL event pump")?;
    let mut canvas_builder = window.into_canvas().accelerated();

    if config.graphics.vsync.unwrap_or(true) {
        canvas_builder = canvas_builder.present_vsync();
    }

    let mut canvas = canvas_builder
        .build()
        .context("failed to initialize SDL canvas")?;
    let texture_creator = canvas.texture_creator();

    let mut gba_frame_texture = texture_creator
        .create_texture_streaming(PixelFormatEnum::BGR555, 240, 160)
        .context("failed to create GBA frame texture")?;
    gba_frame_texture
        .with_lock(None, |buffer: &mut [u8], _: usize| buffer.fill(0))
        .map_err(Error::msg)
        .context("failed to update GBA frame texture")?;
    let screen_buffer = Arc::new(Mutex::new([0u16; gba::SCREEN_PIXEL_COUNT]));

    let gba = pyrite::GbaHandle::new();

    let rom = get_rom_from_args().context("error occurred retrieving ROM from args")?;
    let bios = config.gba.bios_path.and_then(|path| {
        match std::fs::read(&path)
            .with_context(|| format!("failed to read path {}", path.display()))
        {
            Ok(data) => Some(data),
            Err(err) => {
                log::error!("error occurred while loading BIOS: {:?}", err);
                None
            }
        }
    });
    let boot_from_bios = config.gba.boot_from_bios.unwrap_or(true);

    gba.after_frame_wait(move |gba, _| {
        gba.set_gamepak(rom);
        gba.set_bios(bios);
        gba.reset(boot_from_bios);
    });

    let mut buttons = ButtonSet::default();
    let buttons_u16 = Arc::new(AtomicU16::new(buttons.into()));
    let buttons_u16_gba = buttons_u16.clone();

    let frame_duration_accumulate = Arc::new(AtomicU64::new(0));
    let frame_duration_count = Arc::new(AtomicU32::new(0));

    let frame_duration_accumulate_gba = frame_duration_accumulate.clone();
    let frame_duration_count_gba = frame_duration_count.clone();
    let screen = screen_buffer.clone();
    gba.on_frame(move |gba, state| {
        let buttons = buttons_u16_gba.load(atomic::Ordering::Acquire);
        gba.set_buttons(ButtonSet::from(buttons));

        let mut screen = screen.lock().expect("failed to lock screen buffer");
        screen.copy_from_slice(gba.video().screen());

        let frame_duration_us: u64 = state.frame_duration().as_micros().try_into().unwrap();
        frame_duration_accumulate_gba.fetch_add(frame_duration_us, atomic::Ordering::Release);
        frame_duration_count_gba.fetch_add(1, atomic::Ordering::Release);
    });
    gba.set_paused(false);

    let mut fps_counter = FPSCounter::default();

    log::trace!("starting SDL loop...");
    canvas.set_draw_color(Color::RGB(255, 0, 255));
    'main_loop: loop {
        let frame_start_time = Instant::now();

        for event in event_pump.poll_iter() {
            match event {
                Event::Quit { .. }
                | Event::KeyUp {
                    keycode: Some(Keycode::Escape),
                    ..
                } => {
                    log::trace!("exiting SDL loop...");
                    break 'main_loop;
                }

                Event::KeyDown {
                    keycode: Some(keycode),
                    ..
                } => {
                    if let Some(button) = keycode_to_button(keycode) {
                        buttons.set_pressed(button, true)
                    }
                }

                Event::KeyUp {
                    keycode: Some(keycode),
                    ..
                } => {
                    if let Some(button) = keycode_to_button(keycode) {
                        buttons.set_pressed(button, false)
                    }
                }

                _ => (),
            }
        }

        buttons_u16.store(buttons.into(), atomic::Ordering::Release);

        gba_frame_texture
            .with_lock(None, |buffer: &mut [u8], pitch: usize| {
                let screen = screen_buffer.lock().expect("failed to lock screen buffer");
                for y in 0..160 {
                    for x in 0..240 {
                        let offset_src = y * 240 + x;
                        let offset_dst = y * pitch + x * 2;

                        buffer[offset_dst] = screen[offset_src] as u8;
                        buffer[offset_dst + 1] = (screen[offset_src] >> 8) as u8;
                    }
                }
            })
            .map_err(Error::msg)
            .context("failed to update GBA frame texture")?;

        canvas.clear();
        canvas
            .copy(&gba_frame_texture, None, None)
            .map_err(Error::msg)
            .context("failed to copy GBA frame texture to canvas")?;
        canvas.present();

        let frame_end_time = Instant::now();
        let frame_duration = frame_end_time.duration_since(frame_start_time);
        if frame_duration < frame_delay_target && LOCK_FPS {
            let grace = std::time::Duration::from_millis(1);
            std::thread::sleep((frame_delay_target - frame_duration).saturating_sub(grace));
        }

        if let Some(fps) = fps_counter.count(frame_end_time) {
            const GBA_TARGET_FRAME_DURATION_US: f64 = 1000000.0 / 60.0;

            let frame_duration_accumulate_value =
                frame_duration_accumulate.load(atomic::Ordering::Acquire);
            let frame_duration_count_value = frame_duration_count.load(atomic::Ordering::Acquire);
            let average_frame_duration_us =
                frame_duration_accumulate_value as f64 / frame_duration_count_value as f64;
            frame_duration_accumulate
                .fetch_sub(frame_duration_accumulate_value, atomic::Ordering::Release);
            frame_duration_count.fetch_sub(frame_duration_count_value, atomic::Ordering::Release);

            let gba_average_frame_duration = average_frame_duration_us / 1000.0;
            let gba_performance_percentage =
                (GBA_TARGET_FRAME_DURATION_US / average_frame_duration_us) * 100.0;

            let title = format!(
                "Pyrite ({:.1} FPS) (GBA: {:.1} ms | {:.1} %)",
                fps, gba_average_frame_duration, gba_performance_percentage
            );
            canvas
                .window_mut()
                .set_title(&title)
                .context("failed to set window title")?;
        }
    }

    device.pause();
    gba.set_paused(true);
    gba.after_frame_wait(|gba, _| {
        log::debug!(
            "end address: 0x{:08X} (THUMB: {})",
            gba.cpu().next_exec_pc(),
            gba.cpu().registers.getf_t()
        );
    });

    log::info!("exiting...");

    Ok(())
}

#[allow(dead_code)]
pub struct GbaAudio {
    spec: AudioSpec,
}

impl AudioCallback for GbaAudio {
    type Channel = f32;

    fn callback(&mut self, out: &mut [Self::Channel]) {
        out.fill(0.0)
    }
}

fn get_rom_from_args() -> anyhow::Result<Vec<u8>> {
    let rom_path = std::env::args()
        .nth(1)
        .map(std::path::PathBuf::from)
        .ok_or_else(|| anyhow::anyhow!("expected ROM path as first argument"))?;
    std::fs::read(&rom_path)
        .with_context(|| format!("failed to read ROM path `{}`", rom_path.display()))
}

#[derive(Default)]
struct FPSCounter {
    start_time: Option<std::time::Instant>,
    frames: u32,
}

impl FPSCounter {
    pub fn count(&mut self, now: std::time::Instant) -> Option<f64> {
        self.frames += 1;
        if self.start_time.is_none() {
            self.start_time = Some(now);
            return None;
        }

        let elapsed = now.duration_since(self.start_time.unwrap());
        if elapsed < std::time::Duration::from_secs(1) {
            return None;
        }
        self.start_time = Some(now);

        let seconds = elapsed.as_secs_f64();
        let fps = self.frames as f64 / seconds;
        self.frames = 0;

        Some(fps)
    }
}

fn keycode_to_button(keycode: sdl2::keyboard::Keycode) -> Option<Button> {
    match keycode {
        sdl2::keyboard::Keycode::Z => Some(Button::A),
        sdl2::keyboard::Keycode::X => Some(Button::B),
        sdl2::keyboard::Keycode::Left => Some(Button::Left),
        sdl2::keyboard::Keycode::Right => Some(Button::Right),
        sdl2::keyboard::Keycode::Up => Some(Button::Up),
        sdl2::keyboard::Keycode::Down => Some(Button::Down),
        sdl2::keyboard::Keycode::A => Some(Button::L),
        sdl2::keyboard::Keycode::S => Some(Button::R),
        sdl2::keyboard::Keycode::Return => Some(Button::Start),
        sdl2::keyboard::Keycode::Backspace => Some(Button::Select),
        _ => None,
    }
}
<|endoftext|>


<commit_before>use std::env;

use crate::error::*;
use crate::map::Map;
use crate::source::Source;
use crate::value::{Value, ValueKind};

#[derive(Clone, Debug, Default)]
pub struct Environment {
    /// Optional prefix that will limit access to the environment to only keys that
    /// begin with the defined prefix.
    ///
    /// A prefix with a separator of `_` is tested to be present on each key before its considered
    /// to be part of the source environment.
    ///
    /// For example, the key `CONFIG_DEBUG` would become `DEBUG` with a prefix of `config`.
    prefix: Option<String>,

    /// Optional character sequence that separates each key segment in an environment key pattern.
    /// Consider a nested configuration such as `redis.password`, a separator of `_` would allow
    /// an environment key of `REDIS_PASSWORD` to match.
    separator: Option<String>,

    /// Ignore empty env values (treat as unset).
    ignore_empty: bool,

    /// Parses booleans, integers and floats if they're detected (can be safely parsed).
    try_parsing: bool,
}

impl Environment {
    #[deprecated(since = "0.12.0", note = "please use 'Environment::default' instead")]
    pub fn new() -> Self {
        Environment::default()
    }

    pub fn with_prefix(s: &str) -> Self {
        Environment {
            prefix: Some(s.into()),
            ..Environment::default()
        }
    }

    pub fn prefix(mut self, s: &str) -> Self {
        self.prefix = Some(s.into());
        self
    }

    pub fn separator(mut self, s: &str) -> Self {
        self.separator = Some(s.into());
        self
    }

    pub fn ignore_empty(mut self, ignore: bool) -> Self {
        self.ignore_empty = ignore;
        self
    }

    /// Note: enabling `try_parsing` can reduce performance it will try and parse
    /// each environment variable 3 times (bool, i64, f64)
    pub fn try_parsing(mut self, try_parsing: bool) -> Self {
        self.try_parsing = try_parsing;
        self
    }
}

impl Source for Environment {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<Map<String, Value>> {
        let mut m = Map::new();
        let uri: String = "the environment".into();

        let separator = self.separator.as_deref().unwrap_or("");
        let group_separator = self.separator.as_deref().unwrap_or("_");

        // Define a prefix pattern to test and exclude from keys
        let prefix_pattern = self
            .prefix
            .as_ref()
            .map(|prefix| format!("{}{}", prefix, group_separator).to_lowercase());

        for (key, value) in env::vars() {
            // Treat empty environment variables as unset
            if self.ignore_empty && value.is_empty() {
                continue;
            }

            let mut key = key.to_lowercase();

            // Check for prefix
            if let Some(ref prefix_pattern) = prefix_pattern {
                if key.starts_with(prefix_pattern) {
                    // Remove this prefix from the key
                    key = key[prefix_pattern.len()..].to_string();
                } else {
                    // Skip this key
                    continue;
                }
            }

            // If separator is given replace with `.`
            if !separator.is_empty() {
                key = key.replace(separator, ".");
            }

            let value = if self.try_parsing {
                // convert to lowercase because bool parsing expects all lowercase
                if let Ok(parsed) = value.to_lowercase().parse::<bool>() {
                    ValueKind::Boolean(parsed)
                } else if let Ok(parsed) = value.parse::<i64>() {
                    ValueKind::Integer(parsed)
                } else if let Ok(parsed) = value.parse::<f64>() {
                    ValueKind::Float(parsed)
                } else {
                    ValueKind::String(value)
                }
            } else {
                ValueKind::String(value)
            };

            m.insert(key, Value::new(Some(&uri), value));
        }

        Ok(m)
    }
}
<commit_msg>feat: allow providing the environment for testing<commit_after>use std::env;

use crate::error::*;
use crate::map::Map;
use crate::source::Source;
use crate::value::{Value, ValueKind};

#[derive(Clone, Debug, Default)]
pub struct Environment {
    /// Optional prefix that will limit access to the environment to only keys that
    /// begin with the defined prefix.
    ///
    /// A prefix with a separator of `_` is tested to be present on each key before its considered
    /// to be part of the source environment.
    ///
    /// For example, the key `CONFIG_DEBUG` would become `DEBUG` with a prefix of `config`.
    prefix: Option<String>,

    /// Optional character sequence that separates each key segment in an environment key pattern.
    /// Consider a nested configuration such as `redis.password`, a separator of `_` would allow
    /// an environment key of `REDIS_PASSWORD` to match.
    separator: Option<String>,

    /// Ignore empty env values (treat as unset).
    ignore_empty: bool,

    /// Parses booleans, integers and floats if they're detected (can be safely parsed).
    try_parsing: bool,

    /// Alternate source for the environment. This can be used when you want to test your own code
    /// using this source, without the need to change the actual system environment variables.
    ///
    /// ## Example
    ///
    /// ```rust
    /// # use config::{Environment, Config};
    /// # use serde::Deserialize;
    /// # use std::collections::HashMap;
    /// # use std::convert::TryInto;
    /// #
    /// #[test]
    /// fn test_config() -> Result<(), config::ConfigError> {
    ///   #[derive(Clone, Debug, Deserialize)]
    ///   struct MyConfig {
    ///     pub my_string: String,
    ///   }
    ///
    ///   let source = Environment::default()
    ///     .source(Some({
    ///       let mut env = HashMap::new();
    ///       env.insert("MY_STRING".into(), "my-value".into());
    ///       env
    ///   }));
    ///
    ///   let config: MyConfig = Config::builder()
    ///     .add_source(source)
    ///     .build()?
    ///     .try_into()?;
    ///   assert_eq!(config.my_string, "my-value");
    ///
    ///   Ok(())
    /// }
    /// ```
    source: Option<Map<String, String>>,
}

impl Environment {
    #[deprecated(since = "0.12.0", note = "please use 'Environment::default' instead")]
    pub fn new() -> Self {
        Environment::default()
    }

    pub fn with_prefix(s: &str) -> Self {
        Environment {
            prefix: Some(s.into()),
            ..Environment::default()
        }
    }

    pub fn prefix(mut self, s: &str) -> Self {
        self.prefix = Some(s.into());
        self
    }

    pub fn separator(mut self, s: &str) -> Self {
        self.separator = Some(s.into());
        self
    }

    pub fn ignore_empty(mut self, ignore: bool) -> Self {
        self.ignore_empty = ignore;
        self
    }

    /// Note: enabling `try_parsing` can reduce performance it will try and parse
    /// each environment variable 3 times (bool, i64, f64)
    pub fn try_parsing(mut self, try_parsing: bool) -> Self {
        self.try_parsing = try_parsing;
        self
    }

    pub fn source(mut self, source: Option<Map<String, String>>) -> Self {
        self.source = source;
        self
    }
}

impl Source for Environment {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<Map<String, Value>> {
        let mut m = Map::new();
        let uri: String = "the environment".into();

        let separator = self.separator.as_deref().unwrap_or("");
        let group_separator = self.separator.as_deref().unwrap_or("_");

        // Define a prefix pattern to test and exclude from keys
        let prefix_pattern = self
            .prefix
            .as_ref()
            .map(|prefix| format!("{}{}", prefix, group_separator).to_lowercase());

        let collector = |(key, value): (String, String)| {
            // Treat empty environment variables as unset
            if self.ignore_empty && value.is_empty() {
                return;
            }

            let mut key = key.to_lowercase();

            // Check for prefix
            if let Some(ref prefix_pattern) = prefix_pattern {
                if key.starts_with(prefix_pattern) {
                    // Remove this prefix from the key
                    key = key[prefix_pattern.len()..].to_string();
                } else {
                    // Skip this key
                    return;
                }
            }

            // If separator is given replace with `.`
            if !separator.is_empty() {
                key = key.replace(separator, ".");
            }

            let value = if self.try_parsing {
                // convert to lowercase because bool parsing expects all lowercase
                if let Ok(parsed) = value.to_lowercase().parse::<bool>() {
                    ValueKind::Boolean(parsed)
                } else if let Ok(parsed) = value.parse::<i64>() {
                    ValueKind::Integer(parsed)
                } else if let Ok(parsed) = value.parse::<f64>() {
                    ValueKind::Float(parsed)
                } else {
                    ValueKind::String(value)
                }
            } else {
                ValueKind::String(value)
            };

            m.insert(key, Value::new(Some(&uri), value));
        };

        match &self.source {
            Some(source) => source.clone().into_iter().for_each(collector),
            None => env::vars().for_each(collector),
        }

        Ok(m)
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::convert::TryInto;
use std::env;

use crate::error::*;
use crate::source::Source;
use crate::value::{Value, ValueKind};

#[derive(Clone, Debug)]
pub struct Environment {
    /// Optional prefix that will limit access to the environment to only keys that
    /// begin with the defined prefix.
    ///
    /// A prefix with a separator of `_` is tested to be present on each key before its considered
    /// to be part of the source environment.
    ///
    /// For example, the key `CONFIG_DEBUG` would become `DEBUG` with a prefix of `config`.
    prefix: Option<String>,

    /// Optional character sequence that separates each key segment in an environment key pattern.
    /// Consider a nested configuration such as `redis.password`, a separator of `_` would allow
    /// an environment key of `REDIS_PASSWORD` to match.
    separator: Option<String>,

    /// Ignore empty env values (treat as unset).
    ignore_empty: bool,

    /// Parses booleans, integers and floats if they're detected (can be safely parsed).
    try_parsing: bool,
}

impl Environment {
    pub fn new() -> Self {
        Environment::default()
    }

    pub fn with_prefix(s: &str) -> Self {
        Environment {
            prefix: Some(s.into()),
            ..Environment::default()
        }
    }

    pub fn prefix(mut self, s: &str) -> Self {
        self.prefix = Some(s.into());
        self
    }

    pub fn separator(mut self, s: &str) -> Self {
        self.separator = Some(s.into());
        self
    }

    pub fn ignore_empty(mut self, ignore: bool) -> Self {
        self.ignore_empty = ignore;
        self
    }

    pub fn try_parsing(mut self, try_parsing: bool) -> Self {
        self.try_parsing = try_parsing;
        self
    }
}

impl Default for Environment {
    fn default() -> Environment {
        Environment {
            prefix: None,
            separator: None,
            ignore_empty: false,
            try_parsing: false,
        }
    }
}

impl Source for Environment {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        let mut m = HashMap::new();
        let uri: String = "the environment".into();

        let separator = match self.separator {
            Some(ref separator) => separator,
            _ => "",
        };

        let group_separator = match self.separator {
            Some(ref separator) => separator,
            _ => "_",
        };

        // Define a prefix pattern to test and exclude from keys
        let prefix_pattern = self
            .prefix
            .as_ref()
            .map(|prefix| format!("{}{}", prefix.clone(), group_separator));

        for (key, value) in env::vars() {
            // Treat empty environment variables as unset
            if self.ignore_empty && value.is_empty() {
                continue;
            }

            let mut key = key.to_string();

            // Check for prefix
            if let Some(ref prefix_pattern) = prefix_pattern {
                if key
                    .to_lowercase()
                    .starts_with(&prefix_pattern.to_lowercase())
                {
                    // Remove this prefix from the key
                    key = key[prefix_pattern.len()..].to_string();
                } else {
                    // Skip this key
                    continue;
                }
            }

            // If separator is given replace with `.`
            if !separator.is_empty() {
                key = key.replace(separator, ".");
            }

            let value = if self.try_parsing {
                let string_value = Value::new(Some(&uri), ValueKind::String(value.clone()));
                let is_true_false =
                    value.to_lowercase() == "true" || value.to_lowercase() == "false";

                // if the value is "true" or "false" it should be parsed as a bool
                if let (true, Ok(parsed)) = (is_true_false, string_value.clone().into_bool()) {
                    ValueKind::Boolean(parsed)
                } else if let Ok(parsed) = string_value.clone().into_int() {
                    ValueKind::Integer(parsed)
                } else if let Ok(parsed) = string_value.clone().into_float() {
                    ValueKind::Float(parsed)
                } else if let Ok(parsed) = string_value.clone().into_bool() {
                    ValueKind::Boolean(parsed)
                } else {
                    ValueKind::String(value)
                }
            } else {
                ValueKind::String(value)
            };

            m.insert(key.to_lowercase(), Value::new(Some(&uri), value));
        }

        Ok(m)
    }
}
<commit_msg>more predictable env parsing and reduced clones<commit_after>use std::collections::HashMap;
use std::convert::TryInto;
use std::env;

use crate::error::*;
use crate::source::Source;
use crate::value::{Value, ValueKind};

#[derive(Clone, Debug)]
pub struct Environment {
    /// Optional prefix that will limit access to the environment to only keys that
    /// begin with the defined prefix.
    ///
    /// A prefix with a separator of `_` is tested to be present on each key before its considered
    /// to be part of the source environment.
    ///
    /// For example, the key `CONFIG_DEBUG` would become `DEBUG` with a prefix of `config`.
    prefix: Option<String>,

    /// Optional character sequence that separates each key segment in an environment key pattern.
    /// Consider a nested configuration such as `redis.password`, a separator of `_` would allow
    /// an environment key of `REDIS_PASSWORD` to match.
    separator: Option<String>,

    /// Ignore empty env values (treat as unset).
    ignore_empty: bool,

    /// Parses booleans, integers and floats if they're detected (can be safely parsed).
    try_parsing: bool,
}

impl Environment {
    pub fn new() -> Self {
        Environment::default()
    }

    pub fn with_prefix(s: &str) -> Self {
        Environment {
            prefix: Some(s.into()),
            ..Environment::default()
        }
    }

    pub fn prefix(mut self, s: &str) -> Self {
        self.prefix = Some(s.into());
        self
    }

    pub fn separator(mut self, s: &str) -> Self {
        self.separator = Some(s.into());
        self
    }

    pub fn ignore_empty(mut self, ignore: bool) -> Self {
        self.ignore_empty = ignore;
        self
    }

    pub fn try_parsing(mut self, try_parsing: bool) -> Self {
        self.try_parsing = try_parsing;
        self
    }
}

impl Default for Environment {
    fn default() -> Environment {
        Environment {
            prefix: None,
            separator: None,
            ignore_empty: false,
            try_parsing: false,
        }
    }
}

impl Source for Environment {
    fn clone_into_box(&self) -> Box<dyn Source + Send + Sync> {
        Box::new((*self).clone())
    }

    fn collect(&self) -> Result<HashMap<String, Value>> {
        let mut m = HashMap::new();
        let uri: String = "the environment".into();

        let separator = match self.separator {
            Some(ref separator) => separator,
            _ => "",
        };

        let group_separator = match self.separator {
            Some(ref separator) => separator,
            _ => "_",
        };

        // Define a prefix pattern to test and exclude from keys
        let prefix_pattern = self
            .prefix
            .as_ref()
            .map(|prefix| format!("{}{}", prefix.clone(), group_separator));

        for (key, value) in env::vars() {
            // Treat empty environment variables as unset
            if self.ignore_empty && value.is_empty() {
                continue;
            }

            let mut key = key.to_string();

            // Check for prefix
            if let Some(ref prefix_pattern) = prefix_pattern {
                if key
                    .to_lowercase()
                    .starts_with(&prefix_pattern.to_lowercase())
                {
                    // Remove this prefix from the key
                    key = key[prefix_pattern.len()..].to_string();
                } else {
                    // Skip this key
                    continue;
                }
            }

            // If separator is given replace with `.`
            if !separator.is_empty() {
                key = key.replace(separator, ".");
            }

            let value = if self.try_parsing {
                // convert to lowercase because bool parsing expects all lowercase
                if let Ok(parsed) = value.to_lowercase().parse::<bool>() {
                    Value::new(Some(&uri), ValueKind::Boolean(parsed))
                } else if let Ok(parsed) = value.parse::<i64>() {
                    Value::new(Some(&uri), ValueKind::Integer(parsed))
                } else if let Ok(parsed) = value.parse::<f64>() {
                    Value::new(Some(&uri), ValueKind::Float(parsed))
                } else {
                    Value::new(Some(&uri), ValueKind::String(value))
                }
            } else {
                Value::new(Some(&uri), ValueKind::String(value))
            };

            m.insert(key.to_lowercase(), value);
        }

        Ok(m)
    }
}
<|endoftext|>


<commit_before>use vmx;

#[repr(packed)]
struct GDTEntry {
    limit_low: u16,
    base_low: u16,
    base_middle: u8,
    access: u8,
    granularity: u8,
    base_high: u8,
    base_highest: u32,
    reserved0: u32,
}

const GDT: [GDTEntry; 3] = [
    GDTEntry{
        limit_low: 0,
        base_low: 0,
        base_middle: 0,
        access: 0,
        granularity: 0,
        base_high: 0,
        base_highest: 0,
        reserved0: 0,
    },
    GDTEntry{
        limit_low: 0xffff,
        base_low: 0,
        base_middle: 0,
        access: 0b10101001,
        granularity: 0b11110111,
        base_high: 0,
        base_highest: 0,
        reserved0: 0,
    },
    GDTEntry{
        limit_low: 0xffff,
        base_low: 0,
        base_middle: 0,
        access: 0b00101001,
        granularity: 0b11110111,
        base_high: 0,
        base_highest: 0,
        reserved0: 0,
    },
];

pub fn new_host_descriptor() -> vmx::CPUTableDescriptor {
    vmx::CPUTableDescriptor{ limit: 0xffff, base: GDT.as_ptr() as u64 }
}

#[cfg(feature = "runtime_tests")]
pub mod runtime_tests {

    use cli;
    use vmx;
    use super::new_host_descriptor;

    pub fn test_load() {
        cli::ClearLocalInterruptsGuard::new();
        let gdt_desc = new_host_descriptor();
        let mut orig_gdt_desc: vmx::CPUTableDescriptor = Default::default();
        vmx::sgdt(&mut orig_gdt_desc);
        vmx::lgdt(&gdt_desc);
        vmx::lgdt(&orig_gdt_desc);
    }
}
<commit_msg>Forcibly squash dead code warnings<commit_after>use vmx;

#[repr(packed)]
#[allow(dead_code)]
struct GDTEntry {
    limit_low: u16,
    base_low: u16,
    base_middle: u8,
    access: u8,
    granularity: u8,
    base_high: u8,
    base_highest: u32,
    reserved0: u32,
}

const GDT: [GDTEntry; 3] = [
    GDTEntry{
        limit_low: 0,
        base_low: 0,
        base_middle: 0,
        access: 0,
        granularity: 0,
        base_high: 0,
        base_highest: 0,
        reserved0: 0,
    },
    GDTEntry{
        limit_low: 0xffff,
        base_low: 0,
        base_middle: 0,
        access: 0b10101001,
        granularity: 0b11110111,
        base_high: 0,
        base_highest: 0,
        reserved0: 0,
    },
    GDTEntry{
        limit_low: 0xffff,
        base_low: 0,
        base_middle: 0,
        access: 0b00101001,
        granularity: 0b11110111,
        base_high: 0,
        base_highest: 0,
        reserved0: 0,
    },
];

pub fn new_host_descriptor() -> vmx::CPUTableDescriptor {
    vmx::CPUTableDescriptor{ limit: 0xffff, base: GDT.as_ptr() as u64 }
}

#[cfg(feature = "runtime_tests")]
pub mod runtime_tests {

    use cli;
    use vmx;
    use super::new_host_descriptor;

    pub fn test_load() {
        cli::ClearLocalInterruptsGuard::new();
        let gdt_desc = new_host_descriptor();
        let mut orig_gdt_desc: vmx::CPUTableDescriptor = Default::default();
        vmx::sgdt(&mut orig_gdt_desc);
        vmx::lgdt(&gdt_desc);
        vmx::lgdt(&orig_gdt_desc);
    }
}
<|endoftext|>


<commit_before>//!  Decoding of GIF Images
//!
//!  GIF (Graphics Interchange Format) is an image format that supports lossless compression.
//!
//!  # Related Links
//!  * <http://www.w3.org/Graphics/GIF/spec-gif89a.txt> - The GIF Specification
//!
//! # Examples
//! ```rust,no_run
//! use image::gif::{Decoder, Encoder};
//! use image::{ImageDecoder, AnimationDecoder};
//! use std::fs::File;
//! # fn main() -> std::io::Result<()> {
//! // Decode a gif into frames
//! let file_in = File::open("foo.gif")?;
//! let mut decoder = Decoder::new(file_in).unwrap();
//! let frames = decoder.into_frames();
//! let frames = frames.collect_frames().expect("error decoding gif");
//!
//! // Encode frames into a gif and save to a file
//! let mut file_out = File::open("out.gif")?;
//! let mut encoder = Encoder::new(file_out);
//! encoder.encode_frames(frames.into_iter());
//! # Ok(())
//! # }
//! ```
#![cfg_attr(feature = "cargo-clippy", allow(while_let_loop))]

extern crate gif;
extern crate num_rational;

use std::clone::Clone;
use std::io::{self, Cursor, Read, Write};
use std::marker::PhantomData;
use std::mem;

use self::gif::{ColorOutput, SetParameter};
pub use self::gif::{DisposalMethod, Frame};

use animation;
use buffer::{ImageBuffer, Pixel};
use color;
use color::Rgba;
use image::{AnimationDecoder, ImageDecoder, ImageError, ImageResult};
use num_rational::Ratio;

/// GIF decoder
pub struct Decoder<R: Read> {
    reader: gif::Reader<R>,
}

impl<R: Read> Decoder<R> {
    /// Creates a new decoder that decodes the input steam ```r```
    pub fn new(r: R) -> ImageResult<Decoder<R>> {
        let mut decoder = gif::Decoder::new(r);
        decoder.set(ColorOutput::RGBA);

        Ok(Decoder {
            reader: decoder.read_info()?,
        })
    }
}

/// Wrapper struct around a `Cursor<Vec<u8>>`
pub struct GifReader<R>(Cursor<Vec<u8>>, PhantomData<R>);
impl<R> Read for GifReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        if self.0.position() == 0 && buf.is_empty() {
            mem::swap(buf, self.0.get_mut());
            Ok(buf.len())
        } else {
            self.0.read_to_end(buf)
        }
    }
}

impl<'a, R: 'a + Read> ImageDecoder<'a> for Decoder<R> {
    type Reader = GifReader<R>;

    fn dimensions(&self) -> (u64, u64) {
        (self.reader.width() as u64, self.reader.height() as u64)
    }

    fn colortype(&self) -> color::ColorType {
        color::ColorType::RGBA(8)
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        Ok(GifReader(Cursor::new(self.read_image()?), PhantomData))
    }

    fn read_image(mut self) -> ImageResult<Vec<u8>> {
        if self.reader.next_frame_info()?.is_some() {
            let mut buf = vec![0; self.reader.buffer_size()];
            self.reader.read_into_buffer(&mut buf)?;
            Ok(buf)
        } else {
            Err(ImageError::ImageEnd)
        }
    }
}

struct GifFrameIterator<R: Read> {
    reader: gif::Reader<R>,

    width: u32,
    height: u32,

    background_img: ImageBuffer<Rgba<u8>, Vec<u8>>,
    non_disposed_frame: ImageBuffer<Rgba<u8>, Vec<u8>>,
}

impl<R: Read> GifFrameIterator<R> {
    fn new(decoder: Decoder<R>) -> GifFrameIterator<R> {
        let (width, height) = decoder.dimensions();

        // TODO: Avoid this cast
        let (width, height) = (width as u32, height as u32);

        // set the background color to be either the bg_color defined in the gif
        // or a transparent pixel.
        let background_color_option = decoder.reader.bg_color();
        let mut background_color = vec![0; 4];
        let background_pixel = {
            let global_palette = decoder.reader.global_palette();
            match background_color_option {
                Some(index) => {
                    // find the color by looking in the global palette
                    match global_palette {
                        // take the color from the palette that is at the index defined
                        // by background_color_option
                        Some(slice) => {
                            background_color.clone_from_slice(&slice[index..(index + 4)]);
                            Rgba::from_slice(&background_color)
                        }
                        // if there is no global palette, assign the background color to be
                        // transparent
                        None => Rgba::from_slice(&[0, 0, 0, 0]),
                    }
                }
                // return a transparent background color
                None => Rgba::from_slice(&[0, 0, 0, 0]),
            }
        };

        // create the background image to use later
        let background_img = ImageBuffer::from_pixel(width, height, *background_pixel);

        // the background image is the first non disposed frame
        let non_disposed_frame = background_img.clone();

        GifFrameIterator {
            reader: decoder.reader,
            width,
            height,
            background_img,
            non_disposed_frame,
        }
    }
}

impl<R: Read> Iterator for GifFrameIterator<R> {
    type Item = ImageResult<animation::Frame>;

    fn next(&mut self) -> Option<ImageResult<animation::Frame>> {
        // begin looping over each frame
        let (left, top, delay, dispose, f_width, f_height);

        match self.reader.next_frame_info() {
            Ok(frame_info) => {
                if let Some(frame) = frame_info {
                    left = u32::from(frame.left);
                    top = u32::from(frame.top);
                    f_width = u32::from(frame.width);
                    f_height = u32::from(frame.height);

                    // frame.delay is in units of 10ms so frame.delay*10 is in ms
                    delay = Ratio::new(frame.delay * 10, 1);
                    dispose = frame.dispose;
                } else {
                    // no more frames
                    return None;
                }
            },
            Err(err) => return Some(Err(err.into())),
        }

        let mut vec = vec![0; self.reader.buffer_size()];
        if let Err(err) = self.reader.fill_buffer(&mut vec) {
            return Some(Err(err.into()));
        }

        // create the image buffer from the raw frame
        let image_buffer_raw = ImageBuffer::from_raw(f_width, f_height, vec).unwrap();

        // if `image_buffer_raw`'s frame exactly matches the entire image, then
        // use it directly.
        let mut image_buffer = if (left, top) == (0, 0) && (f_width, f_height) == (self.width, self.height) {
            image_buffer_raw
        } else {
            // otherwise, `image_buffer_raw` represents a smaller image.
            // create a new image of the target size and place
            // `image_buffer_raw` within it. the outside region is filled with
            // transparent pixels.
            ImageBuffer::from_fn(self.width, self.height, |x, y| {
                let x = x.wrapping_sub(left);
                let y = y.wrapping_sub(top);
                if x < f_width && y < f_height {
                    *image_buffer_raw.get_pixel(x, y)
                } else {
                    Rgba([0, 0, 0, 0])
                }
            })
        };

        // loop over all pixels, checking if any pixels from the non disposed
        // frame need to be used
        for (x, y, pixel) in image_buffer.enumerate_pixels_mut() {
            let previous_img_buffer = &self.non_disposed_frame;
            let adjusted_pixel: &mut Rgba<u8> = pixel;
            let previous_pixel: &Rgba<u8> = previous_img_buffer.get_pixel(x, y);

            let pixel_alpha = adjusted_pixel.channels()[3];

            // If a pixel is not visible then we show the non disposed frame pixel instead
            if pixel_alpha == 0 {
                adjusted_pixel.blend(previous_pixel);
            }
        }

        let frame = animation::Frame::from_parts(
            image_buffer.clone(), 0, 0, delay
        );

        match dispose {
            DisposalMethod::Any => {
                // do nothing
                // (completely replace this frame with the next)
            }
            DisposalMethod::Keep => {
                // do not dispose
                // (keep pixels from this frame)
                self.non_disposed_frame = image_buffer;
            }
            DisposalMethod::Background => {
                // restore to background color
                // (background shows through transparent pixels in the next frame)
                self.non_disposed_frame = self.background_img.clone();
            }
            DisposalMethod::Previous => {
                // restore to previous
                // (dispose frames leaving the last none disposal frame)
            }
        };

        Some(Ok(frame))
    }
}

impl<'a, R: Read + 'a> AnimationDecoder<'a> for Decoder<R> {
    fn into_frames(self) -> animation::Frames<'a> {
        animation::Frames::new(Box::new(GifFrameIterator::new(self)))
    }
}

/// GIF encoder.
pub struct Encoder<W: Write> {
    w: Option<W>,
    gif_encoder: Option<gif::Encoder<W>>,
}

impl<W: Write> Encoder<W> {
    /// Creates a new GIF encoder.
    pub fn new(w: W) -> Encoder<W> {
        Encoder {
            w: Some(w),
            gif_encoder: None,
        }
    }
    /// Encodes a frame.
    pub fn encode(&mut self, frame: &Frame) -> ImageResult<()> {
        let result;
        if let Some(ref mut encoder) = self.gif_encoder {
            result = encoder.write_frame(frame).map_err(|err| err.into());
        } else {
            let writer = self.w.take().unwrap();
            let mut encoder = gif::Encoder::new(writer, frame.width, frame.height, &[])?;
            result = encoder.write_frame(&frame).map_err(|err| err.into());
            self.gif_encoder = Some(encoder);
        }
        result
    }

    /// Encodes Frames.
    /// Consider using `try_encode_frames` instead to encode an `animation::Frames` like iterator.
    pub fn encode_frames<F>(&mut self, frames: F) -> ImageResult<()>
    where
        F: IntoIterator<Item = animation::Frame>
    {
        for img_frame in frames {
            self.encode_single_frame(img_frame)?;
        }
        Ok(())
    }

    /// Try to encode a collection of `ImageResult<animation::Frame>` objects.
    /// Use this function to encode an `animation::Frames` like iterator.
    /// Whenever an `Err` item is encountered, that value is returned without further actions.
    pub fn try_encode_frames<F>(&mut self, frames: F) -> ImageResult<()>
    where
        F: IntoIterator<Item = ImageResult<animation::Frame>>
    {
        for img_frame in frames {
            self.encode_single_frame(img_frame?)?;
        }
        Ok(())
    }

    fn encode_single_frame(&mut self, img_frame: animation::Frame) -> ImageResult<()> {
        // get the delay before converting img_frame
        let frame_delay = img_frame.delay().to_integer();
        // convert img_frame into RgbaImage
        let rbga_frame = img_frame.into_buffer();

        // Create the gif::Frame from the animation::Frame
        let mut frame = Frame::from_rgba(rbga_frame.width() as u16, rbga_frame.height() as u16, &mut rbga_frame.into_raw());
        frame.delay = frame_delay;

        // encode the gif::Frame
        self.encode(&frame)
    }
}

impl From<gif::DecodingError> for ImageError {
    fn from(err: gif::DecodingError) -> ImageError {
        use self::gif::DecodingError::*;
        match err {
            Format(desc) | Internal(desc) => ImageError::FormatError(desc.into()),
            Io(io_err) => ImageError::IoError(io_err),
        }
    }
}
<commit_msg>Implement the "background" disposal method correctly<commit_after>//!  Decoding of GIF Images
//!
//!  GIF (Graphics Interchange Format) is an image format that supports lossless compression.
//!
//!  # Related Links
//!  * <http://www.w3.org/Graphics/GIF/spec-gif89a.txt> - The GIF Specification
//!
//! # Examples
//! ```rust,no_run
//! use image::gif::{Decoder, Encoder};
//! use image::{ImageDecoder, AnimationDecoder};
//! use std::fs::File;
//! # fn main() -> std::io::Result<()> {
//! // Decode a gif into frames
//! let file_in = File::open("foo.gif")?;
//! let mut decoder = Decoder::new(file_in).unwrap();
//! let frames = decoder.into_frames();
//! let frames = frames.collect_frames().expect("error decoding gif");
//!
//! // Encode frames into a gif and save to a file
//! let mut file_out = File::open("out.gif")?;
//! let mut encoder = Encoder::new(file_out);
//! encoder.encode_frames(frames.into_iter());
//! # Ok(())
//! # }
//! ```
#![cfg_attr(feature = "cargo-clippy", allow(while_let_loop))]

extern crate gif;
extern crate num_rational;

use std::clone::Clone;
use std::io::{self, Cursor, Read, Write};
use std::marker::PhantomData;
use std::mem;

use self::gif::{ColorOutput, SetParameter};
pub use self::gif::{DisposalMethod, Frame};

use animation;
use buffer::{ImageBuffer, Pixel};
use color;
use color::Rgba;
use image::{AnimationDecoder, ImageDecoder, ImageError, ImageResult};
use num_rational::Ratio;

/// GIF decoder
pub struct Decoder<R: Read> {
    reader: gif::Reader<R>,
}

impl<R: Read> Decoder<R> {
    /// Creates a new decoder that decodes the input steam ```r```
    pub fn new(r: R) -> ImageResult<Decoder<R>> {
        let mut decoder = gif::Decoder::new(r);
        decoder.set(ColorOutput::RGBA);

        Ok(Decoder {
            reader: decoder.read_info()?,
        })
    }
}

/// Wrapper struct around a `Cursor<Vec<u8>>`
pub struct GifReader<R>(Cursor<Vec<u8>>, PhantomData<R>);
impl<R> Read for GifReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        if self.0.position() == 0 && buf.is_empty() {
            mem::swap(buf, self.0.get_mut());
            Ok(buf.len())
        } else {
            self.0.read_to_end(buf)
        }
    }
}

impl<'a, R: 'a + Read> ImageDecoder<'a> for Decoder<R> {
    type Reader = GifReader<R>;

    fn dimensions(&self) -> (u64, u64) {
        (self.reader.width() as u64, self.reader.height() as u64)
    }

    fn colortype(&self) -> color::ColorType {
        color::ColorType::RGBA(8)
    }

    fn into_reader(self) -> ImageResult<Self::Reader> {
        Ok(GifReader(Cursor::new(self.read_image()?), PhantomData))
    }

    fn read_image(mut self) -> ImageResult<Vec<u8>> {
        if self.reader.next_frame_info()?.is_some() {
            let mut buf = vec![0; self.reader.buffer_size()];
            self.reader.read_into_buffer(&mut buf)?;
            Ok(buf)
        } else {
            Err(ImageError::ImageEnd)
        }
    }
}

struct GifFrameIterator<R: Read> {
    reader: gif::Reader<R>,

    width: u32,
    height: u32,

    background_color: Rgba<u8>,
    non_disposed_frame: ImageBuffer<Rgba<u8>, Vec<u8>>,
}

impl<R: Read> GifFrameIterator<R> {
    fn new(decoder: Decoder<R>) -> GifFrameIterator<R> {
        let (width, height) = decoder.dimensions();

        // TODO: Avoid this cast
        let (width, height) = (width as u32, height as u32);

        // set the background color to be either the bg_color defined in the gif
        // or a transparent pixel.
        let background_color_option = decoder.reader.bg_color();
        let mut background_color = vec![0; 4];
        let background_pixel = {
            let global_palette = decoder.reader.global_palette();
            match background_color_option {
                Some(index) => {
                    // find the color by looking in the global palette
                    match global_palette {
                        // take the color from the palette that is at the index defined
                        // by background_color_option
                        Some(slice) => {
                            background_color.clone_from_slice(&slice[index..(index + 4)]);
                            Rgba::from_slice(&background_color)
                        }
                        // if there is no global palette, assign the background color to be
                        // transparent
                        None => Rgba::from_slice(&[0, 0, 0, 0]),
                    }
                }
                // return a transparent background color
                None => Rgba::from_slice(&[0, 0, 0, 0]),
            }
        };

        // the background image is the first non disposed frame
        let non_disposed_frame = ImageBuffer::from_pixel(width, height, *background_pixel);

        GifFrameIterator {
            reader: decoder.reader,
            width,
            height,
            background_color: *background_pixel,
            non_disposed_frame,
        }
    }
}

impl<R: Read> Iterator for GifFrameIterator<R> {
    type Item = ImageResult<animation::Frame>;

    fn next(&mut self) -> Option<ImageResult<animation::Frame>> {
        // begin looping over each frame
        let (left, top, delay, dispose, f_width, f_height);

        match self.reader.next_frame_info() {
            Ok(frame_info) => {
                if let Some(frame) = frame_info {
                    left = u32::from(frame.left);
                    top = u32::from(frame.top);
                    f_width = u32::from(frame.width);
                    f_height = u32::from(frame.height);

                    // frame.delay is in units of 10ms so frame.delay*10 is in ms
                    delay = Ratio::new(frame.delay * 10, 1);
                    dispose = frame.dispose;
                } else {
                    // no more frames
                    return None;
                }
            },
            Err(err) => return Some(Err(err.into())),
        }

        let mut vec = vec![0; self.reader.buffer_size()];
        if let Err(err) = self.reader.fill_buffer(&mut vec) {
            return Some(Err(err.into()));
        }

        // create the image buffer from the raw frame
        let image_buffer_raw = ImageBuffer::from_raw(f_width, f_height, vec).unwrap();

        // if `image_buffer_raw`'s frame exactly matches the entire image, then
        // use it directly.
        let mut image_buffer = if (left, top) == (0, 0) && (f_width, f_height) == (self.width, self.height) {
            image_buffer_raw
        } else {
            // otherwise, `image_buffer_raw` represents a smaller image.
            // create a new image of the target size and place
            // `image_buffer_raw` within it. the outside region is filled with
            // transparent pixels.
            ImageBuffer::from_fn(self.width, self.height, |x, y| {
                let x = x.wrapping_sub(left);
                let y = y.wrapping_sub(top);
                if x < f_width && y < f_height {
                    *image_buffer_raw.get_pixel(x, y)
                } else {
                    Rgba([0, 0, 0, 0])
                }
            })
        };

        // loop over all pixels, checking if any pixels from the non disposed
        // frame need to be used
        for (x, y, pixel) in image_buffer.enumerate_pixels_mut() {
            let previous_img_buffer = &self.non_disposed_frame;
            let adjusted_pixel: &mut Rgba<u8> = pixel;
            let previous_pixel: &Rgba<u8> = previous_img_buffer.get_pixel(x, y);

            let pixel_alpha = adjusted_pixel.channels()[3];

            // If a pixel is not visible then we show the non disposed frame pixel instead
            if pixel_alpha == 0 {
                adjusted_pixel.blend(previous_pixel);
            }
        }

        let frame = animation::Frame::from_parts(
            image_buffer.clone(), 0, 0, delay
        );

        match dispose {
            DisposalMethod::Any => {
                // do nothing
                // (completely replace this frame with the next)
            }
            DisposalMethod::Keep => {
                // do not dispose
                // (keep pixels from this frame)
                self.non_disposed_frame = image_buffer;
            }
            DisposalMethod::Background => {
                // restore to background color
                // (background shows through transparent pixels in the next frame)
                for y in top..top + f_height {
                    for x in left..left + f_width {
                        self.non_disposed_frame.put_pixel(x, y, self.background_color);
                    }
                }
            }
            DisposalMethod::Previous => {
                // restore to previous
                // (dispose frames leaving the last none disposal frame)
            }
        };

        Some(Ok(frame))
    }
}

impl<'a, R: Read + 'a> AnimationDecoder<'a> for Decoder<R> {
    fn into_frames(self) -> animation::Frames<'a> {
        animation::Frames::new(Box::new(GifFrameIterator::new(self)))
    }
}

/// GIF encoder.
pub struct Encoder<W: Write> {
    w: Option<W>,
    gif_encoder: Option<gif::Encoder<W>>,
}

impl<W: Write> Encoder<W> {
    /// Creates a new GIF encoder.
    pub fn new(w: W) -> Encoder<W> {
        Encoder {
            w: Some(w),
            gif_encoder: None,
        }
    }
    /// Encodes a frame.
    pub fn encode(&mut self, frame: &Frame) -> ImageResult<()> {
        let result;
        if let Some(ref mut encoder) = self.gif_encoder {
            result = encoder.write_frame(frame).map_err(|err| err.into());
        } else {
            let writer = self.w.take().unwrap();
            let mut encoder = gif::Encoder::new(writer, frame.width, frame.height, &[])?;
            result = encoder.write_frame(&frame).map_err(|err| err.into());
            self.gif_encoder = Some(encoder);
        }
        result
    }

    /// Encodes Frames.
    /// Consider using `try_encode_frames` instead to encode an `animation::Frames` like iterator.
    pub fn encode_frames<F>(&mut self, frames: F) -> ImageResult<()>
    where
        F: IntoIterator<Item = animation::Frame>
    {
        for img_frame in frames {
            self.encode_single_frame(img_frame)?;
        }
        Ok(())
    }

    /// Try to encode a collection of `ImageResult<animation::Frame>` objects.
    /// Use this function to encode an `animation::Frames` like iterator.
    /// Whenever an `Err` item is encountered, that value is returned without further actions.
    pub fn try_encode_frames<F>(&mut self, frames: F) -> ImageResult<()>
    where
        F: IntoIterator<Item = ImageResult<animation::Frame>>
    {
        for img_frame in frames {
            self.encode_single_frame(img_frame?)?;
        }
        Ok(())
    }

    fn encode_single_frame(&mut self, img_frame: animation::Frame) -> ImageResult<()> {
        // get the delay before converting img_frame
        let frame_delay = img_frame.delay().to_integer();
        // convert img_frame into RgbaImage
        let rbga_frame = img_frame.into_buffer();

        // Create the gif::Frame from the animation::Frame
        let mut frame = Frame::from_rgba(rbga_frame.width() as u16, rbga_frame.height() as u16, &mut rbga_frame.into_raw());
        frame.delay = frame_delay;

        // encode the gif::Frame
        self.encode(&frame)
    }
}

impl From<gif::DecodingError> for ImageError {
    fn from(err: gif::DecodingError) -> ImageError {
        use self::gif::DecodingError::*;
        match err {
            Format(desc) | Internal(desc) => ImageError::FormatError(desc.into()),
            Io(io_err) => ImageError::IoError(io_err),
        }
    }
}
<|endoftext|>