<commit_before>//! Check whether a type has (potentially) non-trivial drop glue.

use rustc_data_structures::fx::FxHashSet;
use rustc_hir::def_id::DefId;
use rustc_middle::ty::subst::Subst;
use rustc_middle::ty::subst::SubstsRef;
use rustc_middle::ty::util::{needs_drop_components, AlwaysRequiresDrop};
use rustc_middle::ty::{self, Ty, TyCtxt};
use rustc_session::Limit;
use rustc_span::{sym, DUMMY_SP};

type NeedsDropResult<T> = Result<T, AlwaysRequiresDrop>;

fn needs_drop_raw<'tcx>(tcx: TyCtxt<'tcx>, query: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
    // If we don't know a type doesn't need drop, for example if it's a type
    // parameter without a `Copy` bound, then we conservatively return that it
    // needs drop.
    let adt_has_dtor =
        |adt_def: &ty::AdtDef| adt_def.destructor(tcx).map(|_| DtorType::Significant);
    let res = drop_tys_helper(tcx, query.value, query.param_env, adt_has_dtor).next().is_some();

    debug!("needs_drop_raw({:?}) = {:?}", query, res);
    res
}

fn has_significant_drop_raw<'tcx>(
    tcx: TyCtxt<'tcx>,
    query: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
) -> bool {
    let res =
        drop_tys_helper(tcx, query.value, query.param_env, adt_consider_insignificant_dtor(tcx))
            .next()
            .is_some();
    debug!("has_significant_drop_raw({:?}) = {:?}", query, res);
    res
}

struct NeedsDropTypes<'tcx, F> {
    tcx: TyCtxt<'tcx>,
    param_env: ty::ParamEnv<'tcx>,
    query_ty: Ty<'tcx>,
    seen_tys: FxHashSet<Ty<'tcx>>,
    /// A stack of types left to process, and the recursion depth when we
    /// pushed that type. Each round, we pop something from the stack and check
    /// if it needs drop. If the result depends on whether some other types
    /// need drop we push them onto the stack.
    unchecked_tys: Vec<(Ty<'tcx>, usize)>,
    recursion_limit: Limit,
    adt_components: F,
}

impl<'tcx, F> NeedsDropTypes<'tcx, F> {
    fn new(
        tcx: TyCtxt<'tcx>,
        param_env: ty::ParamEnv<'tcx>,
        ty: Ty<'tcx>,
        adt_components: F,
    ) -> Self {
        let mut seen_tys = FxHashSet::default();
        seen_tys.insert(ty);
        Self {
            tcx,
            param_env,
            seen_tys,
            query_ty: ty,
            unchecked_tys: vec![(ty, 0)],
            recursion_limit: tcx.recursion_limit(),
            adt_components,
        }
    }
}

impl<'tcx, F, I> Iterator for NeedsDropTypes<'tcx, F>
where
    F: Fn(&ty::AdtDef, SubstsRef<'tcx>) -> NeedsDropResult<I>,
    I: Iterator<Item = Ty<'tcx>>,
{
    type Item = NeedsDropResult<Ty<'tcx>>;

    fn next(&mut self) -> Option<NeedsDropResult<Ty<'tcx>>> {
        let tcx = self.tcx;

        while let Some((ty, level)) = self.unchecked_tys.pop() {
            if !self.recursion_limit.value_within_limit(level) {
                // Not having a `Span` isn't great. But there's hopefully some other
                // recursion limit error as well.
                tcx.sess.span_err(
                    DUMMY_SP,
                    &format!("overflow while checking whether `{}` requires drop", self.query_ty),
                );
                return Some(Err(AlwaysRequiresDrop));
            }

            let components = match needs_drop_components(ty, &tcx.data_layout) {
                Err(e) => return Some(Err(e)),
                Ok(components) => components,
            };
            debug!("needs_drop_components({:?}) = {:?}", ty, components);

            let queue_type = move |this: &mut Self, component: Ty<'tcx>| {
                if this.seen_tys.insert(component) {
                    this.unchecked_tys.push((component, level + 1));
                }
            };

            for component in components {
                match *component.kind() {
                    _ if component.is_copy_modulo_regions(tcx.at(DUMMY_SP), self.param_env) => (),

                    ty::Closure(_, substs) => {
                        queue_type(self, substs.as_closure().tupled_upvars_ty());
                    }

                    ty::Generator(def_id, substs, _) => {
                        let substs = substs.as_generator();
                        queue_type(self, substs.tupled_upvars_ty());

                        let witness = substs.witness();
                        let interior_tys = match witness.kind() {
                            &ty::GeneratorWitness(tys) => tcx.erase_late_bound_regions(tys),
                            _ => {
                                tcx.sess.delay_span_bug(
                                    tcx.hir().span_if_local(def_id).unwrap_or(DUMMY_SP),
                                    &format!("unexpected generator witness type {:?}", witness),
                                );
                                return Some(Err(AlwaysRequiresDrop));
                            }
                        };

                        for interior_ty in interior_tys {
                            queue_type(self, interior_ty);
                        }
                    }

                    // Check for a `Drop` impl and whether this is a union or
                    // `ManuallyDrop`. If it's a struct or enum without a `Drop`
                    // impl then check whether the field types need `Drop`.
                    ty::Adt(adt_def, substs) => {
                        let tys = match (self.adt_components)(adt_def, substs) {
                            Err(e) => return Some(Err(e)),
                            Ok(tys) => tys,
                        };
                        for required_ty in tys {
                            let subst_ty =
                                tcx.normalize_erasing_regions(self.param_env, required_ty);
                            queue_type(self, subst_ty);
                        }
                    }
                    ty::Array(..) | ty::Opaque(..) | ty::Projection(..) | ty::Param(_) => {
                        if ty == component {
                            // Return the type to the caller: they may be able
                            // to normalize further than we can.
                            return Some(Ok(component));
                        } else {
                            // Store the type for later. We can't return here
                            // because we would then lose any other components
                            // of the type.
                            queue_type(self, component);
                        }
                    }
                    _ => return Some(Err(AlwaysRequiresDrop)),
                }
            }
        }

        None
    }
}

enum DtorType {
    /// Type has a `Drop` but it is considered insignificant.
    /// Check the query `adt_significant_drop_tys` for understanding
    /// "significant" / "insignificant".
    Insignificant,

    /// Type has a `Drop` implentation.
    Significant,
}

// This is a helper function for `adt_drop_tys` and `adt_significant_drop_tys`.
// Depending on the implentation of `adt_has_dtor`, it is used to check if the
// ADT has a destructor or if the ADT only has a significant destructor. For
// understanding significant destructor look at `adt_significant_drop_tys`.
fn drop_tys_helper<'tcx>(
    tcx: TyCtxt<'tcx>,
    ty: Ty<'tcx>,
    param_env: rustc_middle::ty::ParamEnv<'tcx>,
    adt_has_dtor: impl Fn(&ty::AdtDef) -> Option<DtorType>,
) -> impl Iterator<Item = NeedsDropResult<Ty<'tcx>>> {
    let adt_components = move |adt_def: &ty::AdtDef, substs: SubstsRef<'tcx>| {
        if adt_def.is_manually_drop() {
            debug!("drop_tys_helper: `{:?}` is manually drop", adt_def);
            return Ok(Vec::new().into_iter());
        } else if let Some(dtor_info) = adt_has_dtor(adt_def) {
            match dtor_info {
                DtorType::Significant => {
                    debug!("drop_tys_helper: `{:?}` implements `Drop`", adt_def);
                    return Err(AlwaysRequiresDrop);
                }
                DtorType::Insignificant => {
                    debug!("drop_tys_helper: `{:?}` drop is insignificant", adt_def);

                    // Since the destructor is insignificant, we just want to make sure all of
                    // the passed in type parameters are also insignificant.
                    // Eg: Vec<T> dtor is insignificant when T=i32 but significant when T=Mutex.
                    return Ok(substs.types().collect::<Vec<Ty<'_>>>().into_iter());
                }
            }
        } else if adt_def.is_union() {
            debug!("drop_tys_helper: `{:?}` is a union", adt_def);
            return Ok(Vec::new().into_iter());
        }
        Ok(adt_def
            .all_fields()
            .map(|field| {
                let r = tcx.type_of(field.did).subst(tcx, substs);
                debug!("drop_tys_helper: Subst into {:?} with {:?} gettng {:?}", field, substs, r);
                r
            })
            .collect::<Vec<_>>()
            .into_iter())
    };

    NeedsDropTypes::new(tcx, param_env, ty, adt_components)
}

fn adt_consider_insignificant_dtor<'tcx>(
    tcx: TyCtxt<'tcx>,
) -> impl Fn(&ty::AdtDef) -> Option<DtorType> + 'tcx {
    move |adt_def: &ty::AdtDef| {
        let is_marked_insig = tcx.has_attr(adt_def.did, sym::rustc_insignificant_dtor);
        if is_marked_insig {
            // In some cases like `std::collections::HashMap` where the struct is a wrapper around
            // a type that is a Drop type, and the wrapped type (eg: `hashbrown::HashMap`) lies
            // outside stdlib, we might choose to still annotate the the wrapper (std HashMap) with
            // `rustc_insignificant_dtor`, even if the type itself doesn't have a `Drop` impl.
            Some(DtorType::Insignificant)
        } else if adt_def.destructor(tcx).is_some() {
            // There is a Drop impl and the type isn't marked insignificant, therefore Drop must be
            // significant.
            Some(DtorType::Significant)
        } else {
            // No destructor found nor the type is annotated with `rustc_insignificant_dtor`, we
            // treat this as the simple case of Drop impl for type.
            None
        }
    }
}

