<commit_before>//! Collision detector between two `~Geom`.

use std::unstable::intrinsics::TypeId;
use std::hashmap::HashMap;
use nalgebra::na;
use geom::{AnnotatedPoint, Geom, ConcaveGeom, Cone, Box, Ball, Capsule, Convex, Cylinder, Compound,
           Mesh, Triangle, Segment, Plane};
use implicit::{Implicit, PreferedSamplingDirections};
use contact::Contact;
use narrow::algorithm::simplex::Simplex;
use narrow::algorithm::johnson_simplex::{JohnsonSimplex, RecursionTemplate};
use narrow::{CollisionDetector, ImplicitImplicit, BallBall,
             ImplicitPlane, PlaneImplicit, ConcaveGeomGeomFactory, GeomConcaveGeomFactory};
use OSCMG = narrow::OneShotContactManifoldGenerator;
use math::{N, V, LV, M};

/// Same as the `CollisionDetector` trait but using dynamic dispatch on the geometries.
pub trait GeomGeomCollisionDetector {
    /// Runs the collision detection on two objects. It is assumed that the same
    /// collision detector (the same structure) is always used with the same
    /// pair of object.
    fn update(&mut self,
              &GeomGeomDispatcher,
              &M,
              &Geom,
              &M,
              &Geom);

    /// The number of collision detected during the last update.
    fn num_colls(&self) -> uint;

    /// Collects the collisions detected during the last update.
    fn colls(&self, &mut ~[Contact]);
}

/// Trait to be implemented by collision detector using dynamic dispatch.
///
/// This is used to know the exact type of the geometries.
pub trait DynamicCollisionDetector<G1, G2>: GeomGeomCollisionDetector { }

#[deriving(Clone)]
struct DetectorWithoutRedispatch<D> {
    detector: D
}

impl<D> DetectorWithoutRedispatch<D> {
    pub fn new(d: D) -> DetectorWithoutRedispatch<D> {
        DetectorWithoutRedispatch {
            detector: d
        }
    }
}

impl<D: CollisionDetector<G1, G2>, G1, G2>
CollisionDetector<G1, G2> for DetectorWithoutRedispatch<D> {
    fn update(&mut self, _: &M, _: &G1, _: &M, _: &G2) { unreachable!() }
    fn num_colls(&self) -> uint { unreachable!() }
    fn colls(&self, _: &mut ~[Contact]) { unreachable!() }
    fn toi(_: Option<DetectorWithoutRedispatch<D>>, _: &M, _: &V, _: &N, _: &G1, _: &M, _: &G2) -> Option<N> {
        unreachable!()
    }
}

impl<D: CollisionDetector<G1, G2>, G1, G2>
DynamicCollisionDetector<G1, G2> for DetectorWithoutRedispatch<D> { }

impl<D: CollisionDetector<G1, G2>, G1: 'static, G2: 'static>
GeomGeomCollisionDetector for DetectorWithoutRedispatch<D> {
    #[inline]
    fn update(&mut self,
              _:  &GeomGeomDispatcher,
              m1: &M,
              g1: &Geom,
              m2: &M,
              g2: &Geom) {
        self.detector.update(
            m1,
            g1.as_ref::<G1>().expect("Invalid geometry."),
            m2,
            g2.as_ref::<G2>().expect("Invalid geometry."))
    }

    #[inline]
    fn num_colls(&self) -> uint {
        self.detector.num_colls()
    }

    #[inline]
    fn colls(&self, cs: &mut ~[Contact]) {
        self.detector.colls(cs)
    }
}

/// Collision dispatcher between two `~Geom`.
pub struct GeomGeomDispatcher {
    priv constructors: HashMap<(TypeId, TypeId), ~CollisionDetectorFactory>
}

impl GeomGeomDispatcher {
    /// Creates a new `GeomGeomDispatcher` without the default set of collision detectors
    /// factories.
    pub fn new_without_default() -> GeomGeomDispatcher {
        GeomGeomDispatcher {
            constructors: HashMap::new()
        }
    }

    /// Registers a new dynamic collision detector for two geometries.
    pub fn register_dynamic_detector<G1: 'static + Any,
                                     G2: 'static + Any,
                                     D:  'static + Send + Clone +
                                         DynamicCollisionDetector<G1, G2>>(
                                     &mut self,
                                     d:   D) {
        let key     = (TypeId::of::<G1>(), TypeId::of::<G2>());
        let factory = ~CollisionDetectorCloner::new(d);
        self.constructors.insert(key, factory as ~CollisionDetectorFactory);
    }

    /// Registers a new collision detector for two geometries.
    pub fn register_detector<G1: 'static + Any,
                             G2: 'static + Any,
                             D:  'static + Send + CollisionDetector<G1, G2> + Clone>(
                             &mut self,
                             d:   D) {
        self.register_dynamic_detector(DetectorWithoutRedispatch::new(d));
    }

    /// Unregister the collision detector for a givem pair of geometries.
    pub fn unregister_detector<G1: 'static + Any, G2: 'static + Any>(&mut self) {
        let key = (TypeId::of::<G1>(), TypeId::of::<G2>());
        self.constructors.remove(&key);
    }

    /// Creates a new collision detector adapted for the two given geometries.
    pub fn dispatch(&self, a: &Geom, b: &Geom) -> ~GeomGeomCollisionDetector {
        match self.constructors.find(&(a.get_type_id(), b.get_type_id())) {
            Some(f) => f.build(),
            None    => fail!("Unable to find a collision detector.")
        }
    }
}

impl GeomGeomDispatcher {
    // FIXME: make this a function which has the simplex and the prediction margin as parameters
    /// Creates a new `GeomGeomDispatcher` able do build collision detectors for any valid pair of
    /// geometries supported by `ncollide`.
    pub fn new() -> GeomGeomDispatcher {
        let mut res = GeomGeomDispatcher::new_without_default();

        type Simplex  = JohnsonSimplex<AnnotatedPoint>;
        type Self     = GeomGeomDispatcher;
        type Super    = ~GeomGeomCollisionDetector;

        /*
         * Involving a Plane
         */
        let prediction: &N = &na::cast(0.1);

        // Ball vs. Ball
        let bb: BallBall = BallBall::new(prediction.clone());
        res.register_detector(bb);

        // Plane vs. Implicit
        res.register_default_plane_implicit_detector::<Ball>(false, prediction);
        res.register_default_plane_implicit_detector::<Box>(true, prediction);
        res.register_default_plane_implicit_detector::<Cone>(true, prediction);
        res.register_default_plane_implicit_detector::<Cylinder>(true, prediction);
        res.register_default_plane_implicit_detector::<Capsule>(true, prediction);
        res.register_default_plane_implicit_detector::<Convex>(true, prediction);
        res.register_default_plane_implicit_detector::<Triangle>(true, prediction);
        res.register_default_plane_implicit_detector::<Segment>(true, prediction);

        // Implicit vs. Implicit
        // NOTE: some pair will be registered twice…
        res.register_default_implicit_detectors::<Box>(true, prediction);
        res.register_default_implicit_detectors::<Cone>(true, prediction);
        res.register_default_implicit_detectors::<Cylinder>(true, prediction);
        res.register_default_implicit_detectors::<Capsule>(true, prediction);
        res.register_default_implicit_detectors::<Convex>(true, prediction);
        res.register_default_implicit_detectors::<Triangle>(true, prediction);
        res.register_default_implicit_detectors::<Segment>(true, prediction);

        // FIXME: refactor the three following blocks?
        // Compound vs. Other
        res.register_default_concave_geom_geom_detector::<Compound, Plane>();
        res.register_default_concave_geom_geom_detector::<Compound, Ball>();
        res.register_default_concave_geom_geom_detector::<Compound, Box>();
        res.register_default_concave_geom_geom_detector::<Compound, Cone>();
        res.register_default_concave_geom_geom_detector::<Compound, Cylinder>();
        res.register_default_concave_geom_geom_detector::<Compound, Capsule>();
        res.register_default_concave_geom_geom_detector::<Compound, Convex>();
        res.register_default_concave_geom_geom_detector::<Compound, Triangle>();
        res.register_default_concave_geom_geom_detector::<Compound, Segment>();

        // TriangleMesh vs. Other
        res.register_default_concave_geom_geom_detector::<Mesh, Plane>();
        res.register_default_concave_geom_geom_detector::<Mesh, Ball>();
        res.register_default_concave_geom_geom_detector::<Mesh, Box>();
        res.register_default_concave_geom_geom_detector::<Mesh, Cone>();
        res.register_default_concave_geom_geom_detector::<Mesh, Cylinder>();
        res.register_default_concave_geom_geom_detector::<Mesh, Capsule>();
        res.register_default_concave_geom_geom_detector::<Mesh, Convex>();
        res.register_default_concave_geom_geom_detector::<Mesh, Triangle>();
        res.register_default_concave_geom_geom_detector::<Mesh, Segment>();

        // FIXME: implement a ConcaveGeomConcaveGeom detector?
        res.register_default_concave_geom_geom_detector::<Compound, Compound>();
        res.register_default_concave_geom_geom_detector::<Mesh, Compound>();

        res
    }

