<commit_before>use std::rand::Rng;
use calx::dijkstra::Dijkstra;
use world;
use location::{Location};
use dir6::Dir6;
use flags;
use mob;
use mob::{Status, Intrinsic};
use components::{Kind};
use geom::HexGeom;
use spatial::Place;
use action;
use fov::Fov;
use rng;
use msg;

/// Game object handle.
#[deriving(Copy, PartialEq, Eq, Clone, Hash, Show, Decodable, Encodable)]
pub struct Entity(pub uint);

impl Entity {
    /// Place the entity in a location in the game world.
    pub fn place(self, loc: Location) {
        world::with_mut(|w| w.spatial.insert_at(self, loc));
        self.on_move_to(loc);
    }

    /// Remove the entity from all the game world systems. THE ENTITY VALUE
    /// WILL BE INVALID FROM HERE ON AND USING IT WILL LEAD TO BUGS. THE
    /// CALLER IS RESPONSIBLE FOR ENSUING THAT AN ENTITY WILL NOT BE
    /// USED FROM ANYWHERE AFTER THE DELETE OPERATION.
    pub fn delete(self) {
        // COMPONENTS CHECKPOINT
        // This needs to call every component system.
        world::with_mut(|w| w.descs_mut().remove(self));
        world::with_mut(|w| w.kinds_mut().remove(self));
        world::with_mut(|w| w.map_memories_mut().remove(self));
        world::with_mut(|w| w.mobs_mut().remove(self));
        world::with_mut(|w| w.spawns_mut().remove(self));

        world::with_mut(|w| w.spatial.remove(self));

        world::with_mut(|w| w.ecs.delete(self));
    }

    pub fn blocks_walk(self) -> bool { self.is_mob() }

    /// Return the kind of the entity.
    pub fn kind(self) -> Kind {
        // XXX: Will crash if an entity has no kind specified.
        world::with(|w| *w.kinds().get(self).unwrap())
    }

    pub fn name(self) -> String {
        world::with(|w|
            match w.descs().get(self) {
                Some(desc) => desc.name.clone(),
                None => "".to_string()
            }
        )
    }

    /// Create a cloned entity that uses the current entity as a parent
    /// prototype. Components not defined in the clone entity will be read
    /// from the current entity.
    pub fn clone_at(self, loc: Location) -> Entity {
        world::with_mut(|w| {
            let e = w.ecs.new_entity(Some(self));
            w.spatial.insert_at(e, loc);
            e
        })
    }

// Spatial methods /////////////////////////////////////////////////////

    pub fn can_enter(self, loc: Location) -> bool {
        if self.is_mob() && loc.has_mobs() { return false; }
        if loc.terrain().is_door() && !self.has_intrinsic(Intrinsic::Hands) {
            // Can't open doors without hands.
            return false;
        }
        if loc.blocks_walk() { return false; }
        true
    }

    /// Return whether the entity can move in a direction.
    pub fn can_step(self, dir: Dir6) -> bool {
        let place = world::with(|w| w.spatial.get(self));
        if let Some(Place::At(loc)) = place {
            let new_loc = loc + dir.to_v2();
            return self.can_enter(new_loc);
        }
        return false;
    }

    /// Try to move the entity in direction.
    pub fn step(self, dir: Dir6) {
        let place = world::with(|w| w.spatial.get(self));
        if let Some(Place::At(loc)) = place {
            let new_loc = loc + dir.to_v2();
            if self.can_enter(new_loc) {
                world::with_mut(|w| w.spatial.insert_at(self, new_loc));
                self.on_move_to(new_loc);
            }
        }
    }

    pub fn location(self) -> Option<Location> {
        match world::with(|w| w.spatial.get(self)) {
            Some(Place::At(loc)) => Some(loc),
            Some(Place::In(e)) => e.location(),
            _ => None
        }
    }

    pub fn distance_from(self, other: Entity) -> Option<int> {
        if let (Some(loc1), Some(loc2)) = (self.location(), other.location()) {
            loc1.distance_from(loc2)
        } else {
            None
        }
    }

// Damage and lifetime /////////////////////////////////////////////////

    pub fn damage(self, amount: int) {
        if amount == 0 { return; }

        let (_amount, kill) = world::with_mut(|w| {
            if let Some(ref mut mob) = w.mobs_mut().get(self) {
                mob.hp -= amount;
                if mob.hp <= 0 {
                    // Remember this when we're out of the borrow.
                    return (amount, true);
                }
            } else {
                panic!("Damaging a non-mob");
            };
            (amount, false)
        });

        msg::push_msg(::Msg::Damage(self));
        if kill {
            self.kill();
        }
    }

    /// Do any game logic stuff related to this entity dying violently before
    /// deleting it.
    pub fn kill(self) {
        msgln!("{} dies.", self.name());
        msg::push_msg(::Msg::Gib(self.location().unwrap()));
        self.delete();
    }

// Mob methods /////////////////////////////////////////////////////////

    pub fn is_mob(self) -> bool {
        world::with(|w| {
            if let Some(&Kind::Mob(_)) = w.kinds().get(self) {
                return true;
            }
            return false;
        })
    }

    pub fn mob_spec(self) -> Option<&'static mob::MobSpec> {
        world::with(|w| {
            if let Some(&Kind::Mob(mt)) = w.kinds().get(self) {
                Some(&mob::MOB_SPECS[mt as uint])
            } else {
                None
            }
        })
    }

    /// Return whether this mob is the player avatar.
    pub fn is_player(self) -> bool {
        self.name() == "Player" && self.location().is_some()
    }

    pub fn has_status(self, status: Status) -> bool {
        world::with(|w| {
            if let Some(&mob) = w.mobs().get(self) {
                return mob.has_status(status);
            }
            return false;
        })
    }

    pub fn add_status(self, status: Status) {
        if !self.is_mob() { return; }
        world::with_mut(|w| w.mobs_mut().get(self).unwrap().add_status(status));
        assert!(self.has_status(status));
    }

    pub fn remove_status(self, status: Status) {
        if !self.is_mob() { return; }
        world::with_mut(|w| w.mobs_mut().get(self).unwrap().remove_status(status));
        assert!(!self.has_status(status));
    }

    pub fn has_intrinsic(self, intrinsic: Intrinsic) -> bool {
        if let Some(spec) = self.mob_spec() {
            return spec.intrinsics & intrinsic as int != 0;
        }
        return false;
    }

    /// Return whether this entity is an awake mob.
    pub fn is_active(self) -> bool {
        self.is_mob() && !self.has_status(Status::Asleep)
    }

    /// Return if the entity is a mob that should get an update this frame
    /// based on its speed properties. Does not check for status effects like
    /// sleep that might prevent actual action.
    pub fn ticks_this_frame(self) -> bool {
        if !self.is_mob() { return false; }

        let tick = flags::get_tick();
        // Go through a cycle of 5 phases to get 4 possible speeds.
        // System idea from Jeff Lait.
        let phase = tick % 5;
        match phase {
            0 => return true,
            1 => return self.has_intrinsic(Intrinsic::Fast),
            2 => return true,
            3 => return self.has_status(Status::Quick),
            4 => return !self.has_intrinsic(Intrinsic::Slow)
                        && !self.has_status(Status::Slow),
            _ => panic!("Invalid action phase"),
        }
    }

    /// Return whether the entity is a mob that will act this frame.
    pub fn acts_this_frame(self) -> bool {
        if !self.is_active() { return false; }
        return self.ticks_this_frame();
    }

    /// Return whether the entity is an awake non-player mob and should be
    /// animated with a bob.
    pub fn is_bobbing(self) -> bool {
        self.is_active() && !self.is_player()
    }

    pub fn melee(self, dir: Dir6) {
        let loc = self.location().unwrap() + dir.to_v2();
        if let Some(e) = loc.mob_at() {
            // Attack power
            let p = world::with(|w| w.mobs().get(self).unwrap().power);
            if p == 0 {
                // No fight capacity.
                return;
            }

            // Every five points of power is one certain hit.
            let full = p / 5;
            // The fractional points are one probabilistic hit.
            let partial = (p % 5) as f64 / 5.0;

            let damage = full + if rng::p(partial) { 1 } else { 0 };
            e.damage(damage)
        }
    }

    pub fn hp(self) -> int {
        world::with(|w| w.mobs().get(self).unwrap().hp)
    }

    pub fn max_hp(self) -> int {
        self.mob_spec().unwrap().power
    }