fn adt_drop_tys(tcx: TyCtxt<'_>, def_id: DefId) -> Result<&ty::List<Ty<'_>>, AlwaysRequiresDrop> {
    // This is for the "adt_drop_tys" query, that considers all `Drop` impls, therefore all dtors are
    // significant.
    let adt_has_dtor =
        |adt_def: &ty::AdtDef| adt_def.destructor(tcx).map(|_| DtorType::Significant);
    drop_tys_helper(tcx, tcx.type_of(def_id), tcx.param_env(def_id), adt_has_dtor)
        .collect::<Result<Vec<_>, _>>()
        .map(|components| tcx.intern_type_list(&components))
}

fn adt_significant_drop_tys(
    tcx: TyCtxt<'_>,
    def_id: DefId,
) -> Result<&ty::List<Ty<'_>>, AlwaysRequiresDrop> {
    drop_tys_helper(
        tcx,
        tcx.type_of(def_id),
        tcx.param_env(def_id),
        adt_consider_insignificant_dtor(tcx),
    )
    .collect::<Result<Vec<_>, _>>()
    .map(|components| tcx.intern_type_list(&components))
}

pub(crate) fn provide(providers: &mut ty::query::Providers) {
    *providers = ty::query::Providers {
        needs_drop_raw,
        has_significant_drop_raw,
        adt_drop_tys,
        adt_significant_drop_tys,
        ..*providers
    };
}
<commit_msg>Recurse through query system when checking ADT drop types, hopefully improving perf<commit_after>//! Check whether a type has (potentially) non-trivial drop glue.

use rustc_data_structures::fx::FxHashSet;
use rustc_hir::def_id::DefId;
use rustc_middle::ty::subst::Subst;
use rustc_middle::ty::subst::SubstsRef;
use rustc_middle::ty::util::{needs_drop_components, AlwaysRequiresDrop};
use rustc_middle::ty::{self, Ty, TyCtxt};
use rustc_session::Limit;
use rustc_span::{sym, DUMMY_SP};

type NeedsDropResult<T> = Result<T, AlwaysRequiresDrop>;

fn needs_drop_raw<'tcx>(tcx: TyCtxt<'tcx>, query: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
    // If we don't know a type doesn't need drop, for example if it's a type
    // parameter without a `Copy` bound, then we conservatively return that it
    // needs drop.
    let adt_has_dtor =
        |adt_def: &ty::AdtDef| adt_def.destructor(tcx).map(|_| DtorType::Significant);
    let res =
        drop_tys_helper(tcx, query.value, query.param_env, adt_has_dtor, false).next().is_some();

    debug!("needs_drop_raw({:?}) = {:?}", query, res);
    res
}

fn has_significant_drop_raw<'tcx>(
    tcx: TyCtxt<'tcx>,
    query: ty::ParamEnvAnd<'tcx, Ty<'tcx>>,
) -> bool {
    let res = drop_tys_helper(
        tcx,
        query.value,
        query.param_env,
        adt_consider_insignificant_dtor(tcx),
        true,
    )
    .next()
    .is_some();
    debug!("has_significant_drop_raw({:?}) = {:?}", query, res);
    res
}

struct NeedsDropTypes<'tcx, F> {
    tcx: TyCtxt<'tcx>,
    param_env: ty::ParamEnv<'tcx>,
    query_ty: Ty<'tcx>,
    seen_tys: FxHashSet<Ty<'tcx>>,
    /// A stack of types left to process, and the recursion depth when we
    /// pushed that type. Each round, we pop something from the stack and check
    /// if it needs drop. If the result depends on whether some other types
    /// need drop we push them onto the stack.
    unchecked_tys: Vec<(Ty<'tcx>, usize)>,
    recursion_limit: Limit,
    adt_components: F,
}

impl<'tcx, F> NeedsDropTypes<'tcx, F> {
    fn new(
        tcx: TyCtxt<'tcx>,
        param_env: ty::ParamEnv<'tcx>,
        ty: Ty<'tcx>,
        adt_components: F,
    ) -> Self {
        let mut seen_tys = FxHashSet::default();
        seen_tys.insert(ty);
        Self {
            tcx,
            param_env,
            seen_tys,
            query_ty: ty,
            unchecked_tys: vec![(ty, 0)],
            recursion_limit: tcx.recursion_limit(),
            adt_components,
        }
    }
}

impl<'tcx, F, I> Iterator for NeedsDropTypes<'tcx, F>
where
    F: Fn(&ty::AdtDef, SubstsRef<'tcx>) -> NeedsDropResult<I>,
    I: Iterator<Item = Ty<'tcx>>,
{
    type Item = NeedsDropResult<Ty<'tcx>>;

    fn next(&mut self) -> Option<NeedsDropResult<Ty<'tcx>>> {
        let tcx = self.tcx;

        while let Some((ty, level)) = self.unchecked_tys.pop() {
            if !self.recursion_limit.value_within_limit(level) {
                // Not having a `Span` isn't great. But there's hopefully some other
                // recursion limit error as well.
                tcx.sess.span_err(
                    DUMMY_SP,
                    &format!("overflow while checking whether `{}` requires drop", self.query_ty),
                );
                return Some(Err(AlwaysRequiresDrop));
            }

            let components = match needs_drop_components(ty, &tcx.data_layout) {
                Err(e) => return Some(Err(e)),
                Ok(components) => components,
            };
            debug!("needs_drop_components({:?}) = {:?}", ty, components);

            let queue_type = move |this: &mut Self, component: Ty<'tcx>| {
                if this.seen_tys.insert(component) {
                    this.unchecked_tys.push((component, level + 1));
                }
            };

            for component in components {
                match *component.kind() {
                    _ if component.is_copy_modulo_regions(tcx.at(DUMMY_SP), self.param_env) => (),

                    ty::Closure(_, substs) => {
                        queue_type(self, substs.as_closure().tupled_upvars_ty());
                    }

                    ty::Generator(def_id, substs, _) => {
                        let substs = substs.as_generator();
                        queue_type(self, substs.tupled_upvars_ty());

                        let witness = substs.witness();
                        let interior_tys = match witness.kind() {
                            &ty::GeneratorWitness(tys) => tcx.erase_late_bound_regions(tys),
                            _ => {
                                tcx.sess.delay_span_bug(
                                    tcx.hir().span_if_local(def_id).unwrap_or(DUMMY_SP),
                                    &format!("unexpected generator witness type {:?}", witness),
                                );
                                return Some(Err(AlwaysRequiresDrop));
                            }
                        };

                        for interior_ty in interior_tys {
                            queue_type(self, interior_ty);
                        }
                    }

                    // Check for a `Drop` impl and whether this is a union or
                    // `ManuallyDrop`. If it's a struct or enum without a `Drop`
                    // impl then check whether the field types need `Drop`.
                    ty::Adt(adt_def, substs) => {
                        let tys = match (self.adt_components)(adt_def, substs) {
                            Err(e) => return Some(Err(e)),
                            Ok(tys) => tys,
                        };
                        for required_ty in tys {
                            let required =
                                tcx.normalize_erasing_regions(self.param_env, required_ty);
                            queue_type(self, required);
                        }
                    }
                    ty::Array(..) | ty::Opaque(..) | ty::Projection(..) | ty::Param(_) => {
                        if ty == component {
                            // Return the type to the caller: they may be able
                            // to normalize further than we can.
                            return Some(Ok(component));
                        } else {
                            // Store the type for later. We can't return here
                            // because we would then lose any other components
                            // of the type.
                            queue_type(self, component);
                        }
                    }
                    _ => return Some(Err(AlwaysRequiresDrop)),
                }
            }
        }

        None
    }
}

enum DtorType {
    /// Type has a `Drop` but it is considered insignificant.
    /// Check the query `adt_significant_drop_tys` for understanding
    /// "significant" / "insignificant".
    Insignificant,

    /// Type has a `Drop` implentation.
    Significant,
}

// This is a helper function for `adt_drop_tys` and `adt_significant_drop_tys`.
// Depending on the implentation of `adt_has_dtor`, it is used to check if the
// ADT has a destructor or if the ADT only has a significant destructor. For
// understanding significant destructor look at `adt_significant_drop_tys`.
fn drop_tys_helper<'tcx>(
    tcx: TyCtxt<'tcx>,
    ty: Ty<'tcx>,
    param_env: rustc_middle::ty::ParamEnv<'tcx>,
    adt_has_dtor: impl Fn(&ty::AdtDef) -> Option<DtorType>,
    only_significant: bool,
) -> impl Iterator<Item = NeedsDropResult<Ty<'tcx>>> {
    fn with_query_cache<'tcx>(
        tcx: TyCtxt<'tcx>,
        iter: impl IntoIterator<Item = Ty<'tcx>>,
        only_significant: bool,
    ) -> NeedsDropResult<Vec<Ty<'tcx>>> {
        iter.into_iter().try_fold(Vec::new(), |mut vec, subty| {
            match subty.kind() {
                ty::Adt(adt_id, subst) => {
                    for subty in if only_significant {
                        tcx.adt_significant_drop_tys(adt_id.did)?
                    } else {
                        tcx.adt_drop_tys(adt_id.did)?
                    } {
                        vec.push(subty.subst(tcx, subst));
                    }
                }
                _ => vec.push(subty),
            };
            Ok(vec)
        })
    }

    let adt_components = move |adt_def: &ty::AdtDef, substs: SubstsRef<'tcx>| {
        if adt_def.is_manually_drop() {
            debug!("drop_tys_helper: `{:?}` is manually drop", adt_def);
            Ok(Vec::new())
        } else if let Some(dtor_info) = adt_has_dtor(adt_def) {
            match dtor_info {
                DtorType::Significant => {
                    debug!("drop_tys_helper: `{:?}` implements `Drop`", adt_def);
                    Err(AlwaysRequiresDrop)
                }
                DtorType::Insignificant => {
                    debug!("drop_tys_helper: `{:?}` drop is insignificant", adt_def);

                    // Since the destructor is insignificant, we just want to make sure all of
                    // the passed in type parameters are also insignificant.
                    // Eg: Vec<T> dtor is insignificant when T=i32 but significant when T=Mutex.
                    with_query_cache(tcx, substs.types(), only_significant)
                }
            }
        } else if adt_def.is_union() {
            debug!("drop_tys_helper: `{:?}` is a union", adt_def);
            Ok(Vec::new())
        } else {
            with_query_cache(
                tcx,
                adt_def.all_fields().map(|field| {
                    let r = tcx.type_of(field.did).subst(tcx, substs);
                    debug!(
                        "drop_tys_helper: Subst into {:?} with {:?} gettng {:?}",
                        field, substs, r
                    );
                    r
                }),
                only_significant,
            )
        }
        .map(|v| v.into_iter())
    };

    NeedsDropTypes::new(tcx, param_env, ty, adt_components)
}

