<commit_before>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::rc::Rc;

use attributes;
use arena::TypedArena;
use back::symbol_names;
use llvm::{ValueRef, get_params};
use rustc::hir::def_id::DefId;
use rustc::ty::subst::{Subst, Substs};
use rustc::traits::{self, Reveal};
use abi::FnType;
use base::*;
use build::*;
use callee::{Callee, Virtual, trans_fn_pointer_shim};
use closure;
use common::*;
use consts;
use debuginfo::DebugLoc;
use declare;
use glue;
use machine;
use type_::Type;
use type_of::*;
use value::Value;
use rustc::ty::{self, Ty, TyCtxt, TypeFoldable};

use syntax::ast::Name;
use syntax_pos::DUMMY_SP;

// drop_glue pointer, size, align.
const VTABLE_OFFSET: usize = 3;

/// Extracts a method from a trait object's vtable, at the specified index.
pub fn get_virtual_method<'blk, 'tcx>(bcx: Block<'blk, 'tcx>,
                                      llvtable: ValueRef,
                                      vtable_index: usize)
                                      -> ValueRef {
    // Load the data pointer from the object.
    debug!("get_virtual_method(vtable_index={}, llvtable={:?})",
           vtable_index, Value(llvtable));

    Load(bcx, GEPi(bcx, llvtable, &[vtable_index + VTABLE_OFFSET]))
}

/// Generate a shim function that allows an object type like `SomeTrait` to
/// implement the type `SomeTrait`. Imagine a trait definition:
///
///    trait SomeTrait { fn get(&self) -> i32; ... }
///
/// And a generic bit of code:
///
///    fn foo<T:SomeTrait>(t: &T) {
///        let x = SomeTrait::get;
///        x(t)
///    }
///
/// What is the value of `x` when `foo` is invoked with `T=SomeTrait`?
/// The answer is that it is a shim function generated by this routine:
///
///    fn shim(t: &SomeTrait) -> i32 {
///        // ... call t.get() virtually ...
///    }
///
/// In fact, all virtual calls can be thought of as normal trait calls
/// that go through this shim function.
pub fn trans_object_shim<'a, 'tcx>(ccx: &'a CrateContext<'a, 'tcx>,
                                   method_ty: Ty<'tcx>,
                                   vtable_index: usize)
                                   -> ValueRef {
    let _icx = push_ctxt("trans_object_shim");
    let tcx = ccx.tcx();

    debug!("trans_object_shim(vtable_index={}, method_ty={:?})",
           vtable_index,
           method_ty);

    let sig = tcx.erase_late_bound_regions(&method_ty.fn_sig());
    let sig = tcx.normalize_associated_type(&sig);
    let fn_ty = FnType::new(ccx, method_ty.fn_abi(), &sig, &[]);

    let function_name =
        symbol_names::internal_name_from_type_and_suffix(ccx, method_ty, "object_shim");
    let llfn = declare::define_internal_fn(ccx, &function_name, method_ty);
    attributes::set_frame_pointer_elimination(ccx, llfn);

    let (block_arena, fcx): (TypedArena<_>, FunctionContext);
    block_arena = TypedArena::new();
    fcx = FunctionContext::new(ccx, llfn, fn_ty, None, &block_arena);
    let mut bcx = fcx.init(false);

    let dest = fcx.llretslotptr.get();

    debug!("trans_object_shim: method_offset_in_vtable={}",
           vtable_index);

    let llargs = get_params(fcx.llfn);

    let callee = Callee {
        data: Virtual(vtable_index),
        ty: method_ty
    };
    bcx = callee.call(bcx, DebugLoc::None,
                      &llargs[fcx.fn_ty.ret.is_indirect() as usize..], dest).bcx;

    fcx.finish(bcx, DebugLoc::None);

    llfn
}

/// Creates a returns a dynamic vtable for the given type and vtable origin.
/// This is used only for objects.
///
/// The `trait_ref` encodes the erased self type. Hence if we are
/// making an object `Foo<Trait>` from a value of type `Foo<T>`, then
/// `trait_ref` would map `T:Trait`.
pub fn get_vtable<'a, 'tcx>(ccx: &CrateContext<'a, 'tcx>,
                            trait_ref: ty::PolyTraitRef<'tcx>)
                            -> ValueRef
{
    let tcx = ccx.tcx();
    let _icx = push_ctxt("meth::get_vtable");

    debug!("get_vtable(trait_ref={:?})", trait_ref);

    // Check the cache.
    match ccx.vtables().borrow().get(&trait_ref) {
        Some(&val) => { return val }
        None => { }
    }

    // Not in the cache. Build it.
    let methods = traits::supertraits(tcx, trait_ref.clone()).flat_map(|trait_ref| {
        let vtable = fulfill_obligation(ccx.shared(), DUMMY_SP, trait_ref.clone());
        match vtable {
            // Should default trait error here?
            traits::VtableDefaultImpl(_) |
            traits::VtableBuiltin(_) => {
                Vec::new().into_iter()
            }
            traits::VtableImpl(
                traits::VtableImplData {
                    impl_def_id: id,
                    substs,
                    nested: _ }) => {
                let nullptr = C_null(Type::nil(ccx).ptr_to());
                get_vtable_methods(tcx, id, substs)
                    .into_iter()
                    .map(|opt_mth| opt_mth.map_or(nullptr, |mth| {
                        Callee::def(ccx, mth.method.def_id, &mth.substs).reify(ccx)
                    }))
                    .collect::<Vec<_>>()
                    .into_iter()
            }
            traits::VtableClosure(
                traits::VtableClosureData {
                    closure_def_id,
                    substs,
                    nested: _ }) => {
                let trait_closure_kind = tcx.lang_items.fn_trait_kind(trait_ref.def_id()).unwrap();
                let llfn = closure::trans_closure_method(ccx,
                                                         closure_def_id,
                                                         substs,
                                                         trait_closure_kind);
                vec![llfn].into_iter()
            }
            traits::VtableFnPointer(
                traits::VtableFnPointerData {
                    fn_ty: bare_fn_ty,
                    nested: _ }) => {
                let trait_closure_kind = tcx.lang_items.fn_trait_kind(trait_ref.def_id()).unwrap();
                vec![trans_fn_pointer_shim(ccx, trait_closure_kind, bare_fn_ty)].into_iter()
            }
            traits::VtableObject(ref data) => {
                // this would imply that the Self type being erased is
                // an object type; this cannot happen because we
                // cannot cast an unsized type into a trait object
                bug!("cannot get vtable for an object type: {:?}",
                     data);
            }
            traits::VtableParam(..) => {
                bug!("resolved vtable for {:?} to bad vtable {:?} in trans",
                     trait_ref,
                     vtable);
            }
        }
    });

    let size_ty = sizing_type_of(ccx, trait_ref.self_ty());
    let size = machine::llsize_of_alloc(ccx, size_ty);
    let align = align_of(ccx, trait_ref.self_ty());

    let components: Vec<_> = vec![
        // Generate a destructor for the vtable.
        glue::get_drop_glue(ccx, trait_ref.self_ty()),
        C_uint(ccx, size),
        C_uint(ccx, align)
    ].into_iter().chain(methods).collect();

    let vtable_const = C_struct(ccx, &components, false);
    let align = machine::llalign_of_pref(ccx, val_ty(vtable_const));
    let vtable = consts::addr_of(ccx, vtable_const, align, "vtable");

    ccx.vtables().borrow_mut().insert(trait_ref, vtable);
    vtable
}

pub fn get_vtable_methods<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                    impl_id: DefId,
                                    substs: &'tcx Substs<'tcx>)
                                    -> Vec<Option<ImplMethod<'tcx>>>
{
    debug!("get_vtable_methods(impl_id={:?}, substs={:?}", impl_id, substs);

    let trait_id = match tcx.impl_trait_ref(impl_id) {
        Some(t_id) => t_id.def_id,
        None       => bug!("make_impl_vtable: don't know how to \
                            make a vtable for a type impl!")
    };

    tcx.populate_implementations_for_trait_if_necessary(trait_id);

    let trait_item_def_ids = tcx.trait_item_def_ids(trait_id);
    trait_item_def_ids
        .iter()

        // Filter out non-method items.
        .filter_map(|item_def_id| {
            match *item_def_id {
                ty::MethodTraitItemId(def_id) => Some(def_id),
                _ => None,
            }
        })

        // Now produce pointers for each remaining method. If the
        // method could never be called from this object, just supply
        // null.
        .map(|trait_method_def_id| {
            debug!("get_vtable_methods: trait_method_def_id={:?}",
                   trait_method_def_id);

            let trait_method_type = match tcx.impl_or_trait_item(trait_method_def_id) {
                ty::MethodTraitItem(m) => m,
                _ => bug!("should be a method, not other assoc item"),
            };
            let name = trait_method_type.name;

            // Some methods cannot be called on an object; skip those.
            if !tcx.is_vtable_safe_method(trait_id, &trait_method_type) {
                debug!("get_vtable_methods: not vtable safe");
                return None;
            }

            debug!("get_vtable_methods: trait_method_type={:?}",
                   trait_method_type);

            // the method may have some early-bound lifetimes, add
            // regions for those
            let method_substs = Substs::for_item(tcx, trait_method_def_id,
                                                 |_, _| tcx.mk_region(ty::ReErased),
                                                 |_, _| tcx.types.err);

            // The substitutions we have are on the impl, so we grab
            // the method type from the impl to substitute into.
            let mth = get_impl_method(tcx, method_substs, impl_id, substs, name);

            debug!("get_vtable_methods: mth={:?}", mth);

            // If this is a default method, it's possible that it
            // relies on where clauses that do not hold for this
            // particular set of type parameters. Note that this
            // method could then never be called, so we do not want to
            // try and trans it, in that case. Issue #23435.
            if mth.is_provided {
                let predicates = mth.method.predicates.predicates.subst(tcx, &mth.substs);
                if !normalize_and_test_predicates(tcx, predicates) {
                    debug!("get_vtable_methods: predicates do not hold");
                    return None;
                }
            }

            Some(mth)
        })
        .collect()
}

#[derive(Debug)]
pub struct ImplMethod<'tcx> {
    pub method: Rc<ty::Method<'tcx>>,
    pub substs: &'tcx Substs<'tcx>,
    pub is_provided: bool
}

/// Locates the applicable definition of a method, given its name.
pub fn get_impl_method<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                 substs: &'tcx Substs<'tcx>,
                                 impl_def_id: DefId,
                                 impl_substs: &'tcx Substs<'tcx>,
                                 name: Name)
                                 -> ImplMethod<'tcx>
{
    assert!(!substs.needs_infer());

    let trait_def_id = tcx.trait_id_of_impl(impl_def_id).unwrap();
    let trait_def = tcx.lookup_trait_def(trait_def_id);

    match trait_def.ancestors(impl_def_id).fn_defs(tcx, name).next() {
        Some(node_item) => {
            let substs = tcx.infer_ctxt(None, None, Reveal::All).enter(|infcx| {
                let substs = substs.rebase_onto(tcx, trait_def_id, impl_substs);
                let substs = traits::translate_substs(&infcx, impl_def_id,
                                                      substs, node_item.node);
                tcx.lift(&substs).unwrap_or_else(|| {
                    bug!("trans::meth::get_impl_method: translate_substs \
                          returned {:?} which contains inference types/regions",
                         substs);
                })
            });
            ImplMethod {
                method: node_item.item,
                substs: substs,
                is_provided: node_item.node.is_from_trait(),
            }
        }
        None => {
            bug!("method {:?} not found in {:?}", name, impl_def_id)
        }
    }
}
<commit_msg>clean up `get_vtable`'s doc comment<commit_after>// Copyright 2012 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::rc::Rc;