    /// Registers a `PlaneImplicit` collision detector between a given implicit geometry and a plane.
    pub fn register_default_plane_implicit_detector<I: 'static + Implicit<LV, M>>(
                                                    &mut self,
                                                    generate_manifold: bool,
                                                    prediction:        &N) {
        let p = if generate_manifold { na::zero() } else { prediction.clone() };

        type IP    = ImplicitPlane<I>;
        let d1: IP = ImplicitPlane::new(p.clone());

        type PI    = PlaneImplicit<I>;
        let d2: PI = PlaneImplicit::new(p.clone());

        if generate_manifold {
            self.register_detector_with_contact_manifold_generator(d1, prediction);
            self.register_detector_with_contact_manifold_generator(d2, prediction);
        }
        else {
            self.register_detector(d1);
            self.register_detector(d2);
        }
    }

    /// Register an `ImplicitImplicit` collision detector between two implicit geometries.
    pub fn register_default_implicit_implicit_detector<G1: 'static            +
                                                           Implicit<LV, M> +
                                                           PreferedSamplingDirections<LV, M>,
                                                       G2: 'static            +
                                                           Implicit<LV, M> +
                                                           PreferedSamplingDirections<LV, M>,
                                                       S:  Send + Clone + Simplex<AnnotatedPoint>>(
                                                       &mut self,
                                                       generate_manifold: bool,
                                                       prediction:        &N,
                                                       simplex:           &S) {
        let p = if generate_manifold { na::zero() } else { prediction.clone() };

        type I1I2    = ImplicitImplicit<S, G1, G2>;
        let d1: I1I2 = ImplicitImplicit::new(p.clone(), simplex.clone());

        type I2I1    = ImplicitImplicit<S, G2, G1>;
        let d2: I2I1 = ImplicitImplicit::new(p.clone(), simplex.clone());

        if generate_manifold {
            self.register_detector_with_contact_manifold_generator(d1, prediction);
            self.register_detector_with_contact_manifold_generator(d2, prediction);
        }
        else {
            self.register_detector(d1);
            self.register_detector(d2);
        }
    }

    /// Register an `ConcaveGeomGeom` collision detector between a given concave geometry and a
    /// given geometry.
    pub fn register_default_concave_geom_geom_detector<G1: 'static + ConcaveGeom,
                                                       G2: 'static + Geom>(&mut self) {
        type CGG      = ConcaveGeomGeomFactory<G1, G2>;
        let  f1: CGG  = ConcaveGeomGeomFactory;

        type GCG      = GeomConcaveGeomFactory<G2, G1>;
        let  f2: GCG  = GeomConcaveGeomFactory;

        // FIXME: find a way to factorize that?
        let key     = (TypeId::of::<G1>(), TypeId::of::<G2>());
        self.constructors.insert(key, ~f1 as ~CollisionDetectorFactory);

        let key     = (TypeId::of::<G2>(), TypeId::of::<G1>());
        self.constructors.insert(key, ~f2 as ~CollisionDetectorFactory);
    }

    /// Register a given collision detector and adds it a contact manifold generator (a
    /// `OneShotContactManifoldGenerator`).
    pub fn register_detector_with_contact_manifold_generator<G1: 'static + Any,
                                                             G2: 'static + Any,
                                                             D:  'static + Send +
                                                                 CollisionDetector<G1, G2> +
                                                                 Clone>(
                                                             &mut self,
                                                             d:          D,
                                                             prediction: &N) {
        type OSCMGD   = OSCMG<D>;
        let d: OSCMGD = OSCMG::new(prediction.clone(), d);
        self.register_detector(d);
    }

    /*
    pub fn register_default_compound_any_detector<G, S>(&mut self) {
        type CA = CompoundAABBAny<~Geom>;
    }
    */

    /// Register `ImplicitImplicit` collision detectors between a given geometry and every implicit
    /// geometry supported by `ncollide`.
    pub fn register_default_implicit_detectors<G: 'static + Implicit<LV, M> + PreferedSamplingDirections<LV, M>>(
                                               &mut self,
                                               generate_manifold: bool,
                                               prediction:        &N) {
        type Simplex  = JohnsonSimplex<AnnotatedPoint>;

        // Implicit vs. Implicit
        let rt = RecursionTemplate::new(na::dim::<LV>());
        let js = &JohnsonSimplex::new(rt);

        self.register_default_implicit_implicit_detector::<Ball, G, Simplex>(false, prediction, js);
        self.register_default_implicit_implicit_detector::<Box, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Cone, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Cylinder, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Capsule, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Convex, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Triangle, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Segment, G, Simplex>(generate_manifold, prediction, js);
    }
}

// FIXME: rename that GeomGeomCollisionDetectorFactory ?
/// Trait of structures able do build a new collision detector.
pub trait CollisionDetectorFactory : Send + Freeze {
    /// Builds a new collision detector.
    fn build(&self) -> ~GeomGeomCollisionDetector;
}

/// Cloning-based collision detector factory.
pub struct CollisionDetectorCloner<CD> {
    priv template: CD
}

impl<CD: GeomGeomCollisionDetector + Clone> CollisionDetectorCloner<CD> {
    /// Creates a new `CollisionDetectorCloner`.
    ///
    /// The cloned detector is `CD`.
    fn new(detector: CD) -> CollisionDetectorCloner<CD> {
        CollisionDetectorCloner {
            template: detector
        }
    }
}

impl<CD: 'static + Send + Freeze + GeomGeomCollisionDetector + Clone>
CollisionDetectorFactory for CollisionDetectorCloner<CD> {
    fn build(&self) -> ~GeomGeomCollisionDetector {
        ~self.template.clone() as ~GeomGeomCollisionDetector
    }
}
<commit_msg>Update to the last Rust.<commit_after>//! Collision detector between two `~Geom`.

use std::unstable::intrinsics::TypeId;
use std::hashmap::HashMap;
use nalgebra::na;
use geom::{AnnotatedPoint, Geom, ConcaveGeom, Cone, Box, Ball, Capsule, Convex, Cylinder, Compound,
           Mesh, Triangle, Segment, Plane};
use implicit::{Implicit, PreferedSamplingDirections};
use contact::Contact;
use narrow::algorithm::simplex::Simplex;
use narrow::algorithm::johnson_simplex::{JohnsonSimplex, RecursionTemplate};
use narrow::{CollisionDetector, ImplicitImplicit, BallBall,
             ImplicitPlane, PlaneImplicit, ConcaveGeomGeomFactory, GeomConcaveGeomFactory};
use OSCMG = narrow::OneShotContactManifoldGenerator;
use math::{N, V, LV, M};

/// Same as the `CollisionDetector` trait but using dynamic dispatch on the geometries.
pub trait GeomGeomCollisionDetector {
    /// Runs the collision detection on two objects. It is assumed that the same
    /// collision detector (the same structure) is always used with the same
    /// pair of object.
    fn update(&mut self,
              &GeomGeomDispatcher,
              &M,
              &Geom,
              &M,
              &Geom);

    /// The number of collision detected during the last update.
    fn num_colls(&self) -> uint;

    /// Collects the collisions detected during the last update.
    fn colls(&self, &mut ~[Contact]);
}

/// Trait to be implemented by collision detector using dynamic dispatch.
///
/// This is used to know the exact type of the geometries.
pub trait DynamicCollisionDetector<G1, G2>: GeomGeomCollisionDetector { }

#[deriving(Clone)]
struct DetectorWithoutRedispatch<D> {
    detector: D
}

impl<D> DetectorWithoutRedispatch<D> {
    pub fn new(d: D) -> DetectorWithoutRedispatch<D> {
        DetectorWithoutRedispatch {
            detector: d
        }
    }
}

impl<D: CollisionDetector<G1, G2>, G1, G2>
CollisionDetector<G1, G2> for DetectorWithoutRedispatch<D> {
    fn update(&mut self, _: &M, _: &G1, _: &M, _: &G2) { unreachable!() }
    fn num_colls(&self) -> uint { unreachable!() }
    fn colls(&self, _: &mut ~[Contact]) { unreachable!() }
    fn toi(_: Option<DetectorWithoutRedispatch<D>>, _: &M, _: &V, _: &N, _: &G1, _: &M, _: &G2) -> Option<N> {
        unreachable!()
    }
}

impl<D: CollisionDetector<G1, G2>, G1, G2>
DynamicCollisionDetector<G1, G2> for DetectorWithoutRedispatch<D> { }

impl<D: CollisionDetector<G1, G2>, G1: 'static, G2: 'static>
GeomGeomCollisionDetector for DetectorWithoutRedispatch<D> {
    #[inline]
    fn update(&mut self,
              _:  &GeomGeomDispatcher,
              m1: &M,
              g1: &Geom,
              m2: &M,
              g2: &Geom) {
        self.detector.update(
            m1,
            g1.as_ref::<G1>().expect("Invalid geometry."),
            m2,
            g2.as_ref::<G2>().expect("Invalid geometry."))
    }

    #[inline]
    fn num_colls(&self) -> uint {
        self.detector.num_colls()
    }

    #[inline]
    fn colls(&self, cs: &mut ~[Contact]) {
        self.detector.colls(cs)
    }
}

/// Collision dispatcher between two `~Geom`.
pub struct GeomGeomDispatcher {
    priv constructors: HashMap<(TypeId, TypeId), ~CollisionDetectorFactory>
}

impl GeomGeomDispatcher {
    /// Creates a new `GeomGeomDispatcher` without the default set of collision detectors
    /// factories.
    pub fn new_without_default() -> GeomGeomDispatcher {
        GeomGeomDispatcher {
            constructors: HashMap::new()
        }
    }

    /// Registers a new dynamic collision detector for two geometries.
    pub fn register_dynamic_detector<G1: 'static + Any,
                                     G2: 'static + Any,
                                     D:  'static + Send + Clone +
                                         DynamicCollisionDetector<G1, G2>>(
                                     &mut self,
                                     d:   D) {
        let key     = (TypeId::of::<G1>(), TypeId::of::<G2>());
        let factory = ~CollisionDetectorCloner::new(d);
        self.constructors.insert(key, factory as ~CollisionDetectorFactory);
    }

    /// Registers a new collision detector for two geometries.
    pub fn register_detector<G1: 'static + Any,
                             G2: 'static + Any,
                             D:  'static + Send + CollisionDetector<G1, G2> + Clone>(
                             &mut self,
                             d:   D) {
        self.register_dynamic_detector(DetectorWithoutRedispatch::new(d));
    }

    /// Unregister the collision detector for a givem pair of geometries.
    pub fn unregister_detector<G1: 'static + Any, G2: 'static + Any>(&mut self) {
        let key = (TypeId::of::<G1>(), TypeId::of::<G2>());
        self.constructors.remove(&key);
    }

    /// Creates a new collision detector adapted for the two given geometries.
    pub fn dispatch(&self, a: &Geom, b: &Geom) -> ~GeomGeomCollisionDetector {
        match self.constructors.find(&(a.get_type_id(), b.get_type_id())) {
            Some(f) => f.build(),
            None    => fail!("Unable to find a collision detector.")
        }
    }
}