// AI methods /////////////////////////////////////////////////////////

    /// Top-level method called each frame to update the entity.
    pub fn update(self) {
        if self.is_mob() && !self.is_player() && self.ticks_this_frame() {
            self.mob_ai();
        }
    }

    /// AI routine for autonomous mobs.
    fn mob_ai(self) {
        assert!(self.is_mob());
        assert!(!self.is_player());
        assert!(self.ticks_this_frame());

        if self.has_status(Status::Asleep) {
            if let Some(p) = action::player() {
                // TODO: Line-of-sight, stealth concerns, other enemies than
                // player etc.
                if let Some(d) = p.distance_from(self) {
                    if d < 6 {
                        self.remove_status(Status::Asleep);
                    }
                }
            }

            return;
        }

        if let Some(p) = action::player() {
            let loc = self.location().unwrap();

            let vec_to_enemy = loc.v2_at(p.location().unwrap());
            if let Some(v) = vec_to_enemy {
                if v.hex_dist() == 1 {
                    // Melee range, hit.
                    self.melee(Dir6::from_v2(v));
                } else {
                    // Walk towards.
                    let pathing_depth = 16;
                    let pathing = Dijkstra::new(
                        vec![p.location().unwrap()], |&loc| !loc.blocks_walk(),
                        pathing_depth);

                    let steps = pathing.sorted_neighbors(&loc);
                    if steps.len() > 0 {
                        self.step(loc.dir6_towards(steps[0]).unwrap());
                    } else {
                        self.step(rng::with(|ref mut rng| rng.gen::<Dir6>()));
                        // TODO: Fall asleep if things get boring.
                    }
                }
            }
        }
    }

    /// Return whether this thing wants to fight the other thing.
    fn is_hostile_to(self, other: Entity) -> bool {
        if let Some(p) = action::player() {
            if self.is_mob() && other.is_mob() && self != other &&
                (self == p || other == p) {
                return true;
            }
        }
        return false;
    }

    /// Return whether the mob has hostiles in its immediate vicinity.
    pub fn is_threatened(self) -> bool {
        // XXX: Expensive.
        let range = 12u;
        let loc = self.location().unwrap();
        let seen: Vec<Location> = Fov::new(
            |pt| (loc + pt).blocks_sight(), range)
            .map(|pt| loc + pt)
            .collect();
        for loc in seen.iter() {
            if let Some(m) = loc.mob_at() {
                if m.is_hostile_to(self) {
                    return true;
                }
            }
        }
        false
    }

// Callbacks ///////////////////////////////////////////////////////////

    /// Called after the entity is moved to a new location.
    pub fn on_move_to(self, loc: Location) {
        self.do_fov();

        if self.is_player() {
            if loc.terrain().is_exit() {
                action::next_level();
            }
        }
    }

// FOV and map memory //////////////////////////////////////////////////

    fn has_map_memory(self) -> bool {
        world::with(|w| w.map_memories().get(self).is_some())
    }

    fn do_fov(self) {
        let range = 12u;
        if let Some(loc) = self.location() {
            if self.has_map_memory() {
                let seen: Vec<Location> = Fov::new(
                    |pt| (loc + pt).blocks_sight(), range)
                    .map(|pt| loc + pt)
                    .collect();
                world::with_mut(|w| {
                    if let Some(ref mut mm) = w.map_memories_mut().get(self) {
                        mm.seen.clear();
                        mm.seen.extend(seen.clone().into_iter());
                        mm.remembered.extend(seen.iter().map(|&x| x));
                    } else {
                        panic!("Couldn't bind map memory");
                    }
                });
            }
        }
    }

    pub fn forget_map(self) {
        if self.has_map_memory() {
            world::with_mut(|w| {
                if let Some(ref mut mm) = w.map_memories_mut().get(self) {
                    mm.seen.clear();
                    mm.remembered.clear();
                } else {
                    panic!("Couldn't bind map memory");
                }
            });
        }
    }
}
<commit_msg>Clone entity uses place method<commit_after>use std::rand::Rng;
use calx::dijkstra::Dijkstra;
use world;
use location::{Location};
use dir6::Dir6;
use flags;
use mob;
use mob::{Status, Intrinsic};
use components::{Kind};
use geom::HexGeom;
use spatial::Place;
use action;
use fov::Fov;
use rng;
use msg;

/// Game object handle.
#[deriving(Copy, PartialEq, Eq, Clone, Hash, Show, Decodable, Encodable)]
pub struct Entity(pub uint);

impl Entity {
    /// Place the entity in a location in the game world.
    pub fn place(self, loc: Location) {
        world::with_mut(|w| w.spatial.insert_at(self, loc));
        self.on_move_to(loc);
    }

    /// Remove the entity from all the game world systems. THE ENTITY VALUE
    /// WILL BE INVALID FROM HERE ON AND USING IT WILL LEAD TO BUGS. THE
    /// CALLER IS RESPONSIBLE FOR ENSUING THAT AN ENTITY WILL NOT BE
    /// USED FROM ANYWHERE AFTER THE DELETE OPERATION.
    pub fn delete(self) {
        // COMPONENTS CHECKPOINT
        // This needs to call every component system.
        world::with_mut(|w| w.descs_mut().remove(self));
        world::with_mut(|w| w.kinds_mut().remove(self));
        world::with_mut(|w| w.map_memories_mut().remove(self));
        world::with_mut(|w| w.mobs_mut().remove(self));
        world::with_mut(|w| w.spawns_mut().remove(self));

        world::with_mut(|w| w.spatial.remove(self));

        world::with_mut(|w| w.ecs.delete(self));
    }

    pub fn blocks_walk(self) -> bool { self.is_mob() }

    /// Return the kind of the entity.
    pub fn kind(self) -> Kind {
        // XXX: Will crash if an entity has no kind specified.
        world::with(|w| *w.kinds().get(self).unwrap())
    }

    pub fn name(self) -> String {
        world::with(|w|
            match w.descs().get(self) {
                Some(desc) => desc.name.clone(),
                None => "".to_string()
            }
        )
    }

    /// Create a cloned entity that uses the current entity as a parent
    /// prototype. Components not defined in the clone entity will be read
    /// from the current entity.
    pub fn clone_at(self, loc: Location) -> Entity {
        let ret = world::with_mut(|w| { w.ecs.new_entity(Some(self)) });
        ret.place(loc);
        ret
    }

// Spatial methods /////////////////////////////////////////////////////

    pub fn can_enter(self, loc: Location) -> bool {
        if self.is_mob() && loc.has_mobs() { return false; }
        if loc.terrain().is_door() && !self.has_intrinsic(Intrinsic::Hands) {
            // Can't open doors without hands.
            return false;
        }
        if loc.blocks_walk() { return false; }
        true
    }

    /// Return whether the entity can move in a direction.
    pub fn can_step(self, dir: Dir6) -> bool {
        let place = world::with(|w| w.spatial.get(self));
        if let Some(Place::At(loc)) = place {
            let new_loc = loc + dir.to_v2();
            return self.can_enter(new_loc);
        }
        return false;
    }

    /// Try to move the entity in direction.
    pub fn step(self, dir: Dir6) {
        let place = world::with(|w| w.spatial.get(self));
        if let Some(Place::At(loc)) = place {
            let new_loc = loc + dir.to_v2();
            if self.can_enter(new_loc) {
                world::with_mut(|w| w.spatial.insert_at(self, new_loc));
                self.on_move_to(new_loc);
            }
        }
    }

    pub fn location(self) -> Option<Location> {
        match world::with(|w| w.spatial.get(self)) {
            Some(Place::At(loc)) => Some(loc),
            Some(Place::In(e)) => e.location(),
            _ => None
        }
    }

    pub fn distance_from(self, other: Entity) -> Option<int> {
        if let (Some(loc1), Some(loc2)) = (self.location(), other.location()) {
            loc1.distance_from(loc2)
        } else {
            None
        }
    }

// Damage and lifetime /////////////////////////////////////////////////

    pub fn damage(self, amount: int) {
        if amount == 0 { return; }

        let (_amount, kill) = world::with_mut(|w| {
            if let Some(ref mut mob) = w.mobs_mut().get(self) {
                mob.hp -= amount;
                if mob.hp <= 0 {
                    // Remember this when we're out of the borrow.
                    return (amount, true);
                }
            } else {
                panic!("Damaging a non-mob");
            };
            (amount, false)
        });

        msg::push_msg(::Msg::Damage(self));
        if kill {
            self.kill();
        }
    }

    /// Do any game logic stuff related to this entity dying violently before
    /// deleting it.
    pub fn kill(self) {
        msgln!("{} dies.", self.name());
        msg::push_msg(::Msg::Gib(self.location().unwrap()));
        self.delete();
    }

// Mob methods /////////////////////////////////////////////////////////

    pub fn is_mob(self) -> bool {
        world::with(|w| {
            if let Some(&Kind::Mob(_)) = w.kinds().get(self) {
                return true;
            }
            return false;
        })
    }