use attributes;
use arena::TypedArena;
use back::symbol_names;
use llvm::{ValueRef, get_params};
use rustc::hir::def_id::DefId;
use rustc::ty::subst::{Subst, Substs};
use rustc::traits::{self, Reveal};
use abi::FnType;
use base::*;
use build::*;
use callee::{Callee, Virtual, trans_fn_pointer_shim};
use closure;
use common::*;
use consts;
use debuginfo::DebugLoc;
use declare;
use glue;
use machine;
use type_::Type;
use type_of::*;
use value::Value;
use rustc::ty::{self, Ty, TyCtxt, TypeFoldable};

use syntax::ast::Name;
use syntax_pos::DUMMY_SP;

// drop_glue pointer, size, align.
const VTABLE_OFFSET: usize = 3;

/// Extracts a method from a trait object's vtable, at the specified index.
pub fn get_virtual_method<'blk, 'tcx>(bcx: Block<'blk, 'tcx>,
                                      llvtable: ValueRef,
                                      vtable_index: usize)
                                      -> ValueRef {
    // Load the data pointer from the object.
    debug!("get_virtual_method(vtable_index={}, llvtable={:?})",
           vtable_index, Value(llvtable));

    Load(bcx, GEPi(bcx, llvtable, &[vtable_index + VTABLE_OFFSET]))
}

/// Generate a shim function that allows an object type like `SomeTrait` to
/// implement the type `SomeTrait`. Imagine a trait definition:
///
///    trait SomeTrait { fn get(&self) -> i32; ... }
///
/// And a generic bit of code:
///
///    fn foo<T:SomeTrait>(t: &T) {
///        let x = SomeTrait::get;
///        x(t)
///    }
///
/// What is the value of `x` when `foo` is invoked with `T=SomeTrait`?
/// The answer is that it is a shim function generated by this routine:
///
///    fn shim(t: &SomeTrait) -> i32 {
///        // ... call t.get() virtually ...
///    }
///
/// In fact, all virtual calls can be thought of as normal trait calls
/// that go through this shim function.
pub fn trans_object_shim<'a, 'tcx>(ccx: &'a CrateContext<'a, 'tcx>,
                                   method_ty: Ty<'tcx>,
                                   vtable_index: usize)
                                   -> ValueRef {
    let _icx = push_ctxt("trans_object_shim");
    let tcx = ccx.tcx();

    debug!("trans_object_shim(vtable_index={}, method_ty={:?})",
           vtable_index,
           method_ty);

    let sig = tcx.erase_late_bound_regions(&method_ty.fn_sig());
    let sig = tcx.normalize_associated_type(&sig);
    let fn_ty = FnType::new(ccx, method_ty.fn_abi(), &sig, &[]);

    let function_name =
        symbol_names::internal_name_from_type_and_suffix(ccx, method_ty, "object_shim");
    let llfn = declare::define_internal_fn(ccx, &function_name, method_ty);
    attributes::set_frame_pointer_elimination(ccx, llfn);

    let (block_arena, fcx): (TypedArena<_>, FunctionContext);
    block_arena = TypedArena::new();
    fcx = FunctionContext::new(ccx, llfn, fn_ty, None, &block_arena);
    let mut bcx = fcx.init(false);

    let dest = fcx.llretslotptr.get();

    debug!("trans_object_shim: method_offset_in_vtable={}",
           vtable_index);

    let llargs = get_params(fcx.llfn);

    let callee = Callee {
        data: Virtual(vtable_index),
        ty: method_ty
    };
    bcx = callee.call(bcx, DebugLoc::None,
                      &llargs[fcx.fn_ty.ret.is_indirect() as usize..], dest).bcx;

    fcx.finish(bcx, DebugLoc::None);

    llfn
}

/// Creates a dynamic vtable for the given type and vtable origin.
/// This is used only for objects.
///
/// The vtables are cached instead of created on every call.
///
/// The `trait_ref` encodes the erased self type. Hence if we are
/// making an object `Foo<Trait>` from a value of type `Foo<T>`, then
/// `trait_ref` would map `T:Trait`.
pub fn get_vtable<'a, 'tcx>(ccx: &CrateContext<'a, 'tcx>,
                            trait_ref: ty::PolyTraitRef<'tcx>)
                            -> ValueRef
{
    let tcx = ccx.tcx();
    let _icx = push_ctxt("meth::get_vtable");

    debug!("get_vtable(trait_ref={:?})", trait_ref);

    // Check the cache.
    match ccx.vtables().borrow().get(&trait_ref) {
        Some(&val) => { return val }
        None => { }
    }

    // Not in the cache. Build it.
    let methods = traits::supertraits(tcx, trait_ref.clone()).flat_map(|trait_ref| {
        let vtable = fulfill_obligation(ccx.shared(), DUMMY_SP, trait_ref.clone());
        match vtable {
            // Should default trait error here?
            traits::VtableDefaultImpl(_) |
            traits::VtableBuiltin(_) => {
                Vec::new().into_iter()
            }
            traits::VtableImpl(
                traits::VtableImplData {
                    impl_def_id: id,
                    substs,
                    nested: _ }) => {
                let nullptr = C_null(Type::nil(ccx).ptr_to());
                get_vtable_methods(tcx, id, substs)
                    .into_iter()
                    .map(|opt_mth| opt_mth.map_or(nullptr, |mth| {
                        Callee::def(ccx, mth.method.def_id, &mth.substs).reify(ccx)
                    }))
                    .collect::<Vec<_>>()
                    .into_iter()
            }
            traits::VtableClosure(
                traits::VtableClosureData {
                    closure_def_id,
                    substs,
                    nested: _ }) => {
                let trait_closure_kind = tcx.lang_items.fn_trait_kind(trait_ref.def_id()).unwrap();
                let llfn = closure::trans_closure_method(ccx,
                                                         closure_def_id,
                                                         substs,
                                                         trait_closure_kind);
                vec![llfn].into_iter()
            }
            traits::VtableFnPointer(
                traits::VtableFnPointerData {
                    fn_ty: bare_fn_ty,
                    nested: _ }) => {
                let trait_closure_kind = tcx.lang_items.fn_trait_kind(trait_ref.def_id()).unwrap();
                vec![trans_fn_pointer_shim(ccx, trait_closure_kind, bare_fn_ty)].into_iter()
            }
            traits::VtableObject(ref data) => {
                // this would imply that the Self type being erased is
                // an object type; this cannot happen because we
                // cannot cast an unsized type into a trait object
                bug!("cannot get vtable for an object type: {:?}",
                     data);
            }
            traits::VtableParam(..) => {
                bug!("resolved vtable for {:?} to bad vtable {:?} in trans",
                     trait_ref,
                     vtable);
            }
        }
    });

    let size_ty = sizing_type_of(ccx, trait_ref.self_ty());
    let size = machine::llsize_of_alloc(ccx, size_ty);
    let align = align_of(ccx, trait_ref.self_ty());

    let components: Vec<_> = vec![
        // Generate a destructor for the vtable.
        glue::get_drop_glue(ccx, trait_ref.self_ty()),
        C_uint(ccx, size),
        C_uint(ccx, align)
    ].into_iter().chain(methods).collect();

    let vtable_const = C_struct(ccx, &components, false);
    let align = machine::llalign_of_pref(ccx, val_ty(vtable_const));
    let vtable = consts::addr_of(ccx, vtable_const, align, "vtable");

    ccx.vtables().borrow_mut().insert(trait_ref, vtable);
    vtable
}

pub fn get_vtable_methods<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                    impl_id: DefId,
                                    substs: &'tcx Substs<'tcx>)
                                    -> Vec<Option<ImplMethod<'tcx>>>
{
    debug!("get_vtable_methods(impl_id={:?}, substs={:?}", impl_id, substs);

    let trait_id = match tcx.impl_trait_ref(impl_id) {
        Some(t_id) => t_id.def_id,
        None       => bug!("make_impl_vtable: don't know how to \
                            make a vtable for a type impl!")
    };

    tcx.populate_implementations_for_trait_if_necessary(trait_id);

    let trait_item_def_ids = tcx.trait_item_def_ids(trait_id);
    trait_item_def_ids
        .iter()

        // Filter out non-method items.
        .filter_map(|item_def_id| {
            match *item_def_id {
                ty::MethodTraitItemId(def_id) => Some(def_id),
                _ => None,
            }
        })

        // Now produce pointers for each remaining method. If the
        // method could never be called from this object, just supply
        // null.
        .map(|trait_method_def_id| {
            debug!("get_vtable_methods: trait_method_def_id={:?}",
                   trait_method_def_id);

            let trait_method_type = match tcx.impl_or_trait_item(trait_method_def_id) {
                ty::MethodTraitItem(m) => m,
                _ => bug!("should be a method, not other assoc item"),
            };
            let name = trait_method_type.name;

            // Some methods cannot be called on an object; skip those.
            if !tcx.is_vtable_safe_method(trait_id, &trait_method_type) {
                debug!("get_vtable_methods: not vtable safe");
                return None;
            }

            debug!("get_vtable_methods: trait_method_type={:?}",
                   trait_method_type);

            // the method may have some early-bound lifetimes, add
            // regions for those
            let method_substs = Substs::for_item(tcx, trait_method_def_id,
                                                 |_, _| tcx.mk_region(ty::ReErased),
                                                 |_, _| tcx.types.err);

            // The substitutions we have are on the impl, so we grab
            // the method type from the impl to substitute into.
            let mth = get_impl_method(tcx, method_substs, impl_id, substs, name);

            debug!("get_vtable_methods: mth={:?}", mth);

            // If this is a default method, it's possible that it
            // relies on where clauses that do not hold for this
            // particular set of type parameters. Note that this
            // method could then never be called, so we do not want to
            // try and trans it, in that case. Issue #23435.
            if mth.is_provided {
                let predicates = mth.method.predicates.predicates.subst(tcx, &mth.substs);
                if !normalize_and_test_predicates(tcx, predicates) {
                    debug!("get_vtable_methods: predicates do not hold");
                    return None;
                }
            }

            Some(mth)
        })
        .collect()
}

#[derive(Debug)]
pub struct ImplMethod<'tcx> {
    pub method: Rc<ty::Method<'tcx>>,
    pub substs: &'tcx Substs<'tcx>,
    pub is_provided: bool
}