impl GeomGeomDispatcher {
    // FIXME: make this a function which has the simplex and the prediction margin as parameters
    /// Creates a new `GeomGeomDispatcher` able do build collision detectors for any valid pair of
    /// geometries supported by `ncollide`.
    pub fn new() -> GeomGeomDispatcher {
        let mut res = GeomGeomDispatcher::new_without_default();

        type Simplex  = JohnsonSimplex<AnnotatedPoint>;
        type Self     = GeomGeomDispatcher;
        type Super    = ~GeomGeomCollisionDetector;

        /*
         * Involving a Plane
         */
        let prediction: &N = &na::cast(0.1);

        // Ball vs. Ball
        let bb: BallBall = BallBall::new(prediction.clone());
        res.register_detector(bb);

        // Plane vs. Implicit
        res.register_default_plane_implicit_detector::<Ball>(false, prediction);
        res.register_default_plane_implicit_detector::<Box>(true, prediction);
        res.register_default_plane_implicit_detector::<Cone>(true, prediction);
        res.register_default_plane_implicit_detector::<Cylinder>(true, prediction);
        res.register_default_plane_implicit_detector::<Capsule>(true, prediction);
        res.register_default_plane_implicit_detector::<Convex>(true, prediction);
        res.register_default_plane_implicit_detector::<Triangle>(true, prediction);
        res.register_default_plane_implicit_detector::<Segment>(true, prediction);

        // Implicit vs. Implicit
        // NOTE: some pair will be registered twice…
        res.register_default_implicit_detectors::<Box>(true, prediction);
        res.register_default_implicit_detectors::<Cone>(true, prediction);
        res.register_default_implicit_detectors::<Cylinder>(true, prediction);
        res.register_default_implicit_detectors::<Capsule>(true, prediction);
        res.register_default_implicit_detectors::<Convex>(true, prediction);
        res.register_default_implicit_detectors::<Triangle>(true, prediction);
        res.register_default_implicit_detectors::<Segment>(true, prediction);

        // FIXME: refactor the three following blocks?
        // Compound vs. Other
        res.register_default_concave_geom_geom_detector::<Compound, Plane>();
        res.register_default_concave_geom_geom_detector::<Compound, Ball>();
        res.register_default_concave_geom_geom_detector::<Compound, Box>();
        res.register_default_concave_geom_geom_detector::<Compound, Cone>();
        res.register_default_concave_geom_geom_detector::<Compound, Cylinder>();
        res.register_default_concave_geom_geom_detector::<Compound, Capsule>();
        res.register_default_concave_geom_geom_detector::<Compound, Convex>();
        res.register_default_concave_geom_geom_detector::<Compound, Triangle>();
        res.register_default_concave_geom_geom_detector::<Compound, Segment>();

        // TriangleMesh vs. Other
        res.register_default_concave_geom_geom_detector::<Mesh, Plane>();
        res.register_default_concave_geom_geom_detector::<Mesh, Ball>();
        res.register_default_concave_geom_geom_detector::<Mesh, Box>();
        res.register_default_concave_geom_geom_detector::<Mesh, Cone>();
        res.register_default_concave_geom_geom_detector::<Mesh, Cylinder>();
        res.register_default_concave_geom_geom_detector::<Mesh, Capsule>();
        res.register_default_concave_geom_geom_detector::<Mesh, Convex>();
        res.register_default_concave_geom_geom_detector::<Mesh, Triangle>();
        res.register_default_concave_geom_geom_detector::<Mesh, Segment>();

        // FIXME: implement a ConcaveGeomConcaveGeom detector?
        res.register_default_concave_geom_geom_detector::<Compound, Compound>();
        res.register_default_concave_geom_geom_detector::<Mesh, Compound>();

        res
    }

    /// Registers a `PlaneImplicit` collision detector between a given implicit geometry and a plane.
    pub fn register_default_plane_implicit_detector<I: 'static + Implicit<LV, M>>(
                                                    &mut self,
                                                    generate_manifold: bool,
                                                    prediction:        &N) {
        let p = if generate_manifold { na::zero() } else { prediction.clone() };

        let d1 = ImplicitPlane::<I>::new(p.clone());
        let d2 = PlaneImplicit::<I>::new(p.clone());

        if generate_manifold {
            self.register_detector_with_contact_manifold_generator(d1, prediction);
            self.register_detector_with_contact_manifold_generator(d2, prediction);
        }
        else {
            self.register_detector(d1);
            self.register_detector(d2);
        }
    }

    /// Register an `ImplicitImplicit` collision detector between two implicit geometries.
    pub fn register_default_implicit_implicit_detector<G1: 'static            +
                                                           Implicit<LV, M> +
                                                           PreferedSamplingDirections<LV, M>,
                                                       G2: 'static            +
                                                           Implicit<LV, M> +
                                                           PreferedSamplingDirections<LV, M>,
                                                       S:  Send + Clone + Simplex<AnnotatedPoint>>(
                                                       &mut self,
                                                       generate_manifold: bool,
                                                       prediction:        &N,
                                                       simplex:           &S) {
        let p = if generate_manifold { na::zero() } else { prediction.clone() };

        let d1 = ImplicitImplicit::<S, G1, G2>::new(p.clone(), simplex.clone());
        let d2 = ImplicitImplicit::<S, G2, G1>::new(p.clone(), simplex.clone());

        if generate_manifold {
            self.register_detector_with_contact_manifold_generator(d1, prediction);
            self.register_detector_with_contact_manifold_generator(d2, prediction);
        }
        else {
            self.register_detector(d1);
            self.register_detector(d2);
        }
    }

    /// Register an `ConcaveGeomGeom` collision detector between a given concave geometry and a
    /// given geometry.
    pub fn register_default_concave_geom_geom_detector<G1: 'static + ConcaveGeom,
                                                       G2: 'static + Geom>(&mut self) {
        let  f1 = ConcaveGeomGeomFactory::<G1, G2>;
        let  f2 = GeomConcaveGeomFactory::<G2, G1>;

        // FIXME: find a way to factorize that?
        let key     = (TypeId::of::<G1>(), TypeId::of::<G2>());
        self.constructors.insert(key, ~f1 as ~CollisionDetectorFactory);

        let key     = (TypeId::of::<G2>(), TypeId::of::<G1>());
        self.constructors.insert(key, ~f2 as ~CollisionDetectorFactory);
    }

    /// Register a given collision detector and adds it a contact manifold generator (a
    /// `OneShotContactManifoldGenerator`).
    pub fn register_detector_with_contact_manifold_generator<G1: 'static + Any,
                                                             G2: 'static + Any,
                                                             D:  'static + Send +
                                                                 CollisionDetector<G1, G2> +
                                                                 Clone>(
                                                             &mut self,
                                                             d:          D,
                                                             prediction: &N) {
        let d = OSCMG::<D>::new(prediction.clone(), d);
        self.register_detector(d);
    }

    /*
    pub fn register_default_compound_any_detector<G, S>(&mut self) {
        type CA = CompoundAABBAny<~Geom>;
    }
    */

    /// Register `ImplicitImplicit` collision detectors between a given geometry and every implicit
    /// geometry supported by `ncollide`.
    pub fn register_default_implicit_detectors<G: 'static + Implicit<LV, M> + PreferedSamplingDirections<LV, M>>(
                                               &mut self,
                                               generate_manifold: bool,
                                               prediction:        &N) {
        type Simplex  = JohnsonSimplex<AnnotatedPoint>;

        // Implicit vs. Implicit
        let rt = RecursionTemplate::new(na::dim::<LV>());
        let js = &JohnsonSimplex::new(rt);

        self.register_default_implicit_implicit_detector::<Ball, G, Simplex>(false, prediction, js);
        self.register_default_implicit_implicit_detector::<Box, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Cone, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Cylinder, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Capsule, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Convex, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Triangle, G, Simplex>(generate_manifold, prediction, js);
        self.register_default_implicit_implicit_detector::<Segment, G, Simplex>(generate_manifold, prediction, js);
    }
}

// FIXME: rename that GeomGeomCollisionDetectorFactory ?
/// Trait of structures able do build a new collision detector.
pub trait CollisionDetectorFactory : Send + Freeze {
    /// Builds a new collision detector.
    fn build(&self) -> ~GeomGeomCollisionDetector;
}

/// Cloning-based collision detector factory.
pub struct CollisionDetectorCloner<CD> {
    priv template: CD
}

impl<CD: GeomGeomCollisionDetector + Clone> CollisionDetectorCloner<CD> {
    /// Creates a new `CollisionDetectorCloner`.
    ///
    /// The cloned detector is `CD`.
    fn new(detector: CD) -> CollisionDetectorCloner<CD> {
        CollisionDetectorCloner {
            template: detector
        }
    }
}

impl<CD: 'static + Send + Freeze + GeomGeomCollisionDetector + Clone>
CollisionDetectorFactory for CollisionDetectorCloner<CD> {
    fn build(&self) -> ~GeomGeomCollisionDetector {
        ~self.template.clone() as ~GeomGeomCollisionDetector
    }
}
<|endoftext|>


<commit_before>//! The element nodes.

// SVG, Script, Style element has no Default.
#![allow(clippy::new_without_default)]
// Clippy doesn't like the `.add()` method on elements.
#![allow(clippy::should_implement_trait)]

use std::fmt;

use crate::node::{Attributes, Children, Node, Value};

pub mod path;
pub mod tag;

#[doc(hidden)]
#[derive(Clone, Debug)]
pub struct Element {
    name: String,
    attributes: Attributes,
    children: Children,
}

impl Element {
    pub fn new<T>(name: T) -> Self
    where
        T: Into<String>,
    {
        Element {
            name: name.into(),
            attributes: Attributes::new(),
            children: Children::new(),
        }
    }
}

impl fmt::Display for Element {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "<{}", self.name)?;
        let mut attributes = self.attributes.iter().collect::<Vec<_>>();
        attributes.sort_by_key(|pair| pair.0.as_str());
        for (name, value) in attributes {
            match (value.contains('\''), value.contains('"')) {
                (true, false) | (false, false) => {
                    write!(formatter, r#" {}="{}""#, name, value)?;
                }
                (false, true) => {
                    write!(formatter, r#" {}='{}'"#, name, value)?;
                }
                _ => {}
            }
        }
        if self.children.is_empty() {
            return write!(formatter, "/>");
        }
        write!(formatter, ">")?;
        for child in self.children.iter() {
            write!(formatter, "\n{}", child)?;
        }
        write!(formatter, "\n</{}>", self.name)
    }
}

impl Node for Element {
    #[inline]
    fn append<T>(&mut self, node: T)
    where
        T: Node,
    {
        self.children.push(Box::new(node));
    }