    pub fn mob_spec(self) -> Option<&'static mob::MobSpec> {
        world::with(|w| {
            if let Some(&Kind::Mob(mt)) = w.kinds().get(self) {
                Some(&mob::MOB_SPECS[mt as uint])
            } else {
                None
            }
        })
    }

    /// Return whether this mob is the player avatar.
    pub fn is_player(self) -> bool {
        self.name() == "Player" && self.location().is_some()
    }

    pub fn has_status(self, status: Status) -> bool {
        world::with(|w| {
            if let Some(&mob) = w.mobs().get(self) {
                return mob.has_status(status);
            }
            return false;
        })
    }

    pub fn add_status(self, status: Status) {
        if !self.is_mob() { return; }
        world::with_mut(|w| w.mobs_mut().get(self).unwrap().add_status(status));
        assert!(self.has_status(status));
    }

    pub fn remove_status(self, status: Status) {
        if !self.is_mob() { return; }
        world::with_mut(|w| w.mobs_mut().get(self).unwrap().remove_status(status));
        assert!(!self.has_status(status));
    }

    pub fn has_intrinsic(self, intrinsic: Intrinsic) -> bool {
        if let Some(spec) = self.mob_spec() {
            return spec.intrinsics & intrinsic as int != 0;
        }
        return false;
    }

    /// Return whether this entity is an awake mob.
    pub fn is_active(self) -> bool {
        self.is_mob() && !self.has_status(Status::Asleep)
    }

    /// Return if the entity is a mob that should get an update this frame
    /// based on its speed properties. Does not check for status effects like
    /// sleep that might prevent actual action.
    pub fn ticks_this_frame(self) -> bool {
        if !self.is_mob() { return false; }

        let tick = flags::get_tick();
        // Go through a cycle of 5 phases to get 4 possible speeds.
        // System idea from Jeff Lait.
        let phase = tick % 5;
        match phase {
            0 => return true,
            1 => return self.has_intrinsic(Intrinsic::Fast),
            2 => return true,
            3 => return self.has_status(Status::Quick),
            4 => return !self.has_intrinsic(Intrinsic::Slow)
                        && !self.has_status(Status::Slow),
            _ => panic!("Invalid action phase"),
        }
    }

    /// Return whether the entity is a mob that will act this frame.
    pub fn acts_this_frame(self) -> bool {
        if !self.is_active() { return false; }
        return self.ticks_this_frame();
    }

    /// Return whether the entity is an awake non-player mob and should be
    /// animated with a bob.
    pub fn is_bobbing(self) -> bool {
        self.is_active() && !self.is_player()
    }

    pub fn melee(self, dir: Dir6) {
        let loc = self.location().unwrap() + dir.to_v2();
        if let Some(e) = loc.mob_at() {
            // Attack power
            let p = world::with(|w| w.mobs().get(self).unwrap().power);
            if p == 0 {
                // No fight capacity.
                return;
            }

            // Every five points of power is one certain hit.
            let full = p / 5;
            // The fractional points are one probabilistic hit.
            let partial = (p % 5) as f64 / 5.0;

            let damage = full + if rng::p(partial) { 1 } else { 0 };
            e.damage(damage)
        }
    }

    pub fn hp(self) -> int {
        world::with(|w| w.mobs().get(self).unwrap().hp)
    }

    pub fn max_hp(self) -> int {
        self.mob_spec().unwrap().power
    }

// AI methods /////////////////////////////////////////////////////////

    /// Top-level method called each frame to update the entity.
    pub fn update(self) {
        if self.is_mob() && !self.is_player() && self.ticks_this_frame() {
            self.mob_ai();
        }
    }

    /// AI routine for autonomous mobs.
    fn mob_ai(self) {
        assert!(self.is_mob());
        assert!(!self.is_player());
        assert!(self.ticks_this_frame());

        if self.has_status(Status::Asleep) {
            if let Some(p) = action::player() {
                // TODO: Line-of-sight, stealth concerns, other enemies than
                // player etc.
                if let Some(d) = p.distance_from(self) {
                    if d < 6 {
                        self.remove_status(Status::Asleep);
                    }
                }
            }

            return;
        }

        if let Some(p) = action::player() {
            let loc = self.location().unwrap();

            let vec_to_enemy = loc.v2_at(p.location().unwrap());
            if let Some(v) = vec_to_enemy {
                if v.hex_dist() == 1 {
                    // Melee range, hit.
                    self.melee(Dir6::from_v2(v));
                } else {
                    // Walk towards.
                    let pathing_depth = 16;
                    let pathing = Dijkstra::new(
                        vec![p.location().unwrap()], |&loc| !loc.blocks_walk(),
                        pathing_depth);

                    let steps = pathing.sorted_neighbors(&loc);
                    if steps.len() > 0 {
                        self.step(loc.dir6_towards(steps[0]).unwrap());
                    } else {
                        self.step(rng::with(|ref mut rng| rng.gen::<Dir6>()));
                        // TODO: Fall asleep if things get boring.
                    }
                }
            }
        }
    }

    /// Return whether this thing wants to fight the other thing.
    fn is_hostile_to(self, other: Entity) -> bool {
        if let Some(p) = action::player() {
            if self.is_mob() && other.is_mob() && self != other &&
                (self == p || other == p) {
                return true;
            }
        }
        return false;
    }

    /// Return whether the mob has hostiles in its immediate vicinity.
    pub fn is_threatened(self) -> bool {
        // XXX: Expensive.
        let range = 12u;
        let loc = self.location().unwrap();
        let seen: Vec<Location> = Fov::new(
            |pt| (loc + pt).blocks_sight(), range)
            .map(|pt| loc + pt)
            .collect();
        for loc in seen.iter() {
            if let Some(m) = loc.mob_at() {
                if m.is_hostile_to(self) {
                    return true;
                }
            }
        }
        false
    }

// Callbacks ///////////////////////////////////////////////////////////

    /// Called after the entity is moved to a new location.
    pub fn on_move_to(self, loc: Location) {
        self.do_fov();

        if self.is_player() {
            if loc.terrain().is_exit() {
                action::next_level();
            }
        }
    }

// FOV and map memory //////////////////////////////////////////////////

    fn has_map_memory(self) -> bool {
        world::with(|w| w.map_memories().get(self).is_some())
    }

    fn do_fov(self) {
        let range = 12u;
        if let Some(loc) = self.location() {
            if self.has_map_memory() {
                let seen: Vec<Location> = Fov::new(
                    |pt| (loc + pt).blocks_sight(), range)
                    .map(|pt| loc + pt)
                    .collect();
                world::with_mut(|w| {
                    if let Some(ref mut mm) = w.map_memories_mut().get(self) {
                        mm.seen.clear();
                        mm.seen.extend(seen.clone().into_iter());
                        mm.remembered.extend(seen.iter().map(|&x| x));
                    } else {
                        panic!("Couldn't bind map memory");
                    }
                });
            }
        }
    }

    pub fn forget_map(self) {
        if self.has_map_memory() {
            world::with_mut(|w| {
                if let Some(ref mut mm) = w.map_memories_mut().get(self) {
                    mm.seen.clear();
                    mm.remembered.clear();
                } else {
                    panic!("Couldn't bind map memory");
                }
            });
        }
    }
}
<|endoftext|>


<commit_before>extern crate rand;

use std::io;
use rand::{thread_rng, sample};

struct Stack {
    content: Vec<i32>,
    digits: Vec<i32>,
}

impl Stack {
    fn push(&mut self, d: i32) -> Result<i32, &str> {
        self.content.push(d);
        Ok(d)
    }

    fn pop(&mut self) -> Option<i32> {
        self.content.pop()
    }

    fn consume_digit(&mut self, d: i32) {
        let i = self.digits.iter().position(|&i| i == d).expect("Cannot consume!");
        self.digits.remove(i);
    }

    fn valid_digit(&self, d: i32) -> bool {
        self.digits.contains(&d)
    }

    fn reveal<'a>(&self) -> Result<i32, &'a str> {
        match self.content.len() {
            1 => match self.digits.len() {
                0 => Ok(self.content[0]),
                _ => Err("Did not use all digits!"),
            },
            _ => Err("Invalid formula!"),
        }
    }

    fn add_digit(&mut self, c: char) -> Result<i32, &str> {
        match c.to_digit(10) {
            Some(n) => match self.valid_digit(n as i32) {
                true => {
                    self.consume_digit(n as i32);
                    self.push(n as i32)
                },
                false => Err("Invalid or duplicate digit!"),
            },
            None => Err("Invalid character"),
        }
    }

    fn apply<F>(&mut self, f: F) -> Result<i32, &str>
        where F : Fn(i32, i32) -> i32 {

        let r = self.pop();
        let l = self.pop();

        match l {
            Some(n) => match r {
                Some(i) => self.push(f(n, i)),
                None => Err("Stack empty"),
            },
            None => Err("Stack empty"),
        }
    }
}