/// Locates the applicable definition of a method, given its name.
pub fn get_impl_method<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
                                 substs: &'tcx Substs<'tcx>,
                                 impl_def_id: DefId,
                                 impl_substs: &'tcx Substs<'tcx>,
                                 name: Name)
                                 -> ImplMethod<'tcx>
{
    assert!(!substs.needs_infer());

    let trait_def_id = tcx.trait_id_of_impl(impl_def_id).unwrap();
    let trait_def = tcx.lookup_trait_def(trait_def_id);

    match trait_def.ancestors(impl_def_id).fn_defs(tcx, name).next() {
        Some(node_item) => {
            let substs = tcx.infer_ctxt(None, None, Reveal::All).enter(|infcx| {
                let substs = substs.rebase_onto(tcx, trait_def_id, impl_substs);
                let substs = traits::translate_substs(&infcx, impl_def_id,
                                                      substs, node_item.node);
                tcx.lift(&substs).unwrap_or_else(|| {
                    bug!("trans::meth::get_impl_method: translate_substs \
                          returned {:?} which contains inference types/regions",
                         substs);
                })
            });
            ImplMethod {
                method: node_item.item,
                substs: substs,
                is_provided: node_item.node.is_from_trait(),
            }
        }
        None => {
            bug!("method {:?} not found in {:?}", name, impl_def_id)
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation of `std::os` functionality for unix systems

#![allow(unused_imports)] // lots of cfg code here

use os::unix::prelude::*;

use error::Error as StdError;
use ffi::{OsString, OsStr};
use fmt;
use io::{self, Read, Write};
use iter;
use libc::{self, c_int, c_char, c_void};
use marker::PhantomData;
use mem;
use memchr;
use path::{self, PathBuf};
use ptr;
use slice;
use str;
use sys_common::mutex::Mutex;
use sys::cvt;
use sys::fd;
use vec;

const TMPBUF_SZ: usize = 128;
static ENV_LOCK: Mutex = Mutex::new();

/// Returns the platform-specific value of errno
pub fn errno() -> i32 {
    0
}

/// Gets a detailed string description for the given error number.
pub fn error_string(errno: i32) -> String {
    if let Some(string) = libc::STR_ERROR.get(errno as usize) {
        string.to_string()
    } else {
        "unknown error".to_string()
    }
}

pub fn getcwd() -> io::Result<PathBuf> {
    let mut buf = [0; 4096];
    let count = cvt(libc::getcwd(&mut buf))?;
    Ok(PathBuf::from(OsString::from_vec(buf[.. count].to_vec())))
}

pub fn chdir(p: &path::Path) -> io::Result<()> {
    cvt(libc::chdir(p.to_str().unwrap())).and(Ok(()))
}

pub struct SplitPaths<'a> {
    iter: iter::Map<slice::Split<'a, u8, fn(&u8) -> bool>,
                    fn(&'a [u8]) -> PathBuf>,
}

pub fn split_paths(unparsed: &OsStr) -> SplitPaths {
    fn bytes_to_path(b: &[u8]) -> PathBuf {
        PathBuf::from(<OsStr as OsStrExt>::from_bytes(b))
    }
    fn is_colon(b: &u8) -> bool { *b == b':' }
    let unparsed = unparsed.as_bytes();
    SplitPaths {
        iter: unparsed.split(is_colon as fn(&u8) -> bool)
                      .map(bytes_to_path as fn(&[u8]) -> PathBuf)
    }
}

impl<'a> Iterator for SplitPaths<'a> {
    type Item = PathBuf;
    fn next(&mut self) -> Option<PathBuf> { self.iter.next() }
    fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
}

#[derive(Debug)]
pub struct JoinPathsError;

pub fn join_paths<I, T>(paths: I) -> Result<OsString, JoinPathsError>
    where I: Iterator<Item=T>, T: AsRef<OsStr>
{
    let mut joined = Vec::new();
    let sep = b':';

    for (i, path) in paths.enumerate() {
        let path = path.as_ref().as_bytes();
        if i > 0 { joined.push(sep) }
        if path.contains(&sep) {
            return Err(JoinPathsError)
        }
        joined.extend_from_slice(path);
    }
    Ok(OsStringExt::from_vec(joined))
}

impl fmt::Display for JoinPathsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        "path segment contains separator `:`".fmt(f)
    }
}

impl StdError for JoinPathsError {
    fn description(&self) -> &str { "failed to join paths" }
}

pub fn current_exe() -> io::Result<PathBuf> {
    use io::ErrorKind;
    Err(io::Error::new(ErrorKind::Other, "Not yet implemented on redox"))
}

pub struct Env {
    iter: vec::IntoIter<(OsString, OsString)>,
    _dont_send_or_sync_me: PhantomData<*mut ()>,
}

impl Iterator for Env {
    type Item = (OsString, OsString);
    fn next(&mut self) -> Option<(OsString, OsString)> { self.iter.next() }
    fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
}

/// Returns a vector of (variable, value) byte-vector pairs for all the
/// environment variables of the current process.
pub fn env() -> Env {
    let mut variables: Vec<(OsString, OsString)> = Vec::new();
    if let Ok(mut file) = ::fs::File::open("env:") {
        let mut string = String::new();
        if file.read_to_string(&mut string).is_ok() {
            for line in string.lines() {
                if let Some(equal_sign) = line.chars().position(|c| c == '=') {
                    let name = line.chars().take(equal_sign).collect::<String>();
                    let value = line.chars().skip(equal_sign+1).collect::<String>();
                    variables.push((OsString::from(name), OsString::from(value)));
                }
            }
        }
    }
    Env { iter: variables.into_iter(), _dont_send_or_sync_me: PhantomData }
}

pub fn getenv(key: &OsStr) -> io::Result<Option<OsString>> {
    if ! key.is_empty() {
        if let Ok(mut file) = ::fs::File::open(&("env:".to_owned() + key.to_str().unwrap())) {
            let mut string = String::new();
            file.read_to_string(&mut string)?;
            Ok(Some(OsString::from(string)))
        } else {
            Ok(None)
        }
    } else {
        Ok(None)
    }
}

pub fn setenv(key: &OsStr, value: &OsStr) -> io::Result<()> {
    if ! key.is_empty() {
        let mut file = ::fs::File::open(&("env:".to_owned() + key.to_str().unwrap()))?;
        file.write_all(value.as_bytes())?;
        file.set_len(value.len() as u64)?;
    }
    Ok(())
}

pub fn unsetenv(key: &OsStr) -> io::Result<()> {
    ::fs::remove_file(&("env:".to_owned() + key.to_str().unwrap()))?;
    Ok(())
}

pub fn page_size() -> usize {
    4096
}

pub fn temp_dir() -> PathBuf {
    ::env::var_os("TMPDIR").map(PathBuf::from).unwrap_or_else(|| {
        PathBuf::from("/tmp")
    })
}

pub fn home_dir() -> Option<PathBuf> {
    return ::env::var_os("HOME").map(PathBuf::from);
}

pub fn exit(code: i32) -> ! {
    let _ = libc::exit(code as usize);
    unreachable!();
}
<commit_msg>Add current_exe support<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Implementation of `std::os` functionality for unix systems

#![allow(unused_imports)] // lots of cfg code here

use os::unix::prelude::*;

use error::Error as StdError;
use ffi::{OsString, OsStr};
use fmt;
use io::{self, Read, Write};
use iter;
use libc::{self, c_int, c_char, c_void};
use marker::PhantomData;
use mem;
use memchr;
use path::{self, PathBuf};
use ptr;
use slice;
use str;
use sys_common::mutex::Mutex;
use sys::cvt;
use sys::fd;
use vec;

const TMPBUF_SZ: usize = 128;
static ENV_LOCK: Mutex = Mutex::new();

/// Returns the platform-specific value of errno
pub fn errno() -> i32 {
    0
}

/// Gets a detailed string description for the given error number.
pub fn error_string(errno: i32) -> String {
    if let Some(string) = libc::STR_ERROR.get(errno as usize) {
        string.to_string()
    } else {
        "unknown error".to_string()
    }
}

pub fn getcwd() -> io::Result<PathBuf> {
    let mut buf = [0; 4096];
    let count = cvt(libc::getcwd(&mut buf))?;
    Ok(PathBuf::from(OsString::from_vec(buf[.. count].to_vec())))
}

pub fn chdir(p: &path::Path) -> io::Result<()> {
    cvt(libc::chdir(p.to_str().unwrap())).and(Ok(()))
}

pub struct SplitPaths<'a> {
    iter: iter::Map<slice::Split<'a, u8, fn(&u8) -> bool>,
                    fn(&'a [u8]) -> PathBuf>,
}

pub fn split_paths(unparsed: &OsStr) -> SplitPaths {
    fn bytes_to_path(b: &[u8]) -> PathBuf {
        PathBuf::from(<OsStr as OsStrExt>::from_bytes(b))
    }
    fn is_colon(b: &u8) -> bool { *b == b':' }
    let unparsed = unparsed.as_bytes();
    SplitPaths {
        iter: unparsed.split(is_colon as fn(&u8) -> bool)
                      .map(bytes_to_path as fn(&[u8]) -> PathBuf)
    }
}

impl<'a> Iterator for SplitPaths<'a> {
    type Item = PathBuf;
    fn next(&mut self) -> Option<PathBuf> { self.iter.next() }
    fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
}

#[derive(Debug)]
pub struct JoinPathsError;

pub fn join_paths<I, T>(paths: I) -> Result<OsString, JoinPathsError>
    where I: Iterator<Item=T>, T: AsRef<OsStr>
{
    let mut joined = Vec::new();
    let sep = b':';

    for (i, path) in paths.enumerate() {
        let path = path.as_ref().as_bytes();
        if i > 0 { joined.push(sep) }
        if path.contains(&sep) {
            return Err(JoinPathsError)
        }
        joined.extend_from_slice(path);
    }
    Ok(OsStringExt::from_vec(joined))
}

impl fmt::Display for JoinPathsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        "path segment contains separator `:`".fmt(f)
    }
}

impl StdError for JoinPathsError {
    fn description(&self) -> &str { "failed to join paths" }
}

pub fn current_exe() -> io::Result<PathBuf> {
    use fs::File;

    let mut file = File::open("sys:exe")?;

    let mut path = String::new();
    file.read_to_string(&mut path)?;

    if path.ends_with('\n') {
        path.pop();
    }

    Ok(PathBuf::from(path))
}

pub struct Env {
    iter: vec::IntoIter<(OsString, OsString)>,
    _dont_send_or_sync_me: PhantomData<*mut ()>,
}

impl Iterator for Env {
    type Item = (OsString, OsString);
    fn next(&mut self) -> Option<(OsString, OsString)> { self.iter.next() }
    fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
}

/// Returns a vector of (variable, value) byte-vector pairs for all the
/// environment variables of the current process.
pub fn env() -> Env {
    let mut variables: Vec<(OsString, OsString)> = Vec::new();
    if let Ok(mut file) = ::fs::File::open("env:") {
        let mut string = String::new();
        if file.read_to_string(&mut string).is_ok() {
            for line in string.lines() {
                if let Some(equal_sign) = line.chars().position(|c| c == '=') {
                    let name = line.chars().take(equal_sign).collect::<String>();
                    let value = line.chars().skip(equal_sign+1).collect::<String>();
                    variables.push((OsString::from(name), OsString::from(value)));
                }
            }
        }
    }
    Env { iter: variables.into_iter(), _dont_send_or_sync_me: PhantomData }
}

pub fn getenv(key: &OsStr) -> io::Result<Option<OsString>> {
    if ! key.is_empty() {
        if let Ok(mut file) = ::fs::File::open(&("env:".to_owned() + key.to_str().unwrap())) {
            let mut string = String::new();
            file.read_to_string(&mut string)?;
            Ok(Some(OsString::from(string)))
        } else {
            Ok(None)
        }
    } else {
        Ok(None)
    }
}

pub fn setenv(key: &OsStr, value: &OsStr) -> io::Result<()> {
    if ! key.is_empty() {
        let mut file = ::fs::File::open(&("env:".to_owned() + key.to_str().unwrap()))?;
        file.write_all(value.as_bytes())?;
        file.set_len(value.len() as u64)?;
    }
    Ok(())
}

pub fn unsetenv(key: &OsStr) -> io::Result<()> {
    ::fs::remove_file(&("env:".to_owned() + key.to_str().unwrap()))?;
    Ok(())
}

pub fn page_size() -> usize {
    4096
}

pub fn temp_dir() -> PathBuf {
    ::env::var_os("TMPDIR").map(PathBuf::from).unwrap_or_else(|| {
        PathBuf::from("/tmp")
    })
}

pub fn home_dir() -> Option<PathBuf> {
    return ::env::var_os("HOME").map(PathBuf::from);
}

pub fn exit(code: i32) -> ! {
    let _ = libc::exit(code as usize);
    unreachable!();
}
<|endoftext|>


<commit_before>use std::cmp;
use std::fmt;
use std::marker::PhantomData;
use std::ops::Range;
use std::usize;

pub trait ReprUsize {
    fn from_usize(usize) -> Self;
    fn into_usize(self) -> usize;
}

impl ReprUsize for bool {
    fn from_usize(i: usize) -> bool {
        match i {
            0 => false,
            1 => true,
            _ => unreachable!(),
        }
    }