    #[inline]
    fn assign<T, U>(&mut self, name: T, value: U)
    where
        T: Into<String>,
        U: Into<Value>,
    {
        self.attributes.insert(name.into(), value.into());
    }
}

macro_rules! implement {
    ($(#[$doc:meta] struct $struct_name:ident)*) => ($(
        #[$doc]
        #[derive(Clone, Debug)]
        pub struct $struct_name {
            inner: Element,
        }

        impl $struct_name {
            /// Create a node.
            #[inline]
            pub fn new() -> Self {
                $struct_name {
                    inner: Element::new(tag::$struct_name),
                }
            }
        }

        impl Default for $struct_name {
            fn default() -> Self {
                Self::new()
            }
        }

        node! { $struct_name::inner }
    )*);
}

implement! {
    #[doc = "An [`animate`](https://www.w3.org/TR/SVG/animate.html#AnimateElement) element."]
    struct Animate

    #[doc = "An [`animateColor`](https://www.w3.org/TR/SVG/animate.html#AnimateColorElement) element."]
    struct AnimateColor

    #[doc = "An [`animateMotion`](https://www.w3.org/TR/SVG/animate.html#AnimateMotionElement) element."]
    struct AnimateMotion

    #[doc = "An [`animateTransform`](https://www.w3.org/TR/SVG/animate.html#AnimateTransformElement) element."]
    struct AnimateTransform

    #[doc = "A [`circle`](https://www.w3.org/TR/SVG/shapes.html#CircleElement) element."]
    struct Circle

    #[doc = "A [`clipPath`](https://www.w3.org/TR/SVG/masking.html#ClipPathElement) element."]
    struct ClipPath

    #[doc = "A [`defs`](https://www.w3.org/TR/SVG/struct.html#DefsElement) element."]
    struct Definitions

    #[doc = "A [`desc`](https://www.w3.org/TR/SVG/struct.html#DescElement) element."]
    struct Description

    #[doc = "An [`ellipse`](https://www.w3.org/TR/SVG/shapes.html#EllipseElement) element."]
    struct Ellipse

    #[doc = "A [`filter`](https://www.w3.org/TR/SVG/filters.html#FilterElement) element."]
    struct Filter

    #[doc = "A [`g`](https://www.w3.org/TR/SVG/struct.html#GElement) element."]
    struct Group

    #[doc = "An [`image`](https://www.w3.org/TR/SVG/struct.html#ImageElement) element."]
    struct Image

    #[doc = "A [`line`](https://www.w3.org/TR/SVG/shapes.html#LineElement) element."]
    struct Line

    #[doc = "A [`linearGradient`](https://www.w3.org/TR/SVG/pservers.html#LinearGradientElement) element."]
    struct LinearGradient

    #[doc = "An [`a`](https://www.w3.org/TR/SVG/linking.html#AElement) element."]
    struct Link

    #[doc = "A [`marker`](https://www.w3.org/TR/SVG/painting.html#MarkerElement) element."]
    struct Marker

    #[doc = "A [`mask`](https://www.w3.org/TR/SVG/masking.html#MaskElement) element."]
    struct Mask

    #[doc = "An [`mpath`](https://www.w3.org/TR/SVG/animate.html#MPathElement) element."]
    struct MotionPath

    #[doc = "A [`path`](https://www.w3.org/TR/SVG/paths.html#PathElement) element."]
    struct Path

    #[doc = "A [`pattern`](https://www.w3.org/TR/SVG/pservers.html#PatternElement) element."]
    struct Pattern

    #[doc = "A [`polygon`](https://www.w3.org/TR/SVG/shapes.html#PolygonElement) element."]
    struct Polygon

    #[doc = "A [`polyline`](https://www.w3.org/TR/SVG/shapes.html#PolylineElement) element."]
    struct Polyline

    #[doc = "A [`radialGradient`](https://www.w3.org/TR/SVG/pservers.html#RadialGradientElement) element."]
    struct RadialGradient

    #[doc = "A [`rect`](https://www.w3.org/TR/SVG/shapes.html#RectElement) element."]
    struct Rectangle

    #[doc = "A [`stop`](https://www.w3.org/TR/SVG/pservers.html#StopElement) element."]
    struct Stop

    #[doc = "A [`symbol`](https://www.w3.org/TR/SVG/struct.html#SymbolElement) element."]
    struct Symbol

    #[doc = "A [`text`](https://www.w3.org/TR/SVG/text.html#TextElement) element."]
    struct Text

    #[doc = "A [`textPath`](https://www.w3.org/TR/SVG/text.html#TextPathElement) element."]
    struct TextPath

    #[doc = "A [`title`](https://www.w3.org/TR/SVG/struct.html#TitleElement) element."]
    struct Title

    #[doc = "A [`use`](https://www.w3.org/TR/SVG/struct.html#UseElement) element."]
    struct Use
}

macro_rules! implement {
    (@itemize $i:item) => ($i);
    ($(
        #[$doc:meta]
        struct $struct_name:ident
        [$($pn:ident: $($pt:tt)*),*] [$inner:ident $(,$an:ident: $at:ty)*] $body:block
    )*) => ($(
        #[$doc]
        #[derive(Clone, Debug)]
        pub struct $struct_name {
            inner: Element,
        }

        implement! { @itemize
            impl $struct_name {
                /// Create a node.
                #[inline]
                pub fn new<$($pn: $($pt)*),*>($($an: $at),*) -> Self {
                    #[inline(always)]
                    fn initialize<$($pn: $($pt)*),*>($inner: &mut Element $(, $an: $at)*) $body
                    let mut inner = Element::new(tag::$struct_name);
                    initialize(&mut inner $(, $an)*);
                    $struct_name {
                        inner: inner,
                    }
                }
            }
        }

        node! { $struct_name::inner }
    )*);
}

implement! {
    #[doc = "An [`svg`](https://www.w3.org/TR/SVG/struct.html#SVGElement) element."]
    struct SVG [] [inner] {
        inner.assign("xmlns", "http://www.w3.org/2000/svg");
    }

    #[doc = "A [`script`](https://www.w3.org/TR/SVG/script.html#ScriptElement) element."]
    struct Script [T: Into<String>] [inner, content: T] {
        inner.append(crate::node::Text::new(content));
    }

    #[doc = "A [`style`](https://www.w3.org/TR/SVG/styling.html#StyleElement) element."]
    struct Style [T: Into<String>] [inner, content: T] {
        inner.append(crate::node::Text::new(content));
    }
}

#[cfg(test)]
mod tests {
    use super::{Element, Style};
    use crate::node::Node;

    #[test]
    fn element_display() {
        let mut element = Element::new("foo");
        element.assign("x", -10);
        element.assign("y", "10px");
        element.assign("s", (12.5, 13.0));
        element.assign("c", "green");
        element.append(Element::new("bar"));

        assert_eq!(
            element.to_string(),
            "<foo c=\"green\" s=\"12.5 13\" x=\"-10\" y=\"10px\">\n\
             <bar/>\n\
             </foo>\
             "
        );
    }