fn evaluate<'a>(program: Vec<char>, digits: Vec<i32>) -> Result<i32, &'a str> {
    let mut stack = Stack {
        content: vec![],
        digits: digits,
    };

    for c in program {
        let answer: Result<i32, &str> = match c {
            '*' => stack.apply(|l, r| l * r),
            '+' => stack.apply(|l, r| l + r),
            '-' => stack.apply(|l, r| l - r),
            '/' => stack.apply(|l, r| l / r),
            _ => stack.add_digit(c),
        };

        if answer.is_err() {
            // TODO: Fix this. Work out how to return "answer" without compiler error.
            //return answer;
            println!("Actual: {}", answer.unwrap_err());
            return Err("Program error");
        }
    }

    stack.reveal()
}

fn main() {
    let goal: i32 = 24;
    let mut rng = thread_rng();

    println!("GUESSING GAME!

RULES

  * Using only the given digits, create an expression that evaluates to 24.
  * All digits must be used.
  * Each digit may only be used once.
  * Expressions must be in Reverse-Polish Notation (RPN).
  * Only +, -, * and / are allowed.
  * Order of the digits is not important.\n");

    loop {
        let mut program = String::new();
        let digits = sample(&mut rng, 1..9, 4);

        println!("Your digits are: {}, {}, {} and {}\n", digits[0], digits[1], digits[2], digits[3]);
        println!("Enter an expression:\n");

        io::stdin().read_line(&mut program)
            .ok()
            .expect("Could not read input.");

        let parsed = program
            .trim()
            .chars()
            .filter(|&c| c != ' ')
            .collect();

        match evaluate(parsed, digits) {
            Ok(n) => if n == goal {
                println!("\nWell done!\n");
                break;
            } else {
                println!("\nNice try, but {} is not {}!\n", n, goal);
            },
            Err(s) => {
                println!("\nOops: {}\n", s);
            },
        }
    }
}
<commit_msg>Cleaned up nested matches<commit_after>extern crate rand;

use std::io;
use rand::{thread_rng, sample};

struct Stack {
    content: Vec<i32>,
    digits: Vec<i32>,
}

impl Stack {
    fn push(&mut self, d: i32) -> Result<i32, &str> {
        self.content.push(d);
        Ok(d)
    }

    fn pop(&mut self) -> Option<i32> {
        self.content.pop()
    }

    fn consume_digit(&mut self, d: i32) {
        let i = self.digits.iter().position(|&i| i == d).expect("Cannot consume!");
        self.digits.remove(i);
    }

    fn valid_digit(&self, d: i32) -> bool {
        self.digits.contains(&d)
    }

    fn reveal<'a>(&self) -> Result<i32, &'a str> {
        match (self.content.len(), self.digits.len()) {
            (1, 0) => Ok(self.content[0]),
            (1, _) => Err("Did not use all digits!"),
            _ => Err("Invalid formula!"),
        }
    }

    fn add_digit(&mut self, c: char) -> Result<i32, &str> {
        match c.to_digit(10) {
            Some(n) => match self.valid_digit(n as i32) {
                true => {
                    self.consume_digit(n as i32);
                    self.push(n as i32)
                },
                false => Err("Invalid or duplicate digit!"),
            },
            None => Err("Invalid character"),
        }
    }

    fn apply<F>(&mut self, f: F) -> Result<i32, &str>
        where F : Fn(i32, i32) -> i32 {

        let (r, l) = (self.pop(), self.pop());

        match (l, r) {
            (Some(x), Some(y)) => self.push(f(x, y)),
            _ => Err("Stack empty"),
        }
    }
}

fn evaluate<'a>(program: Vec<char>, digits: Vec<i32>) -> Result<i32, &'a str> {
    let mut stack = Stack {
        content: vec![],
        digits: digits,
    };

    for c in program {
        let answer: Result<i32, &str> = match c {
            '*' => stack.apply(|l, r| l * r),
            '+' => stack.apply(|l, r| l + r),
            '-' => stack.apply(|l, r| l - r),
            '/' => stack.apply(|l, r| l / r),
            _ => stack.add_digit(c),
        };

        if answer.is_err() {
            // TODO: Fix this. Work out how to return "answer" without compiler error.
          //return answer;
          //return Err("Program error");
          println!("Actual: {}", answer.unwrap_err());
          break;
        }
    }

    stack.reveal()
}

fn main() {
    let goal: i32 = 24;
    let mut rng = thread_rng();

    println!("GUESSING GAME!

RULES

  * Using only the given digits, create an expression that evaluates to 24.
  * All digits must be used.
  * Each digit may only be used once.
  * Expressions must be in Reverse-Polish Notation (RPN).
  * Only +, -, * and / are allowed.
  * Order of the digits is not important.\n");

    loop {
        let mut program = String::new();
        let digits = sample(&mut rng, 1..9, 4);

        println!("Your digits are: {}, {}, {} and {}\n", digits[0], digits[1], digits[2], digits[3]);
        println!("Enter an expression:\n");

        io::stdin().read_line(&mut program)
            .ok()
            .expect("Could not read input.");

        let parsed = program
            .trim()
            .chars()
            .filter(|&c| c != ' ')
            .collect();

        match evaluate(parsed, digits) {
            Ok(n) => if n == goal {
                println!("\nWell done!\n");
                break;
            } else {
                println!("\nNice try, but {} is not {}!\n", n, goal);
            },
            Err(s) => {
                println!("\nOops: {}\n", s);
            },
        }
    }
}
<|endoftext|>


<commit_before>extern crate libc;

use libc::funcs::posix88::unistd::getlogin;
use std::c_str::CString;
use std::os;

fn main() {
    unsafe {

        let login_c = getlogin() as *const libc::c_char;
        let login = CString::new(login_c, false);

        match login.as_str() {
            Some(username) => println!("{}", username),
            None => os::set_exit_status(1),
        }

    }
}
<commit_msg>Move code outside of unsafe block<commit_after>extern crate libc;

use libc::funcs::posix88::unistd::getlogin;
use std::c_str::CString;
use std::os;

fn main() {
    let login = unsafe {
        CString::new(
            getlogin() as *const libc::c_char,
            false
        )
    };

    match login.as_str() {
        Some(username) => println!("{}", username),
        None => os::set_exit_status(1),
    }
}
<|endoftext|>


<commit_before>//! This is a binary running in the server environment
//!
//! You have to provide all needed configuration attributes via command line parameters,
//! or you could specify a configuration file. The format of configuration file is defined
//! in mod `config`.
//!
//! *It should be notice that the extented configuration file is not suitable for the server
//! side.*

use std::net::{IpAddr, SocketAddr};

use clap::{clap_app, Arg};
use futures::future::{self, Either};
use log::{error, info};
use tokio::{self, runtime::Builder};

use shadowsocks::{
    acl::AccessControl,
    crypto::CipherType,
    run_manager,
    Config,
    ConfigType,
    ManagerAddr,
    Mode,
    ServerAddr,
};

mod logging;
mod monitor;
mod validator;

fn main() {
    let app = clap_app!(shadowsocks =>
        (version: shadowsocks::VERSION)
        (about: "A fast tunnel proxy that helps you bypass firewalls.")
        (@arg VERBOSE: -v ... "Set the level of debug")
        (@arg UDP_ONLY: -u conflicts_with[TCP_AND_UDP] "Server mode UDP_ONLY")
        (@arg TCP_AND_UDP: -U conflicts_with[UDP_ONLY] "Server mode TCP_AND_UDP")
        (@arg CONFIG: -c --config +takes_value "Specify config file")
        (@arg BIND_ADDR: -b --("bind-addr") +takes_value "Bind address, outbound socket will bind this address")
        (@arg NO_DELAY: --("no-delay") !takes_value "Set no-delay option for socket")
        (@arg MANAGER_ADDRESS: --("manager-address") +takes_value {validator::validate_manager_addr} "ShadowSocks Manager (ssmgr) address, could be ip:port, domain:port or /path/to/unix.sock")
        (@arg ENCRYPT_METHOD: -m --("encrypt-method") +takes_value possible_values(&CipherType::available_ciphers()) requires[SERVER_ADDR PASSWORD] "Encryption method")
        (@group MANAGER_CONFIG =>
            (@attributes +required ... arg[CONFIG MANAGER_ADDRESS])
        )
        (@arg NOFILE: -n --nofile +takes_value "Set RLIMIT_NOFILE with both soft and hard limit (only for *nix systems)")
        (@arg ACL: --acl +takes_value "Path to ACL (Access Control List)")
        (@arg LOG_WITHOUT_TIME: --("log-without-time") "Log without datetime prefix")
    );

    let matches = app
        .arg(
            Arg::with_name("IPV6_FIRST")
                .short("6")
                .help("Resovle hostname to IPv6 address first"),
        )
        .get_matches();

    let debug_level = matches.occurrences_of("VERBOSE");
    logging::init(debug_level, "ssmanager", matches.is_present("LOG_WITHOUT_TIME"));

    let mut config = match matches.value_of("CONFIG") {
        Some(cpath) => match Config::load_from_file(cpath, ConfigType::Manager) {
            Ok(cfg) => cfg,
            Err(err) => {
                error!("{:?}", err);
                return;
            }
        },
        None => Config::new(ConfigType::Manager),
    };

    if let Some(method) = matches.value_of("ENCRYPT_METHOD") {
        config.manager_method = Some(method.parse().expect("encryption method"));
    }

    if let Some(bind_addr) = matches.value_of("BIND_ADDR") {
        let bind_addr = match bind_addr.parse::<IpAddr>() {
            Ok(ip) => ServerAddr::from(SocketAddr::new(ip, 0)),
            Err(..) => ServerAddr::from((bind_addr, 0)),
        };

        config.local_addr = Some(bind_addr);
    }

    if matches.is_present("UDP_ONLY") {
        if config.mode.enable_tcp() {
            config.mode = Mode::TcpAndUdp;
        } else {
            config.mode = Mode::UdpOnly;
        }
    }

    if matches.is_present("TCP_AND_UDP") {
        config.mode = Mode::TcpAndUdp;
    }

    if matches.is_present("NO_DELAY") {
        config.no_delay = true;
    }

    if let Some(m) = matches.value_of("MANAGER_ADDRESS") {
        config.manager_addr = Some(m.parse::<ManagerAddr>().expect("manager bind address"));
    }

    if let Some(nofile) = matches.value_of("NOFILE") {
        config.nofile = Some(nofile.parse::<u64>().expect("an unsigned integer for `nofile`"));
    }

    if let Some(acl_file) = matches.value_of("ACL") {
        let acl = AccessControl::load_from_file(acl_file).expect("load ACL file");
        config.acl = Some(acl);
    }

    if matches.is_present("IPV6_FIRST") {
        config.ipv6_first = true;
    }

    // DONE reading options

    if config.manager_addr.is_none() {
        eprintln!(
            "missing `manager_address`, consider specifying it by --manager-address command line option, \
             or \"manager_address\" and \"manager_port\" keys in configuration file"
        );
        println!("{}", matches.usage());
        return;
    }

    info!("shadowsocks {}", shadowsocks::VERSION);

    let mut builder = Builder::new();
    if cfg!(feature = "single-threaded") {
        builder.basic_scheduler();
    } else {
        builder.threaded_scheduler();
    }
    let mut runtime = builder.enable_all().build().expect("create tokio Runtime");
    let rt_handle = runtime.handle().clone();

    runtime.block_on(async move {
        let abort_signal = monitor::create_signal_monitor();
        let server = run_manager(config, rt_handle);

        tokio::pin!(abort_signal);
        tokio::pin!(server);

        match future::select(server, abort_signal).await {
            // Server future resolved without an error. This should never happen.
            Either::Left((Ok(..), ..)) => panic!("server exited unexpectly"),
            // Server future resolved with error, which are listener errors in most cases
            Either::Left((Err(err), ..)) => panic!("server exited unexpectly with {}", err),
            // The abort signal future resolved. Means we should just exit.
            Either::Right(_) => (),
        }
    })
}
<commit_msg>Fixed wrong requires in ssmanager command line defines<commit_after>//! This is a binary running in the server environment
//!
//! You have to provide all needed configuration attributes via command line parameters,
//! or you could specify a configuration file. The format of configuration file is defined
//! in mod `config`.
//!
//! *It should be notice that the extented configuration file is not suitable for the server
//! side.*

use std::net::{IpAddr, SocketAddr};

use clap::{clap_app, Arg};
use futures::future::{self, Either};
use log::{error, info};
use tokio::{self, runtime::Builder};

use shadowsocks::{
    acl::AccessControl,
    crypto::CipherType,
    run_manager,
    Config,
    ConfigType,
    ManagerAddr,
    Mode,
    ServerAddr,
};

mod logging;
mod monitor;
mod validator;

fn main() {
    let available_ciphers = CipherType::available_ciphers();

    let app = clap_app!(shadowsocks =>
        (version: shadowsocks::VERSION)
        (about: "A fast tunnel proxy that helps you bypass firewalls.")
        (@arg VERBOSE: -v ... "Set the level of debug")
        (@arg UDP_ONLY: -u conflicts_with[TCP_AND_UDP] "Server mode UDP_ONLY")
        (@arg TCP_AND_UDP: -U conflicts_with[UDP_ONLY] "Server mode TCP_AND_UDP")
        (@arg CONFIG: -c --config +takes_value "Specify config file")
        (@arg BIND_ADDR: -b --("bind-addr") +takes_value "Bind address, outbound socket will bind this address")
        (@arg NO_DELAY: --("no-delay") !takes_value "Set no-delay option for socket")
        (@arg MANAGER_ADDRESS: --("manager-address") +takes_value {validator::validate_manager_addr} "ShadowSocks Manager (ssmgr) address, could be ip:port, domain:port or /path/to/unix.sock")
        (@arg ENCRYPT_METHOD: -m --("encrypt-method") +takes_value possible_values(&available_ciphers) "Encryption method")
        (@group MANAGER_CONFIG =>
            (@attributes +required ... arg[CONFIG MANAGER_ADDRESS])
        )
        (@arg NOFILE: -n --nofile +takes_value "Set RLIMIT_NOFILE with both soft and hard limit (only for *nix systems)")
        (@arg ACL: --acl +takes_value "Path to ACL (Access Control List)")
        (@arg LOG_WITHOUT_TIME: --("log-without-time") "Log without datetime prefix")
    );

    let matches = app
        .arg(
            Arg::with_name("IPV6_FIRST")
                .short("6")
                .help("Resovle hostname to IPv6 address first"),
        )
        .get_matches();

    drop(available_ciphers);

    let debug_level = matches.occurrences_of("VERBOSE");
    logging::init(debug_level, "ssmanager", matches.is_present("LOG_WITHOUT_TIME"));

    let mut config = match matches.value_of("CONFIG") {
        Some(cpath) => match Config::load_from_file(cpath, ConfigType::Manager) {
            Ok(cfg) => cfg,
            Err(err) => {
                error!("{:?}", err);
                return;
            }
        },
        None => Config::new(ConfigType::Manager),
    };

    if let Some(method) = matches.value_of("ENCRYPT_METHOD") {
        config.manager_method = Some(method.parse().expect("encryption method"));
    }

    if let Some(bind_addr) = matches.value_of("BIND_ADDR") {
        let bind_addr = match bind_addr.parse::<IpAddr>() {
            Ok(ip) => ServerAddr::from(SocketAddr::new(ip, 0)),
            Err(..) => ServerAddr::from((bind_addr, 0)),
        };

        config.local_addr = Some(bind_addr);
    }

    if matches.is_present("UDP_ONLY") {
        if config.mode.enable_tcp() {
            config.mode = Mode::TcpAndUdp;
        } else {
            config.mode = Mode::UdpOnly;
        }
    }

    if matches.is_present("TCP_AND_UDP") {
        config.mode = Mode::TcpAndUdp;
    }

    if matches.is_present("NO_DELAY") {
        config.no_delay = true;
    }

    if let Some(m) = matches.value_of("MANAGER_ADDRESS") {
        config.manager_addr = Some(m.parse::<ManagerAddr>().expect("manager bind address"));
    }

    if let Some(nofile) = matches.value_of("NOFILE") {
        config.nofile = Some(nofile.parse::<u64>().expect("an unsigned integer for `nofile`"));
    }

    if let Some(acl_file) = matches.value_of("ACL") {
        let acl = AccessControl::load_from_file(acl_file).expect("load ACL file");
        config.acl = Some(acl);
    }

    if matches.is_present("IPV6_FIRST") {
        config.ipv6_first = true;
    }

    // DONE reading options

    if config.manager_addr.is_none() {
        eprintln!(
            "missing `manager_address`, consider specifying it by --manager-address command line option, \
             or \"manager_address\" and \"manager_port\" keys in configuration file"
        );
        println!("{}", matches.usage());
        return;
    }

    info!("shadowsocks {}", shadowsocks::VERSION);

    let mut builder = Builder::new();
    if cfg!(feature = "single-threaded") {
        builder.basic_scheduler();
    } else {
        builder.threaded_scheduler();
    }
    let mut runtime = builder.enable_all().build().expect("create tokio Runtime");
    let rt_handle = runtime.handle().clone();

    runtime.block_on(async move {
        let abort_signal = monitor::create_signal_monitor();
        let server = run_manager(config, rt_handle);

        tokio::pin!(abort_signal);
        tokio::pin!(server);

        match future::select(server, abort_signal).await {
            // Server future resolved without an error. This should never happen.
            Either::Left((Ok(..), ..)) => panic!("server exited unexpectly"),
            // Server future resolved with error, which are listener errors in most cases
            Either::Left((Err(err), ..)) => panic!("server exited unexpectly with {}", err),
            // The abort signal future resolved. Means we should just exit.
            Either::Right(_) => (),
        }
    })
}
<|endoftext|>


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

#![cfg(not(test))]

extern crate log;
extern crate rustfmt;
extern crate toml;
extern crate env_logger;
extern crate getopts;

use rustfmt::{run, Input, Summary};
use rustfmt::config::{Config, WriteMode};

use std::{env, error};
use std::fs::{self, File};
use std::io::{self, ErrorKind, Read, Write};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use getopts::{Matches, Options};

type FmtError = Box<error::Error + Send + Sync>;
type FmtResult<T> = std::result::Result<T, FmtError>;

/// Rustfmt operations.
enum Operation {
    /// Format files and their child modules.
    Format {
        files: Vec<PathBuf>,
        config_path: Option<PathBuf>,
    },
    /// Print the help message.
    Help,
    // Print version information
    Version,
    /// Print detailed configuration help.
    ConfigHelp,
    /// No file specified, read from stdin
    Stdin {
        input: String,
        config_path: Option<PathBuf>,
    },
}

/// Try to find a project file in the given directory and its parents. Returns the path of a the
/// nearest project file if one exists, or `None` if no project file was found.
fn lookup_project_file(dir: &Path) -> FmtResult<Option<PathBuf>> {
    let mut current = if dir.is_relative() {
        try!(env::current_dir()).join(dir)
    } else {
        dir.to_path_buf()
    };

    current = try!(fs::canonicalize(current));

    loop {
        let config_file = current.join("rustfmt.toml");
        match fs::metadata(&config_file) {
            // Only return if it's a file to handle the unlikely situation of a directory named
            // `rustfmt.toml`.
            Ok(ref md) if md.is_file() => return Ok(Some(config_file)),
            // Return the error if it's something other than `NotFound`; otherwise we didn't find
            // the project file yet, and continue searching.
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    return Err(FmtError::from(e));
                }
            }
            _ => {}
        }

        // If the current directory has no parent, we're done searching.
        if !current.pop() {
            return Ok(None);
        }
    }
}