    fn into_usize(self) -> usize { self as usize }
}

impl ReprUsize for char {
    fn from_usize(i: usize) -> char { i as u8 as char }
    fn into_usize(self) -> usize { self as u8 as usize }
}

macro_rules! impl_predefined_type {
    ($ty: ty) => {
        impl ReprUsize for $ty {
            fn from_usize(i: usize) -> $ty { i as $ty }
            fn into_usize(self) -> usize { self as usize }
        }
    };
}

impl_predefined_type!(u8);
impl_predefined_type!(u16);
impl_predefined_type!(u32);
impl_predefined_type!(u64);
impl_predefined_type!(usize);
impl_predefined_type!(i8);
impl_predefined_type!(i16);
impl_predefined_type!(i32);
impl_predefined_type!(i64);
impl_predefined_type!(isize);
impl_predefined_type!(f32);
impl_predefined_type!(f64);

#[derive(Clone, Hash)]
pub struct BitsVec<T: ReprUsize> {
    inner: Vec<usize>,
    units: usize,
    bits: usize,
    max_bits: usize,
    leftover: usize,
    _marker: PhantomData<T>,
}

impl<T: ReprUsize> BitsVec<T> {
    pub fn new(bits: usize) -> BitsVec<T> {
        let max = usize::MAX.count_ones() as usize;
        // We can store more bits, but then we might need BigInt to get them out!
        assert!(bits < max, "cannot hold more than {} bits at a time", max - 1);

        BitsVec {
            inner: vec![0],
            units: 0,
            bits: bits,
            max_bits: max,
            leftover: max,
            _marker: PhantomData,
        }
    }

    pub fn with_capacity(bits: usize, capacity: usize) -> BitsVec<T> {
        let mut vec = BitsVec::new(bits);
        vec.reserve(capacity);
        vec
    }

    pub fn push(&mut self, value: T) {
        let mut value = value.into_usize();
        assert!(value >> self.bits == 0,
                "input size is more than allowed size ({} >= {})", value, 2usize.pow(self.bits as u32));

        let mut idx = self.inner.len() - 1;
        let shift;

        if self.leftover < self.bits {
            let left = self.bits - self.leftover;
            self.inner[idx] |= value >> left;
            if self.leftover != 0 {     // special case, in which masking would result in zero!
                value &= (1 << left) - 1;
            }

            self.inner.push(0);
            self.leftover = self.max_bits - left;
            shift = self.max_bits - left;
            idx += 1;
        } else {
            shift = self.leftover - self.bits;
            self.leftover -= self.bits;
        }

        value <<= shift;
        self.inner[idx] |= value;
        self.units += 1;
    }

    pub fn get(&self, i: usize) -> Option<T> {
        if i >= self.units {
            return None
        }

        let idx = i * self.bits / self.max_bits;
        let bits = (i * self.bits) % self.max_bits;
        let diff = self.max_bits - bits;
        let mut val = self.inner[idx];
        if bits != 0 {
            val &= (1 << diff) - 1;
        }

        if diff >= self.bits {
            Some(T::from_usize(val >> (diff - self.bits)))
        } else {
            let shift = self.bits - diff;
            let out = (val << shift) | (self.inner[idx + 1] >> (self.max_bits - shift));
            Some(T::from_usize(out))
        }
    }

    pub fn set(&mut self, i: usize, value: T) {
        let value = value.into_usize();
        assert!(i < self.units, "index out of bounds ({} >= {})", i, self.units);
        assert!(value >> self.bits == 0,
                "input size is more than allowed size ({} >= {})", value, 2usize.pow(self.bits as u32));

        let idx = i * self.bits / self.max_bits;
        let bits = (i * self.bits) % self.max_bits;
        let diff = self.max_bits - bits;
        let mut val = self.inner[idx];

        if diff >= self.bits {
            let shift = diff - self.bits;
            let last = val & ((1 << shift) - 1);
            let mask = if bits == 0 { 0 } else { ((1 << bits) - 1) << diff };   // prevent overflow
            val &= mask;
            val |= value << shift;
            self.inner[idx] = val | last;
        } else {
            let shift = self.bits - diff;
            self.inner[idx] >>= diff;
            self.inner[idx] <<= diff;
            self.inner[idx] |= value >> shift;
            let last = value & ((1 << shift) - 1);
            self.inner[idx + 1] &= (1 << (self.max_bits - shift)) - 1;
            self.inner[idx + 1] |= last << (self.max_bits - shift);
        }
    }

    pub fn len(&self) -> usize {
        self.units
    }

    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional * self.bits / self.max_bits + 1);
    }

    pub fn inner_len(&self) -> usize {
        self.inner.len()
    }

    pub fn iter(&self) -> Iter<T> {
        Iter { vec: self, range: 0..self.units }
    }

    pub fn into_iter(self) -> IntoIter<T> {
        IntoIter { range: 0..self.units, vec: self }
    }
}

impl<T: ReprUsize + Clone> BitsVec<T> {
    pub fn extend_with_element(&mut self, length: usize, value: T) {
        assert!(length > self.len(), "final length should be greater than the initial length");
        // Three phases (somewhat inefficient, using safe code and all, but much better than `push`)
        let mut remain = length - self.len();
        self.reserve(remain);

        // 1. Slow push until we get to a common multiple of (self.bits, self.max_bits)
        while self.leftover > 0 {
            self.push(value.clone());
            remain -= 1;
            if remain == 0 {
                return
            }
        }

        // 2. Do the same to a new BitsVec
        let mut temp = BitsVec::new(self.bits);
        temp.reserve(cmp::min(remain, self.max_bits));
        temp.push(value.clone());
        while temp.leftover > 0 && remain > 0 {
            temp.push(value.clone());
        }

        if remain == 0 {
            self.units += temp.units;
            self.inner.extend(&temp.inner);
            return
        }

        // 3. Extend from the new BitsVec
        while remain >= temp.units {
            self.inner.extend(&temp.inner);
            self.units += temp.units;
            remain -= temp.units;
        }

        for _ in 0..remain {
            self.push(value.clone());
        }
    }
}

impl<T: ReprUsize + fmt::Debug> fmt::Debug for BitsVec<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_list().entries(self.iter()).finish()
    }
}

pub struct Iter<'a, T: ReprUsize + 'a> {
    vec: &'a BitsVec<T>,
    range: Range<usize>,
}

impl<'a, T: ReprUsize> IntoIterator for &'a BitsVec<T> {
    type Item = T;
    type IntoIter = Iter<'a, T>;

    fn into_iter(self) -> Iter<'a, T> {
        self.iter()
    }
}

impl<'a, T: ReprUsize> Iterator for Iter<'a, T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.range.next().and_then(|i| self.vec.get(i))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.range.size_hint()
    }
}

impl<'a, T: ReprUsize> DoubleEndedIterator for Iter<'a, T> {
    fn next_back(&mut self) -> Option<T> {
        self.range.next_back().and_then(|i| self.vec.get(i))
    }
}

impl<'a, T: ReprUsize> ExactSizeIterator for Iter<'a, T> {}

pub struct IntoIter<T: ReprUsize> {
    vec: BitsVec<T>,
    range: Range<usize>,
}

impl<T: ReprUsize> IntoIterator for BitsVec<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;

    fn into_iter(self) -> IntoIter<T> {
        self.into_iter()
    }
}

impl<T: ReprUsize> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.range.next().and_then(|i| self.vec.get(i))
    }
}

impl<T: ReprUsize> DoubleEndedIterator for IntoIter<T> {
    fn next_back(&mut self) -> Option<T> {
        self.range.next_back().and_then(|i| self.vec.get(i))
    }
}

impl<T: ReprUsize> ExactSizeIterator for IntoIter<T> {}
<commit_msg>Some easy methods from Vec<T><commit_after>use std::cmp;
use std::fmt;
use std::marker::PhantomData;
use std::mem;
use std::ops::Range;
use std::usize;

pub trait ReprUsize {
    fn from_usize(usize) -> Self;
    fn into_usize(self) -> usize;
}

impl ReprUsize for bool {
    fn from_usize(i: usize) -> bool {
        match i {
            0 => false,
            1 => true,
            _ => unreachable!(),
        }
    }

    fn into_usize(self) -> usize { self as usize }
}

impl ReprUsize for char {
    fn from_usize(i: usize) -> char { i as u8 as char }
    fn into_usize(self) -> usize { self as u8 as usize }
}

macro_rules! impl_predefined_type {
    ($ty: ty) => {
        impl ReprUsize for $ty {
            fn from_usize(i: usize) -> $ty { i as $ty }
            fn into_usize(self) -> usize { self as usize }
        }
    };
}

impl_predefined_type!(u8);
impl_predefined_type!(u16);
impl_predefined_type!(u32);
impl_predefined_type!(u64);
impl_predefined_type!(usize);
impl_predefined_type!(i8);
impl_predefined_type!(i16);
impl_predefined_type!(i32);
impl_predefined_type!(i64);
impl_predefined_type!(isize);
impl_predefined_type!(f32);
impl_predefined_type!(f64);

#[derive(Clone, Hash)]
pub struct BitsVec<T: ReprUsize> {
    inner: Vec<usize>,
    units: usize,
    bits: usize,
    max_bits: usize,
    leftover: usize,
    _marker: PhantomData<T>,
}

impl<T: ReprUsize> BitsVec<T> {
    pub fn new(bits: usize) -> BitsVec<T> {
        let max = usize::MAX.count_ones() as usize;
        // We can store more bits, but then we might need BigInt to get them out!
        assert!(bits < max, "cannot hold more than {} bits at a time", max - 1);

        BitsVec {
            inner: vec![0],
            units: 0,
            bits: bits,
            max_bits: max,
            leftover: max,
            _marker: PhantomData,
        }
    }

    pub fn with_capacity(bits: usize, capacity: usize) -> BitsVec<T> {
        let mut vec = BitsVec::new(bits);
        vec.reserve(capacity);
        vec
    }

    pub fn push(&mut self, value: T) {
        let mut value = value.into_usize();
        assert!(value >> self.bits == 0,
                "input size is more than allowed size ({} >= {})", value, 2usize.pow(self.bits as u32));

        let mut idx = self.inner.len() - 1;
        let shift;

        if self.leftover < self.bits {
            let left = self.bits - self.leftover;
            self.inner[idx] |= value >> left;
            if self.leftover != 0 {     // special case, in which masking would result in zero!
                value &= (1 << left) - 1;
            }

            self.inner.push(0);
            self.leftover = self.max_bits - left;
            shift = self.max_bits - left;
            idx += 1;
        } else {
            shift = self.leftover - self.bits;
            self.leftover -= self.bits;
        }

        value <<= shift;
        self.inner[idx] |= value;
        self.units += 1;
    }

    pub fn get(&self, i: usize) -> Option<T> {
        if i >= self.units {
            return None
        }

        let idx = i * self.bits / self.max_bits;
        let bits = (i * self.bits) % self.max_bits;
        let diff = self.max_bits - bits;
        let mut val = self.inner[idx];
        if bits != 0 {
            val &= (1 << diff) - 1;
        }

        if diff >= self.bits {
            Some(T::from_usize(val >> (diff - self.bits)))
        } else {
            let shift = self.bits - diff;
            let out = (val << shift) | (self.inner[idx + 1] >> (self.max_bits - shift));
            Some(T::from_usize(out))
        }
    }