fn adt_consider_insignificant_dtor<'tcx>(
    tcx: TyCtxt<'tcx>,
) -> impl Fn(&ty::AdtDef) -> Option<DtorType> + 'tcx {
    move |adt_def: &ty::AdtDef| {
        let is_marked_insig = tcx.has_attr(adt_def.did, sym::rustc_insignificant_dtor);
        if is_marked_insig {
            // In some cases like `std::collections::HashMap` where the struct is a wrapper around
            // a type that is a Drop type, and the wrapped type (eg: `hashbrown::HashMap`) lies
            // outside stdlib, we might choose to still annotate the the wrapper (std HashMap) with
            // `rustc_insignificant_dtor`, even if the type itself doesn't have a `Drop` impl.
            Some(DtorType::Insignificant)
        } else if adt_def.destructor(tcx).is_some() {
            // There is a Drop impl and the type isn't marked insignificant, therefore Drop must be
            // significant.
            Some(DtorType::Significant)
        } else {
            // No destructor found nor the type is annotated with `rustc_insignificant_dtor`, we
            // treat this as the simple case of Drop impl for type.
            None
        }
    }
}

fn adt_drop_tys(tcx: TyCtxt<'_>, def_id: DefId) -> Result<&ty::List<Ty<'_>>, AlwaysRequiresDrop> {
    // This is for the "adt_drop_tys" query, that considers all `Drop` impls, therefore all dtors are
    // significant.
    let adt_has_dtor =
        |adt_def: &ty::AdtDef| adt_def.destructor(tcx).map(|_| DtorType::Significant);
    // `tcx.type_of(def_id)` identical to `tcx.make_adt(def, identity_substs)`
    drop_tys_helper(tcx, tcx.type_of(def_id), tcx.param_env(def_id), adt_has_dtor, false)
        .collect::<Result<Vec<_>, _>>()
        .map(|components| tcx.intern_type_list(&components))
}
// If `def_id` refers to a generic ADT, the queries above and below act as if they had been handed
// a `tcx.make_ty(def, identity_substs)` and as such it is legal to substitue the generic parameters
// of the ADT into the outputted `ty`s.
fn adt_significant_drop_tys(
    tcx: TyCtxt<'_>,
    def_id: DefId,
) -> Result<&ty::List<Ty<'_>>, AlwaysRequiresDrop> {
    drop_tys_helper(
        tcx,
        tcx.type_of(def_id), // identical to `tcx.make_adt(def, identity_substs)`
        tcx.param_env(def_id),
        adt_consider_insignificant_dtor(tcx),
        true,
    )
    .collect::<Result<Vec<_>, _>>()
    .map(|components| tcx.intern_type_list(&components))
}

pub(crate) fn provide(providers: &mut ty::query::Providers) {
    *providers = ty::query::Providers {
        needs_drop_raw,
        has_significant_drop_raw,
        adt_drop_tys,
        adt_significant_drop_tys,
        ..*providers
    };
}
<|endoftext|>


<commit_before>// A port of the simplistic benchmark from
//
//    http://github.com/PaulKeeble/ScalaVErlangAgents
//
// I *think* it's the same, more or less.

extern mod std;
extern mod zmq;

use dvec::DVec;

fn server(ctx: zmq::Context, ch: oldcomm::Chan<()>, workers: uint) {
    let mut workers = workers;

    let pull_socket = match ctx.socket(zmq::PULL) {
      Ok(move socket) => socket,
      Err(e) => fail e.to_str(),
    };

    let push_socket = match ctx.socket(zmq::PUSH) {
      Ok(move socket) => socket,
      Err(e) => fail e.to_str(),
    };

    match pull_socket.bind("tcp://127.0.0.1:3456") {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    }

    match push_socket.bind("tcp://127.0.0.1:3457") {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    }

    // Let the main thread know we're ready.
    ch.send(());

    let mut count = 0u;
    while workers != 0 {
        match pull_socket.recv(0) {
            Err(e) => fail e.to_str(),
            Ok(msg) => {
                do msg.with_str |s| {
                    if s == "" {
                        workers -= 1;
                    } else {
                        count += uint::from_str(s).get();
                    }
                }
            }
        }
    }

    match push_socket.send_str(uint::str(count), 0) {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    }

    pull_socket.close();
    push_socket.close();

    ch.send(());
}

fn worker(ctx: zmq::Context, count: uint) {
    let push_socket = match ctx.socket(zmq::PUSH) {
      Ok(move socket) => socket,
      Err(e) => fail e.to_str(),
    };

    match push_socket.connect("tcp://127.0.0.1:3456") {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    }

    for count.times {
        match push_socket.send_str(uint::str(100u), 0) {
          Ok(()) => { }
          Err(e) => fail e.to_str(),
        }
    }

    // Let the server know we're done.
    match push_socket.send_str("", 0) {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    }

    push_socket.close();
}

fn run(ctx: zmq::Context, size: uint, workers: uint) {
    // Spawn the server.
    let po = oldcomm::Port();
    let ch = oldcomm::Chan(&po);
    do task::spawn_sched(task::SingleThreaded) {
        server(ctx, ch, workers);
    }

    // Wait for the server to start.
    po.recv();

    // Create some oldcommand/control sockets.
    let push_socket = match ctx.socket(zmq::PUSH) {
      Ok(move socket) => socket,
      Err(e) => fail e.to_str(),
    };

    match push_socket.connect("tcp://127.0.0.1:3456") {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    }

    let pull_socket = match ctx.socket(zmq::PULL) {
      Ok(move socket) => socket,
      Err(e) => fail e.to_str(),
    };

    match pull_socket.connect("tcp://127.0.0.1:3457") {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    }

    let start = std::time::precise_time_s();

    // Spawn all the workers.
    let mut worker_results = ~[];

    for workers.times {
        let po = oldcomm::Port();
        let ch = oldcomm::Chan(&po);

        worker_results.push(po);

        do task::spawn_sched(task::SingleThreaded) {
            worker(ctx, size / workers);
            ch.send(());
        }
    }

    // Block until all the workers finish.
    for worker_results.each |po| { po.recv(); }

    /*
    // Shut down the server.
    push_socket.send_str("stop", 0);
    match push_socket.close() {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    }
    */

    po.recv();

    // Receive the final count.
    let result = match pull_socket.recv(0) {
      Ok(msg) => msg.with_str(|s| uint::from_str(s).get()),
      Err(e) => fail e.to_str(),
    };

    let end = std::time::precise_time_s();
    let elapsed = end - start;

    io::println(fmt!("Count is %?", result));
    io::println(fmt!("Test took %? seconds", elapsed));
    let thruput = ((size / workers * workers) as float) / (elapsed as float);
    io::println(fmt!("Throughput=%f per sec", thruput));
}

fn main() {
    let args = os::args();

    let args = if os::getenv(~"RUST_BENCH").is_some() {
        ~[~"", ~"1000000", ~"10000"]
    } else if args.len() <= 1u {
        ~[~"", ~"10000", ~"4"]
    } else {
        args
    };

    let size = uint::from_str(args[1]).get();
    let workers = uint::from_str(args[2]).get();

    let ctx = match zmq::init(1) {
      Ok(ctx) => ctx,
      Err(e) => fail e.to_str(),
    };

    run(ctx, size, workers);

    match ctx.term() {
      Ok(()) => { }
      Err(e) => fail e.to_str(),
    };
}
<commit_msg>Update msgsend-zmq example to rust incoming<commit_after>// A port of the simplistic benchmark from
//
//    http://github.com/PaulKeeble/ScalaVErlangAgents
//
// I *think* it's the same, more or less.

extern mod std;
extern mod extra;
extern mod zmq;

use std::comm;
use std::io;
use std::os;
use std::task;
use std::uint;

fn server(ctx: zmq::Context, ch: &comm::Chan<()>, workers: uint) {
    let mut workers = workers;

    let mut pull_socket = match ctx.socket(zmq::PULL) {
        Ok(socket) => socket,
        Err(e) => fail!(e.to_str()),
    };

    let mut push_socket = match ctx.socket(zmq::PUSH) {
        Ok(socket) => socket,
        Err(e) => fail!(e.to_str()),
    };

    match pull_socket.bind("tcp://127.0.0.1:3456") {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }

    match push_socket.bind("tcp://127.0.0.1:3457") {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }

    // Let the main thread know we're ready.
    ch.send(());

    let mut count = 0u;
    while workers != 0 {
        match unsafe { pull_socket.recv(0) } {
            Err(e) => fail!(e.to_str()),
            Ok(msg) => {
                do msg.with_str |s| {
                    if s == "" {
                        workers -= 1;
                    } else {
                        count += uint::from_str(s).get();
                    }
                }
            }
        }
    }

    match push_socket.send_str(count.to_str(), 0) {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }

    pull_socket.close();
    push_socket.close();

    ch.send(());
}