/// Resolve the config for input in `dir`.
///
/// Returns the `Config` to use, and the path of the project file if there was
/// one.
fn resolve_config(dir: &Path) -> FmtResult<(Config, Option<PathBuf>)> {
    let path = try!(lookup_project_file(dir));
    if path.is_none() {
        return Ok((Config::default(), None));
    }
    let path = path.unwrap();
    let mut file = try!(File::open(&path));
    let mut toml = String::new();
    try!(file.read_to_string(&mut toml));
    Ok((Config::from_toml(&toml), Some(path)))
}

/// read the given config file path recursively if present else read the project file path
fn match_cli_path_or_file(config_path: Option<PathBuf>,
                          input_file: &Path)
                          -> FmtResult<(Config, Option<PathBuf>)> {

    if let Some(config_file) = config_path {
        let (toml, path) = try!(resolve_config(config_file.as_ref()));
        if path.is_some() {
            return Ok((toml, path));
        }
    }
    resolve_config(input_file)
}

fn update_config(config: &mut Config, matches: &Matches) -> FmtResult<()> {
    config.verbose = matches.opt_present("verbose");
    config.skip_children = matches.opt_present("skip-children");

    let write_mode = matches.opt_str("write-mode");
    match matches.opt_str("write-mode").map(|wm| WriteMode::from_str(&wm)) {
        None => Ok(()),
        Some(Ok(write_mode)) => {
            config.write_mode = write_mode;
            Ok(())
        }
        Some(Err(_)) => {
            Err(FmtError::from(format!("Invalid write-mode: {}",
                                       write_mode.expect("cannot happen"))))
        }
    }
}