    pub fn set(&mut self, i: usize, value: T) {
        let value = value.into_usize();
        assert!(i < self.units, "index out of bounds ({} >= {})", i, self.units);
        assert!(value >> self.bits == 0,
                "input size is more than allowed size ({} >= {})", value, 2usize.pow(self.bits as u32));

        let idx = i * self.bits / self.max_bits;
        let bits = (i * self.bits) % self.max_bits;
        let diff = self.max_bits - bits;
        let mut val = self.inner[idx];

        if diff >= self.bits {
            let shift = diff - self.bits;
            let last = val & ((1 << shift) - 1);
            let mask = if bits == 0 { 0 } else { ((1 << bits) - 1) << diff };   // prevent overflow
            val &= mask;
            val |= value << shift;
            self.inner[idx] = val | last;
        } else {
            let shift = self.bits - diff;
            self.inner[idx] >>= diff;
            self.inner[idx] <<= diff;
            self.inner[idx] |= value >> shift;
            let last = value & ((1 << shift) - 1);
            self.inner[idx + 1] &= (1 << (self.max_bits - shift)) - 1;
            self.inner[idx + 1] |= last << (self.max_bits - shift);
        }
    }

    pub fn len(&self) -> usize {
        self.units
    }

    pub fn is_empty(&self) -> bool {
        self.units == 0
    }

    pub fn reserve(&mut self, additional: usize) {
        self.inner.reserve(additional * self.bits / self.max_bits + 1);
    }

    pub fn clear(&mut self) {
        let bits = self.bits;
        mem::replace(self, BitsVec::new(bits));
    }

    pub fn inner_len(&self) -> usize {
        self.inner.len()
    }

    pub fn iter(&self) -> Iter<T> {
        Iter { vec: self, range: 0..self.units }
    }

    pub fn into_iter(self) -> IntoIter<T> {
        IntoIter { range: 0..self.units, vec: self }
    }
}

impl<T: ReprUsize + Clone> BitsVec<T> {
    pub fn with_elements(bits: usize, length: usize, value: T) -> BitsVec<T> {
        let mut vec = BitsVec::new(bits);
        vec.extend_with_element(length, value);
        vec
    }

    pub fn extend_with_element(&mut self, length: usize, value: T) {
        assert!(length > self.len(), "final length should be greater than the initial length");
        // Three phases (somewhat inefficient, using safe code and all, but much better than `push`)
        let mut remain = length - self.len();
        self.reserve(remain);

        // 1. Slow push until we get to a common multiple of (self.bits, self.max_bits)
        while self.leftover > 0 {
            self.push(value.clone());
            remain -= 1;
            if remain == 0 {
                return
            }
        }

        // 2. Do the same to a new BitsVec
        let mut temp = BitsVec::new(self.bits);
        temp.reserve(cmp::min(remain, self.max_bits));
        temp.push(value.clone());
        while temp.leftover > 0 && remain > 0 {
            temp.push(value.clone());
        }

        if remain == 0 {
            self.units += temp.units;
            self.inner.extend(&temp.inner);
            return
        }

        // 3. Extend from the new BitsVec
        while remain >= temp.units {
            self.inner.extend(&temp.inner);
            self.units += temp.units;
            remain -= temp.units;
        }

        for _ in 0..remain {
            self.push(value.clone());
        }
    }
}

impl<T: ReprUsize + PartialEq> BitsVec<T> {
    pub fn contains(&self, element: &T) -> bool {
        self.iter().find(|ref i| i == &element).is_some()
    }
}

impl<T: ReprUsize + fmt::Debug> fmt::Debug for BitsVec<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_list().entries(self.iter()).finish()
    }
}

impl<T: ReprUsize> PartialEq for BitsVec<T> {
    fn eq(&self, other: &Self) -> bool {
        if self.units != other.units || self.bits != other.bits {
            return false
        }

        self.inner == other.inner
    }
}

pub struct Iter<'a, T: ReprUsize + 'a> {
    vec: &'a BitsVec<T>,
    range: Range<usize>,
}

impl<'a, T: ReprUsize> IntoIterator for &'a BitsVec<T> {
    type Item = T;
    type IntoIter = Iter<'a, T>;

    fn into_iter(self) -> Iter<'a, T> {
        self.iter()
    }
}

impl<'a, T: ReprUsize> Iterator for Iter<'a, T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.range.next().and_then(|i| self.vec.get(i))
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.range.size_hint()
    }
}

impl<'a, T: ReprUsize> DoubleEndedIterator for Iter<'a, T> {
    fn next_back(&mut self) -> Option<T> {
        self.range.next_back().and_then(|i| self.vec.get(i))
    }
}

impl<'a, T: ReprUsize> ExactSizeIterator for Iter<'a, T> {}

pub struct IntoIter<T: ReprUsize> {
    vec: BitsVec<T>,
    range: Range<usize>,
}

impl<T: ReprUsize> IntoIterator for BitsVec<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;

    fn into_iter(self) -> IntoIter<T> {
        self.into_iter()
    }
}

impl<T: ReprUsize> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.range.next().and_then(|i| self.vec.get(i))
    }
}

impl<T: ReprUsize> DoubleEndedIterator for IntoIter<T> {
    fn next_back(&mut self) -> Option<T> {
        self.range.next_back().and_then(|i| self.vec.get(i))
    }
}

impl<T: ReprUsize> ExactSizeIterator for IntoIter<T> {}
<|endoftext|>


<commit_before>// Copyright 2017 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

pub mod msg;
pub mod codec;
pub mod error;

mod session;

use std::io;

use std::rc::Rc;
use std::cell::Cell;
use std::net::ToSocketAddrs;
use std::collections::VecDeque;

use futures::{StartSend, Async, AsyncSink, Poll};
use futures::future::Future;
use futures::stream::Stream;
use futures::sink::Sink;

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core as TokioCore;
use tokio_core::reactor::Handle as TokioCoreHandle;

use tokio_io::AsyncRead;

use futures_ex::StreamEx;

use net::arrow::proto::codec::{ArrowCodec, FromBytes};
use net::arrow::proto::error::ArrowError;
use net::arrow::proto::msg::ArrowMessage;
use net::arrow::proto::msg::control::{
    ACK_NO_ERROR,
    ACK_UNSUPPORTED_PROTOCOL_VERSION,
    ACK_UNAUTHORIZED,
    ACK_INTERNAL_SERVER_ERROR,

    AckMessage,
    ControlMessage,
    ControlMessageType,
    HupMessage,
    RedirectMessage
};
use net::arrow::proto::session::SessionManager;
use net::utils::Timeout;

pub use net::arrow::proto::msg::control::svc_table::{Service, ServiceTable};

/// Currently supported version of the Arrow protocol.
pub const ARROW_PROTOCOL_VERSION: u8 = 1;

const ACK_TIMEOUT: u64 = 20000;

/// Arrow Protocol states.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum ProtocolState {
    Handshake,
    Established
}

/// Arrow Client implementation.
struct ArrowClient {
    tc_handle:     TokioCoreHandle,
    cmsg_factory:  ControlMessageFactory,
    sessions:      SessionManager,
    messages:      VecDeque<ArrowMessage>,
    expected_acks: VecDeque<u16>,
    ack_timeout:   Timeout,
    state:         ProtocolState,
    redirect:      Option<String>,
}

impl ArrowClient {
    /// Create a new Arrow Client.
    fn new<T>(tc_handle: TokioCoreHandle, svc_table: T) -> ArrowClient
        where T: 'static + ServiceTable {
        let cmsg_factory = ControlMessageFactory::new();
        let session_manager = SessionManager::new(
            tc_handle.clone(),
            svc_table,
            cmsg_factory.clone());

        let messages = VecDeque::new();

        // TODO: add REGISTER message into the message queue

        ArrowClient {
            tc_handle:     tc_handle,
            cmsg_factory:  cmsg_factory,
            sessions:      session_manager,
            messages:      messages,
            expected_acks: VecDeque::new(),
            ack_timeout:   Timeout::new(),
            state:         ProtocolState::Handshake,
            redirect:      None,
        }
    }

    /// Get redirect address (if any).
    fn get_redirect(&self) -> Option<&str> {
        self.redirect.as_ref()
            .map(|r| r as &str)
    }

    /// Check if the client has been closed.
    fn is_closed(&self) -> bool {
        self.redirect.is_some()
    }

    /// Insert a given Control Protocol message into the output message queue.
    fn send_control_message(&mut self, msg: ControlMessage) {
        self.messages.push_back(ArrowMessage::new(0, 0, msg))
    }

    /// Process a given Control Protocol message.
    fn process_control_protocol_message(&mut self, msg: ArrowMessage) -> Result<(), ArrowError> {
        let msg = ControlMessage::from_bytes(msg.payload())?
            .expect("unable to decode an Arrow Control Protocol message");

        let header = msg.header();

        match header.message_type() {
            ControlMessageType::ACK             => self.process_ack_message(msg),
            ControlMessageType::PING            => self.process_ping_message(msg),
            ControlMessageType::HUP             => self.process_hup_message(msg),
            ControlMessageType::REDIRECT        => self.process_redirect_message(msg),
            ControlMessageType::GET_STATUS      => self.process_get_status_message(msg),
            ControlMessageType::GET_SCAN_REPORT => self.process_get_scan_report_message(msg),
            ControlMessageType::RESET_SVC_TABLE => self.process_reset_svc_table_message(msg),
            ControlMessageType::SCAN_NETWORK    => self.process_scan_network_message(msg),
            ControlMessageType::UNKNOWN
                => Err(ArrowError::from("unknow control message received")),
            _   => Err(ArrowError::from("unexpected control message received")),
        }
    }

    /// Process a given ACK message.
    fn process_ack_message(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let header = msg.header();

        let expected_ack = self.expected_acks.pop_front();

        if self.expected_acks.is_empty() {
            self.ack_timeout.clear();
        } else {
            self.ack_timeout.set(ACK_TIMEOUT);
        }

        if let Some(expected_ack) = expected_ack {
            if header.msg_id == expected_ack {
                if self.state == ProtocolState::Handshake {
                    self.process_handshake_ack(msg)
                } else {
                    Ok(())
                }
            } else {
                Err(ArrowError::from("unexpected ACK message ID"))
            }
        } else {
            Err(ArrowError::from("no ACK message expected"))
        }
    }

    /// Process handshake ACK.
    fn process_handshake_ack(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let ack = msg.body::<AckMessage>()
            .expect("ACK message expected");

        if ack.err == ACK_NO_ERROR {
            // switch the protocol state into normal operation
            self.state = ProtocolState::Established;

            // TODO: spawn job for periodical sending of UPDATE messages
            // TODO: spawn job for periodical sending of PING messages
            // TODO: get the diagnostic mode state
            let diagnostic_mode = false;

            // report a fake redirect in case of the diagnostic mode
            if diagnostic_mode {
                self.redirect = Some(String::new());
            }

            Ok(())
        } else if ack.err == ACK_UNAUTHORIZED {
            Err(ArrowError::from("Arrow REGISTER failed (unauthorized)"))
        } else if ack.err == ACK_UNSUPPORTED_PROTOCOL_VERSION {
            Err(ArrowError::from("Arrow REGISTER failed (unsupported version of the Arrow Protocol)"))
        } else if ack.err == ACK_INTERNAL_SERVER_ERROR {
            Err(ArrowError::from("Arrow REGISTER failed (internal server error)"))
        } else {
            Err(ArrowError::from("Arrow REGISTER failed (unknown error)"))
        }
    }

    /// Process a given PING message.
    fn process_ping_message(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let header = msg.header();

        self.send_control_message(
            ControlMessage::ack(header.msg_id, 0));

        Ok(())
    }

    /// Process a given HUP message.
    fn process_hup_message(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let hup = msg.body::<HupMessage>()
            .expect("HUP message expected");

        self.sessions.close(
            hup.session_id,
            hup.error_code);

        Ok(())
    }

    /// Process a given REDIRECT message.
    fn process_redirect_message(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let msg = msg.body::<RedirectMessage>()
            .expect("REDIRECT message expected");

        self.redirect = Some(msg.target.clone());

        Ok(())
    }

    /// Process a given GET_STATUS message.
    fn process_get_status_message(&mut self, _: ControlMessage) -> Result<(), ArrowError> {
        // TODO
        Ok(())
    }

    /// Process a given GET_SCAN_REPORT message.
    fn process_get_scan_report_message(&mut self, _: ControlMessage) -> Result<(), ArrowError> {
        // TODO
        Ok(())
    }

    /// Process a given RESET_SVC_TABLE message.
    fn process_reset_svc_table_message(&mut self, _: ControlMessage) -> Result<(), ArrowError> {
        // TODO
        Ok(())
    }

    /// Process a given SCAN_NETWORK message.
    fn process_scan_network_message(&mut self, _: ControlMessage) -> Result<(), ArrowError> {
        // TODO
        Ok(())
    }

    /// Process a given service request message.
    fn process_service_request_message(&mut self, msg: ArrowMessage) -> Result<(), ArrowError> {
        if self.state != ProtocolState::Established {
            return Err(ArrowError::from("cannot handle service requests in the Handshake state"))
        }

        self.sessions.send(msg);

        Ok(())
    }
}