    #[test]
    fn element_display_quotes() {
        let mut element = Element::new("foo");
        element.assign("s", "'single'");
        element.assign("d", r#""double""#);
        element.assign("m", r#""mixed'"#);

        assert_eq!(element.to_string(), r#"<foo d='"double"' s="'single'"/>"#);
    }

    #[test]
    fn style_display() {
        let element = Style::new("* { font-family: foo; }");

        assert_eq!(
            element.to_string(),
            "<style>\n\
             * { font-family: foo; }\n\
             </style>\
             "
        );
    }
}
<commit_msg>Remove comment explaining clippy warnings<commit_after>//! The element nodes.

#![allow(clippy::new_without_default)]
#![allow(clippy::should_implement_trait)]

use std::fmt;

use crate::node::{Attributes, Children, Node, Value};

pub mod path;
pub mod tag;

#[doc(hidden)]
#[derive(Clone, Debug)]
pub struct Element {
    name: String,
    attributes: Attributes,
    children: Children,
}

impl Element {
    pub fn new<T>(name: T) -> Self
    where
        T: Into<String>,
    {
        Element {
            name: name.into(),
            attributes: Attributes::new(),
            children: Children::new(),
        }
    }
}

impl fmt::Display for Element {
    fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
        write!(formatter, "<{}", self.name)?;
        let mut attributes = self.attributes.iter().collect::<Vec<_>>();
        attributes.sort_by_key(|pair| pair.0.as_str());
        for (name, value) in attributes {
            match (value.contains('\''), value.contains('"')) {
                (true, false) | (false, false) => {
                    write!(formatter, r#" {}="{}""#, name, value)?;
                }
                (false, true) => {
                    write!(formatter, r#" {}='{}'"#, name, value)?;
                }
                _ => {}
            }
        }
        if self.children.is_empty() {
            return write!(formatter, "/>");
        }
        write!(formatter, ">")?;
        for child in self.children.iter() {
            write!(formatter, "\n{}", child)?;
        }
        write!(formatter, "\n</{}>", self.name)
    }
}

impl Node for Element {
    #[inline]
    fn append<T>(&mut self, node: T)
    where
        T: Node,
    {
        self.children.push(Box::new(node));
    }

    #[inline]
    fn assign<T, U>(&mut self, name: T, value: U)
    where
        T: Into<String>,
        U: Into<Value>,
    {
        self.attributes.insert(name.into(), value.into());
    }
}

macro_rules! implement {
    ($(#[$doc:meta] struct $struct_name:ident)*) => ($(
        #[$doc]
        #[derive(Clone, Debug)]
        pub struct $struct_name {
            inner: Element,
        }

        impl $struct_name {
            /// Create a node.
            #[inline]
            pub fn new() -> Self {
                $struct_name {
                    inner: Element::new(tag::$struct_name),
                }
            }
        }

        impl Default for $struct_name {
            fn default() -> Self {
                Self::new()
            }
        }

        node! { $struct_name::inner }
    )*);
}

implement! {
    #[doc = "An [`animate`](https://www.w3.org/TR/SVG/animate.html#AnimateElement) element."]
    struct Animate

    #[doc = "An [`animateColor`](https://www.w3.org/TR/SVG/animate.html#AnimateColorElement) element."]
    struct AnimateColor

    #[doc = "An [`animateMotion`](https://www.w3.org/TR/SVG/animate.html#AnimateMotionElement) element."]
    struct AnimateMotion

    #[doc = "An [`animateTransform`](https://www.w3.org/TR/SVG/animate.html#AnimateTransformElement) element."]
    struct AnimateTransform

    #[doc = "A [`circle`](https://www.w3.org/TR/SVG/shapes.html#CircleElement) element."]
    struct Circle

    #[doc = "A [`clipPath`](https://www.w3.org/TR/SVG/masking.html#ClipPathElement) element."]
    struct ClipPath

    #[doc = "A [`defs`](https://www.w3.org/TR/SVG/struct.html#DefsElement) element."]
    struct Definitions

    #[doc = "A [`desc`](https://www.w3.org/TR/SVG/struct.html#DescElement) element."]
    struct Description

    #[doc = "An [`ellipse`](https://www.w3.org/TR/SVG/shapes.html#EllipseElement) element."]
    struct Ellipse

    #[doc = "A [`filter`](https://www.w3.org/TR/SVG/filters.html#FilterElement) element."]
    struct Filter

    #[doc = "A [`g`](https://www.w3.org/TR/SVG/struct.html#GElement) element."]
    struct Group

    #[doc = "An [`image`](https://www.w3.org/TR/SVG/struct.html#ImageElement) element."]
    struct Image

    #[doc = "A [`line`](https://www.w3.org/TR/SVG/shapes.html#LineElement) element."]
    struct Line

    #[doc = "A [`linearGradient`](https://www.w3.org/TR/SVG/pservers.html#LinearGradientElement) element."]
    struct LinearGradient

    #[doc = "An [`a`](https://www.w3.org/TR/SVG/linking.html#AElement) element."]
    struct Link

    #[doc = "A [`marker`](https://www.w3.org/TR/SVG/painting.html#MarkerElement) element."]
    struct Marker

    #[doc = "A [`mask`](https://www.w3.org/TR/SVG/masking.html#MaskElement) element."]
    struct Mask

    #[doc = "An [`mpath`](https://www.w3.org/TR/SVG/animate.html#MPathElement) element."]
    struct MotionPath

    #[doc = "A [`path`](https://www.w3.org/TR/SVG/paths.html#PathElement) element."]
    struct Path

    #[doc = "A [`pattern`](https://www.w3.org/TR/SVG/pservers.html#PatternElement) element."]
    struct Pattern

    #[doc = "A [`polygon`](https://www.w3.org/TR/SVG/shapes.html#PolygonElement) element."]
    struct Polygon

    #[doc = "A [`polyline`](https://www.w3.org/TR/SVG/shapes.html#PolylineElement) element."]
    struct Polyline

    #[doc = "A [`radialGradient`](https://www.w3.org/TR/SVG/pservers.html#RadialGradientElement) element."]
    struct RadialGradient

    #[doc = "A [`rect`](https://www.w3.org/TR/SVG/shapes.html#RectElement) element."]
    struct Rectangle

    #[doc = "A [`stop`](https://www.w3.org/TR/SVG/pservers.html#StopElement) element."]
    struct Stop

    #[doc = "A [`symbol`](https://www.w3.org/TR/SVG/struct.html#SymbolElement) element."]
    struct Symbol

    #[doc = "A [`text`](https://www.w3.org/TR/SVG/text.html#TextElement) element."]
    struct Text

    #[doc = "A [`textPath`](https://www.w3.org/TR/SVG/text.html#TextPathElement) element."]
    struct TextPath

    #[doc = "A [`title`](https://www.w3.org/TR/SVG/struct.html#TitleElement) element."]
    struct Title

    #[doc = "A [`use`](https://www.w3.org/TR/SVG/struct.html#UseElement) element."]
    struct Use
}

macro_rules! implement {
    (@itemize $i:item) => ($i);
    ($(
        #[$doc:meta]
        struct $struct_name:ident
        [$($pn:ident: $($pt:tt)*),*] [$inner:ident $(,$an:ident: $at:ty)*] $body:block
    )*) => ($(
        #[$doc]
        #[derive(Clone, Debug)]
        pub struct $struct_name {
            inner: Element,
        }

        implement! { @itemize
            impl $struct_name {
                /// Create a node.
                #[inline]
                pub fn new<$($pn: $($pt)*),*>($($an: $at),*) -> Self {
                    #[inline(always)]
                    fn initialize<$($pn: $($pt)*),*>($inner: &mut Element $(, $an: $at)*) $body
                    let mut inner = Element::new(tag::$struct_name);
                    initialize(&mut inner $(, $an)*);
                    $struct_name {
                        inner: inner,
                    }
                }
            }
        }

        node! { $struct_name::inner }
    )*);
}

implement! {
    #[doc = "An [`svg`](https://www.w3.org/TR/SVG/struct.html#SVGElement) element."]
    struct SVG [] [inner] {
        inner.assign("xmlns", "http://www.w3.org/2000/svg");
    }

    #[doc = "A [`script`](https://www.w3.org/TR/SVG/script.html#ScriptElement) element."]
    struct Script [T: Into<String>] [inner, content: T] {
        inner.append(crate::node::Text::new(content));
    }

    #[doc = "A [`style`](https://www.w3.org/TR/SVG/styling.html#StyleElement) element."]
    struct Style [T: Into<String>] [inner, content: T] {
        inner.append(crate::node::Text::new(content));
    }
}

#[cfg(test)]
mod tests {
    use super::{Element, Style};
    use crate::node::Node;

    #[test]
    fn element_display() {
        let mut element = Element::new("foo");
        element.assign("x", -10);
        element.assign("y", "10px");
        element.assign("s", (12.5, 13.0));
        element.assign("c", "green");
        element.append(Element::new("bar"));

        assert_eq!(
            element.to_string(),
            "<foo c=\"green\" s=\"12.5 13\" x=\"-10\" y=\"10px\">\n\
             <bar/>\n\
             </foo>\
             "
        );
    }

    #[test]
    fn element_display_quotes() {
        let mut element = Element::new("foo");
        element.assign("s", "'single'");
        element.assign("d", r#""double""#);
        element.assign("m", r#""mixed'"#);

        assert_eq!(element.to_string(), r#"<foo d='"double"' s="'single'"/>"#);
    }

    #[test]
    fn style_display() {
        let element = Style::new("* { font-family: foo; }");

        assert_eq!(
            element.to_string(),
            "<style>\n\
             * { font-family: foo; }\n\
             </style>\
             "
        );
    }
}
<|endoftext|>


<commit_before>use crate::alloc::string::{String, ToString};
use scroll::{ctx, Pread, Pwrite};
use crate::error::{self, Error};
use crate::pe::relocation;
use std::io::Write;

#[repr(C)]
#[derive(Debug, PartialEq, Clone, Default)]
pub struct SectionTable {
    pub name: [u8; 8],
    pub real_name: Option<String>,
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: u32,
}

pub const SIZEOF_SECTION_TABLE: usize = 8 * 5;

// Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L70
// Decodes a string table entry in base 64 (//AAAAAA). Expects string without
// prefixed slashes.
fn base64_decode_string_entry(s: &str) -> Result<usize, ()> {
    assert!(s.len() <= 6, "String too long, possible overflow.");

    let mut val = 0;
    for c in s.bytes() {
        let v = if b'A' <= c && c <= b'Z' { // 00..=25
            c - b'A'
        } else if b'a' <= c && c <= b'z' { // 26..=51
            c - b'a' + 26
        } else if b'0' <= c && c <= b'9' { // 52..=61
            c - b'0' + 52
        } else if c == b'+' { // 62
            62
        } else if c == b'/' { // 63
            63
        } else {
            return Err(())
        };
        val = val * 64 + v as usize;
    }
    Ok(val)
}

impl SectionTable {
    pub fn parse(bytes: &[u8], offset: &mut usize, string_table_offset: usize) -> error::Result<Self> {
        let mut table = SectionTable::default();
        let mut name = [0u8; 8];
        name.copy_from_slice(bytes.gread_with(offset, 8)?);

        table.name = name;
        table.virtual_size = bytes.gread_with(offset, scroll::LE)?;
        table.virtual_address = bytes.gread_with(offset, scroll::LE)?;
        table.size_of_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.characteristics = bytes.gread_with(offset, scroll::LE)?;

        if let Some(idx) = table.name_offset()? {
            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
        }
        Ok(table)
    }

    pub fn name_offset(&self) -> error::Result<Option<usize>> {
        // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
        if self.name[0] == b'/' {
            let idx: usize = if self.name[1] == b'/' {
                let b64idx = self.name.pread::<&str>(2)?;
                base64_decode_string_entry(b64idx).map_err(|_|
                    Error::Malformed(format!("Invalid indirect section name //{}: base64 decoding failed", b64idx)))?
            } else {
                let name = self.name.pread::<&str>(1)?;
                name.parse().map_err(|err|
                    Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
            };
            Ok(Some(idx))
        } else {
            Ok(None)
        }
    }

    pub fn set_name_offset(&mut self, mut idx: usize) -> error::Result<()> {
        if idx <= 9_999_999 { // 10^7 - 1
            self.name = [0; 8];
            self.name[0] = b'/';
            write!(&mut self.name[1..], "{}", idx).unwrap();
            Ok(())
        } else if idx <= 0xfff_fff_fff { // 64^6 - 1
            self.name[0] = b'/';
            self.name[1] = b'/';
            for i in 0..6 {
                let rem = (idx % 64) as u8;
                idx = idx / 64;
                let c = match rem {
                    0..=25 => b'A' + rem,
                    26..=51 => b'a' + rem - 26,
                    52..=61 => b'0' + rem - 52,
                    62 => b'+',
                    63 => b'/',
                    _ => unreachable!(),
                };
                self.name[7 - i] = c;
            }
            Ok(())
        } else {
            Err(Error::Malformed(format!("Invalid section name offset: {}", idx)))
        }
    }

    pub fn name(&self) -> error::Result<&str> {
        match self.real_name.as_ref() {
            Some(s) => Ok(s),
            None => Ok(self.name.pread(0)?)
        }
    }

    pub fn relocations<'a>(&self, bytes: &'a[u8]) -> error::Result<relocation::Relocations<'a>> {
        let offset = self.pointer_to_relocations as usize;
        let number = self.number_of_relocations as usize;
        relocation::Relocations::parse(bytes, offset, number)
    }
}

impl ctx::SizeWith<scroll::Endian> for SectionTable {
    type Units = usize;
    fn size_with(_ctx: &scroll::Endian) -> usize {
        SIZEOF_SECTION_TABLE
    }
}

impl ctx::TryIntoCtx<scroll::Endian> for SectionTable {
    type Error = error::Error;
    type Size = usize;
    fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<Self::Size, Self::Error> {
        let offset = &mut 0;
        bytes.gwrite(&self.name[..], offset)?;
        bytes.gwrite_with(self.virtual_size, offset, ctx)?;
        bytes.gwrite_with(self.virtual_address, offset, ctx)?;
        bytes.gwrite_with(self.size_of_raw_data, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_raw_data, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_relocations, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_linenumbers, offset, ctx)?;
        bytes.gwrite_with(self.number_of_relocations, offset, ctx)?;
        bytes.gwrite_with(self.number_of_linenumbers, offset, ctx)?;
        bytes.gwrite_with(self.characteristics, offset, ctx)?;
        Ok(SIZEOF_SECTION_TABLE)
    }
}

impl ctx::IntoCtx<scroll::Endian> for SectionTable {
    fn into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) {
        bytes.pwrite_with(self, 0, ctx).unwrap();
    }
}

/// The section should not be padded to the next boundary. This flag is obsolete and is replaced
/// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x0000_0008;
/// The section contains executable code.
pub const IMAGE_SCN_CNT_CODE: u32 = 0x0000_0020;
/// The section contains initialized data.
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x0000_0040;
///  The section contains uninitialized data.
pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x0000_0080;
pub const IMAGE_SCN_LNK_OTHER: u32 = 0x0000_0100;
/// The section contains comments or other information. The .drectve section has this type.
/// This is valid for object files only.
pub const IMAGE_SCN_LNK_INFO: u32 = 0x0000_0200;
/// The section will not become part of the image. This is valid only for object files.
pub const IMAGE_SCN_LNK_REMOVE: u32 = 0x0000_0800;
/// The section contains COMDAT data. This is valid only for object files.
pub const IMAGE_SCN_LNK_COMDAT: u32 = 0x0000_1000;
/// The section contains data referenced through the global pointer (GP).
pub const IMAGE_SCN_GPREL: u32 = 0x0000_8000;
pub const IMAGE_SCN_MEM_PURGEABLE: u32 = 0x0002_0000;
pub const IMAGE_SCN_MEM_16BIT: u32 = 0x0002_0000;
pub const IMAGE_SCN_MEM_LOCKED: u32 = 0x0004_0000;
pub const IMAGE_SCN_MEM_PRELOAD: u32 = 0x0008_0000;

pub const IMAGE_SCN_ALIGN_1BYTES: u32 = 0x0010_0000;
pub const IMAGE_SCN_ALIGN_2BYTES: u32 = 0x0020_0000;
pub const IMAGE_SCN_ALIGN_4BYTES: u32 = 0x0030_0000;
pub const IMAGE_SCN_ALIGN_8BYTES: u32 = 0x0040_0000;
pub const IMAGE_SCN_ALIGN_16BYTES: u32 = 0x0050_0000;
pub const IMAGE_SCN_ALIGN_32BYTES: u32 = 0x0060_0000;
pub const IMAGE_SCN_ALIGN_64BYTES: u32 = 0x0070_0000;
pub const IMAGE_SCN_ALIGN_128BYTES: u32 = 0x0080_0000;
pub const IMAGE_SCN_ALIGN_256BYTES: u32 = 0x0090_0000;
pub const IMAGE_SCN_ALIGN_512BYTES: u32 = 0x00A0_0000;
pub const IMAGE_SCN_ALIGN_1024BYTES: u32 = 0x00B0_0000;
pub const IMAGE_SCN_ALIGN_2048BYTES: u32 = 0x00C0_0000;
pub const IMAGE_SCN_ALIGN_4096BYTES: u32 = 0x00D0_0000;
pub const IMAGE_SCN_ALIGN_8192BYTES: u32 = 0x00E0_0000;
pub const IMAGE_SCN_ALIGN_MASK: u32 = 0x00F0_0000;

/// The section contains extended relocations.
pub const IMAGE_SCN_LNK_NRELOC_OVFL: u32 = 0x0100_0000;
/// The section can be discarded as needed.
pub const IMAGE_SCN_MEM_DISCARDABLE: u32 = 0x0200_0000;
/// The section cannot be cached.
pub const IMAGE_SCN_MEM_NOT_CACHED: u32 = 0x0400_0000;
/// The section is not pageable.
pub const IMAGE_SCN_MEM_NOT_PAGED: u32 = 0x0800_0000;
/// The section can be shared in memory.
pub const IMAGE_SCN_MEM_SHARED: u32 = 0x1000_0000;
/// The section can be executed as code.
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;
/// The section can be read.
pub const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;
/// The section can be written to.
pub const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

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

    #[test]
    fn set_name_offset() {
        let mut section = SectionTable::default();
        for (offset, name) in vec![
            (0, b"/0\0\0\0\0\0\0"),
            (1, b"/1\0\0\0\0\0\0"),
            (9_999_999, b"/9999999"),
            (10_000_000, b"//AAmJaA"),
            (0xfff_fff_fff, b"////////"),
        ] {
            section.set_name_offset(offset).unwrap();
            assert_eq!(&section.name, name);
            assert_eq!(section.name_offset().unwrap(), Some(offset));
        }
        assert!(section.set_name_offset(0x1_000_000_000).is_err());
    }
}
<commit_msg>pe: fix set_name_offset compilation for 32-bit (#163)<commit_after>use crate::alloc::string::{String, ToString};
use scroll::{ctx, Pread, Pwrite};
use crate::error::{self, Error};
use crate::pe::relocation;
use std::io::Write;

#[repr(C)]
#[derive(Debug, PartialEq, Clone, Default)]
pub struct SectionTable {
    pub name: [u8; 8],
    pub real_name: Option<String>,
    pub virtual_size: u32,
    pub virtual_address: u32,
    pub size_of_raw_data: u32,
    pub pointer_to_raw_data: u32,
    pub pointer_to_relocations: u32,
    pub pointer_to_linenumbers: u32,
    pub number_of_relocations: u16,
    pub number_of_linenumbers: u16,
    pub characteristics: u32,
}

pub const SIZEOF_SECTION_TABLE: usize = 8 * 5;

// Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L70
// Decodes a string table entry in base 64 (//AAAAAA). Expects string without
// prefixed slashes.
fn base64_decode_string_entry(s: &str) -> Result<usize, ()> {
    assert!(s.len() <= 6, "String too long, possible overflow.");

    let mut val = 0;
    for c in s.bytes() {
        let v = if b'A' <= c && c <= b'Z' { // 00..=25
            c - b'A'
        } else if b'a' <= c && c <= b'z' { // 26..=51
            c - b'a' + 26
        } else if b'0' <= c && c <= b'9' { // 52..=61
            c - b'0' + 52
        } else if c == b'+' { // 62
            62
        } else if c == b'/' { // 63
            63
        } else {
            return Err(())
        };
        val = val * 64 + v as usize;
    }
    Ok(val)
}

impl SectionTable {
    pub fn parse(bytes: &[u8], offset: &mut usize, string_table_offset: usize) -> error::Result<Self> {
        let mut table = SectionTable::default();
        let mut name = [0u8; 8];
        name.copy_from_slice(bytes.gread_with(offset, 8)?);

        table.name = name;
        table.virtual_size = bytes.gread_with(offset, scroll::LE)?;
        table.virtual_address = bytes.gread_with(offset, scroll::LE)?;
        table.size_of_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_raw_data = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.pointer_to_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_relocations = bytes.gread_with(offset, scroll::LE)?;
        table.number_of_linenumbers = bytes.gread_with(offset, scroll::LE)?;
        table.characteristics = bytes.gread_with(offset, scroll::LE)?;

        if let Some(idx) = table.name_offset()? {
            table.real_name = Some(bytes.pread::<&str>(string_table_offset + idx)?.to_string());
        }
        Ok(table)
    }

    pub fn name_offset(&self) -> error::Result<Option<usize>> {
        // Based on https://github.com/llvm-mirror/llvm/blob/af7b1832a03ab6486c42a40d21695b2c03b2d8a3/lib/Object/COFFObjectFile.cpp#L1054
        if self.name[0] == b'/' {
            let idx: usize = if self.name[1] == b'/' {
                let b64idx = self.name.pread::<&str>(2)?;
                base64_decode_string_entry(b64idx).map_err(|_|
                    Error::Malformed(format!("Invalid indirect section name //{}: base64 decoding failed", b64idx)))?
            } else {
                let name = self.name.pread::<&str>(1)?;
                name.parse().map_err(|err|
                    Error::Malformed(format!("Invalid indirect section name /{}: {}", name, err)))?
            };
            Ok(Some(idx))
        } else {
            Ok(None)
        }
    }

    pub fn set_name_offset(&mut self, mut idx: usize) -> error::Result<()> {
        if idx <= 9_999_999 { // 10^7 - 1
            self.name = [0; 8];
            self.name[0] = b'/';
            write!(&mut self.name[1..], "{}", idx).unwrap();
            Ok(())
        } else if idx as u64 <= 0xfff_fff_fff { // 64^6 - 1
            self.name[0] = b'/';
            self.name[1] = b'/';
            for i in 0..6 {
                let rem = (idx % 64) as u8;
                idx = idx / 64;
                let c = match rem {
                    0..=25 => b'A' + rem,
                    26..=51 => b'a' + rem - 26,
                    52..=61 => b'0' + rem - 52,
                    62 => b'+',
                    63 => b'/',
                    _ => unreachable!(),
                };
                self.name[7 - i] = c;
            }
            Ok(())
        } else {
            Err(Error::Malformed(format!("Invalid section name offset: {}", idx)))
        }
    }

    pub fn name(&self) -> error::Result<&str> {
        match self.real_name.as_ref() {
            Some(s) => Ok(s),
            None => Ok(self.name.pread(0)?)
        }
    }

    pub fn relocations<'a>(&self, bytes: &'a[u8]) -> error::Result<relocation::Relocations<'a>> {
        let offset = self.pointer_to_relocations as usize;
        let number = self.number_of_relocations as usize;
        relocation::Relocations::parse(bytes, offset, number)
    }
}

impl ctx::SizeWith<scroll::Endian> for SectionTable {
    type Units = usize;
    fn size_with(_ctx: &scroll::Endian) -> usize {
        SIZEOF_SECTION_TABLE
    }
}

impl ctx::TryIntoCtx<scroll::Endian> for SectionTable {
    type Error = error::Error;
    type Size = usize;
    fn try_into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) -> Result<Self::Size, Self::Error> {
        let offset = &mut 0;
        bytes.gwrite(&self.name[..], offset)?;
        bytes.gwrite_with(self.virtual_size, offset, ctx)?;
        bytes.gwrite_with(self.virtual_address, offset, ctx)?;
        bytes.gwrite_with(self.size_of_raw_data, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_raw_data, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_relocations, offset, ctx)?;
        bytes.gwrite_with(self.pointer_to_linenumbers, offset, ctx)?;
        bytes.gwrite_with(self.number_of_relocations, offset, ctx)?;
        bytes.gwrite_with(self.number_of_linenumbers, offset, ctx)?;
        bytes.gwrite_with(self.characteristics, offset, ctx)?;
        Ok(SIZEOF_SECTION_TABLE)
    }
}

impl ctx::IntoCtx<scroll::Endian> for SectionTable {
    fn into_ctx(self, bytes: &mut [u8], ctx: scroll::Endian) {
        bytes.pwrite_with(self, 0, ctx).unwrap();
    }
}

/// The section should not be padded to the next boundary. This flag is obsolete and is replaced
/// by `IMAGE_SCN_ALIGN_1BYTES`. This is valid only for object files.
pub const IMAGE_SCN_TYPE_NO_PAD: u32 = 0x0000_0008;
/// The section contains executable code.
pub const IMAGE_SCN_CNT_CODE: u32 = 0x0000_0020;
/// The section contains initialized data.
pub const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x0000_0040;
///  The section contains uninitialized data.
pub const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x0000_0080;
pub const IMAGE_SCN_LNK_OTHER: u32 = 0x0000_0100;
/// The section contains comments or other information. The .drectve section has this type.
/// This is valid for object files only.
pub const IMAGE_SCN_LNK_INFO: u32 = 0x0000_0200;
/// The section will not become part of the image. This is valid only for object files.
pub const IMAGE_SCN_LNK_REMOVE: u32 = 0x0000_0800;
/// The section contains COMDAT data. This is valid only for object files.
pub const IMAGE_SCN_LNK_COMDAT: u32 = 0x0000_1000;
/// The section contains data referenced through the global pointer (GP).
pub const IMAGE_SCN_GPREL: u32 = 0x0000_8000;
pub const IMAGE_SCN_MEM_PURGEABLE: u32 = 0x0002_0000;
pub const IMAGE_SCN_MEM_16BIT: u32 = 0x0002_0000;
pub const IMAGE_SCN_MEM_LOCKED: u32 = 0x0004_0000;
pub const IMAGE_SCN_MEM_PRELOAD: u32 = 0x0008_0000;

pub const IMAGE_SCN_ALIGN_1BYTES: u32 = 0x0010_0000;
pub const IMAGE_SCN_ALIGN_2BYTES: u32 = 0x0020_0000;
pub const IMAGE_SCN_ALIGN_4BYTES: u32 = 0x0030_0000;
pub const IMAGE_SCN_ALIGN_8BYTES: u32 = 0x0040_0000;
pub const IMAGE_SCN_ALIGN_16BYTES: u32 = 0x0050_0000;
pub const IMAGE_SCN_ALIGN_32BYTES: u32 = 0x0060_0000;
pub const IMAGE_SCN_ALIGN_64BYTES: u32 = 0x0070_0000;
pub const IMAGE_SCN_ALIGN_128BYTES: u32 = 0x0080_0000;
pub const IMAGE_SCN_ALIGN_256BYTES: u32 = 0x0090_0000;
pub const IMAGE_SCN_ALIGN_512BYTES: u32 = 0x00A0_0000;
pub const IMAGE_SCN_ALIGN_1024BYTES: u32 = 0x00B0_0000;
pub const IMAGE_SCN_ALIGN_2048BYTES: u32 = 0x00C0_0000;
pub const IMAGE_SCN_ALIGN_4096BYTES: u32 = 0x00D0_0000;
pub const IMAGE_SCN_ALIGN_8192BYTES: u32 = 0x00E0_0000;
pub const IMAGE_SCN_ALIGN_MASK: u32 = 0x00F0_0000;

/// The section contains extended relocations.
pub const IMAGE_SCN_LNK_NRELOC_OVFL: u32 = 0x0100_0000;
/// The section can be discarded as needed.
pub const IMAGE_SCN_MEM_DISCARDABLE: u32 = 0x0200_0000;
/// The section cannot be cached.
pub const IMAGE_SCN_MEM_NOT_CACHED: u32 = 0x0400_0000;
/// The section is not pageable.
pub const IMAGE_SCN_MEM_NOT_PAGED: u32 = 0x0800_0000;
/// The section can be shared in memory.
pub const IMAGE_SCN_MEM_SHARED: u32 = 0x1000_0000;
/// The section can be executed as code.
pub const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;
/// The section can be read.
pub const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;
/// The section can be written to.
pub const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

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

    #[test]
    fn set_name_offset() {
        let mut section = SectionTable::default();
        for (offset, name) in vec![
            (0, b"/0\0\0\0\0\0\0"),
            (1, b"/1\0\0\0\0\0\0"),
            (9_999_999, b"/9999999"),
            (10_000_000, b"//AAmJaA"),
            #[cfg(target_pointer_width = "64")]
            (0xfff_fff_fff, b"////////"),
        ] {
            section.set_name_offset(offset).unwrap();
            assert_eq!(&section.name, name);
            assert_eq!(section.name_offset().unwrap(), Some(offset));
        }
        #[cfg(target_pointer_width = "64")]
        assert!(section.set_name_offset(0x1_000_000_000).is_err());
    }
}
<|endoftext|>


<commit_before>use byteread::*;

pub trait BitReadEndian {
    fn peek_val(&mut self, n:usize) -> u64;
    fn merge_val(msp:u64, lsp:u64, msb:usize, lsb:usize) -> u64;
    fn skip_rem(&mut self, n:usize) -> ();
}

pub trait BitReadFill {
    fn can_refill(&self) -> bool;
    fn fill32(&self) -> u64;
    fn fill64(&self) -> u64;
}

pub trait BitReadInternal : BitReadEndian + BitReadFill {
    #[inline]
    fn left(&self) -> usize;
    fn refill32(&mut self) -> ();
    fn refill64(&mut self) -> ();