fn make_opts() -> Options {
    let mut opts = Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("V", "version", "show version information");
    opts.optflag("v", "verbose", "show progress");
    opts.optopt("",
                "write-mode",
                "mode to write in (not usable when piping from stdin)",
                "[replace|overwrite|display|diff|coverage|checkstyle]");
    opts.optflag("", "skip-children", "don't reformat child modules");

    opts.optflag("",
                 "config-help",
                 "show details of rustfmt configuration options");
    opts.optopt("",
                "config-path",
                "Recursively searches the given path for the rustfmt.toml config file. If not \
                 found reverts to the input file path",
                "[Path for the configuration file]");

    opts
}

fn execute(opts: &Options) -> FmtResult<Summary> {
    let matches = try!(opts.parse(env::args().skip(1)));

    match try!(determine_operation(&matches)) {
        Operation::Help => {
            print_usage(&opts, "");
            Ok(Summary::new())
        }
        Operation::Version => {
            print_version();
            Ok(Summary::new())
        }
        Operation::ConfigHelp => {
            Config::print_docs();
            Ok(Summary::new())
        }
        Operation::Stdin { input, config_path } => {
            // try to read config from local directory
            let (mut config, _) = match_cli_path_or_file(config_path, &env::current_dir().unwrap())
                                      .expect("Error resolving config");

            // write_mode is always Plain for Stdin.
            config.write_mode = WriteMode::Plain;

            Ok(run(Input::Text(input), &config))
        }
        Operation::Format { files, config_path } => {
            let mut config = Config::default();
            let mut path = None;
            // Load the config path file if provided
            if let Some(config_file) = config_path {
                let (cfg_tmp, path_tmp) = resolve_config(config_file.as_ref())
                                              .expect(&format!("Error resolving config for {:?}",
                                                               config_file));
                config = cfg_tmp;
                path = path_tmp;
            };
            if let Some(path) = path.as_ref() {
                println!("Using rustfmt config file {}", path.display());
            }

            let mut error_summary = Summary::new();
            for file in files {
                // Check the file directory if the config-path could not be read or not provided
                if path.is_none() {
                    let (config_tmp, path_tmp) = resolve_config(file.parent().unwrap())
                                                     .expect(&format!("Error resolving config \
                                                                       for {}",
                                                                      file.display()));
                    if let Some(path) = path_tmp.as_ref() {
                        println!("Using rustfmt config file {} for {}",
                                 path.display(),
                                 file.display());
                    }
                    config = config_tmp;
                }

                try!(update_config(&mut config, &matches));
                error_summary.add(run(Input::File(file), &config));
            }
            Ok(error_summary)
        }
    }
}

fn main() {
    let _ = env_logger::init();

    let opts = make_opts();

    let exit_code = match execute(&opts) {
        Ok(summary) => {
            if summary.has_operational_errors() {
                1
            } else if summary.has_parsing_errors() {
                2
            } else if summary.has_formatting_errors() {
                3
            } else {
                assert!(summary.has_no_errors());
                0
            }
        }
        Err(e) => {
            print_usage(&opts, &e.to_string());
            1
        }
    };
    // Make sure standard output is flushed before we exit.
    std::io::stdout().flush().unwrap();

    // Exit with given exit code.
    //
    // NOTE: This immediately terminates the process without doing any cleanup,
    // so make sure to finish all necessary cleanup before this is called.
    std::process::exit(exit_code);
}

fn print_usage(opts: &Options, reason: &str) {
    let reason = format!("{}\nusage: {} [options] <file>...",
                         reason,
                         env::args_os().next().unwrap().to_string_lossy());
    println!("{}", opts.usage(&reason));
}

fn print_version() {
    println!("{}.{}.{}{}",
             option_env!("CARGO_PKG_VERSION_MAJOR").unwrap_or("X"),
             option_env!("CARGO_PKG_VERSION_MINOR").unwrap_or("X"),
             option_env!("CARGO_PKG_VERSION_PATCH").unwrap_or("X"),
             option_env!("CARGO_PKG_VERSION_PRE").unwrap_or(""));
}