fn worker(ctx: zmq::Context, count: uint) {
    let mut push_socket = match ctx.socket(zmq::PUSH) {
        Ok(socket) => socket,
        Err(e) => fail!(e.to_str()),
    };

    match push_socket.connect("tcp://127.0.0.1:3456") {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }

    for count.times {
        match push_socket.send_str(100u.to_str(), 0) {
          Ok(()) => { }
          Err(e) => fail!(e.to_str()),
        }
    }

    // Let the server know we're done.
    match push_socket.send_str("", 0) {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }

    push_socket.close();
}

fn run(ctx: zmq::Context, size: uint, workers: uint) {
    // Spawn the server.
    let (po, ch) = comm::stream();
    do task::spawn_sched(task::SingleThreaded) {
        server(ctx, &ch, workers);
    }

    // Wait for the server to start.
    po.recv();

    // Create some command/control sockets.
    let push_socket = match ctx.socket(zmq::PUSH) {
        Ok(socket) => socket,
        Err(e) => fail!(e.to_str()),
    };

    match push_socket.connect("tcp://127.0.0.1:3456") {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }

    let pull_socket = match ctx.socket(zmq::PULL) {
        Ok(socket) => socket,
        Err(e) => fail!(e.to_str()),
    };

    match pull_socket.connect("tcp://127.0.0.1:3457") {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }

    let start = extra::time::precise_time_s();

    // Spawn all the workers.
    let mut worker_results = ~[];

    for workers.times {
        let (po, ch) = comm::stream();

        worker_results.push(po);

        do task::spawn_sched(task::SingleThreaded) {
            worker(ctx, size / workers);
            ch.send(());
        }
    }

    // Block until all the workers finish.
    for worker_results.each |po| { po.recv(); }

    /*
    // Shut down the server.
    push_socket.send_str("stop", 0);
    match push_socket.close() {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    }
    */

    po.recv();

    // Receive the final count.
    let result = match unsafe { pull_socket.recv(0) } {
        Ok(msg) => msg.with_str(|s| uint::from_str(s).get()),
        Err(e) => fail!(e.to_str()),
    };

    let end = extra::time::precise_time_s();
    let elapsed = end - start;

    io::println(fmt!("Count is %?", result));
    io::println(fmt!("Test took %? seconds", elapsed));
    let thruput = ((size / workers * workers) as float) / (elapsed as float);
    io::println(fmt!("Throughput=%f per sec", thruput));
}

fn main() {
    let args = os::args();

    let args = if os::getenv("RUST_BENCH").is_some() {
        ~[~"", ~"1000000", ~"10000"]
    } else if args.len() <= 1u {
        ~[~"", ~"10000", ~"4"]
    } else {
        args
    };

    let size = uint::from_str(args[1]).get();
    let workers = uint::from_str(args[2]).get();

    let ctx = match zmq::init(1) {
        Ok(ctx) => ctx,
        Err(e) => fail!(e.to_str()),
    };

    run(ctx, size, workers);

    match ctx.term() {
        Ok(()) => { }
        Err(e) => fail!(e.to_str()),
    };
}
<|endoftext|>


<commit_before>use clap::{App, Arg, Format, ArgMatches};
use std::io::{self, Write};
use std::process::exit;
use std::fs;
use crowbook::Book;

/// Prints an error on stderr and exit the program
pub fn print_error(s: &str) -> ! {
    writeln!(&mut io::stderr(), "{} {}", Format::Error("error:"), s).unwrap();
    exit(0);
}

/// Gets the book options in a (key, value) list, or print an error
pub fn get_book_options<'a>(matches: &'a ArgMatches) -> Vec<(&'a str, &'a str)> {
    let mut output = vec!();
    if let Some(iter) = matches.values_of("set") {
        let v:Vec<_> = iter.collect();
        if v.len() %2 != 0 {
            print_error("An odd number of arguments was passed to --set, but it takes a list of key value pairs.");
        }

        for i in 0..v.len()/2 {
            let key = v[i * 2];
            let value = v[i * 2 + 1];
            output.push((key, value));
        }
    }
    output
}

/// sets the book options according to command line arguments
/// Also print these options to a string, so it can be used at
/// the creation of a book to check that parameters are OK and t
/// then print them to file
pub fn set_book_options(book: &mut Book, matches: &ArgMatches) -> String {
    let mut output = String::new();
    let options = get_book_options(matches);

    for (key, value) in options {
        let res = book.options.set(key, value);
        if let Err(err) = res {
            print_error(&format!("Error in setting key {}: {}", key, err));
        }
        output.push_str(&format!("{}: {}\n", key, value));
    }
    output
}

/// create a book file with the command line arguments
/// and exit the process at the end
pub fn create_book(matches: &ArgMatches) -> ! {
    let mut f:Box<Write> = if let Some(book) = matches.value_of("BOOK") {
        if fs::metadata(book).is_ok() {
            print_error(&format!("Could not create file {}: it already exists!", book));
        }
        Box::new(fs::File::create(book).unwrap())
    } else {
        Box::new(io::stdout())
    };

    if let Some(values) = matches.values_of("create") {
        if matches.is_present("set") {
            let mut book = Book::new();
            let s = set_book_options(&mut book, matches);
            f.write_all(&s.as_bytes()).unwrap();
        } else {
            f.write_all(b"author: Your name
title: Your title
lang: en

# Uncomment and fill to generate files
# output.html: some_file.html
# output.epub: some_file.epub
# output.pdf: some_file.pdf

# Uncomment and fill to set cover image (for Epub)
# cover: some_cover.png\n").unwrap();
        }
        f.write_all(b"\n# List of chapters\n").unwrap();
        for file in values {
            f.write_all(&format!("+ {}\n", file).as_bytes()).unwrap();
        }
        if let Some(s) = matches.value_of("BOOK") {
            println!("Created {}, now you'll have to complete it!", s);
        }
        exit(0);
    } else {
        unreachable!(); // because Clap takes care of it
    }
}

pub fn create_matches<'a>() -> ArgMatches<'a> {
    let app = App::new("crowbook")
        .version(env!("CARGO_PKG_VERSION"))
        .about("Render a markdown book in Epub, PDF or HTML.")
        .arg_from_usage("-s, --single 'Use a single Markdown file instead of a book configuration file'")
        .arg_from_usage("-v, --verbose 'Print warnings in parsing/rendering'")
        .arg(Arg::from_usage("-d, --debug 'Print debugging information'")
             .hidden(true))
        .arg_from_usage("-c, --create [FILES]... 'Creates a new book with existing markdown files'")
        .arg(Arg::from_usage("-o, --output [FILE] 'Specifies output file'")
             .requires("to"))
        .arg(Arg::from_usage("-t, --to [FORMAT] 'Generate specific format")
             .possible_values(&["epub", "pdf", "html", "tex", "odt"]))
        .arg(Arg::from_usage("--set [KEY_VALUES] 'Sets a list of book options'")
             .min_values(2))
        .arg_from_usage("-l --list-options 'Lists all possible options")
        .arg(Arg::from_usage("--list-options-md 'List all options, formatted in Markdown'")
             .hidden(true))
        .arg_from_usage("--print-template [TEMPLATE] 'Displays the default value of a template'")
        .arg(Arg::with_name("BOOK")
             .index(1)
             .help("File containing the book configuration, or a Markdown file when called with --single"));

    let matches = app.get_matches();

    pre_check(&matches);
    matches
}

/// Pre-check the matches to see if there isn't illegal options not detected by clap
fn pre_check(matches: &ArgMatches) {
    if matches.is_present("files") && !matches.is_present("create") {
        print_error("A list of additional files is only valid with the --create option.");
    }
}
<commit_msg>Fix typo in helpers.rs<commit_after>use clap::{App, Arg, Format, ArgMatches};
use std::io::{self, Write};
use std::process::exit;
use std::fs;
use crowbook::Book;

/// Prints an error on stderr and exit the program
pub fn print_error(s: &str) -> ! {
    writeln!(&mut io::stderr(), "{} {}", Format::Error("error:"), s).unwrap();
    exit(0);
}

/// Gets the book options in a (key, value) list, or print an error
pub fn get_book_options<'a>(matches: &'a ArgMatches) -> Vec<(&'a str, &'a str)> {
    let mut output = vec!();
    if let Some(iter) = matches.values_of("set") {
        let v:Vec<_> = iter.collect();
        if v.len() %2 != 0 {
            print_error("An odd number of arguments was passed to --set, but it takes a list of key value pairs.");
        }

        for i in 0..v.len()/2 {
            let key = v[i * 2];
            let value = v[i * 2 + 1];
            output.push((key, value));
        }
    }
    output
}

/// Sets the book options according to command line arguments
/// Also print these options to a string, so it can be used at
/// the creation of a book to check that parameters are OK and t
/// then print them to file
pub fn set_book_options(book: &mut Book, matches: &ArgMatches) -> String {
    let mut output = String::new();
    let options = get_book_options(matches);

    for (key, value) in options {
        let res = book.options.set(key, value);
        if let Err(err) = res {
            print_error(&format!("Error in setting key {}: {}", key, err));
        }
        output.push_str(&format!("{}: {}\n", key, value));
    }
    output
}