impl Sink for ArrowClient {
    type SinkItem  = ArrowMessage;
    type SinkError = ArrowError;

    fn start_send(&mut self, msg: ArrowMessage) -> StartSend<ArrowMessage, ArrowError> {
        // ignore the message if the client has been closed
        if self.is_closed() {
            return Ok(AsyncSink::Ready)
        }

        let header = msg.header();

        if header.service == 0 {
            self.process_control_protocol_message(msg)?;
        } else {
            self.process_service_request_message(msg)?;
        }

        Ok(AsyncSink::Ready)
    }

    fn poll_complete(&mut self) -> Poll<(), ArrowError> {
        Ok(Async::Ready(()))
    }
}

impl Stream for ArrowClient {
    type Item  = ArrowMessage;
    type Error = ArrowError;

    fn poll(&mut self) -> Poll<Option<ArrowMessage>, ArrowError> {
        if self.is_closed() {
            Ok(Async::Ready(None))
        } else if let Some(msg) = self.messages.pop_front() {
            Ok(Async::Ready(Some(msg)))
        } else {
            self.sessions.poll()
        }
    }
}

/// Control Protocol message factory with shared message ID counter.
#[derive(Clone)]
pub struct ControlMessageFactory {
    counter: Rc<Cell<u16>>,
}

impl ControlMessageFactory {
    /// Create a new Control Protocol message factory.
    pub fn new() -> ControlMessageFactory {
        ControlMessageFactory {
            counter: Rc::new(Cell::new(0)),
        }
    }

    /// Get next message ID and increment the counter.
    fn next_id(&mut self) -> u16 {
        let res = self.counter.get();

        self.counter.set(res.wrapping_add(1));

        res
    }

    /// Create a new ACK message with a given error code.
    pub fn ack(&mut self, error_code: u32) -> ArrowMessage {
        ArrowMessage::from(
            ControlMessage::ack(
                    self.next_id(),
                    error_code))
    }

    /// Create a new HUP message with a given session ID and error code.
    pub fn hup(&mut self, session_id: u32, error_code: u32) -> ArrowMessage {
        ArrowMessage::from(
            ControlMessage::hup(
                self.next_id(),
                session_id,
                error_code))
    }
}

/// Connect Arrow Client to a given address and return either a redirect address or an error.
pub fn connect<T>(addr: &str, svc_table: T) -> Result<String, ArrowError>
    where T: 'static + ServiceTable {
    let mut core = TokioCore::new()?;

    let addr = addr.to_socket_addrs()?
        .next()
        .ok_or(io::Error::new(io::ErrorKind::Other, "unable to resolve a given address"))?;

    let aclient = ArrowClient::new(core.handle(), svc_table);

    let client = TcpStream::connect(&addr, &core.handle())
        .map_err(|err| ArrowError::from(err))
        .and_then(|stream| {
            let framed = stream.framed(ArrowCodec);
            let (sink, stream) = framed.split();

            let messages = stream.pipe(aclient);

            sink.send_all(messages)
                .and_then(|(_, pipe)| {
                    let (_, _, context) = pipe.unpipe();
                    let redirect = context.get_redirect()
                        .expect("connection closed, redirect expected")
                        .to_string();

                    Ok(redirect)
                })
        });

    core.run(client)
}
<commit_msg>SCAN_NETWORK and RESET_SVC_TABLE message handling<commit_after>// Copyright 2017 click2stream, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

pub mod msg;
pub mod codec;
pub mod error;

mod session;

use std::io;

use std::rc::Rc;
use std::cell::Cell;
use std::net::ToSocketAddrs;
use std::collections::VecDeque;

use futures::{StartSend, Async, AsyncSink, Poll};
use futures::future::Future;
use futures::stream::Stream;
use futures::sink::Sink;

use tokio_core::net::TcpStream;
use tokio_core::reactor::Core as TokioCore;
use tokio_core::reactor::Handle as TokioCoreHandle;

use tokio_io::AsyncRead;

use futures_ex::StreamEx;

use net::arrow::proto::codec::{ArrowCodec, FromBytes};
use net::arrow::proto::error::ArrowError;
use net::arrow::proto::msg::ArrowMessage;
use net::arrow::proto::msg::control::{
    ACK_NO_ERROR,
    ACK_UNSUPPORTED_PROTOCOL_VERSION,
    ACK_UNAUTHORIZED,
    ACK_INTERNAL_SERVER_ERROR,

    AckMessage,
    ControlMessage,
    ControlMessageType,
    HupMessage,
    RedirectMessage
};
use net::arrow::proto::session::SessionManager;
use net::utils::Timeout;

pub use net::arrow::proto::msg::control::svc_table::{Service, ServiceTable};

use utils::logger::Logger;

/// Currently supported version of the Arrow protocol.
pub const ARROW_PROTOCOL_VERSION: u8 = 1;

const ACK_TIMEOUT: u64 = 20000;

/// Commands that might be sent by the Arrow Client into a given command queue.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Command {
    ResetServiceTable,
    ScanNetwork,
}

/// Common trait for various implementations of command senders.
pub trait Sender {
    /// Send a given command or return the command back if the send operation
    /// failed.
    fn send(&self, cmd: Command) -> Result<(), Command>;
}

/// Arrow Protocol states.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
enum ProtocolState {
    Handshake,
    Established
}

/// Arrow Client implementation.
struct ArrowClient<L, S> {
    logger:        L,
    cmd_sender:    S,
    tc_handle:     TokioCoreHandle,
    cmsg_factory:  ControlMessageFactory,
    sessions:      SessionManager,
    messages:      VecDeque<ArrowMessage>,
    expected_acks: VecDeque<u16>,
    ack_timeout:   Timeout,
    state:         ProtocolState,
    redirect:      Option<String>,
}

impl<L, S> ArrowClient<L, S>
    where L: Logger,
          S: Sender {
    /// Create a new Arrow Client.
    fn new<T>(
        logger: L,
        cmd_sender: S,
        svc_table: T,
        tc_handle: TokioCoreHandle) -> ArrowClient<L, S>
        where T: 'static + ServiceTable {
        let cmsg_factory = ControlMessageFactory::new();
        let session_manager = SessionManager::new(
            tc_handle.clone(),
            svc_table,
            cmsg_factory.clone());

        let messages = VecDeque::new();

        // TODO: add REGISTER message into the message queue

        ArrowClient {
            logger:        logger,
            cmd_sender:    cmd_sender,
            tc_handle:     tc_handle,
            cmsg_factory:  cmsg_factory,
            sessions:      session_manager,
            messages:      messages,
            expected_acks: VecDeque::new(),
            ack_timeout:   Timeout::new(),
            state:         ProtocolState::Handshake,
            redirect:      None,
        }
    }

    /// Get redirect address (if any).
    fn get_redirect(&self) -> Option<&str> {
        self.redirect.as_ref()
            .map(|r| r as &str)
    }

    /// Check if the client has been closed.
    fn is_closed(&self) -> bool {
        self.redirect.is_some()
    }

    /// Insert a given Control Protocol message into the output message queue.
    fn send_control_message(&mut self, msg: ControlMessage) {
        self.messages.push_back(ArrowMessage::new(0, 0, msg))
    }

    /// Process a given Control Protocol message.
    fn process_control_protocol_message(&mut self, msg: ArrowMessage) -> Result<(), ArrowError> {
        let msg = ControlMessage::from_bytes(msg.payload())?
            .expect("unable to decode an Arrow Control Protocol message");

        let header = msg.header();

        match header.message_type() {
            ControlMessageType::ACK             => self.process_ack_message(msg),
            ControlMessageType::PING            => self.process_ping_message(msg),
            ControlMessageType::HUP             => self.process_hup_message(msg),
            ControlMessageType::REDIRECT        => self.process_redirect_message(msg),
            ControlMessageType::GET_STATUS      => self.process_get_status_message(msg),
            ControlMessageType::GET_SCAN_REPORT => self.process_get_scan_report_message(msg),
            ControlMessageType::RESET_SVC_TABLE => self.process_command(Command::ResetServiceTable),
            ControlMessageType::SCAN_NETWORK    => self.process_command(Command::ScanNetwork),
            ControlMessageType::UNKNOWN
                => Err(ArrowError::from("unknow control message received")),
            _   => Err(ArrowError::from("unexpected control message received")),
        }
    }

    /// Process a given ACK message.
    fn process_ack_message(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let header = msg.header();

        let expected_ack = self.expected_acks.pop_front();

        if self.expected_acks.is_empty() {
            self.ack_timeout.clear();
        } else {
            self.ack_timeout.set(ACK_TIMEOUT);
        }

        if let Some(expected_ack) = expected_ack {
            if header.msg_id == expected_ack {
                if self.state == ProtocolState::Handshake {
                    self.process_handshake_ack(msg)
                } else {
                    Ok(())
                }
            } else {
                Err(ArrowError::from("unexpected ACK message ID"))
            }
        } else {
            Err(ArrowError::from("no ACK message expected"))
        }
    }

    /// Process handshake ACK.
    fn process_handshake_ack(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let ack = msg.body::<AckMessage>()
            .expect("ACK message expected");

        if ack.err == ACK_NO_ERROR {
            // switch the protocol state into normal operation
            self.state = ProtocolState::Established;

            // TODO: spawn job for periodical sending of UPDATE messages
            // TODO: spawn job for periodical sending of PING messages
            // TODO: get the diagnostic mode state
            let diagnostic_mode = false;

            // report a fake redirect in case of the diagnostic mode
            if diagnostic_mode {
                self.redirect = Some(String::new());
            }

            Ok(())
        } else if ack.err == ACK_UNAUTHORIZED {
            Err(ArrowError::from("Arrow REGISTER failed (unauthorized)"))
        } else if ack.err == ACK_UNSUPPORTED_PROTOCOL_VERSION {
            Err(ArrowError::from("Arrow REGISTER failed (unsupported version of the Arrow Protocol)"))
        } else if ack.err == ACK_INTERNAL_SERVER_ERROR {
            Err(ArrowError::from("Arrow REGISTER failed (internal server error)"))
        } else {
            Err(ArrowError::from("Arrow REGISTER failed (unknown error)"))
        }
    }

    /// Process a given PING message.
    fn process_ping_message(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let header = msg.header();

        self.send_control_message(
            ControlMessage::ack(header.msg_id, 0));

        Ok(())
    }

    /// Process a given HUP message.
    fn process_hup_message(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let hup = msg.body::<HupMessage>()
            .expect("HUP message expected");

        self.sessions.close(
            hup.session_id,
            hup.error_code);

        Ok(())
    }

    /// Process a given REDIRECT message.
    fn process_redirect_message(&mut self, msg: ControlMessage) -> Result<(), ArrowError> {
        let msg = msg.body::<RedirectMessage>()
            .expect("REDIRECT message expected");

        self.redirect = Some(msg.target.clone());

        Ok(())
    }

    /// Process a given GET_STATUS message.
    fn process_get_status_message(&mut self, _: ControlMessage) -> Result<(), ArrowError> {
        // TODO
        Ok(())
    }

    /// Process a given GET_SCAN_REPORT message.
    fn process_get_scan_report_message(&mut self, _: ControlMessage) -> Result<(), ArrowError> {
        // TODO
        Ok(())
    }

    /// Send a given command using the underlaying command channel.
    fn process_command(&mut self, cmd: Command) -> Result<(), ArrowError> {
        if let Err(cmd) = self.cmd_sender.send(cmd) {
            log_warn!(self.logger, "unable to process command {:?}", cmd);
        }

        Ok(())
    }

    /// Process a given service request message.
    fn process_service_request_message(&mut self, msg: ArrowMessage) -> Result<(), ArrowError> {
        if self.state != ProtocolState::Established {
            return Err(ArrowError::from("cannot handle service requests in the Handshake state"))
        }

        self.sessions.send(msg);

        Ok(())
    }
}