    fn get_val(&mut self, n:usize) -> u64 {
        let ret = self.peek_val(n);

        self.skip_rem(n);

        return ret;
    }

}

pub trait BitRead<'a>: BitReadInternal+Copy {
    fn new(&'a[u8]) -> Self;
    fn consumed(&self) -> usize;
    fn available(&self) -> usize;

    fn skip_bits(&mut self, size : usize) -> ();

    #[inline]
    fn get_bit(&mut self) -> bool {
        if self.left() <= 0 {
            self.refill64();
        }

        self.get_val(1) != 0
    }

    #[inline]
    fn get_bits_64(&mut self, mut n:usize) -> u64 {
        let mut left = 0;
        let mut ret = 0;

        if n == 0 {
            return 0;
        }

        if self.left() < n {
            n   -= self.left();
            left = self.left();
            ret  = self.get_val(left);
            self.refill64();
        }

        Self::merge_val(self.get_val(n), ret, left, n)
    }

    #[inline]
    fn get_bits_32(&mut self, n:usize) -> u32 {
        if n == 0 {
            return 0;
        }

        if self.left() <= n {
            self.refill32();
        }

        return self.get_val(n) as u32;
    }

    #[inline]
    fn peek_bit(&mut self) -> bool {
        let mut tmp = *self;

        tmp.get_bit()
    }

    #[inline]
    fn peek_bits_32(&mut self, n:usize) -> u32 {
        let mut tmp = *self;

        tmp.get_bits_32(n)
    }

    #[inline]
    fn peek_bits_64(&self, n:usize) -> u64 {
        let mut tmp = *self;

        tmp.get_bits_64(n)
    }

    #[inline]
    fn align_bits(&mut self) -> () {
        let left = self.left() & 7;

        self.skip_bits(left);
    }
}

macro_rules! endian_reader {
    {$name: ident} => {
        #[derive(Debug, Clone, Copy)]
        pub struct $name<'a> {
            buffer : &'a[u8], /// read buffer, 8-bytes padded
            index : usize,
            cache : u64,
            left : usize,
        }
        impl <'a> BitReadInternal for $name<'a> {
            #[inline]
            fn left(&self) -> usize {
                self.left
            }
            #[inline]
            fn refill32(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }
                let val = self.fill32();

                self.cache  = Self::merge_val(val, self.cache,
                                              self.left, 32 - self.left);
                self.index += 4;
                self.left  += 32;
            }
            #[inline]
            fn refill64(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }

                self.cache  = self.fill64();
                self.index += 8;
                self.left   = 64;
            }
        }

        impl <'a> BitRead<'a> for $name<'a> {
            fn new(buffer: &'a[u8]) -> $name<'a> {
                let mut reader = $name {
                    buffer: buffer,
                    index: 0,
                    cache: 0,
                    left: 0
                };

                reader.refill64();
                return reader;
            }
            #[inline]
            fn consumed(&self) -> usize {
                self.index * 8 - self.left
            }

            #[inline]
            fn available(&self) -> usize {
                (self.buffer.len() - self.index) * 8 + self.left
            }

            #[inline]
            fn skip_bits(&mut self, mut n:usize) -> () {
                if self.left < n {
                    n -= self.left;
                    if n > 64 {
                        let skip = n / 8;

                        n -= skip * 8;
                        self.index += skip;
                    }
                    self.skip_rem(n);
                    self.refill64();
                }

                self.skip_rem(n);
            }

        }
    }
}