/// create a book file with the command line arguments
/// and exit the process at the end
pub fn create_book(matches: &ArgMatches) -> ! {
    let mut f:Box<Write> = if let Some(book) = matches.value_of("BOOK") {
        if fs::metadata(book).is_ok() {
            print_error(&format!("Could not create file {}: it already exists!", book));
        }
        Box::new(fs::File::create(book).unwrap())
    } else {
        Box::new(io::stdout())
    };

    if let Some(values) = matches.values_of("create") {
        if matches.is_present("set") {
            let mut book = Book::new();
            let s = set_book_options(&mut book, matches);
            f.write_all(&s.as_bytes()).unwrap();
        } else {
            f.write_all(b"author: Your name
title: Your title
lang: en

# Uncomment and fill to generate files
# output.html: some_file.html
# output.epub: some_file.epub
# output.pdf: some_file.pdf

# Uncomment and fill to set cover image (for Epub)
# cover: some_cover.png\n").unwrap();
        }
        f.write_all(b"\n# List of chapters\n").unwrap();
        for file in values {
            f.write_all(&format!("+ {}\n", file).as_bytes()).unwrap();
        }
        if let Some(s) = matches.value_of("BOOK") {
            println!("Created {}, now you'll have to complete it!", s);
        }
        exit(0);
    } else {
        unreachable!(); // because Clap takes care of it
    }
}

pub fn create_matches<'a>() -> ArgMatches<'a> {
    let app = App::new("crowbook")
        .version(env!("CARGO_PKG_VERSION"))
        .about("Render a markdown book in Epub, PDF or HTML.")
        .arg_from_usage("-s, --single 'Use a single Markdown file instead of a book configuration file'")
        .arg_from_usage("-v, --verbose 'Print warnings in parsing/rendering'")
        .arg(Arg::from_usage("-d, --debug 'Print debugging information'")
             .hidden(true))
        .arg_from_usage("-c, --create [FILES]... 'Creates a new book with existing markdown files'")
        .arg(Arg::from_usage("-o, --output [FILE] 'Specifies output file'")
             .requires("to"))
        .arg(Arg::from_usage("-t, --to [FORMAT] 'Generate specific format")
             .possible_values(&["epub", "pdf", "html", "tex", "odt"]))
        .arg(Arg::from_usage("--set [KEY_VALUES] 'Sets a list of book options'")
             .min_values(2))
        .arg_from_usage("-l --list-options 'Lists all possible options")
        .arg(Arg::from_usage("--list-options-md 'List all options, formatted in Markdown'")
             .hidden(true))
        .arg_from_usage("--print-template [TEMPLATE] 'Displays the default value of a template'")
        .arg(Arg::with_name("BOOK")
             .index(1)
             .help("File containing the book configuration, or a Markdown file when called with --single"));

    let matches = app.get_matches();

    pre_check(&matches);
    matches
}

/// Pre-check the matches to see if there isn't illegal options not detected by clap
fn pre_check(matches: &ArgMatches) {
    if matches.is_present("files") && !matches.is_present("create") {
        print_error("A list of additional files is only valid with the --create option.");
    }
}
<|endoftext|>


<commit_before>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::{self, Write};

use color;

/// The representation of a BMP encoder.
pub struct BMPEncoder<'a, W: 'a> {
    writer: &'a mut W,
}

impl<'a, W: Write + 'a> BMPEncoder<'a, W> {
    /// Create a new encoder that writes its output to ```w```.
    pub fn new(w: &'a mut W) -> Self {
        BMPEncoder { writer: w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```.
    pub fn encode(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        c: color::ColorType,
    ) -> io::Result<()> {
        let bmp_header_size = 14;
        let dib_header_size = 40; // using BITMAPINFOHEADER

        let (raw_pixel_size, written_pixel_size, palette_color_count) = try!(get_pixel_info(c));
        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes

        let image_size = width * height * written_pixel_size + (height * row_pad_size);
        let palette_size = palette_color_count * 4; // all palette colors are BGRA
        let file_size = bmp_header_size + dib_header_size + palette_size + image_size;

        // write BMP header
        try!(self.writer.write_u8(b'B'));
        try!(self.writer.write_u8(b'M'));
        try!(self.writer.write_u32::<LittleEndian>(file_size)); // file size
        try!(self.writer.write_u16::<LittleEndian>(0)); // reserved 1
        try!(self.writer.write_u16::<LittleEndian>(0)); // reserved 2
        try!(
            self.writer
                .write_u32::<LittleEndian>(bmp_header_size + dib_header_size + palette_size)
        ); // image data offset

        // write DIB header
        try!(self.writer.write_u32::<LittleEndian>(dib_header_size));
        try!(self.writer.write_i32::<LittleEndian>(width as i32));
        try!(self.writer.write_i32::<LittleEndian>(height as i32));
        try!(self.writer.write_u16::<LittleEndian>(1)); // color planes
        try!(
            self.writer
                .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)
        ); // bits per pixel
        try!(self.writer.write_u32::<LittleEndian>(0)); // compression method - no compression
        try!(self.writer.write_u32::<LittleEndian>(image_size));
        try!(self.writer.write_i32::<LittleEndian>(0)); // horizontal ppm
        try!(self.writer.write_i32::<LittleEndian>(0)); // vertical ppm
        try!(self.writer.write_u32::<LittleEndian>(palette_color_count));
        try!(self.writer.write_u32::<LittleEndian>(0)); // all colors are important

        // write image data
        match c {
            color::ColorType::RGB(8) | color::ColorType::RGBA(8) => {
                try!(self.encode_rgb(image, width, height, row_pad_size, raw_pixel_size))
            }
            color::ColorType::Gray(8) | color::ColorType::GrayA(8) => {
                try!(self.encode_gray(image, width, height, row_pad_size, raw_pixel_size))
            }
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &get_unsupported_error_message(c)[..],
                ))
            }
        }

        Ok(())
    }

    fn encode_rgb(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                let r = image[pixel_start];
                let g = image[pixel_start + 1];
                let b = image[pixel_start + 2];
                // written as BGR
                try!(self.writer.write_u8(b));
                try!(self.writer.write_u8(g));
                try!(self.writer.write_u8(r));
                // alpha is never written as it's not widely supported
            }

            try!(self.write_row_pad(row_pad_size));
        }

        Ok(())
    }

    fn encode_gray(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        // write grayscale palette
        for val in 0..256 {
            // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
            let val = val as u8;
            try!(self.writer.write_u8(val));
            try!(self.writer.write_u8(val));
            try!(self.writer.write_u8(val));
            try!(self.writer.write_u8(0));
        }

        // write image data
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                // color value is equal to the palette index
                try!(self.writer.write_u8(image[pixel_start]));
                // alpha is never written as it's not widely supported
            }

            try!(self.write_row_pad(row_pad_size));
        }

        Ok(())
    }

    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
        for _ in 0..row_pad_size {
            try!(self.writer.write_u8(0));
        }

        Ok(())
    }
}

fn get_unsupported_error_message(c: color::ColorType) -> String {
    format!(
        "Unsupported color type {:?}.  Supported types: RGB(8), RGBA(8), Gray(8), GrayA(8).",
        c
    )
}

/// Returns a tuple representing: (raw pixel size, written pixel size, palette color count).
fn get_pixel_info(c: color::ColorType) -> io::Result<(u32, u32, u32)> {
    let sizes = match c {
        color::ColorType::RGB(8) => (3, 3, 0),
        color::ColorType::RGBA(8) => (4, 3, 0),
        color::ColorType::Gray(8) => (1, 1, 256),
        color::ColorType::GrayA(8) => (2, 1, 256),
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                &get_unsupported_error_message(c)[..],
            ))
        }
    };

    Ok(sizes)
}

#[cfg(test)]
mod tests {
    use super::super::BMPDecoder;
    use super::BMPEncoder;
    use color::ColorType;
    use image::{DecodingResult, ImageDecoder};
    use std::io::Cursor;

    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ColorType) -> Vec<u8> {
        let mut encoded_data = Vec::new();
        {
            let mut encoder = BMPEncoder::new(&mut encoded_data);
            encoder
                .encode(&image, width, height, c)
                .expect("could not encode image");
        }

        let mut decoder = BMPDecoder::new(Cursor::new(&encoded_data));
        match decoder.read_image().expect("failed to decode") {
            DecodingResult::U8(decoded) => decoded,
            _ => panic!("failed to decode"),
        }
    }

    #[test]
    fn round_trip_single_pixel_rgb() {
        let image = [255u8, 0, 0]; // single red pixel
        let decoded = round_trip_image(&image, 1, 1, ColorType::RGB(8));
        assert_eq!(3, decoded.len());
        assert_eq!(255, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
    }

    #[test]
    fn round_trip_single_pixel_rgba() {
        let image = [255u8, 0, 0, 0]; // single red pixel
        let decoded = round_trip_image(&image, 1, 1, ColorType::RGBA(8));
        assert_eq!(3, decoded.len());
        assert_eq!(255, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
    }

    #[test]
    fn round_trip_3px_rgb() {
        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
        let _decoded = round_trip_image(&image, 3, 3, ColorType::RGB(8));
    }

    #[test]
    fn round_trip_gray() {
        let image = [0u8, 1, 2]; // 3 pixels
        let decoded = round_trip_image(&image, 3, 1, ColorType::Gray(8));
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }

    #[test]
    fn round_trip_graya() {
        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
        let decoded = round_trip_image(&image, 1, 3, ColorType::GrayA(8));
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }
}
<commit_msg>bmp/encoder: encode RGBA using bitfields with alpha mask<commit_after>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::{self, Write};

use color;

const BITMAPFILEHEADER_SIZE: u32 = 14;
const BITMAPINFOHEADER_SIZE: u32 = 40;
const BITMAPV4HEADER_SIZE: u32 = 108;

/// The representation of a BMP encoder.
pub struct BMPEncoder<'a, W: 'a> {
    writer: &'a mut W,
}

impl<'a, W: Write + 'a> BMPEncoder<'a, W> {
    /// Create a new encoder that writes its output to ```w```.
    pub fn new(w: &'a mut W) -> Self {
        BMPEncoder { writer: w }
    }

    /// Encodes the image ```image```
    /// that has dimensions ```width``` and ```height```
    /// and ```ColorType``` ```c```.
    pub fn encode(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        c: color::ColorType,
    ) -> io::Result<()> {
        let bmp_header_size = BITMAPFILEHEADER_SIZE;

        let (dib_header_size, written_pixel_size, palette_color_count) = try!(get_pixel_info(c));
        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes

        let image_size = width * height * written_pixel_size + (height * row_pad_size);
        let palette_size = palette_color_count * 4; // all palette colors are BGRA
        let file_size = bmp_header_size + dib_header_size + palette_size + image_size;

        // write BMP header
        try!(self.writer.write_u8(b'B'));
        try!(self.writer.write_u8(b'M'));
        try!(self.writer.write_u32::<LittleEndian>(file_size)); // file size
        try!(self.writer.write_u16::<LittleEndian>(0)); // reserved 1
        try!(self.writer.write_u16::<LittleEndian>(0)); // reserved 2
        try!(
            self.writer
                .write_u32::<LittleEndian>(bmp_header_size + dib_header_size + palette_size)
        ); // image data offset

        // write DIB header
        try!(self.writer.write_u32::<LittleEndian>(dib_header_size));
        try!(self.writer.write_i32::<LittleEndian>(width as i32));
        try!(self.writer.write_i32::<LittleEndian>(height as i32));
        try!(self.writer.write_u16::<LittleEndian>(1)); // color planes
        try!(
            self.writer
                .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)
        ); // bits per pixel
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            try!(self.writer.write_u32::<LittleEndian>(3)); // compression method - bitfields
        } else {
            try!(self.writer.write_u32::<LittleEndian>(0)); // compression method - no compression
        }
        try!(self.writer.write_u32::<LittleEndian>(image_size));
        try!(self.writer.write_i32::<LittleEndian>(0)); // horizontal ppm
        try!(self.writer.write_i32::<LittleEndian>(0)); // vertical ppm
        try!(self.writer.write_u32::<LittleEndian>(palette_color_count));
        try!(self.writer.write_u32::<LittleEndian>(0)); // all colors are important
        if dib_header_size >= BITMAPV4HEADER_SIZE {
            // Assume BGRA32
            try!(self.writer.write_u32::<LittleEndian>(0xff << 16)); // red mask
            try!(self.writer.write_u32::<LittleEndian>(0xff << 8)); // green mask
            try!(self.writer.write_u32::<LittleEndian>(0xff << 0)); // blue mask
            try!(self.writer.write_u32::<LittleEndian>(0xff << 24)); // alpha mask
            try!(self.writer.write_u32::<LittleEndian>(0x73524742)); // colorspace - sRGB
            // endpoints (3x3) and gamma (3)
            for _ in 0..12 {
                try!(self.writer.write_u32::<LittleEndian>(0));
            }
        }

        // write image data
        match c {
            color::ColorType::RGB(8) => {
                try!(self.encode_rgb(image, width, height, row_pad_size, 3))
            }
            color::ColorType::RGBA(8) => {
                try!(self.encode_rgba(image, width, height, row_pad_size, 4))
            }
            color::ColorType::Gray(8) => {
                try!(self.encode_gray(image, width, height, row_pad_size, 1))
            }
            color::ColorType::GrayA(8) => {
                try!(self.encode_gray(image, width, height, row_pad_size, 2))
            }
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &get_unsupported_error_message(c)[..],
                ))
            }
        }

        Ok(())
    }

    fn encode_rgb(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                let r = image[pixel_start];
                let g = image[pixel_start + 1];
                let b = image[pixel_start + 2];
                // written as BGR
                try!(self.writer.write_u8(b));
                try!(self.writer.write_u8(g));
                try!(self.writer.write_u8(r));
                // alpha is never written as it's not widely supported
            }

            try!(self.write_row_pad(row_pad_size));
        }

        Ok(())
    }

    fn encode_rgba(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                let r = image[pixel_start];
                let g = image[pixel_start + 1];
                let b = image[pixel_start + 2];
                let a = image[pixel_start + 3];
                // written as BGRA
                try!(self.writer.write_u8(b));
                try!(self.writer.write_u8(g));
                try!(self.writer.write_u8(r));
                try!(self.writer.write_u8(a));
            }

            try!(self.write_row_pad(row_pad_size));
        }

        Ok(())
    }

    fn encode_gray(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        // write grayscale palette
        for val in 0..256 {
            // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
            let val = val as u8;
            try!(self.writer.write_u8(val));
            try!(self.writer.write_u8(val));
            try!(self.writer.write_u8(val));
            try!(self.writer.write_u8(0));
        }

        // write image data
        let x_stride = bytes_per_pixel;
        let y_stride = width * x_stride;
        for row in 0..height {
            // from the bottom up
            let row_start = (height - row - 1) * y_stride;
            for col in 0..width {
                let pixel_start = (row_start + (col * x_stride)) as usize;
                // color value is equal to the palette index
                try!(self.writer.write_u8(image[pixel_start]));
                // alpha is never written as it's not widely supported
            }

            try!(self.write_row_pad(row_pad_size));
        }

        Ok(())
    }

    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
        for _ in 0..row_pad_size {
            try!(self.writer.write_u8(0));
        }

        Ok(())
    }
}

fn get_unsupported_error_message(c: color::ColorType) -> String {
    format!(
        "Unsupported color type {:?}.  Supported types: RGB(8), RGBA(8), Gray(8), GrayA(8).",
        c
    )
}

/// Returns a tuple representing: (dib header size, written pixel size, palette color count).
fn get_pixel_info(c: color::ColorType) -> io::Result<(u32, u32, u32)> {
    let sizes = match c {
        color::ColorType::RGB(8) => (BITMAPINFOHEADER_SIZE, 3, 0),
        color::ColorType::RGBA(8) => (BITMAPV4HEADER_SIZE, 4, 0),
        color::ColorType::Gray(8) => (BITMAPINFOHEADER_SIZE, 1, 256),
        color::ColorType::GrayA(8) => (BITMAPINFOHEADER_SIZE, 1, 256),
        _ => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                &get_unsupported_error_message(c)[..],
            ))
        }
    };

    Ok(sizes)
}

#[cfg(test)]
mod tests {
    use super::super::BMPDecoder;
    use super::BMPEncoder;
    use color::ColorType;
    use image::{DecodingResult, ImageDecoder};
    use std::io::Cursor;

    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ColorType) -> Vec<u8> {
        let mut encoded_data = Vec::new();
        {
            let mut encoder = BMPEncoder::new(&mut encoded_data);
            encoder
                .encode(&image, width, height, c)
                .expect("could not encode image");
        }

        let mut decoder = BMPDecoder::new(Cursor::new(&encoded_data));
        match decoder.read_image().expect("failed to decode") {
            DecodingResult::U8(decoded) => decoded,
            _ => panic!("failed to decode"),
        }
    }

    #[test]
    fn round_trip_single_pixel_rgb() {
        let image = [255u8, 0, 0]; // single red pixel
        let decoded = round_trip_image(&image, 1, 1, ColorType::RGB(8));
        assert_eq!(3, decoded.len());
        assert_eq!(255, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
    }

    #[test]
    fn round_trip_single_pixel_rgba() {
        let image = [1, 2, 3, 4];
        let decoded = round_trip_image(&image, 1, 1, ColorType::RGBA(8));
        assert_eq!(&decoded[..], &image[..]);
    }

    #[test]
    fn round_trip_3px_rgb() {
        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
        let _decoded = round_trip_image(&image, 3, 3, ColorType::RGB(8));
    }

    #[test]
    fn round_trip_gray() {
        let image = [0u8, 1, 2]; // 3 pixels
        let decoded = round_trip_image(&image, 3, 1, ColorType::Gray(8));
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }

    #[test]
    fn round_trip_graya() {
        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
        let decoded = round_trip_image(&image, 1, 3, ColorType::GrayA(8));
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }
}
<|endoftext|>


<commit_before>//! Insertion sort algorithm.
//!
//! The `insertion_sort` module contains the simple sorting algorithm "Insertion Sort".
//!
//! Source: https://en.wikipedia.org/wiki/Insertion_sort

use std::fmt::Debug;

pub fn sort<T: Debug + PartialOrd>(input: &Vec<T>) -> Vec<usize> {
    let n = input.len();

    let mut output_index: Vec<usize> = Vec::with_capacity(n);

    // Place the first element's index (0) as the first element of the output_index vector.
    output_index.push(0);

    println!("n = {:?}", n);
    println!("input:        {:?}", input);
    println!("output_index: {:?}", output_index);

    if n > 1 {
        // Loop over input vector, skipping the first element as it's already inserted as the first
        // element of output_index.
        for (i_start_at_0,elem) in input.iter().skip(1).enumerate() {
            let i = i_start_at_0 + 1;
            println!("  i={:?} Finding where to put element {:?} (i={:?}) in output_index: {:?}", i, elem, i, output_index);

            // Verify next element if it is smaller
            // If next element is larger than the last stored index, just place that element at
            // the end.
            if *elem >= input[*output_index.last().unwrap()] {
                // Next element "elem" is already larger than the previous stored index.
                // Just use its index.
                output_index.push(i);

                println!("    1. Pushing element '{:?}' at the end.", elem);
            } else {
                println!("    2. elem ?< input[output_index[0]]");
                println!("           {:?} ?< {:?}", elem, input[output_index[0]]);
                // The next element of the input vector is smaller than the last stored on.
                // Let's find where it should go by looping back on the stored indices.
                // NOTE: we must use a loop on indices as to prevent problems with the borrow-checker.
                if *elem < input[output_index[0]] {
                    println!("        2a. Yes: --> putting i in front");
                    output_index.insert(0, i);
                } else {
                    println!("        2b. No: --> Looping back");
                    for j in (0..i).rev() {
                        println!("            j={:?}   elem={:?}   input[output_index[j]]={:?}", j, elem, input[output_index[j]]);
                        if *elem > input[output_index[j]] {
                            output_index.insert(j+1, i);
                            break;
                        }
                    }
                }
                println!("        3. output_index: {:?}", output_index);
            }

            print!("           output: [");
            for j in output_index.iter() {
                    print!("{:?}, ", input[*j]);
            }
            println!("]");
        }
    }

    println!("");
    println!("END: input:        {:?}", input);
    println!("END: output_index: {:?}", output_index);
    print!("END: output:       [");
    for j in output_index.iter() {
            print!("{:?}, ", input[*j]);
    }
    println!("]");

    assert_eq!(output_index.len(), n);

    return output_index;
}
<commit_msg>Documentation for insertion sort.<commit_after>//! Insertion sort algorithm.
//!
//! The `insertion_sort` module contains the simple sorting algorithm "Insertion Sort".
//!
//! Source: https://en.wikipedia.org/wiki/Insertion_sort

use std::fmt::Debug;

/// Simple sort: insertion sort.
///
/// # Details
///
/// Insertion sort is one of the simplest sorting algorithm. Every elements of the list to sort
/// are picked one at a time and inserted into a the sorted list at the right location.
///
/// In the best case scenario, the list is already sorted and the work is simply to append the
/// element at the end of the list, hence a O(N) scaling.
///
/// In the worst case scenario, the list is in reverse order. Hence, every element is compared to
/// all previous elements already sorted until this new one gets inserted at the beginning of the
/// sorted list, thus making this case O(N^2).
///
/// Insertion sort is still useful as its simplicity (and thus small overhead) makes it ideal for
/// small vectors.
///
/// This function takes an immutable reference vector of any type and returns a vector of `usize`
/// of the same length containing the indices of the initial vector sorted.
///
/// # Examples
///
/// ```
/// let data: Vec<i32> = vec![4, 2, 3, 1];
/// assert_eq!(vec![3, 1, 2, 0], sorting::simplesorts::insertion_sort::sort(&data));
/// ```
///
pub fn sort<T: Debug + PartialOrd>(input: &Vec<T>) -> Vec<usize> {
    let n = input.len();

    let mut output_index: Vec<usize> = Vec::with_capacity(n);

    // Place the first element's index (0) as the first element of the output_index vector.
    output_index.push(0);

    println!("n = {:?}", n);
    println!("input:        {:?}", input);
    println!("output_index: {:?}", output_index);

    if n > 1 {
        // Loop over input vector, skipping the first element as it's already inserted as the first
        // element of output_index.
        for (i_start_at_0,elem) in input.iter().skip(1).enumerate() {
            let i = i_start_at_0 + 1;
            println!("  i={:?} Finding where to put element {:?} (i={:?}) in output_index: {:?}", i, elem, i, output_index);

            // Verify next element if it is smaller
            // If next element is larger than the last stored index, just place that element at
            // the end.
            if *elem >= input[*output_index.last().unwrap()] {
                // Next element "elem" is already larger than the previous stored index.
                // Just use its index.
                output_index.push(i);

                println!("    1. Pushing element '{:?}' at the end.", elem);
            } else {
                println!("    2. elem ?< input[output_index[0]]");
                println!("           {:?} ?< {:?}", elem, input[output_index[0]]);
                // The next element of the input vector is smaller than the last stored on.
                // Let's find where it should go by looping back on the stored indices.
                // NOTE: we must use a loop on indices as to prevent problems with the borrow-checker.
                if *elem < input[output_index[0]] {
                    println!("        2a. Yes: --> putting i in front");
                    output_index.insert(0, i);
                } else {
                    println!("        2b. No: --> Looping back");
                    for j in (0..i).rev() {
                        println!("            j={:?}   elem={:?}   input[output_index[j]]={:?}", j, elem, input[output_index[j]]);
                        if *elem > input[output_index[j]] {
                            output_index.insert(j+1, i);
                            break;
                        }
                    }
                }
                println!("        3. output_index: {:?}", output_index);
            }

            print!("           output: [");
            for j in output_index.iter() {
                    print!("{:?}, ", input[*j]);
            }
            println!("]");
        }
    }

    println!("");
    println!("END: input:        {:?}", input);
    println!("END: output_index: {:?}", output_index);
    print!("END: output:       [");
    for j in output_index.iter() {
            print!("{:?}, ", input[*j]);
    }
    println!("]");

    assert_eq!(output_index.len(), n);

    return output_index;
}
<|endoftext|>


<commit_before>#[derive(Debug, Deserialize, Default, Clone)]
pub struct Mass {
    pub value: f64,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Inertia {
    pub ixx: f64,
    pub ixy: f64,
    pub ixz: f64,
    pub iyy: f64,
    pub iyz: f64,
    pub izz: f64,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Inertial {
    #[serde(default)]
    pub origin: Pose,
    pub mass: Mass,
    pub inertia: Inertia,
}

#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "snake_case")]
pub enum Geometry {
    Box {
        #[serde(with = "urdf_vec3")]
        size: [f64; 3],
    },
    Cylinder { radius: f64, length: f64 },
    Sphere { radius: f64 },
    Mesh {
        filename: String,
        #[serde(default = "default_scale")]
        #[serde(with = "urdf_vec3")]
        scale: [f64; 3],
    },
}

fn default_scale() -> [f64; 3] {
    [1.0f64; 3]
}

fn default_one() -> f64 {
    1.0f64
}

impl Default for Geometry {
    fn default() -> Geometry {
        Geometry::Box { size: [0.0f64, 0.0, 0.0] }
    }
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Color {
    #[serde(with = "urdf_vec4")]
    #[serde(default = "default_rgba")]
    pub rgba: [f64; 4],
}

#[derive(Debug, Deserialize, Clone)]
pub struct Texture {
    pub filename: String,
}

impl Default for Texture {
    fn default() -> Texture {
        Texture { filename: "".to_string() }
    }
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Material {
    #[serde(default)]
    pub name: String,
    #[serde(default)]
    pub color: Color,
    #[serde(default)]
    pub texture: Texture,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Visual {
    #[serde(default)]
    pub name: String,
    #[serde(default)]
    pub origin: Pose,
    pub geometry: Geometry,
    #[serde(default)]
    pub material: Material,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Collision {
    #[serde(default)]
    pub name: String,
    #[serde(default)]
    pub origin: Pose,
    pub geometry: Geometry,
}

/// Urdf Link element
/// See http://wiki.ros.org/urdf/XML/link for more detail.
#[derive(Debug, Deserialize, Clone)]
pub struct Link {
    pub name: String,
    #[serde(default)]
    pub inertial: Inertial,
    #[serde(default)]
    pub visual: Vec<Visual>,
    #[serde(default)]
    pub collision: Vec<Collision>,
}

#[derive(Deserialize, Debug, Clone)]
pub struct Vec3 {
    #[serde(with = "urdf_vec3")]
    pub data: [f64; 3],
}

mod urdf_vec3 {
    use serde::{self, Deserialize, Deserializer};
    pub fn deserialize<D>(deserializer: D) -> Result<[f64; 3], D::Error>
    where
        D: Deserializer,
    {
        let s = String::deserialize(deserializer)?;
        let vec = s.split(' ')
            .filter_map(|x| x.parse::<f64>().ok())
            .collect::<Vec<_>>();
        if vec.len() != 3 {
            return Err(serde::de::Error::custom(
                format!("failed to parse float array in {}", s),
            ));
        }
        let mut arr = [0.0f64; 3];
        for i in 0..3 {
            arr[i] = vec[i];
        }
        Ok(arr)
    }
}

mod urdf_vec4 {
    use serde::{self, Deserialize, Deserializer};
    pub fn deserialize<D>(deserializer: D) -> Result<[f64; 4], D::Error>
    where
        D: Deserializer,
    {
        let s = String::deserialize(deserializer)?;
        let vec = s.split(' ')
            .filter_map(|x| x.parse::<f64>().ok())
            .collect::<Vec<_>>();
        if vec.len() != 4 {
            return Err(serde::de::Error::custom(
                format!("failed to parse float array in {}", s),
            ));
        }
        let mut arr = [0.0f64; 4];
        for i in 0..4 {
            arr[i] = vec[i];
        }
        Ok(arr)
    }
}

#[derive(Debug, Deserialize, Clone)]
pub struct Axis {
    #[serde(with = "urdf_vec3")]
    pub xyz: [f64; 3],
}

impl Default for Axis {
    fn default() -> Axis {
        Axis { xyz: [1.0f64, 0.0, 0.0] }
    }
}

#[derive(Debug, Deserialize, Clone)]
pub struct Pose {
    #[serde(with = "urdf_vec3")]
    #[serde(default = "default_zero3")]
    pub xyz: [f64; 3],
    #[serde(with = "urdf_vec3")]
    #[serde(default = "default_zero3")]
    pub rpy: [f64; 3],
}

fn default_zero3() -> [f64; 3] {
    [0.0f64, 0.0, 0.0]
}

fn default_rgba() -> [f64; 4] {
    [1.0f64, 1.0, 1.0, 1.0]
}

impl Default for Pose {
    fn default() -> Pose {
        Pose {
            xyz: default_zero3(),
            rpy: default_zero3(),
        }
    }
}

#[derive(Debug, Deserialize, Clone)]
pub struct LinkName {
    pub link: String,
}

#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "snake_case")]
pub enum JointType {
    Revolute,
    Continuous,
    Prismatic,
    Fixed,
    Floating,
    Planar,
    Free,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct JointLimit {
    #[serde(default)]
    pub lower: f64,
    #[serde(default)]
    pub upper: f64,
    #[serde(default)]
    pub effort: f64,
    #[serde(default)]
    pub velocity: f64,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Mimic {
    pub joint: String,
    #[serde(default = "default_one")]
    pub multiplier: f64,
    #[serde(default)]
    pub offset: f64,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct SafetyController {
    #[serde(default)]
    pub soft_lower_limit: f64,
    #[serde(default)]
    pub soft_upper_limit: f64,
    #[serde(default)]
    pub k_position: f64,
    pub k_velocity: f64,
}

/// Urdf Joint element
/// See http://wiki.ros.org/urdf/XML/joint for more detail.
#[derive(Debug, Deserialize, Clone)]
pub struct Joint {
    pub name: String,
    #[serde(rename = "type")]
    pub joint_type: JointType,
    #[serde(default)]
    pub origin: Pose,
    pub parent: LinkName,
    pub child: LinkName,
    #[serde(default)]
    pub axis: Axis,
    #[serde(default)]
    pub limit: JointLimit,
    #[serde(default)]
    pub dynamics: Dynamics,
    #[serde(default)]
    pub mimic: Mimic,
    #[serde(default)]
    pub safety_controller: SafetyController,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Dynamics {
    #[serde(default)]
    pub damping: f64,
    #[serde(default)]
    pub friction: f64,
}

/// Top level struct to access urdf.
#[derive(Debug, Deserialize, Clone)]
pub struct Robot {
    #[serde(default)]
    pub name: String,

    #[serde(rename = "link", default)]
    pub links: Vec<Link>,

    #[serde(rename = "joint", default)]
    pub joints: Vec<Joint>,

    #[serde(rename = "material", default)]
    pub materials: Vec<Material>,
}
<commit_msg>add non-standard capsule geometry and spherical joint<commit_after>#[derive(Debug, Deserialize, Default, Clone)]
pub struct Mass {
    pub value: f64,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Inertia {
    pub ixx: f64,
    pub ixy: f64,
    pub ixz: f64,
    pub iyy: f64,
    pub iyz: f64,
    pub izz: f64,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Inertial {
    #[serde(default)]
    pub origin: Pose,
    pub mass: Mass,
    pub inertia: Inertia,
}

#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "snake_case")]
pub enum Geometry {
    Box {
        #[serde(with = "urdf_vec3")]
        size: [f64; 3],
    },
    Cylinder { radius: f64, length: f64 },
    Capsule { radius: f64, length: f64 },
    Sphere { radius: f64 },
    Mesh {
        filename: String,
        #[serde(default = "default_scale")]
        #[serde(with = "urdf_vec3")]
        scale: [f64; 3],
    },
}

fn default_scale() -> [f64; 3] {
    [1.0f64; 3]
}

fn default_one() -> f64 {
    1.0f64
}

impl Default for Geometry {
    fn default() -> Geometry {
        Geometry::Box { size: [0.0f64, 0.0, 0.0] }
    }
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Color {
    #[serde(with = "urdf_vec4")]
    #[serde(default = "default_rgba")]
    pub rgba: [f64; 4],
}

#[derive(Debug, Deserialize, Clone)]
pub struct Texture {
    pub filename: String,
}

impl Default for Texture {
    fn default() -> Texture {
        Texture { filename: "".to_string() }
    }
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Material {
    #[serde(default)]
    pub name: String,
    #[serde(default)]
    pub color: Color,
    #[serde(default)]
    pub texture: Texture,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Visual {
    #[serde(default)]
    pub name: String,
    #[serde(default)]
    pub origin: Pose,
    pub geometry: Geometry,
    #[serde(default)]
    pub material: Material,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Collision {
    #[serde(default)]
    pub name: String,
    #[serde(default)]
    pub origin: Pose,
    pub geometry: Geometry,
}

/// Urdf Link element
/// See http://wiki.ros.org/urdf/XML/link for more detail.
#[derive(Debug, Deserialize, Clone)]
pub struct Link {
    pub name: String,
    #[serde(default)]
    pub inertial: Inertial,
    #[serde(default)]
    pub visual: Vec<Visual>,
    #[serde(default)]
    pub collision: Vec<Collision>,
}

#[derive(Deserialize, Debug, Clone)]
pub struct Vec3 {
    #[serde(with = "urdf_vec3")]
    pub data: [f64; 3],
}

mod urdf_vec3 {
    use serde::{self, Deserialize, Deserializer};
    pub fn deserialize<D>(deserializer: D) -> Result<[f64; 3], D::Error>
    where
        D: Deserializer,
    {
        let s = String::deserialize(deserializer)?;
        let vec = s.split(' ')
            .filter_map(|x| x.parse::<f64>().ok())
            .collect::<Vec<_>>();
        if vec.len() != 3 {
            return Err(serde::de::Error::custom(
                format!("failed to parse float array in {}", s),
            ));
        }
        let mut arr = [0.0f64; 3];
        for i in 0..3 {
            arr[i] = vec[i];
        }
        Ok(arr)
    }
}

mod urdf_vec4 {
    use serde::{self, Deserialize, Deserializer};
    pub fn deserialize<D>(deserializer: D) -> Result<[f64; 4], D::Error>
    where
        D: Deserializer,
    {
        let s = String::deserialize(deserializer)?;
        let vec = s.split(' ')
            .filter_map(|x| x.parse::<f64>().ok())
            .collect::<Vec<_>>();
        if vec.len() != 4 {
            return Err(serde::de::Error::custom(
                format!("failed to parse float array in {}", s),
            ));
        }
        let mut arr = [0.0f64; 4];
        for i in 0..4 {
            arr[i] = vec[i];
        }
        Ok(arr)
    }
}

#[derive(Debug, Deserialize, Clone)]
pub struct Axis {
    #[serde(with = "urdf_vec3")]
    pub xyz: [f64; 3],
}

impl Default for Axis {
    fn default() -> Axis {
        Axis { xyz: [1.0f64, 0.0, 0.0] }
    }
}

#[derive(Debug, Deserialize, Clone)]
pub struct Pose {
    #[serde(with = "urdf_vec3")]
    #[serde(default = "default_zero3")]
    pub xyz: [f64; 3],
    #[serde(with = "urdf_vec3")]
    #[serde(default = "default_zero3")]
    pub rpy: [f64; 3],
}

fn default_zero3() -> [f64; 3] {
    [0.0f64, 0.0, 0.0]
}

fn default_rgba() -> [f64; 4] {
    [1.0f64, 1.0, 1.0, 1.0]
}

impl Default for Pose {
    fn default() -> Pose {
        Pose {
            xyz: default_zero3(),
            rpy: default_zero3(),
        }
    }
}

#[derive(Debug, Deserialize, Clone)]
pub struct LinkName {
    pub link: String,
}

#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "snake_case")]
pub enum JointType {
    Revolute,
    Continuous,
    Prismatic,
    Fixed,
    Floating,
    Planar,
    Free,
    Spherical,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct JointLimit {
    #[serde(default)]
    pub lower: f64,
    #[serde(default)]
    pub upper: f64,
    #[serde(default)]
    pub effort: f64,
    #[serde(default)]
    pub velocity: f64,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Mimic {
    pub joint: String,
    #[serde(default = "default_one")]
    pub multiplier: f64,
    #[serde(default)]
    pub offset: f64,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct SafetyController {
    #[serde(default)]
    pub soft_lower_limit: f64,
    #[serde(default)]
    pub soft_upper_limit: f64,
    #[serde(default)]
    pub k_position: f64,
    pub k_velocity: f64,
}

/// Urdf Joint element
/// See http://wiki.ros.org/urdf/XML/joint for more detail.
#[derive(Debug, Deserialize, Clone)]
pub struct Joint {
    pub name: String,
    #[serde(rename = "type")]
    pub joint_type: JointType,
    #[serde(default)]
    pub origin: Pose,
    pub parent: LinkName,
    pub child: LinkName,
    #[serde(default)]
    pub axis: Axis,
    #[serde(default)]
    pub limit: JointLimit,
    #[serde(default)]
    pub dynamics: Dynamics,
    #[serde(default)]
    pub mimic: Mimic,
    #[serde(default)]
    pub safety_controller: SafetyController,
}

#[derive(Debug, Deserialize, Default, Clone)]
pub struct Dynamics {
    #[serde(default)]
    pub damping: f64,
    #[serde(default)]
    pub friction: f64,
}

/// Top level struct to access urdf.
#[derive(Debug, Deserialize, Clone)]
pub struct Robot {
    #[serde(default)]
    pub name: String,

    #[serde(rename = "link", default)]
    pub links: Vec<Link>,

    #[serde(rename = "joint", default)]
    pub joints: Vec<Joint>,

    #[serde(rename = "material", default)]
    pub materials: Vec<Material>,
}
<|endoftext|>