impl<L, S> Sink for ArrowClient<L, S>
    where L: Logger,
          S: Sender {
    type SinkItem  = ArrowMessage;
    type SinkError = ArrowError;

    fn start_send(&mut self, msg: ArrowMessage) -> StartSend<ArrowMessage, ArrowError> {
        // ignore the message if the client has been closed
        if self.is_closed() {
            return Ok(AsyncSink::Ready)
        }

        let header = msg.header();

        if header.service == 0 {
            self.process_control_protocol_message(msg)?;
        } else {
            self.process_service_request_message(msg)?;
        }

        Ok(AsyncSink::Ready)
    }

    fn poll_complete(&mut self) -> Poll<(), ArrowError> {
        Ok(Async::Ready(()))
    }
}

impl<L, S> Stream for ArrowClient<L, S>
    where L: Logger,
          S: Sender {
    type Item  = ArrowMessage;
    type Error = ArrowError;

    fn poll(&mut self) -> Poll<Option<ArrowMessage>, ArrowError> {
        if self.is_closed() {
            Ok(Async::Ready(None))
        } else if let Some(msg) = self.messages.pop_front() {
            Ok(Async::Ready(Some(msg)))
        } else {
            self.sessions.poll()
        }
    }
}

/// Control Protocol message factory with shared message ID counter.
#[derive(Clone)]
pub struct ControlMessageFactory {
    counter: Rc<Cell<u16>>,
}

impl ControlMessageFactory {
    /// Create a new Control Protocol message factory.
    pub fn new() -> ControlMessageFactory {
        ControlMessageFactory {
            counter: Rc::new(Cell::new(0)),
        }
    }

    /// Get next message ID and increment the counter.
    fn next_id(&mut self) -> u16 {
        let res = self.counter.get();

        self.counter.set(res.wrapping_add(1));

        res
    }

    /// Create a new ACK message with a given error code.
    pub fn ack(&mut self, error_code: u32) -> ArrowMessage {
        ArrowMessage::from(
            ControlMessage::ack(
                    self.next_id(),
                    error_code))
    }

    /// Create a new HUP message with a given session ID and error code.
    pub fn hup(&mut self, session_id: u32, error_code: u32) -> ArrowMessage {
        ArrowMessage::from(
            ControlMessage::hup(
                self.next_id(),
                session_id,
                error_code))
    }
}

/// Connect Arrow Client to a given address and return either a redirect address or an error.
pub fn connect<L, S, T>(
    logger: L,
    cmd_sender: S,
    svc_table: T,
    addr: &str) -> Result<String, ArrowError>
    where L: Logger,
          S: Sender,
          T: 'static + ServiceTable {
    let mut core = TokioCore::new()?;

    let addr = addr.to_socket_addrs()?
        .next()
        .ok_or(io::Error::new(io::ErrorKind::Other, "unable to resolve a given address"))?;

    let aclient = ArrowClient::new(
        logger,
        cmd_sender,
        svc_table,
        core.handle());

    let client = TcpStream::connect(&addr, &core.handle())
        .map_err(|err| ArrowError::from(err))
        .and_then(|stream| {
            let framed = stream.framed(ArrowCodec);
            let (sink, stream) = framed.split();

            let messages = stream.pipe(aclient);

            sink.send_all(messages)
                .and_then(|(_, pipe)| {
                    let (_, _, context) = pipe.unpipe();
                    let redirect = context.get_redirect()
                        .expect("connection closed, redirect expected")
                        .to_string();

                    Ok(redirect)
                })
        });

    core.run(client)
}
<|endoftext|>


<commit_before>#[crate_id = "bignum#0.1.0-pre"];

#[comment = "Bignum library for Rust"];
#[crate_type = "rlib"];

#[feature(macro_rules)];

extern crate gmp;

use gmp::{Mpz, RandState};
use std::fmt;
use std::from_str::FromStr;
use std::num::{One, Zero, ToStrRadix};
use std::rand::Rng;
use std::libc::c_ulong;

#[deriving(Clone, Eq, Ord, TotalEq, TotalOrd, Zero)]
pub struct BigUint {
    data: Mpz
}

impl One for BigUint {
    fn one() -> BigUint {
        BigUint{ data: One::one() }
    }
}

impl FromPrimitive for BigUint {
    fn from_u64(other: u64) -> Option<BigUint> {
        let mpz: Option<Mpz> = FromPrimitive::from_u64(other);
        match mpz {
            Some(mpz) => Some(BigUint{ data: mpz }),
            None      => None
        }
    }

    fn from_i64(other: i64) -> Option<BigUint> {
        if other < 0 {
            None
        } else {
            let mpz: Option<Mpz> = FromPrimitive::from_i64(other);
            match mpz {
                Some(mpz) => Some(BigUint{ data: mpz }),
                None      => None
            }
        }
    }
}

pub trait ToBigUint {
    fn to_biguint(&self) -> Option<BigUint>;
}

macro_rules! impl_to_biguint(
    ($T:ty, $from_ty:path) => {
        impl ToBigUint for $T {
            fn to_biguint(&self) -> Option<BigUint> {
                $from_ty(*self)
            }
        }
    }
)

impl_to_biguint!(int,  FromPrimitive::from_int)
impl_to_biguint!(i8,   FromPrimitive::from_i8)
impl_to_biguint!(i16,  FromPrimitive::from_i16)
impl_to_biguint!(i32,  FromPrimitive::from_i32)
impl_to_biguint!(i64,  FromPrimitive::from_i64)
impl_to_biguint!(uint, FromPrimitive::from_uint)
impl_to_biguint!(u8,   FromPrimitive::from_u8)
impl_to_biguint!(u16,  FromPrimitive::from_u16)
impl_to_biguint!(u32,  FromPrimitive::from_u32)
impl_to_biguint!(u64,  FromPrimitive::from_u64)

impl FromStr for BigUint {
    fn from_str(s: &str) -> Option<BigUint> {
        let data: Option<Mpz> = FromStr::from_str(s);
        match data {
            Some(data) => Some(BigUint{ data: data }),
            None       => None
        }
    }
}

impl ToStrRadix for BigUint {
    fn to_str_radix(&self, radix: uint) -> ~str {
        self.data.to_str_radix(radix)
    }
}

impl fmt::Show for BigUint {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f.buf, "{}", self.to_str_radix(10))
    }
}

impl Add<BigUint, BigUint> for BigUint {
    fn add(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.add(&other.data)}
    }
}

impl Sub<BigUint, BigUint> for BigUint {
    fn sub(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.sub(&other.data)}
    }
}

impl Mul<BigUint, BigUint> for BigUint {
    fn mul(&self, other: &BigUint) -> BigUint {
        BigUint{data: self.data.mul(&other.data)}
    }
}

impl BigUint {
    pub fn is_odd(&self) -> bool {
        self.data.is_odd()
    }

    pub fn is_even(&self) -> bool {
        self.data.is_even()
    }
}

pub trait RandBigInt {
    /// Generate a random `BigUint` of the given bit size.
    fn gen_biguint(&mut self, bit_size: uint) -> BigUint;

    // /// Generate a random BigInt of the given bit size.
    // fn gen_bigint(&mut self, bit_size: uint) -> BigInt;

    /// Generate a random `BigUint` less than the given bound. Fails
    /// when the bound is zero.
    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;

    /// Generate a random `BigUint` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;

    // /// Generate a random `BigInt` within the given range. The lower
    // /// bound is inclusive; the upper bound is exclusive. Fails when
    // /// the upper bound is not greater than the lower bound.
    // fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
}

impl<R: Rng> RandBigInt for R {
    fn gen_biguint(&mut self, bit_size: uint) -> BigUint {
        let mut state = RandState::new();
        let seed: c_ulong = self.gen();
        state.seed_ui(seed);
        let mpz = state.urandom_2exp(bit_size as c_ulong);
        BigUint{ data: mpz }
    }

    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
        // FIXME Add assertion once is_zero is implemented
        // assert!(!bound.is_zero());

        let mut state = RandState::new();
        let seed: c_ulong = self.gen();
        state.seed_ui(seed);
        let mpz = state.urandom(&bound.data);
        BigUint{ data: mpz }
    }

    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
        // FIXME Add assertion once cmp::Ord is implemented
        // assert!(lbound < ubound);

        return *lbound + self.gen_biguint_below(&(*ubound - *lbound));
    }
}

#[cfg(test)]
mod test {
    use super::{BigUint, RandBigInt, ToBigUint};
    use std::{u32,u64};
    use std::from_str::FromStr;
    use std::num::{Zero, One};
    use std::rand::task_rng;

    #[test]
    fn test_clone() {
        let two = 2u.to_biguint().unwrap();
        let also_two = two.clone();
        assert_eq!(two.to_str(), also_two.to_str());
    }

    #[test]
    fn test_from_primitive() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        assert_eq!(two.to_str(), ~"2");
    }

    #[test]
    fn test_from_str() {
        let two: BigUint = FromStr::from_str("2").unwrap();
        assert_eq!(two.to_str(), ~"2");
    }

    #[test]
    fn test_to_biguint() {
        let three = 3u.to_biguint().unwrap();
        assert_eq!(three.to_str(), ~"3");
    }

    #[test]
    fn test_comparisons() {
        let two = 2u.to_biguint().unwrap();
        let also_two = 2u.to_biguint().unwrap();
        let three = 3u.to_biguint().unwrap();
        assert!(two == also_two);
        assert!(two >= also_two);
        assert!(two <= also_two);
        assert!(three > two);
        assert!(three >= two);
        assert!(two < three);
        assert!(two <= three);
    }

    #[test]
    fn test_zero_and_one() {
        let zero: BigUint = Zero::zero();
        let one: BigUint = One::one();
        assert_eq!(zero.to_str(), ~"0");
        assert_eq!(one.to_str(), ~"1");
    }

    #[test]
    fn test_printing() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        println!("{} == 2!", two); // Doesn't fail
    }

    #[test]
    fn test_add() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two.add(&three).to_str(), ~"5");
        assert_eq!((two + three).to_str(), ~"5");
    }

    #[test]
    fn test_simple_sub() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(three.sub(&two).to_str(), ~"1");
        assert_eq!((three - two).to_str(), ~"1");
    }

    #[test]
    fn test_mul() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two.mul(&three).to_str(), ~"6");
        assert_eq!((two * three).to_str(), ~"6");
    }

    #[test]
    fn test_is_odd() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert!(!two.is_odd());
        assert!(three.is_odd());
    }

    #[test]
    fn test_is_even() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert!(two.is_even());
        assert!(!three.is_even());
    }

    #[test]
    fn test_rand_gen_biguint() {
        let mut rng = task_rng();
        let rand1 = rng.gen_biguint(128);
        let rand2 = rng.gen_biguint(128);
        assert!(rand1.to_str() != rand2.to_str());
    }

    #[test]
    fn test_rand_gen_biguint_below() {
        let mut rng = task_rng();
        let max: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        let rand1 = rng.gen_biguint_below(&max);
        let rand2 = rng.gen_biguint_below(&max);
        assert!(rand1.to_str() != rand2.to_str());
    }

    #[test]
    fn test_rand_gen_biguint_range() {
        let mut rng = task_rng();
        let min: BigUint = FromPrimitive::from_u32(u32::MAX).unwrap();
        let max: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        let rand1 = rng.gen_biguint_range(&min, &max);
        let rand2 = rng.gen_biguint_range(&min, &max);
        assert!(rand1.to_str() != rand2.to_str());
    }
}
<commit_msg>Implement a few more operations<commit_after>#[crate_id = "bignum#0.1.0-pre"];