macro_rules! little_endian_reader {
    {$name: ident} => {
        endian_reader!{ $name }

        impl <'a> BitReadEndian for $name<'a> {
            #[inline]
            fn peek_val(&mut self, n:usize) -> u64 {
                let v = self.cache & ((1u64 << n) - 1);

                return v;
            }
            #[inline]
            fn skip_rem(&mut self, n:usize) -> () {
                self.cache = self.cache >> n;
                self.left = self.left.saturating_sub(n);
            }
            #[inline]
            fn merge_val(msp:u64, lsp:u64, msb:usize, _:usize) -> u64 {
                msp << msb | lsp
            }
        }
    }
}

little_endian_reader!{ BitReadLE }

impl <'a> BitReadFill for BitReadLE<'a> {
    #[inline]
    fn can_refill(&self) -> bool {
        self.index + 8 <= self.buffer.len()
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32l(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64l(&self.buffer[self.index..])
    }
}

macro_rules! big_endian_reader {
    {$name: ident} => {
        endian_reader!{ $name }

        impl <'a> BitReadEndian for $name<'a> {
            #[inline]
            fn peek_val(&mut self, n:usize) -> u64 {
                self.cache >> (64 - n)
            }
            #[inline]
            fn skip_rem(&mut self, n:usize) -> () {
                self.cache = self.cache << n;
                self.left = self.left.saturating_sub(n);
            }
            #[inline]
            fn merge_val(msp:u64, lsp:u64, _:usize, lsb:usize) -> u64 {
                msp | lsp << lsb
            }
        }
    }
}

big_endian_reader!{ BitReadBE }

impl <'a> BitReadFill for BitReadBE<'a> {
    #[inline]
    fn can_refill(&self) -> bool {
        self.index + 8 <= self.buffer.len()
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32b(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64b(&self.buffer[self.index..])
    }
}

#[cfg(test)]
mod test {
    pub const CHECKBOARD0101: [u8; 128] = [0b01010101; 128];
    pub const CHECKBOARD0011: [u8; 128] = [0b00110011; 128];