fn determine_operation(matches: &Matches) -> FmtResult<Operation> {
    if matches.opt_present("h") {
        return Ok(Operation::Help);
    }

    if matches.opt_present("config-help") {
        return Ok(Operation::ConfigHelp);
    }

    if matches.opt_present("version") {
        return Ok(Operation::Version);
    }

    // Read the config_path and convert to parent dir if a file is provided.
    let config_path: Option<PathBuf> = matches.opt_str("config-path")
                                              .map(PathBuf::from)
                                              .and_then(|dir| {
                                                  if dir.is_file() {
                                                      return dir.parent().map(|v| v.into());
                                                  }
                                                  Some(dir)
                                              });

    // if no file argument is supplied, read from stdin
    if matches.free.is_empty() {

        let mut buffer = String::new();
        try!(io::stdin().read_to_string(&mut buffer));

        return Ok(Operation::Stdin {
            input: buffer,
            config_path: config_path,
        });
    }

    let files: Vec<_> = matches.free.iter().map(PathBuf::from).collect();

    Ok(Operation::Format {
        files: files,
        config_path: config_path,
    })
}
<commit_msg>rustfmt: Parse options once instead of once per file argument (#944)<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![cfg(not(test))]

extern crate log;
extern crate rustfmt;
extern crate toml;
extern crate env_logger;
extern crate getopts;

use rustfmt::{run, Input, Summary};
use rustfmt::config::{Config, WriteMode};

use std::{env, error};
use std::fs::{self, File};
use std::io::{self, ErrorKind, Read, Write};
use std::path::{Path, PathBuf};
use std::str::FromStr;

use getopts::{Matches, Options};

type FmtError = Box<error::Error + Send + Sync>;
type FmtResult<T> = std::result::Result<T, FmtError>;

/// Rustfmt operations.
enum Operation {
    /// Format files and their child modules.
    Format {
        files: Vec<PathBuf>,
        config_path: Option<PathBuf>,
    },
    /// Print the help message.
    Help,
    // Print version information
    Version,
    /// Print detailed configuration help.
    ConfigHelp,
    /// No file specified, read from stdin
    Stdin {
        input: String,
        config_path: Option<PathBuf>,
    },
}

/// Parsed command line options.
#[derive(Clone, Debug, Default)]
struct CliOptions {
    skip_children: bool,
    verbose: bool,
    write_mode: Option<WriteMode>,
}

impl CliOptions {
    fn from_matches(matches: &Matches) -> FmtResult<CliOptions> {
        let mut options = CliOptions::default();
        options.skip_children = matches.opt_present("skip-children");
        options.verbose = matches.opt_present("verbose");

        if let Some(ref write_mode) = matches.opt_str("write-mode") {
            if let Ok(write_mode) = WriteMode::from_str(write_mode) {
                options.write_mode = Some(write_mode);
            } else {
                return Err(FmtError::from(format!("Invalid write-mode: {}", write_mode)));
            }
        }

        Ok(options)
    }

    fn apply_to(&self, config: &mut Config) {
        config.skip_children = self.skip_children;
        config.verbose = self.verbose;
        if let Some(write_mode) = self.write_mode {
            config.write_mode = write_mode;
        }
    }
}

/// Try to find a project file in the given directory and its parents. Returns the path of a the
/// nearest project file if one exists, or `None` if no project file was found.
fn lookup_project_file(dir: &Path) -> FmtResult<Option<PathBuf>> {
    let mut current = if dir.is_relative() {
        try!(env::current_dir()).join(dir)
    } else {
        dir.to_path_buf()
    };

    current = try!(fs::canonicalize(current));

    loop {
        let config_file = current.join("rustfmt.toml");
        match fs::metadata(&config_file) {
            // Only return if it's a file to handle the unlikely situation of a directory named
            // `rustfmt.toml`.
            Ok(ref md) if md.is_file() => return Ok(Some(config_file)),
            // Return the error if it's something other than `NotFound`; otherwise we didn't find
            // the project file yet, and continue searching.
            Err(e) => {
                if e.kind() != ErrorKind::NotFound {
                    return Err(FmtError::from(e));
                }
            }
            _ => {}
        }

        // If the current directory has no parent, we're done searching.
        if !current.pop() {
            return Ok(None);
        }
    }
}

/// Resolve the config for input in `dir`.
///
/// Returns the `Config` to use, and the path of the project file if there was
/// one.
fn resolve_config(dir: &Path) -> FmtResult<(Config, Option<PathBuf>)> {
    let path = try!(lookup_project_file(dir));
    if path.is_none() {
        return Ok((Config::default(), None));
    }
    let path = path.unwrap();
    let mut file = try!(File::open(&path));
    let mut toml = String::new();
    try!(file.read_to_string(&mut toml));
    Ok((Config::from_toml(&toml), Some(path)))
}

/// read the given config file path recursively if present else read the project file path
fn match_cli_path_or_file(config_path: Option<PathBuf>,
                          input_file: &Path)
                          -> FmtResult<(Config, Option<PathBuf>)> {

    if let Some(config_file) = config_path {
        let (toml, path) = try!(resolve_config(config_file.as_ref()));
        if path.is_some() {
            return Ok((toml, path));
        }
    }
    resolve_config(input_file)
}

fn make_opts() -> Options {
    let mut opts = Options::new();
    opts.optflag("h", "help", "show this message");
    opts.optflag("V", "version", "show version information");
    opts.optflag("v", "verbose", "show progress");
    opts.optopt("",
                "write-mode",
                "mode to write in (not usable when piping from stdin)",
                "[replace|overwrite|display|diff|coverage|checkstyle]");
    opts.optflag("", "skip-children", "don't reformat child modules");

    opts.optflag("",
                 "config-help",
                 "show details of rustfmt configuration options");
    opts.optopt("",
                "config-path",
                "Recursively searches the given path for the rustfmt.toml config file. If not \
                 found reverts to the input file path",
                "[Path for the configuration file]");

    opts
}

fn execute(opts: &Options) -> FmtResult<Summary> {
    let matches = try!(opts.parse(env::args().skip(1)));

    match try!(determine_operation(&matches)) {
        Operation::Help => {
            print_usage(&opts, "");
            Ok(Summary::new())
        }
        Operation::Version => {
            print_version();
            Ok(Summary::new())
        }
        Operation::ConfigHelp => {
            Config::print_docs();
            Ok(Summary::new())
        }
        Operation::Stdin { input, config_path } => {
            // try to read config from local directory
            let (mut config, _) = match_cli_path_or_file(config_path, &env::current_dir().unwrap())
                                      .expect("Error resolving config");

            // write_mode is always Plain for Stdin.
            config.write_mode = WriteMode::Plain;

            Ok(run(Input::Text(input), &config))
        }
        Operation::Format { files, config_path } => {
            let options = try!(CliOptions::from_matches(&matches));
            let mut config = Config::default();
            let mut path = None;
            // Load the config path file if provided
            if let Some(config_file) = config_path {
                let (cfg_tmp, path_tmp) = resolve_config(config_file.as_ref())
                                              .expect(&format!("Error resolving config for {:?}",
                                                               config_file));
                config = cfg_tmp;
                path = path_tmp;
            };
            if let Some(path) = path.as_ref() {
                println!("Using rustfmt config file {}", path.display());
            }

            let mut error_summary = Summary::new();
            for file in files {
                // Check the file directory if the config-path could not be read or not provided
                if path.is_none() {
                    let (config_tmp, path_tmp) = resolve_config(file.parent().unwrap())
                                                     .expect(&format!("Error resolving config \
                                                                       for {}",
                                                                      file.display()));
                    if let Some(path) = path_tmp.as_ref() {
                        println!("Using rustfmt config file {} for {}",
                                 path.display(),
                                 file.display());
                    }
                    config = config_tmp;
                }

                options.apply_to(&mut config);
                error_summary.add(run(Input::File(file), &config));
            }
            Ok(error_summary)
        }
    }
}

fn main() {
    let _ = env_logger::init();

    let opts = make_opts();

    let exit_code = match execute(&opts) {
        Ok(summary) => {
            if summary.has_operational_errors() {
                1
            } else if summary.has_parsing_errors() {
                2
            } else if summary.has_formatting_errors() {
                3
            } else {
                assert!(summary.has_no_errors());
                0
            }
        }
        Err(e) => {
            print_usage(&opts, &e.to_string());
            1
        }
    };
    // Make sure standard output is flushed before we exit.
    std::io::stdout().flush().unwrap();

    // Exit with given exit code.
    //
    // NOTE: This immediately terminates the process without doing any cleanup,
    // so make sure to finish all necessary cleanup before this is called.
    std::process::exit(exit_code);
}

fn print_usage(opts: &Options, reason: &str) {
    let reason = format!("{}\nusage: {} [options] <file>...",
                         reason,
                         env::args_os().next().unwrap().to_string_lossy());
    println!("{}", opts.usage(&reason));
}

fn print_version() {
    println!("{}.{}.{}{}",
             option_env!("CARGO_PKG_VERSION_MAJOR").unwrap_or("X"),
             option_env!("CARGO_PKG_VERSION_MINOR").unwrap_or("X"),
             option_env!("CARGO_PKG_VERSION_PATCH").unwrap_or("X"),
             option_env!("CARGO_PKG_VERSION_PRE").unwrap_or(""));
}

fn determine_operation(matches: &Matches) -> FmtResult<Operation> {
    if matches.opt_present("h") {
        return Ok(Operation::Help);
    }

    if matches.opt_present("config-help") {
        return Ok(Operation::ConfigHelp);
    }

    if matches.opt_present("version") {
        return Ok(Operation::Version);
    }

    // Read the config_path and convert to parent dir if a file is provided.
    let config_path: Option<PathBuf> = matches.opt_str("config-path")
                                              .map(PathBuf::from)
                                              .and_then(|dir| {
                                                  if dir.is_file() {
                                                      return dir.parent().map(|v| v.into());
                                                  }
                                                  Some(dir)
                                              });

    // if no file argument is supplied, read from stdin
    if matches.free.is_empty() {

        let mut buffer = String::new();
        try!(io::stdin().read_to_string(&mut buffer));

        return Ok(Operation::Stdin {
            input: buffer,
            config_path: config_path,
        });
    }

    let files: Vec<_> = matches.free.iter().map(PathBuf::from).collect();

    Ok(Operation::Format {
        files: files,
        config_path: config_path,
    })
}
<|endoftext|>


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

use crate::color;
use crate::error::{ImageError, ImageResult, ParameterError, ParameterErrorKind};
use crate::image::ImageEncoder;

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

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

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

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

        let (dib_header_size, written_pixel_size, palette_color_count) = get_pixel_info(c)?;
        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes
        let image_size = width
            .checked_mul(height)
            .and_then(|v| v.checked_mul(written_pixel_size))
            .and_then(|v| v.checked_add(height * row_pad_size))
            .ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(
                    ParameterErrorKind::DimensionMismatch,
                ))
            })?;
        let palette_size = palette_color_count * 4; // all palette colors are BGRA
        let file_size = bmp_header_size + dib_header_size + palette_size + image_size;

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

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

            // endpoints (3x3) and gamma (3)
            for _ in 0..12 {
                self.writer.write_u32::<LittleEndian>(0)?;
            }
        }

        // write image data
        match c {
            color::ColorType::Rgb8 => {
                self.encode_rgb(image, width, height, row_pad_size, 3)?
            }
            color::ColorType::Rgba8 => {
                self.encode_rgba(image, width, height, row_pad_size, 4)?
            }
            color::ColorType::L8 => {
                self.encode_gray(image, width, height, row_pad_size, 1)?
            }
            color::ColorType::La8 => {
                self.encode_gray(image, width, height, row_pad_size, 2)?
            }
            _ => {
                return Err(ImageError::IoError(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &get_unsupported_error_message(c)[..],
                )))
            }
        }

        Ok(())
    }

    fn encode_rgb(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let width = width as usize;
        let height = height as usize;
        let x_stride = bytes_per_pixel as usize;
        let y_stride = width * x_stride;
        for row in (0..height).rev() {
            // from the bottom up
            let row_start = row * y_stride;
            for px in image[row_start.. row_start + y_stride].chunks_exact(x_stride) {
                let r = px[0];
                let g = px[1];
                let b = px[2];
                // written as BGR
                self.writer.write(&[b, g, r])?;
            }
            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

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

        Ok(())
    }

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

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

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

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

        Ok(())
    }
}