#[comment = "Bignum library for Rust"];
#[crate_type = "rlib"];

#[feature(macro_rules)];

extern crate gmp;

use gmp::{Mpz, RandState};
use std::fmt;
use std::from_str::FromStr;
use std::num::{One, Zero, ToStrRadix};
use std::rand::Rng;
use std::libc::c_ulong;

#[deriving(Clone, Eq, Ord, TotalEq, TotalOrd, Zero)]
pub struct BigUint {
    data: Mpz
}

impl BigUint {
    pub fn is_odd(&self) -> bool {
        self.data.is_odd()
    }

    pub fn is_even(&self) -> bool {
        self.data.is_even()
    }

    pub fn divides(&self, other: &BigUint) -> bool {
        self.data.divides(&other.data)
    }

    pub fn bits(&self) -> uint {
        self.data.bit_length()
    }
}

impl One for BigUint {
    fn one() -> BigUint {
        BigUint{ data: One::one() }
    }
}

impl FromPrimitive for BigUint {
    fn from_u64(other: u64) -> Option<BigUint> {
        let mpz: Option<Mpz> = FromPrimitive::from_u64(other);
        match mpz {
            Some(mpz) => Some(BigUint{ data: mpz }),
            None      => None
        }
    }

    fn from_i64(other: i64) -> Option<BigUint> {
        if other < 0 {
            None
        } else {
            let mpz: Option<Mpz> = FromPrimitive::from_i64(other);
            match mpz {
                Some(mpz) => Some(BigUint{ data: mpz }),
                None      => None
            }
        }
    }
}

pub trait ToBigUint {
    fn to_biguint(&self) -> Option<BigUint>;
}

macro_rules! impl_to_biguint(
    ($T:ty, $from_ty:path) => {
        impl ToBigUint for $T {
            fn to_biguint(&self) -> Option<BigUint> {
                $from_ty(*self)
            }
        }
    }
)

impl_to_biguint!(int,  FromPrimitive::from_int)
impl_to_biguint!(i8,   FromPrimitive::from_i8)
impl_to_biguint!(i16,  FromPrimitive::from_i16)
impl_to_biguint!(i32,  FromPrimitive::from_i32)
impl_to_biguint!(i64,  FromPrimitive::from_i64)
impl_to_biguint!(uint, FromPrimitive::from_uint)
impl_to_biguint!(u8,   FromPrimitive::from_u8)
impl_to_biguint!(u16,  FromPrimitive::from_u16)
impl_to_biguint!(u32,  FromPrimitive::from_u32)
impl_to_biguint!(u64,  FromPrimitive::from_u64)

impl FromStr for BigUint {
    fn from_str(s: &str) -> Option<BigUint> {
        let data: Option<Mpz> = FromStr::from_str(s);
        match data {
            Some(data) => Some(BigUint{ data: data }),
            None       => None
        }
    }
}

impl ToStrRadix for BigUint {
    fn to_str_radix(&self, radix: uint) -> ~str {
        self.data.to_str_radix(radix)
    }
}

impl fmt::Show for BigUint {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f.buf, "{}", self.to_str_radix(10))
    }
}

impl Add<BigUint, BigUint> for BigUint {
    fn add(&self, other: &BigUint) -> BigUint {
        BigUint{ data: self.data.add(&other.data) }
    }
}

impl Sub<BigUint, BigUint> for BigUint {
    fn sub(&self, other: &BigUint) -> BigUint {
        BigUint{ data: self.data.sub(&other.data) }
    }
}

impl Mul<BigUint, BigUint> for BigUint {
    fn mul(&self, other: &BigUint) -> BigUint {
        BigUint{ data: self.data.mul(&other.data) }
    }
}

impl Div<BigUint, BigUint> for BigUint {
    fn div(&self, other: &BigUint) -> BigUint {
        BigUint{ data: self.data.div(&other.data) }
    }
}

impl Rem<BigUint, BigUint> for BigUint {
    fn rem(&self, other: &BigUint) -> BigUint {
        BigUint{ data: self.data.rem(&other.data) }
    }
}

impl BitAnd<BigUint, BigUint> for BigUint {
    fn bitand(&self, other: &BigUint) -> BigUint {
        BigUint{ data: self.data.bitand(&other.data) }
    }
}

impl Shr<uint, BigUint> for BigUint {
    fn shr(&self, rhs: &uint) -> BigUint {
        let shift = *rhs as c_ulong;
        BigUint{ data: self.data.shr(&shift) }
    }
}

pub trait RandBigInt {
    /// Generate a random `BigUint` of the given bit size.
    fn gen_biguint(&mut self, bit_size: uint) -> BigUint;

    // /// Generate a random BigInt of the given bit size.
    // fn gen_bigint(&mut self, bit_size: uint) -> BigInt;

    /// Generate a random `BigUint` less than the given bound. Fails
    /// when the bound is zero.
    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint;

    /// Generate a random `BigUint` within the given range. The lower
    /// bound is inclusive; the upper bound is exclusive. Fails when
    /// the upper bound is not greater than the lower bound.
    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint;

    // /// Generate a random `BigInt` within the given range. The lower
    // /// bound is inclusive; the upper bound is exclusive. Fails when
    // /// the upper bound is not greater than the lower bound.
    // fn gen_bigint_range(&mut self, lbound: &BigInt, ubound: &BigInt) -> BigInt;
}

impl<R: Rng> RandBigInt for R {
    fn gen_biguint(&mut self, bit_size: uint) -> BigUint {
        let mut state = RandState::new();
        let seed: c_ulong = self.gen();
        state.seed_ui(seed);
        let mpz = state.urandom_2exp(bit_size as c_ulong);
        BigUint{ data: mpz }
    }

    fn gen_biguint_below(&mut self, bound: &BigUint) -> BigUint {
        // FIXME Add assertion once is_zero is implemented
        // assert!(!bound.is_zero());

        let mut state = RandState::new();
        let seed: c_ulong = self.gen();
        state.seed_ui(seed);
        let mpz = state.urandom(&bound.data);
        BigUint{ data: mpz }
    }

    fn gen_biguint_range(&mut self, lbound: &BigUint, ubound: &BigUint) -> BigUint {
        // FIXME Add assertion once cmp::Ord is implemented
        // assert!(lbound < ubound);

        return *lbound + self.gen_biguint_below(&(*ubound - *lbound));
    }
}

#[cfg(test)]
mod test {
    use super::{BigUint, RandBigInt, ToBigUint};
    use std::{u32,u64};
    use std::from_str::FromStr;
    use std::num::{Zero, One};
    use std::rand::task_rng;

    #[test]
    fn test_clone() {
        let two = 2u.to_biguint().unwrap();
        let also_two = two.clone();
        assert_eq!(two.to_str(), also_two.to_str());
    }

    #[test]
    fn test_from_primitive() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        assert_eq!(two.to_str(), ~"2");
    }

    #[test]
    fn test_from_str() {
        let two: BigUint = FromStr::from_str("2").unwrap();
        assert_eq!(two.to_str(), ~"2");
    }

    #[test]
    fn test_to_biguint() {
        let three = 3u.to_biguint().unwrap();
        assert_eq!(three.to_str(), ~"3");
    }

    #[test]
    fn test_comparisons() {
        let two = 2u.to_biguint().unwrap();
        let also_two = 2u.to_biguint().unwrap();
        let three = 3u.to_biguint().unwrap();
        assert!(two == also_two);
        assert!(two >= also_two);
        assert!(two <= also_two);
        assert!(three > two);
        assert!(three >= two);
        assert!(two < three);
        assert!(two <= three);
    }

    #[test]
    fn test_zero_and_one() {
        let zero: BigUint = Zero::zero();
        let one: BigUint = One::one();
        assert_eq!(zero.to_str(), ~"0");
        assert_eq!(one.to_str(), ~"1");
    }

    #[test]
    fn test_printing() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        println!("{} == 2!", two); // Doesn't fail
    }

    #[test]
    fn test_bits() {
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();
        assert_eq!(three.bits(), 2);

        let max64: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        assert_eq!(max64.bits(), 64);
        assert_eq!((max64 + three).bits(), 65);
    }

    #[test]
    fn test_add() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two.add(&three).to_str(), ~"5");
        assert_eq!((two + three).to_str(), ~"5");
    }

    #[test]
    fn test_simple_sub() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(three.sub(&two).to_str(), ~"1");
        assert_eq!((three - two).to_str(), ~"1");
    }

    #[test]
    fn test_mul() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two.mul(&three).to_str(), ~"6");
        assert_eq!((two * three).to_str(), ~"6");
    }

    #[test]
    fn test_div() {
        let one: BigUint = FromPrimitive::from_uint(1).unwrap();
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();

        assert_eq!(two / two, one);
        assert_eq!(two.div(&two), one);
    }

    #[test]
    fn test_rem() {
        let one: BigUint = FromPrimitive::from_uint(1).unwrap();
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(three % two, one);
        assert_eq!(three.rem(&two), one);
    }

    #[test]
    fn test_bitand() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert_eq!(two & three, two);
        assert_eq!(two.bitand(&three), two);
    }

    #[test]
    fn test_shr() {
        let one: BigUint = FromPrimitive::from_uint(1).unwrap();
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();

        assert_eq!(two >> 1, one);
        assert_eq!(two.shr(&1), one);
    }

    #[test]
    fn test_is_odd() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert!(!two.is_odd());
        assert!(three.is_odd());
    }

    #[test]
    fn test_is_even() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();

        assert!(two.is_even());
        assert!(!three.is_even());
    }

    #[test]
    fn test_divides() {
        let two: BigUint = FromPrimitive::from_uint(2).unwrap();
        let three: BigUint = FromPrimitive::from_uint(3).unwrap();
        let six: BigUint = FromPrimitive::from_uint(6).unwrap();

        assert!(two.divides(&six));
        assert!(three.divides(&six));
        assert!(!two.divides(&three));
    }

    #[test]
    fn test_rand_gen_biguint() {
        let mut rng = task_rng();
        let rand1 = rng.gen_biguint(128);
        let rand2 = rng.gen_biguint(128);
        assert!(rand1.to_str() != rand2.to_str());
    }

    #[test]
    fn test_rand_gen_biguint_below() {
        let mut rng = task_rng();
        let max: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        let rand1 = rng.gen_biguint_below(&max);
        let rand2 = rng.gen_biguint_below(&max);
        assert!(rand1.to_str() != rand2.to_str());
    }

    #[test]
    fn test_rand_gen_biguint_range() {
        let mut rng = task_rng();
        let min: BigUint = FromPrimitive::from_u32(u32::MAX).unwrap();
        let max: BigUint = FromPrimitive::from_u64(u64::MAX).unwrap();
        let rand1 = rng.gen_biguint_range(&min, &max);
        let rand2 = rng.gen_biguint_range(&min, &max);
        assert!(rand1.to_str() != rand2.to_str());
    }
}
<|endoftext|>