    mod le {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bit());
            assert!(!reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let mut reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 3);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
    mod be {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(!reader.get_bit());
            assert!(reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let b = &CHECKBOARD0101;
            let reader = BitReadBE::new(b);

            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
}
<commit_msg>bitstream: Export the bitreader-building macros<commit_after>use byteread::*;

pub trait BitReadEndian {
    fn peek_val(&mut self, n:usize) -> u64;
    fn merge_val(msp:u64, lsp:u64, msb:usize, lsb:usize) -> u64;
    fn skip_rem(&mut self, n:usize) -> ();
}

pub trait BitReadFill {
    fn can_refill(&self) -> bool;
    fn fill32(&self) -> u64;
    fn fill64(&self) -> u64;
}

pub trait BitReadInternal : BitReadEndian + BitReadFill {
    #[inline]
    fn left(&self) -> usize;
    fn refill32(&mut self) -> ();
    fn refill64(&mut self) -> ();

    fn get_val(&mut self, n:usize) -> u64 {
        let ret = self.peek_val(n);

        self.skip_rem(n);

        return ret;
    }

}

pub trait BitRead<'a>: BitReadInternal+Copy {
    fn new(&'a[u8]) -> Self;
    fn consumed(&self) -> usize;
    fn available(&self) -> usize;

    fn skip_bits(&mut self, size : usize) -> ();

    #[inline]
    fn get_bit(&mut self) -> bool {
        if self.left() <= 0 {
            self.refill64();
        }

        self.get_val(1) != 0
    }

    #[inline]
    fn get_bits_64(&mut self, mut n:usize) -> u64 {
        let mut left = 0;
        let mut ret = 0;

        if n == 0 {
            return 0;
        }

        if self.left() < n {
            n   -= self.left();
            left = self.left();
            ret  = self.get_val(left);
            self.refill64();
        }

        Self::merge_val(self.get_val(n), ret, left, n)
    }

    #[inline]
    fn get_bits_32(&mut self, n:usize) -> u32 {
        if n == 0 {
            return 0;
        }

        if self.left() <= n {
            self.refill32();
        }

        return self.get_val(n) as u32;
    }

    #[inline]
    fn peek_bit(&mut self) -> bool {
        let mut tmp = *self;

        tmp.get_bit()
    }

    #[inline]
    fn peek_bits_32(&mut self, n:usize) -> u32 {
        let mut tmp = *self;

        tmp.get_bits_32(n)
    }

    #[inline]
    fn peek_bits_64(&self, n:usize) -> u64 {
        let mut tmp = *self;

        tmp.get_bits_64(n)
    }

    #[inline]
    fn align_bits(&mut self) -> () {
        let left = self.left() & 7;

        self.skip_bits(left);
    }
}

#[macro_export]
macro_rules! endian_reader {
    {$name: ident} => {
        #[derive(Debug, Clone, Copy)]
        pub struct $name<'a> {
            buffer : &'a[u8], /// read buffer, 8-bytes padded
            index : usize,
            cache : u64,
            left : usize,
        }
        impl <'a> BitReadInternal for $name<'a> {
            #[inline]
            fn left(&self) -> usize {
                self.left
            }
            #[inline]
            fn refill32(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }
                let val = self.fill32();

                self.cache  = Self::merge_val(val, self.cache,
                                              self.left, 32 - self.left);
                self.index += 4;
                self.left  += 32;
            }
            #[inline]
            fn refill64(&mut self) -> () {
                if !self.can_refill() {
                    return;
                }

                self.cache  = self.fill64();
                self.index += 8;
                self.left   = 64;
            }
        }

        impl <'a> BitRead<'a> for $name<'a> {
            fn new(buffer: &'a[u8]) -> $name<'a> {
                let mut reader = $name {
                    buffer: buffer,
                    index: 0,
                    cache: 0,
                    left: 0
                };

                reader.refill64();
                return reader;
            }
            #[inline]
            fn consumed(&self) -> usize {
                self.index * 8 - self.left
            }

            #[inline]
            fn available(&self) -> usize {
                (self.buffer.len() - self.index) * 8 + self.left
            }

            #[inline]
            fn skip_bits(&mut self, mut n:usize) -> () {
                if self.left < n {
                    n -= self.left;
                    if n > 64 {
                        let skip = n / 8;

                        n -= skip * 8;
                        self.index += skip;
                    }
                    self.skip_rem(n);
                    self.refill64();
                }

                self.skip_rem(n);
            }

        }
    }
}

#[macro_export]
macro_rules! little_endian_reader {
    {$name: ident} => {
        endian_reader!{ $name }

        impl <'a> BitReadEndian for $name<'a> {
            #[inline]
            fn peek_val(&mut self, n:usize) -> u64 {
                let v = self.cache & ((1u64 << n) - 1);

                return v;
            }
            #[inline]
            fn skip_rem(&mut self, n:usize) -> () {
                self.cache = self.cache >> n;
                self.left = self.left.saturating_sub(n);
            }
            #[inline]
            fn merge_val(msp:u64, lsp:u64, msb:usize, _:usize) -> u64 {
                msp << msb | lsp
            }
        }
    }
}

little_endian_reader!{ BitReadLE }

impl <'a> BitReadFill for BitReadLE<'a> {
    #[inline]
    fn can_refill(&self) -> bool {
        self.index + 8 <= self.buffer.len()
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32l(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64l(&self.buffer[self.index..])
    }
}

#[macro_export]
macro_rules! big_endian_reader {
    {$name: ident} => {
        endian_reader!{ $name }

        impl <'a> BitReadEndian for $name<'a> {
            #[inline]
            fn peek_val(&mut self, n:usize) -> u64 {
                self.cache >> (64 - n)
            }
            #[inline]
            fn skip_rem(&mut self, n:usize) -> () {
                self.cache = self.cache << n;
                self.left = self.left.saturating_sub(n);
            }
            #[inline]
            fn merge_val(msp:u64, lsp:u64, _:usize, lsb:usize) -> u64 {
                msp | lsp << lsb
            }
        }
    }
}

big_endian_reader!{ BitReadBE }

impl <'a> BitReadFill for BitReadBE<'a> {
    #[inline]
    fn can_refill(&self) -> bool {
        self.index + 8 <= self.buffer.len()
    }
    #[inline(always)]
    fn fill32(&self) -> u64 {
        get_u32b(&self.buffer[self.index..]) as u64
    }
    #[inline(always)]
    fn fill64(&self) -> u64 {
        get_u64b(&self.buffer[self.index..])
    }
}

#[cfg(test)]
mod test {
    pub const CHECKBOARD0101: [u8; 128] = [0b01010101; 128];
    pub const CHECKBOARD0011: [u8; 128] = [0b00110011; 128];

    mod le {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bit());
            assert!(!reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(1) == 1);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let mut reader = BitReadLE {
                buffer: &CHECKBOARD0101,
                index: 0,
                cache: 0,
                left: 0
            };

            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(1) == 1);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 1);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 3);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadLE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
    mod be {
        use super::super::*;
        use super::*;

        #[test]
        fn get_bit() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(!reader.get_bit());
            assert!(reader.get_bit());
        }

        #[test]
        fn get_bits_64() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_64(1) == 0);
            assert!(reader.get_bits_64(2) == 2);
            assert!(reader.get_bits_64(4) == 10);
            assert!(reader.get_bits_64(1) == 1);
            assert!(reader.get_bits_64(8) == 85);
        }

        #[test]
        fn peek_bits_64() {
            let b = &CHECKBOARD0101;
            let reader = BitReadBE::new(b);

            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(1) == 0);
            assert!(reader.peek_bits_64(2) == 1);
            assert!(reader.peek_bits_64(2) == 1);
        }

        #[test]
        fn get_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.get_bits_32(1) == 0);
            assert!(reader.get_bits_32(2) == 2);
            assert!(reader.get_bits_32(4) == 10);
            assert!(reader.get_bits_32(1) == 1);
            assert!(reader.get_bits_32(8) == 85);
        }

        #[test]
        fn peek_bits_32() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(1) == 0);
            assert!(reader.peek_bits_32(2) == 1);
            assert!(reader.peek_bits_32(2) == 1);
        }

        #[test]
        fn skip_bits() {
            let b = &CHECKBOARD0101;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(0);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
            reader.skip_bits(2);
            assert!(reader.peek_bits_32(1) == 0);
        }

        #[test]
        fn align_bits() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.align_bits();
            assert!(reader.get_bits_64(3) == 1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
            reader.skip_bits(1);
            reader.align_bits();
            assert!(reader.get_bits_64(4) == 3);
        }

        #[test]
        fn overread() {
            let b = &CHECKBOARD0011;
            let mut reader = BitReadBE::new(b);

            reader.skip_bits(128*8+2);
            reader.get_bits_64(6);
        }
    }
}
<|endoftext|>