impl<'a, W: Write> ImageEncoder for BMPEncoder<'a, W> {
    fn write_image(
        mut self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: color::ColorType,
    ) -> ImageResult<()> {
        self.encode(buf, width, height, color_type)
    }
}

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

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

    Ok(sizes)
}

#[cfg(test)]
mod tests {
    use super::super::BmpDecoder;
    use super::BMPEncoder;
    use crate::color::ColorType;
    use crate::image::ImageDecoder;
    use std::io::Cursor;

    #[cfg(feature = "benchmarks")]
    use test::{Bencher};

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

        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");

        let mut buf = vec![0; decoder.total_bytes() as usize];
        decoder.read_image(&mut buf).expect("failed to decode");
        buf
    }

    #[cfg(feature = "benchmarks")]
    #[bench]
    fn bench_encode_rgb(b: &mut Bencher) {
        let mut v= Vec::with_capacity(2000 * 2000);
        let mut x = BMPEncoder::new(&mut v);
        let mut im = vec![0; 2000 * 2000];
        b.iter(|| {
            x.encode_rgb(&mut im, 500, 2000, 0, 4);
        });
    }

    #[cfg(feature = "benchmarks")]
    #[bench]
    fn bench_encode_rgba(b: &mut Bencher) {
        let mut v= Vec::with_capacity(2000 * 2000);
        let mut x = BMPEncoder::new(&mut v);
        let mut im = vec![0; 2000 * 2000];
        b.iter(|| {
            x.encode_rgba(&mut im, 500, 2000, 0, 4);
        });
    }

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

    #[test]
    fn huge_files_return_error() {
        let mut encoded_data = Vec::new();
        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
        let mut encoder = BMPEncoder::new(&mut encoded_data);
        let result = encoder.encode(&image, 40_000, 40_000, ColorType::Rgb8);
        assert!(result.is_err());
    }

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

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

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

    #[test]
    fn round_trip_graya() {
        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
        let decoded = round_trip_image(&image, 1, 3, ColorType::La8);
        // should be read back as 3 RGB pixels
        assert_eq!(9, decoded.len());
        assert_eq!(0, decoded[0]);
        assert_eq!(0, decoded[1]);
        assert_eq!(0, decoded[2]);
        assert_eq!(1, decoded[3]);
        assert_eq!(1, decoded[4]);
        assert_eq!(1, decoded[5]);
        assert_eq!(2, decoded[6]);
        assert_eq!(2, decoded[7]);
        assert_eq!(2, decoded[8]);
    }
}
<commit_msg>Use write_all to ensure bytes get written.<commit_after>use byteorder::{LittleEndian, WriteBytesExt};
use std::io::{self, Write};

use crate::color;
use crate::error::{ImageError, ImageResult, ParameterError, ParameterErrorKind};
use crate::image::ImageEncoder;

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

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

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

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

        let (dib_header_size, written_pixel_size, palette_color_count) = get_pixel_info(c)?;
        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes
        let image_size = width
            .checked_mul(height)
            .and_then(|v| v.checked_mul(written_pixel_size))
            .and_then(|v| v.checked_add(height * row_pad_size))
            .ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(
                    ParameterErrorKind::DimensionMismatch,
                ))
            })?;
        let palette_size = palette_color_count * 4; // all palette colors are BGRA
        let file_size = bmp_header_size + dib_header_size + palette_size + image_size;

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

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

            // endpoints (3x3) and gamma (3)
            for _ in 0..12 {
                self.writer.write_u32::<LittleEndian>(0)?;
            }
        }

        // write image data
        match c {
            color::ColorType::Rgb8 => {
                self.encode_rgb(image, width, height, row_pad_size, 3)?
            }
            color::ColorType::Rgba8 => {
                self.encode_rgba(image, width, height, row_pad_size, 4)?
            }
            color::ColorType::L8 => {
                self.encode_gray(image, width, height, row_pad_size, 1)?
            }
            color::ColorType::La8 => {
                self.encode_gray(image, width, height, row_pad_size, 2)?
            }
            _ => {
                return Err(ImageError::IoError(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    &get_unsupported_error_message(c)[..],
                )))
            }
        }

        Ok(())
    }

    fn encode_rgb(
        &mut self,
        image: &[u8],
        width: u32,
        height: u32,
        row_pad_size: u32,
        bytes_per_pixel: u32,
    ) -> io::Result<()> {
        let width = width as usize;
        let height = height as usize;
        let x_stride = bytes_per_pixel as usize;
        let y_stride = width * x_stride;
        for row in (0..height).rev() {
            // from the bottom up
            let row_start = row * y_stride;
            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
                let r = px[0];
                let g = px[1];
                let b = px[2];
                // written as BGR
                self.writer.write_all(&[b, g, r])?;
            }
            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

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

        Ok(())
    }

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

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

            self.write_row_pad(row_pad_size)?;
        }

        Ok(())
    }

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

        Ok(())
    }
}

impl<'a, W: Write> ImageEncoder for BMPEncoder<'a, W> {
    fn write_image(
        mut self,
        buf: &[u8],
        width: u32,
        height: u32,
        color_type: color::ColorType,
    ) -> ImageResult<()> {
        self.encode(buf, width, height, color_type)
    }
}

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

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

    Ok(sizes)
}

#[cfg(test)]
mod tests {
    use super::super::BmpDecoder;
    use super::BMPEncoder;
    use crate::color::ColorType;
    use crate::image::ImageDecoder;
    use std::io::Cursor;

    #[cfg(feature = "benchmarks")]
    use test::{Bencher};

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

        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");

        let mut buf = vec![0; decoder.total_bytes() as usize];
        decoder.read_image(&mut buf).expect("failed to decode");
        buf
    }

    #[cfg(feature = "benchmarks")]
    #[bench]
    fn bench_encode_rgb(b: &mut Bencher) {
        let mut v= Vec::with_capacity(2000 * 2000);
        let mut x = BMPEncoder::new(&mut v);
        let mut im = vec![0; 2000 * 2000];
        b.iter(|| {
            x.encode_rgb(&mut im, 500, 2000, 0, 4);
        });
    }

    #[cfg(feature = "benchmarks")]
    #[bench]
    fn bench_encode_rgba(b: &mut Bencher) {
        let mut v= Vec::with_capacity(2000 * 2000);
        let mut x = BMPEncoder::new(&mut v);
        let mut im = vec![0; 2000 * 2000];
        b.iter(|| {
            x.encode_rgba(&mut im, 500, 2000, 0, 4);
        });
    }

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

    #[test]
    fn huge_files_return_error() {
        let mut encoded_data = Vec::new();
        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
        let mut encoder = BMPEncoder::new(&mut encoded_data);
        let result = encoder.encode(&image, 40_000, 40_000, ColorType::Rgb8);
        assert!(result.is_err());
    }

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

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

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

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